Self-Augmenting Retrieval for Diffusion Language Models
Paul Jünger 1 Justin Lovelace 1 Linxi Zhao 1 Dongyoung Go 1 Kilian Q. Weinberger 1
arXiv:2606.06474v1 [cs.CL] 4 Jun 2026
Abstract
with a test-time-adjustable number of iterations (Austin et al., 2021; Li et al., 2022; Snell et al., 2025; Ye et al., 2025; Lovelace et al., 2023). Recent discrete DLMs such as DREAM-7B (Ye et al., 2025) and LLaDA (Nie et al., 2026) have matured into competitive models at scale, raising the question of how their non-autoregressive structure can be exploited beyond raw throughput.
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 lowconfidence 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× higher throughput. Our code is available at https://github.com/ pauljngr/SARDI.
Denoising unlocks two structural advantages over AR generation in retrieval-augmented settings: one for retrieval and one for decoding. Diffusion Trajectories as Lookahead for Retrieval. 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. This is especially useful for multi-hop question answering, where the evidence needed for later reasoning steps depends on intermediate bridge entities that the question alone does not specify. Figure 1 illustrates this: asked which city is home to the museum that displays the Mona Lisa, a static retriever cannot retrieve the passage “The Louvre is located in Paris” without first identifying the bridge entity Louvre. Intermediate diffusion states surface this entity early in the trajectory, enabling retrieval of the second-hop evidence before the final answer is committed. Diffusion trajectories thus let the model “peek into the future” of its own generation to improve retrieval.
1. Introduction Large language models are increasingly deployed as interfaces to external knowledge, with retrieval-augmented generation grounding responses in retrieved evidence (Lewis et al., 2020; Izacard & Grave, 2021). Yet the dominant decoding paradigm, token-by-token autoregressive (AR) generation, ties latency directly to output length (Pope et al., 2023) and forces retrieval to condition on a left-to-right committed prefix (Trivedi et al., 2023). Diffusion language models (DLMs) offer an alternative. Rather than emitting one token at a time, they iteratively denoise a corrupted sequence, updating many positions in parallel at each step 1 Department of Computer Science, Cornell University. Correspondence to: Paul Jünger <[email protected]>.
RAG grounding promotes parallel decoding. Realizing the speedup potential of parallel decoding has proven challenging in practice: sampling multiple tokens simultaneously risks generating conflicting or incoherent spans, errors that can cascade through the response and degrade output quality (Wu et al., 2026). RAG changes this. When generation is grounded in informative evidence D, many output tokens are copied or paraphrased straight from the context and are therefore conditionally independent given D. We confirm this empirically in Section 5.5: grounding sharply reduces inter-token dependence, promoting parallel decoding.
Proceedings of the 43 rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s).
We propose Self-Augmenting Retrieval for Diffusion Language Models (S ARDI), the first framework to condition retrieval on intermediate diffusion states. S ARDI interleaves 1
Self-Augmenting Retrieval for Diffusion Language Models
Which city is home to the museum that displays the Mona Lisa? (2-hop question) Step 0
[MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK]
retrieve + doc. retrieval (based on question): "The Mona Lisa is displayed at the Louvre [...]" Bridge entity
Step 1 denoise
[MASK] is home to the Louvre, which displays the Mona Lisa.
retrieve
+ doc. retrieval (based on question+partial generation): "The Louvre is located in Paris."
Output
Paris is home to the Louvre, which displays the Mona Lisa.
Figure 1. Static question-only retrieval can fail on multi-hop QA when the question does not specify the bridge entity. Intermediate diffusion states often surface such entities early, enabling retrieval of later-hop evidence before the output is finalized.
retrieval with denoising: at each iteration, it constructs a query from the partially denoised sequence, retrieves fresh evidence, and conditions the next step on the updated context. Central to S ARDI is a separation between retrieval and generation confidence unique to non-autoregressive decoders: speculative future tokens can inform retrieval long before they are stable enough to commit to the output. S ARDI is training-free and works plug-and-play with any discrete diffusion language model that can produce reasoning traces.
restricted to this single-shot paradigm (Yu et al., 2026; Fang et al., 2026).
Dynamic and agentic retrieval in autoregressive LMs. For autoregressive models, this single-shot limitation has been addressed by a line of work that interleaves retrieval with generation, deriving follow-up queries from the tokens produced so far (Trivedi et al., 2023; Jiang et al., 2023; Jeong et al., 2024; Asai et al., 2024). Most relevant to this work is FLARE (Jiang et al., 2023), which makes retrieval forwardlooking: rather than querying only on already-committed text, FLARE first generates a tentative next sentence and, if that sentence contains low-confidence tokens, uses it as a query to retrieve fresh evidence. Then, it regenerates the sentence under the updated context. While anticipating future tokens does help, autoregressive decoding makes this lookahead fragile: the tentative span is generated left-to-right, so a single early error can compound and produce hallucinated queries that retrieve irrelevant documents (see results in Section 5.3). In contrast, S ARDI predicts all tentative tokens in parallel, so errors do not compound.
We empirically validate two properties that make diffusion language models well suited to retrieval. First, we demonstrate that intermediate denoising states surface bridge entities earlier than autoregressive generation, providing a stronger signal for retrieval. Second, grounding generation in retrieved evidence sharply reduces inter-token mutual information, making the response easier to decode in parallel. Together, this lets S ARDI dominate the quality–latency frontier across five multi-hop QA benchmarks, outperforming current training-free diffusion and autoregressive baselines at substantially lower latency.
More recently, agentic systems generate explicit search queries via planning and self-reflection (Yao et al., 2022; Asai et al., 2024; Xu et al., 2026; Li et al., 2025; Jin et al., 2025). While effective, these approaches typically require specialized training via reinforcement learning, increasing engineering complexity and computational overhead. In contrast, S ARDI is plug-and-play and requires no learned retrieval controller or query generator. To the best of our knowledge, S ARDI is the first retrieval framework that conditions retrieval on intermediate diffusion states and refreshes evidence throughout the denoising trajectory.
2. Related Work Retrieval-augmented generation. Early RAG systems (Lewis et al., 2020; Izacard & Grave, 2021; Karpukhin et al., 2020; Guu et al., 2020) perform single-shot retrieval from the input query and keep the retrieved context fixed throughout generation. This constrains their effectiveness on multi-hop tasks, where the evidence needed for later reasoning steps depends on bridge entities that the question does not name (Yang et al., 2018; Ho et al., 2020). All existing RAG approaches for diffusion language models are 2
Self-Augmenting Retrieval for Diffusion Language Models
Step
External corpus
Denoise
query
Prompt Diffusion Language Model
concatenate
Documents Partially Denoised Response
Retrieve top K docs
Retriever Retriever
Prompt New Documents
tokens with confidence
tokens with confi-
M
dence
M
M
Further Denoised Response
Masked tokens
Figure 2. Overview of self-augmenting retrieval for diffusion language models. At step t, the diffusion LM denoises the partially masked response. Tokens with confidence ci ≥ τq then form a query to refresh the retrieved evidence, while only the more confident tokens (ci ≥ τc ) are committed to the next response xt−1 . Speculative tokens can thus inform retrieval before they are stable enough to commit.
3. Background
Parallel decoding and RAG. Sampling positions independently approximates the joint distribution pθ (x | q, D) by the product of its marginals. This breaks down under strong inter-token dependence. For example, the first name “Albert” makes “Einstein” very likely, whereas independent sampling can produce “Albert Curie”. We argue that RAG sharply reduces such dependence: tokens copied or paraphrased from retrieved passages are largely determined by D, leaving adjacent positions nearly independent. We validate this hypothesis in Section 5.5, confirming that RAG is a regime well-suited to parallel decoding.
We consider open-domain (multi-hop) question answering. Given a question q, the goal is to generate a response x = (x1 , . . . , xL ) ∈ V L of length L over a finite vocabulary V. Retrieval-augmented generation (RAG) retrieves K passages D = {d1 , . . . , dK } from an external corpus C via a query s, and conditions generation on both q and D. 3.1. Autoregressive Language Models Autoregressive (AR) language models factor the response likelihood left-to-right: pθ (x | q, D) =
L Y
4. Self-Augmenting Retrieval for Diffusion pθ (xi | x<i , q, D),
(1)
The diffusion trajectory {xt }Tt=0 exposes a sequence of intermediate states at which retrieval can be revisited. Whereas static RAG issues a single query from the question, SelfAugmenting Retrieval for Diffusion (S ARDI) refines the query as the response takes shape, surfacing new evidence at every step. S ARDI interleaves retrieval with denoising (Figure 2). At each step t, the model predicts a token for every masked position with a confidence score ci = maxv∈V pθ (v | xt , q, Dt ). Depending on ci , a position is used to:
i=1
where x<i = (x1 , . . . , xi−1 ) and θ denotes the model parameters. Decoding proceeds one token at a time, so retrieval can only condition on the committed prefix x<i . 3.2. Discrete Diffusion Language Models Discrete diffusion language models (DLMs) process all positions in parallel through iterative denoising. A DLM defines a sequence of states {xt }Tt=0 , where xT = ([MASK], . . . , [MASK]) is fully masked and x0 ∈ V L is the final output. Generation proceeds from t = T to t = 0: at each step, given a partially masked sequence xt ∈ (V ∪ {[MASK]})L , q, and D, the denoiser predicts a per-position distribution pθ (xt−1 | xt , q, D), i
i ∈ {1, . . . , L},
(i) Retrieve. If ci ≥ τq (the query threshold), the predicted token is added to the retrieval query, refreshing the evidence for the next step. (ii) Commit. If ci ≥ τc (the commit threshold), the token prediction is committed to; the rest of the tokens are remasked.
(2)
and selected masked positions are unmasked to their argmax. 3
Self-Augmenting Retrieval for Diffusion Language Models
Algorithm 1 S ARDI: Self-Augmenting Retrieval for Diffusion Language Models
Separating the two thresholds is the central design choice in S ARDI: because τq ≤ τc , tentative tokens can inform retrieval well before they are reliable enough to commit. The following sections detail query construction (Section 4.1), evidence refresh (Section 4.2), and confidence-based commitment (Section 4.3).
Require: Question q, retriever R, denoiser pθ , query threshold τq , commit threshold τc , context size K Ensure: Generated sequence x 1: x ← ([MASK], . . . , [MASK]) {Fully masked initialization} 2: D ← R(q, K) {Initial retrieval from question only} 3: while not fully unmasked do
4.1. Query Construction As denoising progresses, the emerging response surfaces intermediate entities — names, dates, relations — that the question does not contain but that later-hop retrieval needs. We want to feed these to the retriever as early as possible, even while they are still uncertain. This is safe because retrieval and generation tolerate errors very differently: committing an incorrect token can directly corrupt the output, whereas retrieval is robust to noisy queries. We can therefore set the query threshold τq well below the commit threshold τc , exposing tentative tokens to the retriever long before they are reliable enough to commit. At τq = 0 every position enters the query (maximal lookahead); raising τq restricts it to more confident tokens.
4: 5: 6: 7: 8: 9: 10: 11:
// — Retrieve: refresh evidence — Construct proxy x̃ by filling masks with ci ≥ τq via argmax (Equation (3)) 14: s ← q ∥ Detokenize(x̃) {Form retrieval query} 15: D ← R(s, K) {Retrieve new evidence} 16: end while 17: return x 12: 13:
Concretely, let xt ∈ (V ∪ {[MASK]})L be the current sequence and define ci = maxv∈V pθ (v | xt , q, Dt ) as the model confidence at position i. We form the proxy sequence x̃t by using token predictions with confidence at least τq : t xti ̸= [MASK] x i , t x̃i = arg maxv pθ (v | xt , q, Dt ), ci ≥ τq (3) [MASK], otherwise.
4.3. Confidence-Based Unmasking Because S ARDI refreshes evidence at every step, the order in which tokens are committed directly shapes downstream retrieval. We adopt threshold-based unmasking (Wu et al., 2026), which reveals all positions whose confidence exceeds the commit threshold τc :
The proxy is detokenized to rt = Detokenize(x̃t ) (remaining [MASK] tokens are dropped), and the retrieval query concatenates the question with this intermediate response: st = q ∥ rt .
// — Decode: commit high-confidence tokens — Compute pθ (· | x, q, D) for all masked positions ci ← maxv∈V pθ (v | x, q, D) for each masked position i U ← {i : xi = [MASK] ∧ ci ≥ τc } if U = ∅then U ← arg maxi:xi =[MASK] ci {Ensure progress} end if For all i ∈ U : xi ← arg maxv∈V pθ (v | x, q, D)
(4)
U t = i | xti = [MASK] ∧ ci ≥ τc ,
The fixed question q anchors the query when early predictions are noisy, while the evolving rt progressively specializes retrieval. Unless otherwise specified we use τq = 0; Section 5.4 reports an empirical sweep validating this choice.
(6)
where ci = maxv∈V pθ (v | xt , q, Dt ). Each unmasked position commits to its argmax: xt−1 ← arg max pθ (v | xt , q, Dt ). i
(7)
v∈V
4.2. Evidence Refresh
If no position exceeds τc , the single most confident masked position is unmasked to guarantee progress.
Using the query st constructed above, S ARDI retrieves a fresh set of K passages at each step: D
t−1
t
← R(s , K),
This creates a natural curriculum: high-confidence tokens (often text spans that can already be grounded in the current evidence) commit first and inform the next retrieval, while uncertain spans wait for refined evidence. The commit threshold τc therefore controls both decoding parallelism and the granularity of retrieval refinement, offering a single knob to trade accuracy against throughput (Figure 3). Algorithm 1 gives the full procedure.
(5)
where R is the retriever. The new evidence Dt−1 replaces the previous context entirely and conditions the next denoising step. We use BM25 (Robertson et al., 2009) in our experiments for efficiency, but S ARDI is retriever-agnostic and works with any sparse or dense retriever. 4
Self-Augmenting Retrieval for Diffusion Language Models Table 1. Main results on multi-hop QA benchmarks (Exact Match ×100). ± denotes bootstrap standard deviation; Time is mean wall-clock seconds per example. CofCA and SynthWorlds-RM use counterfactual corpora (made-up facts) to prevent data-leakage. Search-R1 is grayed as it requires additional RL training (see Section 5.3). 2WikiMultihopQA Method
HotpotQA
CofCA
MuSiQue
SynthWorlds-RM
EM
Time
EM
Time
EM
Time
EM
Time
EM
Time
Training-free; autoregressive AR W / RET @ STATIC AR W / RET @10 AR W / RET @1 AR W / RET @ ADAPTIVE AdaptiveRAG (Jeong et al., 2024) ReAct (Yao et al., 2022)
44.5±0.6 53.6±0.6 58.8±0.6 46.6±0.6 34.1±0.6 42.7±0.6
0.74 0.80 1.26 1.47 5.37 2.15
40.3±0.8 45.1±0.8 47.4±0.8 37.3±0.8 37.1±0.8 40.1±0.8
0.66 0.75 1.40 1.34 4.07 2.07
41.0±1.6 39.6±1.6 41.2±1.6 32.4±1.6 38.4±1.6 42.9±1.6
0.76 0.83 1.37 1.49 2.88 2.02
12.7±0.7 17.7±0.8 19.8±0.8 17.0±0.8 14.9±0.7 20.9±0.8
0.69 0.74 1.23 1.42 7.38 2.32
16.5±1.1 19.3±1.1 20.4±1.1 15.7±1.1 13.2±1.0 22.2±1.2
0.81 0.95 2.08 1.71 8.94 3.01
Training-free; diffusion DLM W / RET @ STATIC τc =0.9 DLM W / S ARDI (Ours) τc =0.9 DLM W / S ARDI (Ours) τc =0.95
43.7±0.6 57.8±0.6 59.1±0.6
0.46 0.39 0.56
39.9±0.8 48.5±0.8 48.7±0.8
0.55 0.64 0.90
43.4±1.6 45.3±1.6 44.9±1.6
0.79 0.75 1.09
11.1±0.7 20.5±0.8 20.6±0.8
1.01 0.88 1.19
14.4±1.0 21.1±1.2 21.7±1.2
1.49 1.29 1.78
RL-trained; autoregressive AR (Search-R1) (Jin et al., 2025)
52.4±0.6
3.36
50.3±0.8
2.83
44.4±1.6
3.00
26.4±0.9
3.14
26.9±1.3
3.74
5. Experiments
proved correct for AR models, and we expect the same for DLMs as they mature.
5.1. Experimental Setup
For now, to evaluate our retrieval mechanism fairly, we elicit the model to output reasoning via light supervised finetuning on synthetic chain-of-thought traces from gpt-4omini (App. A). To ensure a fair comparison, we apply identical fine-tuning to DREAM-7B and the AR model Qwen2.57B. Crucially, after fine-tuning the two models reach nearidentical EM under static question-only retrieval (K=7) on 2WikiMultiHopQA: DLM W / RET @ STATIC achieves 43.7% and AR W / RET @ STATIC achieves 44.5% (Table 1). To ensure the gains do not come from memorization, we include two counterfactual benchmarks. CofCA and SynthWorlds-RM use corpora with made-up facts that are absent from both pre-training and our fine-tuning data. Finally, we emphasize that S ARDI is not tied to the specific reasoning format induced by the synthetic gpt-4o-mini traces; any reasoning style that includes intermediate entities will work.
Datasets and benchmarks. We evaluate S ARDI on five multi-hop QA benchmarks: 2WikiMultiHopQA (Ho et al., 2020), HotpotQA (Yang et al., 2018), MuSiQue (Trivedi et al., 2022), CofCA (Wu et al., 2025), and SynthWorldsRM (Gu et al., 2025), reporting Exact Match (EM) following Jin et al. (2025). To measure per-query retrieval recall, the corpus must be at the same passage granularity as each benchmark’s gold annotations: HotpotQA provides a full Wikipedia corpus at this granularity (Yang et al., 2018), and for the other four we follow prior work (Prabhu & Anand, 2024; Yu et al., 2026) and concatenate all gold and distractor passages into a single corpus. Unless stated otherwise we use the sparse retrieval method BM25 (Robertson et al., 2009), retrieving K=7 passages per iteration, and additionally evaluate the E5-base-v2 dense retriever in Section 5.6. Latency is measured in wall-clock time unbatched on a single NVIDIA B200 GPU.
5.2. Baselines RAG and reasoning with DLMs. S ARDI relies on the existence of a reasoning trace in order to surface bridge entities. In our analysis, off-the-shelf instruction-tuned diffusion LMs such as DREAM-7B almost never produce reasoning traces in a RAG setting, even when using fewshot prompting (Table 6, App. A). Trivedi et al. (2023) made an analogous observation for AR models three years ago: “IRCoT relies on the base LM to have a zero or few-shot CoT-generation ability [. . . ] not as common for small LMs (under 20B) [. . . ] smaller LMs will likely increasingly acquire such ability.” (Trivedi et al., 2023) Their prediction
We compare S ARDI against retrieval-augmented baselines from both AR and diffusion paradigms. AR baselines use Qwen2.5-7B and diffusion baselines use DREAM-7B, with the same light fine-tuning (Section 5.1) applied to both backbones. The one exception is Search-R1, which we evaluate from the authors’ released RL-trained checkpoint (App. C). All methods share identical prompts and evidence formatting (App. B), isolating the retrieval mechanism. • AR W / RET @ STATIC: Single retrieval from q; D is fixed throughout generation. 5
Self-Augmenting Retrieval for Diffusion Language Models
2WikiMultihopQA
Exact match (%)
60
ret@1
55
0.95 0.9 0.8
ret@10
50
ret@adaptive ret@static
45
0.7
40 0
1
2
3 4 5 Questions per second
6
0.95 0.9
45.0 42.5
ret@static ret@1 ret@10
40.0 37.5
30.0
7
0.8
SARDI ( c [0.60, 0.95]) AR baselines ReAct Search-R1 (RL trained)
0.7
35.0 32.5
0.6
35
CofCA
47.5
SARDI ( c [0.60, 0.95]) AR baselines ReAct Search-R1 (RL trained)
Exact match (%)
65
ret@adaptive 0
1
0.6 2 3 Questions per second
4
5
Figure 3. Accuracy vs. throughput trade-off on 2WikiMultiHopQA and CofCA, traced by adjusting the commit threshold τc . Similar pareto patterns hold on HotpotQA and MuSiQue (Figure 6).
• AR W / RET @N: Retrieval every N ∈ {1, 10} AR tokens.
dominates the quality–latency frontier, running up to 8× faster than AR iterative-retrieval baselines at comparable or better accuracy. We analyze the sources of these gains in the following sections.
• AR W / RET @ ADAPTIVE (Jiang et al., 2023) (FLARE): Confidence-triggered retrieval using low-confidence spans as both triggers and queries.
Comparison to trained search agents. Search-R1 (Jin et al., 2025) is an RL-trained search agent that learns to emit explicit search queries during generation. We include it as a strong reference point but emphasize that S ARDI and Search-R1 occupy different design points: Search-R1 invests in RL-based query generation to maximize accuracy at inference cost, whereas S ARDI is a training-free, plugand-play retrieval mechanism that requires no reward design or policy optimization. Across the five benchmarks S ARDI is the strongest training-free method while running 3–8× faster than Search-R1 at comparable accuracy on 2WikiMultiHopQA, HotpotQA, and CofCA (Table 1). The two approaches are complementary rather than competing, and combining diffusion trajectories with light RL-based query supervision is a natural direction for future work.
• AdaptiveRAG (Jeong et al., 2024): AR LM equipped with a query-complexity router that triggers either singlestep or iterative retrieval. • ReAct (Yao et al., 2022): A training-free agentic loop with Retrieve[query] and Finish[answer] actions. • DLM W / RET @ STATIC: The diffusion counterpart of AR W / RET @ STATIC; one retrieval, fixed D. • DLM W / S ARDI (this work): DREAM-7B equipped with S ARDI. • AR (Search-R1) (Jin et al., 2025): An RL-trained search agent that learns to emit explicit search queries during generation. Included as a strong reference point; discussed in Section 5.3.
5.4. Diffusion Trajectories as Lookahead for Retrieval A central hypothesis of this work is that the uncommitted, low-confidence tokens in the diffusion trajectory are a useful lookahead signal for retrieval. Because a DLM denoises the whole sequence at once, it holds tentative predictions for every position long before they are committed. S ARDI feeds these speculative tokens to the retriever, surfacing later-hop evidence early on in the generation process. The following experiments test this hypothesis.
5.3. S ARDI Is Faster and More Accurate Accuracy and throughput. S ARDI improves substantially over static diffusion retrieval on every benchmark (Table 1): on 2WikiMultiHopQA it raises EM from 44 to 59, with similar gains on HotpotQA (40 → 49), CofCA (43 → 45), and MuSiQue (11 → 21). It also matches or beats every training-free AR baseline, and does so at much lower latency. We can trade off throughput against accuracy via the commit confidence threshold τc : a higher τc commits fewer tokens per denoising step and triggers more retrieval rounds (higher accuracy, lower throughput), while a lower τc does the reverse. Sweeping τc (Figure 3), S ARDI
More lookahead is better. The query threshold τq sets how confident the model must be in a token before exposing it to the retriever (Equation (3)): τq =0 exposes every position, while τq =τc restricts the query to tokens that are 6
Self-Augmenting Retrieval for Diffusion Language Models
2WikiMultihopQA c = 0.9 c = 0.95
58
Per-query gold recall (%)
Exact match (\%)
60
56 54 0.0
0.2 0.4 0.6 0.8 Query confidence threshold q
100
Per-query gold-document recall Oracle (Question+GT reasoning) Question-only (DLM ret@static) SARDI c = 0.9, q = 0 SARDI c = 0.9, q = 0.9 AR ret@1 AR ret@10
90 80 70 60 0%
Figure 4. Sweep of the query threshold τq on 2WikiMultiHopQA (BM25, K=7), at the two commit thresholds τc ∈ {0.9, 0.95}. Similar trend holds on HotpotQA, MuSiQue, and SynthWorldsRM (Table 11).
25% 50% 75% 100% Fraction of generation completed
Figure 5. Per-query gold-document recall (fraction of gold passages in the current retrieved set) as generation progresses (K=7, BM25, 2WikiMQA). We restrict to methods that retrieve at every (or every-N th ) token; agentic methods and FLARE issue specific targeted queries by design, so their per-query recall is low even though cumulative recall is comparable (Appendix D, Table 8).
confident enough to be committed. Sweeping τq on 2WikiMultiHopQA (Figure 4), EM peaks at the most aggressive setting (τq ≈ 0) and falls as the query becomes more conservative; similar trends hold on HotpotQA, MuSiQue, and SynthWorlds-RM (Table 11). This directly supports our hypothesis that speculative future tokens can inform retrieval long before they are stable enough to commit to the output.
Table 2. Question-type breakdown on 2Wiki (EM %). SARDI benefits most from question types requiring multi-step reasoning. Method
1-hop comp.
DLM W / RET @ STATIC 85.6 DLM W / S ARDI(Ours) τc =0.95 84.6-1.0
Multi-hop bridge-comp. compos. 65.3 71.7+6.4
16.6 45.3+28.7
infer. 14.0 37.5+23.5
5.5. RAG Grounding Promotes Parallel Decoding Lookahead surfaces evidence earlier. To further analyze the effect of lookahead, Figure 5 plots per-query document recall (the fraction of gold passages in the current retrieved set) against generation progress. Two horizontal lines provide baselines: the recall of a question-only query (what static retrieval sees) and the recall of a query built from the ground-truth reasoning trace (roughly the best our framework could reach). As expected, early in generation S ARDI sits well above the AR baselines (a +19-point recall gain at 25% of generation) and closes much of the gap toward the oracle line. As expected, the more aggressive the lookahead (lower τq ), the higher the early recall. By the final step S ARDI converges to AR ret@1, so its main advantage is that strong evidence arrives early, letting it make more progress in the opening steps.
We argued in Section 1 that retrieved evidence reduces intertoken dependence, making RAG well-suited to parallel decoding. We test this by measuring the conditional mutual information (CMI) between adjacent tokens in ground-truth reasoning traces on 2WikiMultiHopQA: CMI(xi ; xi+1 | D) = Exi KL p(xi+1 | xi , D) ∥ p(xi+1 | D) . We approximate the expectation over xi by its top-7 values and vary the amount of retrieved evidence in D. Intuitively, CMI is high when several completions are plausible and the model has not settled on one. Multi-token names are the clearest case: until the model commits to which entity it is naming, the tokens must agree to stay coherent (Albert Einstein or Isaac Newton, not Albert Newton), so fixing one sharply shifts the distribution over the other. Strong evidence removes this coupling: the model copies the entity straight from the retrieved passage instead of coordinating across positions.
Gains concentrate on multi-hop reasoning. Table 2 breaks 2WikiMultiHopQA down by question type: S ARDI yields over 2.5× EM gains on inference and compositional questions (which require identifying bridge entities through multi-step reasoning) and leaves single-hop (comparison) questions (where static retrieval already suffices) unchanged.
The results in Table 3 confirm this. With all gold documents present, entity pairs show very low dependence (CMI=0.060), i.e., they can be decoded in parallel. As gold 7
Self-Augmenting Retrieval for Diffusion Language Models Table 3. Conditional mutual information (CMI) between adjacent tokens in DREAM-7B reasoning traces (2WikiMultiHopQA), as gold documents are progressively removed from the context. We separately report entity and non-entity pairs. All gold
Gold−1
Gold−2
No gold
0.060 0.136
0.219 0.214
0.396 0.247
0.588 0.264
CMI (entity) CMI (non-entity)
6. Limitations While S ARDI is training-free in principle, the tested diffusion language models do not yet reliably produce reasoningtraces through prompting alone. This is analogous to what Trivedi et al. (2023) observed for early medium-scale autoregressive models, which has since been resolved as AR models matured; we expect the same to happen for DLMs, removing the need for the fine-tuning step. Additionally, our method refreshes retrieval at every denoising step, which can incur document encoding overhead under naive implementations. Because a large portion of the retrieved documents persists between consecutive steps, integrating Fast-dLLM-style block KV caching is a natural direction for future work. Extending self-augmenting retrieval to latent diffusion language models is another promising direction.
documents are removed, entity-pair CMI rises nearly 10× to 0.588, while non-entity pairs increase from 0.136 to 0.264. Grounding thus removes the inter-token dependence most harmful to parallel decoding: tightly coupled entity spans. This confirms that the RAG setting is well suited to parallel decoding. 5.6. Ablations
7. Conclusion
Retriever choice. S ARDI is retriever-agnostic. Replacing BM25 with the E5-base-v2 dense retriever, S ARDI still outperforms the strongest training-free AR baseline (AR W / RET @1) and stays competitive with the RL-trained SearchR1 (Table 4). Also see more detailed recall results in Appendix D. Thus, the observed behavior does not depend on the lexical retriever.
We introduced Self-Augmenting Retrieval (S ARDI), a training-free, dynamic retrieval framework that exploits two structural opportunities that diffusion language models open up for retrieval-augmented generation. First, the denoising trajectory exposes tentative predictions for the entire response at every step, surfacing salient entities early and turning them into a lookahead signal for retrieval. Second, we have shown that grounding generation in retrieved evidence sharply reduces inter-token dependence, making RAG a regime especially well suited to parallel decoding. S ARDI leverages both: it interleaves retrieval with denoising, refining the query as the response takes shape. On five multi-hop QA benchmarks, S ARDI 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 we expect S ARDI to strengthen as diffusion language models mature.
Table 4. S ARDI and strongest baselines with a dense retriever (E5-base-v2). EM (×100), K=7. Method
2Wiki Hotpot CofCA MuSiQue SynthW.
Search-R1 AR W / RET @1 DLM W / S ARDI (Ours)
41 47 50
49 47 52
44 41 45
26 20 22
21 22 23
Refresh schedule. S ARDI refreshes retrieval at every denoising step, which is wasteful when retrieval is expensive (rerankers, large corpora, tight serving budgets). Table 5 varies the refresh frequency: refreshing every 2 steps costs only 1–2 EM. Moreover, we show in Appendix E that the retrieved set is relatively stable: on average 83–90% of documents persist between consecutive steps. Both findings suggest that the per-step retrieval cost can be substantially amortized – a direction we discuss in the Limitations section.
Acknowledgements We thank Alexander Kreibich and Lucas He for productive discussions that helped motivate this work. JL is supported by a Google PhD Fellowship, LZ by a LinkedIn PhD Fellowship, and DG by an Empire AI Postdoctoral Fellowship. We are grateful to NVIDIA for access to the DGX Station compute platform through their Early Access Program. This work was further supported by the National Science Foundation (NSF) under grants OAC-2118310 and 2530143, through the AI Research Institutes program (Award No. DMR-2433348), by funding from NewYork-Presbyterian for the NYP-Cornell Cardiovascular AI Collaboration, and by arXiv. Finally, we thank the anonymous reviewers, whose feedback strengthened the paper by motivating the analyses of early-generation recall and mutual information.
Table 5. S ARDI accuracy under different retrieval refresh frequencies. EM (×100); subscripts give the change from refreshing every step.
Refresh
2Wiki
HotpotQA
CofCA
MuSiQue
every step every 2 every 4
58 56−2 52−6
48 47−1 46−2
45 45 45
20 20 18−2
8
Self-Augmenting Retrieval for Diffusion Language Models
Impact Statement
Izacard, G. and Grave, E. Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering. In Proceedings of the 16th conference of the european chapter of the association for computational linguistics: main volume, pp. 874–880, 2021.
This work studies retrieval-augmented generation for diffusion language models. Improving evidence grounding can reduce hallucinations in knowledge-intensive settings, but the same retrieval mechanisms could produce more convincing misinformation if paired with untrusted corpora. We recommend deploying S ARDI 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.
Jeong, S., Baek, J., Cho, S., Hwang, S. J., and Park, J. Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL), 2024.
References Jiang, Z., Xu, F. F., Gao, L., Sun, Z., Liu, Q., DwivediYu, J., Yang, Y., Callan, J., and Neubig, G. Active retrieval augmented generation. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 7969–7992, 2023.
Asai, A., Wu, Z., Wang, Y., Sil, A., and Hajishirzi, H. Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Jin, B., Zeng, H., Yue, Z., Yoon, J., Arik, S. O., Wang, D., Zamani, H., and Han, J. Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning. In Second Conference on Language Modeling, 2025.
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and van den Berg, R. Structured Denoising Diffusion Models in Discrete State-Spaces. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17981–17993, 2021.
Karpukhin, V., Oguz, B., Min, S., Lewis, P. S., Wu, L., Edunov, S., Chen, D., and Yih, W.-t. Dense Passage Retrieval for Open-Domain Question Answering. In EMNLP (1), pp. 6769–6781, 2020.
Fang, Y., Yu, H.-T., Joho, H., and Fujita, S. RetrievalAugmented Diffusion Language Model for Generative Commonsense Reasoning. In Jung, H., Wang, T., Toyoda, M., Kwon, H.-Y., and Lee, J.-w. (eds.), Database Systems for Advanced Applications, pp. 341–356, Singapore, 2026. Springer Nature Singapore. ISBN 978-981-920369-7.
Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Küttler, H., Lewis, M., Yih, W.-t., Rocktäschel, T., et al. Retrieval-Augmented Generation for KnowledgeIntensive NLP Tasks. Advances in Neural Information Processing Systems, 33:9459–9474, 2020.
Gu, K., Bhat, A., Merrill, M. A., West, R., Liu, X., McDuff, D., and Althoff, T. SynthWorlds: Controlled Parallel Worlds for Disentangling Reasoning and Knowledge in Language Models. arXiv preprint arXiv:2510.24427, 2025.
Li, X., Dong, G., Jin, J., Zhang, Y., Zhou, Y., Zhu, Y., Zhang, P., and Dou, Z. Search-o1: Agentic SearchEnhanced Large Reasoning Models. In Christodoulopoulos, C., Chakraborty, T., Rose, C., and Peng, V. (eds.), Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, EMNLP 2025, Suzhou, China, November 4-9, 2025, pp. 5420–5438. Association for Computational Linguistics, 2025. doi: 10.18653/V1/2025.EMNLP-MAIN.276.
Guu, K., Lee, K., Tung, Z., Pasupat, P., and Chang, M. Retrieval Augmented Language Model Pre-Training. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, Proceedings of Machine Learning Research, pp. 3929–3938. PMLR, 2020. Ho, X., Nguyen, A. D., Sugawara, S., and Aizawa, A. Constructing A multi-hop QA dataset for comprehensive evaluation of reasoning steps. In Scott, D., Bel, N., and Zong, C. (eds.), Proceedings of the 28th International Conference on Computational Linguistics, COLING 2020, Barcelona, Spain (Online), December 8-13, 2020, pp. 6609–6625. International Committee on Computational Linguistics, 2020. doi: 10.18653/V1/2020. COLING-MAIN.580.
Li, X. L., Thickstun, J., Gulrajani, I., Liang, P., and Hashimoto, T. B. Diffusion-LM Improves Controllable Text Generation. In Advances in Neural Information Processing Systems 35, 2022. Lovelace, J., Kishore, V., Wan, C., Shekhtman, E., and Weinberger, K. Q. Latent Diffusion for Language Generation. Conference on Neural Information Processing Systems (NeurIPS), 2023. 9
Self-Augmenting Retrieval for Diffusion Language Models
Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., ZHOU, J., Lin, Y., Wen, J.-R., and Li, C. Large Language Diffusion Models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2026.
Xu, Z., Liu, Z., Yan, Y., Wang, S., Yu, S., Zeng, Z., Xiao, C., Liu, Z., Yu, G., and Xiong, C. ThinkNote: Enhancing knowledge integration and utilization of large language models via constructivist cognition modeling. In Demberg, V., Inui, K., and Marquez, L. (eds.), Findings of the Association for Computational Linguistics: EACL 2026, pp. 211–229, Rabat, Morocco, March 2026. Association for Computational Linguistics. ISBN 979-8-89176-386-9. doi: 10.18653/v1/2026.findings-eacl.12.
Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently Scaling Transformer Inference. In Song, D., Carbin, M., and Chen, T. (eds.), Proceedings of the Sixth Conference on Machine Learning and Systems, MLSys 2023, Miami, FL, USA, June 4-8, 2023. mlsys.org, 2023.
Yang, Z., Qi, P., Zhang, S., Bengio, Y., Cohen, W., Salakhutdinov, R., and Manning, C. D. HotpotQA: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 conference on empirical methods in natural language processing, pp. 2369–2380, 2018.
Prabhu, V. V. D. and Anand, A. DEXTER: A Benchmark for open-domain Complex Question Answering using LLMs. arXiv preprint arXiv:2406.17158, 2024. Robertson, S., Zaragoza, H., et al. The Probabilistic Relevance Framework: BM25 and Beyond. Foundations and trends® in information retrieval, 3(4):333–389, 2009.
Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K. R., and Cao, Y. ReAct: Synergizing Reasoning and Acting in Language Models. In The Eleventh International Conference on Learning Representations, 2022.
Snell, C. V., Lee, J., Xu, K., and Kumar, A. Scaling LLM Test-Time Compute Optimally Can be More Effective than Scaling Parameters for Reasoning. In The Thirteenth International Conference on Learning Representations, 2025.
Ye, J., Xie, Z., Zheng, L., Gao, J., Wu, Z., Jiang, X., Li, Z., and Kong, L. Dream 7b: Diffusion large language models. arXiv preprint arXiv:2508.15487, 2025. Yu, C., Wang, J., Li, Y., Chang, H., Lan, G., Sun, Q., Li, J., Li, J., and Zhang, Z. Unlocking the Potentials of Retrieval-Augmented Generation for Diffusion Language Models. arXiv preprint arXiv:2601.11342, 2026.
Trivedi, H., Balasubramanian, N., Khot, T., and Sabharwal, A. MuSiQue: Multihop Questions via Single-hop Question Composition. Transactions of the Association for Computational Linguistics, 10:539–554, 2022. Trivedi, H., Balasubramanian, N., Khot, T., and Sabharwal, A. Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions. In Rogers, A., Boyd-Graber, J. L., and Okazaki, N. (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023, Toronto, Canada, July 9-14, 2023, pp. 10014–10037. Association for Computational Linguistics, 2023. doi: 10.18653/V1/2023.ACL-LONG.557. Wang, L., Yang, N., Huang, X., Jiao, B., Yang, L., Jiang, D., Majumder, R., and Wei, F. Text Embeddings by WeaklySupervised Contrastive Pre-training. arXiv preprint arXiv:2212.03533, 2022. Wu, C., Zhang, H., Xue, S., Liu, Z., Diao, S., Zhu, L., Luo, P., Han, S., and Xie, E. Fast-dLLM: Training-free acceleration of diffusion LLM by enabling KV cache and parallel decoding. In The Fourteenth International Conference on Learning Representations, 2026. Wu, J., Yang, L., Wang, Z., Okumura, M., and Zhang, Y. CofCA: A STEP-WISE counterfactual multi-hop QA benchmark. In The Thirteenth International Conference on Learning Representations, 2025. 10
Self-Augmenting Retrieval for Diffusion Language Models
A. Fine-Tuning Analysis Table 6. Output-mode breakdown for base and fine-tuned DREAM-7B and Qwen2.5-7B on 2WikiMultiHopQA (static K=7 RAG, n=500).
DREAM-7B Qwen2.5-7B DREAM-7B (SFT) Qwen2.5-7B (SFT) Correct (with reasoning trace) Correct (no reasoning trace) Wrong
1% 27% 72%
32% 0% 68%
44% 0% 56%
44% 0% 56%
To motivate the fine-tuning stage and isolate its effect from the retrieval mechanism, we present an output analysis of base and fine-tuned models on 500 examples from 2WikiMultiHopQA under static K=7 retrieval. Table 6 reports the proportion of outputs in three categories: correct with a structured reasoning trace, correct without a reasoning trace, and wrong (which folds in format errors such as missing the ### answer marker). Two observations are worth highlighting: 1. Base DREAM-7B cannot produce structured reasoning from prompting alone. Only 1% of base DREAM-7B outputs include any reasoning trace at all—most simply emit a guess. The same-sized AR Qwen2.5-7B reasons in 32% of cases. S ARDI relies on intermediate states surfacing salient entities, so this failure mode makes S ARDI inapplicable without fine-tuning. 2. Fine-tuning equalizes capability. On this n=500 subset both models land at 44% correct-with-reasoning; on the full test set they sit within 1 EM of each other (DLM W / RET @ STATIC=43.7%, AR W / RET @ STATIC=44.5%; Table 1). Any subsequent gap between S ARDI and matched-fine-tuning AR baselines is therefore attributable to the retrieval mechanism rather than the underlying model. This is analogous to what Trivedi et al. (2023) observed for early small AR models: “IRCoT relies on the base LM to have a zero or few-shot CoT-generation ability. While this is commonly available in large LMs (over 100B), it’s not as common for small LMs (under 20B). . . smaller LMs will likely increasingly acquire such ability.” That prediction proved correct for AR models; we expect the same for DLMs. A.1. Fine-Tuning Data We construct supervised fine-tuning data from the training splits of 2WikiMultiHopQA (Ho et al., 2020) and HotpotQA (Yang et al., 2018), augmented with synthetically generated chain-of-thought reasoning traces. Synthetic reasoning traces. For each training example, we prompt gpt-4o-mini with the question q, gold answer a, and supporting documents to generate a step-by-step reasoning trace that (i) identifies intermediate bridge entities, (ii) references specific facts from the supporting documents, and (iii) concludes with the final answer in the format ### [answer]. Example.
The following illustrates a training example from 2WikiMultiHopQA:
Question: Where was the director of film Ronnie Rocket born? Reasoning (synthetically generated): Step 1: Ronnie Rocket is directed by David Lynch. Step 2: David Lynch was born in Missoula, Montana. ### Montana Training context. During training, we present each training example together with its gold documents, reinforcing Dream-7B to generate reasoning traces when solving RAG problems. 11
Self-Augmenting Retrieval for Diffusion Language Models
A.2. Fine-Tuning Configuration We fine-tune DREAM-7B using fully sharded data parallel (FSDP) on 2× NVIDIA B200 GPUs. Table 7 summarizes the hyperparameters. To ensure a fair comparison, we apply the same configuration and training data to all autoregressive baselines (Qwen2.5-7B), isolating differences attributable to the generation paradigm. Table 7. Fine-tuning hyperparameters.
Hyperparameter
Value
Base model Training epochs Learning rate Global batch size Micro batch size per GPU Maximum sequence length Optimizer Hardware
DREAM-7B 3 2 × 10−6 256 16 2048 tokens AdamW 2× NVIDIA B200 GPUs
B. Prompt Templates B.1. RAG Input Format All models receive input in the following format: Use ONLY the provided facts to answer the question. Think step-by-step, then provide the final answer after the "###" marker. Question: {question} Facts: {facts} where {question} is the input question and {facts} contains the concatenated retrieved passages. B.2. Expected Output Format Models are trained to produce outputs in the following format: Step 1: [First reasoning step grounded in documents] Step 2: [Second reasoning step] ... ### [final answer] 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 (###) for evaluation.
C. Additional Experimental Details Evaluation protocol. We extract the final answer by parsing text following the ### marker. Exact Match (EM) is computed after normalizing both predicted and gold answers and checking for string equality. Retrieval configuration. We use BM25 (Robertson et al., 2009) for all experiments unless otherwise specified. For the dense-retriever results in Section 5.6 we use E5-base-v2 (Wang et al., 2022). For S ARDI, retrieval is performed at every denoising iteration using the concatenation of the original question and the current intermediate response as the query, with K=7 passages per iteration unless otherwise specified. 12
Self-Augmenting Retrieval for Diffusion Language Models
Search-R1 setup. The four AR baselines (AR W / RET @ STATIC, AR W / RET @10, AR W / RET @1, AR W / RET @ ADAPTIVE) and the training-free agentic baselines (AdaptiveRAG, ReAct) all share the same Qwen2.5-7B backbone, supervised-fine-tuned on the combined training sets of 2WikiMultiHopQA and HotpotQA (Appendix A); any accuracy gap with S ARDI therefore isolates the retrieval mechanism. Search-R1 is the only baseline that does not share this backbone: we evaluate the authors’ publicly released checkpoint, which they trained via PPO on HotpotQA using a multi-stage reward (Jin et al., 2025). Reproducing this RL training pipeline was not feasible within our compute budget, so we accept this additional confound and treat Search-R1 as an RL-trained reference point rather than a controlled comparison (hence the gray styling in Table 1).
D. Per-Method Recall: Full Tables Figure 5 in the main text reports gold-document recall over generation progress for SARDI and the AR retrieval baselines that fire at every (or every-N th ) token. For completeness, this section reports the full per-method numbers, including the agentic methods (AdaptiveRAG, ReAct, Search-R1) and AR speculative-lookahead (FLARE). We report two complementary recall metrics: • Per-query final recall (Table 8): fraction of gold documents present in the retrieved set D at the final retrieval step, averaged over questions. This is the metric most directly tied to answer accuracy on a single forward pass. • Cumulative recall (Table 9): fraction of gold documents that have appeared in D at any point during generation, i.e., the union over all retrieval steps. This view is fairer to methods that issue few but targeted queries. Caveat: cumulative recall depends on retrieval volume. A method that issues many queries (e.g., S ARDI, AR ret@1) trivially accumulates more unique documents than one that issues few targeted queries (e.g., Search-R1), and cumulative recall therefore is not directly comparable across methods without controlling for total documents touched. Table 9 accordingly reports the average number of unique documents retrieved per question as a subscript next to each cumulativerecall value, so the reader can normalize. Table 8. Per-query gold-document recall at the final retrieval step (×100). Methods that issue specific targeted queries (AdaptiveRAG, ReAct, AR ret@adaptive/FLARE, Search-R1) retrieve few documents per query, so per-query recall is low by design even when their cumulative coverage (Table 9) is comparable. S ARDI matches AR ret@1 across K and retrievers. 2Wiki
MuSiQue
BM25
Dense
BM25
Dense
Method
K=3 K=7 K=10 K=15 K=3 K=7 K=10 K=15 K=3 K=7 K=10 K=15 K=3 K=7 K=10 K=15
AdaptiveRAG ReAct AR ret@adaptive (FLARE) AR ret@10 AR ret@1 Search-R1 S ARDI τc =0.9 S ARDI τc =0.95
5 38 22 65 66 42 66 67
5 48 27 78 80 53 81 81
4 51 29 82 85 57 83 84
3 54 32 86 88 62 86 88
– 27 14 54 54 28 55 56
5 31 17 61 61 31 63 63
4 31 18 62 63 33 62 64
3 33 19 63 64 35 63 65
9 26 14 42 43 29 43 44
10 34 20 53 54 38 55 56
9 37 22 56 58 42 56 59
7 40 25 61 62 45 60 64
– 27 14 38 38 28 39 40
9 31 20 47 48 35 50 50
– 32 21 51 51 37 53 53
– 35 22 54 54 41 55 56
E. Document Persistence Across Refresh Steps Table 5 relies on the observation that retrieved documents change only gradually between consecutive denoising steps, so document-level KV-cache reuse is feasible. We measure the average overlap between consecutive retrieved sets as |Dt ∩ Dt−1 |/|Dt−1 |: Dataset
% docs retained per step 88% 83% 89% 84% 90%
2WikiMultiHopQA HotpotQA CofCA MuSiQue SynthWorlds-RM 13
Self-Augmenting Retrieval for Diffusion Language Models Table 9. Cumulative gold-document recall (×100, main number) and average number of unique documents retrieved per question (subscript ). Cumulative recall measures the union of all gold documents surfaced at any point during generation, but is not directly comparable across methods that retrieve different total document counts: a method that issues many queries (e.g., S ARDI, AR ret@1) trivially accumulates more docs than one that issues few (e.g., Search-R1). The subscripts let the reader normalize. Even controlling for total docs, S ARDI remains competitive with the strongest baselines. 2Wiki
MuSiQue
BM25
Dense
BM25
Dense
Method
K=3 K=7 K=10 K=15 K=3 K=7 K=10 K=15 K=3 K=7 K=10 K=15 K=3 K=7 K=10 K=15
AdaptiveRAG ReAct AR ret@adaptive (FLARE) AR ret@10 AR ret@1 Search-R1 S ARDI τc =0.9 S ARDI τc =0.95
716 737 738 716 737 746 705 715
7913 8115 8220 8315 8520 8516 8314 8314
7714 8221 8429 8622 8929 8823 8520 –
6215 8331 8743 8934 9146 9034 8831 9033
– 577 639 616 649 607 616 637
– 5715 6622 6615 6823 6317 6817 6818
– 5721 6732 6722 6933 6325 6624 7026
– 5831 6849 6833 7050 6337 6737 7140
487 506 519 496 528 577 477 527
5613 6015 6222 6016 6321 6516 6319 6420
5514 6322 6531 6423 6730 6823 6128 6830
4815 6533 6948 6834 7145 – 6544 7346
– 547 529 466 509 597 477 498
– 6016 6022 5515 5922 6517 5819 5920
– 6222 6432 5822 6231 6723 5928 6229
– 6433 6548 6232 6546 7034 6142 6644
Since 83–90% of documents persist between consecutive steps, only a small fraction needs recomputation, making document-level KV caching similar to Fast-dLLM (Wu et al., 2026) an attractive direction for combining S ARDI with diffusion-decoding-acceleration techniques.
F. Accuracy vs. Throughput: Full Datasets Figure 3 in the main text shows the Pareto frontier on 2WikiMultiHopQA and CofCA. Figure 6 reports the same plot across all four benchmarks. S ARDI dominates the autoregressive iterative-retrieval baselines on every dataset and matches Search-R1 at 3–8× lower latency.
2WikiMultihopQA
Exact match (%)
60
ret@1
55
0.95 0.9 0.8
ret@10
50
ret@adaptive ret@static
45
HotpotQA
SARDI ( c [0.60, 0.95]) AR baselines ReAct Search-R1 (RL trained)
Exact match (%)
65
0.7
40 0.6
35 0
1
2
3 4 5 Questions per second
6
7
0
MuSiQue
28 SARDI ( c [0.60, 0.95]) AR baselines 26 ReAct 24 Search-R1 (RL trained) 22 0.95 0.9 20 0.8 ret@1 18 0.7 ret@10 16 ret@adaptive 14 0.6 12 ret@static 0.0 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 Questions per second
1
0.8
SARDI ( c [0.60, 0.95]) AR baselines ReAct Search-R1 (RL trained)
0.7 0.6
2 3 Questions per second
0.95 0.9
45.0 42.5
ret@static ret@1 ret@10
40.0 37.5
0.8
0.7
32.5
ret@adaptive 0
1
0.6 2 3 Questions per second
4
Figure 6. Accuracy vs. throughput trade-off across all four benchmarks, companion to Figure 3.
14
5
SARDI ( c [0.60, 0.95]) AR baselines ReAct Search-R1 (RL trained)
35.0 30.0
4
CofCA
47.5
Exact match (%)
Exact match (%)
52.5 50.0 0.95 0.9 47.5 ret@1 45.0 ret@10 42.5 40.0 ret@static 37.5 ret@adaptive 35.0
5
Self-Augmenting Retrieval for Diffusion Language Models Table 10. Threshold-based unmasking matches fixed-step accuracy at 2–3× the speed. EM (×100) and time per example (s). 2Wiki EM / T
HotpotQA EM / T
CofCA EM / T
MuSiQue EM / T
50 fixed steps 59 / 1.33 47 / 1.46 43 / 1.48 20 / 1.42 τc =0.9 58−1 / 0.39−0.94 48+1 / 0.64−0.82 45+2 / 0.75−0.73 20 / 0.88−0.54 τc =0.95 59 / 0.56−0.77 49+2 / 0.90−0.56 45+2 / 1.09−0.39 21+1 / 1.19−0.23
Table 11. Query-threshold sweep, full numbers. EM (%) on the four datasets with τq sweeps, at commit thresholds τc ∈ {0.9, 0.95}.
Dataset
τc
τq =0 τq =0.1 τq =0.2 τq =0.3 τq =0.4 τq =0.5 τq =0.6 τq =0.7 τq =0.8 τq =0.9
2WikiMultihopQA
0.9 0.95
57.8 59.1
58.1 59.4
57.4 58.7
57.3 58.4
56.6 57.6
55.6 57.1
54.7 56.2
54.5 55.6
54.0 54.8
53.2 54.0
HotpotQA
0.9 0.95
48.3 48.7
48.1 48.4
47.6 47.8
46.4 47.0
45.9 46.3
46.0 45.6
44.8 45.3
44.7 44.9
44.4 44.4
44.1 44.1
MuSiQue
0.9 0.95
20.5 20.6
20.6 20.6
19.7 19.5
19.2 18.9
18.2 17.9
16.7 17.1
16.3 16.3
15.7 16.3
15.7 15.6
15.2 14.7
SynthWorlds
0.9 0.95
21.1 21.7
21.4 21.5
21.1 21.3
19.9 19.8
17.9 19.5
17.5 18.2
17.9 17.5
17.0 17.5
16.8 17.0
17.2 16.9
G. Threshold-Based vs. Fixed-Step Decoding Table 10 compares fixed-step decoding (T =50) against threshold-based unmasking at the two commit thresholds used in the main paper, τc =0.9 and τc =0.95. Threshold decoding matches fixed-step accuracy at up to 3× the speed.
H. Query-Threshold Sweep: Full Numbers Figure 4 in the main text plots the τq sweep on 2WikiMultiHopQA. Table 11 reports full EM numbers across all four benchmarks for which τq sweeps were run, at both commit thresholds τc ∈ {0.9, 0.95} and ten coarse τq values spanning the spectrum from aggressive lookahead (τq =0, the default S ARDI setting) to the conservative endpoint (τq =0.9, querying only on near-committed tokens). The pattern observed on 2WikiMultiHopQA reproduces on HotpotQA, MuSiQue, and SynthWorlds-RM: peak accuracy lands at τq ∈ [0, 0.1] on every (dataset, τc ) cell, and EM generally declines as the query threshold tightens. The conservative endpoint loses 4–6 EM relative to the default, consistent with our hypothesis that retrieval benefits from speculative future tokens that are not yet stable enough to commit.
I. Qualitative Example: Threshold-Based Unmasking A central hypothesis of this work is that RAG exhibits a structure uniquely amenable to parallel decoding: when retrieved evidence is sufficiently informative, many output tokens become conditionally independent given the evidence and can be committed simultaneously. Threshold-based unmasking (Equation (6)) exploits this by committing all high-confidence tokens at once rather than artificially limiting parallelism. The following example illustrates this behavior for τc =0.8, requiring only 2 steps: Where was the place of death of the director of Fight Of The Tertia? x2 = “[M][M][M][M][M][M][M][M][M][M][M][M][M]” D2 = “Fight of the Tertia is a 1952 West German family film directed by Erik Ode [. . . ]” x1 = “(1) Fight of the Tertia is directed by Erik Ode. (2) Erik Ode died in [M][M][M]” D1 = “Erik Ode (born Fritz Erik Signy Odemar, 6 November 1910 in Berlin, died 19 July 1983 in Kreuth-Weißach) was a [. . . ]” x0 = “(1) Fight of the Tertia is directed by Erik Ode. (2) Erik Ode died in Kreuth. Kreuth”
15
###