Skip to content

DeepScaleR:通过扩展强化学习,用 15 亿参数模型超越 O1-Preview

作者:Michael Luo、Sijun Tan、Justin Wong、Xiaoxiang Shi、William Tang、Manan Roongta、Colin Cai、Jeffrey Luo 导师:Tianjun Zhang、Li Erran Li、Raluca Ada Popa、Ion Stoica

摘要

强化学习(RL)的魅力无处不在!我们推出了 DeepScaleR-1.5B-Preview,这是一个基于 Deepseek-R1-Distilled-Qwen-1.5B 通过简单强化学习微调得到的语言模型。它在 AIME2024 上实现了 43.1% 的 Pass@1 准确率(比基础模型提升了 14.3%),仅用 15 亿参数就超越了 OpenAI 的 O1-Preview。我们开源了我们的数据集、代码和训练日志,希望与大家共同探索如何通过强化学习扩展智能。

DeepScaleR-1.5B-Preview 性能对比

模型AIME 2024MATH 500AMC 2023Minerva MathOlympiadBench平均值
DeepScaleR-1.5B-Preview43.1%87.8%73.6%30.2%50.0%57.0%
DeepSeek-R1-Distill-Qwen-1.5B28.8%82.8%62.9%26.5%43.3%48.9%
O1-Preview40.0%81.4%----
img

图 1:随着训练的进行,DeepScaleR 在 AIME2024 上的 Pass@1 准确率。在第 1040 步和 1520 步时,上下文长度分别扩展到 16K 和 24K。

引言:推动强化学习在大语言模型中的普及

最近开源的 Deepseek-R1(一个与 OpenAI 的 O1 相当的模型)标志着推理模型民主化迈出了重要一步。然而,其确切的训练配方、超参数和底层系统仍然不可用。在这项工作中,我们朝着完全开放的强化学习扩展配方迈出了重要一步。

扩展强化学习的最大挑战之一是高昂的计算成本。例如,我们发现直接复制 DeepSeek-R1 的实验(≥32K 上下文,约 8000 步)至少需要 70,000 个 A100 GPU 小时,即使是对于一个 15 亿参数的模型也是如此。为了解决这一问题,我们利用了蒸馏模型,并引入了一种新颖的迭代扩展方案用于强化学习,将计算需求降低到仅 3800 个 A100 GPU 小时——减少了 18.42 倍——同时仅用一个 15 亿参数模型就实现了超越 OpenAI 的 O1-Preview 的性能。

我们的工作表明,通过强化学习开发定制推理模型可以既可扩展又高效。在接下来的博文中,我们将介绍我们的数据集策划和训练方法,展示评估结果,并分享我们发现的关键见解。

DeepScaleR 的配方

数据集策划

对于我们的训练数据集,我们收集了 1984-2023 年的 AIME 问题和 2023 年之前的 AMC 问题,以及来自 Omni-MATH 和 Still 数据集的问题,这些数据集包含来自各种国家和国际数学竞赛的问题。

我们的数据处理流程包括三个关键步骤:

  1. 提取答案:对于 AMC 和 AIME 等数据集,我们使用 gemini-1.5-pro-002 从官方 AoPS 解决方案中提取答案。
  2. 去除重复问题:我们使用 RAG 和 sentence-transformers/all-MiniLM-L6-v2 的嵌入来消除重复问题。为了防止数据污染,我们还检查了训练集和测试集之间的重叠部分。
  3. 过滤不可评分问题:某些数据集(如 Omni-MATH)包含无法使用 sympy 评估的问题,需要使用 LLM 作为评判。由于使用 LLM 评判可能会减慢训练速度并引入嘈杂的奖励信号,我们增加了一个额外的过滤步骤来移除这些不可评分的问题。

经过去重和过滤后,我们的最终训练数据集包含大约 40,000 个独特的问答对。我们将在未来的运行中扩展我们的数据集。

奖励函数

正如 Deepseek-R1 所倡导的,我们采用结果奖励模型(ORM),而不是过程奖励模型(PRM),以避免奖励欺骗。简而言之,我们的奖励函数返回:

  • 如果 LLM 的答案通过基本的 LaTeX/Sympy 检查,则返回 1。
  • 如果 LLM 的答案错误或格式不正确(例如缺少 <think></think> Token),则返回 0。

