Uncertainty-Guided Optimization on Large Language Model Search Trees

2024-07-04

arXiv.org

AI Summary

Conventional tree-search decoding strategies for large language models (LLMs)—such as greedy and beam search—are myopic, neglecting global path structure and probabilistic priors. Method: This paper proposes the first non-myopic search framework that formulates LLM decoding as a Bayesian optimization problem over a tree structure. It introduces transition probability priors and posterior path belief estimation, and designs an uncertainty-aware, sampling-based acquisition function to enable efficient exploration without costly simulations (e.g., MCTS). Results: Experiments across multiple LLMs demonstrate that the method achieves higher sequence likelihood with significantly fewer node expansions, substantially improving search efficiency over mainstream baselines. The core contribution is the first systematic integration of Bayesian optimization into LLM tree search—uniquely balancing probabilistic modeling, uncertainty-guided exploration, and computational efficiency.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Large Multimodal Models (LMMs)
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Tree search algorithms such as greedy and beam search are the standard when it comes to finding sequences of maximum likelihood in the decoding processes of large language models (LLMs). However, they are myopic since they do not take the complete root-to-leaf path into account. Moreover, they are agnostic to prior knowledge available about the process: For example, it does not consider that the objective being maximized is a probability and thereby has specific properties like being bound in the unit interval. Taking a probabilistic approach, we define prior beliefs over LLMs' transition probabilities and obtain posterior beliefs over the most promising paths in each iteration. These beliefs are useful for defining a sample-based, non-myopic acquisition function that allows for a more data-efficient exploration scheme than standard search algorithms on LLMs. Crucially, unlike expensive simulation-based non-myopic methods like the Monte Carlo tree search, our method only requires samples from the beliefs. Our formulation thus views LLM decoding as Bayesian optimization on trees. We discuss how to select the prior and the acquisition function, and demonstrate in experiments with various LLMs that our method achieves higher efficiency than recent baselines: Our method achieves the same or a higher likelihood while expanding fewer nodes.

Problem

• Addresses combinatorial explosion in tree search
• Optimizes sparse reward log-likelihood functions
• Reduces expensive evaluations using uncertainty guidance

Innovation

• Uncertainty-guided probabilistic heuristic for likelihood tree search
• Backtracking and exploration-exploitation trade-off without roll-outs
• Reduces costly evaluations in sparse reward likelihood settings