PaperlayerView exact revision on arXiv

Source-linked AI summary

rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking

Xinyu Guan, Li Lyna Zhang, Yifei Liu, Ning Shang, Youran Sun, Yi Zhu, Fan Yang, Mao Yang

arXiv:2501.04519v1cs.CL

TL;DR

Small language models often lack reliable reasoning data and process-level supervision for difficult mathematics. rStar-Math uses MCTS with code-verified trajectories, Q-value-based process preferences, and four rounds of self-evolution to train and guide these models. It reaches state-of-the-art results, including 90.0% for Qwen2.5-Math-7B on MATH and 53.3% on AIME 2024.

  • Problem

    High-quality math reasoning data is scarce, intermediate steps can be incorrect even when final answers are right, and reliable process reward models remain difficult to train.

  • Method

    rStar-Math combines MCTS-guided test-time search, code-augmented verified trajectories, Q-value-based preference training for a process preference model, and four rounds of self-evolution.

  • Results

    90.0%: rStar-Math raises Qwen2.5-Math-7B from 58.8% on MATH and achieves 53.3% (8/15) on AIME 2024.

  • Takeaways & Limitations

    rStar-Math demonstrates that small language models can achieve state-of-the-art, OpenAI o1-level math reasoning and can self-generate high-quality training data.

  • Takeaways & Limitations

    Small language models frequently fail on difficult problems, produce flawed intermediate steps, and solve fewer challenging problems than advanced models such as GPT-4.

Abstract

from arXiv · show

We present rStar-Math to demonstrate that small language models (SLMs) can rival or even surpass the math reasoning capability of OpenAI o1, without distillation from superior models. rStar-Math achieves this by exercising "deep thinking" through Monte Carlo Tree Search (MCTS), where a math policy SLM performs test-time search guided by an SLM-based process reward model. rStar-Math introduces three innovations to tackle the challenges in training the two SLMs: (1) a novel code-augmented CoT data sythesis method, which performs extensive MCTS rollouts to generate step-by-step verified reasoning trajectories used to train the policy SLM; (2) a novel process reward model training method that avoids naïve step-level score annotation, yielding a more effective process preference model (PPM); (3) a self-evolution recipe in which the policy SLM and PPM are built from scratch and iteratively evolved to improve reasoning capabilities. Through 4 rounds of self-evolution with millions of synthesized solutions for 747k math problems, rStar-Math boosts SLMs' math reasoning to state-of-the-art levels. On the MATH benchmark, it improves Qwen2.5-Math-7B from 58.8% to 90.0% and Phi3-mini-3.8B from 41.4% to 86.4%, surpassing o1-preview by +4.5% and +0.9%. On the USA Math Olympiad (AIME), rStar-Math solves an average of 53.3% (8/15) of problems, ranking among the top 20% the brightest high school math students. Code and data will be available at https://github.com/microsoft/rStar.

1 Introduction

rStar-Math addresses the difficulty of obtaining reliable step-level math reasoning data by combining MCTS, code verification, process preference modeling, and iterative self-evolution. Across challenging benchmarks, this approach substantially improves small language models and can match or surpass OpenAI o1.

  • Approach: rStar-Math uses smaller language models with Monte Carlo Tree Search to create a self-evolutionary System 2 reasoning process without relying on superior models for data synthesis.The policy model searches solution trajectories while a process-oriented reward model guides selection.
  • Approach: Code-augmented MCTS generates step-by-step verified reasoning trajectories by retaining nodes whose Python code executes successfully and assigning Q-values from rollout outcomes.Steps contributing to more correct trajectories receive higher Q-values.
  • Approach: A process preference model learns to distinguish correct from incorrect reasoning steps through Q-value-based preference pairs rather than noisy direct step-level scores.Pairwise ranking loss trains the model to predict reliable process rewards.
  • Self-evolution: Four self-evolution rounds progressively strengthen the policy model and process preference model while expanding training-data coverage from a 747k-problem dataset.Each round generates better trajectories and targets increasingly difficult, competition-level problems.
  • Results: 90.0%: With 64 search trajectories, rStar-Math raises Qwen2.5-Math-7B from 58.8% on MATH and surpasses o1-preview by 4.5%.On AIME 2024, it solves 53.3% (8/15) of problems, exceeding o1-preview by 8.7%.

2 Related Works

Prior math-reasoning work combines distilled chain-of-thought data, test-time compute scaling, and reward models, but each direction has important constraints. Distillation is bounded by the teacher, while reliable process-level supervision remains difficult to obtain.

  • Math Data Synthesis: GPT-distilled chain-of-thought datasets have driven many advances, but they limit reasoning to the capabilities of the teacher model.Hard problems beyond the teacher’s ability may be excluded from the synthesized data.
  • Scaling Test-time Compute: Test-time compute scaling improves reasoning by generating multiple samples and using reward models to select solutions, with proposed methods including random sampling and MCTS.These approaches shift computation from a single inference toward search and selection.
  • Reward Models: Reward models support System 2 reasoning, but process reward models remain difficult to train because collecting step-level annotations is laborious and challenging.The literature includes outcome reward models, process reward models, and LLM-based verification.

