Revealing the Attention Floating Mechanism in Masked Diffusion Models
This provides insights into the internal workings of MDMs, addressing a gap in understanding for researchers in generative AI, though it is incremental as it builds on existing MDM frameworks.
The paper investigates attention behaviors in masked diffusion models (MDMs), revealing the Attention Floating phenomenon where attention anchors shift dynamically, and shows that this mechanism explains MDMs' strong in-context learning, doubling performance compared to autoregressive models in knowledge-intensive tasks.
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.