rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking — part 2/4

Cleaned-text code points [15428, 31179). Bibliography omitted here; source.md includes it.

Code-augmented CoT Generation . Prior MCTS approaches primarily generate natural language (NL) CoTs ( Qi et al., 2024 ; Zhang et al., 2024a ) . However, LLMs often suffer from hallucination, producing incorrect or irrelevant steps yet still arrive at the correct answer by chance ( Lanham et al., 2023 ) . These flawed steps are challenging to detect and eliminate. To address this, we propose a novel code execution augmented CoT. As shown in Fig. 2 , the policy model generates a one-step NL CoT alongside its corresponding Python code, where the NL CoT is embedded as a Python comment. Only generations with successfully executed Python code are retained as valid candidates.

Specifically, starting from the initial root node x x , we perform multiple MCTS iterations through selection , expansion , rollout , and back-propagation . At step i i , we collect the latest reasoning trajectory x ⊕ s 1 ⊕ s 2 ⊕ … ⊕ s i − 1 x\oplus s_{1}\oplus s_{2}\oplus...\oplus s_{i-1} as the current state. Based on this state, we prompt (see Appendix A.3 ) the policy model to generate n n candidates s i , 0 , … , s i , n − 1 s_{i,0},...,s_{i,n-1} for step i i . Python code execution is then employed to filter valid nodes. As shown in Fig. 2 , each generation s i , j s_{i,j} is concatenated with the code from all previous steps, forming s 1 ⊕ s 2 ⊕ … ⊕ s i − 1 ⊕ s i , j s_{1}\oplus s_{2}\oplus...\oplus s_{i-1}\oplus s_{i,j} . Candidates that execute successfully are retained as valid nodes and scored by the PPM, which assigns a Q-value q ⁡ ( s i ) q(s_{i}) . Then, we use the well-known Upper Confidence bounds for Trees (UCT) ( Kocsis and Szepesvári, 2006 ) to select the best node among the n n candidates. This selection process is mathematically represented as: UCT ​ ( s ) = Q ⁡ ( s ) + c ​ ln ⁡ N p ​ a ​ r ​ e ​ n ​ t ​ ( s ) N ⁡ ( s ) ; where Q ⁡ ( s ) = q ⁡ ( s ) N ⁡ ( s ) \displaystyle\text{UCT}(s)=Q(s)+c\sqrt{\frac{\ln N_{parent}(s)}{N(s)}};\quad\text{where}\quad Q(s)=\frac{q(s)}{N(s)} (1) where N ⁡ ( s ) N(s) denotes the number of visits to node s s , and N parent ​ ( s ) N_{\text{parent}}(s) is the visit count of s s ’s parent node. The predicted reward q ⁡ ( s ) q(s) is provided by the PPM and will be updated through back-propagation. c c is a constant that balances exploitation and exploration.

Extensive Rollouts for Q-value Annotation . Accurate Q-value Q ⁡ ( s ) Q(s) annotation in Eq. 1 is crucial for guiding MCTS node selection towards correct problem-solving paths and identifying high-quality steps within trajectories. To improve Q-value reliability, we draw inspiration from Go players, who retrospectively evaluate the reward of each move based on game outcomes. Although initial estimates may be imprecise, repeated gameplay refines these evaluations over time. Similarly, in each rollout, we update the Q-value of each step based on its contribution to achieving the correct final answer. After extensive MCTS rollouts, steps consistently leading to correct answers achieve higher Q-values, occasional successes yield moderate Q-values, and consistently incorrect steps receive low Q-values. Specifically, we introduce two self-annotation methods to obtain these step-level Q-values. Fig. 1 (c) shows the detailed setting in the four rounds of self-evolution.

