RPM-MCTS: Knowledge-Retrieval as Process Reward Model with Monte Carlo Tree Search for Code Generation

AI summary

Overview

• Research area: Large language model code generation, tree-search decoding, and process reward modeling.
• Technical level: Advanced. Familiarity with Monte Carlo Tree Search, chain-of-thought reasoning, retrieval-augmented generation, and pass@k evaluation is assumed.
• Scope: The paper introduces RPM-MCTS, a training-free Monte Carlo Tree Search method that scores intermediate algorithmic steps using knowledge-base retrieval instead of a trained process reward model, and uses sandbox execution feedback to locate and correct erroneous steps during code generation.

What This Paper Is About

Tree-search methods for code generation struggle because they cannot reliably evaluate intermediate algorithmic steps, and when a step is wrong they often cannot locate it in time to correct it, producing incorrect code at high computational cost. The authors' goal is to evaluate and correct intermediate steps without training a separate process reward model, which normally requires dense per-step human annotations. RPM-MCTS substitutes knowledge-base retrieval over previously validated algorithmic steps as the process reward signal, and adds sandbox execution feedback to find and truncate erroneous steps.

Key Contributions

1. Knowledge-retrieval process reward. The paper proposes using retrieval scores over a knowledge base of correct algorithmic steps to evaluate intermediate steps and guide node selection in MCTS, avoiding the need to train a process reward model with dense step-level annotations.
2. Sandbox-driven error localization and correction. During the simulation phase, sandbox feedback on public test cases is used to evaluate generated code, localize the first erroneous algorithmic step, truncate the simulation there, and fold newly verified correct steps back into the tree, reducing computational cost.
3. Efficiency mechanisms for the tree. Expansion uses sampling decoding to generate diverse child steps plus cosine-similarity filtering to discard redundant nodes, and selection combines UCB with the knowledge-base retrieval score.
4. Empirical validation plus data distillation. The authors report improvements over state-of-the-art baselines on four public benchmarks (six test splits) with approximately 15% lower token consumption, and show that fine-tuning a base model on data synthesized by RPM-MCTS improves its code capabilities.

Main Findings

• Best average results across backbones. On Qwen3-8B, RPM-MCTS reaches an average of 64.0 versus 52.1 for the base LLM (a reported average improvement of 11.90%); on Qwen3-235B-A22B, 72.3 versus 64.6 (7.71%); on Claude Sonnet 3.7, 78.1 versus 67.3 (10.86%).
• Largest gains on hard problems. On the two more challenging datasets (APPS-competition and CodeContests), the reported average improvements are 13.3% for Qwen3-8B, 14.67% for Qwen3-235B-A22B, and 18.34% for Claude Sonnet 3.7.
• Works even without the knowledge base. The "Ours w/o KB" variant still outperforms the baselines overall; for Qwen3-235B-A22B it averages 71.3 versus 72.3 with the knowledge base, and the ablation attributes an overall average drop of 1.05% and a 4.67% drop on the two hardest datasets to removing the knowledge base.
• Execution feedback is the most important component. The ablation on Qwen3-235B-A22B reports the largest performance drop when public-test-case execution rewards are removed from the simulation phase, which the authors attribute to the code execution environment being the core of RPM-MCTS reflection.
• Similarity filtering and error localization. Removing similarity filtering hurts performance and raises cost; removing the LDB-style error-locating component has minimal average impact, though it helps in a few cases.
• Token efficiency. RPM-MCTS reduces token consumption by approximately 15% compared with the previous MCTS method on both Qwen3-235B-A22B and Claude Sonnet 3.7, attributed to knowledge-base guidance toward correct nodes, similarity filtering, and truncation after the first erroneous step.
• Strong performance at low rollout. RPM-MCTS performs better than baselines even at a rollout of 1, because it gets knowledge-base guidance in selection and wrong-step truncation with rethink-based regeneration in simulation.
• Distilled data improves the base model. Full fine-tuning Doubao-1.5-pro-32K on 2.4k RPM-MCTS-generated samples with reasoning steps, combined with a 170k-sample foundational dataset, moves SWE-Bench from 37.6 to 38.5 (+0.9), MBPP+ from 75.4 to 76.7 (+1.3), LiveCodeBench from 46.2 to 50.5 (+4.3), Aider from 17.3 to 22.2 (+4.9), and McEval from 57.5 to 61.2 (+3.7).
• Baseline behavior differs by difficulty. LDB improves more on simpler datasets, and the authors attribute this to LLMs often editing code conditions to pass public test cases rather than fixing the underlying logic on harder problems. SRA-MCTS improves on harder datasets but degrades on simpler ones, which the authors attribute to premature search termination when self-evaluation scores are near-perfect.
• Knowledge base can add noise on easy tasks. On a few simpler datasets, performance improves slightly when knowledge-base retrieval scores are not used, which the authors attribute to retrieval of textually similar but logically different historical cases.

Methodology in Plain English

The method is an altered Monte Carlo Tree Search in which the root node is the programming problem and every other node is one algorithmic step of the solution.

First, the authors build a knowledge base from the APPS and CodeContests training sets, which provide coding problems with correct solutions. Using Claude Sonnet 3.7, they decompose each correct solution into algorithmic steps and store every prefix of those steps (problem plus steps 1 through j), giving 11,038 samples and 82,923 steps. The entries are organized into 14 algorithm categories and embedded with the BGE model so similar problems with different algorithms can be distinguished.

The search then runs four phases repeatedly:

