[논문 리뷰] Revealing the Attention Floating Mechanism in Masked Diffusion Models

Xin Dai, Pengcheng Huang외 7명|arXiv (Cornell University)|2026. 01. 12.

Generative Adversarial Networks and Image Synthesis인용 수 0

한 줄 요약

본 논문은 Masked Diffusion Models(MDMs)에서 Attention Floating 현상을 밝혀내고, 얕은 구조 인식에 기반한 심층 콘텐츠 중심의 주의 패턴이 맥락 내 지식 활용과 강건성을 향상시키며, 검색된 컨텍스트를 활용할 때 ARMs보다 우수함을 보여준다.

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.

연구 동기 및 목표

• 잡음 제거 과정에서 MDMs 내부에서의 어텐션이 어떻게 동작하는지 조사한다.
• Attention Floating 현상을 특성화하고 이를 ARMs의 어텐션 싱크와 대비한다.
• MDMs에서 어텐션 다이나믹스가 맥락 내 학습과 지식 활용에 어떻게 기여하는지 이해한다.
• 맥락적 노이즈, 위치 편향, 증거 구성에서 MDMs의 강건성을 검토한다.

제안 방법

• denoising 단계와 계층에 걸쳐 MDMs의 어텐션 패턴을 정의하고 정량화한다.
• 토큰별 어텐션 가중치를 시각화하고 계층별 QK(쿼리-키) 기하학적 분해를 수행한다(노름 곱과 방향 코사인).
• Floating 토큰을 식별하고 이를 구조적 토큰과 어휘적 토큰으로 분류한다.
• 검색(heads)들을 분석하여 맥락 민감한 정보 흐름에서의 역할을 평가한다.
• 추론 중 입력 영역 간에 어텐션 흐름이 어떻게 이동하는지 추적하기 위해 영역 수준의 어텐션 흐름 분석을 수행한다.
• 검색된 컨텍스트가 있는 경우와 없는 경우 모두에서 지식 집약적 작업에서 MDMs와 ARMs을 비교한다.

실험 결과

연구 질문

• RQ1 MDMs에서 denoising 단계와 계층 전반에 걸친 어텐션 분포의 본질은 무엇인가?
• RQ2 ARMs의 싱크와는 다르게 floating 토큰은 어떤 토큰일 가능성이 높은가(구조적 vs 어휘적)?
• RQ3 어텐션 플로팅 메커니즘이 ARMs에 비해 맥락 내 학습과 강건성에 어떻게 기여하는가?

주요 결과

• MDMs는 Attention Floating을 나타내며, 지배적 주의 고정점이 고정된 싱크로 수렴하기보다 위치와 단계 전반에 걸쳐 떠다니는 경향이 있다.
• 얕은 계층은 떠다니는 구조적 토큰에 의지하여 글로벌 프레임워크를 형성하고, 더 깊은 계층은 의미적 콘텐츠를 향해 주의를 이동시킨다.
• 검색 헤드 분석은 깊이가 깊어질수록 더 깊고 콘텐츠 중심의 검색 헤드가 영향력이 커지는 것을 보여주며, 제안된 얕은 구조, 깊은 콘텐츠 메커니즘과 일치한다.
• MDMs는 지식 집약적 과제에서 검색된 컨텍스트로부터 더 큰 이득을 얻고, 여러 시나리오에서 검색이 있는 ARMs을 능가한다.
• MDMs는 맥락 노이즈, 위치 교란, 증거 분포에 대한 강건성을 보이며 스트레스 테스트에서 ARMs를 능가한다.