PaperlayerView exact revision on arXiv

Source-linked AI summary

Uncertainty-Aware Budget Allocation for Efficient LLM Tree Search

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

arXiv:2610.22858v1cs.CLcs.AI

TL;DR

Tree search can be computationally costly, and fixed expansion budgets do not account for the redundancy of continuations from intermediate states. UAA uses uncertainty to determine how many candidates to expand and value to select promising candidates; it reduces token use and latency while keeping accuracy comparable across search algorithms.

  • Problem

    Tree search methods often expand a fixed number of children at each step, despite redundant continuations and substantial computational overhead.

    Source: §1 · p. 1

  • Method

    UAA is a training-free framework that separates value-based candidate selection from uncertainty-based decisions about how many candidates to expand.

    Source: §1 · p. 2

  • Results

    Across Beam Search, MCTS, DVTS, and REBASE, UAA reduces token usage and latency while keeping accuracy comparable across models and white- and black-box uncertainty proxies.

    Source: §1 · p. 2

  • Takeaways & Limitations

    UAA reduces token usage and inference time while maintaining accuracy across Beam Search, MCTS, and multiple benchmarks.

    Source: §7 · p. 8

    Formal guarantees for pruning and principled hyperparameter choices are left for future work; score-variance proxies for closed-source models are tested only on GPT-4o.

    Source: Limitations · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Tree search supports more robust multi-step reasoning than a single CoT chain, but uniform expansion can waste compute on redundant continuations. UAA uses sibling uncertainty to allocate expansion while values select promising candidates, reducing costs while keeping accuracy comparable across search algorithms.

  • Tree search maintains multiple candidate trajectories to address CoT’s lack of backtracking, but fixed expansion budgets waste compute on trivial steps and redundant continuations.The pilot study found most continuations from the same intermediate state had pairwise cosine similarity above 0.95.Source: §1 · p. 1
  • UAA is training-free and uses sibling-action dispersion to predict downstream redundancy, setting a subtree’s budget before expansion.Source: §1 · p. 2
  • The paper reports redundant continuations beyond mathematical reasoning, including on StrategyQA and planning, and beyond one algorithm and model, including MCTS and Qwen3-32B.Source: §1 · p. 2
  • UAA separates value from uncertainty: values identify promising candidates, while uncertainty determines how many candidates merit expansion.Its two mechanisms are uncertainty-aware pruning and uncertainty-aware budgeting.Source: §1 · p. 2
  • Across Beam Search, MCTS, DVTS, and REBASE, UAA cuts token usage and latency while keeping accuracy comparable across three models and white- and black-box uncertainty proxies.Source: §1 · p. 2

2 Preliminaries

The preliminaries model LLM reasoning as a sequence of steps and tree search as exploration over those steps. They describe candidate selection through value functions, including model confidence and process reward models.

2.1 Reasoning as a sequential decision process

  • Reasoning is modeled as a sequence of steps, with each state containing the question and partial trace, and each transition appending the next action.The policy assigns a distribution over the next step, and a complete trajectory ends in a final answer.Source: §2.1 · p. 2

2.2 Tree search for LLM reasoning

  • Tree search navigates reasoning trajectories using a value function, with each node representing a state and each edge an action.Beam Search retains high-value candidates and samples continuations from them.Source: §2.2 · pp. 2–3

2.3 Value functions

  • Process-level value functions score intermediate steps; model confidence averages generated-token log-probabilities, while PRMs assign step scores.PRMs offer more targeted supervision than outcome-only feedback but require a separately trained model.Source: §2.3 · p. 3

3 Motivation

Uniform expansion can spend compute on near-duplicate continuations, motivating budgets that reflect a step’s uncertainty. The pilot study links sibling dispersion to answer diversity and accuracy, while the information-theoretic analysis motivates smaller sampling budgets at low-uncertainty steps.

3 Motivation

  • Fixed expansion budgets waste compute when candidates are near-duplicates, so the section studies redundancy and the budget a step needs.Source: §3 · p. 3

3.1 Empirical pilot study

  • Most sampled continuations had similarity above 0.95, with mean similarity approaching 1, in reward-guided beam search on GSM8K.The study sampled continuations from intermediate states and compared their text embeddings.Source: §3.1 · p. 3
  • Similarity and unique-answer count had a predominantly negative Spearman correlation with mean close to −0.5, while greater sibling dispersion corresponded to higher downstream accuracy.Highly redundant steps showed negligible marginal accuracy gains under the same budget.Source: §3.1 · pp. 3–4

3.2 Information-theoretic analysis

  • The analysis treats residual answer-entropy as the model’s remaining uncertainty about its own answer after an action, and Monte Carlo sample size as the per-step expansion budget.It assumes residual entropy is Lipschitz in action embeddings.Source: §3.2 · p. 4
  • Under the Lipschitz assumption, low dispersion implies lower variance in the entropy estimate, so fewer samples suffice for fixed precision and confidence.The authors present this as a guideline, not a formal guarantee for UAA.Source: §3.2 · p. 4

4 Uncertainty-Aware Allocation

