BG-MCTS — Budget-Guided Tree Search for Fixed Token Budgets in LLM Reasoning

Paper. Sora Miyamoto, Daisuke Oba, Naoaki Okazaki. Aligning Tree-Search Policies with Fixed Token Budgets in Test-Time Scaling of LLMs. ICML 2026. Institute of Science Tokyo. [arXiv]

0. The Picture in One Paragraph

Test-time scaling via tree search requires an exploration strategy, but most MCTS formulations for LLM reasoning ignore a critical practical constraint: the total token budget is fixed. Standard UCT-based MCTS explores aggressively throughout, even when it is about to run out of budget — generating partial expansions that never reach a complete answer and wasting tokens. Budget-Guided MCTS (BG-MCTS) solves this with a single scalar: the budget ratio ρ = 1 − C_used / B (remaining fraction of the budget). As ρ decreases, three behaviors automatically adjust: exploration annealing (less UCT bonus → more exploitation), completion bias (longer completions → higher probability of reaching a final answer before budget runs out), and widen annealing (harder to open new branches → favor depth over breadth). The net effect is a wide-to-deep schedule: explore broadly early, exploit and complete deeply late. BG-MCTS is training-free — it sits atop any frozen LLM + process reward model. On MATH500 Level-5 problems at a 10K-token budget, BG-MCTS achieves 0.662 vs. beam search’s 0.328 — more than double — with consistent gains on AIME24/25 and UGPhysics across Llama and Qwen backbone models.

1. The Problem — Budget-Blind Tree Search

Tree search scales inference compute by exploring multiple reasoning paths and selecting the best via a process reward model (PRM). MCTS with UCT is the standard approach:

UCT score = Q(v) + c · √(ln N(parent) / N(v))

where Q(v) is the estimated value, N(v) is the visit count, and c is the exploration constant. UCT does not know or care about the remaining token budget.

This causes a systematic failure under fixed budget constraints:

The cost is high: unfinished branches produce no answer, meaning budget spent on them yields zero reward. A good tree search policy under a fixed budget should front-load exploration and back-load exploitation and completion.

No prior LLM reasoning tree search method addresses this directly. BG-MCTS is the first training-free formulation that makes budget awareness a first-class component of the search policy.

2. The Budget Ratio — One Number, Three Behaviors

BG-MCTS introduces a single quantity:

where $B$ is the total output token budget and $C_{\text{used}}$ is the number of output tokens spent so far. ρ ∈ [1, 0]: starts at 1 (full budget remaining), hits 0 when the budget is exhausted.

This ratio drives three coupled modifications to the standard BG-PUCT (budget-guided PUCT, the practical UCT variant):

2.1 Exploration Annealing

The UCT exploration bonus is scaled by ρ:

When ρ is high (budget plentiful), the exploration bonus is near its full value — the search explores widely. As ρ decreases, the exploration bonus shrinks, and the search increasingly exploits high-value nodes. At ρ ≈ 0, the search is nearly greedy.

This is automatic curriculum annealing: no hyperparameter schedule is needed. The budget ratio provides the signal.

2.2 Completion Bias

When sampling a continuation from a node, BG-MCTS biases the length of that continuation toward completing the reasoning trace. Specifically, a completion bias term proportional to (1 − ρ) is applied — as the budget depletes, the sampled continuations become longer and more likely to reach a final answer.

The mechanism: at high ρ, short partial continuations are acceptable (there is budget to expand later). At low ρ, partially expanded traces are likely to be abandoned, so the policy should try to complete them now.

2.3 Widen Annealing

MCTS “widening” refers to opening a new child node (a new reasoning branch). In standard MCTS, widening is relatively unconstrained. BG-MCTS gates widening by ρ: when ρ is high, new branches are easy to open; when ρ is low, widening is suppressed.

The effect is the defining behavior of BG-MCTS: a wide-to-deep schedule emerges automatically without any manual parameter tuning.

Early in budget (ρ ≈ 1):   Wide exploration — many branches opened
Late in budget (ρ ≈ 0):    Deep exploitation — existing branches deepened to completion

3. Why the Three Components Work Together

