ReST-MCTS*！强化自训练，让大模型持续「升级」

大语言模型（LLM）的自训练（self-training），是验证 Scaling Law 能否继续奏效的关键方法之一。　

然而，由于「错误或无用的中间奖励信号」，现有自训练方法通常会产生低质量的微调训练集（如不正确的规划或中间推理），尽管这可能不会影响 LLM 在一些任务中的正确性，但却会限制 LLM 微调复杂任务的最终性能。　

语言模型生成的错误或者无用的中间推理步骤，也可能产生正确的解题思路。也就是说，即使推理路径产生正确的解题思路，也不一定意味着整个路径是准确的。　

解决方法之一是使用价值函数或者奖励模型来验证推理路径的正确性，然后作为自训练的学习信号。然而，训练一个可靠的奖励模型来验证推理路径中的每一步，通常依赖于密集的人类标注（每个推理步骤），并不能很好地扩展。　

基于对此问题的考虑，我们提出了强化自训练方法 ReST-MCTS*，通过树搜索 MCTS* 指导过程奖励，既可以自动获取可靠的推理路径，也能有效地利用奖励信号进行验证和 LLM 自训练。　

论文链接：https://arxiv.org/abs/2406.03816　

GitHub 地址：https://github.com/THUDM/ReST-MCTS　

项目地址：https://rest-mcts.github.io　

具体来说，ReST-MCTS* 通过基于树搜索的强化学习规避了用于训练过程奖励的每步人工标注：在给定 Oracle 最终正确答案的情况下，ReST-MCTS* 能够通过估算每一步有助于得出正确答案的概率，推断出正确的过程奖励。这些推断出的奖励既是进一步完善过程奖励模型的价值目标，也有助于为策略模型的自训练选择高质量的轨迹。　

实验结果表明，在相同的搜索预算下，ReST-MCTS* 中的树搜索策略比 LLM 推理基线（如 CoT + Best-of-N 和 Tree-of-Thought）实现了更高的准确率；另外，将这种树搜索策略搜索到的轨迹作为训练数据，可以在多次迭代中持续增强 LLM，优于 Self-Rewarding LM 等其他自训练算法。　

研究方法

为了提高模型自动获取可靠推理路径的能力，该研究旨在探索基于模型自动获取高质量推理路径以优化奖励信号和模型自训练性能过程中的关键难点与挑战。　

图｜推断过程奖励的过程以及如何进行过程奖励引导树搜索（左），过程奖励模型（PRM）和策略模型（Policy Model）的自训练（右）。　

该研究的四个核心组件是：　

• 全新的搜索算法 MCTS* ：在过程奖励模型（Process Reward Model，PRM）的指导下，在足够的次数内执行树搜索。

• 过程奖励模型PRM ：用于评估任何部分解答过程的质量并且指导MCTS*。

• 策略模型 ：可以为每个问题生成多个中间推理步骤。

• LLM 自我训练 ：使用MCTS*收集推理路径，在正样本上训练策略模型，并在所有生成的路径上训练过程奖赏模型。

首先，针对现有方法的局限性，研究团队建立了一个全新的搜索算法 MCTS*，该算法是蒙特卡罗树搜索（MCTS）的一个变体，使用每个推理步骤的质量值作为一个训练好的基于 LLM 的过程奖励模型的价值目标，并为 MCTS 提供指导，能够通过足够的部署次数自动标注每个中间节点的过程奖励。　

此外，我们还提出了一个过程奖励模型（PRM），该模型能够准确地计算每个推理步骤的奖励值。该研究将研究过程奖励引导与树搜索相结合，以探索有效的解空间并合成高质量的轨迹。

该研究首先计算对于每个树节点在每棵树中得到正确答案所需的最小推理步骤数，至少有一条推理路径（包括根）。然后，该研究使用过程值的硬估计方法计算奖励值，这意味着一个推理步骤如果能得到每棵树中正确的答案，就被认为是正确的。使用最小推理距离和奖励值，该研究能够推导出在一个正确推理路径上或附近的每个节点的部分解的值。对于每个节点，将至少有一个正确的路径和相关的前进步骤。当推理距离设置为相同的上一步推理距离时，该研究可以使用部分解答过程的质量值的计算过程推导出质量值和加权奖励。　

　

该研究从根节点开始更新所有奖励和值，并且收集它们形成成对的新数据集。该数据集用于在下一个迭代中训练过程奖励模型。注意，在上述过程中搜索树会被修剪并且验证推理轨迹。　