Terminal-guided annotation . During the first two rounds, when the PPM is unavailable or insufficiently accurate, we use terminal-guided annotation. Formally, let q ​ ( s i ) k q(s_{i})^{k} denote the q value for step s i s_{i} after back-propagation in the k t ​ h k^{th} rollout. Following AlphaGo ( Silver et al., 2017 ) and rStar ( Qi et al., 2024 ) , we score each intermediate node based on its contribution to the final correct answer: q ​ ( s i ) k = q ​ ( s i ) k − 1 + q ​ ( s d ) k ; \displaystyle q(s_{i})^{k}=q(s_{i})^{k-1}+q(s_{d})^{k}; (2) where the initial q value q ​ ( s i ) 0 = 0 q(s_{i})^{0}=0 in the first rollout. If this step frequently leads to a correct answer, its q q value will increase; otherwise, it decreases. Terminal nodes are scored as q ⁡ ( s d ) = 1 q(s_{d})=1 for correct answers and q ⁡ ( s d ) = − 1 q(s_{d})=-1 otherwise, as shown in Fig. 1 .

PRM-augmented annotation . Starting from the third round, we use PPM to score each step for more effective generation. Compared to terminal-guided annotation, which requires multiple rollouts for a meaningful q q value, PPM directly predicts a non-zero initial q q value. PPM-augmented MCTS also helps the policy model to generate higher-quality steps, guiding solutions towards correct paths. Formally, for step s i s_{i} , PPM predicts an initial q ​ ( s i ) 0 q(s_{i})^{0} value based on the partial trajectory: q ​ ( s i ) 0 = P ​ P ​ M ​ ( x ⊕ s 1 ⊕ s 2 ⊕ … ⊕ s i − 1 ⊕ s i ) \displaystyle q(s_{i})^{0}=PPM(x\oplus s_{1}\oplus s_{2}\oplus...\oplus s_{i-1}\oplus s_{i}) (3) This q q value will be updated based on terminal node’s q ⁡ ( s d ) q(s_{d}) value through MCTS back-propagation in Eq. 2 . For terminal node s d s_{d} , we do not use PRM for scoring during training data generation. Instead, we assign a more accurate score based on ground truth labels as terminal-guided rewarding.

### 3.3 Process Preference Model

Process reward models, which provide granular step-level reward signals, is highly desirable for solving challenging math problems. However, obtaining high-quality step-level training data remains an open challenge. Existing methods rely on human annotations ( Lightman et al., 2023 ) or MCTS-generated scores ( Zhang et al., 2024a ; Chen et al., 2024 ) to assign a score for each step. These scores then serve as training targets, with methods such as MSE loss ( Chen et al., 2024 ) or pointwise loss ( Wang et al., 2024c ; Luo et al., 2024 ; Zhang et al., 2024a ) used to minimize the difference between predicted and labeled scores. As a result, the precision of these annotated step-level reward scores directly determines the effectiveness of the resulting process reward model.

Unfortunately, precise per-step scoring remains a unsolved challenge. Although our extensive MCTS rollouts improve the reliability of Q-values, precisely evaluating fine-grained step quality presents a major obstacle. For instance, among a set of correct steps, it is difficult to rank them as best, second-best, or average and then assign precise scores. Similarly, among incorrect steps, differentiating the worst from moderately poor steps poses analogous challenges. Even expert human annotation struggles with consistency, particularly at scale, leading to inherent noise in training labels.

