DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

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

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 3 functions out of this paper's own repositories and ran 1 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

Repositories linked to this paper

• ielab/diffretriever

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

Abstract

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. DiffRetriever Multi-token: Effectiveness ↑ ~similar latency PromptReps Multi-token ~Effectiveness ↑ much higher latency DiffEmbed/RepLLaMA Effectiveness ↓↓ Similar latency Figure 1: BEIR-7 NDCG@10 vs. query encoding plus search latency (ms/query, 100K-document sample). Appendix B.2 reports full latency scaling.

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2605.07210")
get_code_for_paper("2605.07210")
have("2605.07210")

Connect an agent — have() is free.