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은 long-tail query, noisy document, distribution shift가 존재하는 실제 환경에서 이루어지므로 public benchmark만으로는 평가하기 어렵다. 이 보고서는 contextualized representation을 갖춘 diffusion-pretrained multilingual embedding model을 소개하고, public 및 internal web-scale retrieval benchmark 전반에서 강력한 성능을 보이며 deployment efficiency를 위한 native quantization을 지원함을 보인다.

  • 문제

    Public benchmark는 실제 환경에서 long-tail query, noisy document, distribution shift가 수반되는 web-scale retrieval challenge를 완전히 포착하지 못한다.

  • 방법

    이 model은 diffusion-pretrained bidirectional language model과 multi-stage contrastive learning을 사용해 standard multilingual embedding과 document-contextualized multilingual embedding을 생성한다.

  • 결과

    Public 및 internal web-scale benchmark 전반에서 이 model은 강력한 retrieval 성능을 달성한다. PPLXQ2Q Large에서는 pplx-embed-v1-4B가 73.46% R@10에 도달해 Qwen3-Embedding-4B를 5.56 percentage point 앞선다.

  • 시사점 및 한계

    Native quantization-aware training은 강력한 retrieval 성능을 유지하면서 INT8 및 binary embedding을 직접 생성해 실용적인 deployment 이점을 제공한다.

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. 서론

이 보고서는 diffusion-pretrained bidirectional representation과 multi-stage contrastive learning을 결합해 web-scale retrieval을 수행하는 다국어 pplx-embed 모델을 소개한다. 이 제품군에는 표준 embedding과 document-contextualized embedding이 포함되며, 효율적인 INT8 output을 지원하는 두 가지 parameter scale로 공개된다.

  • 동기: Dense embedding은 query와 document를 공유된 semantic space로 매핑해 approximate nearest neighbor search를 통한 효율적인 retrieval을 가능하게 한다.서론에서는 dense textual embedding을 거리로 의미 있는 semantic relationship을 포착하는 점으로 설명한다.
  • 방법: Diffusion pretraining은 causally masked language model backbone을 bidirectional encoder로 변환해 더 풍부한 passage representation을 생성한다.이후 모델은 question-document pair와 triplet data에 multi-stage contrastive learning을 적용해 embedding geometry를 semantic similarity에 맞춘다.
  • 모델: 공개된 제품군에는 표준 retrieval을 위한 pplx-embed-v1과 document-level context를 포함해 passage를 encoding하는 pplx-embed-context-v1이 있다.두 모델 유형 모두 0.6B- 및 4B-parameter scale로 제공된다.
  • 결과: 69.66% 평균 nDCG@10: INT8 quantization을 적용한 pplx-embed-v1-4B는 MTEB(Multilingual, v2)에서 Qwen3-Embedding-4B (69.60%) 및 gemini-embedding-001 (67.71%)과 같거나 더 높은 성능을 보인다.이 제품군은 native quantization-aware training을 사용하며 기본적으로 INT8 embedding을 출력한다.

2. PPLX 임베딩

PPLX 임베딩은 diffusion pretraining, pair 및 contextual contrastive learning, hard-negative triplet refinement, model merging을 결합한 분기형 다단계 curriculum을 통해 retrieval embedding을 학습한다. 양방향 backbone은 mean pooling을 가능하게 하며, quantization은 사후 binarization에 따른 성능 손실을 최소화하면서 compact embedding을 지원한다.

  • Training framework: 이 framework는 분기형 curriculum에서 네 가지 training paradigm을 결합한 뒤 checkpoint merging과 selection을 수행하여 pplx-embed-v1과 pplx-embed-context-v1을 만든다.Pair training은 sequence-level semantic alignment를 확립하고, contextual training은 chunk-level embedding을 생성하며, triplet training은 유사 문서 경계를 정교화하고, spherical linear interpolation은 pplx-embed-v1을 생성한다.
  • Diffusion pretraining: Diffusion continued pretraining은 causal masking을 제거하고 decoder-only transformer backbone에서 bidirectional self-attention을 가능하게 한다.이 model들은 absorbing [MASK] state와 continuous-time diffusion formulation을 사용해 corruption noise process를 역전한다.
  • Pooling and quantization: Bidirectional representation을 통해 mean pooling이 가능해지며, post-hoc binary quantization은 성능 손실을 최소화하면서 embedding size를 줄인다.Quantized embedding은 inference와 contrastive training에서 signed 8-bit integer entry를 사용하며, binary embedding은 각 mean-pooled entry를 −1 또는 1로 매핑한다.
  • Pair training: Pair training은 InfoNCE를 사용해 query와 relevant document를 정렬하는 동시에, 무관하거나 false negative일 가능성이 높은 in-batch sample을 억제한다.이 curriculum은 successive step에서 English, cross-lingual, multilingual pair data를 도입하며, positive pair보다 similarity가 0.1을 초과해 높은 negative를 masking한다.
  • Contextual training: Contextual training은 local chunk semantic과 global document similarity를 jointly model하며, global-loss weight를 β=0.2에서 점차 0.5를 향해 증가시킨다.Duplicate-document masking과 similarity-threshold masking은 false negative를 완화하고, scheduled objective는 chunk-level 및 document-level 정보를 모두 보존한다.

