Tree of Thoughts and the return of search to language models

Chain of thought is a single rollout

Shunyu Yao and colleagues at Princeton and Google DeepMind posted Tree of Thoughts on May 17, and the number that travelled is on Game of 24. Given four numbers, produce an arithmetic expression that equals 24. GPT-4 with chain of thought prompting solves 4 percent of their test problems. With their method and a beam of five it solves 74 percent.

The paper's diagnosis of why chain of thought fails here is the point of the whole design. A chain of thought is one sample from the model, decoded left to right, with no way to notice that the second step made the problem unsolvable and no way to go back. Game of 24 punishes that. Pick the wrong pair of numbers to combine first and no amount of fluent reasoning afterwards rescues it. Sampling more chains and voting, which the paper reports at 9 percent, does not help much either, because the failure is inside each chain rather than across them.

The four design choices

The framework decomposes into four questions, and we find it easier to hold the paper in our head as those questions than as a single algorithm. First, what is a thought. The unit of search has to be small enough that the model can generate diverse candidates and large enough that a candidate can be evaluated. For Game of 24 a thought is one arithmetic step leaving three numbers. For creative writing it is a paragraph plan. For mini crosswords it is a single word placement.

Second, how are thoughts generated. The paper uses two modes. Sample several thoughts independently from the same prompt when the thought space is rich, as with paragraphs. Propose several distinct thoughts in one call when the space is small and you want to avoid duplicates, as with arithmetic steps. Third, how are states evaluated. Either ask the model to value a single partial state, for instance by rating whether the remaining numbers can plausibly reach 24 as sure, likely or impossible, or show it several states and ask it to vote for the most promising. Fourth, which search algorithm. Breadth-first with a beam of b states per level for the shallow tasks, and depth-first with backtracking for crosswords, where a bad early word placement should be abandoned.

On Game of 24 the configuration is three steps deep, propose up to five thoughts per state, value each, and keep the best b. With b equal to one the success rate is 45 percent, already far above chain of thought. With b equal to five it is 74 percent. The paper reports 7.3 percent for direct input-output prompting, so chain of thought is slightly worse than no reasoning at all on this task, which is a detail that deserves more attention than it has received.

The model as a proposal distribution

What we think is new here is the role the language model plays. In chain of thought the model is the whole solver. In Tree of Thoughts the model is two components inside a classical search loop that someone else wrote. It proposes successors, and it estimates values. Expansion, pruning, backtracking and termination are handled by ordinary code. That is the structure of classical planning, with the heuristic and the successor generator replaced by a prompted model.

This matters because it separates two things that chain of thought fused together. The quality of individual steps is one property of the model. The ability to allocate effort across alternatives is a property of the search procedure, and a procedure can be given more compute at test time in a way that a single decode cannot. The 45 to 74 percent move from b equal to one to b equal to five is a test-time compute curve. The paper does not draw it as one. It is the first such curve we have seen for a general-purpose model on a reasoning task.

The obvious criticism is cost, and the authors make it themselves. Each node in the tree costs a generation call and an evaluation call, so a b equal to five search over three levels is dozens of GPT-4 calls per problem. Whether that is acceptable depends entirely on the price of a wrong answer, and for most chat uses it is not. The method is for problems where you would otherwise pay a person to think.

Where it works and where it does not

The other two tasks are less dramatic and more informative about limits. On creative writing, where the goal is a coherent passage that ends with four given sentences, the evaluation is a vote among plans and then among drafts, and the gain over chain of thought is modest and judged partly by GPT-4 and partly by people. On mini crosswords, depth-first search with backtracking reaches 60 percent word-level accuracy and solves 20 percent of full puzzles. That beats the baselines and is still far from solving the task, because the value estimates are noisy, so the search prunes good branches and pursues bad ones.

That is the general weakness. Search is only as good as the heuristic, and the heuristic here is the same model whose individual steps we did not trust. Game of 24 works because plausibility of reaching 24 from three numbers is something GPT-4 judges well. Tasks where the model cannot tell a promising partial state from a hopeless one gain nothing from the tree, and the paper's own crossword results are the demonstration.

What we would build on it

The first thing we want is the missing plot. Run Game of 24 at b from one to ten and at depth-first versus breadth-first, and report accuracy against total tokens spent. That turns the paper's two data points into a curve that can be compared with the cost of simply sampling more chains and voting. The second is to train the evaluator. The paper uses the same frozen model for proposing and valuing, and a small model fine-tuned on whether partial states led to solved problems would be cheaper and probably better calibrated.

The longer bet is that the loop moves inside the model. If a model can be trained to propose, evaluate and backtrack in its own output, the external search code becomes unnecessary and the test-time compute curve becomes a property of the model itself. Nobody has shown that yet. This paper shows what the target looks like.