Diffusion-Pretrained Dense and Contextual Embeddings

TL;DR

The paper tackles efficient multilingual web-scale retrieval by introducing pplx-embed, a family of diffusion-pretrained embeddings trained with multi-stage contrastive learning to capture full-document context. It presents two model variants, v1 for standard retrieval and context-v1 for document-contextualized representations, and demonstrates strong retrieval performance across diverse public benchmarks and internal web-scale datasets while enabling INT8 and binary quantization for deployment. The approach combines continued diffusion pretraining, mean pooling, and several contrastive objectives (pair, contextual, and triplet) over multilingual data, with quantization-aware training to preserve quality at reduced precision. Collectively, the work shows that diffusion-based bidirectional encoders, when coupled with careful pooling and loss design, yield competitive or superior retrieval quality at scale, offering practical benefits for production search systems.

Abstract

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, which focuses on real-world, large-scale search scenarios over tens of millions of documents. These results validate the models' effectiveness in production environments where retrieval quality and efficiency are critical at scale.

Figures (2)

• Figure 1: Training pipeline of pplx-embed-v1 and pplx-embed-context-v1.
• Figure 2: Smoothed training loss using exponential moving average with $\alpha=0.02$.