PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Julianna Piskorz, Cristina Pinneri, Alvaro Correia, Motasem Alfarra, Risheek Garrepalli, Christos Louizos

arXiv:2511.21338v2cs.LG

TL;DR

This paper examines how masked diffusion models use context, an area that had not been comprehensively evaluated despite their potential for global context integration. It systematically analyzes locality and sensitivity to mask tokens, and finds that both affect context comprehension; mask-agnostic fine-tuning is proposed to improve robustness.

  • Problem

    How masked diffusion training affects context comprehension and familiar autoregressive positional biases remained unclear, limiting understanding of how MDLMs use context during inference.

    Source: §1 · p. 1

  • Method

    The paper systematically studies MDLM context use and proposes a mask-agnostic loss that encourages answer predictions to remain consistent across different mask configurations.

    Source: §1 · p. 2, §6.1 · p. 8

  • Results

    MDLMs show strong locality effects and sensitivity to mask configurations, while extra masks can significantly degrade performance, especially for models trained from scratch in long-context settings.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    MDLM evaluations should report mask counts and include mask-sensitivity analysis on long-context tasks; the authors also propose mask-agnostic fine-tuning to improve robustness.

    Source: §1 · p. 2, §7 · p. 9

    The study focuses on open-source MDLMs whose pretraining details, including exact datasets, are not publicly available, limiting controlled comparisons and separation of model-specific from general behaviors.

    Source: §7 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Although masked diffusion models may integrate context more globally than autoregressive models, how they use context during inference remains unclear. This study finds locality and mask-count sensitivity and proposes mask-agnostic fine-tuning to improve robustness.

1. Introduction

  • The study addresses an open question about how MDLMs use context and whether their training objective changes familiar autoregressive inductive biases.MDLMs have shown competitive scaling, perplexity, and speed on standard benchmarks, but their context use during inference remains unclear.Source: §1 · p. 1
  • MDLMs exhibit strong locality effects, relying more on nearby context than uniformly integrating information across the input sequence.The study also finds that mask-token number and placement can alter performance.Source: §1 · p. 1
  • Appending more mask tokens can significantly degrade performance in MDLMs trained from scratch, especially in long-context settings.The degradation also changes context-processing patterns, levelling performance across input locations at a lower base level.Source: §1 · pp. 1–2
  • The authors propose mask-agnostic fine-tuning to encourage prediction invariance across mask counts and improve MDLM robustness.They also recommend reporting mask configurations and evaluating mask sensitivity on long-context tasks.Source: §1 · p. 2

2. Related Works

Prior work documents position biases in language models and presents competing explanations for autoregressive models. A comprehensive evaluation of how masked diffusion training affects context comprehension remains missing.

2. Related Works

  • A comprehensive evaluation of how the masked diffusion objective influences context comprehension was still missing.MDLMs had scaled to larger sizes and achieved competitive perplexity on standard benchmarks.Source: §2 · p. 2
  • Language models do not use input information uniformly; primacy and recency biases can combine into the lost-in-the-middle pattern in autoregressive models.Prior work links primacy mainly to causal attention and recency to training-data distributions and next-token prediction.Source: §2 · p. 2
  • Prior work evaluated LLaDA on needle-in-the-haystack tasks for generalization to unseen context lengths, but the setup was too simple to reveal within-context recency effects; other work studied left-to-right decoding preference.Shansan et al. (2025) examined MDLMs’ preference for left-to-right decoding, while the paper investigates AR-like tendencies in context processing.Source: §2 · p. 2

3. Experimental Setup

The experiments compare open-source MDLMs with autoregressive counterparts under controlled generation settings. The authors use few-shot rule-learning tasks to test sensitivity to information placement within model context limits.

