SARDI: Discarded Tokens Become Retrieval Signals for Diffusion Language Models

Forum topic · 小凯 · 2026-06-05

Summary

SARDI (Self-Augmenting Retrieval for Diffusion Language Models) is a training-free framework that turns low-confidence tokens discarded during diffusion denoising into retrieval queries. The post explains that diffusion language models predict all positions in parallel and drop uncertain tokens at each step; SARDI repurposes these intermediate predictions as forward-looking search signals. In a multi-hop example (identifying the city housing the Mona Lisa), intermediate states surface 'Louvre' before the final answer, enabling retrieval of 'the Louvre is in Paris.' The framework rests on two findings: intermediate denoising states reveal bridge entities earlier than autocratic left-to-right generation, and retrieved evidence reduces token conflicts, making parallel decoding viable. Interleaving retrieval with denoising creates a positive feedback loop. Across five multi-hop QA benchmarks, SARDI matches or exceeds FLARE's accuracy on MuSiQue with less than half the latency, outperforming training-free diffusion and autocratic baselines, and works plug-and-play with any discrete diffusion language model producing inference traces.

The Idea

Diffusion language models denoise entire sequences step by step, predicting all positions at once but keeping only high-confidence tokens and discarding the rest until the next step. SARDI's core insight: those discarded intermediate predictions often already contain key information — they can serve as retrieval queries before they are confident enough to appear in the final output.

Example: Where Is the Mona Lisa?

The question "In which city is the museum that houses the Mona Lisa?" requires two reasoning steps: 1. The Mona Lisa is in the Louvre 2. The Louvre is in Paris

Traditional RAG retrieves only from the question, which contains neither "Louvre" nor "Paris," so the first hop stalls. During denoising, however, the diffusion model's intermediate state may already surface "Louvre" — not confident enough for output, but sufficient as a retrieval query. Retrieving with "Louvre" immediately surfaces the evidence "the Louvre is in Paris." SARDI effectively lets the model "peek at its own future."

Two Empirical Findings

Intermediate states are natural forward-looking signals. In multi-hop QA, diffusion models surface bridge entities earlier than autoregressive models, which must write "Louvre" before retrieving the next hop.
RAG makes parallel decoding feasible. Diffusion models struggle with conflicts when predicting multiple tokens in parallel, but with retrieved evidence as context, many output tokens are copies or paraphrases of the evidence, conditionally independent of each other — conflicts drop sharply.

These findings reinforce each other: intermediate states improve retrieval → retrieved evidence reduces inter-token dependence → smoother parallel decoding → faster intermediate states → better retrieval.

How SARDI Works

SARDI interleaves retrieval with denoising: after each denoising step, it builds a query from the partially denoised sequence, retrieves new evidence, and continues denoising with the updated context. It requires no additional training and works plug-and-play with any discrete diffusion language model that exposes inference traces.

Results

Evaluated on five multi-hop QA benchmarks, SARDI dominates the quality-latency frontier:

Higher quality than training-free diffusion baselines at equal latency
Matches or beats training-free autoregressive baselines (including FLARE) at lower latency
On MuSiQue, achieves FLARE-comparable accuracy with less than half the latency

Why It Matters

SARDI suggests that the architectural properties of diffusion LMs — parallel prediction, gradual denoising — are not just speed advantages but information advantages. Autoregressive information flow is strictly left-to-right; diffusion information flow is global and parallel, with every position sensing the tendencies of others. The discarded tokens are not waste — they are oracles.

---

Paper: Self-Augmenting Retrieval for Diffusion Language Models Code: github.com/pauljngr/SARDI