rStar-Math：通过自我进化的深度思维 小型 LLM 掌握数学推理

25年1月来自微软亚洲研究院的论文“rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking”。

rStar-Math 证明小语言模型 (SLM) 可以与 OpenAI o1 的数学推理能力相媲美甚至超越它，而无需从优秀模型中进行蒸馏。 rStar-Math 通过蒙特卡洛树搜索 (MCTS) 进行“深度思考”来实现这一点，其中数学策略 SLM 在基于 SLM 的过程奖励模型指导下执行测试-时间搜索。 rStar-Math 引入三项创新来应对训练两个 SLM 的挑战：（1）一种新代码增强 CoT 数据合成方法，它执行广泛的 MCTS 部署以生成用于训练策略 SLM 的逐步验证推理轨迹； （2）一种新过程奖励模型训练方法，避免简单的步骤级打分注释，从而产生了更有效的过程偏好模型（PPM）； （3）一种自我进化的方法，其中策略 SLM 和 PPM 从头开始​​构建，并不断迭代进化以提高推理能力。通过 4 轮自我进化，rStar-Math 为 747,000 个数学问题合成数百万个解决方案，将 SLM 的数学推理能力提升到最先进的水平。在 MATH 基准测试中，它将 Qwen2.5-Math-7B 从 58.8% 提高到 90.0%，将 Phi3-mini-3.8B 从 41.4% 提高到 86.4%，分别超过 o1-preview +4.5% 和 +0.9%。在美国数学奥林匹克竞赛 (AIME) 中，rStar-Math 平均解决 53.3% (8/15) 的问题，位列最聪明高中数学学生的前 20%。

最近的研究表明，大语言模型 (LLM) 能够解决数学问题 [Team, 2024a, Yang et al., 2024, OpenAI, 2024, Liu et al., 2024]。然而，传统的方法是让 LLM 通过一次推理生成完整的解决方案——类似于系统 1 思维 [Daniel, 2011]——通常会产生快速但容易出错的结果 [Valmeekam et al., 2023, OpenAI, 2023]。作为回应，测试-时计算规模化 [Snell et al., 2024, Qi et al., 2024] 建议向系统 2 式思维转变，通过更慢、更深入的思维过程模拟人类的推理。在这种范式中，LLM 作为策略模型来生成多个数学推理步骤，然后由另一个作为奖励模型的 LLM 进行评估 [OpenAI，2024]。选择被认为更有可能正确的步骤和解决方案。该过程不断重复，最终得出最终答案。

在测试-时间计算范式中，关键是训练一个强大的策略模型来生成有希望的解决步骤，以及一个可靠的奖励模型来准确评估它们，这两者都依赖于高质量的训练数据。不幸的是，众所周知，现成的高质量数学推理数据很少，而合成高质量的数学数据面临着根本性的挑战。对于策略模型来说，区分错误的推理步骤和正确的推理步骤是一项挑战，这使得消除低质量数据变得复杂。值得注意的是，在数学推理中，正确的最终答案并不能确保整个推​​理轨迹的正确性 [Lanham et al., 2023]。错误的中间步骤会严重降低数据质量。至于奖励模型，过程奖励模型 (PRM) 通过对中间步骤提供细粒度的反馈显示出巨大的潜力 [Lightman et al., 2023]。然而，这方面的训练数据更加稀缺：准确的逐步反馈需要大量的人工标记工作，并且难以规模化，而那些自动注释尝试由于嘈杂奖励分数而显示出有限的收益 [Luo et al.，2024 ，Wang，2024c，Chen，2024]。由于上述挑战，现有的基于蒸馏数据合成方法来训练策略模型，例如扩大 GPT-4 蒸馏的 CoT 数据 [Tang et al., 2024, Huang et al., 2024]，已经显示出收益递减，并且无法超过他们教师模型的能力；与此同时，截至今天，训练可靠的 PRM 进行数学推理仍然是一个悬而未决的问题。

数学数据合成。 LLM 数学推理的进步很大程度上依赖于整理高质量的 CoT 数据，其中大多数领先方法都是经过 GPT 提炼的，使用 GPT-4 等前沿模型进行综合 [Wang et al., 2024b, Gou et al., 2023, Luo，2023]。值得注意的方法包括 NuminaMath [Jia LI & Polu，2024a] 和 MetaMath [Yu et al.，2023b]。虽然有效，但这限制了LLM教师的推理能力。 LLM 老师无法解决的难题被排除在训练集之外。即使是可解的问题也可能包含容易出错的中间步骤，这些步骤很难被发现。虽然拒绝采样方法 [Yuan et al., 2023, Brown et al., 2024] 可以提高数据质量，但它们并不能保证中间步骤的正确性。因此，扩大 CoT 数据的收益会递减，收益接近饱和——例如，尽管数据集大小增加 8 倍，但 OpenMathInstruct-2 [Toshniwal et al., 2024] 在 MATH 上的提升仅为 3.9%。

