Adaptive Guidance for Retrieval-Augmented Masked Diffusion Models

Jaemin Kim; Jong Chul Ye

cs.CLMar 18, 2026

Adaptive Guidance for Retrieval-Augmented Masked Diffusion Models

Authors: Jaemin Kim, Jong Chul Ye

Organizations: Graduate School of AI, KAIST, Republic of Korea

Abstract

Retrieval-Augmented Generation (RAG) improves factual grounding by incorporating external knowledge into language model generation. However, when retrieved context is noisy, unreliable, or inconsistent with the model's parametric knowledge, it introduces retrieval-prior conflicts that can degrade generation quality. While this problem has been studied in autoregressive language models, it remains largely unexplored in diffusion-based language models, where the iterative denoising process introduces unique challenges for integrating retrieved context. In this work, we propose Adaptive Retrieval-Augmented Masked Diffusion (ARAM), a training-free adaptive guidance framework for Masked Diffusion Models (MDMs) in RAG settings. ARAM dynamically calibrates the guidance scale during denoising according to the Signal-to-Noise Ratio (SNR) of the distributional shift induced by retrieved context. Intuitively, the model strengthens guidance when the retrieved context provides reliable corrective evidence and suppresses it when the contextual signal is noisy or non-supportive. Extensive experiments on multiple knowledge-intensive QA benchmarks show that ARAM improves overall QA performance over competitive RAG baselines.

Figures & tables

Figure 1: Overview of ARAM. ARAM mitigates knowledge conflict via step-wise and token-wise adaptive guidance driven by the Signal-to-Noise Ratio.

Method	Positive ( $\uparrow$ )	Negative ( $\downarrow$ )	Cons. Correct	Cons. Wrong
LLaDA
RAG	0.135	0.051	0.040	0.774
Ours	0.234	0.022	0.069	0.675
Dream
RAG	0.259	0.049	0.100	0.592
Ours	0.266	0.038	0.111	0.585

Table 2: Categorization rates of retrieval-prior interaction outcomes.

Method	Gold		Non-Answering		Irrelevant
Method	LLaDA	Dream	LLaDA	Dream	LLaDA	Dream
RAG	24.2	49.7	0.77	1.54	0.00	0.00
Ours	41.7	52.0	1.92	1.92	0.00	0.00

Table 3: Performance comparison across subsets stratified by retrieval context quality.

Panel of Figure 3 — Figure 3: Dynamics of adaptive guidance . (a) Average guidance scale ( $\lambda_{t}$ ) per step. ARAM maintains high guidance for gold documents but suppresses $\lambda_{t}$ for noisy ones. (b) Token-wise guidance heatmap. Black squares mark unmasking positions, demonstrating a right-to-left generation tendency. ARAM assigns larger guidance to positions that align with the unmasking pattern.

$\lambda_{\max}$	$\beta$	LLaDA		Dream
$\lambda_{\max}$	$\beta$	EM	F1	EM	F1
1.0	1.0	24.9	36.1	33.9	43.6
2.0		13.9	27.3	31.4	42.4
3.0		9.1	22.8	26.3	38.9
1.0	0.1	30.4	39.3	33.5	42.0
	0.3	28.4	38.7	34.3	43.4

Table 4: Ablation study on

\lambda_{\max}

and

\beta

. Bold denotes the best score and underlined represents second-best.

Noise Proxy	LLaDA		Dream
Noise Proxy	EM	F1	EM	F1
$\mathrm{Var}_{p_{\text{prior}}}(s)$	20.3	32.5	33.5	43.2
$\mathcal{H}(p_{\text{cond}})$	24.9	36.1	33.9	43.6

Table 5: Comparison of noise proxies. Average EM and F1 scores are reported.

Method	EM	F1
LLaDA 1.5	15.8	22.0
LLaDA 1.5 + RAG	17.6	30.6
LLaDA 1.5 + Ours	24.2	35.9

