DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models
Summary
This paper introduces DiffRetriever, a method that uses diffusion language models to generate multiple representative tokens in parallel for efficient information retrieval, outperforming autoregressive baselines in speed and accuracy.
View Cached Full Text
Cached at: 05/12/26, 02:50 AM
Paper page - DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models
Source: https://huggingface.co/papers/2605.07210
Abstract
DiffRetriever enables efficient multi-token retrieval using diffusion language models by generating representations in parallel rather than sequentially, achieving superior performance over autoregressive methods.
PromptReps showed that anautoregressive language modelcan 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 priormulti-tokenvariants did not reliably improve oversingle-token decoding. We show that the bottleneck is sequential generation, not themulti-tokenidea itself. DiffRetriever is arepresentative-token retrieverfordiffusion language models: it appends Kmasked positionsto the prompt and reads all K in a singlebidirectional forward pass. Across in-domain and out-of-domain evaluation,multi-tokenDiffRetriever substantially improves over single-token on every diffusion backbone we test, while autoregressivemulti-tokenis flat or negative and pays a latency cost that scales with K where diffusion does not. Aftersupervised fine-tuning, DiffRetriever onDreamis the strongestBEIR-7retriever in our comparison, ahead of PromptReps, the encoder-style DiffEmbed baseline on the same diffusion backbones, and the contrastively fine-tuned single-vectorRepLLaMA. A per-query oracle on the frozen base model exceedscontrastive fine-tuningat the same fixed budget, pointing to adaptive budget selection as future work. Code is available at https://github.com/ielab/diffretriever.
View arXiv pageView PDFGitHub2Add to collection
Get this paper in your agent:
hf papers read 2605\.07210
Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash
Models citing this paper0
No model linking this paper
Cite arxiv.org/abs/2605.07210 in a model README.md to link it from this page.
Datasets citing this paper0
No dataset linking this paper
Cite arxiv.org/abs/2605.07210 in a dataset README.md to link it from this page.
Spaces citing this paper0
No Space linking this paper
Cite arxiv.org/abs/2605.07210 in a Space README.md to link it from this page.
Collections including this paper0
No Collection including this paper
Add this paper to acollectionto link it from this page.
Similar Articles
Set Diffusion: Interpolating Token Orderings Between Autoregression and Diffusion for Fast and Flexible Decoding
Set Diffusion introduces a new class of language models that interpolates between autoregressive and diffusion models by factorizing token generation over flexible-position, flexible-length token sets. This enables faster decoding and flexible token ordering, achieving better speed-quality tradeoffs on reasoning, summarization, and unconditional generation tasks.
How To Build A Modern Diffusion Language Model
A Cornell University article/survey explaining how to build modern discrete diffusion language models, covering iterative parallel refinement, fast sampling at thousands of tokens per second, and their potential to improve scaling laws over autoregressive LLMs.
PreDiff-LM: Pretrained Discrete Masked Diffusion Language Modeling with Hybrid Attention
PreDiff-LM proposes a hybrid attention mechanism that preserves causal attention for prompt tokens and bidirectional attention for masked target tokens, enabling adaptation of pretrained autoregressive models for discrete masked diffusion language modeling, achieving improvements in perplexity and downstream tasks over prior diffusion baselines.
ReDiTT: Retrieval Augmented Conditional Diffusion Transformers for Asynchronous Time Series
This paper presents ReDiTT, a retrieval augmented conditional diffusion transformer for asynchronous time series prediction. The model retrieves structurally similar latent sequences as reference conditions to improve long-horizon forecasting and sample diversity, achieving state-of-the-art performance on seven real-world datasets.
Representation-Space MMD for Diffusion Language Models
This paper introduces a post-training method for diffusion language models that minimizes Maximum Mean Discrepancy between generated and reference distributions in the feature space of a frozen pretrained DLM, achieving lower generative perplexity and better accuracy-compute trade-offs on OpenWebText and GSM8K, with improved decoding parallelism on 16B LLaDA-2.0 models.