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

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, appending a large number of mask tokens—required for generation—can significantly degrade context comprehension in models trained from scratch. Through systematic ablations, we find that these masks act as distractors, reducing the model's ability to process relevant information. To address and further study this undesirable behaviour, we introduce the 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, with implications for training, evaluation and deployment.

Diffusion language models are a new type of language model that can generate many words at once, which could make them much faster than systems like ChatGPT, which generate words sequentially, one by one. Because they process the whole input at the same time rather than word by word, in theory they should also be better at using information from across a long piece of text. We wanted to test whether this is actually true. We designed a set of controlled language tasks where the model had to find useful information hidden among many distracting examples. We found that these models still strongly prefer to rely on nearby information and struggle to use relevant details that are placed farther away. We also discovered a second problem: the special “mask” tokens used during generation can distract the model and make its predictions much worse, especially when many masks are added. To address this, we developed a new training method that teaches the model to ignore unnecessary masks. This made the models more stable and improved their ability to use long context. Our findings highlight important limitations of current diffusion language models and suggest ways to make them more reliable and efficient.