优化大模型强化学习的采样分配,提升训练效率
Spend Your Rollouts Where It Counts: Rollout Allocation for Group-Based RL Post-Training
- 按提示词重要性动态分配采样预算,避免浪费
- 在1.5B到14B模型上,采样成本更低且更快达到目标精度
- 适合大规模语言模型后训练的高效强化学习场景
强化学习是大语言模型后训练的主流范式。但在在线、同策略设置下,采样生成占训练计算成本的主要部分。分组策略优化方法对每个提示词使用多个采样,但盲目分配预算给奖励分布坍缩的提示词,造成昂贵采样的浪费。我们发现分组更新在奖励方差高的情况下最有效。由于策略随训练演化,提示词信息量必须在线估计而非预计算,但逐个评估所有提示词计算开销过大。为此,我们提出Pilot-Commit:一种面向分组强化学习后训练的预算感知采样分配框架。该框架将提示词评估与利用解耦:先用部分预算进行初步评估,再将剩余采样资源集中分配给高价值提示词,低信号提示词则跳过。在多个数学推理基准和1.5B至14B参数规模的模型上,Pilot-Commit以显著更低的采样成本达到基线准确率,在累计采样数上比GRPO快1.9倍,比DAPO快4.0倍。
原文摘要 · Abstract (English)
Reinforcement learning (RL) is the dominant paradigm for post-training large language models. However, in the online, on-policy setting, rollout generation dominates the computational cost of training. Group-based policy optimization methods compute advantages from multiple rollouts per prompt, yet they indiscriminately allocate budget to prompts with collapsed reward distributions, wasting expensive rollouts on negligible learning signals. We demonstrate that group-based updates are most effective in regimes of high reward variance. Since the policy evolves throughout training, prompt informativeness must be estimated online rather than precomputed, but exhaustively evaluating every prompt is computationally prohibitive. We introduce Pilot-Commit, a budget-aware rollout allocation framework for group-based RL post-training. Pilot-Commit decouples prompt evaluation from exploitation: a pilot stage estimates per-prompt informativeness using a fraction of the budget, and the remaining rollouts are allocated to high-leverage prompts while low-signal prompts are skipped. Across multiple math reasoning benchmarks and model scales from 1.5B to 14B parameters, Pilot-Commit matches baseline accuracy with significantly lower sampling costs, reaching target accuracy up to $1.9\times$ faster than GRPO and $4.0\times$ faster than DAPO in cumulative rollouts.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。