3. Experimental Setup

  • The study compares LLaDA-8B with Llama3-8B and Dream-7B with Qwen2.5-7B, using greedy decoding for all models.LLaDA was trained from scratch with masked diffusion loss, whereas Dream was initialized from autoregressive Qwen2.5 weights; additional LLaDA-MoE results appear in Appendix A.1.Source: §3 · p. 2
  • Because standard retrieval tasks require longer contexts to reveal positional bias, the study instead tests within the models’ trained context lengths.LLaDA and Dream were trained with context lengths of 4096 and 2048 tokens, respectively.Source: §3 · pp. 2–3
  • The authors construct few-shot multiple-choice tasks that require inferring an abstract rule from examples and vary relevant-information placement.The correct answer spans one token, enabling accuracy-based evaluation and finer analysis without controlling for token decoding order.Source: §3 · p. 3
  • The evaluation suite contains 16 tasks, each with 1,000 test points, built from eight relevant-word tasks and two number-based distractor tasks.The governing rule requires exposure to multiple examples, and relevant and irrelevant examples are randomly mixed to test long-context comprehension.Source: §3 · p. 3

4. Are MDLMs Location-Sensitive?

MDLM performance depends on the proximity of relevant information to the prediction target, despite the models’ global denoising objective. Gradient analyses also find nonuniform context attribution, while suggesting more distributed attribution in MDLMs than in autoregressive models.

4.1. Is the Performance of MDLMs Sensitive to the Location of Relevant Information?

  • MDLM accuracy is highest when relevant examples immediately precede the test question and declines monotonically as they move farther away.Unlike autoregressive models’ commonly U-shaped pattern, MDLMs show no strong primacy effect.Source: §4.1 · pp. 3–4, Figure 1 · p. 4

4.2. Does ‘Locality’ Depend on Mask Placement?

  • Across tested question positions, MDLMs perform best when relevant information is near the masked question, indicating locality rather than a fixed right-edge preference.Performance is consistently lowest when the masked question is at the input beginning; Dream generally performs better with relevant information to the mask’s left.Source: §4.2 · p. 4, Figure 2 · p. 5
  • The authors hypothesize that the locality bias reflects the masked diffusion loss’s 1/p scaling, which weights cases with few masked tokens more heavily.In those cases, nearby context is often sufficient for prediction, resembling next-token prediction.Source: §4.2 · p. 4

4.3. Quantifying the Bias with Gradient Analysis

  • Gradient attribution is nonuniform across models, with local context most influential, but MDLM base models show more evenly distributed scores than ARLMs.The analysis measures the L2 norm of gradients of the predicted answer-token logit with respect to input-token embeddings.Source: §4.3 · p. 4

5. The Distracting Effect of Extra Masks

Extra mask tokens often impair MDLM context comprehension rather than enabling more global reasoning. The section probes this effect and finds it is mask-specific, linked to context length, and partly recoverable by unmasking.

5.1. Do Extra Masks Affect MDLMs’ Performance?

  • LLaDA accuracy declines monotonically as extra masks increase, while Dream is more robust but still drops with approximately 20 added masks.Dream's base and instruct models drop 6 and 8 percentage points, respectively.Source: Figure 3 · p. 5, §5.1 · p. 5

5.3. Is the Distracting Effect Reflected in the Gradients?

  • Added masks receive markedly higher normalized gradients than non-mask tokens, while the nearest non-mask tokens also receive higher scores than non-masks overall.Dream's strong gradient response to masks does not translate into comparable performance degradation, so gradient attribution alone does not explain the training-paradigm difference.Source: Table 1 · p. 6, §5.3 · p. 6

5.2. Do MDLMs Get Equally Distracted on All Tasks?

  • For LLaDA, mask-related degradation generally grows with distractor count and effective context length; tasks benefiting most from added examples are also most vulnerable.The authors interpret this pattern as evidence that extra masks inhibit long-context comprehension.Source: §5.2 · p. 6

5.4. Is the Degradation Caused by Repeated Tokens?

  • Up to 3 and 10 percentage points for LLaDA-Base and LLaDA-Instruct with extra dots, versus 23 and 27 points with masks, implicate masks rather than repetition.For Dream, mask and dot effects are largely similar.Source: Figure 5 · p. 6, §5.4 · p. 6

