Not All Rollouts are Useful: Down-Sampling Rollouts in LLM Reinforcement Learning

要約

検証可能な報酬（RLVR）による強化学習は、大規模な言語モデルの推論能力を強化するための強力なパラダイムとして浮上しています。
ただし、計算およびメモリの要件における基本的な非対称性によって制約されています。ロールアウト生成は恥ずかしいほど並行してメモリライトであり、ポリシーの更新はコミュニケーションが多いメモリ集約型です。
これに対処するために、ポッド（ダウンサンプリングによるポリシーの最適化）を紹介します。
ポッドは並行して多数のロールアウトを生成し、その後、有益なサブセットのみでトレーニングを行い、更新コストを削減しながら学習信号を保存します。
報酬の多様性を最大化し、$ o（n \ log n）$ソリューションを認めることを示す原則的な基準である最大値のダウンサンプリングを備えたポッドをインスタンス化します。
経験的には、グループ相対ポリシー最適化（GRPO）との結合ポッドは、さまざまな推論ベンチマークやハードウェア環境で標準GRPOよりも優れたパフォーマンスを達成します。

要約(オリジナル)

Reinforcement learning with verifiable rewards (RLVR) has emerged as a powerful paradigm for enhancing reasoning capabilities in large language models. However, it is constrained by a fundamental asymmetry in computation and memory requirements: rollout generation is embarrassingly parallel and memory-light, whereas policy updates are communication-heavy and memory-intensive. To address this, we introduce PODS (Policy Optimization with Down-Sampling). PODS produces numerous rollouts in parallel, then trains on only an informative subset, preserving learning signals while slashing update cost. We instantiate PODS with max-variance down-sampling, a principled criterion that maximises reward diversity and show it admits an $O(n\log n)$ solution. Empirically, coupling PODS with Group Relative Policy Optimization (GRPO) achieves superior performance over standard GRPO across different reasoning benchmarks and hardware environments.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google