We introduce a novel training method that trains a process preference model (PPM) by constructing step-level positive-negative preference pairs. As shown in Fig. 1 (b), instead of using Q-values as direct reward labels, we use them to select steps from MCTS tree for preference pair construction. For each step, we select two candidates with the highest Q-values as positive steps and two with the lowest as negative steps. Critically, the selected positive steps must lead to a correct final answer, while negative steps must lead to incorrect answers. For intermediate steps (except the final answer step), the positive and negative pairs share the same preceding steps. For the final answer step, where identical reasoning trajectories rarely yield different final answers, we relax this restriction. We select two correct trajectories with the highest average Q-values as positive examples and two incorrect trajectories with the lowest average Q-values as negative examples. Following ( Ouyang et al., 2022 ) , we define our loss function using the standard Bradley-Terry model with a pairwise ranking loss: ℒ p ​ p ​ m ​ ( θ ) = − 1 2 × 2 ​ E ( x , y i p ​ o ​ s , y i n ​ e ​ g ∈ 𝔻 ) ​ [ l ​ o ​ g ​ ( σ ⁡ ( r θ ​ ( x , y i p ​ o ​ s ) − r θ ​ ( x , y i n ​ e ​ g ) ) ) ] \displaystyle\mathcal{L}_{ppm}(\theta)=-\frac{1}{2\times 2}E_{(x,y_{i}^{pos},y_{i}^{neg}\in\mathbb{D})}[log(\sigma(r_{\theta}(x,y_{i}^{pos})-r_{\theta}(x,y_{i}^{neg})))] (4) when i is not final answer step , y i p ​ o ​ s = s 1 ⊕ … ⊕ s i − 1 ⊕ s i p ​ o ​ s ; y i n ​ e ​ g = s 1 ⊕ … ⊕ s i − 1 ⊕ s i n ​ e ​ g \displaystyle\text{when $i$ is not final answer step},y_{i}^{pos}=s_{1}\oplus...\oplus s_{i-1}\oplus s_{i}^{pos};y_{i}^{neg}=s_{1}\oplus...\oplus s_{i-1}\oplus s_{i}^{neg}\vskip-4.30554pt (5) Here, r θ ​ ( x , y i ) r_{\theta}(x,y_{i}) denotes the output of the PPM, where x x is the problem and y y is the trajectory from the first step to the i t ​ h i^{th} step.

### 3.4 Self-Evolved Deep Thinking

#### 3.4.1 Training with Step-by-Step Verified Reasoning Trajectory

Math Problems Collection . We collect a large dataset of 747k math word problems with final answer ground-truth labels, primarily from NuminaMath ( Jia LI and Polu, 2024a ) and MetaMath ( Yu et al., 2023b ) . Notably, only competition-level problems (e.g., Olympiads and AIME/AMC) from NuminaMath are included, as we observe that grade-school-level problems do not significantly improve LLM complex math reasoning. To augment the limited competition-level problems, we follow ( Li et al., 2024 ) and use GPT-4 to synthesize new problems based on the seed problems in 7.5k MATH train set and 3.6k AMC-AIME training split. However, GPT-4 often generated unsolvable problems or incorrect solutions for challenging seed problems. To filter these, we prompt GPT-4 to generate 10 solutions per problem, retaining only those with at least 3 consistent solutions.

Reasoning Trajectories Collection . Instead of using the original solutions in the 747k math dataset, we conduct extensive MCTS rollouts (Sec. 3.2 ) to generate higher-quality step-by-step verified reasoning trajectories. In each self-evolution round, we perform 16 rollouts per math problem, which leads to 16 reasoning trajectories. Problems are then categories by difficulty based on the correct ratio of the generated trajectories: easy (all solutions are correct), medium (a mix of correct and incorrect solutions) and hard (all solutions are incorrect). For hard problems with no correct trajectories, an additional MCTS with 16 rollouts is performed. After that, all step-by-step trajectories and their annotated Q-values are collected and filtered to train the policy SLM and process preference model.

Supervised Fine-tuning the Policy SLM . Through extensive experiments, we find that selecting high-quality reasoning trajectories is the key for fine-tuning a frontier math LLM. While methods such as GPT-distillation and Best-of-N can include low-quality or erroneous intermediate steps, a more effective approach ensures that every step in the trajectory is of high quality. To achieve this, we use per-step Q-values to select optimal trajectories from MCTS rollouts. Specifically, for each math problem, we select the top-2 trajectories with the highest average Q-values among those leading to correct answers as SFT training data.

Training PPM . The PPM is initialized from the fine-tuned policy model, with its next-token prediction head replaced by a scalar-value head consisting of a linear layer and a tanh function to constrain outputs to the range [-1, 1]. We filter out math problems where all solution trajectories are fully correct or incorrect. For problems with mixed outcomes, we select two positive and two negative examples for each step based on Q-values, which are used as preference pairs for training data.