3. 평가

평가는 공개 multilingual, code, contextual, RAG, tool-retrieval benchmark와 실제 질의·문서를 대상으로 한 web-scale 내부 테스트를 아우른다. pplx-embed-v1은 이러한 설정 전반에서 경쟁력 있는 성능을 보이며, 특히 대규모 corpus retrieval과 quantization에서 강점을 보인다.

  • BERGEN: pplx-embed-v1-4B는 5개 task 중 3개에서 BERGEN 최고 성능을 달성하고, 4개에서 Qwen3-Embedding-4B를 능가하며, 0.6B 모델은 3개 task에서 승리한다.BERGEN은 24.8 million개의 100-word KILT Wikipedia passage를 색인하고, reranking 없이 top-5 passage를 사용해 다섯 개의 retrieval-augmented question-answering 과제를 평가한다.
  • Tool Search: 평균 nDCG@10 44.45%로 pplx-embed-v1-4B는 ToolRet 전체에서 두 번째 순위를 차지하며, INT8 양자화를 적용했음에도 Web에서 nDCG@10 42.07%를 기록한다.ToolRet은 Web, Code, Custom 범주에 걸친 35개 과제를 다루며 nDCG@10, Precision@10, Recall@10, Comprehensiveness@10을 보고한다.
  • 내부 web-scale benchmark: 73.46% R@10과 86.17% R@100으로 pplx-embed-v1-4B는 Large PPLXQ2Q corpus에서 Qwen3-Embedding-4B를 각각 5.56 및 4.21 percentage point 차이로 능가한다.binary variant는 72.41% R@10과 85.21% R@100을 달성해 0.96–1.05 percentage point만 하락하며, 0.6B 모델도 BGE-M3와 Qwen3-Embedding-0.6B를 능가한다.
  • 웹 규모 내부 벤치마크: pplx-embed-v1-4B는 영어에서 88.23%, 다국어에서 91.66%의 Recall@1000을 기록해 Qwen3-Embedding-4B를 능가하며, K=1000에서 첫 단계 검색을 뒷받침한다.이진 변형은 각각 87.13%와 90.67%를 유지하며, 최소한의 성능 저하는 양자화 인지 학습에 기인한다.

4. Diffusion과 Autoregressive Pretraining 비교

Ablation study는 causally masked 및 bidirectional pretrained backbone을 mean pooling 및 last-token pooling과 조합해 비교하며, diffusion pretraining이 retrieval 성능을 향상시키고 mean pooling이 contextual embedding 학습을 가능하게 함을 보인다.

  • Ablation은 두 base model—causally masked Qwen3와 bidirectional diffusion-pretrained backbone—을 mean pooling 또는 last-token pooling과 조합한 네 가지 구성을 평가한다.
  • 각 variant는 English MTEB(MTEB(En, v2)) retrieval task와 MIRACLRetrievalHardNegatives의 English subset에서 비교된다.
  • Diffusion pretraining은 다양한 retrieval task에서 성능을 향상시키며, 평균 점수를 약 1 percentage point 높인다.
  • Mean pooling은 하나의 document에서 많은 chunk-level representation을 계산할 수 있게 하므로 contextual embedding training에 핵심적이다.

5. 관련 연구

