Revealing the Attention Floating Mechanism in Masked Diffusion Models

TL;DR

This paper investigates attention mechanics in Masked Diffusion Models (MDMs) and identifies Attention Floating, a dynamic, dispersed anchoring of attention that shifts across denoising steps and layers, unlike the fixed attention sink observed in autoregressive models. It introduces a Shallow Structure-Aware, Deep Content-Focused attention mechanism, supported by a QK decomposition and retrieval-head analyses that show shallow layers anchor global structure while deep layers focus on semantic content. Empirically, Attention Floating enhances knowledge extraction from context, with MDMs achieving substantially larger gains from retrieved evidence than ARMs and displaying robustness to contextual noise, positional bias, and evidence distribution via region-level attention flow analyses. These findings offer a mechanistic explanation for the superior in-context learning capabilities of MDMs and suggest avenues to leverage floating signals for improved contextual reasoning and retrieval integration.

Abstract

Masked diffusion models (MDMs), which leverage bidirectional attention and a denoising process, are narrowing the performance gap with autoregressive models (ARMs). However, their internal attention mechanisms remain under-explored. This paper investigates the attention behaviors in MDMs, revealing the phenomenon of Attention Floating. Unlike ARMs, where attention converges to a fixed sink, MDMs exhibit dynamic, dispersed attention anchors that shift across denoising steps and layers. Further analysis reveals its Shallow Structure-Aware, Deep Content-Focused attention mechanism: shallow layers utilize floating tokens to build a global structural framework, while deeper layers allocate more capability toward capturing semantic content. Empirically, this distinctive attention pattern provides a mechanistic explanation for the strong in-context learning capabilities of MDMs, allowing them to double the performance compared to ARMs in knowledge-intensive tasks. All codes and datasets are available at https://github.com/NEUIR/Attention-Floating.

Figures (17)

• Figure 1: Comparison of ARMs and MDMs.
• Figure 3: Layer-Wise Attention Absorption Rate in ARMs and MDMs.
• Figure 4: QK Geometric Decomposition across Different Layers in MDM.
• Figure 5: Retrieval Head Analysis in MDMs.
• Figure 6: Performance with Varying Numbers of Distractor Documents.