論文紹介：Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers

Transcript

1. 因果的ヘッドゲート機構： Transformerにおける注意ヘッドの役割を解釈するためのフレームワーク Andrew J. Nam, Henry C. Conklin, Yukang Yang, Thomas L. Griffiths, Jonathan D. Cohen∗, Sarah-Jane Leslie∗ 1
2. ⽬次 o研究背景 o従来⼿法 o提案⼿法 oCHG o実験設定 o実験結果 o拡張 o考察・限界・応⽤可能性 o結論・今後の展望 2
3. 研究背景: LLMのどのヘッドが何をしている? • 解釈可能性を追求する研究分野 • LLMは強⼒だが内部機構は不透明で信頼性に⽋ける • Attention Headは⼈間可読な役割を⽰唆されている • しかし役割の因果性の同定は難しい • 本当にそのヘッドが必要/有害︖ • 仮説を⽴てて検証するスタイルは有効であったが、 ある程度合っている仮説を⽴てることが難しい 3
4. 従来⼿法: 1. 観察ベース • プロービング…隠れベクトルから概念が線形に読めるか︖ • 可視化…Attentionパターンの可視化 2. 介⼊ベース モデルの内部動作に能動的に介⼊する • CMA(因果媒介分析)…ペアの⼊⼒間で内部成分にパッチ →出⼒変化で因果性を確認 • ターゲットアブレーション…仮説ヘッドを無効化して性能差 を測る 4
5. おまけ: CMA(=Causal Mediation Analysis) ⼊⼒Aと出⼒Yの間で特定の内部成分Mが本当に因果的に 機能しているか︖ 1. 対になったプロンプトを⽤意 2. ソース例の内部活性Mをターゲットへパッチ 3. 予測が望む⽅向へ変化したらMが媒介している 5 ソース ターゲット LLM 「フランスの⾸都は︖」 「⽇本の⾸都は︖」 対象ヘッドの活性 パリ LLM
6. 従来⼿法: 限界 1. 観察ベース … 相関⽌まり、 ”使⽤されている”保証がない 2. 介⼊ベース … テンプレ設計の労⼒、多様なタスクではスケール不全 6 仮説なし+次トークン予測(NTP)だけで、 スケーラブルにヘッドの因果的役割を同定する 目的
7. 提案⼿法 CHG(Causal Head Gating): • ⽬的︓ • LLMの各Attention Headがタスクに与える因果的な 役割を、仮説や⼿作業なしで特定したい • 発想︓ • 各Headの出⼒にON/OFFのゲートを取り付け、 そのゲートの開き具合を学習させることでHeadの 重要度を測る 7
8. 提案⼿法 CHG(Causal Head Gating): 具体的な仕組み • ゲートの導⼊︓ • 各ヘッド出⼒ 𝑍!,# は𝑨𝒍,𝒉 𝑽𝒍,𝒉 に学習可能なスカラーゲート 𝑮𝒍,𝒉 ∈ [0,1]を乗算して計算される • モデル本体の凍結︓ • LLM本体の重みは凍結する • 学習対象はヘッドごとに追加したゲートのパラメータ (1head-1pram) 8
9. 提案⼿法 CHG(Causal Head Gating): 具体的な仕組み • ⽬的関数︓ • 次トークン予測(NTP)の性能最⼤化 • 次トークン予測の負の対数尤度(NLL)+ 正則化 • 2⽅向の最適化︓ • 学習を正則化項の係数λの符号を変えることで2つの⽅向で⾏う 1. 保持⽅向 𝑮 + (λ>0) ゲートを開ける圧⼒をかけ性能維持に貢献するヘッドを探る 2. 除去⽅向 𝑮 − (λ<0) ゲートを閉じる圧⼒をかけ性能維持に不可⽋なヘッドを探る 9
10. 役割の因果定義 11 CHG(Causal Head Gating): 役割の分類 • Facilitating(促進)… 𝑮­ が⾼い=抑制圧でも残る＝必要 • Interfering(妨害)…1- 𝑮 + が⾼い=保持圧でも抑え込まれる＝有害 • Irrelevant(無関係)… 𝑮 +⊙(1­𝑮­)が⾼い=性能無影響で両極端へ分岐
11. 実験設定: モデル・データセット 12 • モデル︓Llama 3系（1B, 3B, 3B-instruct, 8B） • タスク︓ • Math（OpenMathInstruct-2） • Syntax（BIG-Bench “syntax”） • Commansense（CommonsenseQA） • 学習︓各モデルxタスクにつき10seedsでCHGを学習 • 最適化とハイパラ︓ • モデル凍結、ゲートのみ学習（1param/head） • Optim︓AdamW、ステップ1000~3000 • λ=0→>0 , λ=0→<0の順に実⾏（スタートラインは同⼀） https://zenn.dev/mutex_inc/articles/2948e2c5934171
12. 実験設定: 評価指標・検証 13 • 逐次アブレーション︓ • 役割スコアが⾼い順にゲート無効化、Δlog − 𝑝𝑟𝑜𝑏で検証 • 分布/相関︓ • モデル・タスクによる役割スコアの分布や相関（54ペア） • ⼀貫性︓ • 10回の学習試⾏のうち… • 同じ役割に分類されたヘッド︓Always • 少なくとも1回はその役割に分類されたヘッド︓Any • ⽐較︓ • CMA（従来⼿法）の結果とCHGのスコアの⼀致度を⽐較
13. 実験結果: 14 ①妥当性 ~逐次アブレーション曲線~ • 促進ヘッドを消すと性能低下/無関係は変化⼩/妨害は改善 • 複数モデル・複数タスクで⼀貫している 縦軸︓ターゲット対数確率の変化 横軸︓0化するヘッドの数
14. 実験結果: 15 ①妥当性 ~逐次アブレーション曲線~ • 促進ヘッドを消すと性能低下/無関係は変化⼩/妨害は改善 • 複数モデル・複数タスクで⼀貫している 縦軸︓ターゲット対数確率の変化 横軸︓0化するヘッドの数 CHGが算出したスコアは、 ヘッドの因果的な役割を正しく捉えていそう
15. 実験結果: 16 ②スコア分布(ECDF) 縦=しきい値以上のヘッド割合 横=スコア。 右に張り付くほど「⾼スコアのヘッドが多い」
16. 実験結果: 17 ②スコア分布(ECDF) • 54ペアのモデル間相関︓min≈94.4%/mean≈99.2% 縦=しきい値以上のヘッド割合 横=スコア。 右に張り付くほど「⾼スコアのヘッドが多い」 促進スコアが広く⾼い ＝計算に広い下位回路が必要? 無関係多数
17. 実験結果: 18 R=妨害ヘッド(1-G+) G=促進ヘッド(G-) ③役割解明に向けて（Layer x Head） • AlwaysとAnyの可視化で安定ヘッド と可変ヘッドを識別 Math(Always)︓ 促進ヘッドが特定の層に集中している Math(Any)︓ 多くのヘッドが促進や妨害の役割を担う →安定した回路と代替可能な周辺回路が 存在する可能性
18. 実験結果: 19 ④定量表: Always/Any • Syntax/Commonsense︓ • Always促進は<5%程度、Always妨害は ≈0% • Math︓ • Always促進が顕著（Max 38.3%） • Anyでは多数ヘッドが⼀度は促進/妨害に該当=低モジュール性 F:促進 N:妨害
19. 実験結果: 20 ⑤CMAとの整合性（Function Vectors） • CMAで有意と判定されたヘッドはCHG促進スコアが有意に⾼い↑ 縦軸︓CHGの促進度 横軸︓CMAの因果効果
20. 実験結果: 21 ⑤CMAとの整合性（ABA/ABB: 抽象規則） 縦軸︓CHGの促進度 横軸︓CMAの因果効果 • CMAで効くと分かったヘッドはCHGでも “残すべきヘッド”と⼀致 • Induction︓xが⼤きくかつy≈1に密集 → 強い因果効果を持つヘッドが、CHGで も強い促進 • Insignificantはx≈0に密集しばらける • CHG（⾼感度に候補を拾う）×CMA（ ⾼精度に確証）の補完関係を⽰す。
21. 実験結果: まとめ 22 • CHGスコアはアブレーションで予測通りの挙動をする • モデル間で⼀貫、ただしタスクにより回路の広がりが異なる • Mathは広い • CMAと整合 • CHG（⾼感度）x CMA（⾼精度）は補完的
22. 拡張: CCHG 23 課題 機能的分解「どのサブ機能を⽀えている︖」 • 複雑タスク（Math/CoT）は複数のサブ機能で成り⽴つ • CHGはヘッドの3分類「効く/邪魔/無関係」はできたが 何を計算しているか直接わからない
23. 拡張: CCHG概要 24 課題 機能的分解「どのサブ機能を⽀えている︖」 • 複雑タスク（Math/CoT）は複数のサブ機能で成り⽴つ • CHGはヘッドの3分類「効く/邪魔/無関係」はできたが 何を計算しているか直接わからない CCHG(Contrastive Causal Head Gating) • 保持データ𝐷& と忘却データ𝐷' を⽤意 • ⽬的関数︓ 𝐿()* 𝐺 = 𝐿 𝐺; 𝐷& , 𝜆 − 𝐿 𝐺; 𝐷' , 𝜆 (𝜆 < 0)を最⼩化 → 𝐿 𝐺; 𝐷& , 𝜆 は⼩さく、 𝐿 𝐺; 𝐷' , 𝜆 は⼤きくなるようにゲート調整 • 結果︓忘却タスクのみ必要なヘッドのゲート→閉じる • 保持タスクのみ必要なヘッドのゲート→開いたまま • 両タスクで共通のヘッド→バランス調整 • 例︓ICL形式 vs Instruction形式のプロンプトで タスク推論回路と実⾏回路を分離
24. 拡張: CCHGの設定と結果 25 実験結果 • 6タスクで双⽅向にCCHGを学習、テストデータで汎化検証 • 忘却側のプロンプトでは精度崩壊≈0%、保持側は概ね維持 • InstructionとICLの回路は分離可能だが部分共有も⽰唆 Default: 何も変更していない元のLLM Gated: ゲートの開閉を学習させたLLM 上の段: Instruction形式の能⼒を保持しK-Shotを忘却 下の段: K-Shot(ICL)形式を保持しInstruction形式を忘却
25. 考察・限界・応⽤可能性 26 考察 • 低モジュール性…ヘッド役割は他ヘッドの組に依存し変動 • 複数の⼗分回路（低重複）＝冗⻑で代替可能 • 形式依存の回路選択（ICL vs Instruction） 限界 • CHGは「何を計算しているか」の特定は弱い • 稀にG+,G-が乖離（=相互作⽤の複雑性） • Run間の変動→複数フィット&Always/Any集計が必要 応⽤可能性 • 研究︓CHGで候補抽出→CMAで意味づけというパイプライン活⽤ • 安全・制御︓妨害ヘッドの選択的抑制/能⼒のピンポイント編集
26. 結論・今後の展望 27 結論 • CHG︓ラベル/テンプレ不要・NTP直結・スケーラブル • 促進/妨害/無関係という因果分類で、回路の分布・⼀貫性を証明 • CCHG︓提⽰形式などでサブ回路の分離が可能 Next Step • 意味づけの⾃動化（CMA・機能検査のハイブリッド） • LLM・多⽤タスクへの拡張 • 編集/制御への接続（オンライン抑制、スイッチング）