使用 GRPO 在 "Temporal Clue" 游戏中击败 o1, o3-mini 和 R1

Brad Hilton, Kyle Corbitt 2025年3月6日 7 分钟 在这篇文章中,我们将讨论如何使用 GRPO 在一个名为“temporal clue”的推理游戏中超越 R1、o1、o3-mini,并达到 Sonnet 3.7 的几个百分点以内,同时在推理时运行成本降低 100 倍以上。我们将分享关于任务设计和超参数的具体经验,以及我们发现有效的技巧。最后,我们将分享我们用于实现这些结果的训练方案,该方案建立在 torchtune 之上。

背景

自从 OpenAI 推出其强大的新型 o 系列推理模型以来,我们看到了使用强化学习 (RL) 训练的 Large Language Models (LLMs) 的快速发展。 Google DeepMindAlibabaDeepSeekAnthropic 等领先组织迅速跟进,并训练了自己的高级模型,以使用长的“思维链 (CoT)”进行推理,并通过对可验证问题进行强化学习来训练。许多先前具有挑战性的基准测试(例如数学和编码领域)现在已接近饱和。

然而,尽管取得了令人印象深刻的进步,但对于当今最好的模型来说,逻辑演绎仍然非常困难。通常,LLM 很难始终如一地关注所有相关细节,保持逻辑上合理的推理链,或可靠地连接多个演绎步骤。即使是最先进的模型,生成的输出长度是原来的 10-100 倍,也经常会引入人类解算器很容易发现的基本错误。

出于对这个未解之谜的好奇,我们戴上了猎鹿帽,开始调查:较小的、开放权重的模型是否可以通过最新的强化学习技术达到前沿水平的演绎性能?我们从明显较弱的模型开始,并在新的演绎任务上迭代地训练它们。随着时间的推移,我们观察到他们侦探能力的明显提高,最终达到甚至超过了一些最强大的专有模型。

现在,我们很高兴分享我们的发现,包括我们的实验训练方案数据集模型权重,所有这些都在 MIT 许可下免费提供,以及关键的实践见解(就在这里)。拿起你的放大镜,侦探们;游戏开始了!

基准测试

为了开始我们的实验,我们首先必须确定一个具有挑战性的推理任务,该任务具有清晰可验证的解决方案和可扩展的复杂性。碰巧的是(并非巧合),作者之一先前创建了一个名为 Temporal Clue 的谜题集,该谜题集完美地满足了这些需求。除了满足基本事实清晰度的标准之外,还可以根据需要创建新的谜题——这是一个很好的附加功能。

Temporal Clue 的灵感来自流行的棋盘游戏 Clue (Cluedo),玩家竞相揭露是谁在宏伟的庄园中杀害了 Boddy 先生。 Temporal Clue 将游戏变成了一个单独的逻辑谜题,该谜题扩展到标准维度——,(用)什么,和哪里——并融入了两个额外的维度:何时(时间)和 为什么(动机)。谜题是随机生成的,并使用 OR-Tools' CP-SAT solver 选择最少但足够的线索。

在一个黑暗的冬夜,富有而神秘的 John Q. Boddy 先生为他的一些最亲密的伙伴举办了一个小型但奢华的晚宴。然而,当 Boddy 先生在凌晨时分被发现死在 Tudor 庄园的一个房间里时,这个夜晚以悲剧告终。以下相关人员已被确定为嫌疑人... (略)

为了确定此演绎任务的当前技术水平,我们对领先的推理模型(包括 DeepSeek R1、OpenAI 的 o1 和 o3-mini 以及 Anthropic 的 Claude Sonnet 3.7)进行了基准测试。此外,我们还对 14B 和 32B 的 Qwen 模型进行了基准测试,我们稍后将使用强化学习对其进行改进,并包含我们最终结果的预览:

组织 | 模型 | 推理工作量 | 平均准确率 | 平均成本 ---|---|---|---|--- DeepSeek | R1 | 默认 | 51.6% | $0.029 OpenAI | o1 | 默认 | 54.9% | $0.901 OpenAI | o3-mini | 中等 | 55.9% | $0.068 OpenAI | o3-mini | 高 | 56.7% | $0.170 Anthropic | Sonnet 3.7 | 无 | 51.7% | $0.017 Anthropic | Sonnet 3.7 | 16k Token Budget | 61.7% | $0.222 Anthropic | Sonnet 3.7 | 64k Token Budget | 69.5% | $0.392 Alibaba | Qwen 2.5 14B Instruct | 无 | 28.1% → 59.4% | $0.001 Alibaba | Qwen 2.5 32B Instruct | 无 | 37.3% → 67.1% | $0.002

