技術要約：Max@K ポリシーグラディエントにおけるアドバンテージ推定について

問題提起

検証可能な報酬（RLVR）を用いた強化学習は、推論タスクにおける大規模言語モデル（LLM）のポストトレーニングにおいて支配的なパラダイムとなっています。しかし、これらの設定は通常、疎な結果ベースの報酬（例：正解か不正解かのバイナリ信号）に依存しており、これが大規模な組合せ空間における探索を困難にしています。エントロピー正則化やカウントベースのボーナスといった手法が存在しますが、これらは注意深いチューニングを必要とする補助的な目的関数を導入することになります。

補完的なアプローチは、推論時の目的関数、すなわち pass@K（$K$ 個のサンプリングされた補完のうち、少なくとも 1 つが正解である確率）およびその連続的な一般化である max@K（$K$ 個のサンプルにおける期待最大報酬）を直接最適化することです。いくつかの最近の研究では、これらの目的関数のための不偏ポリシーグラディエント（PG）推定器が提案されています。しかし、これらの推定器間の関係は依然として不明確です。具体的には、既存の推定器は異なる信号、ベースライン、および正規化を利用しています。本研究で特定された重要な問題は、主要な推定器（例：Expected Improvement に基づくもの）は PG 推定器としては不偏ですが、**非中心的なアドバンテージ（non-centered advantages）**を生み出すという点です。非中心的なアドバンテージ（平均がゼロではないもの）は、ベースラインの過小評価に対応し、最適化を劣最適なアクションへと導き、勾配の分散を増大させる可能性があります。

手法

本論文は、max@k 目的関数のアドバンテージ推定を修正するために、Leave-Two-Out (L2O) ベースラインを導入する手法である MaxPO (Max@K Policy Optimization) を提案しています。

1. 既存の推定器の分析

著者らは、先行研究（例：Walder and Karkhanis [54]）で使用されている Expected Improvement (EI) 推定器を再検討しています。彼らは、EI ベースの推定器が不偏の PG 推定器である一方で、得られるアドバンテージ信号 $s(a)$ が（EI の定義における ReLU 操作により）厳密に非負であることを示しています。その結果、$E[s(a)] > 0$ となり、アドバンテージが中心化されていないことを意味します。この中心化の欠如は、分散減少の余地を残しています。

2. Leave-Two-Out (L2O) ベースライン

正確に中心化されたアドバンテージを実現しつつ、PG の不偏性を維持するために、著者らは L2O ベースラインを導入しています。

• 課題: 標準的な Leave-One-Out (L1O) ベースラインは、max@K 設定においては失敗します。なぜなら、サンプル $i$ の報酬が、他のサンプル $j$ のアドバンテージ計算に暗黙的に影響を与える可能性があるためです（$i$ が max 操作における潜在的な比較対象として機能するため）。この依存性は、不偏 PG 推定に求められる独立条件に違反します。
• 解決策: L2O ベースライン $b_{-i}^{L2O}$ は、ターゲットとなるサンプル $i$ とサンプル $j$ 自体の両方を除外した比較集合上で計算された、他のすべてのサンプルの Expected Improvement スコアの平均として計算されます。
• 特性:
  • 独立性: すべての $j$ について比較集合から $i$ を除外することで、ベースライン $b_{-i}^{L2O}$ は報酬 $r_i$ と独立し、PG 推定器の不偏性を保証します。
  • 中心化: 著者らは、実現されたバッチ平均アドバンテージが、サンプリングされたすべてのバッチに対して厳密にゼロになることを証明しています：$\frac{1}{B}\sum_{i=1}^B (s_i - b_{-i}^{L2O}) = 0$。
  • 効率性: 素朴な計算は $O(B^3)$ となりますが、著者らは ReLU ベースの定式化を用いたベクトル化された実装を導出し、L2O ベースラインを $O(B^2)$ の時間で計算できることを示しました。これにより、GPU 加速された LLM トレーニングへの適用が可能になります。

3. 標準的な有限バッチ・アドバンテージ

