💬 NLP

Aligning Tree-Search Policies with Fixed Token Budgets in Test-Time Scaling of LLMs

The paper proposes Budget-Guided MCTS (BG-MCTS), a tree-search decoding algorithm that dynamically aligns exploration and refinement strategies with the remaining token budget to outperform budget-agnostic baselines in mathematical and physics reasoning tasks.

Original authors: Sora Miyamoto, Daisuke Oba, Naoaki Okazaki

Published 2026-06-05

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Sora Miyamoto, Daisuke Oba, Naoaki Okazaki

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a very tricky riddle, but you have a strict rule: you can only ask a certain number of questions before the clock runs out. This is exactly the challenge Large Language Models (LLMs) face when solving complex problems like math or physics. They have a "token budget"—a limit on how many words or steps they can generate.

The paper introduces a new method called BG-MCTS (Budget-Guided Monte Carlo Tree Search) to help these AI detectives solve problems better within that strict time limit.

Here is how it works, using simple analogies:

The Problem: The "One-Size-Fits-All" Detective

Currently, most AI search methods act like a detective who has a fixed plan regardless of how much time is left.

The Old Way: The detective spends the first half of the day asking 100 different people for clues (broad exploration). Then, in the last 10 minutes, they realize they have no time left to actually follow up on the most promising leads. They might start a new line of questioning right before the clock hits zero, leaving the case unsolved. Or, they might stop too early, wasting the last 10 minutes of their shift.
The Issue: Existing methods treat the time limit (token budget) only as a "stop sign." They don't change their strategy based on how much time is remaining.

The Solution: The "Smart Detective" (BG-MCTS)

The authors propose a detective who constantly checks their watch and changes their strategy based on how much time is left. They call this Budget-Guided MCTS.

Think of the search process as a tree growing from a root:

Early Stage (Plenty of Time Left): When the detective has 100% of their time, they act like a wide-net fisherman. They cast a wide net, exploring many different shallow paths to see where the fish might be. They don't dive deep yet; they just want to see the whole ocean.
Late Stage (Time Running Out): As the clock ticks down (say, to 25% of the budget), the detective stops casting wide nets. Instead, they pick the two or three most promising spots they found earlier and dive deep. They stop starting new lines of questioning and focus entirely on finishing the investigation on the best leads.

How the AI Does It

The paper describes two specific tricks the AI uses to pull this off:

The "Time-Check" Score: When the AI decides which path to follow next, it uses a formula that looks at how much budget is left.
- If there is lots of budget, the formula encourages trying new, unexplored paths.
- If the budget is low, the formula punishes starting new paths and rewards going deeper into paths that already look good.
The "New Branch" Switch: The AI has a special switch that decides whether to grow a new branch on the tree or just go deeper down an existing branch.
- When time is plentiful, the switch is set to "Grow New Branches."
- When time is running out, the switch flips to "Go Deeper," preventing the AI from wasting its last few seconds starting a new branch that it won't have time to finish.

The Results

The researchers tested this "Smart Detective" against other methods on difficult math and physics problems. They found that:

Better Accuracy: The AI solved more problems correctly within the same token limit.
No Wasted Time: Unlike other methods that might stop too early or start too many new paths at the end, BG-MCTS used the entire budget efficiently. It explored broadly at the start and finished strong at the end.
Consistent Performance: This worked well across different types of AI models and different problem difficulties.

The Bottom Line

The paper claims that by making the AI's search strategy "aware" of the remaining budget, we can get much better answers without needing more computing power. It's like teaching a runner not just to run fast, but to know exactly when to sprint and when to conserve energy so they cross the finish line with the best possible time.

Technical Summary: Aligning Tree-Search Policies with Fixed Token Budgets in Test-Time Scaling of LLMs

Problem Statement

Large Language Models (LLMs) can improve response quality through Test-Time Scaling, which allocates additional compute during inference without updating model parameters. While tree-search decoding (e.g., Monte Carlo Tree Search, MCTS) is an effective form of test-time scaling, real-world deployments often impose a fixed per-query token budget.

Existing tree-search policies are largely budget-agnostic. They typically treat the token budget merely as a termination condition rather than a dynamic constraint guiding search behavior. This mismatch leads to two primary failure modes:

Late-stage over-branching: The search continues to expand shallow nodes late in the process, exhausting the budget before it can refine or complete promising candidates.
Premature termination: The search stops early, leaving a portion of the allocated budget unused.

Current early-stopping variants focus on saving tokens but lack a policy to dynamically shift the search strategy from broad exploration to refinement as the remaining budget decreases.

Methodology: Budget-Guided MCTS (BG-MCTS)

