Policy-Guided Search on Tree-of-Thoughts for Efficient Problem Solving with Bounded Language Model Queries

TL;DR

The paper tackles efficient problem solving with Tree-of-Thoughts (ToT) under strict language-model (LM) query budgets by introducing a policy-guided Levin Tree Search (LTS) approach. By treating the LM as a policy over next thoughts and sampling to generate a bounded number of thoughts per expansion, it provides a ToT-specific bound on state expansions and analyzes how the final softmax temperature influences search. Theoretical results extend LTS guarantees to the pruned ToT setting, and empirical evaluations across Blocksworld, PrOntoQA, and Array Sorting with multiple LMs show that LTS matches or outperforms guided DFS and beam search under tight budgets, with favorable performance under time constraints. These findings suggest a cost-effective, latency-aware decoding paradigm for complex reasoning tasks using ToT.

Abstract

Recent studies explored integrating state-space search algorithms with Language Models (LM) to perform look-ahead on the token generation process, the ''Tree-of-Thoughts'' (ToT), generated by LMs, thereby improving performance on problem-solving tasks. However, the affiliated search algorithms often overlook the significant computational costs associated with LM inference, particularly in scenarios with constrained computational budgets. Consequently, we address the problem of improving LM performance on problem-solving tasks under limited computational budgets. We demonstrate how the probabilities assigned to thoughts by LMs can serve as a heuristic to guide search within the ToT framework, thereby reducing the number of thought evaluations. Building on this insight, we adapt a heuristic search algorithm, Levin Tree Search (LTS), to the ToT framework, which leverages LMs as policies to guide the tree exploration efficiently. We extend the theoretical results of LTS by showing that, for ToT (a pruned tree), LTS guarantees a bound on the number of states expanded, and consequently, on the number of thoughts generated. Additionally, we analyze the sensitivity of this bound to the temperature values commonly used in the final softmax layer of the LM. Empirical evaluation under a fixed LM query budget demonstrates that LTS consistently achieves comparable or higher accuracy than baseline search algorithms within the ToT framework, across three domains (Blocksworld, PrOntoQA, Array Sorting) and four distinct LMs. These findings highlight the efficacy of LTS on ToT, particularly in enabling cost-effective and time-efficient problem-solving, making it well-suited for latency-critical and resource-constrained applications.

Figures (5)

• Figure 1: Illustration of the Tree-of-Thoughts (ToT) framework applied to a Rubik's Cube domain. Cube states with red or green backgrounds represent terminal states as determined by the LM, while gray indicates intermediate states. The green background denotes a goal state. States marked with a cross represent invalid states (not known during search) resulting from thoughts that are infeasible for the domain.
• Figure 2: Comparison of states expanded by DFS and LTS. State $s$ denotes the starting state and the green state indicates a goal state. The values on edges denote the probability of the next state being selected given the previous state. DFS uses these probabilities as rewards to guide the search. The value inside each state for LTS represents $cost = g(s)/\pi(s)$. The blue and purple states represent the states expanded by DFS and LTS, respectively.
• Figure 3: Comparison of accuracy of LTS, DFS, and beam search with increasing budgets across 3 different settings.
• Figure 4: Comparison of accuracy of LTS, DFS, and beam search with increasing time budgets across 3 different settings.
• Figure 5: LTS performance with varying temperature $\tau$ values.

Theorems & Definitions (8)

• Theorem 1
• Theorem 2
• Proposition 1
• proof
• Corollary 1
• Definition 1
• Theorem 3
• proof