💬 NLP

TSLM: Tree-Structured Language Modeling for Divergent Thinking

This paper introduces Tree-Structured Language Modeling (TSLM), a novel approach that encodes branching search paths into a single generation process using special tokens and training on complete search trees to enable efficient, systematic divergent thinking without the redundant computation of external search methods.

Original authors: Doyoung Kim, Jaehyeok Doo, Minjoon Seo

Published 2026-02-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Doyoung Kim, Jaehyeok Doo, Minjoon Seo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a very tricky maze.

The Old Way (Sequential Thinking)
Most current AI models are like a person walking through that maze with their eyes closed, taking one step at a time. They pick a path, walk forward, and if they hit a wall, they have to go all the way back to the start and try a different route. If they make a mistake early on, they might keep walking down a dead end for a long time before realizing it. They can only think about one path at a time.

The New Way (TSLM: Tree-Structured Language Modeling)
The paper introduces a new method called TSLM. Imagine instead of walking the maze alone, you are a magical explorer who can instantly see all the possible paths branching out from where you are standing.

TSLM teaches the AI to do this by using special "magic markers" (tokens) in its text. Instead of just writing one sentence, the AI writes a whole map at once:

It writes down Path A and puts a green flag ([SEP]) saying, "This path looks good, keep going!"
It writes down Path B and puts a red flag ([FAIL]) saying, "This path leads to a dead end, stop here."
It writes down Path C and puts a gold star ([GOAL]) saying, "We found the treasure!"

How It Works

Training: The AI is trained by showing it these complete maps, including the dead ends and the winning paths. It learns to see the whole forest, not just one tree. It learns that sometimes you have to try a path just to realize it's a dead end, and that's okay.
Thinking: When the AI solves a problem later, it doesn't just guess one answer. It generates a whole "tree" of possibilities in one go. It sees the branches, checks which ones are dead ends, and follows the ones that look promising.

Why This is a Big Deal
The paper compares TSLM to other methods using some fun analogies:

The "Game of 24" (A Math Puzzle): Imagine a game where you have four numbers and must use math to get to 24.
- Old AI: Tries one combination, fails, tries another. It gets stuck easily.
- TSLM: Spreads out all the math combinations at once. It found the solution 100% of the time in the tests, while the old way only got it right 17% of the time.
The "Gridworld" (A Navigation Puzzle): Imagine a robot navigating a grid with walls and pits.
- Old AI: When the grid gets bigger (from 10x10 to 20x20), the old AI gets completely lost and confused. It's like a person who knows a small room but gets lost in a huge warehouse.
- TSLM: Even when the grid got twice as big, the TSLM robot still found its way 91.5% of the time. It learned the strategy of navigating, not just the specific map.
The "Hallucination" Problem: Sometimes, a problem has no solution at all (like trying to make 24 from the numbers 1, 1, 2, and 3).
- Old AI: Because it's desperate to give an answer, it might make up a fake math trick that doesn't work.
- TSLM: Because it sees the whole tree of possibilities, it can look at the map, see that every path leads to a dead end, and confidently say, "This is impossible." It correctly identified impossible problems 97% of the time.

The Main Takeaway
The paper argues that we don't need to make AI "think harder" by running it many times in parallel (which is slow and expensive). Instead, we can teach the AI to think in a tree structure during its training. By learning to see the whole map—including the wrong turns—it becomes a much better, faster, and more reliable problem solver. It's the difference between a person who only knows how to walk in a straight line and a person who can look at a map, see all the roads, and choose the best one.

Technical Summary: TSLM: Tree-Structured Language Modeling for Divergent Thinking

Problem Statement

Current large language models (LLMs) fundamentally operate as sequential generators, modeling the probability of the next token based on a linear history ( $p(y|x) = \prod p(y_t | x, y_{<t})$ ). While recent reasoning models (e.g., o1, DeepSeek-R1) have demonstrated impressive capabilities through extended reasoning traces, they remain constrained by this sequential nature. This architecture forces a "linear commitment" to a single path, making it difficult to systematically explore multiple alternatives, recover from early mistakes, or decouple irrelevant exploration paths.

Existing solutions to this problem fall into two categories, both with significant limitations:

Sequential Modeling: Models generate a single deterministic path, limiting the exploration of alternatives.
External Search (e.g., Tree-of-Thoughts): These methods use external orchestration to sample multiple independent trajectories. However, this approach suffers from fragmented distributions, redundant computation of shared prefixes, and a lack of coordination between branches. Furthermore, external search methods often fail to scale robustly when problem complexity increases (e.g., dropping significantly in performance on larger Gridworld environments).

The core challenge identified is the inability of models to coherently construct diverse exploration branches within a single generation process, leading to either premature commitment or inefficient, fragmented parallel sampling.

Methodology: Tree-Structured Language Modeling (TSLM)

The authors introduce Tree-Structured Language Modeling (TSLM), a framework that enables language models to natively generate and manage tree-structured reasoning through supervised learning.

Core Mechanism

TSLM modifies the standard language modeling objective to generate complete search trees in a single forward pass (or a coherent sequence of passes) rather than independent trajectories.

