Self-Augmenting Retrieval for Diffusion Language Models

arXiv:2606.06474 · cs.CL, cs.AI, cs.LG · Submitted 2026-06-04 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Self-Augmenting Retrieval for Diffusion Language Models".

Jane: Here's the gist of the paper. 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 Across five multi-hop QA benchmarks, SARDI outperforms current training-free diffusion and autoregressive retrieval baselines at up to 8× higher throughput SARDI is training-free, retriever-agnostic, and applicable to any reasoning-capable discrete diffusion language model Pruning redundant reasoning tokens Diffusion language models refine the entire response at once: at every denoising step they produce tentative predictions for all token positions This trajectory of intermediate predictions surfaces salient entities and relations before the output is finalized Diffusion trajectories thus let the model “peek into the future” of its own generation to improve retrieval The paper proposes SARDI, which interleaves retrieval with denoising by conditioning retrieval on intermediate diffusion states At each step t, the model predicts a token for every masked position with a confidence score ci = maxv∈V pθ(v x t, q, D t) If ci ≥ τq (the query threshold), the predicted token is added to the retrieval query, refreshing the evidence for the next step The central design choice in SARDI is a separation between retrieval and generation confidence unique to non-autoregressive decoders How the training signal is built SARDI relies on the existence of a reasoning trace in order to surface bridge entities To evaluate SARDI fairly, the authors elicit the model to output reasoning via light supervised finetuning on synthetic chain-of-thought traces from gpt-4omini The fine-tuning data is constructed from training splits of 2WikiMultiHopQA and HotpotQA, augmented with synthetically generated chain-of-thought reasoning traces The synthetic reasoning traces prompt gpt-4o-mini to generate a step-by-step reasoning trace that identifies intermediate bridge entities, references specific facts from the supporting documents, and concludes with the final answer in the format

sic: This format serves two purposes: it enables extraction of intermediate entities from partial reasoning traces for selfaugmenting retrieval, and it provides a consistent extraction point (the

sic: ) for evaluation Where the gains come from SARDI improves substantially over static diffusion retrieval on every benchmark On five multi-hop QA benchmarks, SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The gains concentrate on multi-hop reasoning, yielding over 2.5× EM gains on inference and compositional questions (which require identifying bridge entities through multi-step reasoning) Furthermore, grounding generation in retrieved evidence sharply reduces inter-token dependence, making RAG well-suited to parallel decoding This is confirmed by measuring the conditional mutual information (CMI) between adjacent tokens in ground-truth reasoning traces on 2WikiMultiHopQA, showing that strong evidence removes the coupling most harmful to parallel decoding Where the gains come from The paper demonstrates that grounding generation in retrieved evidence sharply reduces inter-token dependence, making RAG well-suited to parallel decoding The results in Table 3 confirm this, showing that with all gold documents present, entity pairs show very low dependence (CMI=0.060), i.e., they can be decoded in parallel This confirms that the RAG setting is well suited to parallel decoding The paper also shows that speculative future tokens can inform retrieval long before they are stable enough to commit to the output, as evidenced by the peak EM occurring at the most aggressive lookahead setting (τq ≈ 0) Where the gains come from The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static diffusion retrieval on every benchmark, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora The impact statement advises deploying SARDI with curated sources and provenance tracking, such as citing retrieved passages alongside generated answers All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.com/pauljngr/SARDI How the gains come from The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.com/pauljngr/SARDI Where the gains come from The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.com/pauljngr/SARDI Where the gains come from The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.

Paper summary: Tom: So to recap the paper "Self-Augmenting Retrieval for Diffusion Language Models", the core thesis is about how discrete diffusion language models generate text by iteratively denoising a whole response in parallel at every step.

Jane: And at each of those steps, they predict tentative tokens for every masked position in the output. But they discard the ones they aren't confident about. The paper claims these discarded tokens are actually a useful lookahead signal for retrieval augmented generation.

Lu: This means that even low-confidence tokens can surface salient entities early in the denoising trajectory, which allows for retrieval of stronger evidence before the final output is finalized.

Meng: So they're using the model's internal uncertainty as a trigger to proactively fetch external information. How do they structure this?

Jane: They exploit this through SARDI, a dynamic RAG framework that uses those lookahead tokens to guide retrieval during the denoising process itself. It’s designed to be training-free and retriever-agnostic.

Tom: The paper shows that across five multi-hop QA benchmarks, SARDI outperforms current training-free diffusion and autoregressive retrieval baselines at up to eight times higher throughput <ref:2606.06474#pg1,across five multi-hop QA benchmarks, SARDI outperforms current training-free diffusion>. That's a big number for inference speed.

Lu: This performance comes from interleaving retrieval with denoising by conditioning the retrieval query on the intermediate diffusion states. The central design choice is separating generation confidence from retrieval confidence, which is unique for non-autoregressive decoders.

Jane: At each step t, the model predicts tokens with a confidence score ci equal to max v in V p theta of v given x t, q, D t. If that confidence ci is greater than or equal to a query threshold tau q then the predicted token gets added to the retrieval query and refreshes the evidence for the next step.

