SRA-MCTS: Self-driven Reasoning Augmentation with Monte Carlo Tree Search for Code Generation

1 Introduction

Large language models (LLMs) excel at generating code for simple questions, but their performance on more complex tasks remains unsatisfactory. [Luo et al., 2024; Li et al., 2023]. Recent studies have predominantly concentrated on enhancing the quality of training datasets to improve model performance. These improvements typically address two fundamental aspects of the training data.

On one hand, improving the quality of the questions enhances the model's understanding of the input. Some methods optimize the question portion of the dataset by leveraging more powerful models or summarizing existing opensource code, thereby enhancing both the depth and breadth of the questions [Ouyang et al., 2022; Luo et al., 2024; Wei et al., 2023]. On the other hand, increasing the complexity of the answers helps guide the model to generate higher-quality codes. BRAINSTORM [Li et al., 2023] introduces a brainstorming approach to generate and select diverse ideas related to a given question, thereby enhancing the model's reasoning capabilities. DolphCoder [Wang et al., 2024b] utilizes GPT-3.5 with multiple prompt templates to obtain diverse outputs for multi-objective instruction finetuning. These methods all aim to leverage additional natural language text to guide LLMs in generating high-quality results.

Furthermore, ScaleAI [Wang et al., 2024a] has demonstrated that providing LLMs with correct solutions in natural language form can significantly enhance model performance, even when these solutions consist of incomplete plans with only a few dozen tokens. This highlights the potential of natural language solutions to guide and inspire LLMs to think in the right direction. Similar to ScaleAI, there have been some works exploring how to improve the intermediate reasoning steps to improve model performance. [Chen et al., 2024; Li et al., 2024b; Long et al., 2024]. However, these studies lack explanations for the natural language solutions and lack an inherent mechanism for error correction, limiting the model's ability to effectively utilize natural language reasoning steps. Therefore, we focus on the intermediate natural language reasoning steps generated by the model, providing it with a reliable and diverse reasoning direction.

In this paper, we refer to the intermediate natural language reasoning steps generated by the model for code generation as "thinking". We propose a Self-driven Reasoning Augmentation method, leveraging Monte Carlo Tree Search (MCTS) [Coulom, 2006] to generate high-quality thinking, which we call SRA-MCTS as shown in Figure 1. Different from ReST-MCTS* [Zhang et al., 2024a], which uses MCTS to generate complete answers, we apply MCTS to the more critical thinking portion. This provides the model with reliable reasoning, enhancing its ability to generate high-quality code. SRA-MCTS iteratively constructs the final reasoning path for a given question. Specifically, each iteration generates multiple potential next steps based on the current reasoning path. During each step selection, the model reflects on and scores the newly generated step, updating the scores of previously generated steps. This process allows the model to re-evaluate and reconsider unselected steps in subsequent iterations. This iterative process enables the model to engage in self-reflection, correcting potentially flawed reasoning steps. Finally, the sequence of reasoning steps generated over multiple iterations is combined into comprehensive thinking that guides the model in generating the final code. The data generated through SRA-MCTS will be used to fine-tune the model, ultimately resulting in a higher-performing code generation model. Our contributions are as follows:

• We propose a plug-and-play reasoning augmentation data generation method, SRA-MCTS. This method is simple and effective, as it enables the LLM to selfgenerate and self-evaluate data during the process. Then the generated data is used for model training, fostering the model's self-improvement.
• We propose a pipeline for the code generation domain, where input questions are processed by SRA-MCTS to generate diverse, high-quality thinking without additional supervision, thereby creating a positive feedback loop for continuous improvement.
• We perform a detailed analysis of the generated data, examining the impact of its quality on model performance. Additionally, by comparing the model's performance on the original benchmark and the derived, more complex benchmarks, we validate the effectiveness of the selfreasoning augmentation approach.

2 RelatedWork

Data Augmentation Evol-Instruct [Luo et al., 2024] enhances CodeAlpaca [Chaudhary, 2023] by integrating heuristic rules, which increases the complexity and diversity of the seed instructions. OSS-Instruct [Wei et al., 2023] leverages open-source code snippets to generate questions that better represent real-world distributions. CodeOcean [Yu et al., 2023] utilizes GPT-3.5 and GPT-4 for question distillation and filtering, resulting in high-quality, diverse data. DolphCoder [Wang et al., 2024b] improves LLM learning by generating step-by-step answers through multiple turns and pseudo code for problem-solving.