测试-时间计算规模化。引入了新的规模化定律，允许 LLM 通过生成多个样本并使用奖励模型来选择最佳解决方案来提高性能 [Snell et al., 2024, Wu et al., 2024, Brown et al., 2024]。已经提出各种测试-时搜索方法[Kang et al., 2024, Wang et al., 2024a]，包括随机抽样[Wang et al., 2023]和树搜索方法[Yao et al., 2024, Hao et al., 2023, Zhang et al., 2024b, Qi et al., 2024]类似于 MCTS。然而，用于测试时间计算规模化的开源方法在数学推理方面表现出有限的收益，这通常是由于策略 LLM 或奖励模型的限制。

奖励模型。这个对于有效的系统 2 推理至关重要，但很难获得。最近的研究包括用于验证的 LLM-as-a-Judge [Zheng et al., 2023, Qi et al., 2024] 和专门的奖励模型，如结果奖励模型 [Yang et al., 2024, Yu et al., 2023a]以及过程奖励模型（PRM）[Lightman ，2024]。虽然 PRM 为复杂推理提供有希望的、密集的、步骤级奖励信号 [Luo et al., 2024, Wang et al., 2024c]，但收集步骤级注释仍然是一个障碍。而 Kang [2024]，Wang [2024a] 依赖于昂贵的人工注释数据集，例如 PRM800k [Lightman et al., 2024]，最近的方法 [Wang et al., 2024c, Luo et al., 2024] 探索通过蒙特卡洛采样或 MCTS 进行自动注释。然而，他们很难获得精确的奖励分数，这限制性能的提升。

这项工作引入 rStar-Math，一种可自我进化的系统 2 式推理方法，它实现了最先进的数学推理，在具有挑战性的数学竞赛基准上与 OpenAI o1 相媲美，有时甚至超越它，模型大小为少至70亿。与依赖卓越 LLM 进行数据合成的解决方案不同，rStar-Math 利用较小的语言模型 (SLM) 和蒙特卡洛树搜索 (MCTS) 建立自我进化过程，迭代生成更高质量的训练数据。

设计选择

MCTS 用于有效系统 2 推理。目标是训练一个数学策略 SLM 和一个过程奖励模型 (PRM)，并将两者整合到蒙特卡洛树搜索 (MCTS) 中，以进行系统 2 深度思考。选择 MCTS 有两个主要原因。首先，它将复杂的数学问题分解为更简单的单步生成任务，与其他 System 2 方法（如 Best-of-N [Brown et al., 2024] 或自洽性 [Wang et al., 2023]，这需要在一次推理中生成完整的解决方案。其次，MCTS 中的逐步生成，自然会为两个模型产生步级训练数据。标准 MCTS 展开会根据每个步骤对最终正确答案的贡献自动为其分配 Q 值，从而无需在过程奖励模型训练中人工生成步骤级的注释。

理想情况下，GPT-4 等高级 LLM 可以集成到 MCTS 中以生成训练数据。然而，这种方法面临两个关键挑战。首先，即使是这些强大的模型也难以持续解决诸如奥林匹克级别的数学等难题。因此，得到的训练数据主要由较简单的可解问题组成，从而限制其多样性和质量。其次，注释每步 Q 值需要大量的 MCTS 推广；树探索不足可能导致错误的 Q 值分配，例如高估次优步骤。鉴于每次展开都涉及多个单步生成，并且这些模型的计算成本很高，因此增加展开会显著增加推理成本。

如图所示：为了减少错误和低质量的中间步骤，引入一种代码增强的 CoT 合成方法，该方法执行广泛的 MCTS 推出以生成逐步验证的推理轨迹，并用 Q 值注释。为了进一步提高 SLM 在解决挑战性问题上的性能，引入四轮自进化方案。在每一轮中，策略 SLM 和奖励模型都会更新为更强大的版本，逐步解决更困难的问题并生成更高质量的训练数据。最后，提出一种新过程奖励模型训练方法，该方法不再需要精确的每个步骤奖励注释，从而产生更有效的过程偏好模型（PPM）。

逐步验证推理轨迹

