PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion-Pretrained Dense and Contextual Embeddings

Sedigheh Eslami, Maksim Gaiduk, Markus Krimmel, Louis Milliken, Bo Wang, Denis Bykov

arXiv:2602.11151v2cs.LGcs.CLcs.IR

TL;DR

Web-scale retrieval is difficult to assess with public benchmarks alone because production settings involve long-tail queries, noisy documents, and distribution shifts. The report introduces diffusion-pretrained multilingual embedding models with contextualized representations and shows strong performance across public and internal web-scale retrieval benchmarks, with native quantization for deployment efficiency.

  • Problem

    Public benchmarks do not fully capture web-scale retrieval challenges involving long-tail queries, noisy documents, and distribution shifts in production.

  • Method

    The models use diffusion-pretrained bidirectional language models and multi-stage contrastive learning to produce standard and document-contextualized multilingual embeddings.

  • Results

    Across public and internal web-scale benchmarks, the models achieve strong retrieval performance; on PPLXQ2Q Large, pplx-embed-v1-4B reaches 73.46% R@10, surpassing Qwen3-Embedding-4B by 5.56 percentage points.

  • Takeaways & Limitations

    Native quantization-aware training provides practical deployment benefits by directly producing INT8 and binary embeddings while maintaining strong retrieval performance.

Abstract

from arXiv · show

In this report, we introduce pplx-embed, a family of multilingual embedding models that employ multi-stage contrastive learning on a diffusion-pretrained language model backbone for web-scale retrieval. By leveraging bidirectional attention through diffusion-based pretraining, our models capture comprehensive bidirectional context within passages, enabling the use of mean pooling and a late chunking strategy to better preserve global context across long documents. We release two model types: pplx-embed-v1 for standard retrieval, and pplx-embed-context-v1 for contextualized embeddings that incorporate global document context into passage representations. pplx-embed-v1 achieves competitive performance on the MTEB(Multilingual, v2), MTEB(Code), MIRACL, BERGEN, and ToolRet retrieval benchmarks, while pplx-embed-context-v1 sets new records on the ConTEB benchmark. Beyond public benchmarks, pplx-embed-v1 demonstrates strong performance on our internal evaluation suite, focusing on real-world, large-scale search scenarios constructed from 1B production web pages. These results validate the models' effectiveness in production environments where retrieval quality and efficiency are critical at scale.

1. Introduction

The report introduces multilingual pplx-embed models for web-scale retrieval, combining diffusion-pretrained bidirectional representations with multi-stage contrastive learning. The family includes standard and document-contextualized embeddings, released at two parameter scales with efficient INT8 outputs.

  • Motivation: Dense embeddings map queries and documents into a shared semantic space, enabling efficient retrieval through approximate nearest neighbor search.The introduction frames dense textual embeddings as points whose distances capture meaningful semantic relationships.
  • Method: Diffusion pretraining converts a causally masked language model backbone into a bidirectional encoder for richer passage representations.The models then use multi-stage contrastive learning on question-document pairs and triplet data to align embedding geometry with semantic similarity.
  • Models: The released family contains pplx-embed-v1 for standard retrieval and pplx-embed-context-v1 for encoding passages with document-level context.Both model types are available at 0.6B- and 4B-parameter scales.
  • Results: 69.66% average nDCG@10: pplx-embed-v1-4B with INT8 quantization matches or exceeds Qwen3-Embedding-4B (69.60%) and gemini-embedding-001 (67.71%) on MTEB(Multilingual, v2).The family uses native quantization-aware training and outputs INT8 embeddings by default.

2. PPLX Embedding

PPLX Embedding trains retrieval embeddings through a branched, multi-stage curriculum combining diffusion pretraining, pair and contextual contrastive learning, hard-negative triplet refinement, and model merging. Its bidirectional backbone enables mean pooling, while quantization supports compact embeddings with minimal performance loss from post-hoc binarization.

  • Training framework: The framework combines four training paradigms in a branched curriculum followed by checkpoint merging and selection for pplx-embed-v1 and pplx-embed-context-v1.Pair training establishes sequence-level semantic alignment; contextual training produces chunk-level embeddings, triplet training refines similar-document boundaries, and spherical linear interpolation produces pplx-embed-v1.
  • Diffusion pretraining: Diffusion continued pretraining removes causal masking and enables bidirectional self-attention in decoder-only transformer backbones.The models reverse a corrupting noise process using an absorbing [MASK] state and continuous-time diffusion formulation.
  • Pooling and quantization: Mean pooling is enabled by bidirectional representations, and post-hoc binary quantization reduces embedding size with minimal performance loss.Quantized embeddings use signed 8-bit integer entries during inference and contrastive training, while binary embeddings map each mean-pooled entry to −1 or 1.
  • Pair training: Pair training uses InfoNCE to align queries with relevant documents while suppressing unrelated and likely false-negative in-batch samples.The curriculum introduces English, cross-lingual, and multilingual pair data in successive steps, and masks negatives whose similarity exceeds the positive pair by more than 0.1.
  • Contextual training: Contextual training jointly models local chunk semantics and global document similarities, gradually increasing the global-loss weight from β=0.2 toward 0.5.Duplicate-document masking and similarity-threshold masking mitigate false negatives while the scheduled objective preserves both chunk-level and document-level information.

