DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

Abstract: PromptReps showed that an autoregressive LLM can be used directly as a retriever by prompting it to generate dense and sparse representations of a query or passage. Extending this to multiple representatives is inefficient for autoregressive models, since tokens must be generated sequentially, and prior multi-token variants did not reliably improve over single-token decoding. We show that the bottleneck is sequential generation, not the multi-token idea itself. DiffRetriever is a representative-token retriever for diffusion LLMs: it appends K masked positions to the prompt and reads all K in a single bidirectional forward pass. Across in-domain and out-of-domain evaluation, multi-token DiffRetriever substantially improves over single-token on every diffusion backbone we test, while autoregressive multi-token is flat or negative and pays a latency cost that scales with K where diffusion does not. After supervised fine-tuning, DiffRetriever on Dream is the strongest BEIR-7 retriever in our comparison, ahead of PromptReps, the encoder-style DiffEmbed baseline on the same diffusion backbones, and the contrastively fine-tuned single-vector RepLLaMA. A per-query oracle on the frozen base model exceeds contrastive fine-tuning at the same fixed budget, pointing to adaptive budget selection as future work. Code is available at https://github.com/ielab/diffretriever.

• The paper introduces DiffRetriever, a novel framework that employs parallel masked token prediction in diffusion language models to enhance multi-token retrieval efficiency.
• It demonstrates significant retrieval quality and latency improvements over autoregressive models by aligning inference with the diffusion pretraining objective.
• Empirical results reveal that fine-tuned diffusion models double performance metrics (e.g., MRR@10 on MS MARCO) while maintaining constant latency regardless of token count.

DiffRetriever: Parallel Retrieval with Diffusion LLMs

Introduction

This paper presents DiffRetriever, a retrieval framework that exploits parallel masked token prediction in diffusion LLMs (DiffLMs) to produce multiple representative embeddings per input for dense and sparse retrieval. By directly aligning retrieval-time inference with the diffusion pretraining objective (K-way masked prediction), DiffRetriever enables efficient parallel extraction of K representative embeddings per query or passage in a single bidirectional forward pass, overcoming the inherent sequential bottleneck of autoregressive (AR) LLMs. The work analyzes the effectiveness and efficiency tradeoffs of multi-token retrieval across diffusion and AR backbones, positioning diffusion-based retrievers as a compelling alternative in both zero-shot and fine-tuned regimes.

Background and Motivation

PromptReps has demonstrated that representative-token prompting can enable zero-shot retrieval with AR LLMs by reading off dense vectors and sparse token logits sequentially, but prior studies observed that extending to K>1 representatives incurs significant inference latency without consistent gains in retrieval quality. The central question addressed here is whether the multi-token retrieval weakness originates from the inefficacy of multi-token representations or from AR models’ sequential inference overhead. Diffusion LLMs, trained to fill masked positions in parallel under bidirectional attention, offer the architectural means to disentangle these factors.

State-of-the-art diffusion LLMs (e.g., Dream, LLaDA) have demonstrated competitive text generation performance with architectural and training advances, but previous use for retrieval paradigms (such as DiffEmbed) treated them as frozen encoders with mean pooling, discarding the native pretraining objective at inference.

Methods

Parallel Masked Token Retrieval with Diffusion LMs

DiffRetriever reformulates retrieval representation extraction by appending $K_{\mathrm{q}}$ or $K_{\mathrm{p}}$ [MASK] positions to the query/passage prompt, which are simultaneously filled in a single bidirectional pass of a diffusion LM. Both dense hidden states and sparse logits are extracted for all predicted [MASK] positions, yielding $K$-way vector and logit sets per text. In AR backbones, K representatives require K serial forward passes, resulting in wall-clock latency that grows linearly with K; in contrast, diffusion backbones maintain constant latency with respect to K.

The pipeline is aligned with PromptReps for prompt structure, score computation (ColBERT-style late interaction for dense, max-pooled interpolated sparse), and supervised contrastive fine-tuning. However, DiffRetriever's representational expressivity and parallel efficiency are strictly a function of using a diffusion backbone as designed—joint masked-token prediction.

Model-Task Assignment and Comparison Schema

Two AR (LLaMA3-8B, Qwen2.5-7B) and two diffusion (Dream-7B, LLaDA-8B) backbones are compared under controlled experimental conditions for both zero-shot and fine-tuned regimes. The Qwen2.5/Dream pairing is especially revealing, as Dream simply reinitializes from Qwen2.5 then continues training with bidirectional masking, providing a near-isolate for the effect of objective and inference alone. Baselines include PromptReps (AR, sequential), DiffEmbed (diffusion, mean-pool), RepLLaMA (contrastive, AR), and BM25.

Empirical Results

Zero-Shot and Fine-Tuned In-Domain Retrieval

