技術的サマリー：スペクトルプローブ回路

問題定義

事前学習済みトランスフォーマーのメカニズム的解釈可能性は、通常、事後分析に依存しています。つまり、モデルが完全に学習された後に、アテンションヘッドをアブレーション（無効化）し、固定された事前指定タスクにおける性能低下を観察することで、アテンションヘッド回路を特定します。このワークフローは計算コストが高く（すべての条件に対して完全に学習済みのネットワークを順方向に通過させる必要がある）、本質的に反応的です。なぜなら、探索を開始する前にターゲット能力を定義しなければならないからです。さらに、既存の手法は、特に「アテンションシンク」（最初のトークンにアテンションを向けるヘッド）が支配的な大規模モデルにおいて、一般のコンテンツ依存計算とタスク固有の因果回路を区別することにしばしば苦労しています。

手法：3 段階のレシピ

本論文は、ヘッドごとに動作し、タスクラベルや帰属勾配を必要とせず、学習中または最終チェックポイントに適用可能な、アテンションヘッド回路の特定のための 3 段階のレシピを提案します。

ステップ 1：ヘッドごとのスペクトル信号（一般指標）

最初のステップでは、ランダム行列理論に由来するスペクトル指標を用いて、持続的なコンテンツ依存計算を実行しているヘッドを特定します。

• 指標: 各トレーニングチェックポイントにおける各ヘッドについて、ヘッドごとのアテンション出力行列の**参加比率（Participation Ratio, PR）**を計算します。PR は、出力分布の有効ランク（次元数）を測定します。
• 軌道特徴: この信号の時間的進化を統合します：$I(L, H) = \sum_t \max(PR(L, H, t) - 1, 0) \cdot \Delta \log(\text{tokens}_t)$。
• 機能: この「PR 積分」は、時間を通じて高い有効ランク（多数の方向をspanする）を維持するヘッドをランク付けします。これは一般的な専門化の指標として機能し、低ランク（コンテンツ非依存）の状態に収束するか、ノイズのままであるヘッドを除外します。重要なのは、これがヘッドが「何」をしているかを区別するのではなく、専門的な計算を行っていることのみを示す点です。

ステップ 2：タスクパターンスクリーニング（タスク固有のフィルタ）

一般的なスペクトル信号は、標準的なターゲットに対するアテンションパターンを測定することで、特定のタスクに特化されます。

• メカニズム: 「全ヘッド」スクリーンは、特定のアテンションパターン（誘導、前トークン、重複トークン、先頭トークン/BOS、自己、局所など）に対する各ヘッドの選択性を測定します。
• 選択性の定義: 選択性は、ターゲット位置へのアテンションと一様ベースラインとの比率です。
• 閾値: ヘッドの選択性が閾値を超えた場合（例：割り当てについては $\ge 30\times$、誘導回路のメンバーシップについては $\ge 50\times$）、そのヘッドは分類されます。
• 機能: このステップは、一般的な「専門化」信号を、特定の能力に対する候補回路に変換します。これは、アテンションシンクが支配的なモデルにおいて、タスク固有のヘッドよりも汎用的なコンテンツ依存ヘッドを上位にランク付けしてしまうというステップ 1 の限界を克服します。

ステップ 3：因果的検証（反証可能な主張）

候補回路は、層の構成のアーティファクトではなく因果的な効果であることを保証するために、厳密な制御を伴うグループアブレーションを通じて検証されます。

• 手順: アテンション出力投影における残差ストリームから、ヘッドごとのスライスをゼロにします。
• 制御:
  1. マッチド・ランダム: スクリーンで選択されなかった、同じ層内の同等数のヘッドをアブレーションします。これにより、「これらの層の任意のヘッドをアブレーションすることが性能低下を引き起こす」という仮説を排除します。
  2. 上限: 選択された層内のすべてのヘッドをアブレーションします。これにより、その層における可能な最大の影響が確立されます。
• 基準: スクリーンで選ばれたヘッドが、マッチド・ランダム制御よりも有意に大きな性能低下を生み出し、かつ上限がゼロ付近で飽和する場合、その回路は因果的であるとみなされます。

主要な結果

1. アーキテクチャ間およびスケール間での再現性

このレシピは、8 倍のパラメータ範囲（51M から 1B アクティブ/7B 総パラメータ）および 3 つのアーキテクチャファミリー（密 GPT-NeoX、密 Llama スタイル、Llama スタイル MoE）で検証されました。これらは TinyStories、FineWeb、Pile、DCLM といった異なるデータパイプラインで学習されました。