本論文は、max@K のための標準的な有限バッチ・アドバンテージを $u_i - v_i$ として導出しています：

• $u_i$: レスポンス $i$ を含むグループの条件付き期待 max@K 値。
• $v_i$: $i$ を除外したグループから計算された leave-one-out 期待 max@K 値。
  著者らは、MaxPO 推定器（EI + L2O）がこの標準的な信号（$s_i - b_{-i}^{L2O} = u_i - v_i$）を正確に復元することを示しています。これは、既存の推定器を「非中心的」（例：EI のみ/PKPO）と「（正規化された）標準的」（例：MaxPO、および特定の解析的な pass@K 変種）に分類する統一的な視点を提供します。

4. LLM トレーニングへの統合

MaxPO は、グループベースの強化学習アルゴリズム（GRPO など）に統合されます。各プロンプトに対して、$G$ 個の出力からなるグループがサンプリングされます。L2O ベースラインは各グループ内で計算され、シーケンスレベルのアドバンテージを生成し、それがトークンレベルの損失へとブロードキャストされます。

主な貢献

1. 非中心的なアドバンテージの特定: 本論文は、主要な max@K PG 推定器（特に EI ベースのもの）が、勾配推定器としては不偏であるものの、中心化されたアドバンテージを生成できず、結果として劣最適な分散特性をもたらすことを特定しました。
2. Leave-Two-Out (L2O) ベースライン: 非線形な max@K 設定において、PG の不偏性を維持しながら、実現されたバッチアドバンテージが平均ゼロになるように設計された新しいベースライン設計。
3. 標準的な統一: 標準的な有限バッチ・アドバンテージ（$u_i - v_i$）の導出と、MaxPO がこの信号を復元するという証明により、既存の pass@K/max@K 推定器を比較・分類するための統一的な枠組みを提供しました。
4. 効率的な実装: L2O ベースラインを計算するための $O(B^2)$ ベクトル化アルゴリズムにより、LLM ポストトレーニングへの実用的な適用を可能にしました。

実験結果

著者らは、制御されたバンディット/迷路実験および LLM 推論タスクを通じて手法を検証しています。

• 分散減少: バンディット環境において、L2O ベースラインは生の EI 推定器と比較して勾配の分散を大幅に減少させ、その効果はアクション空間のサイズが増加するにつれて大きくなります。LLM トレーニング（Llama-3.2-3B-Instruct）において、本手法は PKPO（EI のみのベースライン）と比較して、勾配分散のプロキシを 77.4% 低減しました。
• 性能向上: 5 つの数学推論ベンチマーク（AIME24, AIME25, AMC23, MATH500, Minerva）において、MaxPO は PKPO、GRPO、Entropy-Adv を含む強力なベースラインを上回りました。
  • Qwen2.5-Math-7B において、MaxPO は PKPO に対して pass@256 を 5.2% 向上させました。
  • Llama-3.2-3B-Instruct において、pass@256 を PKPO より 2.4% 向上させました。
• 推論計算のスケーリング: MaxPO の性能向上は、より高い推論計算レベル（pass@k における大きな $k$）において最も顕著であり、この手法がモデルが追加のテスト時サンプルをより良く活用するのを助けることを示唆しています。

重要性と主張

本論文は、MaxPO が max@K のような推論時目的関数のためのポリシーグラディエント推定器を設計・比較するための原理的な基礎を提供すると主張しています。既存の推定器におけるバイアスを特定し修正することで、本研究は、推論能力を持つモデルの効率性と安定性を向上させる手法を提示しています。

著者らは、分析が i.i.d. サンプリング下での有限バッチ、オンポリシー推定に焦点を当てていることを控えめに述べています。オフポリシー学習、適応的サンプリング、および相関のある生成が行われる設定に関する制限を、今後の課題として認めています。さらに、手法の改善は最適化に関するものですが、繰り返しのサンプリングを行う推論目的関数のより良い最適化は、より広範な影響を及ぼし、欺瞞的な支援や自動化された問題解決に悪用される可能性のあるモデルを強化する可能性があることも強調しています。