DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

TL;DR

DiffRetriever introduces a parallel token retrieval method for diffusion language models, significantly enhancing retrieval performance and efficiency over previous autoregressive approaches.

Contribution

It presents a novel parallel token retrieval technique for diffusion models, overcoming sequential bottlenecks and improving retrieval accuracy and speed.

Findings

DiffRetriever outperforms single-token methods across multiple diffusion backbones.

Multi-token DiffRetriever improves retrieval over autoregressive multi-token methods.

Code is available at https://github.com/ielab/diffretriever.

Abstract

PromptReps showed that an autoregressive language model 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 language models: 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.