Source-linked AI summary
Accessing GPT-4 level Mathematical Olympiad Solutions via Monte Carlo Tree Self-refine with LLaMa-3 8B
Di Zhang, Xiaoshui Huang, Dongzhan Zhou, Yuqiang Li, Wanli Ouyang
TL;DR
LLMs can be unreliable on complex mathematical reasoning because they may produce plausible but incorrect outputs. The paper proposes MCTSr, which combines LLM self-refinement and self-evaluation with Monte Carlo Tree Search, and reports improved performance across mathematical benchmarks, while noting that broader applications remain unexplored.
Problem
LLMs face accuracy and trustworthiness challenges in strategic and mathematical reasoning, including plausible but irrelevant or factually incorrect outputs.
Method
MCTSr integrates LLM self-refinement and self-evaluation with MCTS, using dynamic pruning and an improved UCB-based decision process.
Results
MCTSr improves mathematical problem-solving performance across multiple datasets, including Olympic-level mathematical challenges.
Takeaways & Limitations
The work advances LLM application to sophisticated reasoning tasks and supports future AI integration for decision-making and reasoning accuracy.
Takeaways & Limitations
MCTSr remains preliminary, with broader applications such as black-box optimization and self-driven alignment still requiring further exploration.
Abstract
from arXiv · showhide
This paper introduces the MCT Self-Refine (MCTSr) algorithm, an innovative integration of Large Language Models (LLMs) with Monte Carlo Tree Search (MCTS), designed to enhance performance in complex mathematical reasoning tasks. Addressing the challenges of accuracy and reliability in LLMs, particularly in strategic and mathematical reasoning, MCTSr leverages systematic exploration and heuristic self-refine mechanisms to improve decision-making frameworks within LLMs. The algorithm constructs a Monte Carlo search tree through iterative processes of Selection, self-refine, self-evaluation, and Backpropagation, utilizing an improved Upper Confidence Bound (UCB) formula to optimize the exploration-exploitation balance. Extensive experiments demonstrate MCTSr's efficacy in solving Olympiad-level mathematical problems, significantly improving success rates across multiple datasets, including GSM8K, GSM Hard, MATH, and Olympiad-level benchmarks, including Math Odyssey, AIME, and OlympiadBench. The study advances the application of LLMs in complex reasoning tasks and sets a foundation for future AI integration, enhancing decision-making accuracy and reliability in LLM-driven applications.
1 Introduction
The paper addresses accuracy and trustworthiness problems in LLM reasoning by integrating LLMs with Monte Carlo Tree Search and self-refinement. MCTSr adds dynamic pruning and is reported to improve complex mathematical problem solving.
- Motivation: LLMs can produce plausible but irrelevant or incorrect outputs in complex mathematical reasoning, despite rewriting techniques such as Self-Refine.The paper identifies accuracy and trustworthiness as significant hurdles where precision is paramount.
- Approach: MCTSr combines LLMs’ Self-Refine and Self-Evaluation capabilities with MCTS’s systematic exploration for intricate reasoning tasks.The integration is intended to create a more robust framework for mathematical problems current LLMs struggle with.
- Approach: The method adapts MCTS components to LLM outputs and introduces dynamic pruning with an improved UCB formula to balance exploration and exploitation.Traditional MCTS requires adaptation because LLM outputs are stochastic, generative, and span a continuous space of potential actions.
- Contributions: The paper develops and validates an LLM–UCT-MCTS reasoning algorithm on Olympic-level mathematical problems.The authors describe this as a novel reasoning algorithm whose key components are enhanced for LLM integration.
- Results: Extensive experiments report improved performance in complex reasoning tasks, including mathematical benchmarks.The experiments are presented as evidence of synergistic potential between LLMs and MCTS.
- Significance: The work advances LLM application to sophisticated reasoning challenges and supports future integration of AI for decision-making and reasoning accuracy.This conclusion is stated within the paper’s scope of LLM-driven applications.
2 Preliminary
The preliminary section frames MCTS as a staged search procedure and maps MCTSr onto a tree of iteratively refined mathematical answers. Its notation defines answer nodes, refinement actions, self-rewards, termination, value estimates, and exploration control.
- MCTS: MCTS builds a search tree by repeating Selection, Expansion, Simulation or Evaluation, and Backpropagation to estimate action values.The repeated stages refine strategies when directly calculating the best strategy is infeasible because of a vast state space.
- MCTS: Selection navigates from the root toward promising children, while Expansion adds feasible child nodes at a nonterminal leaf.These stages use strategies such as UCT to choose promising paths and represent possible future moves.
- MCTS: Simulation evaluates a newly added node through rollout moves, and Backpropagation sends the outcome back to update traversed-node statistics.The outcome can be a win, loss, or draw, informing later decisions.
- Selection: UCT balances exploration and exploitation by combining an action’s average reward with a visit-based confidence term.The formula uses the father-node visit count, child-node visit count, and a balancing constant.
- MCTSr: MCTSr represents mathematical answers as tree nodes and self-refinement attempts as edges, sampling self-rewards for answer versions.Its workflow follows the general MCTS pattern while using self-reflective improvement and model-based self-rewarding.
- Notation: The notation defines P as the problem, A as answer nodes, M as refinement actions, R as the self-reward function, and T as the termination function.Additional symbols define reward samples, node values, confidence bounds, parent nodes, child nodes, and visit counts.
3 Methodology
MCTSr combines Monte Carlo Tree Search with iterative self-refinement of mathematical answers. It repeatedly evaluates, updates, and selects candidate answers while using reward constraints, child-node values, and UCT criteria to guide search until termination.
- Algorithm structure: The workflow initializes a root answer, selects an under-expanded high-value node, refines it through feedback, evaluates the result, backpropagates updated values, and updates UCT scores.These stages repeat until a termination condition is met.
- Self-evaluation: Self-reward scores range from -100 to 100, with strict scoring prompts, suppression above 95, and repeated sampling designed to improve comparative reliability.A minimum-reward constraint further counteracts the smoothing tendency of the self-reward function.
- Self-evaluation: The Q value combines the minimum and mean sampled rewards, balancing worst-case and average outcomes when estimating answer quality.The resulting estimate is propagated through the tree when child values change.
- Backpropagation: Backpropagation updates a node using its current Q value and the highest Q value among its children, incorporating promising subsequent answers into evaluation.This updated value reflects both direct reward samples and the best child outcome.
- UCT update and selection: Because LLMs can generate infinitely many refinement actions, full expansion is defined by a child-count limit or a child whose Q value exceeds the parent’s.Candidate nodes are then ranked with UCT, balancing exploration and exploitation before greedy or importance sampling selects the next node.
- Termination: Search terminates when improvements diminish, outcomes repeat, rollout or depth limits are reached, or predefined language-model-logit criteria are satisfied.The best answers are then gathered according to Q values or other termination conditions.
4 Evaluation
MCTSr was evaluated across GSM, MATH, and Olympiad-level benchmarks using LLaMA3-8B configurations with different rollout counts. Results generally improved with more rollouts, while performance remained constrained on the hardest problems.
- Evaluation setup: The evaluation compared Zero-Shot CoT, Self-Refine, 4-rollouts MCTSr, and 8-rollouts MCTSr on GSM8K, GSM-Hard, MATH, and Olympiad-level datasets.LLaMA3-8B results were also compared with GPT-4, Claude 3, and Gemini 1.5-Pro.
- MATH benchmark: 90.16% success was reached on MATH level 1 with 8-rollouts MCTSr, while level 5 reached 34.06%, indicating reduced performance on the hardest problems.The level-1 result solved 394 of 437 problems; level 5 solved 451 of 1324.
- MATH benchmark: 58.24% cumulative success was achieved by 8-rollouts MCTSr on MATH, solving 2912 of 5000 problems, versus 24.36% for Zero-Shot CoT.Increasing rollouts consistently correlated with higher success rates across difficulty levels.
- Olympiad-level benchmarks: Olympiad-level success increased from Zero-Shot CoT to 8-rollouts MCTSr on AIME, GAIC Math Odyssey, and OlympiadBench.The reported increases were 2.36% to 11.79% on AIME, 17.22% to 49.36% on GAIC Math Odyssey, and 1.25% to 7.76% on OlympiadBench.
- Cross-benchmark findings: The findings associate more rollouts with higher success rates and describe MCTSr as effective across varying mathematical complexities, including unseen problems.GAIC Math Odyssey was highlighted as evidence of generalization, while the broader results also identify boundaries on more complex tasks.
- Model comparison: MCTSr was reported to enhance small open-source models such as LLaMA-3 to a comparable level with current closed-source large models on the tested benchmarks.The comparison used reported state-of-the-art performance values for closed-source models.
5 Related works
Related work applies MCTS and other learning or refinement strategies to complex planning and mathematical reasoning. These studies motivate combining search with LLM-based reasoning for difficult multi-step problems.
- MCTS applications: MCTS has been applied to complex domains including multi-agent pathfinding and train timetabling, alongside heuristic and learning-based methods.Prior work compared MCTS-based approaches with heuristic search such as A* and combined MCTS with supervised or unsupervised learning.
- LLM reasoning: Recent LLM research improves mathematical reasoning through collective answer refinement, reinforcement learning from Evol-Instruct Feedback, and visual-mathematical benchmarking.The cited examples include multi-LLM discussion, WizardMath, and MathVista.
- Search-enhanced reasoning: MCTS has also been incorporated into mathematical reasoning systems to address logical and numerical errors in multi-step problems without additional fine-tuning.The passage identifies this direction as part of recent efforts to strengthen LLM mathematical reasoning.
6 Limitations
The paper characterizes MCTSr as preliminary and leaves its broader applicability and component choices for future investigation.
- Scope and maturity: MCTSr remains at a preliminary research stage, with applications beyond mathematical tasks still requiring exploration.Suggested areas include black-box optimization and self-driven alignment for large language models.
- Future development: The algorithm’s scalable components require further development and comparison to improve its practical effectiveness.The authors call for broader evaluation of alternative component algorithms.
7 Conclusion
MCTSr improves LLM performance on complex mathematical problems by integrating Monte Carlo Tree Search with LLM-based reasoning. The paper reports higher problem-solving success rates across multiple datasets while noting that broader applicability remains unexplored.
- MCTSr enhances LLM capability for solving complex mathematical problems by integrating Monte Carlo Tree Search.
- The algorithm’s applicability to broader contexts, including black-box optimization and self-driven alignment, remains to be explored.Future work will optimize algorithmic components and test them across varied problems and settings.
A Prompts in Experiment
The experiments use prompts for iterative answer refinement, strict self-evaluation, and handling dummy answers. These prompts operationalize the Self-Refine, Self-Reward, and Dummy Answers components.
- A.1 Self-Refine: Self-Refine asks the model to critique a weak answer and identify every possible flaw before correction.
- A.1 Self-Refine: The refinement prompt instructs the model to revise its answer using feedback and follow specified reasoning, verification, and final-answer markers.
- A.2 Self-Reward: Self-Reward prompts the model to analyze an answer strictly and assign a score from -100 to +100.
- A.2 Self-Reward: The scoring prompt emphasizes harsh grading and discourages awarding full marks.
- A.3 Dummy Answers: Dummy Answers include several standardized responses expressing inability to understand, solve, or answer a question.
- A.3 Dummy Answers: The dummy-answer list includes multiple variants of “I don’t know” and “I can’t help with this question.”