Self-generated Data Data generated by LLMs can effectively mitigate biases in understanding arising from variations in language style. ReCo [Li et al., 2024a] utilizes LLMs to rewrite code from codebases, thereby reducing retrieval accuracy degradation caused by stylistic deviations during the retrieval process. ReST-MCTS* [Zhang et al., 2024a] introduces a reinforced self-training approach that combines process reward guidance with MCTS to collect higher-quality reasoning trajectories and stepwise values, which are subsequently used to train the strategy and reward models.

Exploration and Thinking Methods Similar to how humans think deeply before answering difficult questions, OpenAI o1 [OpenAI, ] employs a series of thought processes to solve problems. Through reinforcement learning, it refines its strategies, corrects errors, and streamlines complex steps. AFLOW [Zhang et al., 2024b] employs MCTS to refine workflows through iterative code modifications and feedback. StepCoder [Dou et al., 2024] utilizes parts of a standardized solution as a prompt, enabling LLMs to explore simple sequences and improve through feedback-based reinforcement learning.

Unlike previous works, our SRA-MCTS not only enhances the thinking process within the data, improving the model's reasoning and thinking capabilities, but also establishes a fully autonomous pipeline that includes self-generation, selfevaluation, and ultimately self-improvement.

3 Method

We propose a pipeline that leverages the model itself to enhance training data and augment reasoning ability, as shown in Figure 1, consisting of three stages: SRA-MCTS, Thinking to code, and Training. First, we use the LLM with SRAMCTS to generate step-by-step natural language thinking for each question. Then, the LLM generates specific codes based on the question and the generated thinking. Finally, we combine the question, the thinking, and the code to create a finetuning dataset for training the LLM.

3.1 SRA-MCTS

We develop a self-driven reasoning augmentation method with MCTS, called SRA-MCTS, to generate natural language thinking for questions. It consists of four phases: Selection, Expansion, Evaluation & Reflection, and Backpropagation, as shown in Figure 2. These phases are carried out on a search tree composed of tree nodes and will be repeated multiple times, with each iteration generating a specific reasoning step.

A tree node has four attributes: 1) State represents the current state, including the question and the steps generated so far; 2) Action represents the next step generated based on the current state; 3) Reward represents the score generated by LLM based on state and action; 4) Reflection guides the generation of the next iteration. The initial search tree consists of only a root node, with its state attribute set to the input question, and the other three attributes are empty.

Selection

The purpose of the selection is to select one step from the model's currently generated steps for the next step generation. The selection phase involves selecting a leaf node Nd from the search tree, where d represents the depth of the leaf node. This node will be expanded in the expansion phase.

We use UCT [Kocsis and Szepesv´ari, 2006], which is calculated based on the node's reward value, to measure the value of each node being selected. The formula is as follows:

where r is the reward value of the current node that is generated by the Evaluation & Reflection phase in the last iteration, c is the exploration constant set to 0.5, N denotes the number of times the parent node has been selected, and n represents the number of times the current node has been selected. UCT is designed to balance exploitation and exploration, represented by the terms on the left and right of the formula respectively. If the exploration term is removed, it degenerates into a greedy selection, which harms diversity.

With the selection phase starting from the root node, the node with the highest UCT value at the current depth is selected. Then, the process continues by choosing the child node with the highest UCT value among all the child nodes of the selected node. This process continues until a leaf node of the search tree is selected. Note that during the first iteration, only the root node which is also a leaf node exists, so the root node's UCT value does not need to be calculated.

Expansion

The purpose of the expansion is to generate the next step based on the existing steps. The expansion phase involves generating several new child nodes based on the selected node Nd, and assigning values to the state and action attributes of the new nodes.

The state attribute is designed to store the input question and steps that the LLM has generated up to the current node. Therefore, for child nodes that originate from the same parent node, their state values are identical and referred to Sd+1, as they share the input question and the same history of step. Specifically, Sd+1 is formed by concatenating the state and action attributes of the node Nd:

where Sd and Ad represent the state and action of the selected node Nd, respectively.

The action attribute is designed to store the next step generated based on the current state and the reflection of the parent node. Since different methods can be used to generate different next steps, the action attribute can have multiple distinct values, leading to the creation of different child nodes. Each node represents a possible path that can be made from the current state, allowing for the exploration of various possibilities in the search tree.

To generate diversity in the steps, the expansion phase adopts the sample decoding strategy, where each next step is generated one by one. Each time a new step is generated, it represents a new node created from the selected node Nd. To prevent the LLM from generating duplicate next steps, we add the next steps that have already been generated to the input. This process can be represented by the following formula:

