近年来,人工智能领域在推理能力上取得了显著进展。去年,OpenAI 展示了大型语言模型 (LLMs) 的强大推理潜力后,Google DeepMind、阿里巴巴、DeepSeek 和 Anthropic 等机构迅速跟进,利用强化学习 (RL) 技术训练出具备“思维链” (CoT) 能力的先进模型。这些模型在数学、编程等领域的许多基准测试中都取得了接近饱和的成绩。
然而,一个不容忽视的事实是:即使是目前最顶尖的模型,在处理逻辑推理问题时,仍然面临着一些难以逾越的障碍。大型语言模型通常难以持续关注所有相关细节,维持逻辑推理链条的连贯性,或者可靠地连接多个推理步骤。即使是那些生成内容长度达到普通模型 10-100 倍的最先进模型,也常常会出现一些人类解题者很容易发现的低级错误。
为了探索这个问题,本文作者 Brad Hilton 和 Kyle Corbitt 开始思考:能否利用最新的强化学习技术,让规模较小、开源的模型在推理性能上达到前沿水平?为了解答这个问题,他们选择了一种名为 Group Relative Policy Optimization (GRPO) 的强化学习算法。
什么是 GRPO?
简单来说,GRPO 是一种优化策略梯度的方法。策略梯度方法通过调整模型的策略(即模型在特定情况下采取行动的概率分布)来改进模型的性能。GRPO 通过比较一组模型响应的相对优势,来更有效地更新策略。与 Proximal Policy Optimization (PPO) 等传统方法相比,GRPO 简化了训练流程,同时仍能保证出色的性能。
为了验证 GRPO 的有效性,他们选择了一个名为“时间线索 (Temporal Clue)”的推理游戏作为实验平台。
什么是“时间线索”游戏?
“时间线索”是一款受经典棋盘游戏“Clue (Cluedo)”启发的解谜游戏。在这个游戏中,玩家需要通过一系列线索,找出凶手、凶器、作案地点、作案时间以及作案动机。与原版游戏不同的是,“时间线索”将问题扩展到了五个维度,增加了游戏的复杂性和挑战性。
作者从相对较弱的模型入手,在“时间线索”游戏上利用 GRPO 算法对它们进行迭代训练。随着时间的推移,他们观察到这些模型的推理能力有了明显的提升,最终甚至可以匹敌甚至超越一些最强大的专有模型。
现在,他们将分享包括实验、训练方案、数据集 和模型权重 在内的研究成果,所有这些都可以在 MIT 许可下免费获得。
基准测试
为了进行实验,首先需要确定一个具有挑战性的推理任务,该任务需要具备清晰可验证的解决方案和可扩展的复杂性。巧合的是,作者之一 Brad Hilton 之前创建了一个名为 Temporal Clue 的谜题集,完美地满足了这些需求。除了满足基本事实清晰的标准外,还可以根据需要创建新的谜题。
Temporal Clue 的灵感来源于流行的棋盘游戏 Clue (Cluedo)。在这款游戏中,玩家们竞相找出谁在 Tudor Mansion 中杀害了 Boddy 先生。Temporal Clue 将这款游戏变成了一个独立的逻辑谜题,它超越了标准的谁、用什么、在哪里,并加入了两个额外的维度:何时(时间)和为什么(动机)。谜题是随机生成的,最少但足够充分的线索是用 OR-Tools 的 CP-SAT 求解器 选出的。
为了确定当前这项推理任务的最新技术水平,他们对领先的推理模型进行了基准测试,包括 DeepSeek R1、Anthropic 的 Claude Sonnet 3.7,以及阿里巴巴的 Qwen 2.5 14B 和 32B Instruct 模型。此外,他们还提供了最终结果的预览:
机构 | 模型 | 推理能力 | 平均准确率 | 平均成本 |
---|---|---|---|---|
DeepSeek | R1 | 默认 | 51.6% | $0.029 |
Anthropic | Sonnet 3.7 | 无 | 51.7% | $0.017 |
Anthropic | Sonnet 3.7 | 16k | 61.7% | $0.222 |
Anthropic | Sonnet 3.7 | 64k | 69.5% | $0.392 |
阿里巴巴 | Qwen 2.5 14B Instruct | 无 | 28.1% | $0.001 |
阿里巴巴 | Qwen 2.5 32B Instruct | 无 | 37.3% | $0.002 |
从这些基准测试中,可以看到,具有 64k token 的 Claude Sonnet 3.7 在这项任务上表现最好,但所有领先的模型都还有改进空间。DeepSeek R1 是一款流行的开源模型,其性能与 Anthropic 的 Claude Sonnet 3.7 51.7% 的平均准确率 接近。然而,未经调整的 Qwen 2.5 Instruct 模型的性能相对较差。最大的问题是:能否将这些规模较小、开源的模型训练到前沿水平?
训练
为了训练一个前沿的推理模型,他们采用了强化学习——一种允许智能体在受控环境中从自身经验中学习的方法。在这里,LLMs 是智能体,而谜题是环境。他们通过让 LLMs 为每个谜题生成多个响应来引导 LLMs 的学习,从而探索问题的格局。他们对导致正确解决方案的推理进行强化,并对使模型误入歧途的推理进行惩罚。
在各种 RL 方法中,他们选择了由 DeepSeek 开发的流行的 Group Relative Policy Optimization (GRPO) 算法。与 Proximal Policy Optimization (PPO) 等更传统的方法相比,GRPO 简化了训练过程,同时仍然提供了强大的性能。为了加快实验速度,他们省略了 Kullback–Leibler (KL) 散度 惩罚,尽管他们的训练方案也支持它。
总的来说,训练循环遵循以下基本步骤:
- 生成模型对谜题任务的响应。
- 对响应进行评分,并估计每组聊天完成的优势 (这就是 GRPO 中的“Group Relative”部分)。
- 使用由这些优势估计指导的裁剪策略梯度对模型进行微调。
- 使用新的谜题和最新版本的模型重复这些步骤,直到达到最佳性能。
为了生成响应,他们使用了流行的 vLLM 推理引擎。他们调整了参数选择,以最大限度地提高吞吐量并最大限度地减少启动时间。前缀缓存尤为重要,因为他们为每个任务采样了许多响应,而缓存提示有助于避免冗余计算。
他们观察到,过多的请求会使 vLLM 不堪重负,从而导致抢占或换出正在进行的请求。为了解决这个问题,他们使用信号量限制了请求,该信号量经过调整以保持高的键值 (KV) 缓存利用率,同时最大限度地减少交换。更高级的调度机制可以在支持灵活生成长度的同时产生更高的利用率。
采样后,使用标准的 HuggingFace Transformers AutoTokenizer 处理完成。它的聊天模板功能将消息对象呈现为提示字符串,包括一个助手掩码,用于确定 LLM 生成了哪些 token。他们在默认模板中缺少必要的 %generation%
标签,因此在标记化步骤中对其进行了修改。生成的助手掩码包含在用于调整的张量字典中,用于识别哪些位置需要计算损失。
在对响应进行分词并获得助手掩码后,他们将数据打包以进行调优。除了在每个打包序列中包含多个提示/响应对之外,他们还识别了共享的提示 token,并为每个 token 分配了一个父 ID 以及标准的组 ID。特别是对于像 Temporal Clue 这样的任务(每个谜题平均超过 1000 个 token),为每个任务生成大量响应并有效地打包张量可以显著减少冗余。一旦打包了所有必要的信息,就可以二维可视化训练数据集,每一行都是一个 token 序列,可能包含多个提示和完成:
有了紧密打包的数据,就可以进行调优了。模型已经过预训练、指令调优,相当智能,并且擅长遵循指令。然而,它们还不能可靠地解决 Temporal Clue 谜题。不过,它们偶尔会成功,这就足够了。通过增加良好推理的概率并降低不良推理的概率,逐渐将模型引导至“神探”状态。这是使用标准的机器学习技术实现的,采用策略梯度方法来计算损失并有利地转移权重。
对于训练,使用了 PyTorch 团队提供的 torchtune 库。Torchtune 具有高效的仅解码器 transformer 实现,适用于流行的模型,包括 Llama、Gemma、Phi 等。虽然主要使用 Qwen 模型进行此项目,但他们也使用 Meta 的 Llama 8B 和 70B 模型进行了实验。Torchtune 还提供了节省内存和提高性能的实用程序,包括:
- 激活检查点
- 激活卸载
- 量化
- 参数高效微调 (PEFT),例如 低秩自适应 (LoRA)
有关支持的优化的完整列表,请参阅此处的自述文件。
此外,Torchtune 支持多设备(以及现在的多节点)训练,使其成为大型模型的理想选择。它同时支持完全分片数据并行 (FSDP) 和张量并行 (TP) 训练,可以结合使用。他们还提供了十几个方案,鼓励用户复制并针对自己的用例进行定制。他们创建了一个修改版的完整微调方案,支持:
- 多设备和单设备训练
- 用于计算 KL 散度的参考模型加载和权重交换
- 使用组 ID 和父 ID 进行高级因果掩码计算
- GRPO 损失集成和组件日志记录
该方案可以在这里看到。未来,他们希望添加张量并行支持,并探索 PEFT 和量化。
RL 训练过程涉及选择大量的超参数。在训练模型时,他们测试了各种配置,并主要确定了以下配置:
- 模型:Qwen 2.5 Instruct 14B 和 32B
- 每次迭代的任务数:32
- 每次迭代每个任务的样本数:50
- 每次迭代的总样本数:32 * 50 = 1600
- 学习率:6e-6
- 微批量大小:14B 模型为 4 个序列,32B 模型为 8 个
- 批量大小:可变,取决于序列的数量
批量大小是可变的,因为在训练期间响应长度可能会发生变化,序列打包效率在每次迭代中都会波动,并且会丢弃零优势的响应(即模型没有给出任何积极或消极反馈的响应)。在一次运行中,尝试了根据批量大小动态调整学习率,但这导致小批量大小的学习率过高,需要设置上限。设置上限的版本与使用恒定学习率没有明显差异,但调整批量大小和学习率仍然是未来实验的一个有趣领域。
他们还进行了简短的实验,增加每次迭代的任务数,同时减少每个任务的样本数,反之亦然,保持每次迭代的总样本数大致相等。在较短的训练时间内,这些变化没有显示出有意义的差异,这表明该方案对于任务数量和每个任务的样本数量之间的不同平衡是稳健的。
结果
在对模型进行 100 多次迭代训练后,达到了前沿的推理水平。
模型迅速改进,然后准确率的提高开始逐渐减少并最终下降,有时甚至急剧下降。在最佳状态下,14B 模型接近 Claude Sonnet 3.7 在 16k token 时的性能,32B 模型几乎与 Sonnet 在 64k token 时的结果相匹配。
在训练过程中,性能提升遵循幂律,在对数-对数图上形成线性关系(在恶化之前)。
他们怀疑,这些模型可能过早地收敛于一开始就有效的贪婪策略,但可能限制了它们的长期前景。下一步,可以探索鼓励多样化响应的方法,或逐步构建能力的方法(如课程学习),或为特别出色的解决方案分配更大的奖励,以激励彻底的探索。
此外,他们还注意到了训练期间输出长度的有趣模式。最初,响应变得更长,然后稳定下来,然后在训练接近结束时出现分歧,14B 模型的响应变得更长,而 32B 模型的响应长度缩短,尤其是在达到峰值性能之后。
为了定性评估逻辑推理的改进,他们让最强的前沿模型 Claude Sonnet 3.7 识别和评估 Qwen 32B 模型在类似谜题上进行的推理的合理性——在训练 100 多次迭代之前和之后。Sonnet 从基础模型中识别出 6 个推论,其中除 1 个外均被判定为错误;相反,它从训练后的模型中识别出 7 个推论,其中除 1 个外均被判定为逻辑上合理。
最后,假设按需部署具有足够的吞吐量,他们根据 Fireworks AI 的无服务器定价层级估算了 Qwen 模型的成本。他们将准确率与每次响应的平均推理成本的自然对数进行了对比,并观察到在未调整的模型中存在清晰的线性帕累托前沿。通过成功地将开源模型训练到前沿的准确率水平,显著改善了成本与准确率之间的权衡。
总结
在本次调查中,他们试图探索规模较小、开源的语言模型是否可以通过强化学习实现前沿的推理能力。在使用精心选择的超参数和 GRPO 方法对 Qwen 14B 和 32B 模型进行具有挑战性的 Temporal Clue 谜题训练后,他们取得了令人印象深刻的性能提升。这些改进将开源模型带到了推理性能的前沿,同时显著降低了成本。研究结果突出了强化学习在复杂推理任务上高效训练开放模型的潜力。
如前所述,数据集、实验、训练方案 和模型权重 (14B, 32B) 均可在 MIT 许可下免费获得。
此外,他们还发现,只需 16 个训练示例 即可实现高达 10-15% 的有意义的性能提升。