ReCo: Reweighting GRPO Against Distributional Concentration

問題

Group Relative Policy Optimization (GRPO) は、数学やコーディングのような検証可能な報酬が存在するドメインにおいて、ポストトレーニングの言語モデルにおける標準的な強化学習手法として台頭しています。GRPOは、サンプリング予算が小さい場合（Pass@$k$ における低い $k$）にはサンプリング効率を向上させますが、近年の観察では、サンプリング予算が大きい場合（高い $k$）にはベースモデルを上回ることができず、しばにベースモデルを下回る性能を示すことが指摘されています。これは、推論パスのカバー率（coverage）が減少していることを示唆しています。本論文は、GRPOがベースモデルが既に高い確率で生成しているレスポンスに確率質量を集中させてしまうことで、推論の多様性を減少させ、ベース分布の低確率領域に存在する正しい解を発見する能力を制限していることを明らかにしています。

手法

著者らは、この集中を緩和するためにGRPOの更新を2つのレベルで修正する手法である ReCo (Reweighting GRPO Against Distributional Concentration) を提案しています。

1. レスポンス・レベルの再重み付け (Response-Level Reweighting):
   標準的なGRPOでは、高確率のレスポンスがロールアウトグループ内でより頻繁にサンプリングされます。グループのアドバンテージはサンプリングされたレスポンス間で等しく平均化されるため、これらの頻出する軌道は、単に繰り返し発生することによって、より大きな累積勾配質量を蓄積してしまいます。ReCoは、各レスポンスの勾配寄与を、グループ内での期待出現回数によって正規化することで、これに対処します。

   • 重みは、期待される出現回数の逆数として定義されます：$w_i = \frac{1}{G \cdot \bar{\pi}_{\theta_{old}}(o_i | q)}$。ここで $\bar{\pi}$ は長さ正規化された尤度です。
   • これにより、頻出するレスポンスがグループレベルの勾配を支配することを防ぎ、古いポリシーの下ですでに高確率となっている軌道を不当に強化しないようにします。

2. トークン・レベルの再重み付け (Token-Level Reweighting):
   標準的なGRPOは、トークンレベルの勾配をスケールさせるために重要度比（$\pi_\theta / \pi_{\theta_{old}}$）を使用します。これは、現在のポリシーの下でより高確率になったトークンがより大きな更新を受けるというフィードバックループを生み出し、飽和した決定点をさらに強化し、代替となるパスを減少させます。

   • ReCoは、標準的な重要度比を分散ベースの比率 $r_{var} = \frac{p_\theta(1 - p_\theta)}{p_{old}(1 - p_{old})}$ に置き換えます。
   • 項 $p(1-p)$ はベルヌーイ分散を表します。この比率は、非飽和な決定点（$p$ が0でも1でもなく、代替となるトークンの選択肢が依然として妥当である地点）に対して大きな更新スケールを割り当て、飽和した選択（$p \approx 1$）の重みを下げます。
   • このメカニズムは、ポリシーが単一の選択肢に崩壊していない決定点での更新を促進し、それによって推論の多様性を維持します。

主な貢献

• GRPOの集中に関する分析: 本論文は、Pass@$k$ の減少を、GRPOの更新における2つの具体的なメカニメント、すなわちグループレベルでの頻出レスポンスによる支配と、トークンレベルでの飽和トークンの増幅に帰着させています。
• ReCoアルゴリズム: 期待出現回数によってレスポンスの寄与を正規化し、トークンの重要度比を分散ベースの指標に置き換える、2レベルの再重み付けスキームの導入。
• 実証的検証: 異なるモデルファミリー（Qwen2.5-Math, Llama-3.1）およびスケール（1.5B, 7B, 8B）を用い、5つの数学的推論ベンチマーク（AIME 2024/2025, AMC 2023, OlympiadBench, MATH 500）における評価。

結果

• Pass@$k$ パフォーマンス: ReCoは、$k$ が大きい値（例：Pass@64）において一貫してGRPOを上回っており、これは正しい推論軌道のカバー率が向上していることを示しています。$k$ が小さい場合、ReCoはGRPOと同等の性能を維持し、サンプリング効率を保っています。
• ベースモデルの回復: 標準的なGRPOがPass@$k$ においてベースモデルの性能を下回ることが観察されたLlama-3.1-8B-Instructバックボーンにおいて、ReCoはこの劣化を大幅に軽減し、特定のベンチマーク（AIME 2024/2025）ではベースモデルをも上回りました。
• 多様性指標: 学習ダイナミクスの分析により、GRPOが急速に多様性を失う（Distinct-2の減少、Self-BLEUの上昇）一方で、ReCoは学習を通じてベースモデルに近いロールアウトの多様性を維持していることが示されました。
• トークン分散: ReCoは、レスポンスの初期および中間セグメントにおいて、GRPOと比較して実質的に高いトークンレベルのベルヌーイ分散を維持しており、これは推論パスがより長く非飽和状態であることを示しています。
• 汎用性: この手法はコード推論ベンチマーク（HumanEval, HumanEval+, LiveCodeBench）におけるPass@$k$ も改善しており、集中問題が数学的推論に限定されないことを示唆しています。
• 互換性: ReCoは、さらなる改善を得るために他のGRPOのバリアント（DAPO, 80/20 Ruleなど）と組み合わせることが可能です。

意義

本論文は、外部の多様性目的関数や正則化を必要とせずに、分布の集中を引き起こすGRPOの内部メカニズムに対処するReCoを提唱しています。出現確率と決定点の飽和に基づいて更新自体を再重み付けすることにより、ReCoはより広範な正しい推論パスを保持します。これにより、モデルは強化学習のサンプリング効率の利点を維持しながら、解空間のカバー範囲を拡大し、標準的なGRPOの学習で見られる「多様性の崩壊」を効果的に緩和できます。本研究は、推論能力を拡張する上でのRLVRの限界が、ポリシー更新自体の最適化ダイナミクスに起因している可能性があり、それはターゲットを絞った再重み付けによって修正可能であることを示唆しています。