5.5. Can the Negative Effect be Fixed by Unmasking?

  • With 40 decoding steps, unmasking recovers accuracy lost to extra masks, and high-confidence selection outperforms random selection; multiple passes add latency.The high-confidence advantage is especially consistent as mask count increases.Source: Figure 6 · p. 7, §5.5 · p. 7

5.6. Do Extra Masks Affect the Locality of the MDLMs?

  • Adding masks makes performance more uniform across information positions, but mainly because results are consistently poor.The section's takeaway is that extra masks can severely impair context comprehension, especially as context length increases.Source: §5.6 · p. 7

6. Reducing the Distracting Effect Through Mask-Agnostic SFT

The section introduces a fine-tuning objective to make predictions invariant to appended mask count, then reports improved mask-count robustness and reduced locality bias. It also explains the loss construction and its experimental validation.

6. Reducing the Distracting Effect Through Mask-Agnostic SFT

  • The motivation is to avoid requiring users to specify answer length in advance and to obtain mask-count robustness without the extra inference cost of repeated decoding.The authors identify potential benefits for low-latency applications and adaptive parallel decoding.Source: §6 · p. 7

6.1. Formulation of the Mask-Agnostic Loss Function

  • The loss compares two inputs with different numbers of appended masks, while evaluating answer-part predictions and applying cross-entropy and total-variation terms.The two mask counts are sampled without replacement from a range bounded by the predefined maximum context length.Source: §6.1 · p. 8
  • The objective combines answer-token cross-entropy with total-variation loss to encourage predictions to remain consistent across different masking configurations.The cross-entropy term matches predictions to ground-truth answers; the total-variation term aligns answer-token probability distributions across masking configurations.Source: §6.1 · p. 8

6.2. Experimental Validation

  • Training uses LoRA fine-tuning on about 1.2k OpenOrca steps, with a cross-entropy-only ablation to test the contribution of total-variation regularization.OpenOrca differs from the in-context-learning evaluation setup; the authors use its diversity to encourage broader changes rather than task-specific overfitting.Source: §6.2 · p. 8
  • Mask-agnostic fine-tuning improves robustness to mask-count variations and reduces locality bias; the authors interpret its success as evidence that mask sensitivity is a training artifact, not an insurmountable architectural flaw.The reported robustness holds with only a few decoding steps.Source: §6.2 · pp. 8–9, Figure 8 · p. 8

7. Discussion, Conclusions and Future Work

The discussion recommends reporting mask counts and testing mask sensitivity on long-context tasks. It also notes that undisclosed pretraining details limit what can be concluded about MDLM behavior across models.

7. Discussion, Conclusions and Future Work

  • Reports should state the number of masks used, and evaluation should include mask-sensitivity analysis on long-context comprehension tasks.The authors say mask count is critical for reproducibility and is often omitted in prior work.Source: §7 · p. 9
  • The study focuses on open-source MDLMs whose pretraining datasets and protocols are not fully public, limiting controlled comparisons and separation of model-specific quirks from general properties.The authors note that differing training data and pipelines likely contribute to observed behaviors.Source: §7 · p. 9
  • Future work could test uniform diffusion models without explicit masks and investigate how objective weighting and noise schedules relate to locality bias.These comparisons could help clarify whether the reported issues are specific to masked variants or more intrinsic to diffusion models.Source: §7 · p. 9

Impact Statement

The impact statement describes the work as identifying MDLM locality bias and mask-configuration sensitivity with implications for deployment. An appendix figure reports different recency-bias patterns for two LLaDA-MoE models.

Impact Statement

  • The impact statement frames the findings—locality bias and sensitivity to mask configurations—as relevant to MDLM deployment and responsible development.It notes possible beneficial uses as well as potential risks if language models are deployed without proper safeguards.Source: Impact Statement · p. 10

Appendix Contents

  • LLaDA-MoE-Instruct shows strong recency bias, while LLaDA-MoE-Base is more agnostic to the location of relevant examples.Source: Figure 10 · p. 15

A.1. Results on LLaDA-MoE

