TSLM: Tree-Structured Language Modeling for Divergent Thinking

TL;DR

This work introduces Tree-Structured Language Modeling (TSLM), a framework that enables language models to generate and manage complete search trees in a single forward pass using token-based serialization. By training on full trees that include both successful and failed explorations, TSLM internalizes systematic exploration and decouples context within branches, achieving superior performance and substantial inference efficiency compared with external search methods. Empirical results show perfect accuracy on the Game of 24 and robust extrapolation in gridworld tasks, while maintaining favorable scaling on test-time budgets and outpacing Tree-of-Thought in both open-ended and structured reasoning settings. The approach challenges the need for reinforcement learning or heavy inference-time search by demonstrating that structured supervised learning on tree traces can yield robust, scalable reasoning capabilities.

Abstract

Language models generate reasoning sequentially, preventing them from decoupling irrelevant exploration paths during search. We introduce Tree-Structured Language Modeling (TSLM), which uses special tokens to encode branching structure, enabling models to generate and selectively expand multiple search paths within a single generation process. By training on complete search trees including both successful and failed attempts, TSLM learns to internalize systematic exploration without redundant recomputation of shared prefixes. TSLM achieves robust performance and superior inference efficiency by avoiding the multiple independent forward passes required by external search methods. These results suggest a new paradigm of inference-time scaling for robust reasoning, demonstrating that supervised learning on complete tree-structured traces provides an efficient alternative for developing systematic exploration capabilities in language models.

Figures (7)

• Figure 1: Sequential vs. Tree-Structured Reasoning.(a) Sequential approaches commit to single paths, limiting exploration of alternatives. (b) Tree-structured approaches systematically explore multiple possibilities, enabling recovery from mistakes and comprehensive solution space coverage. TSLM bridges this gap by teaching language models to generate tree-structured explorations natively.
• Figure 2: TSLM (Tree-Structured Language Model) is a language model designed for hierarchical exploration in sequence generation tasks. Beginning with an initial state (e.g., "A") and progressing toward a goal state (e.g., "I"), TSLM constructs a tree structure where nodes represent states and branches signify possible paths. During training, the model serializes the tree into linear sequences using special tokens to separate branches and mark the start/end of a sequence, allowing it to learn structured expansions effectively. During inference, TSLM generates multiple branching actions to explore diverse sequences (e.g., expanding from "C" to "G" and "H"). These branches are independently expanded in parallel using a stitching process, enabling broad exploration toward the goal while efficiently pruning unwanted paths.
• Figure 3: Comparison of (a) mean accuracy and (b) average inference time when scaling the number of candidates across different methods on GSM8K. Methods: TSLM ($\bullet$), PC ($\triangle$), and SC ($\lozenge$, convergence at ToT@100: - - -). All the cache control during the inference time is done by vLLM kwon2023efficient v.0.6.6.
• Figure 4: Comparison of mean accuracy (%) over number of candidates for different search strategies (BFS: $\bullet$, DFS $\bullet$)
• Figure 5: GSM8K accuracy comparison across different base models and methods.