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

AI summary

Overview

• Research area: Machine learning / natural language processing — specifically masked diffusion language models (MDLMs) and how they process context.
• Technical level: Advanced. The paper combines behavioural benchmarking with gradient attribution analysis and a custom fine-tuning objective.
• Scope: A systematic empirical study of context comprehension in open-source MDLMs (LLaDA and Dream), showing that they carry a strong locality bias and that appended mask tokens act as distractors, plus a mask-agnostic fine-tuning loss that mitigates the effect.

What This Paper Is About

Masked Diffusion Language Models (MDLMs) are presented as an alternative to autoregressive language models (ARLMs) because they denoise tokens across the whole sequence in parallel and use bidirectional attention, which in principle should let them use context more uniformly. This paper tests that assumption: it asks whether MDLMs really use context more evenly than ARLMs, and whether the mask tokens required for generation help or hurt context comprehension. The authors find that MDLMs still favour nearby context, and that adding many mask tokens actively degrades their ability to use the context, then propose a fine-tuning objective to counteract this.

Key Contributions

1. Locality bias (Section 4): The paper provides what it describes as the first systematic evidence that MDLMs exhibit a strong locality bias, prioritising information near the masked token despite their global denoising objective.
2. Inverse scaling with masks (Section 5): It uncovers an inverse scaling law with extra masks: in MDLMs trained from scratch with the masked diffusion objective, additional mask tokens can significantly degrade performance, especially in long-context settings.
3. Mask-agnostic fine-tuning (Section 6): It proposes a mask-agnostic (MA) objective that enforces prediction invariance to mask count, improving robustness of MDLMs.
4. Evaluation guidelines (Section 7): It establishes mask configuration as a critical factor in MDLM evaluations and recommends standardised practices for fair benchmarking.

Main Findings

• Recency, then locality: With 10 relevant examples grouped in one block and 40 distractor examples, with the test example at the right end of the context, both LLaDA and Dream perform best when relevant information sits immediately before the test question. Accuracy declines monotonically as relevant information moves farther away. Unlike ARLMs, which often show a U-shaped curve, MDLMs show no strong primacy effect.
• The bias follows the mask, not the absolute position: When the position of the masked test question is varied, all studied MDLMs perform best when relevant information is near the masked token, regardless of its absolute position in the prompt. Performance is consistently lowest when the masked question appears at the beginning of the input. For Dream, performance is generally better when relevant information is to the left of the mask.
• A proposed origin: The masked diffusion loss is scaled by 1/p, where p is the probability of masking a token, so training places greater weight on cases where only a few tokens are masked — scenarios where nearby context usually suffices.
• Gradient attributions are non-uniform: Computing the L2 norm of gradients of the predicted answer token's logit with respect to input token embeddings (10 relevant and 40 randomly mixed distractors, ordering randomised over 30 seeds, 20 test questions per task for efficiency) yields non-uniform attribution patterns. The attribution plots show a U-shape across all models, but MDLMs — particularly the base models — display more evenly distributed gradient values than ARLMs.
• More masks hurt: Appending additional mask tokens causes consistent, monotonic accuracy degradation as the mask count increases, for LLaDA-Base and LLaDA-Instruct (and LLaDA-MoE models). Dream models are more robust but still drop 6 and 8 percentage points for the base and instruct models respectively when approximately 20 masks are added.
• Degradation grows with context length: For LLaDA, performance degradation from extra masks generally increases as more distractor examples are added. Tasks that benefit most from additional in-context examples are also the most vulnerable to mask-induced degradation.
• MDLMs over-weight mask tokens mechanistically: With 50 appended masks, average normalised gradients attributed to mask tokens are much higher than for non-mask tokens. Dream-Base: 0.282 ± 0.040 for masks versus 0.012 ± 0.007 for the last 50 non-mask tokens and 0.005 ± 0.003 for all non-masks. Dream-Instruct: 0.144 ± 0.031 versus 0.030 ± 0.005 and 0.018 ± 0.002. LLaDA-Base: 0.234 ± 0.021 versus 0.005 ± 0.002 and 0.005 ± 0.002. LLaDA-Instruct: 0.220 ± 0.031 versus 0.057 ± 0.014 and 0.017 ± 0.003. The last 50 non-mask tokens score higher than non-mask tokens generally, reiterating the recency bias.
• It is the masks, not repetition: Replacing extra masks with the string "." repeated multiple times has only a minor impact on LLaDA. Performance decreases by up to 3 and 10 percentage points for the base and instruct models respectively with dots, compared to 23 and 27 percentage points for masks. For Dream, the effect of masks and dots is largely similar.
• Unmasking recovers accuracy: Running 40 decoding steps and unmasking progressively markedly improves accuracy, especially with the high-confidence strategy compared to random unmasking. This adds latency.
• Extra masks flatten the locality bias: With more masks, accuracy becomes less location-dependent, but mainly because it declines across all positions.
• Mask-agnostic fine-tuning works: Fine-tuning LLaDA-Base and LLaDA-Instruct with the MA loss (LoRA adapters, a subset of the OpenOrca dataset, approximately 1.2k gradient descent steps) improves robustness to the number of appended masks. An ablated model trained with the cross-entropy loss only (β = 0.0) does not show a similar effect. The MA loss also reduces model entropy, makes the logits smoother as a function of masks, does not degrade language modelling performance, and reduces the locality bias of LLaDA-Base.

