Adaptive Guidance for Retrieval-Augmented Masked Diffusion Models

2026-03-18

AI Summary

This work addresses the challenge in retrieval-augmented generation (RAG) where noisy or unreliable retrieved contexts often conflict with a model’s internal knowledge, thereby degrading the output quality of diffusion-based language models. To mitigate this issue, the authors propose ARAM, a novel framework that introduces an adaptive guidance mechanism into mask-based diffusion RAG systems for the first time. ARAM dynamically modulates the strength of guidance during the denoising process by estimating the signal-to-noise ratio induced by distributional shifts from retrieved contexts, effectively balancing external evidence and internal knowledge without requiring additional training. Experimental results demonstrate that ARAM significantly outperforms existing RAG baselines across multiple knowledge-intensive question answering benchmarks, achieving substantial gains in answer accuracy.

Technology Category

• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval
• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Retrieval-Augmented Generation
• retrieval-prior conflict
• diffusion-based language models
• noisy context
• factual grounding

Innovation

• Retrieval-Augmented Generation
• Masked Diffusion Models
• Adaptive Guidance
• Signal-to-Noise Ratio
• Distributional Shift