Uncertainty-aware budgeting speeds up large language model tree search
Uncertainty-Aware Budget Allocation for Efficient LLM Tree Search
Computation and LanguageArtificial Intelligence
Summary
Searching through possible answers step-by-step with large language models usually spends equal effort on every choice, even when some options are very similar or repetitive. The authors found that by measuring how different sibling options are at each step, they can predict when further searching is unnecessary or less helpful. Their method, called Uncertainty-Aware Allocation, adjusts how much effort to spend on each branch in the search tree, saving time and computing power. This approach works without extra training and maintains accuracy while using fewer tokens and running faster.
What this means in practice
- •For software engineers: Speed up the execution of LLM-based reasoning systems by dynamically allocating computational effort based on uncertainty estimates during tree search.
- •For ai service operators: Reduce cloud compute costs and latency for large language model services by implementing uncertainty-aware budget management in tree search algorithms.
Authors
Jianguo Huang, Zhiyi Lyu, Fuxiang Zhang, Yanchen Deng, Deheng Ye, Bo An
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 $\textbf{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.