Model tool chain selection as a search tree.

Root: (query, empty_context)
  ├── OCR → (query, ocr_context)
  │   ├── Retrieval → (query, ocr+ret_context)
  │   │   └── STOP ✓  reward=0.82
  │   ├── Code Sandbox → (query, ocr+code_context)
  │   │   └── STOP ✓  reward=0.71
  │   └── STOP ✓  reward=0.65
  ├── Retrieval → (query, ret_context)
  │   ├── STOP ✓  reward=0.73
  │   └── Web Crawl → ...
  └── STOP ✓  reward=0.30

The STOP action is the key insight. MCTS naturally learns when to stop adding tools — when the marginal quality gain from another tool invocation is less than the cost penalty. This directly solves the "over-building context" problem.

Reward function: R = α·quality - β·cost - γ·latency + δ·brevity

Parameter	Weight	Why
quality	1.0	Lightweight value model estimates answer quality from current context
cost	0.3	Penalizes expensive tool invocations
latency	0.1	Penalizes slow tool chains
brevity	0.05	Rewards staying within token budget

Expected impact: 20-30% cost reduction with equal or better quality.

MCTS Reasoning Search (Highest Accuracy Gain)

Problem: The reasoning head is binary (on/off), single-pass.

Solution: For hard queries (AIME, GPQA, LiveCodeBench), tree search over reasoning paths.

Root: "Prove √2 is irrational"
  ├── "Assume √2 = p/q in lowest terms..."
  │   ├── "Then 2q² = p², so p² is even..."
  │   │   └── "Therefore p is even. Let p = 2k..." → ✓ correct
  │   └── "Consider p mod 4..." → dead end, backtrack
  ├── "Use the rational root theorem..."
  │   └── "Apply to x² - 2 = 0..." → ✓ correct (different path)
  └── "By unique prime factorization..." → ✓ correct (third path)

Three completed paths, all agreeing → self-consistency ratio = 1.0 → high confidence.

This is the same approach that powered:

• rStar-Math: 90% AIME with 7B models (vs Interfaze's 90% with frontier LLMs)
• AlphaProof: IMO silver medal via tree search over proof steps
• SWE-Search: 23% relative improvement on SWE-bench

Expected impact: 3-8 point improvement on AIME-2025, GPQA-Diamond, and LiveCodeBench.

MCTS Context Optimization (Cost Efficiency)

Problem: Fixed token budgets with heuristic scoring.

Solution: Search over which observations/entities/relations to include.

Instead of heuristic relevance scoring, MCTS explores context compositions and evaluates them against a value model. Naturally finds the most informative subset within the token budget.

Expected impact: Better context quality within same or lower token budget. Reduces the "dispersed evidence" degradation mentioned in the paper.

---

Projected Benchmark Impact

Benchmark	Current	Projected	Source of Gain
AIME-2025	90.0%	93-95%	Reasoning MCTS (rStar-style search)
GPQA-Diamond	81.3%	84-87%	Reasoning MCTS + better context selection
MMLU-Pro	83.6%	84-85%	Context optimization (less noise)
LiveCodeBench v5	57.8%	62-65%	Reasoning MCTS with code sandbox feedback
Cost per query	baseline	-20-30%	Controller MCTS avoids unnecessary tools
Avg latency	baseline	-15-25%	Fewer tool invocations on easy queries

---

Why MCTS Specifically?

Alternative	Why MCTS wins
Beam search	No exploration/exploitation balance. Gets stuck in local optima.
Best-of-N sampling	No tree structure. Can't reuse partial computations.
RL fine-tuning	Requires expensive training loop. MCTS works at inference time with no training.
Greedy + heuristics	Current approach. Doesn't explore combinatorial space of tool chains.

MCTS advantages:

• Works at inference time — no additional model training required
• Anytime algorithm — can stop early if time budget is hit, still returns best-so-far
• Naturally balances explore vs exploit via UCB/PUCT
• Composable — each integration point uses the same core engine with different expanders/simulators

---

Implementation

Run the demo:

python -m examples.demo

Run tests:

python -m pytest tests/ -v

---

References

1. rStar-Math: "Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking" (arXiv:2501.04519, 2025)
2. ToolTree: "Deliberate Tool Selection for LLM Agent with Monte Carlo Tree Search" (OpenReview, 2025)
3. MASTER: "A Multi-Agent System with LLM Specialized MCTS" (NAACL 2025)
4. SWE-Search: "Enhancing Software Agents with Monte Carlo Tree Search" (OpenReview, 2025)
5. AlphaProof: Silver medal at IMO 2024 via MCTS over proof steps (DeepMind, 2024)