Masks Can Be Distracting: On Context Comprehension in Diffusion Language Models

TL;DR

This work reveals that diffusion language models, despite their global denoising objective, exhibit a pronounced locality bias and that generation-time mask tokens can distract context processing, especially as the context length grows. It shows that extra masks degrade performance in MDLMs and that unmasking can largely recover accuracy, motivating a mask-agnostic fine-tuning approach. The authors formulate a mask-agnostic loss combining cross-entropy with a distributional alignment term, and demonstrate robustness gains via LoRA-based fine-tuning across MDLMs, improving single-step decoding and reducing locality bias. Collectively, the study provides practical evaluation guidelines and actionable methods to strengthen context comprehension in diffusion-based language models, while highlighting limitations and directions for future work in long-context robustness and training dynamics.

Abstract

Masked Diffusion Language Models (MDLMs) have recently emerged as a promising alternative to Autoregressive Language Models (ARLMs), leveraging a denoising objective that, in principle, should enable more uniform context utilisation. In this work, we examine the context comprehension abilities of MDLMs and uncover two key limitations. First, despite their more global training objective and bidirectional attention mechanism, similarly to ARLMS, MDLMs exhibit a strong locality bias: performance is highly sensitive to the position of relevant information within the input, favouring local over distant context. Second, we show that appending a large number of mask tokens--required for generation--can significantly degrade context comprehension. Through systematic ablations, we find that these masks act as distractors, reducing the model's ability to process relevant information. To address this, we introduce a mask-agnostic loss function that encourages predictions to remain invariant to the number of appended masks. Fine-tuning with this objective substantially mitigates the distracting effect of masks, improving robustness of MDLMs. Overall, our findings reveal critical limitations of the current MDLM training paradigm and provide actionable insights for building diffusion-based language models with stronger context comprehension.

Figures (29)

• Figure 1: MDLMs display a recency bias. The performance of both MDLMs (LLaDA and Dream) and ARLMs is sensitive to the placement of relevant information within the context. For MDLMs, the performance degrades significantly when the relevant information is placed far away from the test question, suggesting a recency bias.
• Figure 2: MDLMs prioritise information placed closest to the mask. All studied MDLMs perform best when relevant information is near the masked token, regardless of question position.
• Figure 3: Gradient attribution analysis further illuminates the locality bias of the models. Although all models display the characteristic U-shaped behaviour, MDLMs demonstrate more uniform gradients across different positions, indicating reduced locality bias compared to their ARLM counterparts.
• Figure 4: Performance of LLaDA decreases significantly with added mask tokens, while Dream is more robust to the extra masks.
• Figure 5: For LLaDA, performance degradation becomes more significant as the context length increases. We do not observe a similar effect for Dream, which is robust to the effect of extra masks.