[論文レビュー] The Hidden Attention of Mamba Models

Ameen Ali, Itamar Zimerman+1 more|arXiv (Cornell University)|Mar 3, 2024

Explainable Artificial Intelligence (XAI)被引用数 5

ひとこと要約

論文は Mamba の選択的状態空間層を自己注意の模倣として再フレーミングし、隠れた注意行列を導出し、Mamba をトランスフォーマーと比較して XAI アプリケーションを可能にする Explainability ツールを開発する。

ABSTRACT

The Mamba layer offers an efficient selective state space model (SSM) that is highly effective in modeling multiple domains, including NLP, long-range sequence processing, and computer vision. Selective SSMs are viewed as dual models, in which one trains in parallel on the entire sequence via an IO-aware parallel scan, and deploys in an autoregressive manner. We add a third view and show that such models can be viewed as attention-driven models. This new perspective enables us to empirically and theoretically compare the underlying mechanisms to that of the self-attention layers in transformers and allows us to peer inside the inner workings of the Mamba model with explainability methods. Our code is publicly available.

研究の動機と目的

• Mamba の選択的状態空間層の内部情報フローのダイナミクスを明らかにする。
• Mamba が暗黙的で因果的な自己注意機構として見ることができることを示す。
• 隠れた注意行列に基づく Mamba 用の explainability ツールを開発・適応する。
• 解釈性と表現力の観点から Mamba ベースの注意をトランスフォーマーの注意と比較する。
• 状態空間モデルにおける注意の進化と能力に関する理論的・経験的洞察を提供する。

提案手法

• 選択的状態空間層をデータ制御型の線形作用素として再定式化し、隠れた注意行列を抽出する。
• チャネルごとの隠れた注意行列を導出し、それをデータ依存の作用素を介して従来の注意に関連づけ、解釈性のための簡略化を含む。
• Attention Rollout を Mamba に適用して、レイヤーとチャネルを横断するクラス非依存の説明可能性を実現する。
• Mamba の注意行列とゲーティング信号を用いるよう Transformer-Attribution を適応させ、Mamba 固有の注意ベースのアトリビューションを開発する。
• 視覚と言語タスクにおける隠れた注意行列を可視化し、比較する。
• ImageNet などの関連データセットで撹乱とセグメンテーション指標を用いて explainability 手法を評価する。

実験結果

リサーチクエスチョン

• RQ1 選択的状態空間層（Mamba）は、トランスフォーマーの自己注意に類似した暗黙の注意機構として解釈できるのか。
• RQ2 Mamba における隠れた注意の性質と規模は何であり、それらはトランスフォーマーの注意とどのように比較されるのか。
• RQ3 Mamba でクラス非依存・クラス特異的な説明可能性のために隠れた注意行列をどのように導出・活用できるか。
• RQ4 Mamba ベースの説明可能性手法は、視覚と自然言語処理タスクの両方でトランスフォーマーベースの手法と同等の解釈性を達成するのか。
• RQ5 状態空間モデルにおける注意の表現力と進化を説明する理論的洞察は何か。

主な発見

• Mamba 層はデータ制御型線形作用素を通じて暗黙的な因果的自己注意機構として再定式化できる。
• 単一の Mamba チャンネルは多数の隠れた注意行列を生成し、トランスフォーマーよりもはるかに多くの注意マップを生み出す。
• 隠れた注意行列は、特定の設定でトランスフォーマーをベースとした説明可能性と比較可能な新しい explainability ツールを可能にする。
• Mamba ベースの注意は、類似のモデルサイズで評価した場合、トランスフォーマーに対して競争力のある explainability 指標を示す。
• 理論分析は、データ制御型の非対角混合器が Mamba の表現的注意機能とコンテキスト内学習の潜在能力を支えていることを示唆する。