技術要約: SARA (Sequential Adaptive Rollout Allocation)

問題提起

検証可能な報酬（Verifiable Rewards）を伴う強化学習（RLVR）は、現在、ロールアウト生成のコストによってボトルネックに直面している。Group Relative Policy Optimization (GRPO) のようなグループベースの推定器では、あるプロンプトがポリシー勾配に与える寄与度は、サンプリングされたグループ内の報酬の分散に依存する。もしグループが「飽和（saturated）」状態（すべての回答が正解、またはすべて不正解）である場合、報酬の分散はゼロとなり、正規化されたアドバンテージが消失して学習信号が得られなくなる。

これらの無駄を軽減するための既存手法は、以下のトレードオフに直面している：

1. 評価・フィルタリング（例：Dynamic Sampling/DS）: 大規模な候補プールをオーバーサンプリングし、完全なグループを生成してから、飽和したグループを破棄する。これらは効果的なグループのクリーンなバッチを保証するが、最終的に破棄されるプロンプトの完全な生成に対してもコストを支払うため、一様サンプリングよりも大幅に多い（しばしば4倍以上の）ロールアウトコストを発生させる。
2. 予測・選択（Predict-then-select）: プロンプトをサンプリングする前にプロンプトの難易度を推定し、有望なプロンプトを優先する。これらは追加のロールアウトを回避できるが、予測が急速に変化するポリシーに対して脆弱であり、予測が不正確な場合に汚染されたバッチが生じるリスクがある。

どちらのアプローチも、グループ内部のダイナミクスを観察する前に、プロンプトレベルで決定を下してしまう。しかし、本論文では、グループの有効性は、そのグループ自身のロールアウトのシーケンス内の早い段階で決定されることが多いという事実に着目している。すでに飽和することが判明したプロンプトに対して、フルグループの予算を使い切ることは計算資源の無駄である。

手法: SARA

著者らは、SARA (Sequential Adaptive Rollout Allocation) を提案する。これは、ステップごとのロールアウト収集を、予算制約のある逐次的割当（最適停止）問題として再定義するものである。すべてのプロンプトに対して固定数のロールアウト ($k$) を生成する代わりに、SARAはバッチ化されたラウンド内でプロンプトを探索し、観測された結果に基づいて信念を更新し、意思決定を行う。

コアメカニズム

1. ベイズモデリング: 各プロンプト $q$ について、SARAはその潜在的な成功率 $\gamma_q$ に関するベータ事後分布を保持する。初期状態では一様事前分布が使用される。$n$ 回のロールアウトで $s$ 回の成功が観測されると、事後分布は $\text{Beta}(\alpha_0 + s, \beta_0 + n - s)$ に更新される。
2. 閉形式の有効性予測器: SARAは、現在のプレフィックス（接頭辞）が与えられたとき、サイズ $k$ のグループが「有効（effective）」（混合された結果）になる事後予測確率 ($p_{\text{eff}}$) を計算する。
   • すでに混合状態（$1 \le s \le n-1$）である場合、$p_{\text{eff}} = 1$ となる。
   • すべて失敗、またはすべて成功のプレフィックスである場合、$p_{\text{eff}}$ はベータ関数を用いて解析的に計算される。一様事前分布を用い、すべて失敗のプレフィックスである場合、これは $p_{\text{eff}}(n, 0) = \frac{k-n}{k+1}$ と簡略化される。
3. 2閾値停止ルール: $p_{\text{eff}}$ に基づき、SARAはウォルドの逐次確率比検定 (SPRT) に似た逐次的決定ルールを適用する：
   • COMMIT（確定）: グループが混合（有効）である場合、そのグループは即座にトレーニングバッチに追加される。
   • ABANDON（放棄）: $p_{\text{eff}}$ が下限閾値 $\tau_{\text{low}}$ を下回った場合、そのプロンプトは飽和している可能性が高いと判断される。そのプロンプトに対する残りの予算は解放される。
   • CONTINUE（継続）: それ以外の場合、そのプロンプトにはもう一度ロールアウトが割り当てられる。
4. 予算の再割当: 放棄されたプロンプトから解放された予算は、プール内の新しいプロンプトへと即座に再割当される。これにより、固定された総予算から、一様割当よりも効果的なグループが得られることが保証される。

