技術要約: rePIRL: LLMの推論における逆強化学習を用いたPRMの学習

問題提起

大規模言語モデル（LLM）の推論において、プロセス報酬モデル（PRM）は、トークンレベルの中間報酬を提供するために利用され、学習効率の向上、分散の低減、および報酬ハッキングの防止を目的としている。しかし、効果的なPRMを学習するための既存の手法には、以下のような重大な限界がある：

• オフライン手法: 教師あり学習（Uesato et al., 2022; Lightman et al., 2023）やDPO型の研究は、エキスパートの軌跡からのトークンレベルの注釈や選好ラベルを必要とするが、これらを取得するコストは高い。モンテカルロ木探索（MCTS）に基づく手法（Wang et al., 2024）は、エキスパート・ポリシーを実行して軌跡を生成する必要があり、エキスパート・ポリシーが利用できない場合には計算コストが膨大になり、実行不可能である。
• オンライン手法: PRIME（Cui et al., 2025a）のようなアプローチは、報酬関数と学習対象のポリシーモデルに関する強力な仮定に依存している。内的報酬を用いる手法（Zhao et al., 2025b）は、モデルのエントロピーや確信度を報酬として使用するが、過度なモード探索行動による「エントロピー崩壊」に苦しみ、学習の後半段階で性能が低下する。

核心となる課題は、最小限の仮定（具体的には、トークンレベルの報酬注釈、選好ラベル、エキスパート・ポリシーへのアクセス、または報酬関数の形式に関する強力な仮定を必要としないこと）の下で、効果的なPRMを学習することである。

手法: rePIRL フレームワーク

著者らは、エキスパートの軌跡のみからPRMを学習する、逆強化学習（IRL）に着想を得たフレームワークである rePIRL を提案する。この手法は、マルチステップのLLM推論をトークンレベルのマルコフ決定過程（MDP）としてモデル化し、ポリシーとPRMを交互に更新するデュアル学習プロセスを採用している。

1. 理論的基礎

• 目的関数の構築: 古典的なIRLフレームワークに従い、著者らは、アクションがエキスパート・ポリシーからサンプリングされたかどうかを示す隠れた変数 $o_t \in \{0, 1\}$ を定義する。彼らは、アクションがエキスパートによってサンプリングされた確率を、エネルギーベースの関数を用いてモデル化する：$p(o_t=1|s_t, a_t) \propto \exp(r_\phi(s_t, a_t))$。
• 尤度最大化: 報酬関数 $r_\phi$ は、エキスパートの軌跡の尤度を最大化することによって学習される。これは、分割関数 $z(\phi)$ を含む目的関数へと導かれるが、広大な状態空間とアクション空間を持つLLMにとって、この関数は計算不可能（intractable）である。
• サンプリングによる計算可能性の確保: 分割関数の計算不可能性を解決するために、著者らは勾配トリックと**重要度サンプリング（importance sampling）**を適用する。未知である最適なポリシーからサンプリングする代わりに、現在のポリシー $\pi_\theta$ から軌跡 $\bar{\tau}$ をサンプリングする。勾配は、重要度重み $w = \exp(r_\phi(\bar{\tau})) / \pi_\theta(\bar{\tau})$ を用いて近似される。
• ポリシー最適化: モデルの崩壊（決定論的ポリシーへの陥落）を防ぐため、ポリシーは最大エントロピーRL（Haarnoja et al., 2018）の目的関数を用いて更新される。これは、報酬の最大化とエントロピー正則化のバランスを取るものである。これは、RLOO、PPO、またはGRPOといったアルゴリズムを用いて解かれる。

2. 学習アルゴリズム

学習ループは、PRMの更新とポリシーの更新を交互に行う：

1. ロールアウト: ポリシーによるロールアウトを生成し、エキスパートのデモンストレーションを取得する。
2. 報酬計算: 結果報酬（outcome rewards）と平均プロセス報酬を計算する。
3. PRMの更新: エキスパートの軌跡とポリシーのロールアウトを、重要度サンプリングのスコアによって重み付けして対比させる損失関数を用いて、報酬モデルを更新する。極めて重要な点として、報酬ハッキングを軽減するために、著者らは累積報酬ではなく平均プロセス報酬を使用している。
4. ポリシーの更新: 学習されたプロセス報酬（および利用可能な場合は結果報酬）を用い、最大エントロピー目的関数を用いてポリシーモデルを更新する。