代码增强的 CoT 生成。先前的 MCTS 方法主要生成自然语言 (NL) CoT [Qi et al., 2024, Zhang et al., 2024a]。然而，LLM 经常会出现幻觉，产生不正确或不相关的步骤，但仍然偶然得出正确答案 [Lanham et al., 2023]。这些有缺陷的步骤很难检测和消除。为了解决这个问题，提出一种代码执行增强 CoT。如图所示，策略模型生成一步 NL CoT 及其对应的 Python 代码，其中 NL CoT 作为 Python 注释嵌入。只有成功执行 Python 代码的生成才会被保留为有效候选。

广泛展开的 Q 值注释。准确的 Q 值 Q(s) 注释，对于引导 MCTS 节点选择正确的问题解决路径以及识别轨迹中的高质量步骤至关重要。为了提高 Q 值的可靠性，从围棋选手那里获得灵感，他们根据游戏结果回顾性地评估每一步的奖励。尽管最初的估计可能不准确，但随着时间的推移，反复的游戏会完善这些评估。类似地，在每次展开时，都会根据每个步骤对获得正确最终答案的贡献来更新其 Q 值。经过广泛的 MCTS 展开之后，持续导致正确答案的步骤会获得更高的 Q 值，偶尔成功的步骤会产生中等的 Q 值，而持续错误的步骤会获得较低的 Q 值。具体来说，引入两种自注释方法来获得这些步骤级 Q 值。

过程偏好模型 PPM

过程奖励模型提供细粒度的步骤级奖励信号，对于解决具有挑战性的数学问题非常有价值。然而，获取高质量的阶梯级训练数据仍然是一个悬而未决的挑战。现有方法依赖于人工注释 [Lightman et al., 2023] 或 MCTS 生成的分数 [Zhang et al., 2024a, Chen et al., 2024] 为每个步骤分配分数。然后将这些分数作为训练目标，使用 MSE 损失 [Chen et al., 2024] 或逐点损失 [Wang et al., 2024c, Luo et al., 2024, Zhang et al., 2024a] 等方法来最小化预测分数和标记分数之间的差异。因此，这些注释的步骤级奖励分数的精确度，直接决定最终的过程奖励模型有效性。

不幸的是，每步精确评分仍然是一个尚未解决的挑战。尽管广泛的 MCTS 展开提高 Q 值的可靠性，但精确评估细粒度步骤质量仍然是一个重大障碍。例如，在一组正确的步骤中，很难将它们排序为最佳、次好或平均，然后分配精确的分数。类似地，在错误的步骤中，区分最差的步骤和中等差的步骤也面临着类似的挑战。即使是专业的人工注释也会在一致性方面遇到困难，特别是在规模上，这导致训练标签中存在固有的噪声。

引入一种新训练方法，通过构建步骤级正负偏好对来训练过程偏好模型（PPM）。不使用 Q 值作为直接奖励标签，而是使用它们从 MCTS 树中选择步骤来构建偏好对。对于每个步骤，选择两个 Q 值最高的候选者作为正步骤，选择两个 Q 值最低的候选者作为负步骤。至关重要的是，所选择的正步骤必须导致正确的最终答案，而负步骤必须导致错误的答案。对于中间步骤（除了最后的答案步骤），正对和负对共享相同的前面步骤。对于最终的回答步骤，相同的推理轨迹很少会产生不同的最终答案，因此放宽这个限制。选择平均 Q 值最高的两条正确轨迹作为正例，选择平均 Q 值最低的两条错误轨迹作为负例。按照 [Ouyang et al.，2022] 的说法，用一个成对排名损失的标准 Bradley-Terry 模型来定义损失函数。

自我进化深度思考

数学问题集锦。收集一个包含 747k 个数学应用题的大型数据集，并附有最终答案的真实标签，主要来自 NuminaMath [Jia LI and Polu，2024a] 和 MetaMath [Yu et al.，2023b]。值得注意的是，仅包括 NuminaMath 的竞赛级别问题（例如奥林匹克竞赛和 AIME/AMC），因为小学级别的问题并不能显著提高 LLM 复杂数学推理能力。为了扩充有限的竞赛级问题，遵循 [Li et al., 2024] 并使用 GPT-4 基于 7.5k MATH 训练集和 3.6k AMC-AIME 训练分割中的种子问题合成新问题。然而，对于具有挑战性的种子问题，GPT-4 经常会产生无法解决的问题或错误的解决方案。为了过滤这些问题，提示 GPT-4 每个问题生成 10 个解决方案，只保留至少有 3 个一致解决方案的解决方案。

