마스크 디퓨전 모델의 어텐션 플로팅 메커니즘 분석

Revealing the Attention Floating Mechanism in Masked Diffusion Models

마스크 디퓨전 모델(MDM)은 양방향 어텐션과 디노이징 과정을 활용하여, 자기 회귀 모델(ARM)과의 성능 격차를 줄여나가고 있습니다. 하지만, MDM의 내부 어텐션 메커니즘은 아직 충분히 연구되지 않았습니다. 본 논문에서는 MDM의 어텐션 동작을 분석하여 '어텐션 플로팅(Attention Floating)'이라는 현상을 밝혀냅니다. ARM과는 달리, 어텐션이 고정된 지점으로 수렴하는 반면, MDM은 디노이징 단계와 레이어에 따라 동적이고 분산된 어텐션 앵커를 보입니다. 추가적인 분석 결과, MDM은 얕은 구조 인지 및 깊은 콘텐츠 집중 어텐션 메커니즘을 사용하는 것으로 나타났습니다. 얕은 레이어는 플로팅 토큰을 사용하여 전역적인 구조 프레임워크를 구축하는 반면, 깊은 레이어는 의미 콘텐츠를 캡처하는 데 더 많은 자원을 할당합니다. 실험적으로, 이러한 독특한 어텐션 패턴은 MDM의 강력한 문맥 학습 능력에 대한 메커니즘적인 설명을 제공하며, 지식 집약적인 작업에서 ARM보다 두 배 이상의 성능을 달성할 수 있도록 합니다. 모든 코드와 데이터셋은 https://github.com/NEUIR/Attention-Floating 에서 확인할 수 있습니다.

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.