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

📄 arXiv: 2507.00417v1 📥 PDF

作者: Joongwon Kim, Anirudh Goyal, Liang Tan, Hannaneh Hajishirzi, Srinivasan Iyer, Tianlu Wang

分类: cs.AI, cs.CL

发布日期: 2025-07-01

备注: 36 pages, 23 figures

💡 一句话要点

ASTRO：通过上下文反思与回溯，教导语言模型进行推理

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 语言模型推理 强化学习 蒙特卡洛树搜索 思维链 自我反思 回溯 数学问题求解

📋 核心要点

  1. 现有方法难以有效提升非推理型LLM的推理能力，尤其是在复杂问题上。
  2. ASTRO通过模拟蒙特卡洛树搜索过程，生成包含反思和回溯的训练数据，引导模型学习结构化搜索行为。
  3. 实验表明，ASTRO显著提升了Llama 3在数学问题上的推理性能，尤其是在需要迭代修正的难题上。

📝 摘要（中文）

本文介绍ASTRO，即“自回归搜索训练推理器”，一个训练语言模型像搜索算法一样推理的框架，它在输出中显式地利用了自我反思、回溯和探索。最近，通过强化学习(RL)训练大型语言模型(LLM)已经带来了推理模型的出现，这些模型具有大大增强的推理能力。开源的推理模型复现虽然成功，但建立在已经表现出强大推理能力以及在RL之前就观察到的搜索行为的模型之上。因此，如何提升包括Llama 3在内的其他非推理模型的推理能力尚不清楚。ASTRO通过从数学问题解决轨迹上的蒙特卡洛树搜索(MCTS)导出的合成数据集，教导这些模型内化结构化搜索行为。通过将搜索轨迹转换为自然语言的思维链，捕捉成功和从失败中恢复的过程，ASTRO为模型提供了丰富的先验知识，以便在RL期间进行探索。我们在这些搜索衍生的轨迹上微调我们的模型，并通过具有可验证奖励的RL进一步提高性能。我们将ASTRO应用于Llama 3系列模型，并在MATH-500上实现了16.0%的绝对性能提升，在AMC 2023上实现了26.9%的提升，在AIME 2024上实现了20.0%的提升，尤其是在需要迭代校正的具有挑战性的问题上。我们的结果表明，受搜索启发的训练提供了一种原则性的方法，可以将强大的推理能力灌输到开放LLM中。

🔬 方法详解

问题定义：论文旨在解决如何提升非推理型大型语言模型（如Llama 3）的推理能力，尤其是在解决需要迭代修正的复杂数学问题时。现有方法，如直接使用强化学习，依赖于已经具备一定推理能力的模型，无法有效提升其他模型的推理能力。

核心思路：论文的核心思路是模仿蒙特卡洛树搜索（MCTS）的搜索过程，生成包含成功和失败案例的训练数据，让模型学习如何在推理过程中进行自我反思、回溯和探索。通过将搜索轨迹转化为自然语言的思维链，模型可以学习到结构化的搜索行为，从而提升推理能力。

技术框架：ASTRO框架包含以下几个主要阶段：1) 使用MCTS在数学问题上生成搜索轨迹；2) 将搜索轨迹转化为自然语言的思维链，构建合成数据集；3) 在合成数据集上微调语言模型；4) 使用强化学习和可验证的奖励进一步提升模型性能。整体流程是从搜索算法中提取知识，然后通过监督学习和强化学习将这些知识迁移到语言模型中。

关键创新：ASTRO的关键创新在于它使用搜索算法（MCTS）生成训练数据，从而让模型学习到结构化的搜索行为。与以往依赖人工标注或直接使用强化学习的方法不同，ASTRO提供了一种更有效、更可控的方式来提升语言模型的推理能力。此外，将搜索轨迹转化为自然语言的思维链也是一个重要的创新，它使得模型能够更好地理解和学习搜索过程。

关键设计：ASTRO的关键设计包括：1) 使用MCTS生成高质量的搜索轨迹，确保训练数据的多样性和覆盖性；2) 设计合适的自然语言模板，将搜索轨迹转化为易于理解的思维链；3) 使用可验证的奖励函数，引导强化学习过程，确保模型能够学习到正确的推理策略。具体的参数设置和网络结构细节可能需要根据具体的模型和任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ASTRO在Llama 3系列模型上取得了显著的性能提升。在MATH-500数据集上，ASTRO实现了16.0%的绝对性能提升；在AMC 2023上，提升了26.9%；在AIME 2024上，提升了20.0%。这些结果表明，ASTRO能够有效提升语言模型在复杂数学问题上的推理能力，尤其是在需要迭代修正的难题上。

🎯 应用场景

ASTRO框架可应用于提升各种语言模型的推理能力，尤其是在需要复杂推理和迭代修正的领域，如数学问题求解、代码生成、逻辑推理等。该研究有助于构建更智能、更可靠的AI系统，并可能在教育、科研和工程等领域产生广泛影响。

📄 摘要（原文）

We introduce ASTRO, the "Autoregressive Search-Taught Reasoner", a framework for training language models to reason like search algorithms, explicitly leveraging self-reflection, backtracking, and exploration in their outputs. Recently, training large language models (LLMs) via reinforcement learning (RL) has led to the advent of reasoning models with greatly enhanced reasoning capabilities. Open-source replications of reasoning models, while successful, build upon models that already exhibit strong reasoning capabilities along with search behavior observed even before RL. As a result, it is yet unclear how to boost the reasoning capabilities of other non-reasoner models including Llama 3. ASTRO teaches such models to internalize structured search behavior through a synthetic dataset derived from Monte Carlo Tree Search (MCTS) over mathematical problem-solving trajectories. By converting search traces into natural language chain-of-thoughts that capture both successes and recoveries from failure, ASTRO bootstraps models with a rich prior for exploration during RL. We finetune our models on these search-derived traces and further improve performance via RL with verifiable rewards. We apply ASTRO to the Llama 3 family of models and achieve absolute performance gains of 16.0% on MATH-500, 26.9% on AMC 2023, and 20.0% on AIME 2024, especially improving upon challenging problems that require iterative correction. Our results demonstrate that search-inspired training offers a principled way to instill robust reasoning capabilities into open LLMs.