Research map / Large language models / Reasoning and chain of thought
Policy Guided Tree Search for Enhanced LLM Reasoning
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 PGTS, a PPO-trained graph-transformer (GPS) policy that chooses expand, branch, backtrack, or terminate actions over an LLM reasoning tree, aiming for CoT/MCTS-level accuracy at lower token cost.
Problem
LLM reasoning on multi-step math, logic and planning tasks relies on predefined heuristics and expensive tree search (e.g., MCTS) with costly self-evaluation rewards.
Approach
A learned RL policy over a tree-search MDP decides among expand, branch, backtrack (1 to D-1 steps), and terminate, implemented as a GPS graph transformer over LLM hidden-state node features and reward edge features, with invalid actions masked by a constraint vector.
Key results
- On MATH with LLaMA3.1-8B, PGTS reaches 41.00% versus CoT's 34.40%, while using about 5.28x CoT tokens versus 16.25x for MCTS.
- PGTS (SC8) with LLaMA3.1-8B reaches 89.99% on GSM8K and 52.20% on MATH, and 36.84% on Blocksworld (4 steps) versus CoT-SC8's 26.32%.
- PGTS underperforms MCTS on GPQA (18.69% vs 34.34% for LLaMA3.1-8B), which the authors attribute to task complexity and limited training data.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use PGTS with self-consistency voting over 8 chains instead of MCTS (Oracle) | LLaMA3.1-8B · GSM8K, MATH, AQUA, StrategyQA, PrOntoQA, GPQA, Blocksworld (4), Blocksworld (8) | accuracy: GSM8K 88.78 -> 89.99; MATH 51.40 -> 52.20; AQUA 64.96 -> 66.54; StrategyQA 84.40 -> 85.70; PrOntoQA 74.80 -> 77.40; GPQA | Mixed |
| add self-consistency voting over 8 chains to PGTS | PGTS with LLaMA3.1-8B · GSM8K | accuracy: 85.29 -> 89.99 | Helped |
| add self-consistency voting over 8 chains to PGTS | PGTS with LLaMA3.1-8B · MATH | accuracy: 41.00 -> 52.20 | Helped |
| add edge features to the GPS policy (ablation) | PGTS · AQUA | accuracy: 50.79 -> 60.63 | Helped |
| add edge features to the GPS policy (ablation) | PGTS · GSM8K | accuracy: 81.96 -> 85.29 | Helped |
| add global attention to the GPS policy (ablation) | PGTS · AQUA | accuracy: 54.72 -> 60.63 | Helped |
| add global attention to the GPS policy (ablation) | PGTS · GSM8K | accuracy: 84.15 -> 85.29 | Helped |
| add local MPNN to the GPS policy (ablation) | PGTS · AQUA | accuracy: 58.76 -> 60.63 | Helped |
| add local MPNN to the GPS policy (ablation) | PGTS · GSM8K | accuracy: 84.78 -> 85.29 | Helped |
| replace GPS layers with SAN (ablation) | PGTS · AQUA and GSM8K | accuracy: AQUA 60.63 -> 60.24; GSM8K 85.29 -> 82.56 | Hurt |
| replace graph policy with distilbert-base-uncased policy (ablation) | PGTS · AQUA and GSM8K | accuracy: AQUA 60.63 -> 53.94; GSM8K 85.29 -> 82.87 | Hurt |
| replace graph policy with an LLM agent policy (ablation) | PGTS · AQUA and GSM8K | accuracy: AQUA 60.63 -> 55.51; GSM8K 85.29 -> 83.32 | Hurt |
Surprising findings
- PGTS is below CoT on PrOntoQA for LLaMA3.1-8B (68.20 vs 69.40) and is far below MCTS on GPQA.
- Removing edge features from the GPS policy drops AQUA accuracy from 60.63 to 50.79.
Limitations
- Underperforms MCTS on GPQA and lacks gains in some settings.
- Intermediate reward is only log-likelihood of each step, not a stronger self-evaluation or PRM.
- Depth and breadth limits and action costs are hand-set, with no dataset-specific tuning.
Experimental setup
Datasets: GSM8K MATH500 AQUA StrategyQA PrOntoQA GPQA Blocksworld
Benchmarks: GSM8K MATH500 AQUA StrategyQA PrOntoQA GPQA Blocksworld
Metrics: accuracy generated token count
Baselines: CoT CoT (SC4) CoT (SC8) MCTS MCTS (Oracle) SAN policy SLM policy LLM Agent policy
Builds on: Chain-of-Thought Monte Carlo Tree Search RAP GPS graph transformer PPO
Closest papers by meaning
- Plan Then Action: High-Level Planning Guidance Reinforcement Learning for LLM Reasoning (ICML 2026 · similarity 0.92)
- Language Agent Tree Search Unifies Reasoning, Acting, and Planning in Language Models (ICML 2024 · similarity 0.91)
- ReJump: A Tree-Jump Representation for Analyzing and Improving LLM Reasoning (ICML 2026 · similarity 0.90)
- DeepSearch: Overcome the Bottleneck of Reinforcement Learning with Verifiable Rewards via Monte Carlo Tree Search (ICLR 2026 · similarity 0.90)
- STRIDE: Learnable Stepwise Language Feedback for LLM Reasoning (NeurIPS 2026 · similarity 0.90)
- Rectifying LLM Thought from Lens of Optimization (ICLR 2026 · similarity 0.90)
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.