技術要約: Pass@K ポリシー最適化 (PKPO)

問題提起

大規模言語モデル（LLM）に適用される強化学習（RL）アルゴリズムは、単一の問題に対して複数の解の試行（$n > 1$）をサンプリングすることが多いが、個々のサンプルのパフォーマンス（pass@1）に基づいて最適化を行う。このアプローチは、集合全体の集合的な有用性よりも、単一の高報酬サンプルの有用性を優先するため、サンプリング能力を実質的に使いこなせていない。その結果、標準的な方策勾配法は、特に単一の正解サンプルの確率が低い困難なタスクにおいて、解空間を十分に探索できない可能性がある。推論時探索手法は存在するものの、推論とモデル更新の間の結合関係は、正解性は検証可能であるものの解が事前に未知であるタスクにおいて、単純に期待単一サンプル報酬を最適化することは、最適ではない可能性を示唆している。

手法

著者らは、Pass-at-k ポリシー最適化 (PKPO) を提案する。これは、ポリシーから独立して抽出された $k$ 個のサンプルの最大報酬の期待値である pass@k 指標（またはその連続的な類似体である maxg@k）を直接最適化するフレームワークである。

コアとなる理論的貢献

主な貢献は、pass@k 目標およびその勾配に対する、新規で低分散かつ不偏な推定器の導出である。これらの推定器は、$n$ 個のサンプルからなるバッチ内のすべての $\binom{n}{k}$ 個のサブセットに対して単純な推定器を平均化することによって機能する。

1. バイナリ報酬 (Section 2):

   • バイナリの正誤ラベル（$f \in \{0, 1\}$）に対して、著者らは、正解数 $c$ を用いた pass@k の不偏推定器 $\rho(n, c, k)$ を導出する。
   • 不偏推定器は $n \ge k$ の場合にのみ存在することを証明している。
   • 対応する不偏勾配推定器も導出されており、正誤および成功するサブセットに含まれる組合せ論的な確率に基づいてサンプルに重みを割り当てる。

2. 連続報酬 (Section 3):

   • 本フレームワークは、maxg@k 目標 $E[\max_{i=1}^k g(x_i)]$ を通じて、連続報酬 $g(x)$ へと一般化される。
   • 報酬がソートされている（$g_1 \le g_2 \le \dots \le g_n$）と仮定することで、著者らは二項係数を用いてソートされた報酬に重みを付ける不偏推定器 $\rho^{(g)}(n, c, k)$ を導出する。
   • サンプルのための重み $s_i$ は、報酬の線形結合として計算され、効率的に $O(k + n \log n)$ 時間で構築される不偏勾配推定器が構成される。

3. 分散減少 (Section 4):

   • バイアスを導入することなく勾配推定の分散を低減するために、著者らは Leave-One-Out (LOO) ベースラインを提案する。
   • 現在のサンプルを除いた $k-1$ 個のサブセットの maxg@k から、$k$ 個のサブセットの maxg@k を差し引く特定のベースライン $s^{(loo-1)}_i$ を導入する。これは、$n=k$ のケースに限定されていた先行研究（具体的には [TZSM25]）を一般化したものである。
   • この定式化により、$n > k$ の場合（LLMの学習において大きなバッチがサンプリングされる一般的なシナリオ）でも、堅牢な最適化が可能になる。

実装

本手法は、任意の方策勾配アルゴリズム（例：PPO）における標準的な報酬変換の「ドロップイン・リプレイスメント（そのまま置き換え可能なもの）」として設計されている。$n$ 個のサンプルとその生の報酬を含むバッチが与えられると、提供されたコード内の sloo_minus_one 関数を用いて報酬が変換され、その後、方策更新ステップに投入される。

主な結果

著者らは、GEMMA2 (2B および 9B パラメータ) と LLAMA3.1 (8B パラメータ) を用いたトイ実験および実世界のベンチマークを通じて PKPO を検証している。

1. トイ実験:

   • 1次元ガウス方策の実験において、PKPO はより高い $k$ に対して方策の平均をより「リスク許容的」な方向へシフトさせることに成功し、これによりモデルは確率が低いが潜在的な報酬が高い領域を探索できるようになる。
   • 分散分析により、提案された sloo_minus_one 推定器が、ナイーブな分割やベースラインのない推定器と比較して、勾配の分散を大幅に低減することが確認された。

2. 実世界ベンチマーク (MATH, Coding, ARC-AGI-1):

   • 選択的最適化: 最適化ターゲット $k_{opt}$ を評価指標 $k_{eval}$ と等しく設定すると、最良のパフォーマンスが得られる。例えば、pass@1 を最適化する場合と比較して、pass@8 を最適化すると pass@8 の評価スコアが大幅に向上する。
     精緻な探索が必要な難しい問題において、従来の pass@1 最適化は停滞し、累積的な解決率の向上に失敗する。対照的に、$k_{opt} > 1$ を用いた PKPO は学習を「アンブロック（解禁）」し、すべての $k$ においてより高い解決率と pass@k スコアを達成する。
   • アニーリング戦略: 訓練中に $k$ をアニーリングする（探索を促すために高い $k_{opt}$ から開始し、その後 $k_{opt}=1$ に減少させる）ことで、pass@k の大幅な向上とともに強力な pass@1 パフォーマンスを達成できることを示し、pass@k の向上によって pass@1 が低下するという典型的なトレードオフを回避できることを実証している。
   • 比較: PKPO は、$n=k$ の特殊なケース（[TZSM25] のように）やエントロピー正則化ベースラインよりも優れた性能を示し、特に探索が極めて重要な困難なタスクにおいて顕著である。

意義と主張

論文は、PKPO が、単一のサンプルまたは平均ではなく、一連のサンプルの集合的な有用性を最適化するための、堅牢で効率的かつ理論的に裏付けられた手法を提供すると主張している。

• 探索のアンブロック: 著者らは、個々のサンプルの有用性よりも、集合的な有用性（$k$ 個のサンプルのうち少なくとも 1 つが成功する確率）を優先することで、探索を改善すると主張している。これは、解空間において「信号」（正しい解）が疎である困難な問題を解く上で極めて重要である。
• 柔軟性: $k=n$（バッチ全体のベストを最適化する）または $k=1$ に制限されていた従来の手法とは異なり、PKPO は任意の $k \le n$ に対して堅牢な最適化を可能にする。
• 実用的な影響: 本手法は、実践者がパラメータ $k_{opt}$ を通じて探索と利用のトレードオフを調整できるようにし、標準的な pass@1 最適化では到達不可能な、より複雑な推論やコーディングタスクをモデルが解決できるようにする。

本研究は、既存の順序統計量方策勾配法の一般化および分散減少として位置づけられており、深い探索を必要とするタスクに対する LLM の RL トレーニングをスケールアップするための実用的な経路を提供している。