Token-Based Serialization: The model uses special tokens to encode the tree structure into a linear sequence compatible with standard Transformer architectures:
- [SEP]: Indicates a viable action that can be further expanded.
- [FAIL]: Indicates a non-viable action (dead end).
- [GOAL]: Marks the desired goal state.
- [BOS]/[EOS]: Mark sequence boundaries for child expansions.
Conditional Dependencies: Unlike standard sequential models where every token depends on the entire prefix, TSLM employs selective context decoupling. A node's generation is conditioned only on its ancestral path and previously generated siblings at the same node, not on unrelated subtrees. This allows the model to learn the branching structure systematically without redundant recomputation.
Training Objective: The model is trained on complete search trees that include both successful paths and failed attempts. The loss function is a standard cross-entropy loss applied to the serialized sequence, teaching the model to:
1. Generate multiple actions at each decision point.
2. Assign appropriate viability markers ([SEP], [FAIL], [GOAL]).
3. Structure the exploration systematically.

Data Construction

Structured Tasks: For tasks with predefined search trees (e.g., board games), the model is trained to directly reproduce the tree structure.
Open-Ended Tasks: For tasks without explicit trees (e.g., GSM8K, ProntoQA), the authors employ a bootstrapping method. They use a supervision model (e.g., Tree-of-Thoughts sampling) to generate pseudo search trees, incorporating known gold trajectories as high-priority branches and ordering remaining branches using a reward function.

Inference Procedure

During inference, TSLM generates a coherent tree structure. It then parses structural tokens to identify viable branches ([SEP]) and forks these branches into independent sequences for parallel expansion. This allows for systematic exploration (e.g., via Breadth-First Search) while maintaining the complete search topology internally, avoiding the redundant sampling inherent in external search methods.

Key Contributions

Novel Framework: Introduction of TSLM, a token-based serialization framework enabling sequential LMs to learn tree-structured reasoning via supervised learning on cohesive search traces.
Performance Gains: Demonstration of consistent performance improvements across diverse tasks:
- Game of 24: 100% accuracy (vs. 17% for sequential baselines).
- Gridworld (Extrapolation): 91.5% accuracy on unseen 20×20 grids (vs. 42.7% for Tree-of-Thoughts), demonstrating robust extrapolation beyond training complexity.
- Open-Ended Reasoning: Enhanced performance on ProntoQA and GSM8K.
Emergent Capabilities: Discovery of new capabilities arising from learning complete exploration patterns:
- Unsolvable Problem Detection: TSLM can identify and terminate on unsolvable instances (e.g., 97/100 unsolvable Game of 24 cases) without hallucinating, a capability absent in sequential baselines.
- Robust Extrapolation: The ability to generalize to larger environments and higher branching factors without retraining.
New Scaling Paradigm: A shift from "parallel sampling $k$ independent trajectories" to "exploring $k$ branches within a single coherent tree." This approach offers superior inference efficiency and higher convergence accuracy.

Experimental Results

The authors evaluated TSLM against several baselines: Sequence Cloning (SC), Procedure Cloning (PC), GRPO (Reinforcement Learning), and Tree-of-Thoughts (ToT).

Structured Planning (Game of 24): TSLM achieved 100% pass@1 accuracy, significantly outperforming SC (17%), PC (47%), and GRPO (15%). Even compared to ToT@100 (pass@100 metric), TSLM (pass@1) achieved superior results.
Gridworld Scaling: In in-domain (10×10) tests, TSLM reached 100%. Crucially, in out-of-domain (20×20) tests, TSLM maintained 91.5% accuracy, whereas ToT collapsed to 42.7%. This highlights TSLM's ability to internalize search procedures that generalize better than external search algorithms.
Open-Ended Reasoning: On GSM8K, TSLM (61.6% pass@1) matched or slightly exceeded ToT (62.3% pass@100) and PC (55.9%), while being significantly more computationally efficient.
Efficiency: TSLM demonstrated significantly faster inference times compared to ToT, which suffers from "time explosion" due to redundant sampling and lack of computation sharing. TSLM shares computation for common prefixes within the single generated tree.

Significance and Claims

The paper challenges the prevailing assumption that robust reasoning capabilities in language models require reinforcement learning (RL) or complex inference-time search orchestration. Instead, the authors argue that properly structured supervised learning of tree-formatted reasoning traces provides a more direct and efficient path to enhanced reasoning and planning.

Key implications claimed by the authors include:

Internalization of Search: By training on complete trees, models internalize systematic exploration strategies, leading to better generalization than external search methods.
Efficiency: TSLM offers a new test-time scaling paradigm that explores more branches within a single coherent structure, avoiding the computational redundancy of independent parallel sampling.
Robustness: The approach enables models to handle unsolvable problems and scale to increased complexity without the catastrophic performance drops seen in external search methods.

The authors position TSLM not just as a performance improvement, but as a fundamental shift in how language models can be trained to reason, suggesting that the "bottleneck" is not inference compute but the systematic construction of the search space during training.