PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

Shuai Wang, Yu Yin, Shengyao Zhuang, Bevan Koopman, Guido Zuccon

arXiv:2605.07210v2cs.IRcs.CL

TL;DR

Existing DLM retrievers use single mean-pooled vectors instead of the masked-position interface, limiting retrieval representations. DiffRetriever directly uses one or more jointly predicted masked positions, achieving the strongest aggregate matched-comparison effectiveness while retaining parallel encoding and exposing adaptive-allocation headroom.

  • Problem

    Existing DLM retrievers use BERT-style single mean-pooled vectors rather than the masked-position prediction interface that DLMs are trained to use.

  • Method

    DiffRetriever appends one or more masked positions to each prompted query or passage and uses their jointly produced outputs as retrieval representations.

  • Results

    DiffRetriever achieves the strongest aggregate effectiveness within the matched comparison across zero-shot and fine-tuned evaluation, outperforming DiffEmbed, PromptReps, and RepLLaMA.

  • Takeaways & Limitations

    Parallel masked-position prediction supports multi-representation retrieval with little additional encoding latency compared with sequential autoregressive generation.

  • Takeaways & Limitations

    Fixed query and passage masked-position counts transfer well in aggregate but leave substantial per-query headroom for adaptive selection.

Abstract

from arXiv · show

This paper shows how diffusion language models (DLMs) can be used as effective and efficient retrievers. Existing DLM-based retrievers (e.g., DiffEmbed) follow BERT-style encoding, representing each query or passage as a single mean-pooled vector. This ignores how DLMs are trained to generate responses through masked-position prediction under bidirectional attention, a capability that can provide stronger retrieval signals. We propose DiffRetriever, which uses the DLM's native masked-position prediction directly for retrieval. For each query or passage, DiffRetriever appends one or more masked positions, using the outputs as retrieval representations in a single forward pass. With one masked position, single-representation DiffRetriever already improves over DiffEmbed on the same backbones. DiffRetriever also naturally extends to multi-representation retrieval: DLMs process multiple masked positions jointly, enabling ColBERT-style fine-grained matching with little additional encoding latency. In autoregressive LLM retrievers, the same multi-representation strategy requires sequential decoding and therefore incurs much higher latency. DiffRetriever obtains the strongest aggregate effectiveness within our matched comparison, outperforming DiffEmbed, PromptReps, and RepLLaMA. Masked-position counts selected on training data transfer well across datasets, while per-query variation suggests headroom for adaptive allocation. Code is available at https://github.com/ielab/diffretriever.

1 Introduction

DiffRetriever uses diffusion models’ native masked-position prediction to replace single-vector pooling and support efficient multi-representation retrieval. This interface improves retrieval effectiveness while avoiding the sequential generation cost of autoregressive alternatives.

  • DLMs jointly predict preallocated masked positions under bidirectional attention, unlike autoregressive models that generate tokens strictly left to right.
  • Existing DLM retrievers such as DiffEmbed mean-pool each query or passage, failing to exploit masked-position prediction for retrieval.
  • DiffRetriever appends one or more masked positions and uses their outputs as retrieval representations, increasing representation count without sequential generation.
  • Multi-representation retrieval enables fine-grained matching, but autoregressive implementations incur sequential generation cost and inconsistent effectiveness gains.
  • Across zero-shot and fine-tuned evaluations, DiffRetriever improves over DiffEmbed, gains further from multiple representations, and achieves the strongest aggregate effectiveness in the matched comparison.

2 Related Work

Related work includes DLM retrievers that use encoder-style pooling, autoregressive embedding and multi-representation methods, and the DiffRetriever design that instead exploits parallel masked positions.

  • DiffEmbed and pplx-embed treat diffusion models as BERT-style encoders that produce one mean-pooled representation.
  • DiffuRank uses diffusion likelihood to rerank candidates from a first-stage retriever rather than directly producing retrieval representations.
  • DiffRetriever differs by producing multiple dense or sparse retrieval signals from jointly processed masked positions in one bidirectional pass.
  • Autoregressive retrievers such as RepLLaMA, E5-Mistral, and GTE-Qwen use contrastive fine-tuning to produce single-vector representations.

3 Method

