[論文レビュー] Revealing the Attention Floating Mechanism in Masked Diffusion Models

Xin Dai, Pengcheng Huang+7 more|arXiv (Cornell University)|Jan 12, 2026

Generative Adversarial Networks and Image Synthesis被引用数 0

ひとこと要約

論文はMasked Diffusion Models (MDMs)におけるAttention Floating現象を明らかにし、浅い構造意識と深い内容志向のアテンションパターンを示して文脈内知識の利用と頑健性を高め、 retrieved context を活用した場合には autoregressive models (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.

研究の動機と目的

• Masked Diffusion Models (MDMs) のノイズ除去中のアテンション挙動を調査する。
• Attention Floating 現象を特徴づけ、それを ARMs のアテンションシンクと対比する。
• MDMs におけるアテンションダイナミクスが文脈内学習と知識利用にどう寄与するかを理解する。
• 文脈的ノイズ、位置バイアス、エビデンス配置下での MDMs の頑健性を検証する。

提案手法

• MDMs の denoising ステップと層全体でのアテンションパターンを定義・定量化する。
• トークンごとのアテンションウェイトを可視化し、層ごとの QK (クエリ-キー) 幾何分解を行う（ノーム積 vs. 方向性余弦）。
• 浮遊トークンを識別し、それらを構造的トークンと語彙的トークンに分類する。
• retrieval heads を分析して文脈依存の情報流れにおける役割を評価する。
• 推論時に入力領域間でアテンションがどのように移動するかを追跡する領域レベルのアテンションフロー分析を実施する。
• retrieved context を用いた場合の知識集約タスクで MDMS と ARMs を比較する。

実験結果

リサーチクエスチョン

• RQ1 MDMs における denoising ステップと層全体でのアテンション分布の性質は何か？
• RQ2 浮遊トークンは ARMs の sink とどう異なり、どのようなトークン（構造的 vs 語彙的）になりやすいか？
• RQ3 アテンション・フローティング機構は ARMs と比べて文脈内学習と頑健性にどのように寄与するか？

主な発見

• MDMs は Attention Floating を示し、支配的なアテンションのアンカーが位置とステップ間で固定的な sink に収束することなく漂う。
• 浅い層は浮遊する構造トークンを用いてグローバルな枠組みを形成し、深い層は意味的内容へアテンションをシフトする。
• retrieval-head の分析から、深い層ほど内容志向の retrieval heads が影響力を増し、提案された浅い構造・深い内容の機構と整合する。
• MDMs は知識集約的タスクで retrieved context からの利得が大きく、 retrieved を用いた場合に ARMs を複数のシナリオで上回る。
• MDMs は文脈的ノイズ、位置の乱れ、エビデンス分布に対して頑健であり、ストレステストで ARMs を上回る。