从这些基准测试中,我们看到具有 64k token thinking budget 的 Claude Sonnet 3.7 在我们的任务中表现最佳,但所有领先的模型都显示出改进的空间。 DeepSeek R1,一种流行的开放权重模型,表现几乎与 OpenAI 的 o1 和 o3-mini 一样好。但是,相比之下,未经调整的 Qwen 2.5 Instruct 模型的性能令人印象深刻。最大的问题是:我们能否训练这些较小的、开放权重的模型达到前沿水平的性能?简单来说,我亲爱的读者——我们只需要正确的方法。

训练

为了训练前沿水平的演绎模型,我们转向了强化学习——这种方法允许智能体从其在受控环境中的自身经验中学习。在这里,LLM 是我们的智能体,而谜题是我们的环境。我们通过让 LLM 为每个谜题生成多个响应来引导 LLM 的学习,从而探索问题范围。我们加强了导致正确解决方案的演绎,并惩罚了使模型误入歧途的推理。

在各种 RL 方法中,我们选择了由 DeepSeek 开发的流行的 Group Relative Policy Optimization (GRPO) 算法。与更传统的方法(如 Proximal Policy Optimization (PPO))相比,GRPO 简化了训练过程,同时仍能提供强大的性能。为了加快我们的实验,我们省略了 Kullback–Leibler (KL) divergence 惩罚,尽管我们的训练方案支持感兴趣的读者使用它。

从高层次上讲,我们的训练循环遵循以下基本步骤:

对于生成响应,我们使用了流行的 vLLM 推理引擎。我们调整了参数选择,以最大限度地提高吞吐量并最大限度地缩短启动时间。前缀缓存尤为重要,因为我们为每个任务采样了许多响应,并且缓存提示有助于避免冗余计算。

我们观察到,用过多的请求淹没 vLLM 会强制抢占或换出正在进行的请求。为了解决这个问题,我们使用了一个信号量来限制请求,该信号量经过调整以保持较高的键值 (KV) 缓存利用率,同时最大限度地减少交换。更高级的调度机制可以产生更高的利用率,同时仍支持灵活的生成长度。

采样后,我们使用标准的 HuggingFace Transformers AutoTokenizer 处理完成。其聊天模板功能将消息对象呈现为提示字符串,包括一个用于确定 LLM 生成哪些 token 的助理掩码。我们发现模型在其默认模板中缺少必要的 % generation % 标签,因此我们在 token 化步骤中对其进行了修改。生成的助理掩码包含在用于调整的张量字典中,用于标识哪些位置需要损失计算。

在 token 化响应并获得助理掩码后,我们打包数据以进行调整。除了在每个打包序列中包含多个提示/响应对之外,我们还标识了共享提示 token,并将每个 token 分配一个父 ID 以及标准组 ID。特别是对于像 Temporal Clue 这样的任务(每个谜题平均超过 1,000 个 token),为每个任务生成大量响应并有效地打包张量显着减少了冗余。一旦用所有必要的信息打包好,我们就可以二维地可视化我们的训练数据集,每一行都是一个 token 序列,可能包含多个提示和完成:

有了紧密打包的数据,我们就可以继续进行调整。我们的模型已经经过预训练、指令调整,相当智能,并且擅长遵循指令。但是,它们还不能可靠地解决 Temporal Clue 谜题。尽管如此,它们偶尔会成功,这就足够了。通过增加良好推理的概率并降低“不好”推理的概率,我们逐渐引导模型走向 Master Detective 状态。我们使用标准的机器学习技术来实现这一点,采用策略梯度方法来计算损失并有益地改变权重。

对于训练,我们使用了 PyTorch 团队提供的 torchtune 库。 Torchtune 具有高效的仅解码器 transformer 实现,适用于流行的模型,包括 Llama、Gemma、Phi 等。尽管我们主要将 Qwen 模型用于此项目,但我们也运行了 8B 和 70B Llama 模型的实验。 Torchtune 还提供了节省内存和提高性能的实用程序,包括:

有关支持的优化的完整列表,请参见此处的 README

此外,Torchtune 支持多设备(现在是多节点)训练,使其成为较大模型的理想选择。它支持完全分片数据并行 (FSDP) 和张量并行 (TP) 训练,两者可以组合。他们还提供了十几个以上的方案,鼓励用户复制并自定义它们以用于他们的用例。我们创建了完整微调方案的修改版本,支持:

