ASTRO: Teaching Language Models to Reason by Reflecting and Backtracking In-Context

我们提出了ASTRO，即“自回归搜索教学推理器”，这是一种训练语言模型的框架，使其像搜索算法一样进行推理，在输出中明确利用自我反思、回溯和探索。近期，通过强化学习（RL）训练大语言模型（LLM）已经催生出推理能力大幅提升的推理模型。尽管开源的推理模型复制工作取得了成功，但这些模型所基于的基础模型本身就已经具备较强的推理能力和在强化学习之前就已展现出的搜索行为。因此，如何提升包括Llama 3在内的其他非推理模型的推理能力仍不清楚。ASTRO通过一个从蒙特卡洛树搜索（MCTS）生成的数学问题解决路径中构建的合成数据集，教会这类模型内化结构化的搜索行为。通过将搜索轨迹转化为自然语言的思维链，既捕捉成功的过程，也记录从失败中的恢复过程，ASTRO为强化学习中的模型提供了丰富的探索先验知识。我们在这些由搜索生成的轨迹上对模型进行微调，并通过可验证奖励的强化学习进一步提升了性能。我们将ASTRO应用于Llama 3系列模型，在MATH-500上实现了16.0%的绝对性能提升，在AMC 2023上提升了26.9%，在AIME 2024上提升了20.0%，尤其在那些需要迭代修正的难题上表现突出。我们的结果表明，受搜索启发的训练方法为向开源大语言模型注入强大的推理能力提供了一种有原则的途径。

论文旨在解决如何提升基础语言模型（如Llama 3）的推理能力这一问题，特别是通过模仿搜索算法的行为（如自我反思、回溯和探索）来增强模型在数学问题求解任务中的表现。该问题在当前研究中具有一定挑战性，尤其是在不依赖已有强推理能力模型的前提下。

ASTRO框架通过利用蒙特卡洛树搜索（MCTS）生成的合成数据集，将搜索轨迹转换为自然语言形式的思维链（chain-of-thought），从而训练语言模型内部形成结构化的搜索行为。这种方法不同于传统强化学习仅依赖外部奖励信号，而是通过模拟搜索过程，使模型学会主动探索与错误恢复，显著提升了其推理能力。

1. ASTRO将MCTS生成的数学问题解决路径转化为自然语言推理路径，用于训练语言模型进行结构化思考。 2. 实验结果显示，在MATH-500、AMC 2023和AIME 2024等数学推理任务上分别取得了16.0%、26.9%和20.0%的绝对性能提升。 3. 特别对需要迭代修正的复杂问题效果显著，表明模型具备更强的错误恢复和探索能力。 4. ASTRO提供了一种基于搜索启发的通用训练范式，适用于其他非推理型语言模型的增强训练。 5. 论文可能开源代码或数据集，有助于推动开放大模型的推理能力研究。

1. AlphaGeometry: Solving Geometry Problems via Language Models and Symbolic Reasoning 2. Chain-of-Thought Prompting Elicits Reasoning in Large Language Models (CoT) 3. Training Verifiable Reasoners via Selective Backtracking 4. Self-Taught Reasoner (STR): Learning to Reason from Scratch through Self-Play 5. Monte Carlo Tree Search for Policy Search and Reinforcement Learning in Continuous Spaces