DiffRetriever appends masked positions to prompted queries and passages, extracts dense and sparse signals from their outputs, and combines them for retrieval and fine-tuning. Its parallel bidirectional processing supports multiple representations without sequential decoding.

  • Parallel Masked-Position Decoding: The method appends K masked positions to a retrieval prompt, with K=1 for single-representation and K>1 for multi-representation retrieval.
  • Parallel Masked-Position Decoding: DiffRetriever obtains all masked-position representations in one bidirectional pass, unlike PromptReps, which generates representations one token at a time.
  • Retrieval Prompts: The prompt uses a system role and a user request for a textual retrieval representation, with “one word” or “a few words” controlling representation count.
  • Parallel Masked-Position Decoding: Closing quote and termination tokens make the masked span a complete, length-specified assistant response compatible with DLM chat formatting.
  • Scoring: Each masked position yields a hidden state and logit vector, allowing dense, sparse, and hybrid retrieval signals from jointly produced representations.
  • Decoding: Main experiments use one denoising step, and iterative denoising with more steps is significantly worse zero-shot across tested DLM backbones.
  • Scoring: Dense scoring matches each query representation to its closest passage representation and averages the similarities, while sparse scoring applies transformed-logit max-pooling.
  • Scoring: Dense and sparse scores are combined by equal-weight linear interpolation after min-max normalization within each retriever’s top-1000 list.

4 Experimental Setup

The experiments compare similarly sized autoregressive and diffusion backbones across in-domain and out-of-domain retrieval benchmarks. Masked-position budgets are selected on MS MARCO training data and reused across evaluations.

  • Models: The study compares LLaMA3 and Qwen2.5 autoregressive models with Dream and LLaDA diffusion models at roughly 7–8B parameters.
  • Models: Dream and Qwen2.5 share architecture and initialization, while LLaDA is paired with LLaMA3 as a complementary size-matched comparison.
  • Baselines: Baselines include BM25, PromptReps, DiffEmbed, and RepLLaMA, covering sparse retrieval, autoregressive retrieval, encoder-style DLM retrieval, and single-vector retrieval.
  • Datasets and Metrics: Evaluation covers MS MARCO with MRR@10, TREC DL 2019/2020 with NDCG@10, and BEIR-7 with NDCG@10 for out-of-domain transfer.
  • Latency and Storage: Latency comparisons measure query encoding plus search on a 100K-document MS MARCO sample, while multi-representation indexing increases passage storage roughly with Kp.
  • Selecting Kq and Kp: Budget selection on MS MARCO training data chooses (Kq, Kp)=(4,16) for Dream and (4,4) for LLaDA, then reuses each pair across evaluations.
  • Fine-tuning: Fine-tuning uses augmented MS MARCO triples with one positive passage and 15 sampled hard negatives, keeping selected masked-position counts fixed during training and inference.

5 Results

Across zero-shot, fine-tuned, and BEIR-7 transfer evaluations, multi-representation DiffRetriever delivers strong effectiveness with single-forward-pass encoding, while results depend on backbone, scoring mode, and supervision. Its latency advantage over autoregressive multi-representation retrieval remains substantial.

  • DiffRetriever’s evaluation spans in-domain zero-shot, in-domain fine-tuned, and out-of-domain BEIR-7 transfer settings.
  • Zero-shot retrieval: In zero-shot retrieval, PromptReps leads in single representation, but DiffRetriever’s multi-representation setting improves across datasets and scoring modes.For DiffRetriever, 16 of 18 multi-representation cells significantly exceed the corresponding single-representation cells, whereas PromptReps shows no consistent gains.
  • Zero-shot retrieval: DiffRetriever’s multi-representation gains arise from both multiple masked-position representations and MaxSim late-interaction scoring.Mean-pooling the same multiple masked-position representations already improves over a single representation, with MaxSim adding further gains.
  • Efficiency: PromptReps requires 275–300 ms per query for multiple representations, compared with DiffRetriever’s 16–20 ms, a roughly 15× latency gap.The comparison uses zero-shot evaluation with PromptReps capped at N = 20.
  • Fine-tuned retrieval: After fine-tuning, multi-representation Dream-based DiffRetriever is strongest within the matched comparison, winning four of nine columns despite smaller margins.It exceeds multi-representation PromptReps, same-backbone DiffEmbed, and RepLLaMA on MS MARCO dense retrieval.
  • Out-of-domain transfer: On BEIR-7, multi-representation DiffRetriever achieves the highest average in both zero-shot and fine-tuned comparisons, with fine-tuned Dream reaching .671.Zero-shot LLaDA-based DiffRetriever reaches .539, while fine-tuned Dream-based DiffRetriever reaches .671; Qwen2.5-family systems transfer better after fine-tuning.

6 Analysis