where Aid+1 represents the action value of the i-th child node at depth d+1, Rd represents the reflection attribute value of the selected node Nd, which is generated in the last iteration. Adding the already generated steps to the input helps the LLM explore different paths in the search space.

To manage computational costs, we set the number of expanded nodes to 3. If a duplicate step is detected, the generation process is repeated, with a maximum of 5 retries At this stage, code generation is not involved, only the natural language reasoning process is generated.

Evaluation & Reflection

The third phase consists of two parts: Evaluation and Reflection, which support the selection and expansion phases, respectively.

Evaluation The evaluation part assigns the reward attribute values to the new nodes generated in the expansion phase. These reward values will be then used in the selection phase to calculate the UCT values.

We use the LLM itself, which generates the steps during the expansion phase, as the evaluator to score the new nodes and explore their potential for self-improvement through selfevaluation. 1 The reward value of the i-th new node N id+1 is as follows:

We use a progressive scoring method, making judgments sequentially from four aspects, the process is shown in Figure 3. The order of judgment is as follows: Single-step correctness, Solution coherence, Solution completeness, and Solution correctness. If the criteria for the current aspect are met, the corresponding score is given directly, and the other aspects are ignored.

Reflection The reflection part aims to guide the next generation direction for the current node, and assigns the reflection attribute values to the new nodes generated in the expansion phase. The reflection attribute is a brief thought generated by the LLM and will be used in the next iteration's expansion phase.

We use the reflection [Zhang et al., 2024a] mechanism to allow the model to think before generating the next action, thus preventing performance degradation. The reflection value of the i-th new node N id+1 is as follows:

Additionally, the reflection value is used to assess whether all the generated steps have solved the question. If the generated steps from state and action at this node are deemed sufficient to solve the question, the reflection value will be an < end > tag to indicate that the SRA-MCTS process has concluded. Otherwise, the LLMwill generate a reflection for the next iteration's reasoning. The reflection value is shown as follows:

1 Note: The use of the model itself here is aimed at exploring the potential for end-to-end self-improvement. Larger open-source models and closed-source model APIs, which may have more accurate evaluation capabilities, can easily replace this component.

Backpropagation

Backpropagation aims to update the reward values of all parent nodes of the current node, making the reward values of the parent nodes more accurate, thereby affecting the generation of the thinking. For the new nodes generated during the expansion phase {N 1 d+1 , N 2 d+1 , ..., N I d+1 }, they collectively update their parent node, which is the node selected in the selection phase Nd. Afterward, the parent node Nd continues to update its parent node Nd-1, and this process is repeated until the root node is updated. The update process of the parent node's reward value rp is as follows:

where vic represents the number of times the i-th child node is selected, ric represents the reward score of the i-th child node, r∆ represents the weighted increment of the parent node's reward value, and α is a hyperparameter to balance the original reward of the parent node and the increment from the child nodes.

The updated reward value of the parent node will influence the selection of child nodes in subsequent iterations. Specifically, in the next selection phase, SRA-MCTS will use the updated reward values to evaluate the relative merits of different nodes, thereby deciding which branch to explore. This dynamic reward adjustment helps gradually refine the strategy across the entire tree, making the search more efficient and progressively converging toward the optimal solution.

SRA-MCTS repeats the four phases above until the < end > tag is output, or the iteration limit is reached.

3.2 Thinking to Code

SRA-MCTS generates multiple steps based on the question, and these steps together form a thinking. In this stage, we have the LLM act as a function implementer, strictly converting the thinking into code.

3.3 Training

After the first two stages, we collect the question, thinking, and code triples, with the question as the input and the latter two as the output, forming the training dataset. This dataset is then used for supervised fine-tuning.

4 Experiment

4.1 Training Data Generation

Question To validate whether SRA-MCTS can indeed generate higher-quality thinking, we need to collect questions with a certain degree of difficulty. Therefore, we select medium and hard-level questions from the LeetCode dataset [Greengerong, 2023], excluding easy-level questions. We retain only the question descriptions and discard the answers, allowing different methods to generate both thinking and code. To prevent any overlap between the training and test sets, we perform decontamination on the training data. Specifically, we conduct 10-gram level duplicate detection on each question in the dataset, setting the threshold at 0.3. If the similarity between a training question and any test set question exceeds this threshold, we remove the corresponding training data. As a result, we obtain a final training set of 1,819 unique questions.