Meng: So they're dynamically updating what they're searching for based on what they think is important at that precise moment during generation. That sounds like it might be expensive to run, though.

Tom: The authors address that by showing how this works using synthetic chain-of-thought traces generated from models like gpt-4o-mini. This fine-tuning creates the reasoning trace needed to surface bridge entities effectively.

Lu: They use these traces, which they generate by prompting gpt-4o-mini for a step-by-step reasoning trace identifying intermediate bridge entities and references specific facts.

Jane: The evaluation setup involves using this format to extract intermediate entities from partial reasoning traces for self-augmenting retrieval, and it also provides a consistent extraction point with the sic notation for evaluation.

Tom: They show where the gains come from by comparing SARDI against static diffusion retrieval on every benchmark. It substantially improves performance across the board.

Lu: Specifically, they see substantial improvements on multi-hop reasoning tasks, yielding over two point five times EM gains on inference and compositional questions that require identifying bridge entities through multi-step reasoning.

Jane: They also confirm that grounding generation in retrieved evidence reduces inter-token dependence, which makes RAG well-suited to parallel decoding because entity pairs show very low dependence, with a conditional mutual information of zero point zero six zero on 2WikiMultiHopQA data.

Meng: So the speed boost isn't just about the RAG part; it's about making sure the generation stays coherent even when you have to pull in external documents constantly during parallel decoding.

Tom: And they show that speculative future tokens can inform retrieval long before they are stable enough to commit, with peak EM occurring at a very aggressive lookahead setting where tau q is approximately zero <ref:2606.06474#pg2,speculative future tokens can inform retrieval long before they are stable enough>.

Conclusion: Jane: So we've covered the basics of "Self-Augmenting Retrieval for Diffusion Language Models" and how it works by interleaving retrieval with denoising to refine the query as the response takes shape.

Tom: The authors are Paul Junger, Justin Lovelace, Linxi Zhao, Dongyoung Go, and Kilian Q. Weinberger. They developed SARDI, which is a dynamic RAG framework that uses lookahead tokens to guide retrieval during diffusion.

Lu: The paper suggests deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora. That's a practical piece of advice for real-world application.

Meng: From an engineering standpoint, the authors also suggest integrating Fast-dLLM-style block KV caching as a natural direction for future work because the retrieved documents persist between consecutive steps.

Jane: Overall, what this paper shows is that denoising trajectories offer a natural signal for dynamic retrieval, and they expect SARDI to get stronger as diffusion language models mature. It substantially improves over static retrieval and matches or beats every training-free autoregressive baseline in terms of throughput.

Tom: The main thing to remember is that SARDI leverages both diffusion trajectories and retrieval benefits by interleaving them to refine the query while the response is forming.

Lu: And they are betting that this approach will become a stronger way for diffusion language models to handle complex reasoning tasks in the future because it uses an intrinsic signal from their generation process.

Cornell University

cs.CL, cs.AI, cs.LG

Submitted: 2026-06-04

Updated: 2026-10-29

Comments: ICML 2026

Code: https://github.com/pauljngr/SARDI

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 85/100

The gist: 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

Key concepts

Discrete Diffusion Language Models
These models create text by gradually refining a response in parallel through an iterative denoising process. At each step, they predict tokens for all masked positions simultaneously. The key is that the discarded, low-confidence predictions contain valuable information about what the model intends to generate next.
Self-Augmenting Retrieval for Diffusion Language Models (SARDI)
SARDI is a dynamic retrieval framework that integrates retrieval directly into the denoising process. It uses tentative tokens generated at each step to decide whether to perform a retrieval query. This allows the model to dynamically fetch relevant evidence precisely when it needs it during generation.
Reasoning Trace Evaluation
To train SARDI, researchers use synthetic chain-of-thought traces generated by GPT-4o-mini. These traces detail step-by-step reasoning, identifying intermediate entities and facts. This provides the necessary 'reasoning trace' signal that SARDI exploits to surface salient entities for retrieval.
Inter-token Dependence Reduction
SARDI improves generation quality by grounding it in retrieved evidence, which reduces how much one token depends on the previous one. Measuring this dependence shows that strong evidence allows tokens to be decoded in parallel, confirming that RAG is well-suited for the parallel decoding nature of diffusion models.

Terminology

Summary

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 Across five multi-hop QA benchmarks, SARDI outperforms current training-free diffusion and autoregressive retrieval baselines at up to 8× higher throughput SARDI is training-free, retriever-agnostic, and applicable to any reasoning-capable discrete diffusion language model

Pruning redundant reasoning tokens

Diffusion language models refine the entire response at once: at every denoising step they produce tentative predictions for all token positions This trajectory of intermediate predictions surfaces salient entities and relations before the output is finalized Diffusion trajectories thus let the model “peek into the future” of its own generation to improve retrieval The paper proposes SARDI, which interleaves retrieval with denoising by conditioning retrieval on intermediate diffusion states At each step t, the model predicts a token for every masked position with a confidence score ci = maxv∈V pθ(v x t, q, D t) If ci ≥ τq (the query threshold), the predicted token is added to the retrieval query, refreshing the evidence for the next step The central design choice in SARDI is a separation between retrieval and generation confidence unique to non-autoregressive decoders