UAA allocates search effort according to sibling uncertainty rather than uniformly: value ranks promising actions, while uncertainty controls pruning and expansion budgets. It reduces branching at low-uncertainty steps while preserving full expansion at high-uncertainty steps.

Uncertainty-aware budgeting

  • For retained low-uncertainty actions, UAB reduces next-step expansions from B to bB, while high-uncertainty steps retain the full budget.The reduced budget satisfies 0 < bB < B.Source: Figure 2 · p. 5, §4 · p. 5

Uncertainty-aware pruning

  • At low uncertainty, UAP retains top-bk candidates by the original value function; otherwise, it keeps all candidates to preserve reasoning diversity.Value scores determine which candidates survive; low uncertainty indicates semantic redundancy among siblings.Source: §4 · p. 5

Combined allocation

  • UAP reduces the number of expandable parents and UAB reduces children per retained low-uncertainty parent, multiplying their overall reduction in expansions.Source: §4 · pp. 4–5

Uncertainty approximation

  • Sibling-action semantic dispersion supplies an uncertainty proxy; embeddings estimate pairwise cosine dissimilarity, while confidence or PRM scores substitute in black-box settings.The proxy is estimated from candidate actions and is an unbiased estimate of step uncertainty.Source: §4 · p. 5

5 Experiments

Experiments evaluate UAA with Beam Search and MCTS across multiple datasets and verifier settings. UAA reduces token use and inference time while maintaining comparable accuracy, and ablations examine the threshold and the two allocation components.

5.1 Experiment Setup

  • Experiments use Llama3.1-8B-Instruct and Qwen3-32B on GSM8K, MATH-500, and AIME24; Beam Search, MCTS, and LiteSearch use Conf. or PRM and are measured by accuracy, tokens, and time.Source: §5.1 · p. 6

5.2 Main results

  • On GSM8K Beam Search with Conf., tokens fall from 3.12k to 2.58k and time from 53.32s to 34.23s with comparable accuracy.Source: §5.2 · p. 6
  • Across Beam Search and MCTS, UAA reduces token use and inference time with comparable accuracy; Qwen3-32B MCTS token consumption falls up to 90% on MATH-500 with Conf.Source: §5.2 · pp. 6–7
  • Increasing λ monotonically reduces token use while accuracy remains essentially flat; under Conf., UAB-only uses 2.63k tokens versus 2.69k for UAP-only.Combining UAP and UAB achieves the fewest tokens with comparable accuracy under both verifiers.Source: §5.2 · p. 7, Figure 3 · p. 7, Table 2 · p. 7

6 Discussion

UAA reduces token use and latency while preserving comparable accuracy across reasoning tasks, models, and search algorithms. Results also include improved accuracy on StrategyQA and transfer to GPT-4o.

Generalization beyond mathematical reasoning

  • On StrategyQA, Beam Search accuracy rises from 64.33% to 72.00% on Qwen3-32B with fewer tokens (2.02k→1.64k).The results also report reduced token consumption and improved accuracy on Blocksworld.Source: §6 · p. 7, Table 3 · p. 7

UAA under black-box settings

  • On GPT-4o, UAA reduces token usage by about 40% across GSM8K, MATH-500, and AIME24 with comparable accuracy.PRM-score variance serves as the uncertainty proxy when embeddings are unavailable.Source: §6 · p. 7, Table 4 · p. 8

UAA improves advanced tree search algorithms

  • Across DVTS and REBASE, UAA reduces token usage and latency while keeping accuracy comparable across benchmarks.The reported results support combining UAA with these existing search strategies.Source: §6 · pp. 7–8, Table 6 · p. 8

7 Conclusion

The conclusion presents UAA as a training-free way to improve LLM tree-search efficiency. It attributes the savings to pruning redundant actions and reducing budgets at low-uncertainty steps while maintaining accuracy.

7 Conclusion

  • UAA is a training-free framework for improving the efficiency of LLM tree-search reasoning.The conclusion reports maintaining accuracy while reducing token usage and inference time across Beam Search, MCTS, and multiple benchmarks.Source: §7 · p. 8
  • The conclusion identifies equal sampling budgets for low- and high-uncertainty steps as redundant.UAA prunes redundant actions and shrinks budgets at low-uncertainty steps.Source: §7 · p. 8
  • Extending uncertainty-guided budgeting to LLM-powered agent planning is proposed as future work.The stated aim is to improve planning efficiency.Source: §7 · p. 8

Limitations

The paper’s analysis relies on a Lipschitz assumption and an empirical premise that sibling dispersion predicts downstream redundancy. It leaves formal pruning guarantees and principled hyperparameter selection for future work.

Limitations

  • Formal guarantees for pruning and a principled choice of UAA hyperparameters are left for future work.The analysis is motivated under a Lipschitz assumption and an empirical dispersion-persistence premise.Source: Limitations · p. 8

A Related work

  • UAA separates value-based candidate selection from uncertainty-based expansion budgets and can layer onto Beam Search, MCTS, DVTS, and REBASE.The related-work discussion contrasts this training-free, plug-and-play approach with methods that use additionally trained models.Source: A Related work · p. 12