1. Selection. Starting at the root, the algorithm walks down to a leaf by maximizing a score that combines UCB (which balances exploitation of high-reward actions against exploration of rarely tried ones) with a knowledge-base retrieval score computed as the maximum cosine similarity between the current problem-plus-steps text and any stored knowledge entry. Ties are broken stochastically.
2. Expansion. The chosen leaf gets up to b new child steps generated by sampling decoding, with all previously generated steps given as context to discourage repetition. After expanding, the b children are embedded and compared by cosine similarity; nodes above a threshold are treated as redundant and filtered out.
3. Evaluation and reflection. The LLM completes all remaining steps and then generates code that strictly follows those steps. The code is run in a sandbox against public test cases, and the LLM also analyzes the full steps given the sandbox feedback. The node's value is a weighted sum of the public-test-case pass rate and the LLM score. If the code fails, the code is decomposed into blocks and each block is debugged sequentially with public test inputs to isolate the first erroneous step; correct steps before that point are kept and added to the tree.
4. Backpropagation. Reward values are propagated backward from the leaf to the root, updating value estimates along the path.

The process stops when the solution passes all public test cases and receives a high LLM evaluation score, or when the maximum iteration count is reached — in which case the highest-value leaf's path is used to generate the final code.

The attention analysis motivating step-level units uses a token-level attention heatmap: attention sinks to the begin-of-sequence token (with reported weights exceeding 30% or even 80% in some heads) and also peaks at algorithmic step boundaries, which the authors use to argue that algorithmic step blocks are better basic units than individual tokens.

Experimental setup: baselines are a base LLM prompt, LDB, ToT, and SRA-MCTS; backbones are Qwen3-235B-A22B, Claude Sonnet 3.7, and Qwen3-8B; metric is pass@1; test sets are APPS (150 samples each for introductory, interview, and competition), CodeContests test (150), HumanEval+ (164), and MBPP+ (378). Hyperparameters: rollout of 5, branching factor b of 3, UCB exploration constant β of 0.5, knowledge-base weight α of 0.5, and similarity filtering threshold of 0.85.

Why This Matters

The work shows that a reusable knowledge base of previously validated algorithmic steps can serve as a process reward signal, sidestepping the expensive step-level human annotation that training a reward model normally requires. It also demonstrates that execution feedback is not just a final filter but a mechanism for locating the first wrong step and pruning the search, which changes how search-based code generation can be budgeted.

Real-world applications:

• Automated competitive-programming and coding-assistant tools that must solve hard problems under a token or latency budget, where the reported approximately 15% token reduction matters.
• Synthetic training-data generation for code models, since the paper shows RPM-MCTS output used for distillation improved a base model across SWE-Bench, MBPP+, LiveCodeBench, Aider, and McEval.
• Repository and agentic coding pipelines where execution sandboxes and test suites already exist, letting the reflection mechanism plug into existing CI-style feedback.
• Educational or debugging tools that need to point at the specific algorithmic step that went wrong rather than only reporting that the code failed.

Industry relevance: the method requires no reward-model training and no change to the base LLM weights, and it reuses two assets enterprises already have — a corpus of solved problems with correct solutions, and sandboxed execution of public test cases. The distillation results suggest direct use in data-generation pipelines for in-house code models.

Future Directions

• Adaptive reward weighting. The conclusion proposes dynamically adjusting the weights of the knowledge-base and sandbox external rewards during MCTS evaluation based on LLM uncertainty, rather than the fixed weights used here.
• Fixing the step-decomposition limitation. The authors note that code solvable in a single line may be split across multiple lines by the step-by-step approach, which does not affect correctness but is a stated limitation.
• Deciding when retrieval helps versus hurts. The finding that knowledge-base rewards slightly hurt on some simpler datasets, and that effectiveness depends on the balance between task difficulty and LLM evaluation confidence, raises the question of when retrieval should be applied at all.
• Scaling the reflection and distillation approach. Since the paper reports distillation on Doubao-1.5-pro-32K with a 2.4k-sample synthesized set combined with 170k foundational samples, a natural next step is testing the recipe at larger data volumes and other model families.

Target Audience

Researchers and practitioners working on LLM code generation, tree-search decoding, and process reward modeling. It is most useful to those building search-based or agentic code systems that already have sandboxed execution and a corpus of solved problems available, and to teams constructing synthetic reasoning-step data for fine-tuning code models. Readers without a background in MCTS, UCB, and retrieval-based scoring will need to consult the cited prior work (ToT, ReST-MCTS*, SRA-MCTS, LDB) to follow the method details.

Authors’ abstract

Tree search-based methods have made significant progress in enhancing the code generation capabilities of large language models. However, due to the difficulty in effectively evaluating intermediate algorithmic steps and the inability to locate and timely correct erroneous steps, these methods often generate incorrect code and incur increased computational costs. To tackle these problems, we propose RPM-MCTS, an effective method that utilizes Knowledge-Retrieval as Process Reward Model based on Monte Carlo Tree Search to evaluate intermediate algorithmic steps. By utilizing knowledge base retrieval, RPM-MCTS avoids the complex training of process reward models. During the expansion phase, similarity filtering is employed to remove redundant nodes, ensuring diversity in reasoning paths. Furthermore, our method utilizes sandbox execution feedback to locate erroneous algorithmic steps during generation, enabling timely and targeted corrections. Extensive experiments on four public code generation benchmarks demonstrate that RPM-MCTS outperforms current state-of-the-art methods while achieving an approximately 15% reduction in token consumption. Furthermore, full fine-tuning of the base model using the data constructed by RPM-MCTS significantly enhances its code capabilities.