技術概要：PAIR（多ターンエージェント最適化のためのプレフィックス認識型内部報酬モデル）

1. 問題定義

現在の大規模言語モデル（LLM）は、一連の動作（例：API 呼び出し、ツール呼び出し）と環境からのフィードバックを必要とする複雑な多段階エージェントタスクに苦慮している。グループ相対方策最適化（GRPO）は、その単純性と効率性からこれらのタスクにおける強化学習（RL）アルゴリズムの最先端を走っているが、希薄な結果報酬に悩まされている。標準的な GRPO では、単一のスカラー報酬が軌道の完了時のみ割り当てられるため、中間ステップ間でのクレジット割当が困難である。この希薄さは、収束の遅延とサンプル効率の低下をもたらす。

報酬を高密度化しようとする既存の試みは、重大な限界に直面している：

• 完全ロールアウト： 結果報酬を後方へ伝播させるには完全な軌道の完了が必要であり、すでに失敗している経路からの無駄なロールアウトを生成する。
• 外部ジャッジ： 各ステップを評価するために外部 LLM（例：GPT-4）を使用すると、高い API コストと遅延が発生し、大規模なオンライン RL には実用的ではない。
• 内在的報酬： 正解（教師強制ログ確率など）に依存する手法は、実行時にゴールドラベルへのアクセスを必要とするが、これはしばしば利用不可能か、運用コストが高い。
• 外部報酬モデル： 個別に訓練されたプロセス報酬モデル（PRM）は、訓練と維持のオーバーヘッドをもたらす。

本研究が取り組む核心的な課題は以下の通りである：外部モデル、正解への依存、または完全軌道ロールアウトなしに、エージェント自身の内部状態のみを使用して、高密度なステップレベルの報酬信号を生成することは可能か？

2. 手法

2.1 実証的動機：単純なプローブの失敗

著者はまず、隠れ状態からステップの正しさを予測するように訓練された内部正しさプローブが報酬として機能しうるか調査する。彼らは、エージェントが誤った中間ステップを生成し、後続のステップのためのプレフィックスを汚染するプレフィックス汚染の概念を導入する。

Qwen-2.5-7B-Instruct をバックボーンとした体系的な分析を通じて、彼らは 3 つの重要な発見を明らかにした：

1. 隠れ状態の劣化： 隠れ状態プローブ（例：ラストトークン、平均プーリング）は、プレフィックス汚染下で著しく劣化する。これらは正解に基づく正しさ（実際のタスク有用性）ではなく、信念一貫性（汚染されたプレフィックスとの整合性）を追跡する。その結果、誤りを永続させるターンに報酬を与え、修復ターンを罰することが多い。
2. アテンションの頑健性： アテンションに基づく特徴は、汚染に対して比較的頑健である。これらはプレフィックスとの意味的整合性ではなく、構造的な動作パターン（例：モデルがタスク関連トークンに注意を向けるかどうか）を捉える。
3. 相補性： 隠れ状態はクリーンなプレフィックスでは優れているが、汚染されたものでは失敗する。一方、アテンション特徴は頑健だが、クリーンなプレフィックスでは性能が劣る。両方の信号を単純に連結しても失敗する。なぜなら、隠れ状態からの逆転した信号が、アテンションからの正しい信号を相殺してしまうからである。

2.2 提案手法：PAIR

これらの信号の相補性を活用するために、著者は**PAIR（プレフィックス認識型内部報酬モデル）**を提案する。これは 2 段階のアーキテクチャである：

• ステージ 1：信念一貫性推定器（凍結）： 事前訓練されたロジスティック回帰プローブが、隠れ状態特徴（特にラストトークンの隠れ状態）を信念一貫性スコア（$s_{bc}$）にマッピングする。このプローブはクリーンなデータで訓練され、RL 訓練中は凍結される。これは、プレフィックスがクリーンな場合の標準的なプローブの正しさ検出能力を捉える。
• ステージ 2：アテンションベースの修正： 軽量なロジスティック回帰ヘッドが、$s_{bc}$ と連結されたアテンション特徴（ピーク強度、広がり、プレフィックス比率、自己比率などの統計量）を入力として受け取る。これは最終スコア（$s_{final}$）を予測する。
  • メカニズム： クリーンなプレフィックスでは、修正ヘッドは $s_{final} \approx s_{bc}$ となるように出力を学習する。汚染されたプレフィックスでは、アテンション特徴が構造的な異常を検知し、ヘッドはバイアスのかかった $s_{bc}$ から正解に基づく正しさへと $s_{final}$ を調整するように学習する。
