Preprint
Reinforcement Learning

Monte Carlo Tree Self-refine

June 1, 2024

0

Citations

0

Influential Citations

Venue

2024

Year

Abstract

Integrates LLMs with Monte Carlo Tree Search to enhance performance in complex mathematical reasoning tasks, leveraging systematic exploration and heuristic self-refine mechanisms to improve decision-making frameworks.

Analysis

Why This Paper Matters

Large language models (LLMs) have shown remarkable fluency but often falter on tasks requiring multi-step logical reasoning, especially in mathematics. Traditional prompting methods like chain-of-thought (CoT) provide only a single pass, leaving no room for iterative improvement or error correction. The Monte Carlo Tree Self-refine (MCTSr) algorithm bridges this gap by integrating the structured exploration of Monte Carlo Tree Search (MCTS) with the generative power of LLMs. This approach allows models to systematically explore solution spaces, refine answers through self-feedback, and backpropagate quality signals—much like how AlphaGo improved its gameplay through simulation. The paper demonstrates that even a relatively small 8B parameter model can achieve competitive results on challenging mathematical benchmarks, making advanced reasoning more accessible.

How It Works

MCTSr adapts the classic MCTS framework to the domain of mathematical problem-solving. The algorithm operates on a search tree where each node represents a candidate answer to a given problem. Edges correspond to self-refine actions—modifications guided by the LLM's own feedback.

Figure 1

The process begins with initialization: a root node is created using a naive model-generated answer and a dummy response ("I don't know") to reduce overfitting. The selection phase uses a value function Q to rank all unexpanded nodes, picking the highest-valued one via a greedy strategy. This node then undergoes self-refine: the LLM generates feedback m and produces an improved answer a'.

Self-evaluation follows, where the refined answer is scored using a reward function R. To enhance reliability, the algorithm applies several constraints: a prompt constraint ensures strict scoring standards; full score suppression reduces any reward above 95 by a constant; and repeated sampling collects multiple reward samples for each node. The Q value is computed as:

Figure 3

where Q(a) = min(Ra) + (1/|Ra|) * sum(Ra). This formula combines the worst-case and average rewards to produce a conservative estimate.

Backpropagation then propagates the refined answer's value upward to parent nodes, updating their Q values using:

Figure 4

where Q'(a) = max(Q(a), max_{i in Children(a)} Q(i)). This ensures that a node's quality reflects both its own rewards and the best among its descendants.

Finally, the UCT update recalculates selection scores for all nodes:

Figure 5

where U(a) = Q(a) + c * sqrt(log(N(father(a)) + 1) / (N(a) + epsilon)). This balances exploitation (high Q) and exploration (low visit count).

The search terminates when improvements diminish, rollout limits are reached, or logit-based criteria are met. The best answers are then gathered from nodes with highest Q values.

Results

The paper evaluates MCTSr on three benchmark families: GSM (grade-school math), MATH (competition-level), and Olympiad-level datasets (AIME, GAIC Math Odyssey, OlympiadBench).

Performance of MCTSr on the GSM Dataset.

On GSM8K, success rates increase monotonically with rollouts, showing clear improvement from iterative refinement. However, GSM-Hard exhibits a ceiling effect, indicating that current self-refine strategies may not handle the most complex grade-school problems.

Performance of MCTSr on the MATH Dataset

On MATH, 8-rollout MCTSr achieves 58.24% overall success rate, more than doubling the 24.36% of Zero-Shot CoT. Level-1 problems reach 90.16%, while Level-5 problems reach 34.06%—a substantial gain over baseline.

Performance of MCTSr on Olympiad-level Datasets

Closed-source LLM performance on mathematical datasets

On Olympiad benchmarks, MCTSr consistently outperforms Zero-Shot CoT. On AIME, success rate jumps from 2.36% to 11.79%; on GAIC Math Odyssey, from 17.22% to 49.36%; on OlympiadBench, from 1.25% to 7.76%. These results show that iterative self-refinement can unlock significant reasoning capability even in small models.

Significance

MCTSr represents a practical fusion of search algorithms and LLMs, enabling systematic exploration of solution spaces without requiring external verifiers or human feedback. By using self-evaluation and backpropagation, the algorithm creates a virtuous cycle of improvement. The open-source implementation in LangChain lowers the barrier for practitioners to adopt this technique. While limitations exist—particularly on very hard problems—the approach points toward a future where smaller, more efficient models can rival larger ones through smarter inference-time computation. This has implications for deploying advanced reasoning in resource-constrained environments, from edge devices to educational tools.