技術要約：因果回路発見のためのマルチグラニュラリティ・ノード・プルーニング

1. 問題提起

因果回路の発見（Causal circuit discovery）は、大規模言語モデル（LLM）内で特定の振る舞いを因果的に駆動する最小限のサブネットワークを特定することを目的としている。既存の手法はメカニスティックな解釈可能性において進展しているものの、主に以下の2つの限界に直面している：

1. 計算のスケーラビリティ： 現在のアプローチは主にエッジ・プルーニング（コンポーネント間の接続の削除）に依存している。トランスフォーマーにおけるエッジの数はノード数に対して二次関数的（$O(|N|^2)$）に増加するため、エッジレベルのプルーニングはモデルのサイズが増大するにつれて計算量的に実行不可能となる。さらに、これらの手法は、分離された残差ストリームの表現や中間活性化を保持する必要があり、メモリ使用量が極めて高くなる。
2. グラニュラリティ（粒度）の制限： 既存の手法は通常、より粗い粒度で動作し、アテンションヘッド全体やMLPブロックを原子的な単位として扱う。これは、特定のニューロン（ノード）がターゲットとなる振る舞いに責任を持っている可能性を見落としており、結果として回収された回路が微細な構造を見逃したり、不要なコンポーネントを保持したりすることになる。

2. 手法

著者らは、単一の統一された最適化ステップ内で、複数の粒度レベルにわたって最小限の回路を発見するマルチグラニュラリティ・ノード・プルーニング・フレームワークを提案している。

2.1 階層的グラニュラリティ設計

エッジをプルーニングする代わりに、本フレームワークは以下の4つの異なるレベルにおけるノードに対するマスクを学習する：

• ブロック・レベル： アテンション・ブロック全体およびMLPブロック。
• ヘッド・レベル： 個々のアテンション・ヘッド。
• ニューロン・レベル： アテンション層内の個々のニューロン、およびMLPブロックの2つの線形層（MLP1/KeysおよびMLP2/Values）。

2.2 二系統因果プルーニング設定

単純なゼロ・アブレーション（zero-ablation）に伴うリスク（分布シフトの発生）を回避するため、本手法は二系統のフォワードパスを採用している：

• クリーン・ストリーム（Clean Stream）： 特定の振る舞いを引き出すように設計された入力（例：西暦が50年より大きいと予測させる入力）を処理する。
• 破損ストラーム（Corrupted Stream）： 異なる予測を引き出すために、入力を最小限に摂動させたバージョン（例：ジェンダー代名詞の反転や開始年の変更）を処理する。
• 補間（Interpolation）： 各ノード $n_i$ について、学習可能なマスク $m_i \in [0, 1]$ が、クリーンな活性化（$h^{clean}$）と破損した活性化（$h^{corrupted}$）を補間する：
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  最適化は、正しい出力（クリーンな振る舞いに一致するもの）を生成させつつ、スパース性を最大化するようなマスクを見出すことを目的とする。

2.3 最適化とマスクのパラメータ化

• ハード・コンクリート分布（Hard Concrete Distribution）： マスクは、勾配ベースの最適化を可能にしつつ、近似的にバイナリ（0または1）の挙動を誘導するために、ハード・コンクリート分布を用いてパラメータ化される。
• 統一された目的関数： 学習目的関数は、以下の3つの要素を組み合わせている：
  1. 忠実度損失（$L_{faith}$）： プルーニングされた回路の出力分布がベースモデルと一致することを保証するKLダイバージェンス。
  2. タスク損失（$L_{task}$）： 回路が正解（ground-truth）とディストラクター（distractor）を正しく区別することを保証するマージンベースのランキング損失。
  3. スパース性正則化項（$L_0$）： マスクをゼロにすることを促し、回路サイズを最小化する。
• 階層的一貫性： 学習後、親ノード（例：MLPブロック）がプルーニングされた場合、その子ノード（例：そのブロック内のニューロン）もすべて無効化されるという制約の下で、マスクを二値化する。

3. 主な貢献

