Research map / Large language models / Agents, code and math
AlphaZero-Like Tree-Search can Guide Large Language Model Decoding and Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TS-LLM, an AlphaZero-like tree-search framework that uses a learned LLM value function and outcome reward model to guide LLM decoding at inference and to iteratively improve the policy and value via tree-search-augmented training.
Problem
Prompt-based value functions in tree-search methods (ToT, RAP) are unreliable, depend on strong LLMs, and work only for shallow trees; they also cannot use in-domain training to improve the LLM.
Approach
Replace prompted self-evaluation with a learned value network and outcome reward model trained on the LLM's own rollouts, and use AlphaZero-style MCTS variants (MCTS-α, MCTS-Rollout) at sentence or token level; tree-search outputs then serve as policy-improvement data for distillation and value learning in an iterative loop.
Key results
- On Game24 with LLaMA-SFT policy, learned LLaMA-V value gives 64.8% Path@1 vs 21.0% with GPT-3.5 prompted value (ToT).
- MCTS-α reaches 2.221 reward on RLHF versus 0.318 for CoT-greedy, and MCTS-Rollout reaches 98.76 win-rate-style score on chess endgame versus 58.14 for CoT-greedy.
- On GSM8K one iteration of training lifts MCTS-α from 51.9% to 56.5%, though tree-search costs many more tokens than CoT.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use learned-value BFS-V instead of CoT-SC with ORM aggregation | LLaMA2-7B · GSM8K | accuracy (%): 52.3 -> 52.5 | No clear effect |
| Use MCTS-Rollout instead of CoT-SC with ORM aggregation | LLaMA2-7B · Game24 | accuracy (%): 50.6 -> 71.3 | Helped |
| Use MCTS-α instead of CoT-SC with ORM aggregation | LLaMA2-7B · PrOntoQA | accuracy (%): 83.2 -> 99.4 | Helped |
| Use MCTS-α instead of CoT-SC with ORM aggregation | GPT-2-small (125M) · RLHF | reward: 1.517 -> 2.221 | Helped |
| Use MCTS-Rollout instead of CoT-SC with ORM aggregation | GPT-2-small (125M) · Chess Endgame | win rate: 73.80 -> 98.76 | Helped |
| Replace prompt-based GPT-3.5 value evaluation with the learned LLaMA-V value function | GPT-3.5 policy · GSM8K | accuracy (%): 72.7 -> 74.0 | Helped |
| Replace prompt-based GPT-3.5 value evaluation with the learned LLaMA-V value function | LLaMA-SFT policy · Game24 | accuracy (%): 21.0 -> 64.8 | Helped |
| Add policy distillation from tree-search trajectories | LLaMA2-7B with greedy decoding · GSM8K | accuracy (%): 41.4 -> 47.9 | Helped |
| Update both policy and value/ORM through tree-search training | LLaMA2-7B with MCTS-α · GSM8K | accuracy (%): 51.9 -> 56.5 | Helped |
| Update both policy and value/ORM through tree-search training | GPT-2-small (125M) with MCTS-α · RLHF | reward: 2.22 -> 2.67 | Helped |
| Use mixed checkpoints rather than pure last-checkpoint data to train the ORM (ablation) | CoT-SC ORM-vote · GSM8K | accuracy (%): 55.3 ± 0.5 -> 55.9 ± 0.7 | Helped |
| Use full pure data instead of one-third pure data to train the value function (ablation) | BFS · GSM8K | accuracy (%): 50.0 ± 0.3 -> 52.7 ± 0.8 | Helped |
Surprising findings
- Under equal token budgets, CoT-SC with ORM aggregation matches or beats TS-LLM on GSM8K.
- Value functions trained on mixed or accumulated data outperform those continually fine-tuned from the previous value model, contrary to typical RL practice.
Limitations
- Tree search incurs large extra computation from node expansion and value evaluation, limiting scalability.
- Aggregation of multiple tree searches benefits less than CoT-SC on small-scale problems.
- Only one iteration of training was run in the main experiments.
Experimental setup
Datasets: GSM8k Game24 PrOntoQA RLHF synthetic data Chess endgame
Benchmarks: GSM8k Game24 PrOntoQA RLHF (synthetic) Chess Endgame
Metrics: Path@1 accuracy Path@N accuracy token count reward win rate
Baselines: CoT greedy CoT-SC ToT-GPT3.5 RFT PPO BFS-V DFS-V MCTS
Builds on: Tree-of-Thought Reasoning via Planning AlphaZero rejection sampling fine-tuning
Closest papers by meaning
- TSLM: Tree-Structured Language Modeling for Divergent Thinking (ICLR 2026 · similarity 0.90)
- Chain of Preference Optimization: Improving Chain-of-Thought Reasoning in LLMs (NeurIPS 2024 · similarity 0.90)
- Advancing Tool-Augmented Large Language Models: Integrating Insights from Errors in Inference Trees (NeurIPS 2024 · similarity 0.90)
- Aligning Tree-Search Policies with Fixed Token Budgets in Test-Time Scaling of LLMs (ICML 2026 · similarity 0.89)
- Language Agent Tree Search Unifies Reasoning, Acting, and Planning in Language Models (ICML 2024 · similarity 0.89)
- Policy Guided Tree Search for Enhanced LLM Reasoning (ICML 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.