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

AI summary

Overview

Research area: Machine learning / large language model (LLM) reasoning, specifically the intersection of state-space heuristic search with LLM decoding (the Tree-of-Thoughts framework).

Technical level: Advanced. The paper is readable at a high level, but its central contributions include a new expansion bound and a temperature-sensitivity theorem with full derivations, which require comfort with probability, softmax, and heuristic search notation.

One-sentence scope: The paper adapts Levin Tree Search (LTS) to the Tree-of-Thoughts (ToT) framework so that the language model itself acts as the search policy, and shows under a fixed query budget that this matches or beats guided depth-first search (DFS) and beam search across three domains and four LLMs.

What This Paper Is About

Tree-of-Thoughts and related methods improve LLM problem solving by generating a tree of intermediate "thoughts" and searching it, but every thought generation and every state evaluation costs an expensive LM query. Existing search algorithms used on that tree (guided DFS, beam search, MCTS) largely ignore this cost, and guided DFS and beam search need extra LM calls just to score which state to expand next. This paper asks how to get better answers out of an LM when the number of thought generations is strictly capped, and answers it by using the LM's own next-thought probabilities as a search policy rather than as something to be re-queried.

Key Contributions

1. Adapting Levin Tree Search to the Tree-of-Thoughts framework. The LM p is used directly as the policy π in LTS, with a cost of cost(s) = g(s)/π(s), where g(s) is the depth of state s and π(s) is the product of probabilities along the path from the start state. Because the LM guides the search, no additional LM queries are needed for state evaluation. The adaptation samples b_max thoughts per expansion, removes duplicate thoughts, and excludes state-cuts (meaningless on a tree with no cycles). Propose-prompt-based thought generation is explicitly ruled out for LTS because those thoughts are generated sequentially and their probabilities are conditioned on previously generated thoughts, so they cannot serve as independent heuristic scores.

2. A new theoretical bound for pruned/sampled ToT trees. The original LTS expansion bound assumes a generic search tree and does not directly apply to ToT, whose structure is pruned by sampling. Proposition 1 shows that for a sampled subtree T with terminal states H', LTS bounds expansions by min_{s in H'} g(s)/π(s), and Corollary 1 converts this into a bound on the number of thoughts generated, M(T,H') <= b_max * min_{s in H'} g(s)/π(s). The paper notes that sampling can prune the best states, so Theorem 1 and Theorem 2 from prior work do not hold as originally stated for ToT.

3. A sensitivity analysis of the bound to the sampling temperature. Defining a "sufficiently accurate" LM via a margin condition Δ+ >= 0 on the solution-path tokens, Theorem 3 bounds the derivative of the thought-generation bound with respect to temperature τ as b_max * g(s_g)^2 / π(s_g)^2 * (k Δ+ / τ^2).

4. Empirical comparison under a fixed LM query budget. LTS is compared against guided DFS and beam search across Blocksworld, PrOntoQA, and Array Sorting (Sort), using Llama 3 instruct (1B, 3B, 8B) and Qwen 2 (7B). The paper reports that LTS is consistently comparable to or better than the baselines under tight budgets.

Main Findings

• DFS suffers from an "initial-commitment" problem. In the paper's illustrative binary-tree example, DFS commits to the left subtree first and exhausts it before backtracking, even when the goal lies to the right. In the same example, LTS assigns a cost of 3/(0.51*0.5) ≈ 11.76 to the children of the root's left successor, roughly 6.12 to the root's right successor, and about 8.16 to the goal state, so it switches to the right subtree and never explores the rest of the left subtree. The paper connects this to its experimental observation that DFS may reach a terminal state faster than LTS but often reaches a non-goal terminal state, lowering overall accuracy.

• A bound on expansions and thoughts generated for ToT. Proposition 1 gives N(T,H') <= min_{s in H'} g(s)/π(s) and Corollary 1 gives M(T,H') <= b_max * min_{s in H'} g(s)/π(s), where M counts thoughts generated before reaching any terminal state in the sampled subtree.

• The original LTS guarantees do not transfer directly. Because sampling can prune the best states with respect to the cost function, states and their descendants may never be expanded, so LTS does not expand states in strict best-first order within the ToT framework.

• Temperature affects the query count in a specific direction. The bound's gradient depends inversely on the square of π(s_g) and on τ. Increasing τ makes probabilities more uniform and, when the path to s_g is the most likely path, decreases π(s_g). The main inference drawn is that the upper bound is non-negative, so with a sufficiently accurate LM, increasing τ leads to an increase in the number of LM queries. The number of expansions is described as somewhat robust to temperature change if the product π(s_g) * τ stays relatively stable. The paper notes that in experiments the ToT often had multiple goal nodes, contrary to the theory's single-goal assumption, so the results should be read as guidance for temperature tuning rather than a strict bound.

• Renormalizing sampled thoughts is discouraged. If the children of a state are uniform, renormalization would give each thought a 1/b_max probability, overestimating these relative to states where probability density is concentrated in the top b_max thoughts, and this can mislead the search across states with different parents.

