Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Routing, pruning and distillation for reasoning · 0 citations (OpenAlex)

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 Budget-Guided MCTS (BG-MCTS), which conditions MCTS selection and widening on the remaining token budget so search shifts from broad exploration to refinement and answer completion, outperforming budget-agnostic tree-search baselines on MATH500, AIME24/25, and UGPhysics.

Problem

Tree-search decoding for LLM test-time scaling under a fixed per-query output-token budget; existing tree-search policies are budget-agnostic, treating the budget only as a stopping condition, leading to late over-branching or premature termination.

Approach

Condition the PUCT selection and widening decisions on the budget sufficiency ratio rho = 1 - C_used/B: anneal the exploration bonus by rho, add a depth-biased completion term scaled by (1-rho), and add a virtual generative child whose score (mean plus rho-scaled variance of child values) decides when to widen, yielding a budget-aligned wide-to-deep schedule.

Key results

• On MATH500 Lv.5 with Qwen2.5-7B-Instruct, BG-MCTS averages 0.691 vs 0.645 for MCTS and 0.666 for Repeated Sampling across 10K/20K/30K budgets.
• On AIME24/25 with Llama-3.1-8B-Instruct at 30K budget, BG-MCTS reaches 0.106 vs 0.089 for MCTS; averaged 0.087 vs 0.067.
• On UGPhysics with Qwen3-32B, BG-MCTS averages 0.357 vs 0.326 for MCTS.

Measured effects reported in the paper

Surprising findings

• AB-MCTS-M underperforms simple Repeated Sampling in the fixed-budget setting, which the authors attribute to its design for richer feedback and larger compute.
• BG-MCTS produces fewer answered trees at budget exhaustion yet more correct answered nodes, indicating a trade-off between answer coverage and candidate quality.

Limitations

• Budget modeling counts only output tokens; input tokens and reward-model compute are excluded.
• GenPRM-7B scores are highly saturated and near-binary, limiting fine-grained node ranking.
• Physics evaluation is limited to 120 sampled questions and uses a rollout-based evaluator rather than a PRM.

Experimental setup

Datasets: MATH500 AIME24/25 UGPhysics

Benchmarks: MATH500 Level 5 AIME24/25 UGPhysics

Metrics: accuracy answered-tree rate answered-node precision average max depth and width

Baselines: Repeated Sampling Sequential Refinement MCTS AB-MCTS-M LiteSearch-Incre. LiteSearch-Batch Beam Search

Builds on: Monte Carlo Tree Search PUCT AB-MCTS LiteSearch dynamic widening

Closest papers by meaning

• Every Rollout Counts: Optimal Resource Allocation for Efficient Test-Time Scaling (NeurIPS 2025 · similarity 0.90)
• Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Search (NeurIPS 2025 · similarity 0.90)
• Provable Test-Time Scaling for Beam Search in LLM Reasoning (NeurIPS 2026 · similarity 0.90)
• Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for LLM Problem-Solving (ICLR 2025 · similarity 0.90)
• Plan and Budget: Effective and Efficient Test-Time Scaling on Reasoning Large Language Models (ICLR 2026 · similarity 0.89)
• AlphaZero-Like Tree-Search can Guide Large Language Model Decoding and Training (ICML 2024 · similarity 0.89)