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

2026-01-02 · Source: cs.AI updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Robotics & Autonomous Systems · Depth: Expert, extended

Summary

Budget-Guided MCTS (BG-MCTS) is a novel tree-search decoding algorithm designed for Large Language Models (LLMs) operating under fixed per-query token budgets. Unlike traditional budget-agnostic policies that often lead to inefficient search patterns, BG-MCTS dynamically aligns its search strategy with the remaining budget. It initiates with broad exploration and progressively shifts towards refining and completing promising candidates as the budget depletes, effectively reducing late-stage over-branching. Evaluated on mathematical reasoning benchmarks MATH500 and AIME24/25, using open-weight LLMs like Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct, BG-MCTS consistently outperformed budget-agnostic baselines across budgets of 10k, 20k, and 30k tokens, achieving superior accuracy and precision by optimizing budget utilization for higher-quality answers.

Key takeaway

For MLOps Engineers deploying LLMs with fixed per-query token budgets, you should consider implementing budget-aware tree-search decoding like BG-MCTS. This approach optimizes answer quality by dynamically shifting from broad exploration to focused refinement as your budget depletes, preventing wasteful late-stage branching. It ensures more reliable solutions within your cost constraints, offering a predictable accuracy-cost trade-off.

Key insights

Budget-Guided MCTS dynamically adjusts LLM tree-search from broad exploration to focused refinement based on remaining token budget.

Principles

• Align search policy with remaining token budget.
• Suppress late-stage widening from shallow nodes.
• Prioritize depth-first search as budget depletes.

Method

BG-MCTS modifies PUCT selection by annealing exploration and adding a depth-biased completion bonus, alongside a budget-guided widening mechanism, all conditioned on the budget sufficiency ratio.

In practice

• Use BG-MCTS for LLM inference in resource-constrained settings.
• Apply budget-conditioned policies for high-quality data synthesis.

Topics

• Large Language Models
• Tree-Search Decoding
• Monte Carlo Tree Search
• Test-Time Scaling
• Token Budget Optimization
• Mathematical Reasoning Benchmarks

Code references

• huggingface/lighteval
• RyanLiu112/Awesome-Process-Reward-Models

Best for: AI Engineer, NLP Engineer, Research Scientist, AI Scientist, Machine Learning Engineer, MLOps Engineer

Related on AIssential

• The Evolution of LLM Inference: Decoding algorithms — Part 1 — LLM inference speed hinges on reducing sequential decoding steps by predicting and verifying multiple tokens in parallel.
• Quantifying the Accuracy and Cost Impact of Design Decisions in Budget-Constrained Agentic LLM Search — Budgeting search depth and retrieval quality offers more accuracy gains than expanding token limits for agentic RAG.
• 3x Faster Search: Parallel Test-Time Scaling with Instructed-Retriever-1 — Parallel test-time scaling with a specialized model significantly boosts retrieval speed and quality by executing search stages concurrently.
• Reasoning Budgets vs. Structured CoT: Controlling LLM Thinking Tokens — Controlling LLM reasoning traces via budget or grammar can reduce token cost but risks accuracy.
• Can I Have Your Order? Monte-Carlo Tree Search for Slot Filling Ordering in Diffusion Language Models — MCTS-guided slot ordering in diffusion models significantly improves reasoning task performance by balancing sequential and non-sequential generation.
• Using GPT-4o-mini as an Entity Resolution Judge: 95% Precision for $0.04 — LLMs can act as cost-effective precision filters in entity resolution pipelines, especially for complex product data.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.AI updates on arXiv.org.