Self-Augmenting Retrieval for Diffusion Language Models

Discrete diffusion language models generate text by iteratively denoising an entire response in parallel. At each step, they predict tentative tokens for every masked position, committing the confident predictions to the output and discarding the unconfident ones. We show that the discarded tokens are in fact a useful lookahead signal for retrieval-augmented generation: even low-confidence tokens often surface salient entities early in the denoising trajectory, enabling retrieval of stronger evidence before the output is finalized. We exploit this through Self-Augmenting Retrieval for Diffusion Language Models (SARDI), a dynamic RAG framework that uses these lookahead tokens to guide retrieval during denoising. SARDI is training-free, retriever-agnostic, and applicable to any reasoning-capable discrete diffusion language model. Across five multi-hop QA benchmarks, SARDI outperforms current training-free diffusion and autoregressive retrieval baselines at up to $8\times$ higher throughput.

When AI answers a question, it usually looks up relevant documents first, the way a student checks references before writing. But some questions hide what to look up. Take "which city is home to the museum that displays the Mona Lisa?" You first need to find the museum before you can find the city, and that clue, the name of the museum, doesn't appear in the question. Most AI writes one word at a time, so it can only search using what it has already written. We studied a newer kind of AI that drafts the whole answer at once and gradually sharpens it. We found that its rough early drafts often point to that missing clue, the museum, well before the answer is settled. Our method, SARDI, feeds these early glimpses to the search engine and hands the fresh results back to the AI to improve its next draft. SARDI needs no training, works with any search tool, and runs several times faster than comparable methods while matching or beating their accuracy.