The authors propose Budget-Guided MCTS (BG-MCTS), a tree-search decoding algorithm that explicitly conditions its search policy on the remaining token budget. The core mechanism relies on a budget sufficiency ratio $\rho \in [0, 1]$ , defined as:
$\rho = 1 - \frac{C_{used}}{B}$
where $C_{used}$ is the cumulative tokens generated and $B$ is the total budget. As $\rho$ decreases (approaching 0), the policy shifts from exploration to exploitation and completion.

BG-MCTS introduces two budget-conditioned control mechanisms:

1. Budget-Guided Selection (BG-PUCT)

Standard PUCT selection is modified to anneal the exploration bonus and introduce a completion bias:
$\text{BG-PUCT}(p, s, \rho) = \frac{\tilde{W}(s, \rho)}{m_s} + \rho \cdot c \cdot P(s|p) \sqrt{\frac{\ln(m_p)}{m_s}}$

Exploration Annealing: The exploration term is multiplied by $\rho$ . As the budget depletes ( $\rho \to 0$ ), the incentive to explore new branches diminishes.
Completion Bias: The accumulated value $\tilde{W}(s, \rho)$ is corrected by a term proportional to node depth $d(x)$ and $(1-\rho)$ . This biases selection toward deeper nodes that are closer to generating a final answer, without artificially boosting nodes that already contain a completed answer.

2. Budget-Guided Tree Widening

To prevent wasteful late-stage branching, BG-MCTS treats "widening" (generating a new child from a parent) as a selectable action alongside existing children. A virtual generative child $s_{gen}(p)$ is added to the selection set. Its score is:
$E_{gen}(p, \rho) = \mu(p) + \lambda \cdot \rho \cdot \sigma^2(p)$

$\mu(p)$ and $\sigma^2(p)$ represent the mean and variance of existing children's scores.
The term is scaled by $\rho$ . When the budget is ample, widening is encouraged (especially if children disagree). As the budget runs low, the incentive to widen is suppressed, forcing the search to deepen existing promising trajectories.

Unified Algorithm

The algorithm follows a standard MCTS loop but makes both selection and widening decisions dependent on $\rho$ . This induces a wide-to-deep schedule: broad exploration early in the search, followed by focused refinement and answer completion as the budget is exhausted.

Key Contributions

Budget-Aware Policy Design: The paper introduces the first tree-search decoding framework that dynamically aligns search behavior (exploration vs. exploitation vs. widening) with the remaining token budget, rather than treating the budget as a static stop signal.
Novel Mechanisms: The proposal of BG-PUCT (with budget-conditioned exploration and depth-biased value correction) and Virtual Generative Children for budget-aware widening.
Empirical Validation: Comprehensive evaluation across mathematical reasoning benchmarks (MATH500, AIME24/25) and a physics reasoning benchmark (UGPhysics) using open-weight LLMs (Llama-3.1-8B, Qwen2.5-7B, Qwen3-32B).

Experimental Results

The authors evaluated BG-MCTS against budget-agnostic baselines (Standard MCTS, AB-MCTS-M, LiteSearch) and scaling baselines (Repeated Sampling, Sequential Refinement) under fixed budgets of 10k, 20k, and 30k tokens.

Performance: BG-MCTS consistently outperformed all baselines across all models and budgets. For example, on MATH500 with Qwen2.5-7B, BG-MCTS achieved 71.1% accuracy at 30k tokens, compared to 65.9% for the next best tree-search baseline (MCTS).
Budget Utilization: Unlike LiteSearch, which often stops prematurely, BG-MCTS continues to improve accuracy as it approaches budget exhaustion.
Qualitative Behavior: Analysis of search trees revealed that BG-MCTS maintains broad width early on but shifts to deepening specific branches late in the search. In contrast, standard MCTS often continues widening shallow nodes until the budget is exhausted, failing to complete promising paths.
Physics Reasoning: The method's effectiveness extended to UGPhysics, demonstrating that the budget-aligned wide-to-deep behavior is beneficial beyond pure mathematical reasoning.

Significance and Claims

The paper claims that conditioning tree-search decisions on the remaining budget is more effective than treating the budget merely as a stopping condition.

Predictable Scaling: BG-MCTS offers a simple interface for predictable test-time scaling, where users specify a token budget, and the policy automatically adapts its exploration-to-completion schedule within that cost envelope.
Data Synthesis Implications: The authors note that BG-MCTS tends to produce fewer answered nodes per problem instance but with significantly higher precision (correctness). This "fewer-instances, better-answers" regime is highlighted as particularly valuable for data synthesis, where high-quality reasoning traces are more critical than raw coverage.
Robustness: The method remains effective even when using imperfect process reward models (PRMs) that provide weak or saturated signals, suggesting that budget-aware structural changes can compensate for noisy evaluation signals.

The work concludes that aligning search policies with fixed token constraints is a crucial step for deploying LLMs in resource-constrained, real-world settings.