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

AI summary

Overview

Research area: Test-time scaling and tree-search decoding for large language models, specifically budget-constrained inference.

Technical level: Intermediate. Readers should be comfortable with the basics of Monte Carlo Tree Search (MCTS), PUCT selection, and LLM inference, but the core idea is intuitive: adapt the search policy to how much token budget remains.

Scope: The paper introduces Budget-Guided MCTS (BG-MCTS), a tree-search decoding algorithm that conditions node selection and tree widening on the remaining token budget, and evaluates it against budget-agnostic baselines on math and physics reasoning benchmarks.

What This Paper Is About

Deploying LLMs in practice usually comes with a fixed ceiling on how many tokens a single query may generate, and that ceiling varies by product and setting. Existing tree-search decoding methods mostly ignore this constraint, using fixed hyperparameters and treating the budget only as a stopping condition. This causes two failures: branching too aggressively late in the search and running out of tokens before refinement, or stopping early and leaving budget unused. The goal is to design a search policy that actively adapts to the remaining budget so the same token allowance yields more reliable final answers.

Key Contributions

1. Budget sufficiency ratio as a conditioning signal. The paper defines ρ = 1 − C_used/B, where C_used is the number of output tokens generated so far and B is the total budget. This single scalar drives every budget-aware decision in the algorithm.

2. Budget-Guided MCTS (BG-MCTS), a budget-conditioned tree-search decoder. It modifies standard PUCT in two ways: the exploration bonus is annealed by multiplying with ρ, and a depth-based completion bias is added to node values, scaled by (1 − ρ), so deeper nodes become more attractive as budget depletes.

3. Budget-guided tree widening via a virtual generative child. Each internal node gets a virtual child representing the action of generating a new branch. Its score combines the mean child value and child-value variance, with the variance term multiplied by ρ, so widening is encouraged early and suppressed late.

4. Empirical validation across models, benchmarks, and budgets. BG-MCTS is evaluated on MATH500, AIME24/25, and UGPhysics with Llama-3.1-8B-Instruct, Qwen2.5-7B-Instruct, and Qwen3-32B at token budgets of 10K, 20K, and 30K, plus an ablation isolating each component.

Main Findings

• Consistent gains over budget-agnostic tree search. BG-MCTS achieves the best average accuracy in every model/benchmark combination tested. On MATH500 (Level 5) with Qwen3-32B, it averages 0.798 versus 0.758 for standard MCTS and 0.765 for repeated sampling. On AIME24/25 with the same model, it averages 0.311 versus 0.278 for MCTS.

• Peak performance near budget exhaustion. Accuracy-vs-tokens curves show BG-MCTS continuing to improve toward the budget limit, while baselines often plateau or degrade, indicating the budget is being used to finish strong rather than wasted on late branching.

• Each component contributes independently. Ablations on MATH500 removing explore annealing (Eq. 3), exploit shaping (Eq. 5), or widen annealing (Eq. 6) generally degrade performance, and no ablated variant consistently matches the full method.

• Fewer but higher-quality answered nodes. On AIME24/25 with Llama-3.1-8B at 30K tokens, BG-MCTS produces 899 answered nodes of which 413.7 are correct, versus MCTS's 916.3 answered nodes with only 143.7 correct. The method trades answer coverage for precision.

• Lower tree-level answer reach, higher final accuracy. BG-MCTS reaches answered nodes later in the search and does so on fewer instances at exhaustion, yet achieves higher final accuracy: a "fewer-instances, better-answers" regime.

• Measurable wide-to-deep schedule. Depth and width analyses show BG-MCTS keeps the tree broad early and progressively narrows while deepening, matching the intended policy shift.

• Cross-domain transfer to physics. On UGPhysics, BG-MCTS reaches 0.246 average with Qwen2.5-7B and 0.357 with Qwen3-32B, outperforming MCTS (0.237 and 0.326) and LiteSearch variants.

• Robustness to a weak reward signal. The GenPRM-7B process reward model produces highly saturated, near-binary scores, yet budget-aware selection still yields gains, suggesting the benefit does not depend on a finely calibrated verifier.

Methodology in Plain English

The researchers start from standard MCTS with PUCT, where a search tree is built over partial LLM generations. Each node is a prefix of reasoning; expanding it generates continuations; a reward model scores each new node; and scores are backpropagated up the tree.

