ASTRO: Teaching Language Models to Reason by Reflecting and Backtracking In-Context (2507.00417v1)

Abstract: We introduce ASTRO, the "Autoregressive Search-Taught Reasoner", a framework for training LLMs to reason like search algorithms, explicitly leveraging self-reflection, backtracking, and exploration in their outputs. Recently, training 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.

• The paper introduces the ASTRO framework that integrates Monte Carlo Tree Search, supervised fine-tuning, and reinforcement learning to embed search-like reasoning in LLMs.
• Empirical results show ASTRO improves performance on math benchmarks, with up to a 26.9% absolute gain on AMC 2023.
• The study highlights the critical role of in-context self-reflection and backtracking for enhancing complex reasoning in open-source language models.

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

The paper "ASTRO: Teaching LLMs to Reason by Reflecting and Backtracking In-Context" (2507.00417) presents a systematic framework for instilling robust, search-like reasoning capabilities into LLMs, with a particular focus on open-source models such as Llama 3. The core contribution is the ASTRO framework, which operationalizes the principles of search—specifically, self-reflection, backtracking, and exploration—within the autoregressive generation process of LLMs. This is achieved through a combination of synthetic data generation, supervised fine-tuning (SFT), and reinforcement learning (RL) with verifiable rewards.

Methodology

ASTRO is structured as a three-stage pipeline:

1. Search Trajectory Generation: The authors employ Monte Carlo Tree Search (MCTS) to explore the solution space of mathematical problems. Each node in the search tree represents a discrete reasoning step, and the tree is annotated with Q-values derived from verifier-based rewards. The search traces, including both successful and failed attempts, are linearized into natural language chain-of-thoughts (CoTs) that explicitly encode self-reflection and backtracking.
2. Supervised Fine-Tuning (SFT): The LLM is fine-tuned on the synthetic dataset of search-derived CoTs. This stage is designed to bootstrap the model with priors for exploration, reflection, and recovery from failure, even when starting from a base model lacking such behaviors.
3. Reinforcement Learning (RL): The fine-tuned model is further optimized using RL with verifiable rewards, employing a variant of Group Relative Policy Optimization (GRPO). The RL stage leverages the search priors established during SFT, enabling the model to internalize and autonomously execute search-like reasoning during inference.

Empirical Results

The application of ASTRO to the Llama 3 family yields substantial improvements on challenging mathematical reasoning benchmarks:

• MATH-500: +16.0% absolute gain (from 65.8% to 81.8% pass@1)
• AMC 2023: +26.9% absolute gain (from 37.5% to 64.4% pass@1)
• AIME 2024: +20.0% absolute gain (from 10.0% to 30.0% pass@1)

Notably, the ASTRO-trained model (Llama-3.1-70B-ASTRO-RL) outperforms both the Llama-3.3-70B-Instruct baseline and other recent post-training techniques such as Step-KTO and spontaneous self-correction (SPOC), even when those methods are applied to larger or more advanced Llama 3 variants.

Ablation studies demonstrate that the explicit inclusion of self-reflection and backtracking priors is critical: models trained on direct solutions without these priors underperform their ASTRO counterparts, both after SFT and RL. Furthermore, the number of backtracks performed by the model during inference is strongly correlated with evaluation performance, indicating that the search-inspired behaviors are not merely superficial but are functionally beneficial for complex reasoning tasks.

Implementation Considerations

• Data Generation: The MCTS-based data generation is computationally intensive, requiring multiple rollouts per node and careful reward assignment via external verifiers. The authors report generating 20.7K search trees and curating 105K CoT solutions, with a final SFT dataset of 36.1K high-quality examples.
• Model Training: SFT is performed for a single epoch to avoid overfitting, using the AdamW optimizer and a maximum sequence length of 8,192 tokens. RL is conducted with a constant learning rate, large batch sizes, and sequence lengths up to 15,360 tokens. Training is distributed across 128 NVIDIA H100 GPUs for both SFT and RL, with RL runs taking approximately 10 days.
• Inference: The ASTRO-trained models generate significantly longer CoTs (up to 6,000 tokens on average during RL), reflecting their iterative, search-like reasoning process. This has implications for inference cost and latency, particularly in production settings.

Theoretical and Practical Implications

The ASTRO framework provides a principled approach for endowing LLMs with structured, search-based reasoning capabilities, independent of pre-existing reflective behaviors. By internalizing the search process, models trained with ASTRO can autonomously explore, reflect, and backtrack within a single inference pass, obviating the need for external scaffolding or post-hoc self-correction mechanisms.

From a theoretical perspective, ASTRO bridges the gap between algorithmic search and neural sequence modeling, demonstrating that LLMs can be taught to emulate search algorithms in natural language. The positive correlation between backtracking frequency and task performance suggests that explicit modeling of exploration and recovery is beneficial for complex, multi-step reasoning tasks.

Practically, ASTRO offers a scalable recipe for improving the reasoning robustness of open-source LLMs, particularly in domains where verifiable rewards are available (e.g., mathematics, code synthesis). The framework is compatible with existing RLHF pipelines and can be integrated into broader post-training regimes.

Future Directions

Several avenues for future research are suggested by this work:

• Generalization Beyond Mathematics: While the current instantiation of ASTRO is focused on mathematical reasoning, the framework is applicable to any domain where solution trajectories can be verified and search traces can be constructed (e.g., program synthesis, scientific discovery).
• Efficiency and Scalability: The computational demands of MCTS-based data generation and long-sequence RL highlight the need for more efficient search and training algorithms, as well as methods for distilling search behaviors into smaller or more efficient models.
• Interpretability: The mapping of CoTs to directed graphs of reasoning steps offers opportunities for enhanced interpretability and debugging of LLM outputs, potentially enabling more transparent and controllable AI systems.
• Integration with External Tools: Combining ASTRO-trained models with external symbolic solvers or verification tools could further enhance their reliability and applicability in high-stakes domains.

In summary, ASTRO establishes a robust methodology for teaching LLMs to reason via in-context reflection and backtracking, yielding strong empirical gains and providing a foundation for future advances in neural reasoning systems.