Masks Can Be Distracting: On Context Comprehension in Diffusion Language Models
Julianna Piskorz, Cristina Pinneri, Alvaro Correia, Motasem Alfarra, Risheek Garrepalli, Christos Louizos
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, 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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 093b7034-102b-4e20-b6ca-e267cc6da797Cited by top-tier papers1
Ask how each one uses itBuilds on23
- Structured Denoising Diffusion Models in Discrete State-SpacesJacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow et al.NeurIPS 2021 · 2,256 citations
- Simple and Effective Masked Diffusion Language ModelsSubham S. Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan et al.NeurIPS 2024 · 929 citations
- Discrete Diffusion Modeling by Estimating the Ratios of the Data DistributionAaron Lou, Chenlin Meng, Stefano ErmonICML 2024 · 473 citations
- Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel DecodingChengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu et al.ICLR 2026 · 428 citations
- Function Vectors in Large Language ModelsEric Todd, Millicent L. Li, Arnab Sen Sharma, Aaron Mueller et al.ICLR 2024 · 229 citations
Related papers
- Mitigating Mask Prior Drift and Positional Attention Collapse in Large Diffusion Vision-Language ModelsSujung Hong, Chanyong Yoon, Seong Jae HwangICML 2026
- Unifying Masked Diffusion Models with Various Generation Orders and BeyondChunsan Hong, Sanghyun Lee, Jong Chul YEICML 2026
- Dual-objective Language Models: Training Efficiency Without OverfittingDavid Samuel, Lucas Georges Gabriel CharpentierICLR 2026
- Any-Order GPT as Masked Diffusion Model: Decoupling Formulation and ArchitectureShuchen Xue, Tianyu Xie, Tianyang Hu, Zijin Feng et al.ICML 2026 · 15 citations
- Diffusion Beats Autoregressive in Data-Constrained SettingsMihir Prabhudesai, Mengning Wu, Amir Zadeh, Katerina Fragkiadaki et al.NeurIPS 2025 · 69 citations