The analysis separates the gains from multiple masked-position representations and late interaction, then examines budget transfer, oracle headroom, and predictors of query-specific budgets. Fixed train-selected budgets transfer well in aggregate, but per-query adaptation remains promising and practically constrained.

  • Where do the gains come from?: Mean-pooling multiple masked-position representations improves DiffRetriever on both diffusion backbones before late interaction is applied.ColBERT-style MaxSim adds further gains, so late interaction contributes but does not solely explain the improvement.
  • Where do the gains come from?: PromptReps consistently declines when multiple generated representations are mean-pooled, unlike DiffRetriever's improvement from multiple masked positions.This aligns with the higher sequential decoding cost of autoregressive multi-representation retrieval.
  • The impact of number of masked tokens: Dream's train-selected budget is (Kq, Kp) = (4, 16), while LLaDA's is (4, 4); both are at or near the best test-set cells on MS MARCO dev and BEIR-7.The per-backbone preference persists across distributions.
  • The impact of number of masked tokens: On MS MARCO dev, the grid spans Δ=.14 for Dream versus Δ=.07 for LLaDA, so the cost of a poor fixed budget is backbone-dependent.Aggregate robustness does not imply one fixed budget is optimal for every query.
  • Per-query oracle headroom and predictability: The joint per-query oracle exceeds fixed-budget fine-tuning by .05 MRR@10 on MS MARCO dev for both backbones and by .03 or .07 NDCG@10 on BEIR-7.The oracle selects budgets using retrieval labels, so these gains are upper bounds rather than deployable-system results.
  • Per-query oracle headroom and predictability: Query length and entropy correlate positively with oracle-preferred query budgets, with length stronger and dense-scoring correlations cleaner than sparse-scoring correlations.Length reaches ρ=+0.31 for Dream and ρ=+0.29 for LLaDA under dense scoring, but the correlations are modest.

7 Conclusion

DiffRetriever reformulates DLM retrieval around native masked-position prediction rather than encoder-style pooling. It supports jointly produced multi-representations, achieves the strongest aggregate effectiveness in the matched comparison, and motivates adaptive allocation.

  • Conclusion: DiffRetriever uses DLM masked positions as retrieval representations, producing multiple representations jointly in one forward pass.This avoids the sequential generation cost associated with autoregressive multi-representation retrieval.
  • Conclusion: DiffRetriever achieves the strongest aggregate effectiveness within the matched comparison in both zero-shot and fine-tuned settings.The conclusion also reports that different queries benefit from different numbers of masked positions.

Limitations

The evaluation is limited to 7–8B-scale models and a specified set of retrieval benchmarks, leaving other scales, languages, and long-document settings for future study. Fixed representation budgets transfer well overall, but adaptive allocation and storage remain open systems concerns.

  • Model and evaluation scope: The experiments compare two diffusion and two autoregressive backbones at 7–8B scale, so whether the interface advantage holds at other scales remains open.Dream is initialized from Qwen2.5, providing the closest available pairing.
  • Model and evaluation scope: The evaluation covers MS MARCO, TREC DL 2019/2020, and seven BEIR datasets, while multilingual and long-document retrieval are left for future work.The BEIR tasks span open-domain, multi-hop, scientific, biomedical, financial, argument, and duplicate-question retrieval.
  • Budget selection: Fixed query and passage masked-position counts selected on MS MARCO training data transfer well in aggregate, but per-query oracle analysis indicates headroom for adaptive selection.Correlations with query length and entropy suggest a lightweight learned predictor as a possible next step.
  • Systems tradeoffs: Multi-representation retrieval increases passage-index storage roughly in proportion to Kp, trading lower sequential query-encoding cost for higher passage-side storage.Compression-aware training and learned per-passage budget allocation are identified as future directions.

A.1 Multi-Step Denoising

The appendix evaluates iterative denoising against the single-pass design and finds that zero-shot retrieval strongly favors S=1, while fine-tuned effects depend on the diffusion backbone. A prompt-phrasing ablation further shows backbone-dependent differences for autoregressive multi-representation retrieval.

  • Multi-Step Denoising: S=1 is the single-pass variant used in the main text, whereas S=2 iteratively denoises masked positions over two forward passes.The multi-step procedure unmasks positions over confidence-based rounds and stores hidden states and logits in a frozen buffer.
  • Multi-Step Denoising: At zero-shot, S=2 is significantly worse than S=1 on every benchmark and BEIR-7 dataset for both backbones.Dream MS MARCO hybrid drops from .218 to .172, while Dream and LLaDA BEIR-7 dense averages drop from .427 to .309 and .497 to .385.
  • Multi-Step Denoising: After fine-tuning, S=2 has no consistent benefit for Dream, while LLaDA’s BEIR-7 dense average rises from .645 to .662 but its MS MARCO dense score drops from .427 to .414.The backbone-dependent result supports retaining S=1 for consistent single-pass inference.
  • Prompt phrasing: On LLaMA3-8B, the “a few words” prompt wins 8 of 9 benchmark–scoring cells, including dense MS MARCO MRR@10 increasing from .094 to .151.On Qwen2.5-7B, “three words” wins 5 of 9 cells, mostly for sparse and hybrid scoring.
  • Implementation details: The fine-tuning recipe uses the same optimization, batching, and LoRA settings across all four backbones and retrained baselines.The trainable-parameter fraction remains within 0.52–0.53%.