Thinking and Code For thinking and code generation, we use SRA-MCTS to generate the corresponding thinking for the questions in the training set. Each step in the generation process is performed in a zero-shot manner. We prompt the model to generate each step according to a specified format, ensuring that the steps can be easily extracted. If the thinking contains code, we use a regular expression to remove it, leaving only the natural language reasoning content. This ensures that the reasoning process is kept separate from the code, enabling a clearer focus on the thought process behind the code generation. After generating the final thinking, we concatenate the question with its corresponding thinking and prompt the model to generate the final code.

4.2 Baseline

We choose gemma-2-2b-it [Team et al., 2024], Meta-Llama3.1-8B-Instruct [Dubey et al., 2024], and Qwen2.5-14BInstruct [Team, 2024] as our backbone models. The rationale behind this selection is their strong instruction-following capabilities and the fact that they have been pre-trained on extensive code data. These models demonstrate significant potential to activate proactive thinking for solving programming tasks through reasoning augmentation techniques. We directly use the backbone models to perform reasoning on the test set, which serves as the baseline performance score for comparison.

We apply the Chain-of-Thought (CoT) [Wei et al., 2022] and Tree-of-Thought (ToT) [Yao et al., 2023] reasoning methods to each model as baseline methods. For CoT, we use the conventional prompt, "Let's think step by step", to generate the thinking. For ToT, we use depth-first search to generate the thinking, setting the maximum depth to 4 and the branching factor to 3, which aligns with the number of nodes expanded during the expansion phase of SRA-MCTS. Similar to SRA-MCTS, the thinking generated by these methods is processed using regular expressions to remove any code content. Finally, the generated thinking is concatenated with the question and fed into the model to produce the final code.

4.3 Test Set

Benchmark We use commonly adopted benchmarks in the code generation field, including Human-Eval [Chen et al., 2021] and MBPP [Austin et al., 2021]. Additionally, we utilize Human-Eval+ and MBPP+ within the EvalPlus [Liu et al., 2023] framework, where the test cases are several times larger than the original versions. As for the task type, HumanEval involves completing a given function, while MBPP requires the model to generate the function from scratch.

Test Set Split To evaluate the model's reasoning ability across different difficulty levels, we categorize the test set questions into different difficulty levels. Specifically, we use GPT-4o [Hurst et al., 2024] to classify the questions into easy, medium, and hard categories. All questions across the three difficulty levels form the Full split and the medium and hard questions are grouped into the Complex split.

Evaluation Metric In code generation field, pass@k is a widely used metric that measures the probability of generating a correct solution in at least one of k attempts [Chen et al., 2021; Kulal et al., 2019]. It reflects the model's success rate across multiple attempts, which is particularly valuable given the inherent uncertainty in the generation process.

4.4 Training Setup

In SRA-MCTS, the total iteration limit is set to 5, the generation temperature in the expansion phase is set to 0.9, with a top-p sampling value of 0.98, and the α value in the backpropagation phase is manually set to 0.5.

We use the models mentioned in the baseline as the foundation for training, setting the learning rate to 1e-4, training for 2 epochs, with a warmup ratio of 0.05 and a batch size of 4. For each model scale, we train using three different datasets: 1) Self-generated CoT data; 2) Self-generated ToT data; 3) SRA-MCTS generated data. We use Llama-factory [Zheng et al., 2024] for LoRA [Hu et al., 2022] fine-tuning, with the models trained on two RTX A6000 GPUs.

4.5 Main Results

SRA-MCTS demonstrates the best overall performance and excels in solving complex questions. As shown in Table 1, our method achieves the highest average performance improvement across all models, with a significantly larger gain on the Complex split of most benchmarks compared to the Full split, particularly with an increase of approximately 2-15 points on MBPP and MBPP+. However, for the HumanEval-related benchmarks, all methods experience a decline in several metrics. While our method still shows improvement in some metrics, it is not obvious compared to the other baseline methods. We attribute this phenomenon to the differences between the task types in the training and test sets. As discussed in Section 4.3, tasks in the MBPP benchmark, like those in the training set, require the model to generate code from scratch based on a given question, whereas Human-Eval tasks ask the model to complete a function header provided in the prompt. Due to the consistency in task style between MBPP and the training set, our method achieves more substantial improvements compared to other methods. This highlights the critical role that the style of training data plays in influencing downstream task performance.

The larger the model, the more pronounced the improvement from our method. As we observe in Table 1, the average performance improvement increases with the model size, and this trend is even more evident on the Complex split. We attribute this to the fact that larger models possess better generation and evaluation capabilities. With stronger generation and evaluation capabilities enabled by larger model scales, these models can produce higher-quality data and more effectively identify errors and redundancies in reasoning steps. SRA-MCTS effectively leverages these enhanced capabilities, leading to significant performance gains.