In the single-token condition, AR backbones outperform diffusion backbones on MS MARCO and TREC DL, reflecting the maturity of left-to-right next-token objectives for producing a singular summary token. However, multi-token retrieval is markedly more effective for diffusion backbones: multi-token DiffRetriever doubles the retrieval quality for Dream on MS MARCO (MRR@10 0.112 → 0.218), statistically outperforming AR baseline systems. Importantly, multi-token AR variants do not yield gains and are subject to a severe latency penalty (up to 15x at zero-shot settings with K≤20).

Supervised fine-tuning benefits all models, compressing performance differences but not eliminating them. Fine-tuned Dream (DiffRetriever, multi-token) achieves top dense and hybrid scores across most benchmarks, while AR models still trail on multi-token efficiency and rarely surpass their own single-token configurations.

Out-of-Domain (BEIR-7) Generalization

Zero-shot transfer to BEIR-7 reconfirms the finding: single-token diffusion falls short, but multi-token DiffRetriever variants achieve the highest or near-highest NDCG@10 on the aggregate and in strong positions per dataset. LLaDA-multi-token achieves the strongest zero-shot BEIR-7 average (NDCG@10 0.539). Following contrastive fine-tuning, Dream-multi-token (DiffRetriever) achieves 0.671, the top reported value, and maintains robust per-dataset consistency.

Latency Analysis

Autoregressive multi-token retrieval incurs linear scaling of encoding latency in K (e.g., 275-300 ms for K=20), while diffusion models hold steady at 16-20 ms for the same K, maintaining near single-token cost with a multiplicative increase in representation power. End-to-end retrieval latency (encoding + search) is overwhelmingly dominated by the encoding step at corpus sizes up to $10^6$ passages; thus, the parallel advantage is practically significant for large-scale or real-time systems.

Budget Selection and Adaptive Routing

The necessity to fix (Kq, Kp) budgets a priori in diffusion backbones is empirically validated: grid search on train data generalizes well to test and out-of-domain data, but per-query optimal K selection can outperform any fixed setting and even contrastive fine-tuning at those settings. Oracle K selection per query (using ground truth labels, not deployable) outperforms fixed-budget models by up to 0.21 NDCG. The paper observes moderate positive correlation between query length/entropy and optimal Kq, suggesting tractability for future adaptive budget routers, but hand-crafted features alone are insufficient to close the gap to oracle selection.

Training and Denoising Steps

Fine-tuning with single-pass (S=1) inference is generally optimal for zero-shot and fine-tuned Dream, while iterative denoising (S>1) can offer marginal gains for LLaDA out-of-domain. Thus, the recommended deployment for top-tier retrieval is parallel readout in a single denoising step.

Notable Claims and Contributions

• Multi-token masked representative extraction is highly effective with diffusion LMs, providing statistically significant gains over single-token extraction at constant latency, in stark contrast to AR backbones where multi-token representations are not only more expensive but less effective.
• Retrieval-time alignment with the diffusion pretraining objective is critical: using masked-token prediction at inference outperforms mean-pooled encoder-style usage (i.e., DiffEmbed) on identical backbones.
• The fixed budget constraint in parallel masked extraction—once considered a limitation—can be compensated for or even exceeded by adaptive policies, as per-query oracle selection can surpass models trained with contrastive supervision.
• The efficiency improvement is large in practical deployments, especially with long sequences or large-scale retrieval corpora.

Implications and Future Directions

This work demonstrates that diffusion LLMs, when queried in a manner consistent with their pretraining (parallel masked position filling), offer both a scalable and robust alternative for retrieval tasks—displacing the AR bias in current LLM-based retrievers, especially for multi-token late interaction paradigms. While currently top-line performance is obtained at the 7–8B scale, the results suggest substantial headroom for larger or more deeply tuned models.

The adaptive budget selection analysis establishes a novel research direction: designing data-driven routers for per-query (Kq, Kp) selection, potentially leveraging query diagnostics or calibrating on cheap surface features for efficiency-effectiveness tradeoff optimization. Such approaches promise even further gains, potentially surpassing what can be obtained with extensive fine-tuning.

The analysis also draws attention to the limitations of mean-pooling diffusion LMs for retrieval, arguing for representation extraction via tasks that match the pretraining denoising objective.

Conclusion

DiffRetriever demonstrates that parallel representative-token retrieval with diffusion LMs realizes substantial effectiveness and latency gains in multi-token retrieval, reversing previous findings from AR LMs and establishing a new state-of-the-art retriever under both zero-shot and supervised regimes. Central to these gains is direct architectural and objective alignment between pretraining and retrieval-time inference. Future developments in model routing, larger-scale evaluation, and hybrid adaptive systems represent promising directions for further exploration in LLM-based information retrieval.

Reference: "DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion LLMs" (2605.07210).