在此基础上，我们结合了监督学习的微调技术，构建了一个新的自训练方法，旨在生成树搜索的过程奖励和高质量的多步推理路径，用于实现过程奖励模型和语言模型相辅相成的自训练。具体而言，在初始化策略模型和过程奖励模型之后，该研究迭代地使用它们，并且利用过程中生成的搜索树为特定的科学或数学问题生成高质量的解决方案，并进行一个自我提升过程。　

实验结果

研究团队从以下三方面验证了 ReST-MCTS* 的有效性。　

首先，研究团队使用生成样本并进行多次迭代评估的自训练方法，如 ReST^EM 和 Self-Rewarding，在三个 LLM backbone 下的分布内和分布外基准上进行了评估。结果显示，ReST-MCTS* 在每次迭代中都优于现有方法，且能够通过自身生成的数据不断自我完善，如下表。　

表｜通过对策略和价值模型进行多次迭代训练得出的初步结果。对于每个 backbone，都分别采用不同的自训练方法。这意味着每种方法都有自身生成的训练数据和相应的奖励（价值）模型。　

然后，我们在 GSM8K 和 MATH500 上对比了 MATH-SHEPHERD（MS）和 SC + MS 等 SOTA 过程奖励模型。结果表明，ReST-MCTS* 学习到了一个很好的过程奖励模型，我们的奖励模型实现了更高的准确率，如下表。　

表｜不同验证器在 GSM8K 测试集和 MATH500 上的准确率。SC: Self-Consistency，MS: MATH-SHEPHERD。验证基于 256 个输出。　

我们还在 MATH 和 SciBench 上进行了相同搜索预算的评估。结果表明，尽管预算不足，ReST-MCTS* 仍明显优于其他基线。　

最后，我们基于 ReST-MCTS*、CoT 和 ToT 树搜索策略分别对比了三个 LLM 在大学水平科学推理基准上的表现，如下表。　

表｜代表性模型在 SciBench 评测集上的总体性能比较。　

　

不足与展望

当然，这项研究也存在一定的局限性。　

例如，研究团队还需要证明 ReST-MCTS* 可以推广到数学以外（如编码、agent 等）的其他推理任务，以及没有 ground-truth（如对话、SWE-Bench 等）的任务。另外，我们还需要扩展所提出的价值模型，并进一步改进数据过滤技术。一个潜在的想法是结合在线 RL 算法，帮助价值模型和策略模型进行更好的自训练。具体如下：　

1. 泛化性有待验证 　

与许多其他自训练研究相似，ReST-MCTS* 也依赖于监督数据集中的 ground-truth oracle 标签来过滤响应；未来，研究团队需要证明 ReST-MCTS* 可以泛化到数学以外（如编码、智能体、对话等）的其他推理任务，以及没有 ground-truth（如对话、SWE-Bench 等）的任务。 　

此外，对于那些需要多步规划和推理的非常复杂的任务（如实施整个软件，如 SWE-Agent 等），没有 ground-truth 答案，还需要一种更好的方法来收集奖励反馈（来自少数人类标签和符号执行或求解器），并训练一个可通用的奖励模型，使其能够在更广泛的任务中发挥作用并提供帮助。　

2. 价值模型的规模和多样性仍需扩展

虽然研究团队基于 Mistral7B: MetaMATH 训练出的价值模型比 SOTA 价值模型 MATH-SHEPHERD 性能更好，但要更好地进行过程奖励模型训练，仍然需要更大规模的价值模型 backbone。此外，训练过程奖励模型的初始训练集是由 SciGLM 生成的，该模型侧重于数学和科学推理任务，但仍然缺乏通用性。虽然目前的过程奖励模型在 MATH 和 SciBench 等多个数学和科学推理任务上取得了 SOTA，但仍然值得探索更多样化的训练集，以便将来扩展到代码生成和智能体规划等多个领域。　

3. 进一步改进自训练数据过滤

推理轨迹的质量影响着自训练的效果，而生成高质量的训练集起着重要作用。因此，研究团队通过训练迭代过程奖励模型来引导树搜索方向，从而获得高质量的轨迹。另一方面，由于价值模型可以帮助筛选出过程值最高的 top-k 生成轨迹，因此我们也希望将更强、更大的 LLM 模型作为价值模型的 backbone。