Methodology in Plain English

The authors study open-source MDLMs so they can control generation settings, comparing LLaDA-8B against Llama3-8B and Dream-7B against Qwen2.5-7B, the ARLM used to initialise Dream-7B. LLaDA was trained from scratch with the masked diffusion loss; Dream was initialised from ARLM weights, making it an interesting midpoint. Greedy decoding is used throughout. LLaDA and Dream were trained with context lengths of 4096 and 2048 tokens respectively.

Rather than needle-in-a-haystack or multi-document QA, which the authors note are relatively easy and require longer contexts than these models support, they build a suite of few-shot multiple-choice tasks where an abstract rule must be inferred from examples. There are 8 relevant word tasks (for example, choose the adjective, verb, or fruit) and 2 distractor number tasks (choose the largest or smallest number); all combinations give 16 tasks, each with 1000 test points. The correct answer is always a single token, which lets the authors use accuracy and also analyse gradients and prediction entropy without worrying about decoding order. Shaded regions in figures mark 95% confidence intervals computed over the 16 datasets. Additional results are reported on HotPotQA, GSM8k, and a multidimensional classification dataset.

To probe position sensitivity, they keep example order within the relevant and distractor groups fixed and vary only where the relevant block sits. To separate "right edge" from "near the mask", they also vary the position of the masked test question. For mechanistic evidence they compute gradient attributions — the L2 norm of the gradient of the predicted answer's logit with respect to input token embeddings — and they test alternative explanations for degradation (such as repeated non-mask tokens). Finally, they fine-tune with a mask-agnostic loss combining a cross-entropy term scaled by 1/p and a total-variation term scaled by p, both divided by the number of masked tokens, so predictions stay consistent across different numbers of appended masks.

Why This Matters

• Impact on research: The paper argues that mask configuration is a critical, often unreported variable in MDLM evaluation and that models trained from scratch with the masked diffusion objective carry a systematic weakness. It frames the problem as a training artifact rather than an architectural flaw, since the mask-agnostic loss corrects it. It also highlights a "mask tax" on parallel decoding: the large number of masks needed for fast generation distorts context comprehension before token-dependency issues even arise.
• Real-world applications:
  • Low-latency text generation on hardware-constrained devices, where few decoding steps are required.
  • Adaptive parallel decoding, which aims to generate multiple tokens in parallel.
  • Distillation and model-compression pipelines, where minimising generation steps is critical.
  • Long-context tasks such as retrieval, multi-hop reasoning (HotPotQA), and grade-school maths word problems (GSM8k), where the mask count may silently change reported performance.
• Industry relevance: The findings come from Qualcomm AI Research with an academic collaborator, and the practical framing is about deploying diffusion language models efficiently. The paper explicitly recommends that benchmark reports state the number of mask tokens used and that mask-sensitivity analysis on long-context tasks become standard, which directly affects reproducibility and fair comparison of decoding strategies and post-training methods.

Future Directions

• Examine uniform diffusion models, which avoid explicit masks and apply noise more evenly across the input, to see whether they are more robust to context placement and decoding configurations — this would clarify whether the issues are intrinsic to diffusion or specific to masked variants.
• Investigate the mechanisms behind the locality bias, which the paper states remain unclear, via deeper analysis of the masked diffusion objective.
• Obtain access to pre-training protocol details, including exact datasets, to disentangle model-specific quirks from general properties of MDLMs; the paper notes that pre-training data and pipelines are not publicly released for the studied models.
• Evaluate whether novel decoding strategies and post-training methods change robustness to mask count, an analysis the authors want incorporated into standard MDLM evaluation pipelines.

Target Audience

Researchers and engineers working on diffusion language models, efficient inference, and long-context language modelling, as well as practitioners who benchmark or deploy MDLMs and need to know which generation settings affect reported accuracy. Readers need some familiarity with masked diffusion objectives, attention, and fine-tuning to follow the loss formulation and gradient analysis.

Authors’ 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.