Effective problem solving with Large Language Models (LLMs) can be enhanced when they are paired with external search algorithms. By viewing the space of diverse ideas and their follow-up possibilities as a tree structure, the search algorithm can navigate such a search space and guide the LLM toward better solutions more efficiently. While the search algorithm enables an effective balance between exploitation and exploration of a tree-structured space, the need for an external component can complicate the overall problem-solving process. We therefore pose the following question: Can LLMs or their underlying Transformer architectures approximate a search algorithm? To answer this question, we first introduce a simplified framework in which tree extensions and feedback signals are externally specified, allowing for controlled evaluation of search capabilities. We call this setting unknown tree search with bandit feedback. Within this setting, we show that Transformers are theoretically expressive enough to implement distinct search strategies and can be trained from scratch to approximate those strategies. Our Transformer models exhibit the possibility of generalizing to unseen conditions such as longer horizons or deeper trees. Furthermore, we demonstrate that continued task-focused training unlocks the complete capabilities of a pretrained LLM, by fine-tuning the LLM on search trajectories.
Figures & tables
Figure 1: Our perspective on effective problem solving as an iterative process of three phases. Given a prompt that describes a simple problem, i.e., a game of 24, an LLM generates several possible steps, selects the next step, and finally evaluates its potential. Repeating this cycle constructs a tree-structured search space, where each branch represents a potential path of problem solving. This example is generated by GPT-5, and paraphrased to clearly illustrate our definition of effective problem solving. Under this perspective, existing problem formulations and our problem formulation (highlighted in gray) are summarized in Figure 1(b) . Specifically, our formulation assumes that next step selection is carried out by an LLM or a Transformer while state expansion and state evaluation are externally given. Figure 1(c) shows the results on multi-reward tree search with binary trees of depth 6 and 8 different goal states; refer to Sections 5 and E.1 for the details of the metric and the experiment, respectively. Existing LLMs are inferior to some of the established algorithms, and Qwen3-8B Thinking is even worse than uniform leaf sampling, which is a naïve strategy.
Figure 2: Two environments investigated in this work, where darker cells represent higher reward values and red cells denote cells that are impassable.
Figure 3: Behavior cloning results on the multi-reward tree search problem, where each binary tree of depth 6 has 8 different goals and a search step budget is 50.
Figure 4: Behavior cloning results on the multi-reward navigation problem, where the size of each problem is 4 × 4, the wall density of each problem is 0.4, and a search step budget is 50.
Figure 5
Figure 7: Generalization analysis over the numbers of test goal states on the binary tree search problem of depth 6 with a step budget of 50. The number of test goal states smaller than 8 is unseen. A gray dashed line indicates the setting used in training.
Figure 8: Generalization analysis over the numbers of test steps on the multi-reward navigation problem of size 4 × 4 with a wall density of 0.4. The number of test steps larger than 50 is unseen in a training phase. Transformer-based models adequately follow their reference algorithms in both in-distribution and out-of-distribution settings. A gray dashed line indicates the setting used in training.
Performance Metric
Qwen3-8B
Qwen3-8B
FT Qwen3-8B
FT Qwen3-8B
+ UCT Policy
+ UCT Policy
Max-Reward Hit Rate, Highest Reward, & Cumulative Reward
0.262 (0.069)
0.286 (0.071)
0.500 (0.078)
0.548 (0.078)
Discounted Cumulative Gain
0.153 (0.046)
0.188 (0.052)
0.262 (0.050)
0.285 (0.047)
Normalized Path Length
0.262 (0.069)
0.286 (0.071)
0.449 (0.075)
0.497 (0.076)
Table 1: Comparison of Qwen3-8B and fine-tuned Qwen3-8B models, with and without UCT policy-guided sampling. FT stands for fine-tuned.
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Construction of a tree-structured space from a particular instance of multi-reward navigation. s and g1 are a start node and a goal node, respectively. In addition, v1 , v2 , v3 , and v4 are visitable nodes.
Figure 10: Visualization of our tokenization scheme. Same colors indicate same state transitions.
Figure 11: Trace example of multi-reward tree search problems.
Figure 12: Trace example of multi-reward navigation problems.
Figure 13: System prompt used to drive the Tree - Search Assistant.
Figure 14: System prompt used to score the path in the Academic Paper Search problem.
Figure 15: System prompt used to select the candidate in the Academic Paper Search problem.
Register
Vα
#
?
%
Sk
value
α
0
0
0
0
#?
0
+ 1
− 1
0
0
isValue
1
0
0
0
0
is#?
0
1
1
0
0
isState
0
0
0
1
0
id i
0
0
0
0
\mathds1[i=k]
Appendix
Table 2: Initial embeddings. Each row is a register (coordinate), and each column gives its initial value for the token types Vα (value), “ # ? % ” separators, and state tokens Sk .
Register
Vα
#
?
>
[bos]
%
Sk
value
α
0
0
0
0
0
0
#?
0
+ 1
− 1
0
0
0
0
isValue
1
0
0
0
0
0
0
is#?
0
1
1
0
0
0
0
is>
0
0
0
1
0
0
0
is?
0
0
1
0
0
0
0
Appendix
Table 3: Initial embeddings for the MCTS-Transformer: static registers encode each token’s fixed features, while dynamic registers are initialized to zero and updated by the network conditioned on the trace.
Start Paper
Target Paper
Start Paper
Target Paper
W2585547904
W3196632732
W2390034408
W2357505547
W2272018300
W574723991
W3196681629
W2780169483
W4252577446
W3213173784
W133467238
W3138433919
W3023894403
W3186777337
W2887992654
W3035676168
W2610360324
W2254883584
W28024976
W4230887124
W375670644
W2645342909
W4403312580
W2360871073
Appendix
Table 4: Start–target paper pairs in the Academic Paper Search problem.
Figure 16: Behavior cloning results on the multi-reward tree search problem, where each binary tree of depth 8 has 8 different goals and a search step budget is 50.
Figure 17: Behavior cloning results on the multi-reward tree search problem, where each quaternary tree of depth 4 has 8 different goals and a search step budget is 50.
Figure 18: Behavior cloning results on the multi-reward tree search problem, where each binary tree of depth 6 has 4 different goals and a search step budget is 50.
Figure 19: Generalization analysis over test wall densities on the multi-reward navigation problem of size 4 × 4 with a step budget of 50. The wall density smaller than 0.4 is unseen in a training phase. A gray dashed line indicates the setting used in training.
Figure 20: Generalization analysis over test tree depths on the binary tree search problem with 8 different goal states and a step budget of 50. The test tree depth larger than 6 is unseen in training. A gray dashed line indicates the setting used in training.
Tree search is a central abstraction behind many language-agent reasoning and decision-making tasks: agents must explore actions, remember failures, and backtrack toward promising alternatives. Yet, we lack a theoretical understanding of how transformer-based policies acquire such search capabilities from the training dynamics of reinforcement learning (RL). We study this question in a stochastic k-ary tree environment, where an agentic transformer observes only its trajectory history through interaction and receives a terminal reward for reaching a hidden leaf goal node. We first construct a two-head transformer that implements randomized depth-first search (DFS): one head tracks previous actions, while the other detects failure outcomes and triggers backtracking. We then analyze the training dynamics of policy gradient under a depth-wise curriculum, showing that this same DFS mechanism emerges in stages from sparse reinforcement feedback without expert demonstrations. The resulting policy exhibits depth generalization: after training only on depth-1 and depth-2 trees, it succeeds on deeper full trees. We further show that, under imbalanced goal distributions, discounting the return leads to a ranked DFS policy that prioritizes higher-probability branches. Overall, our results identify a mechanistic normal form for transformer-based search, in which attention heads specialize and cooperate to extract decision-relevant traces from context and convert them into agentic action selection via RL training.
Tong Yang, Yu Huang, Yingbin Liang +1
Department of Electrical and Computer Engineering, Carnegie Mellon University · Department of Statistics and Data Science, Wharton School, University of Pennsylvania · Department of Electrical and Computer Engineering, The Ohio State University +1
As pretraining scaling laws approach saturation, Test-Time Scaling (TTS) has emerged as an important direction for improving reasoning by allocating inference-time compute to a fixed model prior. Viewed at a high level, TTS reframes inference as search over a space of partial reasoning states. While Chain-of-Thought (CoT) exposes intermediate steps, common instantiations rely on single-trajectory decoding, limiting recovery from early errors and exploration. This survey systematizes recent progress in tree-search-based reasoning, viewing inference as instance-specific optimization rather than decoding. We trace the evolution from uninformed search to Monte Carlo Tree Search (MCTS), highlighting how sampling-based control supports principled exploration-exploitation trade-offs. To unify a fragmented literature, we introduce a Unified Design Space spanning search topology, evaluation signals, and control dynamics, and advocate a standardized compute-reporting abstraction to make compute-accuracy trade-offs explicit and comparable.
Jiaqi Wei, Xiang Zhang, Yuejin Yang +10
1Zhejiang University · 2Shanghai AI Laboratory · 4Fudan University +3
Search-augmented language models can use external evidence to compensate for limitations in parametric knowledge, but search is not uniformly beneficial: models may call search for questions they can already answer, or rely on noisy evidence when correction, clarification, or abstention would be more appropriate. We formulate this as an instance-level search-routing problem: deciding whether search is needed to improve task success relative to a no-search execution. To derive supervision, we compare no-search and forced-search outcomes for the same question and construct an oracle over NO SEARCH, SEARCH, and UNSOLVED based on task-specific success. Using this oracle as both an evaluation criterion and a learning signal, we train search-routing policies with supervised fine-tuning and preference optimization, improving routing macro-F1 on oracle-eligible examples from 0.7082 to 0.8235 for Gemma E2B and from 0.7053 to 0.8365 for Qwen3.5-4B. Further analysis shows that the learned policies reduce model-specific routing failures: Gemma primarily learns no-search restraint, while Qwen further reduces missed search; residual UNSOLVED cases reveal heterogeneous bottlenecks involving model capacity, retrieval budget, evidence use, and policy behavior.