How the training signal is built

SARDI relies on the existence of a reasoning trace in order to surface bridge entities To evaluate SARDI fairly, the authors elicit the model to output reasoning via light supervised finetuning on synthetic chain-of-thought traces from gpt-4omini The fine-tuning data is constructed from training splits of 2WikiMultiHopQA and HotpotQA, augmented with synthetically generated chain-of-thought reasoning traces The synthetic reasoning traces prompt gpt-4o-mini to generate a step-by-step reasoning trace that identifies intermediate bridge entities, references specific facts from the supporting documents, and concludes with the final answer in the format [sic] This format serves two purposes: it enables extraction of intermediate entities from partial reasoning traces for selfaugmenting retrieval, and it provides a consistent extraction point (the [sic]) for evaluation

Where the gains come from

SARDI improves substantially over static diffusion retrieval on every benchmark On five multi-hop QA benchmarks, SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The gains concentrate on multi-hop reasoning, yielding over 2.5× EM gains on inference and compositional questions (which require identifying bridge entities through multi-step reasoning) Furthermore, grounding generation in retrieved evidence sharply reduces inter-token dependence, making RAG well-suited to parallel decoding This is confirmed by measuring the conditional mutual information (CMI) between adjacent tokens in ground-truth reasoning traces on 2WikiMultiHopQA, showing that strong evidence removes the coupling most harmful to parallel decoding

The paper demonstrates that grounding generation in retrieved evidence sharply reduces inter-token dependence, making RAG well-suited to parallel decoding The results in Table 3 confirm this, showing that with all gold documents present, entity pairs show very low dependence (CMI=0.060), i.e., they can be decoded in parallel This confirms that the RAG setting is well suited to parallel decoding The paper also shows that speculative future tokens can inform retrieval long before they are stable enough to commit to the output, as evidenced by the peak EM occurring at the most aggressive lookahead setting (τq ≈ 0)

The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static diffusion retrieval on every benchmark, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora The impact statement advises deploying SARDI with curated sources and provenance tracking, such as citing retrieved passages alongside generated answers All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.com/pauljngr/SARDI

How the gains come from

The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.com/pauljngr/SARDI

The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.com/pauljngr/SARDI

The authors conclude that SARDI leverages both diffusion trajectories and retrieval benefits, interleaving retrieval with denoising to refine the query as the response takes shape SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster More broadly, denoising trajectories offer a natural signal for dynamic retrieval, and the authors expect SARDI to strengthen as diffusion language models mature The paper recommends deploying SARDI with curated sources and provenance tracking to mitigate risks associated with using untrusted corpora All experiments use publicly available QA benchmarks and open retrieval corpora The paper suggests that integrating Fast-dLLM-style block KV caching is a natural direction for future work due to the persistence of retrieved documents between consecutive steps The conclusion is that SARDI substantially improves over static retrieval, matches or beats every training-free AR baseline, and runs up to 8× faster The authors expect SARDI to strengthen as diffusion language models mature The paper is available at https://github.

Improvements for AI systems

  1. Retrieval-augmented generation for multi-hop question answering can be significantly improved by grounding generation in intermediate diffusion states, as SARDI surfaces salient entities and relations before the output is finalized enabling retrieval of later-hop evidence before the final answer is committed.

  2. Discrete diffusion language models can be leveraged to enable parallel decoding by exploiting the property that grounding sharply reduces inter-token dependence, promoting parallel decoding, which allows for generation speeds up to 8 times faster than autoregressive iterative-retrieval baselines.

  3. The system can dynamically adjust accuracy and latency trade-offs by controlling the commit confidence threshold tau c, as a higher tau c commits fewer tokens per denoising step and triggers more retrieval rounds (higher accuracy, lower throughput), while a lower tau c does the reverse.

  4. The model can be made effective on reasoning tasks without requiring specialized training by utilizing speculative future tokens to inform retrieval, since speculative future tokens can inform retrieval long before they are stable enough to commit.

  5. For complex reasoning questions requiring multi-step analysis, the system will show over 2.5 times exact match gains on inference and compositional questions because SARDI benefits most from question types requiring multi-step reasoning.

Abstract

Diffusion language models generate text by iteratively denoising an entire response in parallel, holding a noisy draft of the answer at every step. We show that these intermediate states surface salient entities early in the denoising trajectory, making them a useful lookahead signal for retrieval-augmented generation (RAG). We further demonstrate that grounding generation in retrieved evidence sharply reduces inter-token dependence, making RAG particularly well suited to parallel decoding. We exploit both through Self-Augmenting Retrieval for Diffusion Language Models (SARDI), a dynamic RAG framework that uses speculative token predictions 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 training-free diffusion and autoregressive retrieval baselines at up to 4x higher throughput.

Sources

Related papers