GRPO能高效10倍吗?Kwai AI的SRPO给出肯定答案
Kwai AI提出的SRPO框架将LLM强化学习后训练步骤减少90%,同时在数学和代码领域达到DeepSeek-R1水平。该两阶段强化学习方法结合历史重采样,克服了GRPO的局限性。
OpenAI的o1系列和DeepSeek-R1的巨大成功明确展示了大规模强化学习(RL)在激发复杂推理行为和显著提升大语言模型(LLM)能力方面的潜力。然而,这些突破性推理模型的核心训练方法往往在其技术报告中保持神秘。最近的社区工作主要集中在数学推理上,跨域泛化的挑战在很大程度上未被探索。此外,标准的GRPO算法存在性能瓶颈、样本利用效率低以及混合领域数据集难以培养专门推理技能等问题,阻碍了RL方法在LLM上的有效扩展。
为解决这些限制,来自快手Kwaipilot团队的研究人员引入了一种新颖的强化学习框架:两阶段历史重采样策略优化(SRPO)。该方法从多个维度系统性地应对上述训练挑战。团队公开发布了技术报告,详细阐述了其训练方法的复杂性,并开源了SRPO-Qwen-32B模型。
值得注意的是,这项工作首次在数学和代码领域同时实现了DeepSeek-R1-Zero级别的性能。通过使用与DeepSeek相同的基础模型(Qwen2.5-32B),并采用纯强化学习训练方法,SRPO在AIME24(50)和LiveCodeBench(41.6)基准测试上取得了令人印象深刻的成绩,超越了DeepSeek-R1-Zero-32B。更令人惊叹的是,SRPO仅用了R1-Zero所需训练步骤的十分之一就达到了这一性能水平。
Vanilla GRPO面临的挑战
在初步探索中,Kwaipilot团队尝试了标准GRPO算法。然而,他们很快遇到了瓶颈,模型无法达到期望的R1-Zero性能水平。这些问题包括:
- 跨领域优化冲突(数学与代码):数学问题倾向于产生更长、更详细的推理轨迹(Long CoT),而代码数据对此倾向较弱。直接混合这两种数据类型导致了冲突,使得两个领域的表现都不理想。
- 相似组奖励导致训练效率降低:GRPO算法依赖于采样组内非零奖励的方差来计算优势。当组内的rollout产生几乎相同的奖励值时,计算出的优势趋近于零。如果训练批次中有很大一部分出现这种现象,有效的梯度贡献就变得极小,大大降低了训练效率。
- 过早的性能饱和:GRPO训练在基准评估中遇到了早期性能平台和奖励饱和。这个问题部分归因于数据质量不足。当训练数据缺乏足够的复杂性或多样性,尤其是包含过多简单问题时,模型倾向于保守地维持其在简单任务上的表现,阻碍了其开发复杂深入推理能力。
两阶段训练
为了解决数学和代码领域之间固有的响应长度冲突,Kwaipilot团队实施了两阶段训练范式:
- 第一阶段:激发推理能力:这一初始训练阶段专注于高难度的数学数据。主要目标是充分激励模型的测试时扩展,培养诸如反思性暂停、回溯和逐步分解等能力。
- 第二阶段:技能整合:在这一阶段,代码数据被引入训练过程。基于第一阶段建立的推理基础,此阶段旨在进一步增强编码能力,同时逐步强化程序思维、递归和工具调用能力。
不同训练数据策略对响应长度的影响分析揭示了以下见解:混合训练导致响应长度增长有限且基准性能不佳;纯数学训练带来稳定的响应长度增长和优异的数学基准性能,并培养了强大且可泛化的推理能力;纯代码训练在代码基准上表现提升,但显式推理行为发展不足;而两阶段训练在数学和编程领域均取得了优异结果,模型能生成详细的逐步推理,甚至自发利用代码辅助数学推理。
历史重采样
Kwaipilot团队观察到,在训练的中后期,批次中近50%的采样组产生了相同的奖励。这通常发生在模型在简单问题上持续成功时,导致奖励方差极小和梯度更新无效。为了解决这个问题并提高梯度信号质量,他们引入了历史重采样。在训练过程中,他们记录每个epoch内所有rollout的奖励结果。在epoch结束时,他们根据以下标准重建下一个epoch的数据集:
- 过滤过于简单的样本:所有rollout都正确的样本被排除,因为它们不能提供有信息量的信号用于策略改进。
- 保留信息量丰富的样本:具有多样化结果(既有正确也有错误)或全部错误的样本被保留。这些样本产生正的奖励方差,确保非零优势和有效的梯度信号。此外,当前epoch中所有rollout都错误的困难样本也被保留。其理由是,这些最初具有挑战性的问题对于更新后的策略可能变得相对容易,从而在后续训练中产生有效的梯度。这种策略符合课程学习的原则,逐渐使模型平均面对更具挑战性的样本,以提高训练效率。
与DAPO中提出的动态采样方法相比,历史重采样显著提高了计算效率,并带来了更稳定的响应长度增长。
数据
Kwaipilot团队对公开的Code&Math数据集进行了细致的数据清洗和过滤。他们应用启发式规则过滤掉无关URL、格式噪声,并确保原始数据中核心字段(问题和答案真实值)的完整性。遵循PRIME对数学数据的数据清洗方法,他们移除了多部分问题、纯证明题以及需要图片或表格理解的问题。对于代码数据,他们排除了依赖特定环境、文件I/O或网络交互的问题,专注于算法逻辑。
在数据导入前,他们对数学和代码问题进行了正确性验证,确保答案的准确性和可解性,丢弃了不正确或模糊的解答。随后,他们评估了每个问题的难度,根据通过率(Pass@k)将其分为简单、中等和困难三个级别。
实验结果
本节详细介绍使用SRPO方法获得的实验结果。Kwaipilot团队重点关注训练过程中奖励和响应长度等指标的变化。
在训练过程中,初始奖励增长进入平台期后,训练进入第二阶段。第二阶段开始时,由于模型之前缺乏代码训练,整体奖励下降,随后在后续训练中稳步上升。集成代码数据并未显著增加响应长度,这符合他们的预期。同时,基准测试结果表明,模型在数学和编码能力上持续稳定提升,证明了新方法的有效性。
具体来说,历史重采样确保了每个训练步骤的梯度更新保持有效,直接提高了信息量丰富的梯度比例。这种增强的采样效率带来了稳定的奖励增长,清晰展示了重采样策略带来的训练效率提升。
推理行为
Kwaipilot团队识别了三种代表性的反思模式:重新检查、犹豫和探索。他们对包含这些模式的响应进行了统计分析,并记录了每种模式的平均响应长度。在RL训练过程中,他们观察到模型自我反思、纠正和回溯的频率逐渐增加,表明“自我验证”能力正在出现。他们认为,类似于人类认知过程的“反思”在RL训练中出现在模型中,是策略优化过程产生的适应性行为。
训练早期,模型几乎没有主动检查和反思之前的推理步骤;但随着训练进行,模型表现出显著的反思和回溯行为,形成了诸如逐步推理、数值代入、逐步验证和自我优化等响应模式。有趣的是,他们还发现模型学会了在解决数学问题时自发使用程序代码进行验证。它会先通过数学推理提供解决过程,然后主动编写代码验证答案的正确性。这些实例表明模型能够利用程序思维进行自我纠正和多次尝试,进一步表明在训练后期,模型已经掌握了广泛的思维和基于代码的推理方法的综合应用。
论文地址:arXiv
模型地址:HuggingFace