技術要約：統合ポリシーグラジエントによるLLM推論の解釈と制御

問題提起

大規模言語モデル（LLM）は洗練された推論能力を示しているが、これらの挙動を駆動する内部メカニズムは依然として不透明である。既存のメカニスティック・インタープリタビリティ（機械論的解釈可能性）のアプローチは、推論に対して以下の2つの主要な限界に直面している：

1. 共起関係への依存： 多くの手法は、特定のテキストパターン（例：「待って」などのトークン）や短期的効果（例：算術演算）との相関のみに基づいてコンポーネント（ニューロンやSAE特徴量など）を特定する。これらのアプローチは、複雑で多段階の推論挙動を因果的に駆動するコンポーネントを分離することには失敗している。
2. 逐次的な影響の軽視： 現在の手法は、多くの場合、即時的かつ短期的な効果のみを捉えており、最終的な結果（例：回答の正解性）が中間決定のシーケンスに依存するような、多段階の推論タスクに必要な累積的かつ長期的な影響を無視している。さらに、制御ベクトルを導出するために人間がアノテーションした対照ペアに依存する手法は、意図したターゲット状態が真に引き出されていることを保証するのが困難である。

したがって、アウトカム指向（最終的な推論結果への貢献に基づいてコンポーネントを評価する）であり、かつ逐次的影響を考慮（推論プロセス全体における内部コンポーネントの累積的な影響を考慮する）したフレームワークが切実に求められている。

手法：統合ポリシーグラジエント (Integrated Policy Gradient: IPG)

著者らは、長期的なアウトカムに基づいて内部コンポーネントに挙動を帰属させることで、LLMの推論を解釈および制御するために設計された、新しい学習不要のフレームワークである統合ポリシーグラジエント (IPG) を提案する。

コア原理

• アウトカム指向： 内部コンポーネントの貢献度は、最終的な推論結果（例：回答の正解性や推論チェーンの長さ）に対して評価される。
• 逐次的影響を考慮： この手法は、孤立したステップではなく、全生成軌跡にわたるコンポーネントの累積的な影響を捉える。

技術的フレームワーク

1. 表現空間ポリシーグラジエント：
   標準的なポリシーグラディエント手法がモデルパラメータ（$\theta$）を更新するのに対し、IPGは表現空間へとその視点を拡張する。IPGは、アウトカムに基づく信号（例：検証器から派生した報酬）を、推論軌跡を通じて内部コンポーネント（$h_t$、例えばニューロンやSparse Autoencoder (SAE) 特徴量）へと逆方向に伝播させる。
   表現空間におけるポリシーグラジエントは以下のように定式化される：
   $$\frac{\partial J}{\partial h} = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=1}^T \frac{\partial}{\partial h_t} \log \pi_\theta(a_t | s_t; h_t) \cdot A_\pi(s_t, a_t) \right]$$
   ここで、$J$ は信号関数（例：二値の正解性）、$\tau$ は推論軌跡、$A_\pi$ はトークン選択の長期的利益を推定するアドバンテージ関数である。