• 報酬統合： 最終スコアは、温度クリップ付きシグモイド関数を通じて変換され、モーメントベースの報酬メカニズムと組み合わせられる。これは、軌道の走行平均に対する現在のステップのスコアを対比させ、GRPO 優位性推定に不可欠なグループ相対分散を回復させる。

3. 主要な貢献

1. 実証的発見： 本論文は、多ステップ設定において内部正しさプローブがプレフィックス汚染下で失敗し、正解に基づく正しさではなく信念一貫性を追跡することを示す、最初の体系的な研究を提供する。
2. メカニズム的洞察： 隠れ状態が信念一貫性をエンコードし、アテンションパターンが構造的なルーティング動作をエンコードすることを確立する。これらの信号は相補的であり、単純な連結が失敗する理由を説明し、適応型アーキテクチャの設計を導く。
3. 方法論的革新： PAIR は、隠れ状態とアテンション信号を適応的に組み合わせる 2 段階モデルとして提案される。これは、汚染下でのバイアスを修正しつつ、クリーンな軌道において高い性能を維持する。
4. 実用的効率性： PAIR は、外部 LLM 呼び出し、実行時の正解アクセス、または個別の報酬モデル訓練なしに動作する。方策のフォワードパス中に既に計算された活性化を再利用することで、プローブレベルの無視できる推論コスト（~0.2 ms/ステップ）を実現する。

4. 実験結果

著者は、Qwen-2.5-7B-Instruct を方策モデルとして使用し、2 つの多ステップエージェントベンチマーク（GTA（General Tool Agents）とToolBench）で PAIR を評価した。

• ステップレベル予測： PAIR は、隠れ状態プローブ、アテンションプローブ、教師なし手法（例：Lookback Lens、Head Entropy）を含むすべての内部プローブベースラインの中で、最高の AUROC（GTA で 0.9217、ToolBench で 0.8129）を達成した。
• RL 性能： GRPO 訓練において、PAIR は両方のデータセットで最高成功率を達成した（GTA で 0.2489、ToolBench で 0.4498）。ToolBench では 2 番目に優れたベースライン（LLM-as-a-Judge）を上回り、GTA ではそれと同等の性能を示した。
• クロスドメイン転移： GTA で訓練し、微調整なしで HotpotQA で評価した際、PAIR は 16 の手法の中で 1 位にランクインし、ロールアウトベースまたは正解依存の手法と比較して優れた転移性を示した。
• 効率性： PAIR はプローブレベルの遅延（~0.2 ms/ステップ）で動作し、LLM-as-a-Judge API 呼び出しよりも約 2,500 倍高速であり、木探索手法よりも桁違いに高速である。

5. 意義と主張

本論文は、PAIR が既存の高密度報酬メカニズムの 4 つの主要な限界を同時に解決すると主張している：

1. 完全ロールアウト不要： 報酬は軌道完了を待たずにステップごとに計算される。
2. 外部呼び出し不要： エージェントの内部表現のみに依存し、API コストと遅延を排除する。
3. 正解依存なし： プローブはオフラインで訓練され、各ステップにゴールド回答を必要とせずに実行時に適用される。
4. 個別報酬モデル不要： 大規模な報酬ネットワークの個別訓練を必要としない。

著者は、適応型かつアテンション修正アーキテクチャを通じてプレフィックス汚染下での隠れ状態プローブの失敗モードを解決することで、PAIR を多ターンエージェント RL 訓練における高密度なステップレベル報酬信号を可能にする実用的な解決策として位置づけている。PAIR は、複雑なエージェント動作の最適化において、外部報酬源に対する代替手段として、内部状態分析を実用的かつ効率的なものにする。