迭代上下文扩展:先短后长

img

图 2:随着训练的进行,DeepScaleR 的平均响应长度和训练奖励的变化。曲线显示了以 100 为窗口大小的移动平均值。

在扩展强化学习用于推理任务时,选择最优的上下文窗口用于训练是一个关键挑战。推理工作负载高度计算密集,因为它们生成的输出比标准任务长得多,这会减慢轨迹采样和策略梯度更新的速度。上下文窗口大小翻倍会使训练计算量至少增加 2 倍。

这引入了一个基本的权衡:更长的上下文为模型提供了更多的思考空间,但会显著减慢训练速度;而较短的上下文可以加速训练,但可能会限制模型解决更复杂问题的能力,这些问题需要更长的上下文。因此,平衡效率和准确性至关重要。

总结来说,我们的训练配方采用 Deepseek 的 GRPO 算法,分为两个步骤:

  1. 首先,我们使用 8K 最大上下文进行强化学习训练,以实现更有效的推理和高效的训练。
  2. 接下来,我们将训练扩展到 16K 和 24K 上下文,以便模型能够解决更具挑战性、之前未能解决的问题。

用 8K 上下文启动有效的 CoT

在启动完整的训练运行之前,我们在 AIME2024 上评估了 Deepseek-R1-Distilled-Qwen-1.5B,并分析了轨迹统计信息。平均而言,错误回答包含的Token数量是正确回答的三倍(20,346 对 6,395)。这表明较长的回答往往会导致错误的结果。因此,立即使用长上下文窗口进行训练可能是低效的,因为大多数Token实际上是浪费的。

此外,我们在评估日志中观察到,较长的回答表现出重复的模式,表明它们并没有对有效的链式推理(CoT)做出有意义的贡献。

项目基础模型DeepScaleR-1.5B-8K变化
AIME Pass@128.9%33.9%+5%
正确回答的平均Token数6396.03661.2-2734.8
错误回答的平均Token数20346.36976.8-13369.5
总平均Token数16335.65850.9-10484.7

鉴于这一见解,我们以 8K 上下文启动训练,初始 AIME2024 准确率为 22.9%——仅比原始模型低 6%。这一策略被证明是有效的:在训练过程中,平均训练奖励从 46% 增加到 58%,而平均回答长度从 5500 个Token减少到 3500 个Token(见图 2)。

更重要的是,将输出限制在 8K Token内使模型更有效地利用上下文。如表所示,我们的模型为正确和错误回答生成的Token数显著减少,同时比基础模型的 AIME 准确率高出 5%——并且仅使用了三分之一的Token。

在转折点扩展到 16K 上下文

大约在 1000 步之后,我们的 8K 运行出现了一个有趣的变化:回答长度再次开始增加。然而,这导致了收益递减——准确率趋于平稳,最终下降。与此同时,回答截断比例从 4.2% 上升到 6.5%,表明越来越多的回答在上下文限制处被截断。

img

图 3:在 1000 步之后,回答长度再次上升,但我们的 8K 运行的训练奖励最终下降。

img

图 4:在 1000 步之后,8K 上下文运行的回答截断比例上升。

这些结果表明,模型试图通过“思考更长时间”来提高训练奖励。然而,随着它生成更长的回答,它越来越频繁地遇到 8K 上下文窗口的上限,从而限制了进一步的改进。

认识到这是一个自然的过渡点,我们决定“打开笼子,让鸟儿自由飞翔”。我们取第 1040 步的Checkpoint——回答长度开始上升的地方——并重新启动训练,将上下文窗口扩大到 16K。与从一开始就使用 16K 训练相比,这种两阶段方法要高效得多:8K 的引导训练将平均回答长度保持在 3000 个Token,而不是 9000 个Token,使这一阶段的训练至少快了 2 倍。

在转换之后,我们观察到训练奖励、回答长度和 AIME 准确率的稳步提升。在额外的 500 步之后,平均回答长度从 3500 增加到 5500 个Token,AIME2024 Pass@1 准确率达到 38%。

用 24K 魔法超越 O1-Preview

在 16K 上下文上额外训练 500 步后,我们注意到性能开始趋于平稳——平均训练奖励收敛于 62.5%,AIME Pass@1 准确率徘徊在 38% 左右,回答长度再次开始下降。与此同时,最大回答截断比例上升到 2%。

