Learning to Execute Tree Search over Continuous Thoughts

Tree-structured reasoning lets a language model explore several intermediate thoughts, but the branching, scoring and backtracking are usually controlled by an algorithm outside the model. We ask whether a Transformer can execute that search internally, as one autoregressive chain of thought over the continuous thoughts it generates and feeds back into its own context, and whether standard training delivers this ability. We study search on depth-$H$ trees with the search history represented as a sequence of vector-valued tokens in the Transformer's context. The Transformer predicts the next token to select, backtrack or terminate, while a fixed front end supplies candidates with noisy rewards. For the search itself, we construct a binary instance in which single-path success has an upper bound that decreases polynomially with depth, while depth-first search succeeds with probability $1-H^{-3}$, so backtracking is essential to the search we execute. On the representation side, we construct a Transformer with one attention layer that executes depth-first search exactly by predicting successive tokens, up to the first reward gap below a threshold, for every fixed branching factor. On the learning side, we give two answers to whether a Transformer trained by empirical risk minimization (ERM) executes it. First, because the correct next action is determined by the prefix, on-policy imitation learning with ERM learns a single executor without quantizing thoughts, with failure $\tilde{O}(D/N)$ for $D$ proportional to the trainable parameter count and $N$ independent instances per round. However, this approach requires repeated collection and labeling of the learner's own rollouts. Second, under teacher forcing with a quantized front end, the guarantee depends on what is written back. With raw feedback, our guarantee can require exponentially small training errors as the horizon grows. Our counterexample class contains near-optimal models that leave the teacher's trace. Writing the quantized state back restores a constant tolerance, under which ERM over the executor's sparse template reaches the goal within $5H$ selections with probability $1-H^{-3}-\tilde{O}(H/N)$. Together, our results show that a Transformer can execute search over its own continuous thoughts and clarify how guarantees of learned execution differ between on-policy imitation learning and teacher forcing, and depend on how predictions are fed back into context.