3. 既存手法の統一

主要な理論的貢献は、特定の追加的な仮定を導入することで、rePIRLが様々な既存のPRM学習手法（DPO, DQO, Math-Shepherd, PRIME）を単一のフレームワークの下に統一できることを示した点である。これにより、rePIRLは最も汎用的な手法として位置付けられる。

主な貢献

1. 新規フレームワーク: トークンレベルの注釈、選好ラベル、またはエキスパート・ポリシーへのアクセスを必要とせず、エキスパートの軌跡からPRMを学習する、IRLに着想を得たフレーム沢の導入。
2. 統一理論: rePIRLが、既存の最先端のオフラインおよびオンラインPRM手法（DPO, DQO, PRIME, およびMCTSベースのアプローチを含む）を単一の学習フレームワークの下に統一できることを示す理論的分析を行い、それらの基礎となる仮定を明確化した。
3. スケーラブルなアルゴリズム: 分割関数の計算不可能性を扱うための重要度サンプリングと、学習の安定性を確保するための最大エントロピーRLを活用した、LLMに特化したスケーラブルな学習アルゴリズムの開発。
4. 実証的検証: 数学およびコーディングのベンチマークにおける包括的な評価により、ベースラインに対する優れた性能を実証。

実験結果

著者らは、7つの標準的なベンチマーク（AIME-2024, AMC, Math-500, MinervaMath, OlympiadBench, Leetcode, LiveCodeBench）において、Qwen2.5-3B-Instruct および Qwen3-4B-Base を用いてrePIRLを評価した。

• 性能: rePIRLは、3Bおよび4Bスケールのモデルの中で最先端の結果を達成した。例えば、Qwen3-4B-Baseにおいて、数学ベンチマークで41.6%、コーディングベンチマークで**39.8%**の平均精度に達し、PRIME、MCTS、およびバニラなRLOOといった強力なベースラインを上回った。
• 比較: Qwen3-4B-Baseを用いたバニラなRLOOと比較して、絶対的な改善幅は**3.2%（数学）および4.3%（コーディング）**であった。
• 応用:
  • テスト時トレーニング（Test-time Training）: 結果報酬なしでPRM信号のみで学習した場合でも、rePIRLは検証可能な結果報酬を用いて学習されたモデルの性能に匹敵した。
  • テスト時スケーリング（Test-time Scaling）: より良いロールアウトを選択するための検証器としてPRMを使用することで、多数決（majority voting）を上回る単調な性能向上を実現した。
  • 困難な問題: 結果報酬が疎な困難な問題において、PRMは有益な初期信号を提供し、結果報酬のみの学習と比較して高速な収束をもたらした。
• 効率性: rePIRLは、結果ベースのRL（例：RLOO）と比較して、学習時間のオーバーヘッドは約1.4倍に抑えられており、かつ別途の価値ネットワーク（value networks）を必要としない。

意義と主張

本論文は、LLMの推論における逆強化学習を一般化した最初の研究であり、既存のPRM学習手法を単一のフレームワークに統合した最初の研究であると主張している。

• 最小限の仮定: 著者らは、rePIRLが最小限の必要な仮定の下で動作することを強調しており、トークンレベルの注釈や特定のエキスパート・ポリシーへのアクセスを必要とする手法よりも汎用性が高い。
• 安定性: フレームワークは、最大エントロピーRLを活用することで、内的報酬手法に共通するエントロピー崩壊の問題に対処し、安定した学習ダイナミクスを確保している。
• 有用性: 学習されたPRMは、再利用可能なスタンドアロンのモジュールとして提示されており、結果報酬が利用できない、あるいは疎であるシナリオにおいて、テスト時トレーニングやスケーリングなどのダウンストリームアプリケーションに利益をもたらす。

著者らは、rePIRLが統一的な理論的視点と実用的な有用性を提供する一方で、今後の課題として、内的報酬手法との結合、オンラインRLヒューリスティクスの体系的な比較、およびエージェンティックなタスクやより大規模なモデルへのスケーリングが挙げられると結論付けている。