Table 6: Performance on LLaDA 1.5.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Noise Proxy	LLaDA		Dream
Noise Proxy	EM	F1	EM	F1
$\mathrm{Var}_{p_{\text{prior}}}(s)$	20.3	32.5	33.5	43.2
$\mathcal{H}(p_{\text{prior}})$	20.5	32.4	33.6	42.8
$\|\mathcal{H}(p_{\text{prior}})-\mathcal{H}(p_{\text{cond}})\|$	23.7	35.3	33.4	42.8
$\mathcal{H}(p_{\text{cond}})$	24.9	36.1	33.9	43.6

Appendix

Table 7: Comparison of noise proxies. Average EM and F1 scores are reported.

Signal Proxy	LLaDA		Dream
Signal Proxy	EM	F1	EM	F1
- (Raw SNR)	21.5	33.5	33.5	43.4
Tanh	24.9	36.1	33.9	43.6

Appendix

Table 8: Comparison of stability function. Average EM and F1 scores are reported.

Method	NFE / Step	LLaDA	Dream
No-RAG	1	0.93	0.94
RAG	1	1.55	1.56
CFG	2	2.02	3.06
ARAM (Ours)	2	2.07	3.29

Appendix

Table 9: Average inference time per query (seconds).

Panel of Figure 4 — Figure 4: Prompt template for LLM-based retrieval context quality annotation.

Panel of Figure 5 — Figure 5: Prompt template used for LLaDA.

Panel of Figure 6 — Figure 6: Few-shot prompt template used for the Dream and other AR backbones.

Adaptive RAG

Adaptive retrieval-augmented generation (RAG) adjusts when, what, and how much evidence a language model retrieves for each query.

See all 59 papers

Self-Augmenting Retrieval for Diffusion Language Models
Jun 4, 2026Paul Jünger, Justin Lovelace, Linxi Zhao +2Adaptive RAG Retrieval-Augmented Generation
Department of Computer Science, Cornell University.
Does More Retrieved Evidence Help Visual Retrieval-Augmented Generation with Diffusion Language Models?
Aug 7, 2026Jiankun Wang, Yisen Gao, Ziwei Zhang +3Multimodal RAG Visual Question Answering
School of Computer Science and Engineering, Beihang University · Department of Computer Science and Engineering, The Hong Kong University of Science and Technology · School of Computer Science and Engineering, Guangxi Normal University +2
Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models
Jun 15, 2026Yanming Zhang, Yihan Bian, Jingyuan Qi +3Masked Diffusion Models Diffusion Model Fine-Tuning
University of Maryland · Virginia Tech · Intuit +2

Jun 4, 2026J/K move · Enter open · S save

Self-Augmenting Retrieval for Diffusion Language Models

Paul Jünger, Justin Lovelace, Linxi Zhao, Dongyoung Go+1

Department of Computer Science, Cornell University.

Adaptive RAG Retrieval-Augmented Generation Masked Diffusion Models Multi-Hop QA Diffusion Language Models Diffusion Language Model Inference+1

Open paper AskPDF arXiv

Discrete diffusion language models generate text by iteratively denoising an entire response in parallel. At each step, they predict tentative tokens for every masked position, committing the confident predictions to the output and discarding the unconfident ones. We show that the discarded tokens are in fact a useful lookahead signal for retrieval-augmented generation: even low-confidence tokens often surface salient entities early in the denoising trajectory, enabling retrieval of stronger evidence before the output is finalized. We exploit this through Self-Augmenting Retrieval for Diffusion Language Models (SARDI), a dynamic RAG framework that uses these lookahead tokens to guide retrieval during denoising. SARDI is training-free, retriever-agnostic, and applicable to any reasoning-capable discrete diffusion language model. Across five multi-hop QA benchmarks, SARDI outperforms current training-free diffusion and autoregressive retrieval baselines at up to

8\times

higher throughput.

Abstract

Adaptive RAG

Explore similar work

Self-Augmenting Retrieval for Diffusion Language Models

Does More Retrieved Evidence Help Visual Retrieval-Augmented Generation with Diffusion Language Models?

Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models