GRPOの効率を10倍にできるか？Kwai AIのSRPOが「はい」と答える

OpenAIのo1シリーズとDeepSeek-R1の目覚ましい成功は、大規模強化学習（RL）が高度な推論行動を引き出し、大規模言語モデル（LLM）の能力を大幅に向上させる力を明確に示しています。しかし、これらの革新的な推論モデルの核となる訓練手法は、技術レポートの中でしばしば明らかにされていません。最近のコミュニティの取り組みは主に数学的推論に焦点を当てており、ドメイン横断の汎化という課題はほとんど未探索のままです。さらに、標準的なGRPOアルゴリズムは、性能のボトルネック、サンプル効率の低さ、混合ドメインデータセットでの専門的推論スキルの育成の困難など、一般的な問題に悩まされており、LLM向けRL手法の効果的なスケーリングを複雑にしています。

これらの制限に対処するため、KuaishouのKwaipilotチームの研究者は、新しい強化学習フレームワークである「2段階ヒストリーリサンプリング方策最適化（SRPO）」を導入しました。この革新的なアプローチは、複数の次元にわたって前述の訓練課題に体系的に取り組みます。チームは訓練手法の詳細を説明したテクニカルレポートを公開し、SRPO-Qwen-32Bモデルをオープンソース化しました。

注目すべきは、この研究が数学とコードの両方の領域でDeepSeek-R1-Zeroレベルの性能を同時に達成した初めての例であることです。DeepSeekと同じベースモデル（Qwen2.5-32B）を使用し、純粋な強化学習訓練手法を採用することで、SRPOはAIME24（50）とLiveCodeBench（41.6）のベンチマークで印象的な結果を達成し、DeepSeek-R1-Zero-32Bを上回りました。さらに驚くべきことに、SRPOはR1-Zeroに必要な訓練ステップのわずか10分の1でこの性能を達成しています。

バニラGRPOの課題

初期の探索で、Kwaipilotチームは標準のGRPOアルゴリズムを試しました。しかし、すぐにボトルネックに直面し、モデルは望ましいR1-Zero性能レベルに達しませんでした。これらの問題には以下が含まれます：

• ドメイン間の最適化競合（数学 vs コード）：数学の問題はより長く詳細な推論軌跡（Long CoT）を引き出す傾向がありますが、コードデータはその傾向が弱いです。この2つのデータタイプを直接混合すると競合が生じ、両方のドメインで最適以下の性能となりました。
• 類似グループ報酬による訓練効率の低下：GRPOアルゴリズムは、サンプリングされたグループ内の非ゼロ報酬の分散に依存してアドバンテージを計算します。グループ内のロールアウトがほぼ同じ報酬値を生成する場合、計算されるアドバンテージはゼロに近づきます。訓練バッチのかなりの部分がこの現象を示すと、効果的な勾配の寄与が最小限になり、訓練効率が大幅に低下します。
• 早期の性能飽和：GRPO訓練では、ベンチマーク評価において早期の性能プラトーと報酬飽和が発生しました。この問題は、データ品質の不足に部分的に起因していました。訓練データが十分な複雑さや多様性を欠いている場合、特に簡単な問題が多すぎると、モデルは簡単なタスクでの性能を保守的に維持する傾向があり、難しい問題に必要な複雑で深い推論の開発が妨げられます。

2段階訓練

数学とコードのドメイン間の応答長さの競合に対処するため、Kwaipilotチームは2段階訓練パラダイムを実装しました：

• 第1段階：推論能力の引き出し：この初期訓練フェーズは、難易度の高い数学データにのみ焦点を当てます。主な目標は、モデルのテスト時拡張を完全に促進し、内省的休止、バックトラック、段階的分解などの能力を育成することです。
• 第2段階：スキル統合：このフェーズでは、コードデータが訓練プロセスに導入されます。第1段階で確立された推論基盤に基づき、このフェーズはさらにコーディング能力を強化し、手続き的思考、再帰、ツール呼び出し能力を徐々に強化することを目的としています。

異なる訓練データ戦略が応答長さに与える影響の分析から、以下の洞察が得られました：混合訓練では応答長さの伸びが限定的でベンチマーク性能が低い；数学のみの訓練では応答長さが安定して増加し数学ベンチマークで優れた性能を示し、強力で汎化可能な推論能力を育成；コードのみの訓練ではコードベンチマークの性能は向上するが明示的な推論行動の発達は最小限；段階的訓練は数学とプログラミングの両方で優れた結果をもたらし、モデルは詳細な段階的推論を生成し、数学的推論にコードを自発的に活用するなどの複雑な行動を示しました。

ヒストリーリサンプリング