该方案可以在此处查看。将来,我们希望添加张量并行支持并探索 PEFT 和量化。

RL 训练过程涉及选择大量的超参数。在训练我们的模型时,我们测试了各种配置,并在很大程度上确定了以下内容:

批量大小是可变的,因为响应长度在训练期间可能会有所不同,序列打包效率在每次迭代中都会波动,并且会丢弃具有零优势的响应。对于一次运行,我们尝试动态调整与批量大小成反比的学习率,但这导致小批量大小的学习率过高,需要上限。上限版本与使用恒定学习率没有显着差异,但是调整批量大小和学习率仍然是未来实验的一个有趣领域。

我们还运行了简短的实验,增加了每次迭代的任务数,同时减少了每个任务的样本数(反之亦然),使每次迭代的总样本数大致相等。在较短的训练范围内,这些变化没有显示出明显的差异,这表明该方案对任务数量和每个任务的样本数量之间的不同平衡具有鲁棒性。

结果

在训练我们的模型 100 多次迭代后,我们达到了前沿水平的演绎。 我们的模型迅速改进,然后准确率增益开始逐渐减少,最终下降,有时会急剧下降。在最佳状态下,14B 模型接近 Claude Sonnet 3.7 在 16k 个 token 时的性能,而 32B 模型几乎与 Sonnet 在 64k budget 下的结果相匹配。

在训练过程中,性能增益遵循幂律,在 log-log 图上形成线性关系(在恶化之前)。 我们怀疑这些模型可能过早地收敛于一开始就奏效的贪婪策略,但这可能会限制它们的长期前景。一个合理的下一步是探索鼓励多样化响应的方法,或者逐步建立能力的方法(例如通过课程学习),或者为特别出色的解决方案分配更大的奖励以激励彻底的探索。

此外,我们还注意到训练期间输出长度的有趣模式。最初,响应变长,然后稳定,然后在训练结束时发散,14B 模型的响应变长,而 32B 模型的响应长度崩溃,尤其是在达到最佳性能之后。 为了定性地评估逻辑推理的改进,我们要求最强大的前沿模型 Claude Sonnet 3.7 识别并评估 Qwen 32B 模型在相似谜题上进行 100 多次迭代训练前后的演绎的合理性。Sonnet 从基本模型中识别出 6 个演绎,其中除一个以外的所有演绎都被判断为错误;相反,它从训练后的模型中识别出 7 个演绎,其中除一个以外的所有演绎都被判断为逻辑上合理的。

最后,假设 Fireworks AIserverless pricing tiers. 具有 按需部署足够吞吐量,我们估计了 Qwen 模型的成本。我们绘制了准确率与每个响应的平均推理成本的自然对数的关系图,并观察到未经调整的模型之间存在清晰的线性 Pareto 前沿。通过成功地训练开放权重模型以达到前沿水平的准确率,我们极大地改善了成本-准确率的权衡。

在这里,在为一项出色的工作分享满意的点头后,我们欢呼着一辆双轮马车,然后返回 Baker Street——这是思考我们发现的完美场所。

结论

在我们的调查中,我们着手探索较小的、开放权重的语言模型是否可以通过强化学习实现前沿水平的演绎推理。在使用精心选择的超参数和 GRPO 方法在具有挑战性的 Temporal Clue 谜题上训练 Qwen 14B 和 32B 模型后,我们取得了令人印象深刻的性能提升。这些改进使开放权重模型达到了推理性能的前沿,并大大降低了成本。我们的发现突出了强化学习在有效地训练开放模型以执行复杂的演绎任务方面的有希望的潜力。

如前所述,数据集实验训练方案 和模型权重(14B32B)在 MIT 许可下免费提供。我们鼓励您尝试重现和改进我们的结果。

此外,我们在最后保留了一个特别令人兴奋的发现。我们发现,仅需 16 个训练示例 即可实现高达 10-15% 的有意义的性能改进。这意味着您不需要大量数据即可开始;只需要对您想要解决的问题有一些直觉即可。

您是否有兴趣使用强化学习来训练自己的模型,或者需要一些入门帮助?请随时通过 OpenPipe 与我们联系——我们很乐意聊天!

现在,亲爱的读者,请随时准备好您的猎鹿帽和放大镜;还有更多需要探索。游戏仍然非常激烈。