DiffuGR: Generative Document Retrieval with Diffusion Language Models

TL;DR

DiffuGR reframes document retrieval as discrete diffusion-based DocID generation, enabling parallel token updates and iterative refinement through denoising. It introduces two DocID constructions (learnable via residual quantization and linguistic via titles/leading tokens) and trains with a masked diffusion objective that optimizes an upper bound on the negative log-likelihood, $-\\log p_{ heta}(\bm{x_0}) \le \int_0^1 \frac{1}{t} \\mathbb{E}_{q(\bm{x_t}|\bm{x_0})}[ \\sum_{\\{i| \\bm{x_t^i}=m\\}} -\\log p_{\\theta}(\\bm{x_0^i}|\\bm{x_t})] \, dt$. Inference generates DocIDs in parallel for a fixed budget of $T$ denoising steps, balancing quality and latency. Empirical results on NQ320K and MS MARCO show competitive performance against autoregressive generative retrievers, with linguistic DocIDs delivering the strongest gains and pseudo beam search offering robustness enhancements. Limitations include the lack of beam search for top-$k$ generation and constrained decoding, motivating future work toward top-$k diffusion generation and larger diffusion models.

Abstract

Generative retrieval (GR) re-frames document retrieval as a sequence-based document identifier (DocID) generation task, memorizing documents with model parameters and enabling end-to-end retrieval without explicit indexing. Existing GR methods are based on auto-regressive generative models, i.e., the token generation is performed from left to right. However, such auto-regressive methods suffer from: (1) mismatch between DocID generation and natural language generation, e.g., an incorrect DocID token generated in early left steps would lead to totally erroneous retrieval; and (2) failure to balance the trade-off between retrieval efficiency and accuracy dynamically, which is crucial for practical applications. To address these limitations, we propose generative document retrieval with diffusion language models, dubbed DiffuGR. It models DocID generation as a discrete diffusion process: during training, DocIDs are corrupted through a stochastic masking process, and a diffusion language model is learned to recover them under a retrieval-aware objective. For inference, DiffuGR attempts to generate DocID tokens in parallel and refines them through a controllable number of denoising steps. In contrast to conventional left-to-right auto-regressive decoding, DiffuGR provides a novel mechanism to first generate more confident DocID tokens and refine the generation through diffusion-based denoising. Moreover, DiffuGR also offers explicit runtime control over the qualitylatency tradeoff. Extensive experiments on benchmark retrieval datasets show that DiffuGR is competitive with strong auto-regressive generative retrievers, while offering flexible speed and accuracy tradeoffs through variable denoising budgets. Overall, our results indicate that non-autoregressive diffusion models are a practical and effective alternative for generative document retrieval.

Figures (5)

• Figure 1: In contrast to auto-regressive generative retrieval (above), which generates the DocID from left to right, DiffuGR (below) efficiently generates multiple tokens in parallel and then re-masks uncertain tokens. The generation is progressively refined over multiple denoising steps.
• Figure 2: An overview of DiffuGR, consisting DocID construction, model training, and inference. DiffuGR involves two kinds of DocIDs, i.e., linguistic DocIDs and learnable DocIDs. During training, the model is optimized to recover randomly masked DocID tokens. For inference, DiffuGR generates DocID tokens in parallel and refines them across multiple denoising steps.
• Figure 3: Quality-speed comparison on NQ320K for Dream 7B and Qwen2.5-7B. By adjusting the denoising steps, the performance of DiffuGR can be flexibly tuned towards either speed or quality.
• Figure 4: Retrieval performance of DiffuGR on NQ320K under different model scales. Larger models consistently yield higher retrieval accuracy, demonstrating that increased model sizes enhance the retrieval process. However, larger models also incur substantially higher training costs, revealing a clear tradeoff between effectiveness and efficiency.
• Figure 5: Quality-speed tradeoff on MS MARCO.