• 誘導回路: 合成誘導タスクにおいて、すべてのテストモデルで 3〜6 ヘッドの小さな誘導回路が因果的に必要であることが判明しました。
• 亜線形スケーリング: 総ヘッド数はモデルサイズに対して線形に増加しますが、誘導回路のサイズは亜線形（3〜11 ヘッド）のままです。これは、特定の能力がモデル容量に対して線形にスケーリングしないことを示しています。

2. 保存された専門化ヘッドの割合

8 倍のスケール範囲（124M から 410M）において、識別可能な専門的計算を実行するヘッドの割合は、**約 17〜19%**で保存されています。

• この割合は、(1) トップ K のスペクトル選択を能力クラスに分類する方法と、(2) 空分布のノイズフロア以上の誘導選択性を持つヘッドを数える方法の 2 つの手法を用いて独立して回復されました。
• より大規模な 1B クラスのモデルでは、この割合はより大きく変動し（4.7% から 27.3%）、モデルの規模ではなく BOS 吸引子（アテンションシンク）の支配性と相関しており、保存性は特定のアーキテクチャ/データファミリー内の性質であることを示唆しています。

3. 予測力とシード特異性

• シード特異性: 6 つの異なるランダムシードで学習された 51M パラメータモデルにおいて、このレシピは同じタスクに対して異なるシード固有の回路を特定しました。シード A に対して特定された回路をアブレーションしても、シード B のモデルには影響を与えませんでした。これは、この手法が普遍的なパターンではなく、モデルの実際の実装を見つけていることを確認しました。
• 予測タイミング: スペクトル信号（ステップ 1）は、能力の選択性が機能的閾値を越える前に上昇します。このレシピは、能力の完全な出現に先立ち、トレーニングトークンのわずか一部（0.3〜2%）のみを使用して回路を特定できます。

4. データシフトに対する頑健性

FineWeb で学習し、その後 OpenWebText で継続学習されたモデルにおいて、FineWeb 終点で特定された誘導回路は、データシフトを通じて不変（Jaccard 係数 = 1.00）でした。ヘッドは選択性と PR 軌道を維持しており、この手法がコーパス固有の統計ではなく構造的専門化を特定していることを示唆しています。

5. 構成的依存関係

この手法は、「前トークン」回路（8〜9 ヘッド）を正常に特定し、その誘導における因果的役割を実証しました。前トークン回路をアブレーションすると誘導性能が破壊され、前トークンヘッドが誘導ヘッドに使用されるキーベクトルを構築するという標準的な構成的構造が確認されました。

意義と主張

本論文は、3 論文からなるプログラムの方法論的アンカーとして位置づけられています。その主な主張は控えめであり、レシピ自体の有効性に焦点を当てています。

1. レシピはレシピであり、万能スクリーンではない: 著者は明示的に、ステップ 1（スペクトル信号）はタスク固有の検出器ではなく、一般的な専門化指標であると述べています。ステップ 2（スクリーン）がタスク固有性を提供し、ステップ 3（アブレーション）が因果的な反証可能性を提供します。この組み合わせが手法の頑健性を生み出します。
2. 因果的主張の誠実な枠組み: 論文は、「マッチド・ランダム」の差分が因果的主張を誠実にするものであることを強調しています。この制御なしには、特定の回路と一般的な層効果を区別することはできません。
3. 発達的先行: スペクトル信号は能力選択性の出現に先行しており、事前学習中の回路形成を監視するための潜在的な窓を提供します（これは関連論文で詳述されています）。
4. 限界: 著者は、構成されたタスク（間接目的語の識別など）の場合、主要なアテンションパターンが異なるモデルファミリー間でタスク因果構造から分離する可能性があり、より広範な候補スクリーンのセットが必要であると指摘しています。さらに、この手法は固定評価バッチに依存しており、特定の回路ヘッドは再学習によって変動する可能性があります（ただし、専門化ヘッドの割合は安定しているようです）。

要約すると、本論文は、小さく因果的に必要なアテンションヘッド回路を分離するための再現性のあるラベルフリーのフレームワークを提供し、モデル容量の保存された割合が、モデルサイズに対して亜線形にスケーリングする専門的計算に割り当てられていることを実証しています。