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

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.

Large language models often give better answers when they can spend more computation at test time, for example by generating more reasoning before answering. This approach, known as test-time scaling, is like giving the model more thinking time. In real applications, however, this thinking time costs money and must fit within a fixed budget. One way to use extra computation is tree search, where the model explores multiple possible solution paths before choosing an answer. Existing tree-search methods usually treat the budget only as a stopping rule. As a result, they may keep exploring new paths too late, or stop before using the budget effectively. We propose Budget-Guided MCTS (BG-MCTS), a search method that adapts to the remaining budget. It explores broadly early on, then focuses on refining and completing promising answers as the budget runs low. This helps language models produce more reliable answers under a fixed cost, making test-time reasoning easier to control.