LLM实现自回归搜索！MIT哈佛等提出「行动思维链」COAT，推理能力大提升

新智元报道
编辑：英智
【新智元导读】本研究提出了一种创新的自回归搜索方法，通过两阶段训练框架，小规模格式调优和大规模自我优化，开发出了Satori，一个在数学推理和跨领域任务中均表现优异的7B参数模型。Satori通过自我反思和探索策略，展现了强大的迁移能力和自我纠错能力。
OpenAI o1发布后，为提升LLM的推理能力，研究者尝试了多种方法。
比如用强大的教师模型进行知识蒸馏、采用蒙特卡洛树搜索（MCTS），以及基于奖励模型的引导搜索。
近日，来自MIT、新加坡科技设计大学、哈佛大学等机构的华人研究者探索了全新的方向：让LLM拥有自回归搜索能力。通过自我反思和探索新策略，提升LLM推理能力。
研究者引入了行动-思维链（COAT）机制，使LLM在解决问题时能够执行多种元动作，并提出了一种创新的两阶段训练框架：

小规模格式调优阶段：让LLM熟悉并掌握COAT推理格式。
大规模自我优化阶段：运用重启与探索（RAE）技术，通过RL进行优化。



通过这种方法，成功开发出Satori，在数学推理任务中，成绩优异。
Satori具有以下核心特点：

无需外部指导，即可自我反思与探索。
主要依靠自我改进（RL），实现了最先进的推理性能。
展现出强大的迁移能力，可应用于数学以外的领域。



论文地址：https://arxiv.org/pdf/2502.02508
开源项目：https://github.com/satori-reasoning/Satori
Satori关键设计
研究者把LLM的推理过程看作一个顺序决策问题，其中推理就是逐步构建并完善答案的过程。
具体来说，LLM从输入上下文（初始状态）开始，生成一个推理步骤（动作），并更新上下文（下一个状态）。
LLM会重复这个过程，直到得出最终答案。根据最终答案与真实答案的匹配程度，给予LLM奖励。
通过这种方式，用RL来训练LLM进行推理，旨在让LLM生成一系列推理步骤，以最大化期望奖励。
行动-思维链推理（COAT）
实现自回归搜索时，关键挑战在于让LLM能够在没有外部干预的情况下，判断何时进行反思、继续推理，或是探索替代方案。
为解决这个问题，研究者引入了几种特殊的元行动tokens，来引导LLM的推理过程：

继续推理（）：鼓励LLM依据当前的推理思路，生成下一个中间步骤。
反思（）：提醒模型暂停下来，验证之前的推理步骤是否正确。
探索替代解决方案（）：提示模型识别推理中的关键漏洞，并探索新的解决方案。

这种推理方式称为行动-思维链（COAT）推理。每个COAT推理步骤都是一个tokens序列，并从其中一个元行动tokens开始。
标准LLM无法执行COAT推理，将RL应用于推理面临两个关键挑战：
[ol]
对元动作tokens缺乏认知：如果没有经过训练，LLM在遇到特殊的元动作tokens时，不会意识到需要反思或者寻找替代解决方案。
长期决策与奖励稀疏：推理涉及长期决策，而奖励仅在最终阶段给出。这意味着LLM必须在得到奖励之前，连续做出多个正确的推理步骤，一旦出错，就只能从初始状态重新开始。因为奖励非常稀缺，而奖励对于RL至关重要，这大大增加了学习难度。
[/ol]
一开始，模型对元动作tokens没有认知。为解决这个问题，研究者设置了一个格式调优阶段。
具体做法是，在一个有少量推理轨迹示例的小数据集上对预训练的LLM进行微调。通过这一步，模型就能熟悉元动作tokens的使用，并且做出相应反应。
另外，推理存在决策时间长、奖励少的问题。为解决这个难题，借鉴Go-Explore的思路，提出重启与探索（RAE）策略。
模型会从之前推理过程中的中间步骤重新开始，包括那些推理失败的节点，这样它就能专注于改正错误，而不用每次都从头开始。
同时，还增设了探索奖励，鼓励模型进行更深入的思考，从而提高得出正确答案的可能性。
通过模仿学习进行格式调优
这个阶段的目的是对预训练的基础LLM进行微调，让它能模仿符合COAT推理格式的示范推理轨迹。
为了合成包含试错过程的COAT推理轨迹，研究者提出多代理数据合成框架，通过三个LLM来完成这项任务：

生成器：给定一个输入问题，生成器会运用经典的链式思维（CoT）技术，生成多个推理路径。
Critic：负责评估生成器生成的推理路径是否正确，同时提供反馈以优化推理过程，修正不合理的步骤。
奖励模型：对优化后的推理路径打分，挑选出最有效的路径，作为最终的示范轨迹。

这三个模型相互配合，共同构建出高质量的示范轨迹。仅需10K条示范轨迹，就能让基础LLM学会遵循COAT推理格式。
通过RL进行自我提升
通过格式调优，LLM已经掌握了COAT推理风格，但遇到新问题时，仍然很难泛化。
RL阶段的目标，就是让LLM通过自我反思，提升推理能力。
以完成格式调优的LLM为基础，用经典的PPO算法进一步优化，同时引入两个关键策略：
重启与探索（RAE）：受Go-Explore算法启发，训练LLM时，不仅让它从问题本身出发进行推理，还让它从过去的推理过程中，采样中间步骤来进行推理。
此外，增设了探索奖励，鼓励LLM进行更深入的自我反思，从而增加它找到正确答案的可能性。