3. Evaluations

The evaluations span public multilingual, code, contextual, RAG, and tool-retrieval benchmarks alongside web-scale internal tests over real-world queries and documents. pplx-embed-v1 performs competitively across these settings, with especially strong large-corpus retrieval and quantization results.

  • BERGEN: pplx-embed-v1-4B achieves the best BERGEN results in three of five tasks and outperforms Qwen3-Embedding-4B in four, while its 0.6B model wins three tasks.BERGEN indexes 24.8 million 100-word KILT Wikipedia passages and evaluates five retrieval-augmented question-answering tasks without reranking, using top-5 passages.
  • Tool Search: 44.45% average nDCG@10 ranks pplx-embed-v1-4B second overall on ToolRet, including 42.07% nDCG@10 on Web despite INT8 quantization.ToolRet covers 35 tasks across Web, Code, and Custom categories and reports nDCG@10, Precision@10, Recall@10, and Comprehensiveness@10.
  • Internal web-scale benchmarks: 73.46% R@10 and 86.17% R@100 let pplx-embed-v1-4B surpass Qwen3-Embedding-4B by 5.56 and 4.21 percentage points on the Large PPLXQ2Q corpus.The binary variant reaches 72.41% R@10 and 85.21% R@100, losing only 0.96–1.05 percentage points; the 0.6B model also exceeds BGE-M3 and Qwen3-Embedding-0.6B.
  • Internal web-scale benchmarks: 88.23% English and 91.66% Multilingual Recall@1000 for pplx-embed-v1-4B surpass Qwen3-Embedding-4B, supporting first-stage retrieval at K=1000.Binary variants retain 87.13% and 90.67%, respectively, with minimal degradation attributed to quantization-aware training.

4. Diffusion vs. Autoregressive Pretraining

The ablation study compares causally masked and bidirectional pretrained backbones with mean and last-token pooling, finding that diffusion pretraining improves retrieval performance while mean pooling enables contextual embedding training.

  • The ablation evaluates four configurations combining two base models—including causally masked Qwen3 and a bidirectional diffusion-pretrained backbone—with mean or last-token pooling.
  • The variants are compared on English MTEB(MTEB(En, v2)) retrieval tasks and the English subset of MIRACLRetrievalHardNegatives.
  • Diffusion pretraining improves performance across a range of retrieval tasks, increasing the average score by approximately 1 percentage point.
  • Mean pooling is crucial for contextual embedding training because it enables many chunk-level representations to be computed from a single document.

5. Related Work

Prior work spans diffusion language models, contrastive training, quantization-aware embedding methods, and contextual embeddings. This work studies diffusion-pretrained retrieval embeddings, quantization-aware training for web-scale retrieval, and multi-stage contextual embedding objectives within these lines of research.

  • Diffusion Language Models: Diffusion language models offer an alternative to autoregressive generation, while bidirectional attention is crucial for encoding global context in long and complex documents.Zhang et al. (2025a) systematically studied diffusion language models for text embeddings; this work examines continued diffusion-language-model pretraining for text retrieval.
  • Contrastive Training of Text Embeddings: InfoNCE-based contrastive learning predominates in text embedding training, with recent work emphasizing higher-quality synthesized pair data.The cited approaches align semantically similar texts and distinguish dissimilar ones.
  • Contrastive Training of Text Embeddings: Recent quantization research includes contrastive quantization, quantization-aware finetuning, and post-training quantization evaluations for retrieval-augmented generation.In contrast, this work targets quantization-aware training of embeddings for web-scale retrieval.
  • Contextual Embeddings: Contextual embedding methods use neighboring documents, single-pass chunk processing, or global-context evaluation through ConTEB and in-sequence training.This work builds on these foundations with multi-stage contrastive training and a specialized dual-loss objective while maintaining compatibility with MTEB and ConTEB.

6. Conclusion · Appendix

The report introduces pplx-embed, an embedding-model family built on diffusion-based language models with bidirectional attention to capture global document context. Its multi-stage pipeline supports semantic alignment, contextualized chunk encoding, and fine-grained relevance distinctions across four model variants.

  • 6. Conclusion: pplx-embed uses diffusion-based language models with bidirectional attention to better capture global document context.
  • 6. Conclusion: The multi-stage training pipeline progressively shapes text representations for semantic alignment.
  • 6. Conclusion: The pipeline trains contextualized chunk encodings relative to full documents.
  • 6. Conclusion: The pipeline also develops fine-grained relevance distinctions in text representations.
  • 6. Conclusion: The report provides four pplx-embed variants spanning standard and contextualized model families.The variants are pplx-embed-v1 and pplx-embed-context-v1.
  • 6. Conclusion: Each model family is offered at 0.6B and 4B parameter scales.

A. Details on Continued Pretraining

