1/10训练数据超越GPT-4o！清华等提出隐式过程奖励模型PRIME，在线刷SOTA

文章转载于新智元
打破数据墙，我们还能做些什么？


近日，来自清华UIUC等机构的研究者提出了PRIME（Process Reinforcement through IMplicit REwards）：通过隐式奖励来进行过程强化。


GitHub地址：https://github.com/PRIME-RL/PRIME
这是一种带有过程奖励的在线RL开源解决方案，可以提高语言模型的推理能力，超越了SFT（监督微调）或者蒸馏等方法。


对比SFT，PRIME让模型在重要基准测试上实现了巨大提升：平均提高了16.7%，在AMC和AIME中都提高了20%以上。
Eurus-2-7B-PRIME与Qwen2.5-Math-7B-Instruct，使用了相同的base model（Qwen-2.5-Math-7B），但在上表的6项测试中，5项都超越了instruct版本，同时也超越了GPT-4o。
而这个成绩只用了Qwen Math 1/10的数据资源（230K SFT + 150K RL）！


作者发布了本研究中使用的所有模型和数据，感兴趣的读者请见文后链接。
1
过程奖励模型
热身阶段（SFT）
如前所述，选择Qwen2.5-Math-7B-Base作为起点，然后上点难度，采用竞赛级别的数学和编程基准，包括AIME 2024、AMC、MATH-500、Minerva Math、OlympiadBench、LeetCode和LiveCodeBench（v2）。
首先对基础模型进行监督微调，以获得RL的入门模型（教模型学习某些推理模式）。
为此，研究人员设计了一个以动作为中心的链式推理框架，策略模型在每个步骤中选择7个动作中的一个，并在执行每个动作后停止。
为了构建SFT数据集，研究者从几个开源数据集中收集了推理指令。
值得注意的是，对于许多具有真实答案的数据集，作者选择将其保留用于之后的RL训练，目的是让SFT和RL使用不同的数据集，以使RL中的探索多样化，并且作者认为在PL中真实标签更加重要。
作者用LLaMA-3.1-70B-Instruct来回答指令，并使用系统提示要求模型执行以动作为中心的思维链。
隐式PRM

下面接入过程奖励模型（PRM），这里采用隐式PRM，只需要在响应级别标签上训练ORM。


过程奖励模型简单理解就是对每个推理步骤进行评分，举个例子：


PRM是以这种粒度来评价响应的。
在本文的隐式PRM中，可以使用以下方式免费获得过程奖励：


通过简单地收集响应水平数据和训练ORM来获得PRM，而无需注释步骤标签。
这与ORM训练目标的具体选择无关，比如使用交叉熵损失来实例化隐式PRM，就可以替换成：


强化学习

本文的目标是广泛利用强化学习（RL）来提高推理能力。针对这种资源有限的情况，作者总结了一些最佳实践：
从Ground Truth验证器和高质量数据开始：作者进行了严格的数据收集和清理，以获得可验证的RL数据，并发现仅使用结果验证器足以构建强大的基线。 作者比较了不同的RL算法得出结论，无价值模型的REINFORCE类方法足够有效。 使用「mid-difficulty」问题进行稳定训练：作者提出了一种名为在线提示过滤器的机制，通过过滤掉困难和简单的问题，在很大程度上稳定了RL训练。使用PRM进行强化学习

将PRM集成到在线强化学习中并非易事，这里有几个需要解决的关键挑战。
如何为强化学习提供密集奖励？
奖励稀疏性一直是强化学习中长期存在的问题。到目前为止，我们仍然没有特别好的解决方案来为LLM的在线强化学习构建密集奖励。
以前的方法主要是为密集奖励建立一个额外的价值模型，众所周知，这样的模型很难训练，而且性能提升不大。
根据前文对隐式PRM的介绍，使用
可以免费从隐式PRM中获得token级别的过程奖励。
这种方式可以直接取代PPO中的价值模型，非常容易与任何优势估计函数和结果奖励相结合。在实践中，作者将过程奖励与REINFORCE、RLOO、GRPO、ReMax和PPO集成在一起，并进行了细微的修改。
如何设置一个好的PRM来启动RL？
即使我们找到了在RL中使用过程奖励的途径，训练好的PRM也并非易事：需要收集大规模（过程）奖励数据（很贵），并且模型应该在泛化和分布偏移之间取得良好的平衡。
隐式PRM本质上是一种语言模型。因此从理论上讲，可以使用任何语言模型作为PRM。在实践中，作者发现最初的策略模型本身就是的一个很好的选择。
如何在线更新PRM以防止奖励黑客攻击？
在线RL中，避免RM被过度优化或被黑客入侵至关重要，这需要RM与策略模型一起不断更新。然而，鉴于步骤标签的成本很高，在RL训练期间很难更新PRM，——可扩展性和泛化问题。
但是，本文的隐式PRM仅要求更新结果标签。也就是说，使用结果验证器即可在训练期间轻松更新PRM。
此外，还可以进行双重转发：首先使用策略部署更新PRM，然后使用更新的PRM重新计算过程奖励，从而提供更准确的奖励估算。
PRIME算法

下图表示PRIME算法的整个循环：


策略模型和PRM都使用SFT模型进行初始化。对于每个RL迭代，策略模型首先生成输出。然后，隐式PRM和结果验证器对输出进行评分，隐式PRM在输出时通过结果奖励进行更新。最后，将结果奖励ro和过程奖励rp组合在一起，用于更新策略模型。
以下是算法的伪代码：


实验

默认情况下，使用SFT模型初始化隐式PRM，并保留SFT模型作为参考对数探测器。超参数方面，策略模型的学习率固定为5e-7，PRM学习率为1e-6，使用AdamW优化器，mini batchsize大小为256，micro batchsize为8。
rollout阶段收集256个提示，每个提示采样4个响应。PRM训练时β=0.05，所有实验中将KL系数设置为0。


将PRIME与仅带有结果验证器（OV）的RLOO进行比较，与稀疏奖励相比，PRIME将RL训练加速了2.5倍，并将最终奖励提高了6.9%，且方差更低。在下游任务上，PRIME的性能也始终优于OV。


下面展示PRM在线更新的重要性。比较两种设置：在线PRM使用Eurus-2-7B-SFT初始化，离线PRM使用EurusPRM-Stage1初始化。
从下图中可以看出，在线PRM在训练集和测试集上的性能都大大优于离线PRM。




参考资料：
https://curvy-check-498.notion.site/Process-Reinforcement-through-Implicit-Rewards-15f4fcb9c42180f1b498cc9b2eaf896f
https://the-decoder.com/ai-learns-math-better-with-new-approach-that-uses-a-fraction-of-the-data/