推理轨迹收集。没有使用 747k 数学数据集中的原始解决方案，而是进行广泛的 MCTS 展开以生成更高质量的逐步验证推理轨迹。在每个自我进化轮中，对每个数学问题进行 16 次展开，从而产生 16 条推理轨迹。然后根据生成轨迹的正确比例，按难度对问题进行分类：简单（所有解决方案都是正确的）、中等（正确和不正确的解决方案混合）和困难（所有解决方案都是不正确的）。对于没有正确轨迹的难题，将执行额外具有 16 个展开的 MCTS。之后，收集并过滤所有逐步轨迹及其带注释的 Q 值，以训练策略 SLM 和过程偏好模型。

监督微调策略SLM。通过大量实验，发现选择高质量的推理轨迹是微调前沿数学LLM的关键。虽然 GPT-蒸馏和 Best-of-N 等方法可能包含低质量或错误的中间步骤，但更有效的方法可确保轨迹中的每一步都是高质量的。为了实现这一点，用每步 Q 值从 MCTS 展开中选择最佳轨迹。具体来说，对于每个数学问题，在导致正确答案的轨迹中选择平均 Q 值最高的前 2 条轨迹作为 SFT 训练数据。

训练 PPM。 PPM 由微调策略模型初始化，其下一个token预测头由由线性层和 tanh 函数组成的标量-值头取代，以将输出限制在范围 [-1, 1] 内。过滤掉所有解决轨迹完全正确或不正确的数学问题。对于结果混合的问题，根据 Q 值为每个步骤选择两个正例和两个反例，用作训练数据的偏好对。

由于SLM的能力较弱，进行四轮MCTS深度思考，以逐步生成更高质量的数据，并使用更具挑战性的数学问题来扩展训练集。

每一轮使用MCTS生成逐步验证的推理轨迹，然后用于训练新策略SLM和PPM。然后在下一轮中应用新模型来生成更高质量的训练数据。

第一轮：自举初始强策略 SLM-r1。为了使 SLM 能够自行生成合理良好的训练数据，执行一轮自举轮来微调初始强策略模型，表示为 SLM-r1。使用 DeepSeek-Coder-V2-Instruct (236B) 运行 MCTS 来收集 SFT 数据。由于此轮没有可用的奖励模型，对 Q 值使用终端引导注释，并将 MCTS 限制为 8 次展开以提高效率。对于正确的解决方案，选择平均 Q 值最高的前 2 条轨迹作为 SFT 数据。也训练 PPM-r1，但是有限的展开产生不可靠的 Q 值，影响了 PPM-r1 的有效性。

第二轮：训练可靠的 PPM-r2。在这一轮中，随着策略模型更新至 7B SLM-r1，进行广泛的 MCTS 展开，以获得更可靠的 Q 值注释，并训练第一个可靠的奖励模型 PPM-r2。具体来说，对每个问题执行 16 次 MCTS 展开。由此产生的逐步验证推理轨迹，显示质量和 Q 值精度都有显著的提高。PPM-r2 明显比自举轮更有效。此外，策略 SLM-r2 也如预期继续改善。

第三轮：PPM 增强 MCTS，显著提高数据质量。借助可靠的 PPM-r2，在本轮中执行 PPM 增强 MCTS 来生成数据，从而获得明显更高质量的轨迹，涵盖训练集中更多的数学和奥林匹克级别的问题。然后使用生成的推理轨迹和自注释 Q 值来训练新策略 SLM-r3 和 PPM-r3，两者都显示出显着的改进。

第四轮：解决具有挑战性的数学问题。第三轮之后，虽然小学和数学题目的成功率很高，但只有 62.16% 的奥林匹克级别题目被纳入训练集。这不仅仅是因为 SLM 推理能力弱，因为许多奥林匹克问题仍然未被 GPT-4 或 o1 解决。为了提高覆盖率，采用一种直接的策略。对于 16 次 MCTS 展开后仍未解决的问题，会额外进行 64 次展开，如果需要，则增加到 128 次。还使用不同的随机种子进行多次 MCTS 树展开。这使得奥林匹克级别题目的成功率提高到80.58%。

经过四轮自我进化，747k道数学题中的90.25%成功被覆盖到训练集中。在剩余的未解问题中，相当一部分是合成问题。手动审查 20 个问题的随机样本，发现其中 19 个问题的答案被错误地标记。基于此得出结论，剩余的未解决的问题质量较低，因此在第四轮终止自我进化。

