Tree Search for Language-Model Reasoning

Why It Matters for Modern Models

Search can solve problems that a single left-to-right pass misses: on the Game of 24, GPT-4 with chain-of-thought prompting solved 4% of problems and with Tree of Thoughts 74% (Yao et al., 2023).
LATS reached 92.7% pass@1 on the HumanEval programming benchmark with GPT-4 by combining tree search with the model's own value estimates and reflections (Zhou et al., 2024).
Search spends extra computation at inference time, and it is only as good as the scores that guide it. MCTS itself comes from game playing and planning, where its behavior is well studied (Browne et al., 2012).

Questions to Think About

When the same language model both proposes steps and scores them, which mistakes can search correct, and which will it repeat?
For a fixed compute budget, when is it better to search a tree than to sample many independent solutions and take a majority vote?

Interactive Demo

The interactive demo for this concept is in the newer notebook.

Open the notebook demo

Core Math

If you prefer intuition first, try the demo in the newer notebook, then come back for the math.

Key Equation

a^* = \arg\max_a \left( Q(s,a) + c\sqrt{\frac{\ln N(s)}{N(s,a)}} \right)

Instead of writing one chain of reasoning from left to right, tree search treats partial solutions as nodes of a tree. A node $s$ is the problem plus the steps written so far, an action $a$ is a next step proposed by the language model, and a value estimate says how promising a node is.

Tree of Thoughts has the model propose several candidate next steps and rate the resulting states, and explores the tree breadth-first or depth-first, backtracking from dead ends (Yao et al., 2023).

Monte Carlo tree search (MCTS) repeats four phases: select a path down the tree, expand a leaf with new candidate steps, evaluate the new node (by a rollout, a learned value or a verifier), and pass that value back up to every node on the path. The UCT rule selects the action with the best balance between estimated value and uncertainty (Kocsis & Szepesvári, 2006):

a^* = \arg\max_a \left( Q(s, a) + c\sqrt{\frac{\ln N(s)}{N(s, a)}} \right)

where $Q(s, a)$ is the average value found after taking $a$ in $s$ , $N$ counts visits, and $c$ sets how much to explore. LATS applies MCTS to language agents, using the model to propose actions, score states and reflect on failed attempts (Zhou et al., 2024).