LLaDA-MoE broadly reproduces the studied MDLMs’ locality and mask-sensitivity patterns, though Base performance is more uniform across example positions at the evaluated lengths. Gradient attribution and few-shot task comparisons provide additional evidence of locality bias and greater mask degradation on tasks benefiting from longer contexts.

A.1. Results on LLaDA-MoE

  • A one-mask versus two-mask setup discrepancy in LLaDA-MoE results reflects sensitivity to small changes in mask count.Figure 10 uses one mask followed by end-of-sentence, whereas Figure 15 uses two separate masks.Source: §A.1 · p. 15
  • MA-loss fine-tuning reduces extra-mask harm in LLaDA-MoE, while Base performance is mostly location-agnostic at the evaluated task lengths.The authors hypothesize that Base’s 8k context fine-tuning may explain its position-agnostic performance; attribution results still show recency patterns.Source: §A.1 · p. 15, Figure 11 · p. 16, Figure 16 · p. 17

A.2. Gradient Attribution Analysis of MDLMs

  • Normalized gradients are consistently larger near the masked question, adding evidence of locality bias across middle, left, and right question placements.The analysis also reports similar attribution patterns for LLaDA-MoE.Source: §A.2 · p. 18, Figure 17 · p. 18, Figure 20 · p. 19

A.3. Correlation Between Mask Degradation and Context Significance

  • R = −0.15 for LLaDA-Base and R = −0.16 for LLaDA-Instruct; tasks benefiting more from longer contexts tended to degrade more, by up to 60%.The correlations are described as modest, and the authors note that other factors likely also affect degradation.Source: §A.3 · p. 20

A.4. Robustness Analysis: Decoding with Few Steps on the LLaDA Models

Across low-step decoding configurations, CE+TV fine-tuning reduces sensitivity to mask count, with reported degradation reductions of 38–49% and minimal accuracy trade-offs. The appendix also reports little language-modeling impact, more uniform position performance, and smoother confidence and entropy under fine-tuning.

A.4. Robustness Analysis: Decoding with Few Steps on the LLaDA Models

  • The fine-tuned model maintains higher performance and degrades less as mask count rises across the tested configurations, with minimal accuracy trade-offs.The results include LLaDA models and LLaDA-MoE-Base under random unmasking at 2, 4, and 6 decoding steps.Source: §A.4 · p. 21, Figure 23 · p. 22
  • 38–49%: CE+TV fine-tuning reduces relative degradation across decoding-step configurations compared with the base model.For LLaDA-Base-8B, degradation falls from 15.5% to 7.9%; for LLaDA-Instruct-8B, it falls from 27.5% to 17.0%.Source: §A.4 · pp. 21–22, Figure 24 · p. 22

A.5. Robustness Analysis: Language Modelling Performance

  • MA-loss fine-tuning has little to no significant effect on language-modeling performance in the evaluated high-mask infilling regime.The table compares LLaDA models before and after fine-tuning, averaging 250 samples per dataset.Source: Table 3 · p. 23, §A.5 · p. 23

A.7. Confidence and Entropy as a Function of Masks

  • MA-loss fine-tuning increases Base-model confidence and decreases entropy, while making both smoother as mask count varies.The confidence and entropy results are reported for both LLaDA models.Source: Figure 26 · p. 24, §A.7 · p. 24

A.6. Additional Results for the Fine-Tuned LLaDA-Instruct

  • After fine-tuning, LLaDA-Instruct is more robust to extra masks and performs more uniformly across relevant-information positions.The reported position-uniformity improvement is shown under different numbers of masks.Source: Figure 25 · p. 24, §A.6 · p. 24

A.8. Experiments on the GSM8k dataset

Additional evaluations report mask sensitivity in GSM8k and HotPotQA, and locality bias in multidimensional classification. Mask-agnostic fine-tuning improves robustness in these settings, while GSM8k interpretation is constrained by limited details about the models’ pretraining and fine-tuning procedures.

