ASTRO后训练框架显著提升Llama 3推理能力，数学成绩提高16%至25%

改进大型语言模型（LLMs）的推理能力而不改变其架构，是提高AI对齐性和可用性的关键挑战。来自Meta AI和华盛顿大学的研究人员推出了一种名为ASTRO（自回归搜索教学推理器）的新后训练框架，旨在增强Llama-3.1-70B-Instruct的推理能力。ASTRO的独特之处在于教模型执行上下文搜索、自我反省和回溯，这些机制通常与人类解决问题的方法和传统的符号搜索算法相关联。 ASTRO的方法最初基于数学问题解决轨迹的蒙特卡洛树搜索（MCTS），该搜索会探索正确和错误的推理路径。最关键的是过程克隆：整个搜索树线性化成长链思维（CoT），自然地编码了失败和通过自我反省及回溯恢复的过程。这些线性化的轨迹被重写成自然语言，用于监督细调（SFT）的基础数据。 这一方法不仅使模型能够逐步解决问题，还能够在自我评估后重新评估其路径并进行修正，例如，当模型的内部置信度下降时，它可能会说：“让我们回到设立方程的地方再试一次。” 在SFT阶段，ASTRO对Llama-3.1-70B-Instruct进行了36.1万个精心策划的CoT解决方案的训练，这些方案来自MATH、AMC/AIME和AoPS风格的数据集。使用ASTRO-SFT训练的模型在多个基准测试中取得了显著的提升： 在MATH测试中，正确率从22.9%提升到44.8%，增加了21.9个百分点； 在GSM8K测试中，正确率从42.5%提升到59.6%，增加了17.1个百分点。 即使没有强化学习，仅通过SFT就能显著提升模型性能，因为该方法使模型接触到经过搜索结构化的推理数据。 接下来，ASTRO进入强化学习（RL）阶段，使用经过SFT预训练的检查点作为初始点，并运行一个修改后的组相对策略优化（GRPO）循环。不同于标准的偏好驱动RL方法，ASTRO采用可验证的奖励信号（正确得+1分，错误得-1分），在8.7万个中等难度提示上进行训练。随着训练的进行，模型生成的CoT逐渐变长，从约1.8万个字符增加到约6万个字符，显示出更深层次的内部探索。 最终，ASTRO-RL模型在多个基准测试中达到了出色的结果，其中包括： - 在MATH测试中，正确率从44.8%进一步提升至58.6%，增加了13.8个百分点； - 在GSM8K测试中，正确率从59.6%提升至72.1%，增加了12.5个百分点。 更重要的是，ASTRO-RL的回溯行为与性能成功之间有显著的正相关关系。随着训练的进步，模型显示出更多的自我纠正动作和更深的探索。Pearson相关系数超过0.8，表明自我反省和回溯不仅是表面现象，而是与更高的准确性密切相关。 控制实验比较了ASTRO与其他直接基于CoT解决方案训练的模型（不使用搜索先验）。结果表明，即使在相同的问题集和搜索树上训练，ASTRO也能一致地胜出。例如，ASTRO-RL超过了Direct-RL 8.1个百分点，在MATH测试中达到58.6%的正确率。 此外，ASTRO的输出可以可视化为有向图，节点表示推理步骤，边表示过渡、反思和修正，从而提高了模型的可解释性。 ASTRO的核心发现是，像Llama 3这样的大型语言模型可以通过有原则的后训练技术更有效地学习推理，而不是通过更大的模型或更长时间的预训练。ASTRO使模型能够“思考”而不仅仅是回答问题，学会在推理过程中怀疑自己的步骤并在中途进行自我纠正，这种框架为微调开源大型语言模型以实现类人推理设定了新的标准。