#### 3.4.2 Recipe for Self-Evolution

Due to the weaker capabilities of SLMs, we perform four rounds of MCTS deep thinking to progressively generate higher-quality data and expand the training set with more challenging math problems. Each round uses MCTS to generate step-by-step verified reasoning trajectories, which are then used to train the new policy SLM and PPM. The new models are then applied in next round to generate higher-quality training data. Fig. 1 (c) and Table 2 detail the models used for data generation in each round, along with the identifiers of the trained policy model and PPM. Next, we outline the details and specific improvements targeted in each round.

Round 1: Bootstrapping an initial strong policy SLM-r1 . To enable SLMs to self-generate reasonably good training data, we perform a bootstrap round to fine-tune an initial strong policy model, denoted as SLM-r1. As shown in Table 2 , we run MCTS with DeepSeek-Coder-V2-Instruct (236B) to collect the SFT data. With no available reward model in this round, we use terminal-guided annotation for Q-values and limit MCTS to 8 rollouts for efficiency. For correct solutions, the top-2 trajectories with the highest average Q-values are selected as SFT data. We also train PPM-r1, but the limited rollouts yields unreliable Q-values, affecting the effectiveness of PPM-r1 ( Table 4 ).

Round 2: Training a reliable PPM-r2 . In this round, with the policy model updated to the 7B SLM-r1, we conduct extensive MCTS rollouts for more reliable Q-value annotation and train the first reliable reward model, PPM-r2. Specifically, we perform 16 MCTS rollouts per problem. The resulting step-by-step verified reasoning trajectories show significant improvements in both quality and Q-value precision. As shown in Table 4 , PPM-r2 is notably more effective than in the bootstrap round. Moreover, the policy SLM-r2 also continues to improve as expected (Table 3 ).

Round 3: PPM-augmented MCTS to significantly improve data quality . With the reliable PPM-r2, we perform PPM-augmented MCTS in this round to generate data, leading to significantly higher-quality trajectories that cover more math and Olympiad-level problems in the training set (Table 2 ). The generated reasoning trajectories and self-annotated Q-values are then used to train the new policy SLM-r3 and PPM-r3, both of which show significant improvements.

Round 4: Solving challenging math problems . After the third round, while grade school and MATH problems achieve high success rates, only 62.16% of Olympiad-level problems are included in the training set. This is NOT solely due to weak reasoning abilities in our SLMs, as many Olympiad problems remain unsolved by GPT-4 or o1. To improve coverage, we adopt a straightforward strategy. For unsolved problems after 16 MCTS rollouts, we perform an additional 64 rollouts, and if needed, increase to 128. We also conduct multiple MCTS tree expansions with different random seeds. This boosts the success rate of Olympiad-level problems to 80.58%.

After four rounds of self-evolution, 90.25% of the 747k math problems are successfully covered into the training set, as shown in Table 2 . Among the remaining unsolved problems, a significant portion consists of synthetic questions. We manually review a random sample of 20 problems and find that 19 are incorrectly labeled with wrong answers. Based on this, we conclude that the remaining unsolved problems are of low quality and thus terminate the self-evolution at round 4.

## 4 Evaluation

### 4.1 Setup

Evaluation Datasets . We evaluate rStar-Math on diverse mathematical benchmarks. In addition to the widely-used GSM8K ( Cobbe et al., 2021 ) , we include challenging benchmarks from multiple domains: (i) competition and Olympiad-level benchmarks, such as MATH-500 ( Lightman et al., 2023 ) , AIME 2024 ( AI-MO, 2024a ) , AMC 2023 ( AI-MO, 2024b ) and Olympiad Bench ( He et al., 2024 ) . Specifically, AIME is the exams designed to challenge the brightest high school math students in American, with the 2024 dataset comprising 30 problems from AIME I and II exams; (ii) college-level math problems from College Math ( Tang et al., 2024 ) and (iii) out-of-domain math benchmark: GaoKao (Chinese College Entrance Exam) En 2023 ( Liao et al., 2024 ) .