MIT等机构提出COAT，使LLM具备自回归搜索能力，通过强化学习和自我反思显著提升推理能力，并在数学和跨领域任务中表现出色。
原文标题：LLM实现自回归搜索！MIT哈佛等提出「行动思维链」COAT，推理能力大提升
原文作者：数据派THU
冷月清谈：
怜星夜思：
2、文章提到COAT方法借鉴了Go-Explore算法的重启与探索（RAE）策略，那么这种策略在提升LLM推理能力上的具体优势是什么？为什么从失败的节点重新开始反而有助于找到正确答案？
3、文章最后提到了通过知识蒸馏将Satori的推理能力迁移到较弱的基础模型，这种方法对于LLM的应用和发展有哪些潜在的价值和意义？能否进一步降低LLM的使用门槛？
原文内容
OpenAI o1发布后，为提升LLM的推理能力，研究者尝试了多种方法。
比如用强大的教师模型进行知识蒸馏、采用蒙特卡洛树搜索（MCTS），以及基于奖励模型的引导搜索。近日，来自MIT、新加坡科技设计大学、哈佛大学等机构的华人研究者探索了全新的方向：让LLM拥有自回归搜索能力。通过自我反思和探索新策略，提升LLM推理能力。研究者引入了行动-思维链（COAT）机制，使LLM在解决问题时能够执行多种元动作，并提出了一种创新的两阶段训练框架：
-
小规模格式调优阶段：让LLM熟悉并掌握COAT推理格式。
-
大规模自我优化阶段：运用重启与探索（RAE）技术，通过RL进行优化。
通过这种方法，成功开发出Satori，在数学推理任务中，成绩优异。Satori具有以下核心特点：
-
无需外部指导，即可自我反思与探索。
-
主要依靠自我改进（RL），实现了最先进的推理性能。
-
展现出强大的迁移能力，可应用于数学以外的领域。
Satori关键设计
研究者把LLM的推理过程看作一个顺序决策问题，其中推理就是逐步构建并完善答案的过程。具体来说，LLM从输入上下文（初始状态）开始，生成一个推理步骤（动作），并更新上下文（下一个状态）。LLM会重复这个过程，直到得出最终答案。根据最终答案与真实答案的匹配程度，给予LLM奖励。通过这种方式，用RL来训练LLM进行推理，旨在让LLM生成一系列推理步骤，以最大化期望奖励。
行动-思维链推理（COAT）
实现自回归搜索时，关键挑战在于让LLM能够在没有外部干预的情况下，判断何时进行反思、继续推理，或是探索替代方案。为解决这个问题，研究者引入了几种特殊的元行动tokens，来引导LLM的推理过程：
-
继续推理（<|continue|>）：鼓励LLM依据当前的推理思路，生成下一个中间步骤。
-
反思（<|reflect|>）：提醒模型暂停下来，验证之前的推理步骤是否正确。
-
探索替代解决方案（<|explore|>）：提示模型识别推理中的关键漏洞，并探索新的解决方案。
这种推理方式称为行动-思维链（COAT）推理。每个COAT推理步骤都是一个tokens序列，并从其中一个元行动tokens开始。标准LLM无法执行COAT推理，将RL应用于推理面临两个关键挑战：
-
对元动作tokens缺乏认知：如果没有经过训练，LLM在遇到特殊的元动作tokens时，不会意识到需要反思或者寻找替代解决方案。
-
长期决策与奖励稀疏：推理涉及长期决策，而奖励仅在最终阶段给出。这意味着LLM必须在得到奖励之前，连续做出多个正确的推理步骤，一旦出错，就只能从初始状态重新开始。因为奖励非常稀缺，而奖励对于RL至关重要，这大大增加了学习难度。
一开始，模型对元动作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）数据，用于训练同样的两个模型。之后，在所有数学基准测试数据集上，对这些模型的平均测试准确率进行评估，结果如图所示。
实验表明，经过蒸馏训练的模型，性能比仅经过格式调优的模型更好。这为提升较弱基础模型的推理能力，提供了一种新的高效方法：