A.8. Experiments on the GSM8k dataset

  • GSM8k results cannot be precisely attributed because the model papers and codebases provide limited pretraining and fine-tuning details.The authors specifically note unavailable dataset-level information, such as question-answer length distributions.Source: §A.8 · p. 25
  • GSM8k performance is mask-sensitive, particularly for instruction-tuned models; training with MA loss on GSM8k reduces that sensitivity while increasing overall performance.Fine-tuning on OpenOrca flattens performance across mask counts but lowers it overall; the authors hypothesize this may reflect deterioration of mathematical abilities.Source: §A.8 · p. 25, Figure 28 · p. 26

A.9. Experiments on the HotPotQA dataset

  • On HotPotQA, MA fine-tuning reduces Base-model sensitivity to extra masks; Instruct is more robust even before fine-tuning, with additional gains when few answer tokens are available.The reported small-token setting provides 1–2 available tokens, where the original Instruct model fails to provide a coherent answer.Source: Figure 29 · p. 27, §A.9 · p. 28

A.10. Experiments on the Multi-dimensional Classification Dataset

  • MDLM performance drops when relevant examples are far from the masked question, confirming locality bias that is weaker than in ARLMs.The reported pattern is especially evident for Dream, initialized from ARLM weights.Source: Figure 30 · p. 29, §A.10 · p. 30
  • In multidimensional classification, MA loss prevents Base-model mask-related degradation particularly under random ordering and when relevant examples are far from the test question.For Instruct, fine-tuning improves mask robustness and prevents significant degradation with small mask counts.Source: Figure 31 · p. 30, §A.10 · p. 31

B. Details of the Supervised Fine-Tuning Pipeline

The fine-tuning pipeline computes a mask-agnostic loss from paired inputs with differing appended-mask counts. It combines a total-variation term with cross-entropy and uses a curriculum for appended-mask lengths.

B. Details of the Supervised Fine-Tuning Pipeline

  • Training starts with minimal appended-mask lengths and reaches 600 masks over 5,000 gradient-descent steps.The inputs are padded with end-of-sequence tokens to equalize their lengths.Source: B. Details of the Supervised… · p. 32
  • The selected CE+TV hyperparameters differ by model: Base uses β = 1.0 and LR = 10−5, whereas Instruct uses β = 100.0 and LR = 5 × 10−7.Source: B. Details of the Supervised… · p. 32
  • For each input pair, the algorithm creates two masked-answer inputs with independently sampled appended-mask lengths, then computes the MA loss from their model outputs.The inputs are padded to equal length before model evaluation.Source: B. Details of the Supervised… · p. 32
  • The MA loss combines total variation and cross-entropy, weighted by α and β.For categorical distributions, the total-variation distance is computed as the L1 distance between post-softmax probabilities.Source: B. Details of the Supervised… · p. 32

C. Experimental Details

Appendix C describes the models, quantisation check, few-shot task construction, and prompt formats used in the experiments.

C.1. Models

  • Experiments use LLaDA and LLaDA-MoE trained from scratch with masked diffusion, Dream initialized from Qwen-2.5-7B, and Qwen-2.5-7B and Llama3-8B as autoregressive models.Source: §C.1 · p. 33

C.2. Quantisation

  • 4-bit quantisation had no significant effect on performance under varying mask counts or on model locality in the tested pattern-recognition task.The comparison used one task and is shown in Figures 32 and 33.Source: §C.2 · p. 34, Figure 32 · p. 34, Figure 33 · p. 34

C.3. Details of the Few-Shot Learning Dataset

  • Few-shot tasks pair category-based word selection with number distractors (choose smallest or largest from integers 1–1000), using three answer options.Source: §C.3 · p. 35

C.4. Formatting of the In-Context Learning Examples

  • In-context examples are fixed across test items, with answers bracketed; instruct models use paired user/assistant messages before the final test question.Source: §C.4 · p. 36
  • For locality tests, the masked answer uses one model-specific mask inside brackets; mask-count tests instead append the specified masks after the answer prefix.Source: §C.4 · p. 36
  • The locality experiment places the masked question at the beginning or end of in-context examples, while autoregressive models decode one token and are scored by string matching.Source: §C.4 · p. 36

Related papers

Loading 2511.21338v2…