3 Methodology

rStar-Math combines MCTS with code-verified step-by-step reasoning, PPM-guided search, and iterative self-evolution to train small math models despite unreliable per-step scoring and weak self-generated data.

  • MCTS for System 2 Reasoning: MCTS decomposes complex problems into single-step generation and produces step-level training data for both the policy model and reward model.This reduces the policy SLM’s generation burden compared with methods that generate complete solutions in one inference.
  • Recipe for Self-Evolution: Four self-evolution rounds repeatedly use MCTS-generated verified trajectories to train stronger policy SLMs and PPMs, progressively expanding coverage to harder problems.After four rounds, 90.25% of 747k math problems were correctly solved and included in the training set; the remaining problems were judged largely low quality after manual review.
  • Step-by-Step Verified Reasoning Trajectory: Code-augmented CoT retains only candidate reasoning steps whose accumulated Python code executes successfully, filtering erroneous intermediate steps.The policy model generates one-step natural-language reasoning alongside corresponding Python code, with prior code concatenated at later steps.
  • Q-Value Annotation: Extensive MCTS rollouts annotate intermediate steps with Q-values propagated from terminal outcomes, where correct terminal answers receive 1 and incorrect answers receive −1.Repeated rollouts improve the reliability of Q-values for identifying reasoning paths that lead to correct answers.
  • Process Preference Model: The PPM is trained from Q-value-derived positive-negative preference pairs rather than precise direct step scores, using high-Q correct steps and low-Q incorrect steps.This addresses the difficulty and inconsistency of assigning fine-grained numerical scores to reasoning steps, including at scale for expert annotators.

4 Evaluation

rStar-Math is evaluated across diverse mathematical benchmarks, model sizes, and inference settings, showing strong gains over base models and competitive performance against frontier reasoning systems. Ablations indicate that self-evolution, verified trajectories, and process preference modeling contribute to these results.

  • 4.1 Setup: rStar-Math is evaluated on competition-, Olympiad-, college-, and out-of-domain mathematics benchmarks using SLM policy models from 1.5B to 7B parameters.The benchmark suite includes GSM8K, MATH-500, AIME 2024, AMC 2023, Olympiad Bench, College Math, and GaoKao En 2023.
  • 4.2 Main Results: 90.0% on MATH, up from 58.8% for Qwen2.5-Math-7B, matches o1-mini and exceeds o1-preview and Claude 3.5 Sonnet.With 64 trajectories, Qwen2.5-Math-7B reaches 90.0% Pass@1 accuracy.
  • 4.2 Main Results: 53.3% (8/15) on AIME 2024 places the 7B model among the top 20% of the brightest high school math students.The model solved 8/15 problems in both AIME I and II; eight unsolved problems were geometry-based, requiring visual understanding that rStar-Math does not support.
  • 4.2 Main Results: With only 4 trajectories, rStar-Math exceeds o1-preview and approaches o1-mini, while increasing test-time compute improves accuracy across benchmarks.Performance saturates or improves slowly at 64 trajectories on Math, AIME, and Olympiad Bench, but continues improving steadily on College Math.
  • 4.3 Ablation Study and Analysis: Round 2 onward, rStar-Math surpasses GPT-4o as progressively stronger policy models and PPMs improve System 2 reasoning.Round 1 mainly benefits from supervised fine-tuning, while later rounds gain from stronger PPM-guided MCTS and improved trajectories.
  • 4.3 Ablation Study and Analysis: Step-by-step verified trajectories outperform other SFT datasets, while PPM and PQM outperform ORM through denser step-level reward signals.PPM avoids relying on imprecise Q-values as direct scores and enables the 7B policy model to achieve comparable or superior performance to o1-mini across benchmarks.

5 Findings and Discussions

The experiments reveal intrinsic self-reflection in MCTS-driven reasoning, identify the PPM as a key performance determinant, and show potential generalization beyond math word problems.

  • Intrinsic self-reflection: MCTS-driven deep thinking can recognize low-quality reasoning, backtrack, and solve a problem through a simpler correct approach without self-reflection training data.The observed behavior suggests that advanced System 2 reasoning can foster intrinsic self-reflection.
  • PPM shapes the reasoning boundary: Reward models primarily determine final System 2 performance once the policy model reaches a reasonably strong capability level.Across policy models with different sizes and Pass@1 accuracies, final reasoning accuracy converges after applying the PPM.
  • PPM shapes the reasoning boundary: A 7B PPM enables rStar-Math to outperform a 72B policy model paired with a 72B outcome reward model, despite lower SFT accuracy for the 7B policy.This comparison illustrates the PPM’s role in setting the upper performance limit of System 2 reasoning.
  • PPM spots theorem-application steps: The PPM assigns high reward scores to critical theorem-application steps, guiding the policy model toward relevant mathematical reasoning.Examples include Fermat’s little theorem, Vieta’s formulas, AM-GM, the Pythagorean theorem, and the Shoelace Theorem.
  • Generalization discussions: rStar-Math’s methodology may generalize to theorem proving, code reasoning, and commonsense reasoning when suitable rollout feedback mechanisms are available.The current focus is on word problems because of dataset limitations, while other domains require feedback such as tests, human labels, or mutual verification.