Continued pretraining uses an ELBO objective for diffusion to an absorbing [MASK] state, modeled by a bidirectional transformer. Training combines mixed-precision optimization, scheduled learning-rate decay, and multilingual data proportions based on FineWeb2.

  • Loss: Diffusion pretraining minimizes the standard ELBO for tokens independently decaying to an absorbing [MASK] state under a linear noise schedule.The reverse process is modeled with a bidirectional transformer.
  • Implementation: Because the transformer uses a left-shift operation without a [BOS] token, predictions and the loss sum exclude the first masked-sequence token.
  • Hyperparameters: Training uses bfloat16 mixed precision, FlashAttention-2, AdamW with weight decay 0.01, β1 = 0.9, β2 = 0.98, and gradient clipping.The clipping threshold is unspecified because it depends on implementation details such as loss scaling.
  • Hyperparameters: The learning rate warms up linearly for 6,000 steps and follows cosine decay over the final 12,000 steps.
  • Data: Non-English languages in the pretraining data follow their relative word-count prevalence in the FineWeb2 corpus.

B. Details on Contrastive Training

The models use staged contrastive training spanning pair, contextual, and triplet objectives, with dataset-aware batching, reduced-precision optimization, and hard-negative mining. Contextual training further combines global and local learning across multiple embedding dimensions.

  • Training Data Sampling: Each batch contains data from one dataset, whose selection probability is proportional to dataset size to improve in-batch negative quality.
  • Pair Training: Pair training uses InfoNCE with all in-batch negatives, temperature 0.02, and INT8 tanh quantization applied from training start.Both model sizes use global batch size 16,384 for 50,000 steps at sequence length 256; learning rates are 2×10−4 for 0.6B and 5×10−5 for 4B.
  • Contextual Training: Contextual training initializes from pair-trained checkpoints and uses contextual retrieval datasets with chunk-level annotations, including synthetic MLDR, MLDR, NarrativeQA, and SQuAD.Queries are synthesized for MLDR chunks lacking associated queries, and quantization is applied during both training and inference.
  • Contextual Training: Contextual training combines global document-level InfoNCE with a local loss through a hybrid Matryoshka objective across dimensions [128, 256, 512, 1024, 2048, 2560].The dual-objective weight β is cosine-annealed from 0.2 to 0.5, with masking strategies used to mitigate false negatives.
  • Triplet Training: Triplet training augments in-batch negatives with 3 mined hard negatives, using sequence length 512, global batch size 512, and temperature 0.03.Both models are fine-tuned for 2,000 steps, with learning rates of 5 × 10−5 for 0.6B and 10−5 for 4B.

C. Details on MTEB Evaluation

This section reports per-task nDCG@10 scores for pplx-embed-v1 on the MTEB(Multilingual, v2) and MTEB(Code) retrieval benchmarks, with evaluation settings and task-specific handling described.

  • MTEB Evaluation: pplx-embed-v1 uses sequence length 1024 on all tasks except LEMBPasskeyRetrieval, which uses sequence length 16,384.This exception applies specifically to LEMBPasskeyRetrieval.
  • MTEB Evaluation: For SyntheticText2SQL, small random noise is injected into embeddings because the corpus contains duplicate documents.The noise breaks symmetry between duplicate documents.
  • MTEB Evaluation: Per-task nDCG@10 scores are provided for MTEB(Multilingual, v2) retrieval tasks in Table 11.The table reports task-level results for the multilingual benchmark.
  • MTEB Evaluation: Per-task nDCG@10 scores are also provided for MTEB(Code) retrieval tasks in Table 12.The table reports task-level results for the code benchmark.

D. Details on ConTEB Evaluation

The ConTEB evaluation uses two contextual embedding strategies selected according to document length: ContextualEmbedder for standard-length documents and FixedContextualEmbedder for exceptionally long documents.

  • Evaluation process: ConTEB evaluation selects between two contextual embedding strategies based on document characteristics.The ContextualEmbedder incorporates surrounding chunk information during encoding, while the FixedContextualEmbedder uses fixed partitioning for exceptionally long documents.
  • Standard-length documents: For standard-length documents, ContextualEmbedder generates contextual embeddings using surrounding chunk information during encoding.
  • Exceptionally long documents: For exceptionally long documents, FixedContextualEmbedder applies fixed partitioning, including to ESG Reports documents exceeding 30,000 tokens.

E. Details on BERGEN Evaluation

The BERGEN evaluation uses official implementations of pplx-embed-v1, Qwen3-Embedding, and BGE-M3 in a custom retriever for fair comparison. End-to-end RAG retrieves 100 passages but supplies only the top 5 to the generator.

  • Embedding Models: The custom BERGEN retriever evaluates pplx-embed-v1, Qwen3-Embedding, and BGE-M3 using their official sentence-transformers implementations.Queries for Qwen3-Embedding use prompt_name="query".
  • RAG Configuration: RAG retrieves the top 100 passages from the KILT dump but presents only the top 5 passages to the generator.The evaluation uses the vllm_qwen-25-32b-instruct generator with max_length=32768.
  • RAG Configuration: The end-to-end RAG command specifies the dataset, retrieves 100 passages, and sets the generator's maximum length to 32768.The generator is vllm_qwen-25-32b-instruct.
Loading 2602.11151v2…