5 Analysis

5.1 Ablation Study of the Thinking

Effect of Existence We additionally train a model on a dataset without thinking, containing only the question and code, and compare it with the dataset that includes the question, thinking, and code triples. For a fair comparison, the code in the dataset without thinking is directly extracted from the training set generated by SRA-MCTS. As shown in Figure 4, nearly all cases without thinking perform worse than those with thinking, with this gap being more obvious on the Complex split. We believe this is due to simpler questions not requiring much reasoning from the model to provide the correct solution, allowing a code-only training set to solve these questions effectively. It is only when the question difficulty increases that the model benefits from higher-quality thinking to guide it toward generating the correct code.

Effect of Variants We extract the thinkings generated by CoT, ToT, and SRA-MCTS from their respective training sets and concatenate them with the question and ground truth code to create three new training sets, each differing only in the thinking content. We then retrain the backbone models on these new datasets, and the results are shown in Figure 5. Our method generates the most stable thinking among the three approaches. This is demonstrated by the fact that even when the code is replaced, our method maintains strong performance, whereas the other methods experience noticeable performance degradation. This indicates that the thinking generated by our method not only contributes to the final performance improvement but also provides robustness, leading to a more stable model.

Effect of Quality We use GPT-4o to evaluate the quality of the thinking, categorizing each instance into low, medium, and high quality. The ratios of medium and high-quality thinking in each training set are used to represent the quality score of the dataset. As shown in Figure 6, there is generally a linear relationship between the quality of the thinking in the training set and the model's performance, meaning that higher-quality thinking leads to better final performance. However, for our method, there are cases in the gemma and Qwen models where the thinking quality is relatively low, yet the performance is higher than expected. We hypothesize that, during the evaluation of thinking, the evaluator may overemphasize unnecessary overthinking steps, complicating the model's ability to generate code based on overly complex reasoning processes. This not only suggests that the evaluation of larger models may not always be reliable, but also highlights the superiority of our method.

5.2 Further analysis

Comparison with Competitive Methods. We compared our approach with other competitive methods using models of similar sizes, as shown in Table 2. For 2B and 13B models, our method not only excels in individual metrics but also achieves the best overall average performance. The only exception is the 8B Magicoder-S-CL method, where our approach falls slightly behind. We attribute this to the fact that our fine-tuning dataset is significantly smaller and less diverse compared to Magicoder. The Magicoder dataset is approximately 41 times larger than ours, and its larger scale and potential diversity likely contribute to its superior performance. This comparison underscores the efficiency of our approach. Despite using a small dataset, we are able to generate results that are of sufficiently high quality and diversity. This highlights the effectiveness of our method, which can achieve competitive results even with much more limited data.

Effectiveness of Self-generation. We use Meta-Llama-370B [Dubey et al., 2024] as an external model to generate a training dataset, which is then used to train smaller models. We compare the performance of models trained with selfgenerated data versus those trained with data distilled from the larger model. The results in Table 3 demonstrate that our self-generation approach outperforms the data distillation method across all model sizes and nearly all benchmarks. Notably, SRA-MCTS shows substantial improvements on the Human-Eval and Human-Eval+ datasets. In contrast, the distillation method only outperforms on the MBPP benchmark when using a 14B model. This result strongly supports the feasibility of using small models to self-generate data for training, demonstrating that self-generated data can achieve competitive performance without relying on large-scale data distillation.

6 Conclusion

We propose a self-improvement pipeline where SRA-MCTS performs both self-generation and self-evaluation for given questions. The high-quality thinking text generated through this process serves to guide code generation, with both the thinking and the code subsequently used for fine-tuning. Experimental results demonstrate that SRA-MCTS outperforms both CoT and ToT on more complex tasks, achieving an 11point improvement on the MBPP Complex split.

Acknowledgements

This work was supported by the Major Research Plan of the National Natural Science Foundation of China (Grant No. 92370110).

Contribution Statement

• Bin Xu: Developed the idea; conceived and designed the analysis; performed the analysis; wrote the paper.
• Yiguan Lin (Equal Contribution): Conceived and designed the analysis; performed the analysis; wrote the paper.
• Yinghao Li: Conceived and designed the analysis; performed the analysis; wrote the paper
• YangGao (CorrespondingAuthor): Provided computing resources and funding.