1. マルチグラニュラリティ・フレームワーク： トランスフォーマーのブロック、アテンション・ヘッド、MLPブロック、および個々のニューロンにわたって、同時にマスクを学習する統一された最適化フレームワークの導入。
2. スケーラビリティと効率性： エッジ・プルーニングの二次的なスケーリングや、すべてのエッジに対して中間活性化を保持する必要性を回避する。これにより、GPUメモリ使用量を従来手法と比較して3倍から11倍削減し、単一の40GB GPUを使用してLlama 3.1-8B上での回路発見を可能にした。
3. 微細な発見： 本アプローチは、ベースラインよりも大幅にコンパクトな回路を特定する。具体的には、従来のメソッドがほとんど達成できなかったMLPブロックおよびニューロンレベルにおいて、高いスパース性を実現している。

4. 実験結果

本手法は、GPT-2、GPT-2-XL、およびLlamaモデルを用い、3つの標準的な回路発見タスク（間接目的語識別（IOI）、ジェンダー化された代名詞（GP）、および「Greater Than（GT）」）で評価された。

4.1 スパース性とコンパクトネス

• MLPおよびニューロンのプルーニング： 提案手法は、エッジ・プルーニングのベースライン（EAP、EP）および非構造化重みプルーニング（DiscoGP）を大幅に上回る性能を示した。最も不利な設定においても、従来の手法よりも33.34%多いMLPブロック、および59.8%多いニューロンをプルーニングした。
• タスク依存の構造： 発見された回路は、タスクに基づいて明確なトポロジーを示す：
  • IOI： 分散型のMLP計算が支配的であり、最終層を除いてアテンションは大部分がプルーニングされる。
  • ジェンダー化された代名詞 (GP)： 高度に局所的かつ不連続である。MLP計算は少数の層に限定され、アテンションは疎である。
  • Greater Than (GT)： 不活性な層の連続したブロックと「スキップのような」実行パターンを伴う、極端なスパース性を示す。

4.2 忠実度と一貫性

• パフォーマンス： プルーニングされた回路は、高いタスク精度（しばしば>95%）と、フルモデルに対する低いKLダイバージェンスを維持している。
• 一貫性： 5つのランダムシードにわたる実験により、発見された回路（忠実度指標およびスパース性プロファイル）は非常に安定していることが示された。
• 定性的分析： Llama 3.2-1BにおけるIOIタスクの生存アテンション・ヘッドは、GPT-2で手動発見された回路と同様の特定の機能的役割（IO-Movers、S-Inhibitors）を担っていることが判明しており、本手法がメカニスティックな洞察を回収する能力を検証している。

4.3 計算効率

• メモリ： 本手法は、エッジ・プルーニングのベースライン（例：GPT-2におけるEAPの72.79 GBに対し、6.27 GB）よりも少なくとも3倍少ないメモリを必要とする。
• 時間： 学習は6分未満で完了し、エッジ・プルーニング手法（約50分）よりも大幅に速い。ただし、メモリ使用量を犠牲にして高速化を図った単一パスのEAP法よりはわずかに遅い。

5. 意義と主張

本論文の主な意義は、現在の回路発見におけるスケーラビリティとグラニュラリティのボトルネックを克服することにある。エッジ・プルーニングからマルチグラニュラリティ・ノード・プルーニングへと移行することで、著者らは以下を実証している：

• 粗い粒度の手法によって保持されている多くのニューロンは不要であり、パフォーマンスへの影響を最小限に抑えつつ削除可能であること。
• より大きなモデル（例：Llama 8B）は、より小さなモデルと同等、あるいはより強力なプルーニングをサポートしており、これはモデルのスケールが増大しても、タスク固有の回路がより密になるのではなく、むしろ冗長性が導入されることを示唆していること。
• 本手法は、大規模モデルにおいて微細な因果回路を発見するための計算可能な経路を提供し、異なる振る舞いがアテンション、MLP、およびニューロン活動の異なるパターンによって符号化されていることを明らかにしている。

著者らは、一つの制限事項として、本フレームワークはどのノードが必要であるかは特定するが、生存したノード間の**相互作用構造（エッジ）**を明示的に回収するものではないことを認めている。彼らは、ノード選択とエッジ回収を組み合わせることで、相互作用構造を推論するという将来の研究方向を示唆している。