Their change is to make every selection decision depend on how much of the token budget is left. They track a simple ratio: remaining budget divided by total budget. When this ratio is near 1 (early), the algorithm behaves much like ordinary MCTS, exploring broadly and being willing to open new branches. As the ratio shrinks toward 0 (late), two things happen. First, the exploration term in the selection formula is scaled down, so the search leans more on accumulated node values. Second, a depth bonus is added to node values, making deeper nodes more attractive so the search finishes promising trajectories instead of starting new ones. Deeper nodes already containing a final answer get no bonus, since their value is already determined by answer quality.

Widening is handled with a trick: each internal node gets a virtual "generate a new child" option that competes with its real children. Its score is the mean value of existing children plus a variance term, with the variance term scaled by the remaining-budget ratio. Early on, this encourages branching from uncertain or promising nodes; late, it fades out, discouraging last-minute shallow branching.

They run experiments under a strict fixed output-token budget per problem. Search stops when the budget is consumed, and the highest-scoring completed answer in the tree is returned. Baselines include greedy decoding, repeated sampling, sequential refinement, standard MCTS, AB-MCTS-M, and LiteSearch.

For the physics benchmark, since no suitable reward model exists for partial physics solutions, they replace it with a rollout evaluator: five rollouts per node, scored by match rate against the reference answer, with majority voting to pick the final answer.

Why This Matters

The paper reframes the token budget from a termination rule into a control signal for decoding. This is a small conceptual shift with practical consequences for any deployment where inference is metered or capped.

Impact on research: It opens a direction for budget-conditioned policy design in test-time scaling, complementing work on verifiers, process reward models, and search heuristics. It also provides a clean diagnostic framework (depth/width curves, answer-reach rates) for characterizing how search methods spend compute.

Real-world applications:

• Token-metered API services, where per-query spending must be predictable and output tokens have direct cost, making a policy that adapts within a fixed envelope attractive.
• Synthetic data generation for reasoning, where higher-precision completed traces can matter more than coverage, and successful/failed trajectories can feed preference optimization.
• Resource-constrained on-premise deployment, where model scaling is infeasible and test-time compute is the practical lever, as the paper notes for its open-weight model choices.
• Agentic or multi-step pipelines with per-call token ceilings, where predictable cost profiles simplify orchestration and budgeting.

Industry relevance: Any product that serves LLM reasoning under fixed cost or latency constraints—search assistants, coding helpers, tutoring systems—stands to benefit from policies that reliably convert a token allowance into finished, correct answers rather than half-explored alternatives.

Future Directions

• Better-calibrated reward models. The authors note that saturated process reward scores limit node-ranking granularity. Stronger process evaluators could amplify the gains from budget-aware selection.

• Adaptive or learned budget schedules. The current annealing uses fixed functional forms controlled by κ and λ. Learning the exploration-to-completion schedule, possibly per task or per model, is a natural extension.

• Broader task coverage. Evaluation is limited to math and physics reasoning. Whether the wide-to-deep schedule transfers to open-ended generation, code, or multi-turn dialogue is untested.

• Interaction with other budget dimensions. The paper fixes output-token budgets only. Jointly conditioning on latency, memory, or monetary cost, and combining with parallel scaling strategies, remains open.

Target Audience

Researchers and engineers working on LLM inference efficiency, test-time compute scaling, and tree-search decoding. It is most useful to practitioners who need to serve reasoning models under strict token or cost budgets, and to those studying search policies and process reward models. Readers unfamiliar with MCTS or PUCT will need background reading in those areas; the algorithmic core is otherwise accessible.

Authors’ abstract

Tree-search decoding is an effective form of test-time scaling for large language models (LLMs), but real-world deployment often imposes a fixed per-query token budget that varies across settings. Existing tree-search policies are largely budget-agnostic, treating the budget merely as a termination condition, thereby risking late-stage over-branching or premature termination. We propose Budget-Guided MCTS (BG-MCTS), a tree-search decoding algorithm that aligns its search policy with the remaining token budget: it starts with broad exploration, then prioritizes refinement and answer completion as the remaining budget decreases while reducing late-stage branching from shallow nodes. BG-MCTS consistently outperforms budget-agnostic tree-search baselines across inference budgets on mathematical reasoning benchmarks and an additional physics reasoning benchmark with open-weight LLMs.