6 Conclusion

The paper presents rStar-Math as a self-evolved System 2 approach that boosts small LLMs to state-of-the-art math reasoning performance. Experiments across model sizes and challenging benchmarks also reveal intrinsic self-reflection and effective PPM guidance.

  • Conclusion: rStar-Math significantly boosts small LLMs’ math reasoning capabilities, achieving state-of-the-art, OpenAI o1-level performance.The evaluation spans four differently sized SLMs and challenging math benchmarks.
  • Conclusion: The approach demonstrates that SLMs can self-generate high-quality training data for frontier-level math reasoning.The paper also identifies intrinsic self-reflection and PPM recognition of critical intermediate steps as key findings.
  • Conclusion: Further improvement may come from collecting more challenging math problems, which the paper leaves as future work.

A.1 Additional Experiments and Details

Additional experiments examine training, inference costs, Pass@N behavior, and the relationship between policy models and PPM-augmented MCTS. They also document implementation choices across self-evolution rounds.

  • Data Generation Details: Each self-evolution round begins by generating step-by-step verified trajectories for 747k math word problems.The default setup uses depth 16, 16 MCTS rollouts per problem, 8 candidate nodes per step, and exploration constant c=2.
  • Training Details: Synthetic problems whose trajectories achieve less than 50% accuracy are removed to reduce noise, and the policy model is fine-tuned from the initial base model each round.Training also collects verified trajectories for both policy fine-tuning and PPM training.
  • Training the ORM and PQM: ORM training selects trajectories with the highest average Q-values as positive examples and those with the lowest average Q-values as negative examples.The ORM and Q-value-based PQM share the PPM’s architecture and training parameters.
  • Self-evolution Inference Costs: The bootstrap data-generation round takes approximately two weeks, while rounds 2–4 each take three days with a fine-tuned 7B SLM policy model.The bootstrap uses a 236B policy model on 10 nodes of 8×80GB H100 GPUs; later rounds use 15 nodes of 4×40GB A100 GPUs.
  • Inference Setting: Inference generates 32 candidate nodes per step, scores them with the PPM, uses 4 rollouts per step to update Q-values, and selects the highest-scoring trajectory.Table 9 reports the average number of generated tokens required per question.
  • Pass@N: Under the same PPM guidance, four policy models of varying sizes show convergent capabilities in sampling correct solutions.Figure 6 separately compares random-sampling Pass@N accuracy across policy models.
  • Pass@N: Qwen2.5-Math-7B-Instruct achieves 4.2% higher MATH accuracy than the Qwen2.5-Math-7B base policy under Pass@1, but Pass@1 alone does not fully reflect System 2 capability.Figures 6 and 7 therefore compare Pass@N under random sampling and PPM-augmented MCTS.

A.2 Examples of Reasoning Trajectories

The appendix presents code-augmented reasoning trajectories for coordinate geometry, number theory, polynomial roots, and area computation. These examples combine symbolic or numerical code with intermediate mathematical steps and final verified answers.

  • Parallelogram coordinates: The coordinate-geometry trajectory uses Python code to print the final value of x+y after self-correcting the parallelogram construction.The displayed output is 17.
  • Parallelogram coordinates: The parallelogram example applies vector differences, corrects an initial coordinate choice, and obtains x=8 under the condition x>7.The trajectory’s final answer reports x+y=17.
  • Vieta’s formulas and AM-GM: The polynomial example uses Vieta’s formulas to relate roots and coefficients, then applies AM-GM and verifies that the minimum value is a=3.The trajectory checks attainability using three equal roots r1=r2=r3=−1 and prints 3.
  • Shoelace Theorem: The Shoelace Theorem trajectory assigns coordinates to D, B, E, and F, computes the quadrilateral area, and outputs 8.The area formula is evaluated from the listed coordinates.

A.3 Prompt Examples

The prompt examples specify code-augmented, step-by-step mathematical problem solving, with executable Python code, printed outputs, and concise final answers. Demonstrations illustrate generating and verifying solutions through intermediate calculations and symbolic solving.

  • The required solution format contains code, output, and answer sections, with the final numerical answer presented concisely and boxed.
  • The code-generation instructions require executable Python, concise reasoning in comments, imported packages, and avoidance of common syntax and variable errors.
  • A weight-lifting example computes total weight, weight per repetition, and required repetitions in sequential Python steps, producing 16.0.
  • A symbolic-equation example uses SymPy to define and solve an equation, printing the solution [20/3].
  • The MCTS prompt asks the policy SLM to solve problems step by step with code while incorporating the trajectory of previous steps.
  • A combinatorics example generates products of two or more distinct set members with combinations, deduplicates them, and reports 15 unique products.
Loading 2501.04519v1…