Uncertainty-Aware Tree Search for Efficient LLM Reasoning

Jianguo Huang, Zhiyi Lyu, Fuxiang Zhang, Yanchen Deng, Deheng Ye, Bo An

Keywords: LLM Efficient Inference, Test-time Scaling

Abstract: Tree search is an important class of methods for improving the multi-step reasoning capability of Large Language Models (LLMs) by explicitly exploring the intermediate steps on a search tree guided by a value function. However, the existing tree-search methods often devote equal computational budgets across different reasoning branches regardless of the associated uncertainty, causing significantly high token consumption. In this paper, we first conduct a pilot study to reveal the ubiquitous semantic redundancy of reasoning trajectories starting from an intermediate reasoning step. Such highly certain reasoning steps will ultimately reduce the diversity of the final answers. Further, we theoretically show that under a probabilistic guarantee, the sampling budget required to maintain fixed generation quality grows proportionally with the step uncertainty. Built on top of this, we propose \textbf{Uncertainty-Aware Allocation} (UAA), a plug-and-play framework that allocates search budget adaptively according to the step-wise uncertainty. In particular, UAA detects for highly certain reasoning steps and incorporates (i) uncertainty-aware pruning (UAP) to keep only a high-quality subset of candidate actions, and (ii) uncertainty-aware budgeting (UAB) to shrink the next-step expansion budget. Extensive empirical evaluations demonstrate that UAA can significantly reduce token consumption and wall-clock time without hurting accuracy when applied to Beam Search and MCTS.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Submission Number: 23846