2. パス積分帰属 (Integrated Policy Gradient)：
   生の勾配に伴うノイズと局所的な感度に対処するため、IPGはパス積分アプローチを採用する。これは、ベースライン（例：ゼロベクトル）からターゲットとなる活性化値への線形補間経路に沿って、ポリシーグラジエントを累積するものである。これにより、帰属スコアが、ポリシー出力の全変化にわたるコンポーネントのグローバルな貢献を考慮することを保証し、「完全性（completeness）」の特性を満たす。
   $$\text{IPG}(i; x) = \sum_{t=1}^T (h_{t,i} - h'_{t,i}) \int_0^1 \mathbb{E}_{\tau \sim \pi_\theta} \left[ \frac{\partial}{\partial h_{t,i}} \log \pi_\theta(a_t | s_t; h'_{t} + \alpha(h_t - h'_t)) \cdot A_\pi(s_t, a_t) \right] d\alpha$$

3. コンポーネントの選択と介入：

   • 特定： IPGはデータセット全体で帰属スコアを計算し、それらを集計することで、最も影響力のある上位$p$個のコンポーネント（ニューロンまたはSAE特徴量）を特定する。
   • 制御： 特定された後、これらのコンポーネントに対して乗法的スケーリング係数 $\gamma$ を適用することで、因果的な介入を行う。$\gamma > 1$ に設定するとコンポーネントの影響を強化し、$0 \le \gamma < 1$（または $\gamma=0$）に設定すると抑制する。これにより、推論能力（正確性）や推論強度（連鎖の長さ）といった推論挙動の変調が可能になる。

主な貢献

1. 新しいフレームワーク (IPG)： テキストパターンの相関や短期的効果の分析という限界を克服し、長期的なアウトカムに基づいて推論挙動を直接内部コンポーネントに帰属させる、学習不要のフレームワークを導入した。
2. 解釈と制御のための統一されたツール： このアプローチは、個々のニューロンやSAE特徴量を含む様々な内部ユニットに拡張可能であり、異なるモデルアーキテクチャやスケールにわたって推論挙動を解釈・変調するための統一されたメカニズムを提供する。
3. 実証された有効性と転移性： 実証的な証拠により、IPGがベースライン（例：Control Vectors, Reasoning Neurons）と比較して優れた制御性能を達成することが示された。極めて重要な点として、特定されたコンポーネントは強い転移性を備えており、プロンプトによって誘発されたモデルで発見された推論メカニズムを、同一モデルファミリー内の学習によって誘発された（蒸留された）バリアントに直接適用できる。

実験結果

著者らは、代表的なオープンソースLLM（Qwen2.5-Math-1.5B-Instruct および Llama3.1-8B-Instruct）を用い、多様な推論ベンチマーク（GSM8K, MATH-500, AIME-2024, GPQA-Diamond）においてIPGを評価した。

• 精密な局在化と制御： IPGは、推論能力の強化および抑制の両方において、ベースラインよりも一貫して高い精度を達成した。例えば、Qwen2.5-Math-1.5Bにおいて、IPGはGSM8Kの精度を84.61%（元の82.41%から向上）に向上させ、抑制下では10.69%まで低下させた。これは精密な制御が可能であることを示している。
• きめ細かな変調： この手法は、異なる推論挙動を正常に変調することに成功した。介入によって、一方を犠牲にすることなく、推論能力（正確性）を強化したり、推論強度（シーケンス長）を調整したりすることができた。これは、特定の挙動が個別のコンポーネントにエンコードされていることを示唆している。
• データセットおよびモデル間の転移： あるデータセット（例：GSM8K）で特定されたコンポーネントは、他のデータセット（例：MATH-500）へ効果的に転移した。また、プロンプトによって誘発されたモデルで特定されたメカニズムは、推論蒸留モデル（例：DeepSeek-R1-Distilled-Qwen-1.5B）の性能を向上させた。
• 粒度の高い洞察： 分析により、SAE由来の特徴量はニューロンよりも分離が進んでおり、しばしば単一の推論側面（例：計算 vs 分解）に影響を与える一方で、ニューロンは頻繁に多義性（polysemanticity）を示し、複数の次元に同時に影響を与えることが明らかになった。

意義と主張

本論文は、IPEが表面的な相関から因果的かつアウトカム駆動型の帰属へと焦点を移すことで、LLMの推論におけるメカニスティック・インタープリタビリティに重要な進歩をもたらすと主張している。内部コンポーネントの逐次的影響を効果的に捉えることで、IPGは再学習を必要とせずに推論挙動の信頼性の高いステアリング（操舵）を可能にする。

著者らは、IPGを複雑な推論の「ブラックボックス」を解明するための堅牢なソリューションとして位置づけており、以下の道筋を提示している：

• 推論システムの真実性、安全性、および制御可能性の向上。
• モデルバリアント間で転移可能な推論メカニズムを活用することによる、冗長な解釈可能性への取り組みの削減。
• 特定の推論側面（例：問題の理解、タスクの分解）に関する微細な洞察を提供し、LLMがどのように推論をエンコードし実行するかについての理論的理解を深める。

結論として、現在の手法は多段階の推論の複雑さに苦慮しているが、IPGはポリシーグラジエントとパス積分を統合することで、これら重要なモデル挙動を解釈および制御するための、原理的かつ効果的なアプローチを提供している。