Self-Augmenting Retrieval for Diffusion Language Models
Paul Jünger, Justin Lovelace, Linxi Zhao, Dongyoung Go, Kilian Q. Weinberger
Discrete diffusion language models iteratively denoise masked sequences in parallel rather than generating tokens left-to-right. SARDI repurposes the low-confidence tokens discarded at each denoising step as a lookahead signal for retrieval — those uncertain tokens often surface relevant named entities early enough to fetch better evidence before generation finalizes. On five multi-hop QA benchmarks, this approach beats training-free autoregressive and diffusion RAG baselines while achieving up to 8× higher throughput, with no retriever dependencies or fine-tuning required.
No production traction yet. The GitHub repos referencing this paper are all arXiv aggregators and paper-tracking tools, not implementations. Zero citations and no dedicated codebase. Worth watching if you're building on discrete diffusion models like MDLM or similar — the throughput claim is the interesting lever, since multi-hop QA at 8× the speed of autoregressive RAG would matter for latency-sensitive pipelines — but nothing is shipping today.
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.