In-Context Anchor Guidance: Mitigating Temporal Attention Dilution in Masked Diffusion Language Models

Masked Diffusion Language Models (MDLMs) are a practical, non-autoregressive alternative to standard causal language models, demonstrating In-Context Learning (ICL) abilities. However, while autoregressive (AR) models typically suffer from spatial attention dilution, MDLMs face a distinct challenge: temporal attention dilution. As the generation process transitions from uniform noise to structured text over time, the attention mechanism gradually assigns lower weights to the initial prompt during the final denoising steps. This causes context drift and reduces the model's ability to follow instructions. To address this, we propose In-Context Anchor Guidance (ICAG), a training-free inference method with minimal computational overhead. ICAG dynamically adjusts pre-softmax attention logits by applying a time-dependent scaling factor specifically to the prompt tokens. As an analytical contribution, we adapt the spatial guidance method for AR models, LaPA$^2$, to the MDLM framework, indicating experimentally that static spatial adjustments are less effective for temporal diffusion tasks. Evaluations on the IFEval and GSM8K benchmarks show that ICAG reduces context drift, improving both instruction adherence and complex reasoning accuracy.