技術的サマリー：スパースオートエンコーダの機能マッチングと回路圧縮のためのセマンティック最適輸送

1. 問題定義

スパースオートエンコーダ（SAE）は、潜在空間を解釈可能な機能に分解することにより、大規模言語モデル（LLM）の解釈における主要なツールとして登場しました。しかし、SAE 分析の拡張において、2 つの重要な課題が残されています：

1. 層間機能マッチング：LLM の異なる層にまたがって、意味的に類似した機能を特定すること。既存の手法は通常、デコーダベクトルの比較に依存しています。このアプローチは、非連続的な層間で機能をマッチングする際に失敗します。なぜなら、デコーダベクトルは異なる活性化分布を再構成するように訓練されており、直接のベクトル比較が不正確になるためです。
2. 機能回路の圧縮：層を跨ぐ機能を接続する因果グラフである機能回路は、数百のノードを含むことが多く、手動による検査は非現実的です。現在の圧縮手法は、意味的に類似したノードを手動で「スーパーノード」にグループ化することに依存しており、このプロセスは拡張性がありません。

著者が特定した核心的な問題は、異なる活性化多様体に存在する機能間のセマンティック距離の推定です。既存のアプローチは、機能を単一のベクトル（例えばデコーダ重み）として扱うため、機能が「どこで」「どの程度強く」活性化するかという豊富な文脈を捨象し、情報損失と多様な多様体間比較の処理不能を招いています。

2. 手法

著者は、機能比較を**最適輸送（OT）**問題として再定義する分布フレームワークを提案します。

2.1 分布としての機能表現

単一のデコーダベクトルで機能を表現する代わりに、この手法は層 $\ell$ における各機能 $i$ を、それが活性化する隠れ状態上の活性化重み付き経験分布として表現します。

• トークン位置のコーパス $T$ が与えられた場合、機能 $i$ が最も強く活性化するトップ-$K$ のトークン位置を選択します。
• 選択された各トークン $t$ には、その活性化の大きさに比例する重み $w_{i,t}$ が割り当てられます。
• 機能は分布 $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$ として表現されます。ここで、$x_t$ はトークン $t$ における隠れ状態です。

2.2 共有参照空間

異なる層（異なる多様体に存在する）の機能を比較するために、この手法はこれらの層固有の分布を共有参照空間 $\mathcal{R}$ に射影します。

• 層間マッチングの場合、参照空間は通常、ターゲット層の隠れ状態空間です。ソース層からの活性化は、ターゲット層の対応するトークン位置に射影されます。
• 回路圧縮の場合、参照空間はすべての層にわたる隠れ状態の連結です。

2.3 ワーシュタイン距離によるセマンティック距離

共有空間に射影された後、2 つの機能間のセマンティック距離はワーシュタイン距離（アースムーバーズ距離）を用いて計算されます。

• 距離 $W_c(\hat{\mu}, \hat{\nu})$ は、ある機能の活性化分布の確率質量を他方の機能へ輸送するための最小コストを測定します。
• 基底コスト $c(z, z')$ は、共有空間内の 2 つの文脈表現をマッチングするための費用を測定します。
• このアプローチは、単なる点ごとのベクトル類似性ではなく、活性化パターンの幾何学的構造を捉えます。

2.4 下流タスク

• 機能マッチング：ターゲット機能に対して、この手法は最小ワーシュタイン距離（最隣接近傍）を持つソース機能を検索します。
• 回路圧縮：この手法は回路ノードを分布空間内の点として扱い、クラスタリング（例：凝集クラスタリング）を適用して、ペアごとのワーシュタイン距離が小さいノードをスーパーノードにグループ化し、クラスター内のセマンティック距離を最小化します。

3. 主要な貢献と理論的保証

本論文は、提案されたフレームワークに対して 3 つの主要な理論的保証を提供します：

1. 活性化スケーリング不変性：射影された機能分布は、活性化値のグローバルなスケーリングに対して不変です。表現は活性化の相対的な順位と比例分布（トップ-K 選択と正規化された重み）に依存するため、活性化の大きさにおける任意のスカラーシフトはセマンティック距離を変化させません。
2. 摂動に対する安定性：機能分布間のワーシュタイン距離は、分布の滑らかな摂動に対して安定しています。距離の変化は、基底コストのリプシッツ定数と、摂動された分布および元の分布間のワーシュタイン距離によって制限されます。これにより、トークン選択または SAE の再訓練におけるノイズに対する堅牢性が保証されます。
3. マージンに基づく回復：著者は、経験的な最隣接近傍マッチングとセマンティックグループ化が真の（オラクル）マッチを完全に回復する条件を確立しました。「マージン」（最良のマッチと 2 番目に良いマッチの間のギャップ）が推定誤差の 2 倍を超えれば、経験的な決定は保証されて正しいことになります。

4. 実験結果

この手法は、GPT-2 および Gemma-2-2B で訓練された SAE 上で評価され、SAE Match（デコーダベクトル L2）、Feature Flow（デコーダベクトルコサイン）、アトリビューションパッチング、LLM ベースの選択などのベースラインと比較されました。

• 機能マッチング：提案手法は、連続層、中間から最後、最初から最後の層設定におけるLLM 評価（セマンティック類似性評価）とマッチング精度において、すべてのベースラインを上回りました。重要なのは、クロスエントロピー損失や説明分散などの下流指標においても競争力のある性能を維持しており、単に類似した活性化トークンだけでなく、類似した機能的役割を持つ機能を特定していることを示した点です。特に、他の手法が失敗した微妙な「数字の加算」機能のマッチングに成功しました。
• 回路圧縮：この手法は、手動グループ化や他のベースラインと比較して、機能回路を直感的なスーパーノードに圧縮する精度が向上しました。異なる回路全体で一貫して同じ機能的役割を果たすモジュール的な機能グループ（例：「インダクションヘッド」や「トークン検出」）を正常に特定しました。
• アブレーション研究：実験により、機能を単一の重心ベクトルに縮約するよりも、分布として表現する方が、同じ多様体上に射影された場合でも、より多くのセマンティック情報を保持することが確認されました。

5. 意義と主張

本論文は、層間機能マッチングと回路圧縮の両方を、多様体間でのセマンティック距離の推定という同じ根本的な問題のインスタンスとして解決する統合セマンティックメトリックを提供すると主張しています。

• スケーラビリティ：機能をスーパーノードにグループ化するプロセスを自動化することで、人間のラベル付けへの依存を排除し、大規模な機能回路の分析を実行可能にします。
• 堅牢性：分布アプローチは機能活性化の「文脈」を捉えるため、層を跨いで機能が著しく進化した場合（非類似の分布）でも正確なマッチングを可能にします。これはベクトルベースの手法が失敗するシナリオです。
• 解釈可能性：複雑な回路を解釈可能なスーパーノードに自動的に圧縮する能力は、モデルのメカニズムの理解を深めます。「Dallas（地理）」「韻を踏む」「反対」といった回路のケーススタディがこれを示しています。

著者は、この研究を、ベクトル中心の機能比較の限界を超えた、スケーラブルで自動化された機械的解釈可能性への基礎的な一歩として位置づけています。