迭代自我提升：训练过程中，LLM的策略可能会陷入局部最优解。
借鉴Kickstarting的思路，在每一轮RL训练结束后，通过监督微调，把当前教师策略的知识传递给基础模型。以微调后的LLM为起点，再开展下一轮RL训练。
评估结果
大量实验结果显示，Satori在数学推理基准测试中取得了最佳成绩，在不同领域的任务上也有很强的泛化能力。
研究者选择Qwen-2.5-Math-7B作为基础模型，因为它在数学方面能力很强。训练数据来源于公开的数学指令数据集，包括OpenMathInstruct-2和NuminaMathCoT。
在多智能体数据合成框架中，生成器需生成高质量的逐步推理轨迹，因此选用Qwen-2.5-MathInstruct。而评论者需要有很强的指令跟随能力，于是选了Llama3.1-70B-Instruct。
表中展示了数学基准测试的结果，Satori-Qwen-7B在所有小规模基线模型中表现最佳。
尽管Satori-Qwen-7B使用了与Qwen-2.5-Math-7B-Instruct相同的基础模型，其性能明显优于后者，所需的SFT数据显著减少，并更多依赖于自我改进。


同时在数学领域之外的广泛基准测试上进行了评估，包括逻辑推理（FOLIO、BGQA）、代码推理（CEUXEval）、常识推理（StrategyQA）、表格推理（TableBench）以及特定领域推理（MMLUPro的STEM子集），覆盖物理、化学、计算机科学、工程学、生物学和经济学。
尽管Satori-Qwen-7B只在数学领域的数据集上训练过，但它的推理能力同样适用于其他领域。
表中展示了Satori-Qwen-7B在跨领域基准测试中的表现。


和在数学领域的表现类似，Satori-Qwen-7B在多个基准测试里成绩优异，超过了Qwen-2.5-Math-7B-Instruct。
特别是在难度较高的BoardgameQA推理基准测试中，Satori-Qwen-7B的表现优于所有同规模的基线模型。
这些结果表明，Satori-Qwen-7B不仅掌握了数学解题技能，还具备了通用的推理能力。
最后一行展示了Satori第二轮训练的结果。与Satori-Qwen-7B相比，Satori-Qwen-7B（Round 2）在大多数领域表现出持续的性能提升。
这表明迭代自我改进在提升LLM推理性能方面具有显著的潜力。
Satori展现自我纠错能力
研究者观察到Satori在推理过程中经常自我反思，主要出现这两种情形：一是在推理的中间步骤，二是完成问题后，通过自我反思发起第二次常识。
对第二种情况做定量评估，以衡量Satori的自我纠错能力。
具体做法是，找出那些自我反思前后最终答案不一样的回答，然后计算其中正向（从错误修正为正确）自我纠错或负向（从正确改为错误）的比例。
表中呈现了Satori在领域内数据集（MATH500和Olympiad）以及领域外数据集（MMLUPro）上的评估结果。


与没有经过RL训练阶段的Satori-Qwen-FT相比，Satori-Qwen的自我纠错能力更强。
这种自我纠错能力在领域外任务（MMLUPro-STEM）中同样存在。
这些结果说明，RL对于提升模型实际的推理能力起着关键作用。
RL使Satori具备测试时扩展能力
接下来，讨论RL如何激励Satori进行自回归搜索。
首先，从图中可以看到，随着RL训练计算量的增多，Satori策略的准确率不断上升，同时生成内容的平均token长度也在增加。这表明Satori学会了花更多时间去推理，从而更准确地解决问题。


一个有趣的现象是，响应长度在前0到200步时先减少，然后再增加。
通过深入分析模型的响应，发现在早期阶段，Satori还未学会自我反思能力。
在这个阶段，RL优化可能会先引导模型寻找捷径来解决问题，减少不必要的思考，所以响应长度会暂时变短。
到了后期，模型慢慢学会通过反思来自我纠错，找到更好的解法，因此响应长度随之增加。
此外，研究人员在不同难度的MATH数据集上，对Satori的测试准确率和响应长度做了评估。
经过RL训练，Satori在测试时会自动把更多计算资源，用在解决更难的问题上。与只经过格式调优的模型相比，Satori的性能不断提高。


蒸馏实现从弱到强的泛化能力
最后，我们探究能否借助蒸馏更强的推理模型，提升较弱基础模型的推理能力。
具体做法是，用Satori-Qwen-7B生成24万条合成数据，以此训练Llama-3.1-8B和Granite-3.1-8B这两个基础模型。
作为对比，研究者还合成了24万条格式调优（FT）数据，用于训练同样的两个模型。
之后，在所有数学基准测试数据集上，对这些模型的平均测试准确率进行评估，结果如图所示。


实验表明，经过蒸馏训练的模型，性能比仅经过格式调优的模型更好。
这为提升较弱基础模型的推理能力，提供了一种新的高效方法：
[ol]
通过小规模的格式调优与大规模RL相结合，训练出像Satori-Qwen-7B这样的强推理模型。
运用蒸馏的方式，将这个强推理模型的能力转移到较弱的基础模型中。
[/ol]
由于RL训练只需答案标签作为监督信号，所以这种方法合成数据的成本很低，既不需要多智能体数据合成框架，也无需昂贵的人工标注。
参考资料：
https://x.com/gan_chuang/status/1886990694327238824
https://satori-reasoning.github.io/blog/satori/