Self-Augmenting Retrieval for Diffusion Language Models

Standard retrieval-augmented generation (RAG) often struggles with multi-step questions. The model cannot "see" the necessary information coming until it has already committed to a path. Current state-of-the-art systems typically rely on autoregressive (AR) models. These models generate text one token at a time. This forces retrieval to condition only on the left-to-right prefix that has already been finalized. This creates a bottleneck. If a question requires finding a "bridge entity"—an intermediate piece of information connecting two facts—the retriever is blind to it.

A new paper introduces SARDI (Self-Augmenting Retrieval for Diffusion Language Models). This framework turns a weakness into a strength. It leverages the "speculative" nature of diffusion language models (DLMs). DLMs generate text by iteratively denoising (refining a noisy sequence into a clean one) an entire sequence in parallel. They produce tentative, low-confidence predictions for every position at every step. The authors show these discarded tokens act as a lookahead signal. This allows the system to fetch the right evidence long before the final answer is stabilized.

The blindness of the left-to-right prefix

In the standard autoregressive paradigm, retrieval is reactive. If you ask a model which city is home to the museum that displays the Mona Lisa, a static retriever only sees the question. It might find the Mona Lisa. However, it won't find the specific location of the Louvre unless the model first generates the word "Louvre." As shown in, the model is essentially flying blind through the first hop of the reasoning chain.

Existing attempts to solve this in the AR space, such as FLARE, try to be proactive. They generate a tentative sentence and use it as a query. However, the authors note that this is fragile. Because AR models generate left-to-right, a single early error compounds. This creates a hallucinated query that pulls in irrelevant noise. There is a fundamental tension here. You want to look ahead to improve retrieval, but looking ahead in an AR model often means looking at garbage.

Decoupling speculation from commitment

The core innovation of SARDI is the architectural decoupling of retrieval confidence from generation confidence. In a discrete diffusion model, the engine doesn't just pick one token. It predicts a probability distribution for every masked position in the sequence simultaneously. This allows the authors to implement two distinct thresholds: $\tau_q$ (the query threshold) and $\tau_c$ (the commit threshold).

The mechanism operates in a continuous loop as illustrated in :

1. Denoising and Prediction: At each step $t$, the DLM predicts tokens for all masked positions.
2. Speculative Querying: The system looks at predicted tokens where the confidence $c_i$ meets the lower query threshold ($\tau_q$). These tokens build a "proxy sequence" combining the original question with these tentative predictions. This proxy refreshes the retrieved evidence $D$.
3. Threshold-based Commitment: Only tokens exceeding the much higher commit threshold ($\tau_c$) are "unmasked" (revealed in the output) and finalized.

By setting $\tau_q \le \tau_c$, SARDI ensures the retriever sees "hints" of the future. It catches salient entities and relations well before those tokens are stable enough for the final response.

High throughput via parallel grounding

The results suggest this fundamentally changes the decoding math. The authors report that SARDI outperforms existing training-free diffusion and autoregressive baselines across five multi-hop QA benchmarks. It achieves up to $8\times$ higher throughput (the number of questions processed per second).

On the 2WikiMultiHopQA benchmark, SARDI reaches an Exact Match (EM) of 59.1% (at $\tau_c=0.95$). This is a significant jump from the static diffusion baseline of 43.7%. When looking at the trade-off between accuracy and speed in, SARDI dominates the Pareto frontier.

You can tune the commit threshold $\tau_c$ to move along this curve. A higher $\tau_c$ increases accuracy by triggering more retrieval rounds. A lower $\tau_c$ prioritizes throughput.

The authors provide a theoretical reason for this efficiency. They measure the Conditional Mutual Information (CMI)—a metric representing the statistical dependence between adjacent tokens. They found that grounding generation in retrieved evidence sharply reduces this dependence .

When the model copies an entity directly from a document, the tokens become nearly independent. This is ideal for efficient parallel decoding in a diffusion model.

Implementation hurdles and hidden costs

There are two significant caveats for deployment. First, the "plug-and-play" claim involves the backbone model. The authors admit that off-the-shelf instruction-tuned DLMs like DREAM-7B struggle to produce necessary reasoning traces. Without these traces, the lookahead fails. To achieve these results, they had to perform light supervised fine-tuning on synthetic traces. If your base DLM does not naturally "think" in steps, SARDI has nothing to look ahead to.

Second, the computational footprint of the retrieval loop is substantial. SARDI refreshes the evidence at every single denoising step. The authors note in Section 5.6 that retrieved document sets are relatively stable. They show an 83–90% overlap between steps. However, a naive implementation would involve massive redundant encoding overhead. To make this viable at scale, you would likely need to implement document-level KV caching. This would avoid re-processing the same passages repeatedly.

The verdict

SARDI is an effective framework, but it is not a universal solution. Its success depends heavily on the model's ability to generate intermediate reasoning steps. If you are working with a model that can already produce coherent reasoning traces, this is a massive win for both accuracy and latency.

For engineers building high-throughput RAG systems, the takeaway is clear. The non-autoregressive nature of diffusion models offers a structural advantage for retrieval. The ability to separate "what we think might happen" from "what we are saying" allows for a more aggressive retrieval strategy. If you can fine-tune your backbone for reasoning, this is worth a prototype. Code is available at https://github.com/pauljngr/SARDI.

Figures from the paper