アルゴリズムの特性

• 直交性: SARAはロールアウト収集の段階で動作するため、あらゆるプロンプト選択戦略（例：Dynamic Samplingとの組み合わせ）と互換性がある。
• 追加ロールアウトなし: 難易度を推定するための補助的なモデル呼び出しを必要とする予測的手法とは異なり、SARAは最適化器が本来生成するはずのロールアウトのみを使用する。
• 同期: アルゴリズムは推論のスループットを維持するためにラウンド同期バッチで実行され、通常、1ステップあたり2〜4回の同期ラウンドのみを必要とする。

主な貢献

1. 問題の再定義: 著者らは「グループの有効性が早期に決定可能である」という性質を特定し、ロールアウト収集を、プロンプトレベルの選択とは異なる、逐次的割当問題として再構成した。
2. SARAアルゴリズム: ベータ・二項予測器と2閾値停止ルールを導出し、既存のGRPOパイプラインに統合可能な、予測モデルや追加ロールアウトを必要としないアロケーターを作成した。
3. 理論的保証:
   • 放棄の信頼性: 有効なグループを誤って放棄する確率は、閾値 $\tau_{\text{low}}$ によって抑えられる。
   • ロールアウトの節約: 1プロンプトあたりに費やされる期待ロールアウト数は、Dynamic Samplingで使用される固定数 $k$ よりも厳密に少なく、グループサイズ $k$ が大きくなるにつれて節約効果が増大する。
   • 収量の優位性: 固定予算において、SARAは一様割当と比較して、より高い、あるいは同等の数の有効なグループの収量を保証する。
   • 勾配との関連: 有効なグループの収量を最大化することは、GRPOの期待二乗勾配ノルムの下限を直接的に最大化することに直結する。
4. 実証的検証: 1.5Bおよび3Bモデルを用い、数学的推論およびプランニングのタスクにおける広範な実験を実施した。

実験結果

R1-Distill-Qwen-1.5BおよびQwen2.5-3Bモデルを用い、MATH、AIME24、Countdownなどのデータセットを用いて単一GPUで評価を行った：

• 効率性 vs. Dynamic Sampling (DS): SARAは、Dynamic Sampling（オラクルを使用して飽和したグループをフィルタリングするもの）と同等の精度を実現しながら、22%少ないロールアウトで動作した。
• 予測的選択との結合: SARAとDynamic Samplingを組み合わせた場合（SARA+DPS）は、最高の精度を示し、DSのオラクルをわずかに上回る一方で、DSよりも67%少ないロールアウトを実現した。
• トークン節約: 放棄された「すべて失敗」のトレースは長い傾向にあるため、トークンの節約量はロールアウトの節約量よりも顕著である。
• 堅牢性: 予測的選択とは異なり、ポリシーが変化しても精度が低下することなく、SARAはインサンプル検証に依存することで、トレーニング全体を通じてほぼ100%の有効なバッチ比率を維持した。
• 互換性: SARAは、一様ロールアウト収集の代替として、様々なRLアルゴリズム（PPO, GRPO, RLOO, Reinforce++）の性能を向上させる。

意義と主張

本論文は、高価なオーバーサンプリング（DSなど）を排除しつつ、予測の脆弱性を回避することで、SARAが「両取り」のソリューションを提供すると主張している。ロールアウトグループ内の統計的証拠を利用することで、SARAは補助的なモデル呼び出しを行うことなく、高いトレーニング効率を実現する。

著者らは、SARAをRLVRにおける基本的な効率化レバーとして位置づけており、特にグループサイズが分散低減のために増加する場合に重要であるとしている。また、本手法はバイナリの検証可能な報酬と、グループ内の独立同一分布（i.i.d.）なロールアウトを仮定しているものの、コアとなる逐次的割当のロジックはプロンプト選択や長さ制御の手法とは直交しており、連続的な報酬や木構造のロールアウトへの将来的な拡張が可能であると述べている。この研究は、推論能力の高いLLMのポストトレーニングにおいて、単なるプロンプトの精査ではなく、最適停止戦略を通じて大幅な計算資源の節約が可能であることを示している。