为了最终实现 O1 级别的性能,我们决定施展“24K 魔法”——将上下文窗口扩大到 24K。我们取第 480 步的Checkpoint,并重新启动训练运行,将上下文窗口扩大到 24K。

在扩展上下文窗口后,模型终于突破了限制。大约在 50 步之后,我们的模型最终突破了 40% 的 AIME 准确率,并在第 200 步达到 43%。24K 魔法发挥了巨大作用!

总体而言,我们的训练运行包括大约 1750 步。初始的 8K 阶段在 8 个 A100 GPU 上进行训练,而 16K 和 24K 阶段则扩展到 32 个 A100 GPU。总共,训练大约花费了 3800 个 A100 GPU 小时,相当于在 32 个 A100 上大约 5 天的时间,计算成本约为 4500 美元。

评估

我们在竞赛级别的数学基准测试中评估了我们的模型,包括 AIME 2024、AMC 2023、MATH-500、Minerva Math 和 OlympiadBench。下面报告了 Pass@1 准确率,每个问题平均了 16 个样本的结果。我们运行的基线用于验证分数已用下划线标出。

ModelAIME 2024MATH 500AMC 2023Minerva MathOlympiadBenchAvg.
Qwen-2.5-Math-7B-Instruct13.379.850.634.640.743.8
rStar-Math-7B26.778.447.5-47.1-
Eurus-2-7B-PRIME26.779.257.838.642.148.9
Qwen2.5-7B-SimpleRL26.782.462.539.743.350.9
DeepSeek-R1-Distill-Qwen-1.5B28.882.862.926.543.348.9
Still-1.5B32.584.466.729.045.451.6
DeepScaleR-1.5B-Preview43.187.873.630.250.057.0
O1-Preview40.081.4----

我们将 DeepScaleR 与我们使用的 DeepSeek 基础模型以及最近探索强化学习用于推理任务的学术作品进行了比较。DeepScaleR 在所有基准测试中显著优于基础模型,AIME2024 上实现了 14.4% 的绝对提升,总体提升了 8.1%。此外,DeepScaleR 还超越了最近的学术作品,如 rSTAR、Prime 和 SimpleRL,这些作品都是从 7B 模型微调而来的。如图 5 所示,DeepScaleR 仅用 15 亿参数就实现了与 O1-Preview 相当的性能——这是一个显著的效率提升。

关键要点

  • 强化学习扩展也可以在小模型中实现:Deepseek-R1 表明,直接在小模型上应用强化学习不如蒸馏有效。他们的消融实验表明,强化学习在 Qwen-32B 上实现了 47% 的 AIME 准确率,而仅蒸馏就达到了 72.6%。一个常见的误解是强化学习扩展只对大模型有益。然而,有了从大模型蒸馏出的高质量 SFT 数据,小模型也可以通过强化学习更有效地学习推理。我们的结果证实了这一点:强化学习扩展将 AIME 准确率从 28.9% 提高到 43.1%!这些发现表明,单独的 SFT 或 RL 都是不够的。相反,通过结合高质量的 SFT 蒸馏和强化学习扩展,我们才能真正释放 LLM 的推理潜力。
  • 迭代扩展实现更有效的长度扩展:先前的研究 [1, 2] 表明,直接在 16K 上下文上进行强化学习训练与 8K 相比没有显著提升,可能是因为模型没有足够的计算能力来充分利用扩展的上下文。最近的一项工作 [3] 表明,更长的回答长度包含冗余的自我反思,导致错误的结果。我们的实验结果与这些发现一致。通过首先在较短的上下文(8K)中优化推理,我们使后续 16K 和 24K 运行的训练更快、更有效。这种迭代方法使模型在扩展到更长上下文之前,先形成有效的思考模式,从而使基于 RL 的长度扩展更高效。

结论

我们的工作旨在揭示强化学习对 LLM 的扩展效应,并使其可供每个人使用。DeepScaleR-1.5B-Preview 是我们在这方面努力的第一个模型,AIME Pass@1 准确率达到 43.1%。我们坚信,普及强化学习扩展是一个社区努力的过程,欢迎开源贡献和赞助!让我们共同努力,拓展强化学习在 LLM 推理中的边界!

Maintained by Robin