기존 연구는 diffusion language models, contrastive training, quantization-aware embedding methods, contextual embeddings를 아우른다. 본 연구는 이러한 연구 흐름에서 diffusion-pretrained retrieval embeddings, web-scale retrieval을 위한 quantization-aware training, multi-stage contextual embedding objectives를 다룬다.

  • Diffusion Language Models: Diffusion language models는 autoregressive generation의 대안을 제공하며, bidirectional attention은 길고 복잡한 문서의 전역 문맥을 인코딩하는 데 핵심적이다.Zhang et al. (2025a)은 텍스트 임베딩을 위한 diffusion language models를 체계적으로 연구했으며, 본 연구는 텍스트 검색을 위한 continued diffusion-language-model pretraining을 검토한다.
  • 텍스트 임베딩의 Contrastive Training: InfoNCE-based contrastive learning은 텍스트 임베딩 학습에서 주류를 이루며, 최근 연구는 더 높은 품질의 합성 pair data를 강조한다.인용된 접근법은 의미적으로 유사한 텍스트를 정렬하고 서로 유사하지 않은 텍스트를 구분한다.
  • 텍스트 임베딩의 Contrastive Training: 최근 quantization 연구에는 contrastive quantization, quantization-aware finetuning, retrieval-augmented generation을 위한 post-training quantization 평가가 포함된다.이와 달리 본 연구는 web-scale retrieval을 위한 임베딩의 quantization-aware training을 목표로 한다.
  • Contextual Embeddings: Contextual embedding 방법은 인접 문서, single-pass chunk processing, 또는 ConTEB와 in-sequence training을 통한 전역 문맥 평가를 활용한다.본 연구는 MTEB와 ConTEB의 호환성을 유지하면서 multi-stage contrastive training과 특화된 dual-loss objective를 통해 이러한 기반 위에 구축된다.

6. 결론 · 부록

이 보고서는 전역 문서 문맥을 포착하기 위해 bidirectional attention을 사용하는 diffusion 기반 language model에 구축된 embedding-model family인 pplx-embed를 소개한다. 다단계 파이프라인은 semantic alignment, contextualized chunk encoding, 그리고 네 가지 model variant에 걸친 세밀한 relevance 구분을 지원한다.

  • 6. 결론: pplx-embed는 전역 문서 문맥을 더 잘 포착하기 위해 bidirectional attention을 사용하는 diffusion 기반 language model을 활용한다.
  • 6. 결론: 다단계 training pipeline은 semantic alignment를 위해 text representation을 점진적으로 형성한다.
  • 6. 결론: 이 파이프라인은 full document를 기준으로 contextualized chunk encoding을 학습한다.
  • 6. 결론: 또한 이 파이프라인은 text representation에서 세밀한 relevance 구분을 발전시킨다.
  • 6. 결론: 이 보고서는 standard 및 contextualized model family를 아우르는 네 가지 pplx-embed variant를 제공한다.variant는 pplx-embed-v1과 pplx-embed-context-v1이다.
  • 6. 결론: 각 model family는 0.6B 및 4B parameter scale로 제공된다.

A. 지속 사전학습 세부 사항

지속 사전학습은 bidirectional transformer로 모델링된 absorbing [MASK] 상태로의 diffusion에 ELBO 목적함수를 사용한다. 학습에는 mixed-precision 최적화, 스케줄된 learning-rate decay, FineWeb2에 기반한 다국어 데이터 비율을 결합한다.

  • 손실: Diffusion 사전학습은 linear noise schedule에 따라 토큰이 독립적으로 absorbing [MASK] 상태로 감쇠할 때의 표준 ELBO를 최소화한다.Reverse process는 bidirectional transformer로 모델링한다.
  • 구현: Transformer는 [BOS] 토큰 없이 left-shift 연산을 사용하므로, prediction과 loss sum에서 첫 번째 masked-sequence token을 제외한다.
  • 하이퍼파라미터: 학습에는 bfloat16 mixed precision, FlashAttention-2, weight decay 0.01을 사용하는 AdamW, β1 = 0.9, β2 = 0.98, gradient clipping을 사용한다.Clipping threshold는 loss scaling과 같은 구현 세부 사항에 따라 달라지므로 명시하지 않는다.
  • 하이퍼파라미터: Learning rate는 6,000 steps 동안 선형으로 warm up하고, 마지막 12,000 steps에 걸쳐 cosine decay를 따른다.
  • 데이터: 사전학습 데이터의 비영어권 언어는 FineWeb2 corpus에서의 상대적 word-count prevalence를 따른다.

B. Contrastive Training 세부사항

