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

TL;DR

This paper introduces BG-MCTS, a tree-search decoding algorithm for large language models that dynamically adapts its search strategy based on remaining token budgets, improving performance across various settings.

Contribution

The paper presents BG-MCTS, a novel budget-guided tree-search method that aligns search policies with token budgets, addressing limitations of existing budget-agnostic approaches.

Findings

BG-MCTS outperforms baseline methods across different token budgets.

It reduces late-stage over-branching and premature termination.

Demonstrates improved accuracy on MATH500 and AIME datasets.

Abstract

Tree-search decoding is an effective form of test-time scaling for large language models (LLMs), but real-world deployment imposes a fixed per-query token budget that varies across settings. Existing tree-search policies are largely budget-agnostic, treating the budget as a termination condition, which can lead to 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 budget depletes while reducing late-stage branching from shallow nodes. BG-MCTS consistently outperforms budget-agnostic tree-search baselines across different budgets on MATH500 and AIME24/25 with open-weight LLMs.