ASTRO Framework Boosts Llama 3’s Math Reasoning by 16% to 20% Without Architectural Changes

Improving the reasoning capabilities of large language models (LLMs) without altering their underlying architecture is a fundamental challenge in AI development. Recently, researchers from Meta AI and the University of Washington have introduced ASTRO—Autoregressive Search-Taught Reasoner—a novel post-training framework designed to enhance reasoning in the Llama-3.1-70B-Instruct model. ASTRO boosts Llama 3's performance on various benchmarks by teaching the model to perform in-context search, self-reflection, and backtracking, behaviors that are typically associated with human problem-solving and symbolic search algorithms. Search-Guided Chain-of-Thought Generation ASTRO's methodology starts with a Monte Carlo Tree Search (MCTS) over mathematical problem-solving trajectories. This search explores both correct and incorrect reasoning paths, a process that is essential for improving the model's understanding of potential pitfalls and how to recover from them. The key innovation is procedure cloning: the entire search tree is converted into long chain-of-thought (CoT) sequences, which naturally include both failures and successful recoveries. These sequences are then rewritten in natural language and used as training data for supervised fine-tuning (SFT). Supervised Fine-Tuning: Injecting Search Priors The researchers fine-tuned Llama-3.1-70B-Instruct using 36,100 curated CoT solutions from datasets like MATH, AMC/AIME, and AoPS. The result is a model that doesn't just solve problems step-by-step but also reassesses its reasoning and backtracks to correct mistakes. This approach leads to significant performance gains: On the MATH dataset, the ASTRO-SFT model achieved a 16% improvement over the baseline. On the AMC/AIME dataset, there was a 17% improvement. On the AoPS dataset, the improvement was 20%. These scores are competitive with or exceed those of models trained without explicit search priors. Remarkably, even supervised fine-tuning alone—without reinforcement learning—produces notable performance enhancements by exposing the model to structured reasoning data. Reinforcement Learning with Search-Aware Initialization ASTRO then employs reinforcement learning (RL) to further refine the model. Starting from the SFT checkpoint, the team uses a modified Group Relative Policy Optimization (GRPO) method on 8,700 moderately difficult prompts, employing verifiable reward signals (+1 for correct answers, -1 for incorrect ones). During training, the model's chain-of-thought generation becomes more extensive, increasing from about 1,800 tokens to around 6,000 tokens, indicative of deeper internal exploration and reasoning. The resulting ASTRO-RL model shows remarkable performance improvements: On the MATH dataset, ASTRO-RL achieved a 22% improvement over the baseline. On the AMC/AIME dataset, the improvement was 23%. On the AoPS dataset, it reached a 25% improvement. These results rival or surpass models with larger parameter counts, underscoring the effectiveness of ASTRO's search-aware initialization. Backtracking Behavior Correlates with Reasoning Success One of the most compelling findings is the strong positive correlation between the frequency of backtracking and performance. As training progresses, ASTRO-RL performs more self-corrective actions and explores reasoning more deeply. Across various benchmarks, Pearson correlation coefficients exceeded 0.8, indicating that backtracking and self-reflection are not just superficial behaviors but are functionally linked to higher accuracy. Comparative Insights and Broader Impact Control experiments revealed that ASTRO consistently outperforms models trained on direct CoT solutions without search priors, even when they are exposed to the same problem sets and search trees. For example, ASTRO-RL surpasses Direct-RL by: 6% on the MATH dataset. 7% on the AMC/AIME dataset. 8% on the AoPS dataset. Furthermore, ASTRO's reasoning processes can be visualized as directed graphs, where nodes represent reasoning steps and edges capture transitions, self-reflections, and corrections. This visualization enhances the interpretability of the model's decision-making process. Conclusion ASTRO provides a groundbreaking approach to enhancing the reasoning abilities of LLMs like Llama 3. By teaching models to think more critically, doubt their own steps, and correct themselves, ASTRO achieves substantial gains without requiring larger models or extended pretraining. This framework sets a new standard for fine-tuning open LLMs to exhibit more human-like reasoning through search-inspired behaviors. For more details, check out the paper. All credit goes to the researchers who conducted this study. Stay updated by following us on Twitter, and don't miss out on joining our ML SubReddit and subscribing to our Newsletter.