Kwaipilotチームは、訓練の中後期では、バッチ内のサンプリンググループの約50％が同一の報酬を生成することを観察しました。これは、モデルが簡単な問題で一貫して成功した場合に頻繁に発生し、報酬の分散が小さくなり、勾配更新が無効になります。この非効率性に対処し、勾配信号の質を向上させるために、彼らはヒストリーリサンプリングを導入しました。訓練中、各エポック内のすべてのロールアウトの報酬結果を記録します。エポックの終わりに、以下の基準に基づいて次のエポックのデータセットを再構築します：

• 過度に単純なサンプルのフィルタリング：すべてのロールアウトが正解だったサンプルは除外されます。それらはポリシー改善のための情報シグナルを提供しないからです。
• 情報量の多いサンプルの保持：結果が多様（正解と不正解の両方を含む）またはすべて不正解のサンプルが保持されます。これらのサンプルは正の報酬分散を生み出し、非ゼロのアドバンテージと効果的な勾配信号を保証します。さらに、現在のエポックで全てのロールアウトが不正解だった困難なサンプルも保持されます。その理由は、これらの当初難しかった問題が更新されたポリシーにとって相対的に簡単になる可能性があり、その後の訓練で効果的な勾配を生成するからです。この戦略はカリキュラム学習の原理に沿っており、平均的にモデルを徐々に挑戦的なサンプルにさらし、訓練効率を高めます。

DAPOで提案された動的サンプリングと比較して、ヒストリーリサンプリングは計算効率を大幅に改善し、より安定した応答長さの成長をもたらしました。

データ

Kwaipilotチームは、公開されているCode&Mathデータセットに対して綿密なデータクリーニングとフィルタリングを実施しました。ヒューリスティックルールを適用して無関係なURLやフォーマットノイズを除去し、元データの核心フィールド（質問と回答の正解）の完全性を確保しました。PRIMEの数学データに対するデータクリーニングアプローチに従い、マルチパート問題、純粋な証明問題、画像やテーブルの理解を必要とする問題を削除しました。コードデータについては、特定の環境、ファイルI/O、ネットワーク相互作用に依存する問題を除外し、アルゴリズムロジックに焦点を当てました。

データ投入前に、数学とコードの問題の正しさを検証し、回答の正確性と解決可能性を確認し、不正確または曖昧な解決策を持つものを破棄しました。その後、各問題の難易度を評価し、通過率（Pass@k）に基づいて簡単、中程度、難しいの3つのレベルに分類しました。

実験結果

このセクションでは、SRPO法を用いた実験結果を詳述します。Kwaipilotチームは、訓練中の報酬や応答長さなどの指標の変化に焦点を当てました。訓練過程では、初期の報酬成長がプラトーに達した後、訓練は第2段階に移行しました。第2段階の開始時には、モデルが以前コードで訓練されていなかったため全体的な報酬が低下し、その後の訓練で着実に増加しました。コードデータの統合は応答長さを有意に増加させず、これは期待通りでした。同時に、ベンチマーク結果は、モデルの数学的およびコーディング能力が継続的かつ安定的に向上していることを示し、新しい方法の有効性を実証しました。

具体的には、ヒストリーリサンプリングにより、各訓練ステップで勾配更新が効果的に行われ、情報量の多い勾配の割合が直接増加しました。このサンプリング効率の向上により報酬の成長が安定し、リサンプリング戦略による訓練効率の向上が明確に示されました。

推論行動

Kwaipilotチームは3つの代表的な内省パターン（再チェック、ためらい、探索）を特定しました。これらのパターンを含む応答を統計的に分析し、各パターンの平均応答長さを記録しました。RL訓練中、モデルの自己内省、修正、バックトラックの頻度が徐々に増加し、「自己検証」能力が出現していることを示しました。彼らは、人間の認知プロセスに類似した「内省」がRL訓練中にモデルに現れるのは、ポリシー最適化プロセスの結果としての適応行動であると仮説を立てています。

訓練初期には、モデルは以前の推論ステップの積極的なチェックと内省をほとんど示しませんでしたが、訓練が進むにつれて、顕著な内省とバックトラック行動を示し、段階的推論、数値代入、段階的検証、自己最適化などの応答パターンを形成しました。興味深いことに、モデルは数学問題を解く際に、自発的にプログラムコードを使用して検証することを学習しました。最初に数学的推論で解法を提供し、次に能動的にコードを書いて答えの正しさを検証するのです。これらの事例は、モデルが手続き的思考を利用して自己修正と複数回の試行を行う能力を示しており、訓練後期にはモデルが広範な思考とコードベースの推論アプローチの統合的応用を習得したことを示しています。

論文：arXiv

モデル：HuggingFace