技術要約：相互作用を考慮したバックプロパゲーションによる勾配ベースの回路局在化の修正

問題提起

回路局在化（Circuit localization）は、大規模言語モデル（LLM）における特定の振る舞いに責任を持つ特定のコンポーネント（例：アテンションヘッド、ニューロン、接続）のサブセットを特定することを目的としています。既存のほとんどの手法は、各コンポーネントが独立して作用するという仮定の下で動作しており、各コンポーネントを単独で摂動させることで重要性を推定しています。

しかし、ニューラルネットワークのコンポーネントは相互に作用しており、これらの相互作用を無視することは、コンポーネントの重要性の系統的な誤推定につながります。著者らは、**アテンションの自己修復（attention self-repair）**と呼ばれる、具体的かつ問題のある相互作用を特定しました。この現象では、複数のポジションが大きなアテンション重みを持ち、かつ類似した情報（バリューベクトル）を含んでいる場合、単一のアテンションスコアをアブレーション（除去）しても出力への影響は最小限に留まります。これは、ソフトマックス活性化関数が、残りの冗長なポジションへと重みを再分配することで補償を行うために発生します。その結果、標準的な摂動ベースおよび勾配ベースの説明手法は、これらのアテンションスコアに寄与するコンポーネントの重要性を過小評価し、しばしば影響力のあるスコアに対して勾配消失を引き起こします。

手法：勾配相互作用修正（Gradient Interaction Modifications: GIM）

本論文では、複数の相互作用するコンポーネントを同時に摂動させた際の影響を推定するために設計された、勾配ベースの属性手法である**勾配相互作用修正（GIM）**を導入しています。GIMは、以下の3つの補完的なテクニックを通じて、標準的なバックプロパゲーションを修正します。

1. 温度調整済みソフトマックス勾配（Temperature-adjusted Softmax Gradients: TSG）:

   • 動機: アテンションの自己修復への対処。標準的な勾配では、あるソフトマックス重みを高めることによる利得が、他の重みを抑制することによる損失によって相殺されるため、勾配が消失します。
   • メカニズム: バックプロパゲーション中、TSGは勾配を計算する前に、高い温度（$T > 1$）を用いてソフトマックスを再計算します。これにより確率分布が平坦化され、大きな重みが縮小し、小さな重みが拡大します。これにより「重要な」ポジションの集合を拡大することで、これらのポジション間で寄与が同一であるという仮定を緩和します。これにより、ペアごとの勾配の差（$g_j - g_k$）が消失しないようにし、勾配信号を生存させ、同時摂動の効果を近似することを可能にします。

2. レイヤー正規化のフリーズ（Layernorm Freeze）:

   • 動機: レイヤー正規化はゼロサムのダイナミクスを導入します。つまり、一つのコンポーネントをアブレーションすると活性化の大きさが減少し、正規化係数が残りのコンポーネントを増幅させます（自己修復の一形態）。
   • メカニズム: 正規化係数（$\sigma$）をバックプロパゲーション中に定数として扱います。これにより、正規化による再スケーリングが発生する前の、各コンポーネントの直接的な寄与を測定し、補償的な増幅によって引き起こされる重要性の過小評価を防ぎます。

3. 勾配ノルム（Gradient Norm）:

   • 動機: 乗算操作（例：$A = QK^T$ であり、$Q$ と $K$ の両方が入力 $X$ に依存する場合）は、標準的な一次の勾配において二重カウントを引き起こします。$X$ を摂動させると $Q$ と $K$ の両方に同時に影響を与えますが、標準的な連鎖律はそれらの個別の効果を合算するため、変化を過大評価してしまいます。
   • メカニズム: 乗算相互作用の勾配を、乗算入力の数（例：$Q \times K$ の場合は2で割る）で除算します。これにより、「一様則（uniform rule）」に適合させ、アテンション・バリュー乗算、クエリ・キー乗算、およびMLPゲート投影乗算における寄与の過剰カウントを修正します。

主な貢献

1. GIMフレームワーク: 特徴量の相互作用を明示的に考慮し、メカニスティック解釈可能性ベンチマーク（MIB）の回路局在化トラックにおいて最先端の性能を達成した勾配ベースの属性手法の導入。
2. アテンション自己修復の特定: アテンションの自己修復（ORゲートのような相互作用）が、影響力のあるアテンションスコアの勾配を消失させ、コンポーネントの重要性の系統的な過小評価を招くことを形式的に特定し、証明したこと。
3. 特徴量属性の優位性: GIMが、多様なタスクとLLMアーキテクチャにおいて、既存の勾配ベースの手法を上回る特徴量属性を実現することを実証。

結果

• 回路局在化: MIBにおいて、GIMは複数のモデル（GPT-2, Qwen-2.5, Gemma-2, Llama-3.1）およびタスク（間接目的語識別、算術、MCQA、ARC）にわたり、既存の手法（EAP, EAP-IG, Hybrid-ens, IFRを含む）を大幅に上回りました。特に多段階の推論を必要とするタスクにおいて、最も高い平均回路性能比（CPR）を達成しました。
• 特徴量属性: 7つの指示チューニング済みLLMと6つのデータセット（QA、事実検証、感情分析、ヘイトスピーチ検出をカバー）において、GIMは一般的にGradientXInput, Integrated Gradients, DeepLIFT, TransformerLRPと比較して、より高い包括性（comprehensiveness）および十分性（sufficiency）スコアを達成しました。
• 自己修復の検証: 実証的分析により、アテンションの自己修復が実際のLLMにおいて普及しており、入力あたり平均65〜1200回発生することが確認されました。2つの大きなアテンションスコアを共同でアブレーションすると、それらの個別の合計よりも有意に大きな効果が生じることが示され、理論的分析が検証されました。
• アブレーション研究: GIMからいずれかの修正（TSG, Layernorm Freeze, または Gradient Norm）を取り除くと、ほとんどのモデルとデータセットの組み合わせにおいて性能が低下し、各コンポーネントが手法の忠実性に意味のある貢献をしていることが確認されました。

意義

本論文は、GIMが特徴量の相互作用によって引き起こされるコンポーネントの重要性の系統的な過小評価を修正することにより、LLMのより忠実なメカニズム分析を可能にすると主張しています。コンポーネントがどのように相互作用するか（具体的には、ソフトマックスによる再分配と正規化による再スケーリングを通じて）を明示的にモデル化することで、GIMはモデルの振る舞いを駆動する内部メカニズムのより正確な姿を提供します。著者らは、GIMを単なる漸進的な改善ではなく、現在の解釈可能性ツールに蔓延している「独立性の仮定」に対する必要な修正であり、これまで消失勾配によって隠されていた回路を特定することを可能にするものとして位置づけています。この手法は、より広範な採用を促進するために、オープンソースのPythonパッケージとして公開されています。