The three modifications address three distinct failure modes of standard MCTS under a fixed budget:

Each component is independently necessary: the ablation results confirm that removing any one of the three degrades performance. Together they implement a coherent policy: explore cheap paths early, lock in value late.

4. Results

4.1 MATH500 Level-5 (Hardest Math Problems)

At a 10K-token budget:

More than 2× improvement over beam search on hard math problems — entirely training-free, using the same frozen LLM and PRM.

4.2 Competition Math (AIME24/25) and Science (UGPhysics)

BG-MCTS consistently outperforms beam search across:

• AIME 2024 and AIME 2025 competition problems
• UGPhysics (undergraduate-level physics)

Gains are consistent across backbone models:

• Llama-3.1-8B-Instruct
• Qwen2.5-7B-Instruct
• Qwen3-32B

This suggests BG-MCTS’s benefit is not specific to a particular model family or reasoning style.

4.3 Ablation Results

All three components are individually necessary; removing any one degrades performance on at least one benchmark. No component is redundant.

5. Relationship to Prior Test-Time Scaling Methods

BG-MCTS sits within the inference-time compute scaling literature but occupies a distinct niche:

PRA (review) and Reasoning Cache (review) also improve reasoning under fixed compute budgets, but via trained components (reward agents, replay buffers). BG-MCTS adds value through inference-time policy adjustment alone — no training required.

6. Why It Matters

Three reasons:

1. Budget-awareness closes the gap between MCTS theory and practical deployment. In practice, LLM inference has a hard cost ceiling — you cannot run MCTS until convergence. BG-MCTS is the first tree search formulation that treats the budget constraint as a first-class design element, not an afterthought.
2. Training-free means drop-in compatibility. Any existing LLM + PRM stack can adopt BG-MCTS without any model modification or additional training data. The budget ratio is computable from token counters already present in any inference system.
3. The wide-to-deep scheduling principle is general. The insight that exploration should decrease and completion should increase as budget depletes is not specific to LLM reasoning — it likely applies to any tree search operating under a token or compute budget.

7. Limitations Worth Knowing

• Output-only budget tracking. ρ counts output tokens only. Input tokens from retrieved documents, long system prompts, or tool outputs are not included. For tasks with variable-length context (RAG, tool use), the true budget consumption is underestimated.
• PRM dependency. BG-MCTS requires a process reward model to score intermediate steps. PRM quality is a ceiling on BG-MCTS performance — a poorly calibrated PRM can cause the search to exploit bad paths.
• Fixed-budget assumption. The budget ratio ρ is only meaningful when the total budget B is fixed in advance. Adaptive budgeting (allocate more budget to harder problems) is not addressed.
• No wall-clock latency analysis. The paper reports token counts, not wall-clock time. Sequential MCTS expansion may have higher latency than parallel best-of-N at the same token budget.

8. The Takeaway for a First Reader

If you remember three things:

1. BG-MCTS adds a budget ratio ρ = 1 − C_used/B to standard MCTS, using it to drive three automatic adaptations: exploration annealing (less UCT bonus as budget depletes), completion bias (longer completions late in the budget), and widen annealing (fewer new branches late). Together these produce a wide-to-deep schedule — no training, no manual schedule required.
2. All three components are necessary. Ablations confirm that removing any single component reduces performance — the three modifications address distinct failure modes of budget-blind MCTS.
3. On MATH500 Level-5 at 10K tokens, BG-MCTS scores 0.662 vs. beam search’s 0.328 — more than 2× improvement. Gains hold across AIME, UGPhysics, and three backbone model families (Llama-3.1-8B, Qwen2.5-7B, Qwen3-32B).

References

• Miyamoto, S., Oba, D., & Okazaki, N. (2026). Aligning Tree-Search Policies with Fixed Token Budgets in Test-Time Scaling of LLMs. ICML 2026. arXiv:2602.09574.
• Related on this site: MCTD review — MCTS embedded in the diffusion denoising process, a complementary use of tree search for sequential generation; Reasoning Cache — iterative summarize-and-restart as an alternative extrapolation strategy under fixed budgets; PRA — online process reward signals within beam search for knowledge-intensive reasoning.