Diffusion-Pretrained Dense and Contextual Embeddings for Web-Scale Retrieval

Paper: 2602.11151 Authors: Sedigheh Eslami, Maksim Gaiduk, Markus Krimmel, Louis Milliken, Bo Wang, Denis Bykov Categories: cs.LG, cs.CL, cs.IR

Abstract

This paper presents pplx-embed, a novel family of multilingual embedding models designed for web-scale retrieval tasks. The key innovation lies in leveraging diffusion-based pretraining to create language model backbones with full bidirectional attention, which are then fine-tuned using multi-stage contrastive learning. Unlike traditional encoder-only models, the diffusion pretraining approach enables comprehensive bidirectional context capture within passages, allowing for effective mean pooling and a late chunking strategy that preserves global document context across long texts. The authors release two variants: pplx-embed-v1 for standard dense retrieval and pplx-embed-context-v1 for contextualized embeddings that incorporate document-level context into passage representations. Extensive evaluation demonstrates competitive performance on major benchmarks including MTEB, MIRACL, and BERGEN, with pplx-embed-context-v1 achieving state-of-the-art results on the ConTEB contextualized embedding benchmark.

Key Contributions

• Introduction of diffusion-pretrained language models as embedding backbones, enabling true bidirectional attention for improved context capture
• Development of a late chunking strategy that preserves global document context when processing long documents
• Release of two model variants: pplx-embed-v1 for standard retrieval and pplx-embed-context-v1 for contextualized embeddings
• State-of-the-art performance on ConTEB benchmark for contextualized embeddings
• Validation on real-world, large-scale search scenarios with tens of millions of documents

Technical Methodology

The pplx-embed architecture builds upon diffusion-pretrained language models, which differ fundamentally from traditional masked language models or causal decoders. Diffusion pretraining involves training models to denoise corrupted text, naturally encouraging bidirectional attention patterns that capture both forward and backward context simultaneously. This bidirectional capability is crucial for embedding quality, as it allows the model to understand each token in the full context of its surrounding passage.

The training process employs multi-stage contrastive learning, where the model learns to distinguish relevant passages from irrelevant ones through carefully constructed positive and negative pairs. The contrastive objective can be formalized as maximizing the similarity between query $q$ and relevant passage $p^+$ while minimizing similarity with negative passages $p^-$:

$\mathcal{L} = -\log \frac{\exp(\text{sim}(q, p^+)/\tau)}{\exp(\text{sim}(q, p^+)/\tau) + \sum_{i} \exp(\text{sim}(q, p^-_i)/\tau)}$

where $\tau$ is a temperature parameter and $\text{sim}(\cdot, \cdot)$ represents cosine similarity.

For long document processing, the late chunking strategy splits documents after embedding computation rather than before, allowing each chunk to benefit from the full document context during encoding. This approach contrasts with traditional early chunking methods that process each segment independently, losing valuable cross-segment information.

Experimental Results and Benchmarks

pplx-embed-v1 demonstrates competitive performance across multiple established benchmarks. On MTEB (Multilingual, v2), the model achieves strong results in retrieval tasks spanning diverse languages and domains. The MTEB (Code) benchmark evaluates code search capabilities, where pplx-embed-v1 shows robust performance in matching natural language queries to code snippets. On MIRACL, a multilingual retrieval benchmark covering 18 languages, the model maintains consistent quality across linguistic boundaries.

The contextualized variant, pplx-embed-context-v1, specifically targets scenarios where passage understanding requires document-level context. On the ConTEB benchmark, which evaluates contextualized embeddings, this model achieves state-of-the-art results, demonstrating the effectiveness of incorporating global document information into passage representations. This is particularly valuable for tasks like question answering over long documents, where understanding a passage often requires knowledge of surrounding content.

Beyond public benchmarks, the authors report strong performance on internal evaluation suites designed to simulate real-world search scenarios at scale. These evaluations involve retrieval over tens of millions of documents, testing both quality and efficiency in production-like environments. The results validate that pplx-embed models can maintain high retrieval quality while meeting the latency and throughput requirements of large-scale deployment.

Implications for Retrieval Systems

The introduction of diffusion-pretrained embeddings represents a significant architectural shift in dense retrieval. Traditional embedding models typically rely on BERT-style masked language model pretraining or GPT-style causal pretraining, both of which have limitations for bidirectional context modeling. Diffusion pretraining offers a compelling alternative that naturally encourages comprehensive bidirectional understanding.

The late chunking strategy addresses a long-standing challenge in document retrieval: how to effectively represent long documents while preserving coherence across segments. By maintaining document-level context during embedding computation, this approach enables more semantically meaningful chunk representations. This has practical implications for applications like legal document search, scientific literature retrieval, and enterprise knowledge bases where documents often exceed typical context window limits.

The dual-model release strategy—offering both standard and contextualized variants—acknowledges that different retrieval scenarios have different requirements. Standard dense retrieval benefits from passage-level embeddings that can be precomputed and indexed efficiently. Contextualized embeddings, while potentially more computationally expensive, provide superior quality when document context is crucial for understanding.

Takeaways

1. Diffusion-based pretraining provides a novel and effective approach for creating embedding model backbones with true bidirectional attention capabilities
2. Late chunking strategies that preserve global document context during embedding computation significantly improve long document retrieval quality
3. pplx-embed-context-v1 achieves state-of-the-art results on contextualized embedding benchmarks, demonstrating the value of document-level context integration
4. The models demonstrate strong performance not only on public benchmarks but also in real-world, large-scale search scenarios with tens of millions of documents
5. The release of both standard and contextualized variants provides flexibility for different retrieval use cases and computational constraints