• Budget accounting favors LTS. The budget is measured as thoughts generated, and each thought is counted twice for DFS and beam search because they issue additional LM queries for state evaluation; this assumes thought generation and state evaluation cost roughly the same. For the Sort domain, DFS relied only on LM log probabilities, so each thought counted once there.

• Empirical outcome. Under a fixed LM query budget, LTS consistently achieves comparable or higher accuracy than the baseline search algorithms within the ToT framework across the three domains and four LMs. Numeric accuracy, budget, or timing values are not reported in the provided paper content, which is truncated before the results tables.

Methodology in Plain English

The researchers start from the standard Tree-of-Thoughts setup: a language model is asked to produce several candidate intermediate steps ("thoughts") from a given partial solution, forming a tree, and a search algorithm decides which branch to pursue. Guided DFS picks the best-looking branch by a score and commits to it. Beam search keeps the best three states at each depth. Both of those scores require the LM to be queried again about how promising a state is.

The alternative the authors propose is to stop asking the LM for separate scores and instead let the probabilities it already assigned when generating thoughts do the guiding. Each state is given a cost equal to how deep it is divided by the probability of reaching it along its path. The search always expands whichever frontier state has the lowest cost, which is the Levin Tree Search rule. Practically, the algorithm samples a small number of thoughts (up to a branching factor of 3) per expansion, discards duplicates, and stops as soon as a terminal state is reached rather than generating many terminal states and picking the best among them — a deliberate choice to avoid extra queries in a low-budget setting.

For the experiments, thought generation used a temperature of 0.8 with top-k sampling at k=50 and a maximum branching factor of 3, following prior literature. LTS's separate state-evaluation temperature was 1 unless stated otherwise, and a tiny value of about 10^-14 was added to the probability in the denominator to avoid division by zero. All algorithms stopped at the first terminal state or when the predefined budget was exceeded.

Three domains were used: Blocksworld (arranging blocks with stack, unstack, put, and pickup actions, evaluated at search depths 2, 4, 6, 8, 10, and 12), PrOntoQA (the "true" ontology setting, "random" rule ordering, examples requiring 3 to 5 reasoning hops, the same 500 sampled instances as prior work, capped at search depth 10), and Array Sorting (Sort, arrays of exactly 5 elements sorted by pairwise swaps, elements possibly negative or containing a single decimal point, 100 randomly generated instances, capped at depth 10). Four open-weight LMs were tested: Llama 3 instruct at 1B, 3B, and 8B, and Qwen 2 at 7B.

Several domains from prior ToT work were deliberately excluded: Game of 24 and Crosswords because preliminary experiments with Llama-3.2-1B showed high rates of invalid action generation (over 40% for Game of 24), StrategyQA because prior work noted its bottleneck is external knowledge rather than reasoning, and creative writing because its search trees are shallow. MCTS was excluded as a baseline due to its known sample inefficiency.

Why This Matters

The paper reframes inference-time reasoning as a budgeted search problem and shows that the choice of search algorithm, not just the model or prompt, determines how much reasoning you get per LM call. It also provides theoretical grounding for a widely used empirical technique (reusing model probabilities as heuristics) and does so without any additional training, unlike prior work such as Policy-Guided Tree Search, and without needing handcrafted or retrieved cost functions, unlike A*-based approaches for tree navigation.

Real-world applications include:

• Deployments on edge devices or other resource-constrained hardware where the number of LM calls is the binding constraint.
• Latency-critical assistants that must answer within a fixed time or token budget.
• Planning and scheduling tools where a sequence of actions must be derived from start and goal configurations, as in the Blocksworld setting.
• Logical and multi-hop question answering, as in the PrOntoQA setting, where intermediate reasoning steps must be validated.

Industry relevance: reducing the number of LM queries per solved task directly reduces serving cost and latency. The findings are about small, open-weight models (1B to 8B parameters), which is exactly the regime where inference economics dominate and where many on-device or high-throughput products operate. The code is released at the repository linked in the paper.

Future Directions

• Batch thought generation. Propose-prompt methods that generate several thoughts in one LM call are currently incompatible with LTS because their probabilities are conditioned on one another. Finding a principled way to incorporate them could further cut query counts.
• Multiple goal states. The temperature theorem assumes a unique minimizing goal state, but the authors observed multiple goal nodes in practice. Extending the analysis to the multi-goal case is an open theoretical question.
• Tighter bounds under sampling. Because sampling breaks the best-first ordering that the original LTS theorems rely on, a bound that accounts for the probability of pruning good states would sharpen the guarantees.
• Broader domains and models. Game of 24, Crosswords, StrategyQA, and creative writing were excluded for reasons of invalid action rates, external knowledge dependence, or shallow trees; larger LMs and task-specific evaluators were not tested, leaving the generalizability of the accuracy comparison open.

Target Audience

This paper is most useful to researchers and engineers working on test-time compute scaling, LLM reasoning, and inference efficiency — particularly those who need principled ways to allocate a fixed query budget across a search tree. It is also relevant to the heuristic search and planning community interested in how learned or probabilistic policies interact with classic search guarantees, and to practitioners deploying small open-weight models under tight latency or cost constraints. Readers without a background in search theory can still follow the motivation and experimental framing, but the theoretical sections assume familiarity with policy-guided search and softmax parameterization.

Authors’ 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.