Uncertainty-Aware Budget Allocation for Efficient LLM Tree Search

arXiv:2610.22858 · cs.CL, cs.AI · Submitted 2026-10-16 · Read on arXiv

cs.CL, cs.AI

Submitted: 2026-10-16

Updated: 2026-10-16

License: http://creativecommons.org/licenses/by/4.0/

The gist: Tree search enhances multi-step reasoning in large language models (LLMs), yet prevailing methods expand every step with a fixed budget, irrespective of the marginal utility of expansion.

Abstract

Tree search enhances multi-step reasoning in large language models (LLMs), yet prevailing methods expand every step with a fixed budget, irrespective of the marginal utility of expansion. We first show, across existing search algorithms, that continuations sampled from a shared intermediate state are largely duplicates and collapse to a few distinct answers. Critically, we find that the semantic dispersion of sibling actions predicts the redundancy of their downstream continuations, permitting a subtree's budget to be fixed before expansion. An information-theoretic analysis motivates this: treating the per-step budget as the Monte Carlo sample size for estimating a step's residual answer entropy, we derive a sufficient budget that shrinks with sibling dispersion. We accordingly propose Uncertainty-Aware Allocation (UAA), a training-free framework that decouples value from uncertainty: values determine which candidates are promising, while uncertainty, estimated from sibling dispersion at no additional model cost, determines how many merit expansion, via uncertainty-aware pruning and uncertainty-aware budgeting. Across mathematical reasoning, logical reasoning, and planning benchmarks, UAA reduces token usage by up to 90% and latency by up to 9 times at comparable accuracy, generalizing across four tree search algorithms and both open- and closed-source LLMs.

Related papers