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

📄 arXiv: 2504.13818v3 📥 PDF

作者: Yixuan Even Xu, Yash Savani, Fei Fang, J. Zico Kolter

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-04-18 (更新: 2025-10-01)

备注: 17 pages, 8 figures

💡 一句话要点

提出PODS以解决LLM强化学习中的计算与内存不对称问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 大型语言模型 策略优化 下采样 计算效率 奖励多样性 推理能力

📋 核心要点

  1. 现有的强化学习方法在回滚生成和策略更新之间存在计算与内存的显著不对称性，导致效率低下。
  2. 本文提出PODS，通过对回滚进行战略性下采样，解耦回滚生成与策略更新，降低计算成本。
  3. 实验结果表明，结合PODS的GRPO在不同基准测试中比传统方法快至少1.7倍，显著提升了效率。

📝 摘要（中文）

强化学习与可验证奖励（RLVR）已成为提升大型语言模型推理能力的主要方法。然而，该方法面临计算和内存的不对称性：回滚生成过程并行且内存占用少，而策略更新则通信密集且内存需求高。为此，本文提出PODS（基于下采样的策略优化），通过仅对战略性选择的回滚子集进行训练，解耦回滚生成与策略更新，保持学习质量的同时显著降低更新成本。我们提出了一种原则性的子集选择标准——最大方差下采样，旨在最大化奖励多样性，并提供了高效的$O(n ext{log} n)$实现。实验证明，结合PODS的群体相对策略优化（GRPO）在不同推理基准和硬件配置下，测试准确率至少比传统GRPO快1.7倍。

🔬 方法详解

问题定义：本文旨在解决强化学习中回滚生成与策略更新之间的计算与内存不对称问题。现有方法在生成回滚时效率高，但在策略更新时却面临高内存和通信成本的挑战。

核心思路：论文提出PODS，通过仅对经过战略性选择的回滚子集进行训练，来解耦回滚生成与策略更新，从而在保持学习质量的同时降低更新成本。

技术框架：整体架构包括回滚生成模块和策略更新模块。回滚生成模块负责并行生成大量回滚，而策略更新模块则基于下采样的回滚子集进行高效的策略优化。

关键创新：最重要的技术创新在于提出了最大方差下采样的选择标准，旨在最大化奖励的多样性，这与传统方法的全量回滚训练方式形成了本质区别。

关键设计：在实现中，采用了高效的$O(n ext{log} n)$算法来进行下采样，确保在选择回滚子集时能够快速且有效地最大化奖励多样性。

📊 实验亮点

实验结果显示，结合PODS的GRPO在不同推理基准测试中，测试准确率比传统GRPO快至少1.7倍，展示了显著的性能提升。这一结果表明PODS在降低计算成本的同时，能够有效保持学习质量。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统以及其他需要推理能力的AI应用。通过提高强化学习的效率，PODS可以帮助开发更智能的语言模型，推动AI在复杂任务中的应用和发展，具有重要的实际价值和未来影响。

📄 摘要（原文）

Reinforcement learning with verifiable rewards (RLVR) has emerged as the leading approach for enhancing reasoning capabilities in large language models. However, it faces a fundamental compute and memory asymmetry: 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), which decouples rollout generation from policy updates by training only on a strategically selected subset of rollouts, maintaining learning quality while dramatically reducing update costs. We propose a principled subset selection criterion, max-variance down-sampling, that maximizes reward diversity, and provide an efficient $O(n\log n)$ implementation. Empirically, Group Relative Policy Optimization (GRPO) with PODS achieves the peak test accuracy of vanilla GRPO at least $\mathbf{1.7\times}$ faster across the different reasoning benchmarks and hardware configurations we tested.