모델은 pair, contextual, triplet objective에 걸친 단계적 contrastive training을 사용하며, dataset-aware batching, reduced-precision optimization, hard-negative mining을 적용한다. Contextual training에서는 여러 embedding dimension에 걸쳐 global 및 local learning을 추가로 결합한다.

  • Training Data Sampling: 각 batch에는 하나의 dataset에서 가져온 데이터만 포함되며, dataset 선택 확률은 in-batch negative 품질을 높이기 위해 dataset size에 비례한다.
  • Pair Training: Pair training은 모든 in-batch negative를 사용하는 InfoNCE와 temperature 0.02를 적용하며, training 시작 시점부터 INT8 tanh quantization을 적용한다.두 model size 모두 sequence length 256에서 global batch size 16,384로 50,000 steps를 학습하며, learning rate는 0.6B에서 2×10−4, 4B에서 5×10−5다.
  • Contextual Training: Contextual training은 pair-trained checkpoint에서 초기화하며, synthetic MLDR, MLDR, NarrativeQA, SQuAD를 포함한 chunk-level annotation이 있는 contextual retrieval dataset을 사용한다.연관된 query가 없는 MLDR chunk에 대해서는 query를 합성하고, training과 inference 모두에서 quantization을 적용한다.
  • Contextual Training: Contextual training은 dimension [128, 256, 512, 1024, 2048, 2560]에 걸친 hybrid Matryoshka objective를 통해 global document-level InfoNCE와 local loss를 결합한다.dual-objective weight β는 0.2에서 0.5로 cosine-annealing하며, false negative를 완화하기 위해 masking strategy를 사용한다.
  • Triplet Training: Triplet training은 3개의 mined hard negative로 in-batch negative를 보강하며, sequence length 512, global batch size 512, temperature 0.03을 사용한다.두 model 모두 2,000 steps 동안 fine-tuning하며, learning rate는 0.6B에서 5 × 10−5, 4B에서 10−5다.

C. MTEB 평가 세부 사항

이 절에서는 MTEB(Multilingual, v2) 및 MTEB(Code) retrieval benchmark에서 pplx-embed-v1의 task별 nDCG@10 점수를 보고하며, 평가 설정과 task별 처리를 설명한다.

  • MTEB 평가: pplx-embed-v1은 LEMBPasskeyRetrieval을 제외한 모든 task에서 sequence length 1024를 사용하며, LEMBPasskeyRetrieval에서는 sequence length 16,384를 사용한다.이 예외는 LEMBPasskeyRetrieval에만 적용된다.
  • MTEB 평가: SyntheticText2SQL에서는 corpus에 duplicate document가 포함되어 있으므로 embedding에 small random noise를 주입한다.이 noise는 duplicate document 간 대칭성을 깨뜨린다.
  • MTEB 평가: MTEB(Multilingual, v2) retrieval task의 task별 nDCG@10 점수는 Table 11에 제시한다.해당 표는 multilingual benchmark의 task 수준 결과를 보고한다.
  • MTEB 평가: MTEB(Code) retrieval task의 task별 nDCG@10 점수도 Table 12에 제시한다.해당 표는 code benchmark의 task 수준 결과를 보고한다.

D. ConTEB 평가 세부사항

ConTEB 평가는 문서 길이에 따라 두 가지 contextual embedding 전략을 선택해 사용한다. 일반적인 길이의 문서에는 ContextualEmbedder를, 예외적으로 긴 문서에는 FixedContextualEmbedder를 사용한다.

  • 평가 과정: ConTEB 평가는 문서 특성에 따라 두 가지 contextual embedding 전략 중 하나를 선택한다.ContextualEmbedder는 인코딩 중 주변 chunk 정보를 통합하는 반면, FixedContextualEmbedder는 예외적으로 긴 문서에 고정 분할을 적용한다.
  • 일반적인 길이의 문서: 일반적인 길이의 문서에서 ContextualEmbedder는 인코딩 중 주변 chunk 정보를 사용해 contextual embedding을 생성한다.
  • 예외적으로 긴 문서: 예외적으로 긴 문서에서 FixedContextualEmbedder는 30,000 tokens를 초과하는 ESG Reports 문서를 포함해 고정 분할을 적용한다.

E. BERGEN 평가 세부 사항

BERGEN 평가는 공정한 비교를 위해 pplx-embed-v1, Qwen3-Embedding, BGE-M3의 공식 구현을 custom retriever에서 사용한다. End-to-end RAG는 100개 passage를 검색하지만 generator에는 상위 5개만 제공한다.

  • Embedding Models: custom BERGEN retriever는 공식 sentence-transformers 구현을 사용해 pplx-embed-v1, Qwen3-Embedding, BGE-M3를 평가한다.Qwen3-Embedding에 대한 query에는 prompt_name="query"를 사용한다.
  • RAG Configuration: RAG는 KILT dump에서 상위 100개 passage를 검색하지만 generator에는 상위 5개 passage만 제공한다.평가에는 max_length=32768인 vllm_qwen-25-32b-instruct generator를 사용한다.
  • RAG Configuration: End-to-end RAG command는 dataset을 지정하고 100개 passage를 검색하며 generator의 maximum length를 32768로 설정한다.generator는 vllm_qwen-25-32b-instruct다.
Loading 2602.11151v2…