rStar-Math：小型语言模型通过自我进化的深度思考掌握数学推理
作者: Xinyu Guan1, Li Lyna Zhang12, Yifei Liu, Ning Shang, Youran Sun, Yi Zhu, Fan Yang, Mao Yang
摘要 我们提出了rStar-Math，证明了小型语言模型（SLMs）无需从更强大的模型中进行蒸馏，便能够在数学推理能力上与OpenAI o1匹敌甚至超越。rStar-Math通过蒙特卡洛树搜索（MCTS）进行“深度思考”，其中数学策略SLM在基于SLM的过程奖励模型的指导下进行test-time search。rStar-Math引入了三项创新来解决训练这两个SLM的挑战：
- 一种新颖的代码增强的CoT数据合成方法，通过广泛的MCTS 滚动生成逐步验证的推理轨迹，用于训练策略SLM；
- 一种新颖的过程奖励模型训练方法，避免了简单的步骤级评分标注，生成了更有效的过程偏好模型（PPM）；
- 一种自我进化配方，其中策略SLM和PPM从零开始构建，并通过迭代进化提高推理能力。通过4轮自我进化，生成了747k个数学问题的数百万个合成解决方案，rStar-Math将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%。代码和数据将在https://github.com/microsoft/rStar上提供。
1 引言
最近的研究表明，大型语言模型（LLMs）能够解决数学问题[16, 27, 6]。然而，传统的让LLMs在单次推理中生成完整解决方案的方法——类似于系统1思维[14]——通常会产生快速但容易出错的结果[25, 6]。作为回应，test-time compute 扩展[24, 10]提出了一种向系统2思维模式的转变，该模式通过更慢、更深层次的思考过程模拟人类推理。在这种模式下，LLM作为策略模型生成多个数学推理步骤，然后由另一个LLM作为奖励模型进行评估[6]。被认为更有可能正确的步骤和解决方案被选中。该过程重复迭代，最终得出最终答案。
在test-time compute 范式中，关键是训练一个强大的策略模型来生成有希望的解决方案步骤，以及一个可靠的奖励模型来准确评估它们，这两者都依赖于高质量的训练数据。不幸的是，众所周知，现成的高质量数学推理数据稀缺，合成高质量数学数据面临根本性挑战。对于策略模型来说，区分错误的推理步骤和正确的步骤具有挑战性，这使得消除低质量数据变得复杂。值得注意的是，在数学推理中，正确的最终答案并不能确保整个推理轨迹的正确性（Lanham等，2023）。错误的中间步骤会显著降低数据质量。对于奖励模型，过程奖励建模（PRM）通过提供对中间步骤的细粒度反馈显示出巨大的潜力（Lightman等，2023）。然而，这方面的训练数据更加稀缺：准确的逐步反馈需要大量的人工标注工作，并且难以扩展，而那些自动标注尝试由于噪声奖励分数而显示出有限的增益（Luo等，2024；Wang等；Chen等，2024）。由于上述挑战，现有的基于蒸馏的数据合成方法（例如，扩展GPT4蒸馏的CoT数据[Tang等，2024；Huang等，2024]）已经显示出收益递减，并且无法超越其教师模型的能力；同时，迄今为止，训练可靠的PRM用于数学推理仍然是一个未解决的问题。
在这项工作中，我们引入了rStar-Math，这是一种可自我进化的系统2推理方法，能够在具有7亿参数的模型上实现最先进的数学推理，与OpenAI o1在具有挑战性的数学竞赛基准上匹敌甚至超越。与依赖更强大的LLM进行数据合成的解决方案不同，rStar-Math利用较小的语言模型（SLMs）和蒙特卡洛树搜索（MCTS）建立了一个自我进化过程，迭代生成更高质量的训练数据。为了实现自我进化，rStar-Math引入了三项关键创新。
首先，一种新颖的代码增强的CoT数据合成方法，通过广泛的MCTS滚动生成逐步验证的推理轨迹，并带有自注释的MCTS Q值。具体来说，数学问题解决被分解为MCTS内的多步生成。在每一步中，作为策略模型的SLM对候选节点进行采样，每个节点生成一步CoT和相应的Python代码。为了验证生成质量，仅保留成功执行Python代码的节点，从而减少中间步骤中的错误。此外，广泛的MCTS滚动自动为每个中间步骤分配一个Q值，基于其贡献：导致更多正确答案轨迹的步骤被赋予更高的Q值，并被认为是更高质量的。这确保了由SLM生成的推理轨迹由正确、高质量的中间步骤组成。
其次，一种新颖的方法，训练一个SLM作为过程偏好模型，即PPM，以实现所需的PRM，可靠地预测每个数学推理步骤的奖励标签。PPM利用了这样一个事实：尽管使用广泛的MCTS滚动后，Q值仍然不足以精确评分每个推理步骤，但Q值可以可靠地区分正（正确）步骤和负（无关/错误）步骤。因此，训练方法基于Q值为每个步骤构建偏好对，并使用成对排序损失（Ouyang等，2022）来优化PPM对每个推理步骤的评分预测，从而实现可靠的标注。这种方法避免了传统方法直接使用Q值作为奖励标签（Luo等，2024；Chen等，2024），这些方法在逐步奖励分配中本质上是噪声和不精确的。
最后，一个四轮自我进化配方，逐步从零开始构建前沿策略模型和PPM。我们首先从公开来源策划了一个包含747k个数学问题的数据集。在每一轮中，我们使用最新的策略模型和PPM执行MCTS，使用上述两种方法生成越来越高质量的训练数据，以训练下一轮更强的策略模型和PPM。每一轮都实现了逐步改进：1）更强的策略SLM，2）更可靠的PPM，3）通过PPM增强的MCTS生成更好的推理轨迹，4）提高训练数据覆盖率以应对更具挑战性甚至竞赛级别的数学问题。
在四个SLM（1.5B-7B）和七个数学推理任务上的广泛实验证明了rStar-Math的有效性。值得注意的是，rStar-Math改进了所有四个SLM，在具有挑战性的数学基准上与OpenAI o1匹敌甚至超越。在MATH基准测试中，使用8个搜索轨迹，rStar-Math将Qwen2.5-Math-7B从58.8%提升到89.4%，将Qwen2.5-Math-1.5B从51.2%提升到87.8%。使用64个轨迹，分数分别上升到90%和88.4%，分别超过了o1-preview的4.5%和2.6%，并与o1-mini的90%持平。在奥林匹克级别的AIME 2024中，rStar-Math平均解决了53.3%（8/15）的问题，超过了o1-preview的8.7%和其他所有开源的LLM。我们进一步进行了全面的实验，验证了逐步验证的推理轨迹相对于最先进的数据合成基线的优越性，以及PPM相对于结果奖励模型和基于Q值的PRM的有效性。最后，我们展示了rStar-Math深度思考的关键发现，包括内在的自我反思能力和PPM对定理应用中间步骤的偏好。
2 相关工作
数学数据合成。LLM数学推理的进展主要依赖于策划高质量的CoT数据，大多数领先的方法都是基于GPT蒸馏的，使用GPT-4等前沿模型进行合成[16, 9, 17]。值得注意的工作包括NuminaMath[10]和MetaMath[22]。虽然有效，但这将推理限制在教师LLM的能力范围内。教师LLM无法解决的难题被排除在训练集之外。即使是可解决的问题也可能包含容易出错的中间步骤，这些步骤难以检测。尽管拒绝采样方法[22, 14]可以提高数据质量，但它们不能保证正确的中间步骤。因此，扩展CoT数据的收益递减，增益接近饱和——例如，OpenMathInstruct-2[21]尽管数据集大小增加了8倍，但在MATH上仅看到了3.9%的提升。
**扩展test-time compute **引入了新的扩展定律，允许LLM通过生成多个样本并使用奖励模型进行最佳解决方案选择来提高性能[20, 23, 14]。已经提出了各种 test-time search 方法[19, 24]，包括随机采样[24]和树搜索方法[25, 11, 26, 12]，如MCTS。然而，开源的test-time compute 方法在数学推理中显示出有限的增益，通常是由于策略LLM或奖励模型的限制。rStar-Math通过迭代进化策略LLM和奖励模型来解决这个问题，实现了与OpenAI o1[10]相当的系统2数学推理性能。
奖励模型对于有效的系统2推理至关重要，但难以获得。最近的工作包括用于验证的LLM-as-a-Judge[27, 12]和专门的奖励模型，如结果奖励模型[25, 22]和过程奖励模型（PRM）[18]。虽然PRM提供了有希望的密集逐步奖励信号用于复杂推理[17, 24]，收集逐步注释仍然是一个障碍。依赖于昂贵的人工注释数据集[19, 24]，如PRM800k[18]，最近的方法[24, 17]探索了通过蒙特卡洛采样或MCTS进行自动注释。然而，它们难以生成精确的奖励分数，这限制了性能增益。rStar-Math引入了一种新颖的过程偏好奖励（PPM），消除了对准确逐步奖励分数注释的需求。
3 方法论
设计选择
MCTS用于有效的系统2推理。我们的目标是训练一个数学策略SLM和一个过程奖励模型（PRM），并将两者集成到蒙特卡洛树搜索（MCTS）中，以实现系统2深度思考。选择MCTS有两个关键原因。首先，它将复杂的数学问题分解为更简单的单步生成任务，与其他系统2方法（如Best-of-N[Brown等，2024]或self-consistency[Wang等，2023]）相比，降低了策略SLM的难度，这些方法需要在一次推理中生成完整的解决方案。其次，MCTS中的逐步生成自然地为两个模型提供了逐步训练数据。标准的MCTS滚动自动为每个步骤分配Q值，基于其对最终正确答案的贡献，消除了对人工生成的逐步注释的需求。
理想情况下，可以将GPT-4等先进的LLM集成到MCTS中以生成训练数据。然而，这种方法面临两个关键挑战。首先，即使是这些强大的模型也难以一致地解决难题，如奥林匹克级别的数学问题。因此，生成的训练数据将主要由较简单的可解决问题组成，限制了其多样性和质量。其次，注释每步Q值需要广泛的MCTS滚动；不足的树探索可能导致虚假的Q值分配，例如高估次优步骤。鉴于每次滚动涉及多个单步生成，并且这些模型在计算上昂贵，增加滚动会显著提高推理成本。
概述。为此，我们探索使用两个7B SLM（策略SLM和PRM）生成更高质量的训练数据，其较小的规模允许在可访问的硬件（例如，4×40GB A100 GPU）上进行广泛的MCTS滚动。然而，自生成数据对SLM提出了更大的挑战，因为它们的能力较弱。SLM经常无法生成正确的解决方案，即使最终答案正确，中间步骤也常常存在缺陷或质量低下。此外，与GPT-4等先进模型相比，SLM解决的难题较少。
本节介绍了我们的方法，如图1所示。为了减少错误和低质量的中间步骤，我们引入了一种代码增强的CoT合成方法，通过广泛的MCTS滚动生成逐步验证的推理轨迹，并带有Q值注释。为了进一步提高SLM在难题上的表现，我们引入了一个四轮自我进化配方。在每一轮中，策略SLM和奖励模型都更新为更强的版本，逐步解决更困难的问题并生成更高质量的训练数据。最后，我们提出了一种新颖的过程奖励模型训练方法，消除了对精确每步奖励注释的需求，生成了更有效的过程偏好模型（PPM）。
逐步验证的推理轨迹
我们首先介绍生成带有每步Q值注释的逐步验证推理轨迹的方法。给定一个问题$x$和一个策略模型$M$，我们运行标准的MCTS以逐步构建一个搜索树，用于逐步解决方案探索。如图1(a)所示，根节点表示问题$x$，而子节点对应于由$M$生成的中间步骤$s$。从根节点到终端节点$s_{d}$的路径形成了一条轨迹${t}=x\oplus s_{1}\oplus s_{2}\oplus...\oplus s_{d}$，每个步骤$s_{i}$被分配一个Q值$Q(s_{i})$。从搜索树$\mathcal{T}$中，我们提取解决方案轨迹$\mathbb{T}={ {t}^{1},{t}^{2},...,{t}^{n}}(n\geq 1)$。我们的目标是从$\mathcal{T}$中选择高质量的轨迹来构建训练集。为此，我们引入了代码增强的CoT合成方法，以过滤掉低质量的生成，并通过广泛的滚动提高Q值准确性的可靠性。
代码增强的CoT生成。之前的MCTS方法主要生成自然语言（NL）CoT（Oi等，2024；Zhang等）。然而，LLM经常出现幻觉，生成错误或不相关的步骤，但仍然偶然得出正确答案（Lanham等，2023）。这些有缺陷的步骤难以检测和消除。为了解决这个问题，我们提出了一种新颖的代码执行增强的CoT。如图2所示，策略模型生成一步NL CoT及其相应的Python代码，其中NL CoT嵌入为Python注释。仅保留成功执行Python代码的生成作为有效候选。
具体来说，从初始根节点$x$开始，我们通过选择、扩展、滚动和反向传播执行多次MCTS迭代。在第$i$步，我们收集最新的推理轨迹$x\oplus s_{1}\oplus s_{2}\oplus...\oplus s_{i-1}$作为当前状态。基于此状态，我们提示（见附录A.3）策略模型生成$n$个候选步骤$s_{i,0},...,s_{i,n-1}$。然后使用Python代码执行来过滤有效节点。如图2所示，每个生成$s_{i,j}$与之前所有步骤的代码连接，形成$s_{1}\oplus s_{2}\oplus...\oplus s_{i-1}\oplus s_{i,j}$。成功执行的候选节点被保留为有效节点，并由PPM评分，分配一个Q值$q(s_{i})$。然后，我们使用著名的树的上置信界（UCT）[Kocsis和Szepesvari，2006]从$n$个候选中选择最佳节点。此选择过程数学表示为：
$$ \text{UCT}(s)=Q(s)+c\sqrt{\frac{\ln N_{parent}(s)}{N(s)} };\quad\text{其中}\quad Q(s)=\frac{q(s)}{N(s)} $$
其中$N(s)$表示节点$s$的访问次数，$N_{\text{parent} }(s)$是$s$的父节点的访问次数。预测的奖励$q(s)$由PPM提供，并通过反向传播更新。$c$是一个常数，用于平衡利用和探索。
广泛的滚动用于Q值注释。准确的Q值$Q(s)$注释在公式1中至关重要，用于引导MCTS节点选择正确的解题路径，并识别轨迹中的高质量步骤。为了提高Q值的可靠性，我们从围棋玩家那里获得灵感，他们根据游戏结果回顾性地评估每一步的奖励。尽管初始估计可能不精确，但重复的游戏会随着时间的推移改进这些评估。同样，在每次滚动中，我们根据每个步骤对实现最终正确答案的贡献更新其Q值。经过广泛的MCTS滚动后，始终导致正确答案的步骤获得更高的Q值，偶尔成功的步骤获得中等Q值，而始终错误的步骤获得低Q值。具体来说，我们引入了两种自注释方法来获得这些步骤级Q值。图1(c)显示了四轮自我进化中的详细设置。
终端引导的注释。在前两轮中，当PPM不可用或不够准确时，我们使用终端引导的注释。形式上，让$q(s_{i})^{k}$表示在第$k$次滚动中反向传播后的步骤$s_{i}$的q值。遵循AlphaGo[Silver等，2017]和rStar[Qi等，2024]，我们根据每个中间节点对最终正确答案的贡献进行评分：
$$ q(s_{i})^{k}=q(s_{i})^{k-1}+q(s_{d})^{k}; $$
其中初始q值$q(s_{i})^{0}=0$在第一次滚动中。如果此步骤经常导致正确答案，其$q$值将增加；否则，它将减少。终端节点被评分为$q(s_{d})=1$表示正确答案，$q(s_{d})=-1$表示错误答案，如图1所示。
PRM增强的注释。从第三轮开始，我们使用PPM对每个步骤进行评分，以进行更有效的生成。与终端引导的注释相比，后者需要多次滚动才能获得有意义的$q$值，PPM直接预测一个非零的初始$q$值。PRM增强的MCTS还有助于策略模型生成更高质量的步骤，引导解决方案走向正确的路径。形式上，对于步骤$s_{i}$，PPM基于部分轨迹预测初始$q(s_{i})^{0}$值：
$$ q(s_{i})^{0}=PPM(x\oplus s_{1}\oplus s_{2}\oplus...\oplus s_{i-1}\oplus s_{i}) $$
此$q$值将通过MCTS反向传播根据终端节点的$q(s_{d})$值更新。对于终端节点$s_{d}$，我们在训练数据生成期间不使用PRM进行评分。相反，我们根据ground truth标签分配更准确的分数作为终端引导的奖励。
过程偏好模型
过程奖励模型提供了细粒度的步骤级奖励信号，对于解决具有挑战性的数学问题非常理想。然而，获得高质量的步骤级训练数据仍然是一个未解决的挑战。现有方法依赖于人工注释[Lightman等，2023]或MCTS生成的分数[Zhang等；Chen等，2024]来为每个步骤分配分数。这些分数随后作为训练目标，使用MSE损失[Chen等，2024]或点损失[Wang等；Luo等，2024；Zhang等]来最小化预测分数和标注分数之间的差异。因此，这些标注的步骤级奖励分数的精度直接决定了生成的过程奖励模型的有效性。
不幸的是，精确的每步评分仍然是一个未解决的挑战。尽管我们广泛的MCTS滚动提高了Q值的可靠性，但精确评估细粒度的步骤质量仍然是一个主要障碍。例如，在一组正确的步骤中，很难将它们排名为最佳、次佳或平均，然后分配精确的分数。同样，在错误的步骤中，区分最差和中等差的步骤也面临类似的挑战。即使是专家人工注释在一致性方面也存在困难，特别是在大规模情况下，导致训练标签中存在固有的噪声。
我们引入了一种新颖的训练方法，通过构建步骤级的正负偏好对来训练过程偏好模型（PPM）。如图1(b)所示，我们不使用Q值作为直接奖励标签，而是使用它们从MCTS树中选择步骤以构建偏好对。对于每个步骤，我们选择两个具有最高Q值的候选作为正步骤，选择两个具有最低Q值的候选作为负步骤。关键的是，所选择的正步骤必须导致正确的最终答案，而负步骤必须导致错误的答案。对于中间步骤（除了最终答案步骤），正负对共享相同的前置步骤。对于最终答案步骤，其中相同的推理轨迹很少产生不同的最终答案，我们放宽了这一限制。我们选择两个具有最高平均Q值的正确轨迹作为正例，选择两个具有最低平均Q值的错误轨迹作为负例。遵循[Ouyang等，2022]，我们使用标准的Bradley-Terry模型定义我们的损失函数，并使用成对排序损失：
$$ \mathcal{L}{ppm}(\theta)=-\frac{1}{2\times 2}E{i},y^{neq}\oplus )}[log(\sigma(r_{\theta}(x,y^{pos}{i})-r{\theta}(x,y^{neq}_{i})))] $$
当$i$不是最终答案步骤时，$y^{pos}{i}=s\oplus...\oplus s_{i-1}\oplus s^{pos}{i}$；$y^{neq}=s_{1}\oplus...\oplus s_{i-1}\oplus s^{neq}_{i}$（5）
这里，$r_{\theta}(x,y_{i})$表示PPM的输出，其中$x$是问题，$y$是从第一步到第$i$步的轨迹。
自我进化的深度思考
3.4.1 使用逐步验证的推理轨迹进行训练
数学问题收集。我们收集了一个包含747k个数学问题的数据集，主要来自NuminaMath[Jia LI和Polu]和MetaMath[Yu等]。值得注意的是，仅包括来自NuminaMath的竞赛级别问题（例如，奥林匹克和AIME/AMC），因为我们观察到小学级别的问题对LLM复杂数学推理的提升不大。为了增加有限的竞赛级别问题，我们遵循[Li等，2024]并使用GPT-4基于7.5k MATH训练集和3.6k AMC-AIME训练集中的种子问题合成新问题。然而，GPT-4经常为具有挑战性的种子问题生成无法解决的问题或不正确的解决方案。为了过滤这些问题，我们提示GPT-4为每个问题生成10个解决方案，仅保留至少3个一致解决方案的问题。
推理轨迹收集。我们没有使用747k数学数据集中的原始解决方案，而是进行了广泛的MCTS滚动（第3.2节）以生成更高质量的逐步验证的推理轨迹。在每轮自我进化中，我们对每个数学问题执行16次滚动，生成16条推理轨迹。然后根据生成的轨迹的正确率将问题分类为：简单（所有解决方案都正确）、中等（混合正确和错误的解决方案）和困难（所有解决方案都错误）。对于没有正确轨迹的困难问题，额外执行16次滚动的MCTS。之后，所有逐步轨迹及其注释的Q值被收集并过滤，用于训练策略SLM和过程偏好模型。
监督微调策略SLM。通过广泛的实验，我们发现选择高质量的推理轨迹是微调前沿数学LLM的关键。虽然诸如GPT蒸馏和Best-of-N等方法可能包含低质量或错误的中间步骤，但更有效的方法确保轨迹中的每一步都是高质量的。为此，我们使用每步Q值从MCTS滚动中选择最佳轨迹。具体来说，对于每个数学问题，我们选择前2条具有最高平均Q值的轨迹作为SFT训练数据。
训练PPM。PPM从微调的策略模型初始化，其下一个令牌预测头被替换为由线性层和tanh函数组成的标量值头，以将输出限制在[-1, 1]范围内。我们过滤掉所有解决方案轨迹完全正确或完全错误的数学问题。对于混合结果的问题，我们根据Q值为每个步骤选择两个正例和两个负例，作为训练数据的偏好对。
3.4.2 自我进化配方
由于SLM的能力较弱，我们进行了四轮MCTS深度思考，逐步生成更高质量的数据，并扩展训练集以包含更具挑战性的数学问题。
每轮使用MCTS生成逐步验证的推理轨迹，然后用于训练新的策略SLM和PPM。新模型随后应用于下一轮，以生成更高质量的训练数据。图1(c)和表2详细说明了每轮用于数据生成的模型，以及训练的策略模型和PPM的标识符。接下来，我们概述每轮的细节和具体改进目标。
第1轮：引导初始强策略SLM-r1。为了使SLM能够自生成合理良好的训练数据，我们执行了一轮引导，以微调一个初始强策略模型，记为SLM-r1。如表2所示，我们使用DeepSeek-Coder-V2-Instruct（236B）运行MCTS以收集SFT数据。由于此轮没有可用的奖励模型，我们使用终端引导的注释进行Q值标注，并将MCTS限制为8次滚动以提高效率。对于正确的解决方案，选择具有最高平均Q值的前2条轨迹作为SFT数据。我们还训练了PPM-r1，但有限的滚动导致Q值不可靠，影响了PPM-r1的有效性（表4）。
第2轮：训练可靠的PPM-r2。在这一轮中，随着策略模型更新为7B SLM-r1，我们进行了广泛的MCTS滚动以获得更可靠的Q值注释，并训练了第一个可靠的奖励模型PPM-r2。具体来说，我们对每个问题执行16次MCTS滚动。生成的逐步验证的推理轨迹在质量和Q值精度上都有显著提高。如表4所示，PPM-r2比引导轮显著更有效。此外，策略SLM-r2也如预期继续改进（表3）。
第3轮：PPM增强的MCTS显著提高数据质量。在第三轮中，我们使用可靠的PPM-r2执行PPM增强的MCTS以生成数据，从而生成显著更高质量的轨迹，覆盖了训练集中更多的数学和奥林匹克级别问题（表2）。生成的推理轨迹和自注释的Q值随后用于训练新的策略SLM-r3和PPM-r3，两者都显示出显著改进。
第4轮：解决具有挑战性的数学问题。在第三轮之后，虽然小学和MATH问题达到了高成功率，但训练集中仅包括62.16%的奥林匹克级别问题。这并非仅仅由于我们SLM的推理能力较弱，因为许多奥林匹克问题仍然无法被GPT-4或o1解决。为了提高覆盖率，我们采用了一个简单的策略。对于16次MCTS滚动后仍未解决的问题，我们额外执行64次滚动，如果需要，增加到128次。我们还使用不同的随机种子进行多次MCTS树扩展。这将奥林匹克级别问题的成功率提高到80.58%。
经过四轮自我进化，747k个数学问题中的90.25%成功覆盖到训练集中，如表2所示。在剩余的未解决问题中，很大一部分是合成问题。我们手动审查了20个问题的随机样本，发现19个问题被错误地标记为错误答案。基于此，我们得出结论，剩余的未解决问题质量较低，因此在第4轮终止自我进化。
4 评估
设置
评估数据集。我们在多样化的数学基准上评估rStar-Math。除了广泛使用的GSM8K[Cobbe等，2021]，我们还包括来自多个领域的具有挑战性的基准：
- 竞赛和奥林匹克级别基准，如MATH-500[Lightman等，2023]、AIME 2024[AI-MO]、AMC 2023[AI-MO]和Olympiad Bench[He等，2024]。具体来说，AIME是为挑战美国最优秀的高中数学学生而设计的考试，2024年数据集包括来自AIME I和II考试的30个问题；
- 大学级别的数学问题，来自College Math[Tang等，2024]；
- 领域外的数学基准：中国高考数学考试（Gaokao）En 2023[Liao等，2024]。
基础模型和设置。rStar-Math是一种适用于各种LLM的通用方法。为了展示其有效性和通用性，我们使用不同大小的SLM作为基础策略模型：Qwen2.5-Math-1.5B[Qwen]、Phi3-mini-Instruct（3B）[Microsoft，2024，Abdin等，2024]、Qwen2-Math-7B[Qwen]和Qwen2.5-Math-7B[Qwen]。其中，Phi3-mini-Instruct是一个通用SLM，没有专门用于数学推理。
由于GPU资源有限，我们仅在Qwen2.5-Math-7B上进行了4轮自我进化，生成了4个进化的策略SLM（表3）和4个PPM（表4）。对于其他3个策略LLM，我们使用从Qwen2.5-Math-7B的第4轮生成的逐步验证轨迹对它们进行微调。此轮的最终PPM随后用作3个策略SLM的奖励模型。
基线。rStar-Math是一种系统2方法。我们将其与代表系统1和系统2方法的三个强基线进行比较：
- 前沿LLM，包括GPT-4o、最新的Claude、OpenAI o1-preview和o1-mini。我们测量了它们在AMC 2023、Olympiad Bench、College Math、Gaokao和GSM8K上的准确性，其他基准的准确性数据来自公开的技术报告[Team]。
- 开源的高级推理模型，包括DeepSeek-Coder-v2-Instruct、Mathstral[Team]、NuminaMath-72B[Jia LI和Polu]和LLaMA3.1[Dubey等，2024]，它们代表了当前改进LLM数学推理的主流系统1方法。
- 基础模型的系统1和系统2性能，包括原始模型团队开发的Instruct版本（例如，Qwen2.5-Math-7B-Instruct）和Best-of-N（例如，Qwen2.5-Math-72B-Instruct+Qwen2.5-Math-RM-72B）。值得注意的是，用于三个Qwen基础模型的奖励模型是一个72B ORM，显著大于我们的7B PPM。
评估指标。我们报告所有基线的Pass@1准确性。对于系统2基线，我们使用默认的评估设置，例如o1-mini和o1-preview的默认思考时间。对于具有Best-of-N的Qwen模型，我们重新评估了MATH-500、AIME/AMC的准确性；其他基准的结果来自它们的技术报告。为了公平比较，rStar-Math运行MCTS以生成与Qwen相同数量的解决方案。具体来说，对于AIME/AMC，我们生成16条轨迹，对于其他基准生成8条轨迹，使用PPM选择最佳解决方案。我们还报告了使用64条轨迹增加test-time compute 的性能，记为rStar-Math$^{64}$。
主要结果
在多样化的具有挑战性的数学基准上的结果。表5显示了rStar-Math与最先进的推理模型的比较结果。我们强调了三个关键观察：
- rStar-Math显著提高了SLM的数学推理能力，在模型大小显著较小（1.5B-7B）的情况下，实现了与OpenAI o1相当甚至超越的性能。例如，Qwen2.5-Math-7B在MATH上的准确性从58.8%显著提高到90.0%，超过了o1-preview和Claude 3.5 Sonnet，并与o1-mini持平。在College Math基准上，rStar-Math超过了o1-mini的2.7%。在AIME 2024上，rStar-Math得分为53.3%，仅次于o1-mini，7B模型解决了AIME I和II中的8/15个问题，跻身于最优秀的高中数学学生前20%。值得注意的是，8个未解决的问题是基于几何的，需要视觉理解，这是rStar-Math目前不支持的能力。
- 尽管使用较小的策略模型（1.5B-7B）和奖励模型（7B），rStar-Math显著优于最先进的系统2基线。与使用相同基础模型（Qwen2-Math-7B、Qwen2.5-Math-1.5B/7B）但使用10倍大的奖励模型（Qwen2.5-Math-RM-72B）的Qwen Best-of-N基线相比，rStar-Math一致地将所有基础模型的推理准确性提高到最先进的水平。即使与使用10倍大的Qwen2.5-Math-72B-Instruct策略模型的Best-of-N相比，rStar-Math在除GSM8K之外的所有基准上都超过了它，使用相同数量的采样解决方案。
- 除了众所周知的基准，如MATH、GSM8K和AIME，这些基准可能面临过度优化的风险，rStar-Math在其他具有挑战性的数学基准上表现出强大的泛化能力，包括Olympiad Bench、College Math和中国高考数学考试（Gaokao），设定了新的最先进分数。如第3.4节所述，我们的训练集主要来自公共数据集，没有针对这些基准进行特定优化。
**扩展test-time compute **。rStar-Math使用MCTS增强策略模型，在PPM的指导下搜索解决方案。通过增加test-time compute ，它探索了更多的轨迹，可能提高性能。在图3中，我们通过比较官方Qwen Best-of-N在不同数量的采样轨迹上的准确性，展示了test-time compute 扩展的影响。仅采样一条轨迹对应于策略LLM的Pass@1准确性，表明回退到系统1推理。我们强调了两个关键观察：
- 仅使用4条轨迹，rStar-Math显著优于Best-of-N基线，超过了o1-preview并接近o1-mini，展示了其有效性。
- 扩展test-time compute 提高了所有基准上的推理准确性，尽管趋势不同。在Math、AIME和Olympiad Bench上，rStar-Math在64条轨迹时显示出饱和或缓慢改进，而在College Math上，性能继续稳步提高。
消融研究和分析
我们消融了我们三项创新的有效性。对于系统2风格的推理，Pass@1准确性是在AIME和AMC上使用16条轨迹测量的，其他基准使用8条轨迹。
自我进化的有效性。表5中的令人印象深刻的结果是在4轮rStar-Math自我进化的深度思考后实现的。表6显示了每轮的数学推理性能，展示了准确性的持续改进。在第1轮中，主要改进来自于对基础模型应用SFT。第2轮带来了显著提升，应用了更强的PPM在MCTS中，这解锁了系统2深度推理的全部潜力。值得注意的是，从第2轮开始，rStar-Math超过了GPT-4o。第3轮和第4轮显示了进一步的改进，通过更好的策略SLM和PPM推动了更强的系统2推理。
逐步验证的推理轨迹的有效性。rStar-Math生成逐步验证的推理轨迹，消除了错误的中间步骤，并进一步扩展了训练集以包含更具挑战性的问题。为了评估其有效性，我们使用第4轮生成的数据作为SFT训练数据，并将其与三个强基线进行比较：
- GPT蒸馏，包括使用GPT-4合成的开源CoT解决方案，如MetaMath[26]、NuminaMath-CoT[14]；
- 自生成的随机采样，使用相同的策略模型（即策略SLM-r3）随机生成轨迹；
- 拒绝采样，其中从策略模型中随机采样32条轨迹，并由我们训练的ORM对高质量解决方案进行排名（附录A.1）。为了公平起见，我们在基线（ii）和（iii）中为每个数学问题选择两条正确的轨迹。所有SFT实验使用相同的训练配方。
表7显示了Qwen2.5-Math-7B在不同数据集上微调后的数学推理准确性。我们强调了两个观察：
- 使用我们的逐步验证轨迹进行微调显著优于所有其他基线。这主要是由于我们的PPM增强的MCTS用于代码增强的CoT合成，在数学解决方案生成过程中提供了更密集的验证。它比随机采样（缺乏验证）和拒绝采样（仅提供稀疏验证）更有效。
- 即使是从我们的SLM随机采样的代码增强的CoT解决方案，其性能也与GPT-4合成的NuminaMath和MetaMath数据集相当或更好。这表明我们的策略SLM在经过几轮自我进化后，能够生成高质量的数学解决方案。这些结果展示了我们的方法在不依赖高级LLM蒸馏的情况下自生成更高质量推理数据的巨大潜力。
PPM的有效性。我们训练了一个强大的ORM和一个基于Q值的PRM（PQM）进行比较。为了确保公平评估，我们使用最高质量的训练数据：第4轮生成的逐步验证轨迹，选择的数学问题与用于PPM训练的问题匹配。类似于PPM，我们使用步骤级Q值为每个数学问题选择正负轨迹。ORM使用成对排序损失（Ouyang等，2022）进行训练，而PQM遵循[Chen等，2024；Zhang等]使用Q值作为奖励标签，并使用MSE损失进行优化。详细的训练设置见附录A.1。
表8比较了ORM、PQM和PPM在使用我们最终轮策略模型进行系统2推理时的性能。ORM仅在问题解决结束时提供奖励信号，因此我们使用Best-of-N方法，而PRM和PPM利用MCTS驱动的搜索。如表8所示，PQM和PPM通过提供更密集的步骤级奖励信号，在复杂数学推理任务上优于ORM。然而，PQM在更具挑战性的基准（如MATH和Olympiad Bench）上表现不佳，这是由于Q值固有的不精确性。相比之下，PPM构建了步骤级偏好数据进行训练，使我们的7B策略模型在所有基准上实现了与o1-mini相当或更优的性能。
5 发现与讨论
内在自我反思能力的出现。OpenAI o1的一个关键突破是其内在的自我反思能力。当模型出错时，它能够识别错误并自我纠正，得出正确答案（Noam Brown和Lightman，2024）。然而，它一直被发现在开源的LLM中基本无效。社区积极探索了各种方法，包括自我纠正[Huang等，2023；Kumar等，2024]、自我反思[Renze和Guven，2024；Shinn等，2024]，以明确训练或提示LLM发展这种能力。
在我们的实验中，我们意外地观察到我们的MCTS驱动的深度思考在问题解决过程中表现出自我反思。如图4所示，模型在前三个步骤中使用SymPy形式化了一个方程，这将导致错误的答案（左分支）。有趣的是，在第四步（右分支），策略模型认识到其早期步骤的低质量，并避免继续沿着初始的问题解决路径前进。相反，它回溯并使用一种新的、更简单的方法解决了问题，最终得出了正确答案。附录A.2中提供了另一个自我纠正的示例。值得注意的是，没有包含自我反思的训练数据或提示，这表明先进的系统2推理可以促进内在的自我反思。
PPM在系统2深度思考中塑造推理边界。策略模型和奖励模型对于系统2深度推理都至关重要。我们的实验表明，一旦策略模型达到了合理的能力水平（见附录A.1），PPM就成为性能上限的关键决定因素。图5总结了不同大小的策略模型的准确性，以及应用奖励模型后的改进。尽管由于训练策略、数据集和模型规模的差异，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识别定理应用步骤。在解决具有挑战性的数学问题时，识别和应用相关定理或关键结论通常是成功解决问题的基石[]。在我们的实验中，我们发现，在rStar-Math问题解决过程中，我们的PPM有效地识别了策略模型深度思考过程中的关键定理应用中间步骤。这些步骤被预测为高奖励分数，引导策略模型生成正确的解决方案。附录A.2提供了PPM成功识别关键定理的示例，如费马小定理[]、Vieta公式[]、AM-GM不等式[]、毕达哥拉斯定理[]和鞋带定理[]等。
6 结论
在这项工作中，我们提出了rStar-Math，一种自我进化的系统2深度思考方法，显著提升了小型LLM的数学推理能力，达到了最先进的OpenAI o1水平。我们的方法展示了SLM可以自生成高质量的训练数据，用于前沿级别的数学推理。在四个不同大小的SLM和具有挑战性的数学基准上的广泛实验证明了rStar-Math的优越性，实现了领先的结果，同时超越了现有的数学推理LLM和Best-of-N基线。我们还揭示了关键发现，包括自我反思的出现和PPM在识别关键中间步骤（如定理应用步骤）方面的有效性。最后，rStar-Math可以通过收集更多具有挑战性的数学问题实现进一步改进，我们将此留作未来的工作。
致谢
在早期阶段，由于GPU资源有限和GPT-4 API访问受限，我们面临了重大挑战。我们非常感谢Qiufeng Yin和Chengmin Chi在收集数学问题和提供GPT-4资源以合成新数学问题方面的帮助。特别感谢我的同事Lingxiao Ma、Ying Cao、Baotong Lu、Jing Liu、Jiahang Xu、Chengruidong Zhang、Siyuan Wang、Gaokai Zhang、Yujian Li和Yang Wang慷慨地与我们分享了他们的GPU配额。