B Details of pilot study

  • The pilot study measures continuation redundancy by embedding independently sampled continuations and comparing their pairwise cosine similarity.It also counts distinct solutions to quantify final-answer diversity.Source: B Details of pilot study · p. 12

B.1 Computation of text similarity

  • Candidate-path consistency is measured by averaging off-diagonal cosine similarities between their embeddings.A higher value indicates more similar candidates in embedding space.Source: §B.1 · p. 12

B.2 Details of motivation experiment

  • The motivation experiment uses reward-guided beam search on GSM8K with LLaMA3.1-8B-Instruct and samples 20 continuations per active node.It estimates semantic uncertainty as one minus semantic similarity.Source: §B.2 · p. 13

B.3 Empirical analysis between current reasoning steps and continuations

  • Sibling-node similarity is positively correlated with the similarity of their downstream continuations across multiple steps.The paper treats current-step semantic redundancy as predictive of downstream redundancy, motivating smaller budgets at low-uncertainty steps.Source: §B.3 · p. 13

C More empirical analysis

This passage identifies an empirical analysis of the assumption across MCTS, large LLMs, and StrategyQA. Figure 4 plots the Spearman correlation between continuation dispersion and sibling-node dispersion.

C More empirical analysis

  • The analysis examines whether the dispersion relationship generalizes across MCTS, large LLMs, and StrategyQA.Figure 4 plots Spearman correlation between continuation dispersion and sibling-node dispersion.Source: C More empirical analysis · p. 13, Figure 4 · p. 13

Appendix B.2.

Empirical studies find redundant reasoning trajectories across tasks, search algorithms, and model sizes, with similar trajectories often yielding less answer diversity. The proof setup treats residual answer entropy as a bounded, Lipschitz function of action embeddings and assumes independently sampled actions.

Appendix B.2.

  • On StrategyQA, beam trajectories had textual similarities mostly above 0.95, and similarity was negatively correlated with final-answer diversity.The study used LLaMA3.1-8B-Instruct and Qwen3-32B with LLM confidence as the guidance signal.Source: Appendix B.2 · p. 13
  • MCTS on GSM8K produced substantial trajectory redundancy, with semantically similar reasoning steps contributing little to diverse answer coverage.The pilot used LLaMA3.1-8B-Instruct and Qwen3-32B.Source: Appendix B.2 · p. 13
  • On GSM8K with Qwen3-32B, both Beam Search and MCTS showed near-identical reasoning paths that contributed minimally to diversifying final predictions.Similar results were reported for StrategyQA with Qwen3-32B.Source: Appendix B.2 · p. 14

D Proofs of theoretical analysis

  • The proof setup models the final answer as a random variable over a finite set and represents actions with fixed unit-norm embeddings.Candidate actions at a state are assumed to be sampled independently from the policy.Source: D Proofs of theoretical analysis · p. 14

D.1 Proof of Proposition 1

The proof links low action dispersion to a smaller sufficient sampling budget under a Lipschitz residual-entropy assumption, while noting conditions that limit this reduction. UAA is incorporated into four search algorithms by using uncertainty to adjust pruning or expansion budgets while leaving their value- or reward-based ranking mechanisms unchanged.

D.1 Proof of Proposition 1

  • Under the Lipschitz assumption, the sufficient budget can shrink with dispersion, but reductions are covered only when dispersion is below R^2/(4L^2).If the assumption fails or L is large, including for steps differing by a digit or sign, the bound reverts to the worst-case budget.Source: §D.1 · p. 15

E.1.1 Beam search

  • Beam Search applies uncertainty-aware pruning to remove redundant candidates and uncertainty-aware budgeting to allocate expansion across those retained.The integration modifies candidate filtering and expansion steps.Source: §E.1.1 · p. 16

E.1.2 Monte Carlo Tree Search (MCTS)

  • In MCTS, uncertainty changes the branching factor through pruning and budgeting, while value ranking and UCB selection remain unchanged.Low-uncertainty nodes retain the top-bk candidates and use bB children; otherwise the default branching is retained.Source: §E.1.2 · p. 17

E.1.3 Diverse Verifier Tree Search (DVTS)

  • In DVTS, uncertainty-aware pruning and budgeting operate independently within each subtree, while the existing PRM selection is preserved.Low-uncertainty nodes use the reduced budget bB; high-uncertainty nodes keep M expansions.Source: §E.1.3 · p. 17

E.1.4 Reward Balanced Search (REBASE)

  • In REBASE, uncertainty prunes low-ranked siblings or caps the parent budget, while reward ranking and reward-temperature allocation remain unchanged.Low-uncertainty parents use bB in place of Bi; high-uncertainty parents keep Bi.Source: §E.1.4 · p. 17

F Depth analysis of the search trees

UAA preserves effective reasoning depth while reducing token use without prematurely truncating exploration. In the GSM8K Beam Search ablation, moderate pruning and reduced-budget settings offered a favorable balance between efficiency and accuracy.

F Depth analysis of the search trees

G Ablation study on key hyperparameters

Related papers

Loading 2610.22858v1…