Unlocking the Potentials of Retrieval-Augmented Generation for Diffusion Language Models

TL;DR

This work addresses the challenge of grounding Diffusion Language Models (DLMs) with Retrieval-Augmented Generation (RAG) by identifying Response Semantic Drift (RSD) as a core bottleneck that degrades precision. It introduces SPREAD, a query-relevance-guided denoising framework, which uses contextualized embeddings to compute token-query relevance and steer the denoising process toward query semantics, thereby reducing drift. Across six QA datasets and two DLM baselines, SPREAD significantly improves precision and suppresses RSD with only modest computational overhead. The results suggest that query-aware denoising can unlock the potential of RAG for DLMs, enabling more reliable, semantically faithful generation grounded in retrieved knowledge.

Abstract

Diffusion Language Models (DLMs) have recently demonstrated remarkable capabilities in natural language processing tasks. However, the potential of Retrieval-Augmented Generation (RAG), which shows great successes for enhancing large language models (LLMs), has not been well explored, due to the fundamental difference between LLM and DLM decoding. To fill this critical gap, we systematically test the performance of DLMs within the RAG framework. Our findings reveal that DLMs coupled with RAG show promising potentials with stronger dependency on contextual information, but suffer from limited generation precision. We identify a key underlying issue: Response Semantic Drift (RSD), where the generated answer progressively deviates from the query's original semantics, leading to low precision content. We trace this problem to the denoising strategies in DLMs, which fail to maintain semantic alignment with the query throughout the iterative denoising process. To address this, we propose Semantic-Preserving REtrieval-Augmented Diffusion (SPREAD), a novel framework that introduces a query-relevance-guided denoising strategy. By actively guiding the denoising trajectory, SPREAD ensures the generation remains anchored to the query's semantics and effectively suppresses drift. Experimental results demonstrate that SPREAD significantly enhances the precision and effectively mitigates RSD of generated answers within the RAG framework.

Figures (5)

• Figure 1: Precision and recall of DLMs equipped with RAG. While both LLaDA and Dream achieve high recall, their precision is consistently relatively low.
• Figure 2: An illustration comparing different denoising strategies in DLMs. (a) Context-Based Confidence. The model relies on its internal confidence score, which is potentially biased by the global retrieved context, causing the model to generate irrelevant information. (b) Query-Based Relevance. An ideal approach maintains a persistent focus on the query's semantics, which allows the model to prioritize relevant facts and generate a concise and relevant response.
• Figure 3: An overview and illustrative comparison of our proposed method: (a) Confidence-based Strategy: The existing methods rely on local model confidence for token selection, lacking guidance from the global query semantics and leading to semantic drift during generation. (b) Our SPREAD strategy: We assess the relevance of each candidate token to the query and directly guides the token selection, actively steering the generation process to maintain continuous alignment with the query semantics.
• Figure 4: A comparative analysis of Copy Rate and Redundancy across different denoising strategies.
• Figure 5: The results of multi-dimensional comparison of different masking strategies.