B.2 Latency Scaling

DiffRetriever’s multi-representation encoding remains close to its single-representation encoding, unlike autoregressive multi-representation retrieval, but search cost and index storage grow with representation budgets. Compression can substantially reduce the storage burden.

  • Latency scaling: DiffRetriever multi-representation encoding remains near 15–30 ms, while autoregressive multi-representation retrieval takes 60–80 ms across the measured input lengths.All systems scale roughly linearly with input length; autoregressive single-representation retrieval sits at 15–30 ms.
  • Latency scaling: Multi-representation search costs more because each query performs Kq lookups and each document stores Kp vectors.At 1M documents, Dream’s Kp=16 configuration is slower than LLaDA’s Kp=4 configuration because it stores four times as many passage vectors.
  • Latency scaling: At the fine-tuned N=4 cap, autoregressive multi-representation encoding remains 2–3× slower than autoregressive single-representation encoding, while DiffRetriever’s multi-representation encoding stays close to its single-representation cost.Within the measured range up to 1M documents, encoding remains the dominant contributor to end-to-end latency.
  • Index storage: Index storage grows approximately linearly with corpus size and passage budget Kp, making single-representation systems the smallest-footprint option.Multi-representation indexing retains the passage-side cost of storing multiple dense vectors.
  • Index storage: At MS MARCO scale, LLaDA multi-representation storage is about 213 GiB versus 54 GiB single-representation, while Dream requires about 742 GiB versus 47 GiB.The different totals reflect the train-selected passage budgets: Kp=4 for LLaDA and Kp=16 for Dream.
  • Index storage: Compression shrinks storage by roughly 6× with limited effectiveness loss across backbones, datasets, and budgets.Examples include Dream storage falling from 742 GiB to 118 GiB and LLaDA storage from 213 GiB to 34 GiB.

C Supplementary Results

Supplementary results clarify scoring-mode choices, provide broader baseline context, and examine representation-budget selection. They show a zero-shot hybrid advantage, a fine-tuned dense advantage, and strong matched-comparison performance alongside weaker out-of-domain results than broadly trained embedding models.

  • Scoring modes: In zero-shot BEIR-7 evaluation, hybrid scoring exceeds dense scoring for every system by 0.020 to 0.120 NDCG@10.In fine-tuned evaluation, dense exceeds hybrid for every system by 0.019 to 0.028 NDCG@10.
  • Scoring modes: The scoring-mode results support reporting hybrid scores for zero-shot transfer and dense scores for fine-tuned transfer.The per-dataset breakdown also identifies datasets favoring sparse-heavy behavior.
  • Additional baseline comparison: On MS MARCO dev, Dream multi-representation DiffRetriever achieves .433, exceeding NV-Embed-v2 (.390), Qwen3-Embedding-8B (.369), and Qwen3-Embedding-4B (.357).These dedicated embedding models are reported as landscape context rather than controlled matched baselines.
  • Additional baseline comparison: On BEIR-7, DiffRetriever substantially outperforms BERT-era dense retrievers, including ColBERT-v2 (.619 versus Dream multi .671), but broadly trained dedicated embedding models achieve higher averages.The broader models use substantially more extensive retrieval-oriented training data than DiffRetriever’s MS MARCO-only fine-tuning.
  • Budget selection: The supplementary analyses identify adaptive representation-budget selection as a promising direction beyond fixed counts selected on MS MARCO training data.Per-query budget variation is analyzed as potential headroom for adaptive allocation.

D.1 Per-Dataset (Kq, Kp) Landscape

The (Kq, Kp) landscape shows stable train-selected budgets in-domain but more varied per-dataset optima out-of-domain. Aggregate train-selected settings remain strong, while query- and dataset-specific preferences leave room for adaptive allocation.

  • In-domain landscape: Train-selected settings are optimal or tied for optimal across every in-domain panel.For Dream, (4, 16) is best on every in-domain dataset; for LLaDA, (4, 4) is best on DL19 and lies on tied plateaus elsewhere.
  • Backbone patterns: Dream favors passage-heavy budgets, whereas LLaDA is more isotropic around the symmetric center in every in-domain panel.This per-backbone asymmetry is reported consistently across the in-domain landscapes.
  • Out-of-domain landscape: Out-of-domain per-dataset peaks scatter across the (Kq, Kp) grid despite train-selected cells being best on the BEIR-7 average for both backbones.HotpotQA peaks at (8, 16) for both backbones, while several LLaDA datasets favor fewer passage representations.
  • Adaptive allocation: A fixed train-selected budget works well in aggregate but leaves effectiveness on the table when queries or datasets prefer different budgets.This observation is consistent with the per-query oracle analysis.
Loading 2605.07210v2…