内在自我反思能力的出现。 OpenAI o1 的一个关键突破是其内在的自我反思能力。当模型犯错时，它会识别错误并能自我纠正并给出正确的答案 [Noam Brown & Lightman，2024]。然而，人们一直发现它在开源 LLM 中基本上是无效的。社区积极探索各种方法，包括自我纠正 [Huang et al., 2023, Kumar et al., 2024]、自我反思 [Renze and Guven, 2024, Shinn et al., 2024]，以明确训练或促使 LLM 培养这种能力。

在实验中意外发现，由 MCTS 驱动的深度思考在解决问题的过程中表现出自我反思。如图所示，该模型在前三个步骤中首先使用 SymPy 将方程形式化，这会导致错误的答案（左分支）。有趣的是，在第四步（右分支）中，策略模型认识到其早期步骤的质量较低，并停止继续沿着最初的问题解决路径前进。相反，它回溯并使用一种新的、更简单的方法解决问题，最终得出正确答案。值得注意的是，其中没有包含自我反思训练数据或提示，这表明高级系统 2 推理可以促进内在的自我反思。

PPM 塑造系统 2 深度思考中的推理边界。策略和奖励模型对于系统 2 深度推理都至关重要。实验表明，一旦策略模型达到相当强的能力水平，PPM 就会成为决定性能上限的关键因素。如图总结不同规模策略模型的准确性，以及通过奖励模型取得的改进。尽管由于训练策略、数据集和模型规模的差异，Pass@1 的准确率有所不同，但奖励模型被证明是系统 2 推理中的主导因素。例如，虽然 rStar-Math-7B 的 SFT 准确度低于 Qwen2.5-Math-72B-Instruct，但将其与 7B PPM 配对可以使 rStar-Math 的表现优于使用 Qwen 72B ORM 的 72B 策略模型。此外，尽管在三个策略 SLM 大小中 Pass@1 准确度有所不同，但在应用 PPM 后，最终的推理准确度会收敛。

PPM 识别定理-应用步骤。在解决具有挑战性的数学问题时，识别和应用相关定理或关键结论通常是成功解决问题的基石 [Xin et al.，2024]。在实验中，在 rStar-Math 问题解决过程中，PPM 有效地识别策略模型深度思考过程中的关键定理应用中间步骤。这些步骤以高奖励分数预测，指导策略模型生成正确的解决方案。PPM 成功识别关键定理的示例，例如 Fermat 小定理 [Weisstein, a]、Vieta 公式 [Weisstein, b]、AM-GM 不等式 [amg]、勾股定理 [pyt] 和鞋带定理 （shoeless theorem）[sho]。

泛化讨论。 rStar-Math 提供一种适用于各个领域改进 LLM 推理能力的通用方法。首先，rStar-Math 可以泛化到更具挑战性的数学任务，例如定理证明，尽管由于数据集限制，它目前的重点是文字问题。尽管如此，rStar-Math 展示证明数学陈述的潜力。例如，它成功证明一道涉及费马小定理的奥林匹克级别问题，并通过其深度推理过程提供逐步正确的证明。其次，rStar-Mathcan 可以泛化到其他领域，例如代码和常识推理。值得注意的是，为一般推理合成经过逐步验证的训练轨迹需要一种机制来提供反馈，以确定给定轨迹在 MCTS 推出结束时是否达到所需的输出。例如，在代码推理中，这可能涉及设计大量测试用例；在一般推理中，可以通过人工标记或与另一个 LLM 相互验证来获得反馈 [Qi et al.，2024]。

训练细节。在每一轮中，收集经过验证的逐步轨迹，以微调策略 LLM 并训练 PPM。为了减少合成数学问题中的噪声干扰（例如，GPT-4 标记的错误真实答案），删除轨迹准确率低于 50% 的合成问题。根据广泛的实验，策略 LLM 在每一轮中都从初始基础模型进行微调，而不是在上一轮的模型上逐步训练。所有策略 SLM 都经过 2 个 epoch 的训练，序列长度为 4096 个 token，批处理大小为 128。使用 AdamW 优化器和线性学习率调度程序，将 Qwen 模型的初始学习率设置为 7e-6，将 Phi3-mini-Instruct 的初始学习率设置为 5e-6。PPM 经过 1 个 epoch 的训练，批处理大小为 512，初始学习率为 7e-6。

推理设置。在评估中，运行多个 MCTS 来生成候选解决方案轨迹。对于每个问题，在每个步骤生成 32 个候选节点，并使用 PPM 对每个节点进行评分。由于 PPM 有效地提供步骤级质量评估，将 MCTS 限制为每步仅 4 次展开以更新 Q 值。完成 MCTS 后，选择 PPM 得分最高的轨迹作为最终答案。