Transformers in the Dark: Navigating Unknown Search Spaces via Bandit Feedback
Authors: Jungtaek Kim
Thomas Zeng
Ziqian Lin
Minjae Lee
Chungpa Lee
Jy-yong Sohn
Hyung Il Koo
Kangwook Lee
Venue: Accepted by TMLR ()

Abstract: 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.

Jungtaek Kim Published in Transactions on Machine Learning Research (03/2026)
Transformers in the Dark:
Navigating Unknown Search Spaces via Bandit Feedback
jungtack. kim Swiis.edu
University of Wisconsin-Madden
Thomas Zeng
University of Wisconsin Madison
Ziqian Lin
University of Wisconsin Cattonein-MMaiioon
Madison
Minjae Lee PuriosaAI
Chungpa Lee Yonsei University
Jy-yong Sohn Yonsei University
Hyung Il Koo PuriosaAT
Kangwook Lee
Sangwookkeee@krattoo.ooom
University of Wisconsin Madison, KRAFTON AI, Ludo Robotics
Reviewed on OpenReview: https: // oppnrreveew, net/ forum? id= Jay
Abstract
Effective problem solving with Large Language Models (LLMs) can be enhanced when they are paired with external search algorithms. By viewing the space diverse ideas and their follow-up possibilities as tree structure, the search algorithm can navigate such search space and guide the LLM toward better solutions more efficiently. While the search algo- rithm enables an effective balance between exploitation and exploration of treo-structured space, the need for external component complicate the overall problem-solving pro- coss. We therefore pose the following question: Can LLMs or their underlying Transformer architectures approzimale a search algorithm? To answer this question, we first introduce simplified framework in which tree extensions and feedback signals are externally specified, allowing for controlled evaluation of search capabilities. We call this setting unknoun 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 Lask-focused training unlocks the complete capabilities of a pretrained LLM, by fine-tuning the LLM on search trajectories
'Our implementation and training configurations publicly available https://github. cow.huu..cco..u///aaa/eea Saa//UN-Madison-LLeee in- the-tarke

---

Published in Transactions on Machine Learning Research (03/2026)
1 Introduction
Effective problem solving often follows an iterative process of (i) generating diverse idea candidates, (ii) se- lecting the most promising one to pursue, and (iii)
evaluating its potential. Large Language Models Abdin et al. 2024), including reasoning-focused OpenAl 2025b), suggest that these systems may implicitly implement such process, as visualized in Figure la Moreover, as discussed in previous litera- ture (Yao et al.y 2023a, Hao et al. 2023 "houu et al., 2024h, this problem-solving process can be enhanced when LLMs are paired with external search algorithms. By treating the space of diverse ideas and their Tollow-up possibilities as a tree-structured space, the search algorithm navigales this search space, allowing it efficiently guide the LLM toward better solutions. While the search algorithm lets the model exploit and explore effectively, relying on an external component can complicate the overall problem-solving process.
(LLMs) (Brown et al. 2020;
Llama Team, AT @ Mcta 2024
variants (DeepSeek-A] 2025: Gemini Team, Google
2025
We therefore raise the following research question:
Can LLMs or their underlying Transformer architectures approzimate a search algorithm?
To address this question, we introduce unknown free search with bandil feedback, simplified framework where search space expansions are provided externally and rewards are given as bandit feedback, as presented in the table shown in Figure [ By removing self-generated structures, this setup isolates the model's ability to balance exploitation and exploration during selection, enabling controlled evaluation of LLM behavior under uncertainty. Thus, this provides a rigorous foundation for assessing whether L.LSs genuinely implement structured search and how their strategies operate in such settings.
Within this setting, show that, theoretically, Transformers range of search strategies, and empirically, they can be trained from scratch to imitate such strategies and to exhibit the possibility of generalizing beyond the training distribution, for example to longer horizons or deeper rees. However, existing LLMs are still limited in our simple setting, as shown in Figure Tcl We compare them to the uniform and greedy leaf sampling strategies and the variant of Monte Carlo Tree Search (MCTS) (Sutton Barto 2018) terms of max-reward hit rate; see Sections and 5 for the details of the search algorit hins and the metric, respectively. The LLMs tested in this setting underperform compared to the established search algorithms. In particular, Qwen3-8B is on par with the uniform leaf sampling, which is naïve search algorithm, while Qwen3-8B Thinking performs even worse.
expressive enough to represent wide
The observed performance gap indicates that current LLMs still have limited search capabilities, making them less effective as problem-solving agents when search the core challenge. To fully realize their potential, targeted training is required. Our experiment with fine-tuned Qwen3-8B (Qwen Team 2025) shows that training specifically designed for search under uncertainty significantly enhances LLM effectiveness compared to L.M--only method.
We summarize our contributions as follows:
We introduce unknown tree search with bandit feedback, a simplified and controlled setting that captures the essence of problem solving by externalizing both expansions and feedback;
We provide a theoretical analysis showing that Transformers are expressive enough to represent distinct strategies and conduct empirical studies demonstrating that Transformers trained with these strategies can perform the approximation of behavior cloning and generalization to unseen conditions;
We demonstrate that additional training explicitly designed for search under uncertainty improves LLM performance, mitigating the gap relative to specialized algorithms.
2
Related Work
In this section, briefly review the background and related work relevant to this paper.
2

---

Published in Transactions on Machine Learning Research (03/2026)
Prompt
You are given the numbers (4, 7, B). Using i. - i, combine all numbers unaclly once
Expansion of Next Steps
2.8,8
throgh your sleps explore
diflavent possibiilies.
several dtterent moves.
(28,
20) oan't reach 24,
1188 1 (4, givea 36 or S, not heloful.
Selection of the Next Step
-8 = 70, then (20, meults in P8
Response
28 8 fractional unpromising.
8
7 4
9-4-4
8:8-1
8-4w (4.7.
Then 7+42 +4
(11 -
[4. 16
(28, By
(4.
-7x4-28 28,
Evaluation Its Potential
(4, 7. B.
-7-1s
could combina futther.
-7
leads not 24.
Most these wther overshoot but = swems the mooe
8+
7 x
8-424
8-8=1
promising explore luther.
(4.7,
(28, 8,
(4, 14.7.8
447.0
(a) Effective problem solving with an LLM where prompt on a game of 24 is given
LLM-Guided
Tree Search
Unknown Tree Search with 0.80
Tree Search with External
Search
Bandit Feedback
Expansion
LLM
LLM External search
Externally given LLM or Transfrrmer
Selection
LLM
Evaluation
LLM
LLM
Externally given
Uniionn And
Goods
Lnd Yush
DCT
OPT-1 GPT
Conini 2.5
Geunt 155 Plob Tanslle
Qunadelf Qundaal Thooktis
(b) Comparison of our problem formulation to existing (c) Results of reference search algorithms (blue text) and problem formulations
existing LLMss (red text.)
Our perspective on effective problem solving as an iterative process of three phases. Given a prompt that describes a simple problem, iee., a game of 24, an LLM generates several possible steps, selects the next step, and finally evaluates its potential. Repeating this cycle constructs tree-structured search space, where each branch represents potential path of problem solving. This example is generated by GPT-5, and paraphrased to clcarly illustrate our definition of effective problem solving. Under this perspective, existing problem formulations and our problem formulation (highlighted in gray) are summarized in Figure 1b 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 [Iesshws the results on multi-reward tree search with binary trees of depth and 8 different goal states; refer Sections 5 and E. for the details of the metric and the experiment, respectively. Existing LLMs inferior to some of the established algorithms, and Qwen3-88 8B Thinking is even worse than uniform leaf sampling, which is naïve strategy.
Figure 1:
Blind and Uninformed Search Blind or uninformed search refers algorithms that explore scarch space without. heuristic guidance prior knowledge of its structure (Rssslll & Norvig 2010). Classical methods such Breadth-First Search (BFS), Depth-First Search (DFS), and uniform-cost search system- atically enumerate nodes, guaranteeing completeness and, in some cases, optimality, but scale poorly in large or sparse environments due to their inability to incorporate feedback during search (Russell & Norvig 2010p. Modern extensions introduce adaptivity via stochastic sampling and online feedback. MCTS (Sutton
& Barto 2018 particularly upper confidence bounds applied to trees (UCT) (Kocsis Szepesvári 2006))
balances exploration and exploitation by applying bandit principles tree expansion and has proven highly effective in diverse domains such as Go (Silver et al. 20166) and Atari games (Mnih al., 20133). Bandit- based tree scarch has been further studied in the context of regret minimization and sample complexity (Bubeck
et al. 2011
Learning-Bate Based Planning denotes training models acquire planning abilities from data, allow- ing them to generate action sequences or policies for solving tasks under uncertainty without hand-crafted search. MuZero learns both dynamics model and search policy directly from interaction, achieving strong
3

---

Published in Transactions on Machine Learning Research (03/2026)
performance without access to game rules (Schittwieser et al. 2020j. Valmeekam et al. (2023a) propose benchmarks for evaluating LLMss on planning and reasoning, and given existing pretrained L.MSs such as GPT-3.5 and GPT-4 (OpenAl 2023), Valmeekam et al. (023b) analyze their planning abilities. Chen et al. (2021) prodict actions given the history of states, actions, and rewards, using the Transformer architecture. In addition, historical episodes are collected from a source RL algorithm and then they are used to autore- gressively train a Transformer model (Laskin et al. 2023), which is called algorithmic distillation. Recent studies (Lehnert et al.] 2024 Nolte et al., [024]] Su et al. 2025) propose intriguing methods to train the Transformer-baseee model on the traces of A search. In particular, the method by Nolte et al. (2024) is capable of controlling Transformers' outputs using either fast or slow mode. This line of research provides the entire environment to the models, unlike our problem formulation. Beyond a framework with reasoning and acting (Yao et. al., 2023b) Zhou el al. (2024) make use of the MCTS algorithm in decision-making with LLMs. Kambhampati et al. 2024) introduce the LLM-Modulo framework, in which LLMs generate ideas and serve as external critics.
Reasoning Abilities of LLMs Popular LLMs (DeepSekk A1
2025 Gemini Team, Google 2025 OpenAl
2025b) enhance its performance using a reasoning mechanism. Lightman et al. (2024) compare outcome- supervised and process-supervised reward models for guiding a pretrained LLM's chain-of-thouglt reasoming, showing that step-level supervision yields more reliable reasoning. Dziri et al. (2023) find that Transformers encounter limitations in compositional multi-step reasoning, leading to performance degradation as task complexity increascs. [Sun et al. (2024) show that the reward models trained on casier problems can generalize to supervise harder ones, ena enabling scalable alignment without direct. human leedback. [Seell et al. (2025) investigate that allocating inference-time compute effectively improve LLM performance, comparing larger models under equal budgets. Han et al. (2025) introduce framework that dynamically considers ofan--of-thougt token budgets relative to the complexity of each problem, reducing the mumber of tokens with minimal loss in accuracy. Khalifa et al. (2025) propose ThinkPRM, a generative chain-of-thought process reward model that, with limited step-level supervision, surpasses both LLM-as-a-judge and discriminative PRMs on several benchmarks.
3 Problem Formulation and Model Interfaces
Each problem instance is defined by finite, rooted search tree T= NS with maximum depth D, where S is finite set of states and N : S + 2 is a successor function mapping each state to its children. We assume
a bounded reward function r : S + [0, 1] and define the set of goal states as Sgoal = {s E S : r(s) > 0} requiring each goal state to be a leaf node in T. Goal states are assumed to be sparse, i.c., |Sgoal| < (S], with most states yielding zero reward, reflecting realistic scenarios where solutions are rare. Importantly, the tree T constitutes the underlying search space, which remains hidden from the scarch agent.
We define the value of a state under uniformly random rollout policy as follows:
V(s) = Ex (r(tett)
Sstart s),
(1)
where Tunif denotes uniform random traversal policy through the tree, and Steaf represents the leaf state reached after traversing from the initial state Sstart = 8. To simulate real-world scenarios where exact feedback is unavailable, we approximate V(s) by a random variable ((s). Each roalization of ((s) is obtained via a Monte Carlo or heuristic estimate of (s). For example, the Monte Carlo estimate is generated by performing k independent rollouts using kunif- On the other hand, in our experiments with pretrained or fine-tuned LLMs on "real-world" search trees, we adopt more practical formulation of ((s), which requires internal estimation. Specifically, we can prompt the LLM to directly estimate state values, e.g., by producing score between 0 and 1. These realizations serve as noisy bandit feedback signals received by the search agent.
Given the above setup, the agent-cnvironmenn interaction occurs over T steps, where T < (S). Starting from the root state sp, at each step t, the environment reveals the children N(ss) and a sampled rollout value Uz ~ (st).. We define the frontier of unvisited child states as follows:
f F = = (((()) \ (50, $1 (10,81..... st....,thh sch,
(2)
1=0
4

---

Published in Transactions on Machine Learning Research (03/2026)
the search agent selects the state for the next step, 8141 F,, according to its policy:
a(- Su, vb. (su),, $1 01, ((s)............................ (,
(3)
After T' steps, we obtain a completed search trajectory
[so, co ((ss))), ($1, u1, Ms1)) (sr, (s,,er, vr, N(sr))..
(4)
We evaluate r by the reward achieved at its best (highest-reward) state, r($best) max r(s;). For later convenience, we collect the individual state rewards along the trajectory into the set:
R = {r(s;) (s;, U, N(s))) ET.
(5)
With this notation, the previous expression can be written simply as r(8best) max R. In Section we provide additional diverse metrics, which are computed as functions of R, to provide a more nuanced understanding of each search strategy.
3.1 Specific Instances with Sree-Structured Search Spaces
8U
92
91
$1
$2
93
ST
Bg
$10
f11
812
$15
811
8
(a)
Multi-reward tree search
(b) Multi-reward navigation
Figure 2: Two environments investigated in this work, where darker cells represent higher reward values and red cells denote cells that are impassable.
Given the general sctup, we detail two concrete problems with treo-structurod scarch spaces; their illustrations are shown in Figure 2 The settings of these synthetic problems can be easily controlled to adjust their difficulty, which enables us to conduct the analysis of our Transformer models.
Tree Search This problem is straightforward implementation of our setup, by directly generating a randomized search tree. It is defined with three parameters: (i) a branching factor at each intermediate state B; (ii) a tree depth D; (iii) the number goal states K. The number of accessible states is (BD+1_. - BB - 1), where a root state 8p excluded since initially given. All goal states are located in the leaf nodes of the tree. In Figure 2a So, $1, .... $14 are all states this binary tree, and sg and $11 are goal states where r(ss) r(81)).
Multi-Reward Navigation We define multi-reward maze traversal task within our problem setup. The objective of this task is to find paths from a designated start vertex Ustart to one of several goal vertices. Specifically, consider an undirected graph g = (V. E) representing a two-dimensional grid of size w x h, whore V wh and the edges E connect neighboring vertices (i.e., up, down, left, right). Some wh vertices, which are called walls, cannot be traversed. Here, we define a wall density Twall/ulh /wh, where Awall is the number of wall vertices. The graph contains a designated start vertex Ustart V and set of goal vertices G (g1..... } C V, each associated with strictly decreasing positive reward:
R(g)) > R(gz) > ... R(ga) > 0,
(6)
and R(v) 0 for all non-goal vertices v & G. The branching factor B is naturally bounded by due to the grid structure.
5

---

Published in Transactions on Machine Learning Research (03/2026)
To formulate this problem as search problem within our framework, we define the search space S the set of all valid paths p (vo,...,, in 9, satisfying the following constraints: (i) the initial vertex is fixed as up = Ustart, (ii) each path has length at most T, and (ii) each path visits at most one goal vertex, which if visited, must be the terminal vertex of the path. The successor function is defined by N(e) op E v E V), where p o o denotes appending vertex v to the path p. Finally, the reward function defined as r(p) = R(ve), assigning the reward of the terminal vertex dy the entire path. This construction naturally yields a representation suitable for our general search framework; see Figures 2b and g for the detailed visualization of this construction.
More specifically, in Figure 2b an agent is started from a start vertex 8 and then it can select one of neighbor vertices (i.e., up, down, left, and right vertices of the vertex of interest) sequentially. Red cells indicate impassible vertices. Similar to the binary tree example, r(gı) > r(92) > r(gs). It is noteworthy that we express states as paths from to particular vertex which is shown in Figure 2b so that it can straightforwardly create trco-structured search space.
3.2 Reference Search Algorithms
Here, we introduce several reference search algorithms to solve the problem discussed above.
Uniform Leaf Sampling This method chooses one of Fi-1 uniformly at random at step 8, ~ (FF-1). where U uniform distribution; refer to Algorithm for details.
Greedy Leaf Sampling This algorithm selects one of Ft-1 whose parent state has the highest reward value: & = arg max,e 0(s), where (s)) the estimated value of the parent state of 8; refer to Algorithm 3 for details.
Uniform Path Sampling This strategy uniformly samples the next node at each depth level starting from a root state Sy until meets an instance in F- see Algorithm for details.
Policy-Guided Path Sampling These methods traverse an underlying tree using one of tree traversal policies: pure exploration, greedy, and UCT policies; see Algorithm 5 for details. It is analogous to the conventional MCTS algorithm (Sutton & Barto, 2018), but differs in that it excludes fully-explored sub- trees; refer to Algorithm I and its description.
These search strategies are employed to analyze the abilities of Transformers navigate structured yet unknown environments. The pseudocode of these algorithms is shown in Section B For the sake of brevity, uniform leaf sampling and greedy leaf sampling are categorized as leaf-based sampling, and uniform path sampling and policy-guidee path sampling are categorized as path-based sampling.
3.3 Model Interfaces to Perform Search
Our problem setup can be viewed as a multi-turn interaction between an agent and an enviromment: At each step, the agent sends the environment a message containing the node s it wants to visit, and the environment replies with the node's value v and neighbors (ss) (or unvisited states F).
This formulation allows Transformers potentially any autoregressive next-token prediction model to serve scarch policies by predicting the agent's messages in the conversation conditioned on the chat history. We analyze this search capability through three complementary lenses:
Theoretically demonstrating the existence of Transformers' parameters that implement search algorithms (sce Section
Training Transformers on algorithmic traces (sce Section 5
Improving the performance of pretrained LLMs through targeted fine-tuning (see Section
6

---

Published in Transactions on Machine Learning Research (03/2026)
Since tokenization schemes essential components of both the thcoretical and empirical analyses, we specify them in the respective sections and Section C which provides detailed implementation information. Moreover, our theoretical analysis examines how Transformers' parameters represent search strategies, while our empirical analysis investigates their optimizability and generalizability. Together, these findings enhance our comprehensive understanding of Transformers' representational and performance capabilities search.
4 Theoretical Analysis of Transformers' Search Capability
In this section, we identify the theoretical search capabilities of LLMs by demonstrating that Transform- ers (Vaswani et al, 2017) can implement the search algorithms introduced in Section by simulating each step of the algorithms. For any given trajectory of past search steps, each search strategy will define a prob- ability distribut ion over the next state to visit. We thus show that when the trajectory provided as input (encoded as sequence of tokens), the Transformer can output a distribution over next states that exactly matches the distribution prescribed by the target search strategy. To do this, provide explicit weight constructions for both leaf-based and path-based sampling methods, using Transformers with constant depth and embedding dimension linear in search budget T and branching factor B. Moreover, in our theoretical analysis, specific Transformer model is assumed: (i) layer normalization is omitted; (ii) the fully connected layer replaced with any arbitrary token-wise function; (iii) the conventional softmax attention mechanism is replaced with hard attention. This specific model differs from the actual implementation of Transformers for empirical analysis. Soe Section D.1 for the details of our thcoretical Transformer model.
Given search step budget T and a maximum branching factor B, the trace of the search trajectory can contain at most TB + 1 distinct states there is also the root node. For simplicity, assume exactly TB +1 states in total. Then, we define two tokenization schemes, Leaf- Bas-Based Tokenization and Tree-Basee Tokenization.
Leaf-Based Tokenization
We define three token types for leaf-based search methods:
State tokens S, for each unique state é;
Continuous value tokens Va, where a E [0, 1] is stored in a dedicated coordinate in the embedding:
Structural markers %, #, and
Let a trajectory be
T= ((s0, vo, N(so)),, (81, U1, N(s)), (ST, vr, N(r))))
(7)
where N(ss)) = {8t,1 ... } is the set s child states. We then construct the tokenized trace as follows:
Sxp % Vi # S S ? Ssy % Vo, # Sa,, Say,
for 0, Here, the following special tokens are defined:
scparates the selected state from its value embodding.
precedes the list of child states.
terminates each search iteration.
Note that % is not required for the leaf-basod sampling thcoretical construction and is added ensure consistency with the tree-based sampling methods below, We will construct Transformer model predict the next state token Ss, immediately after each ?.

---

Published in Transactions on Machine Learning Research (03/2026)
Tokenization For path-based search algorithms, extend the leaf-based trace with the start-of-sequncee marker [BOS] and the path separator >. At each iteration, we encode the full tree-policy path leading to the chosen state. The resulting sequence is as follows:
[BOS] Trree Vo S Sea.2
Trree % Vo, Son, Salp ...
where the following:
Se0 S > S..
(8)
tree
denotes the sequence of states selected by the tree policy, whose terminal state S., chosen for expansion at step t. Our constructed model will generate the complete path (1) at each step, conditioned the previous trace. The Transformer terminates the generation by outputting the % token, which also indicates
that the immediately preceding state token is the state that should be used for expansion in the next step.
More details of these tokenization schemes are described in Section
We now present the following thcorems:
Theorem (Leaf-Based Search Policies) There crist 3-layer Transformers with embedding dimension d 10 TB that exactly implement the uniform and greedy leaf sampling policies when using sequential encoding of the trajectory using Leaf-Based Tokenization.
Theorem (Path-Based Based Search Policies) There erist 12-layer Tronsformers with embedding dimension d = 26 7TB that exactly implement the uniform path sampling policy and greedy, pure exploration and UCT based path sampling policies when using a sequential encoding of the trajectory using Treo-Basod Tokenization.
The full proofs of Theorems and 2 are provided in Appendices D.2 and D.3 respectively. Here, we present a high-level overview of the proof for Theorem T specifically on how to construct a Transformer that implements the greedy leaf policy. The construction proceeds in three layers:
(i) Token marking: The first layer uses the input embeddings as lookup table to identify token types (e.g,, state tokens). It then marks whether each state token was selected by the model (a visited node) or given by the environment, (a child node of a visited node).
(ii) Frontier identification and value association: The second layer uses the markings from the previous layer to isolate the frontier states (the unvisited children of visited nodes). For each frontier state, it identifies and retrieves the rollout value from its parent node.
(iii) Selection: The third layer applies a maximum operation over the inherited rollout values associated with the frontier states to select the next state to visit.
The proof for Theorem 2 is similar in flavor, using the same proof techniques. The additional complexities arise from needing to first simulate the path policy before doing selection.
As a final note on the theory, we observe that the leaf-based tokenization linear in the number of iterations, while the tree-based one is at worst quadratic. Thus both are polynomial in the number of search iterations. We also emphasize that the theoretical significance of these two theorems lies in providing guarantees with the following properties: (i) the Transformer constructions have constant depth; (ii) the guarantees are in the form of exact implementations. These results ofler a stronger guarantee than existing universal approximation theorems (e.g., the work by Yun et al. (2020)), which only ensure approximate implementation over compact domains and require non-constant depth.
Finding 1: Transformers possess a sulficient architectural capacity to implement fundamental search algorithms including leaf- based and path-based search policies.
Along with the above finding, one caveat however merits emphasis: While these existence results establish architectural feasibility, they do not address whether such implementations can be learned through standard
8

---

Published in Transactions on Machine Learning Research (03/2026)
training paradigms. Specifically, the Transformers' parameters proved in Theorems 1 and can be thought of as ideal parameters. The practical ability of Transformers to implement these search algorithms depends on whether they can learn the correct parameterizations solely from traces of the reference algorithms, which we empirically explore in the following section.
5 Empirical Analysis on Transformers' Search Abilities
Following the theoretical analysis presented in Section 4 we empirically test Transformers' behavior cloning abilities. We begin by introducing the additional metrics we use for evaluation.
Experimental Details The architecture of Llama models (Touvron et al., 2023 Llama Team, AT @ Meta 2024) adopted our Transformer models that are trained from scratch. Unless otherwise specified, the Transformer models defined with 8 blocks, 8 heads, and 512 hidden dimensions. The mumber of intermediate dimensions is set as 1,024. Minimum and maximum learning rates are 5 x 10-5 and 5 x 10-4 Tt uses a learning rate scheduling technique with a warm-up. The AdamW optimizer (oshchilov & Hutter 2018) with 31 = 0.9 and 32 = 0.99 is utilized, and the rotary positional embedding (Su et al. 2024) with 10000 is used.
As offline training and validation datasets, we generate 200 different problem instances and 100 traces per instance, for a total of 20,000 traces. These datasets are split into 70% training and 30% validation. On the other hand, we report online test performance for all experiments. These test experiments use 10 different problem instances and 100 traces per instance, for total of 1,000 traces. These results with 1,000 traces are used to calculate the mcan and standard deviation of cach experiment sct. We verify that training and test problem instances are mutually exclusive. Moreover, we plot 1.96 standard errors for all figures.
Evaluation Metrics We use the following evaluation metrics to analyze search algorithms:
Max-reward hit rate: The fraction of runs that ever attain the global-best reward r(Sbea) calculates the probability of achieving r(Sbest) over multiple runs (i.e., the hit rate): are the number of runs that hit Sbest and the total number of runs, respectively;
where Nui and Niotal
Discounted cumulative gain: This metrie, based on discounted cumulative gain (Jrreelin & Kekäläinen 2002), quantifies how quickly search algorithm finds Sbest in each run: (L/Nrotal) Novcaa 1/ /ogg2(11 (ix I), where is the 1-indexed max-reward hit iteration index at run k. If particular run fails, is 00;
Normalized path length: This metric regarding the shortest path to Sbest computes a metric value of (1/Norat) Eker exp(Ee,tru) - bx), where fk,truch the ground-truth shortest. path length to 8hest and Eg is the shortest path length predicted by the Transformer model. If it fails find Sbest 4 oc, which makes exp(Ex,truth 4 zero;
Highest/cummultte reward: The highest reward is calculated by (1/Niotal) ENaatal max(Rx) and mulative reward is calculated by (1/Ntotai) Ek= Compared to the max-reward hit rate, these account for both suboptimal and global-best rewards.
Nosial E,ER, r, where R& is the rewards achieved at run
Normalized jump distance: This is the arithmetic mean of jump distances where jump distance is defined by (jump();-1,8t) + jump(st, se+1)) 1))/2, which has been proposed in the work by Zeng et al. (025). Note that jump(s,, 8y) indicates the shortest distance between s; and sy on a tree-structured space.
These metrics are selected to evaluate search algorithms across criteria such as efficiency, robustness, and solution quality. Note that higher values indicate better performance across all metrics except for normalized jump distance.
Transformers' Behavior Cloning Abilitics ble of effectively performing behavior cloning in the two environments presented in Section 3.1 Full details
We show that Transformers trained from scratch are capa-
of the tokenization method are in Section 2
9

---

Published in Transactions on Machine Learning Research (03/2026)
Uniform Loaf Sampling
Unitonn Path Sampling
Croedy Leaf Sampling
1.0
Reference
Transformer
Reference
Transformer
Reference
Transformer
Puth Sampling w/ Pure Exp.
Poth Sampling w/ Groody
Path Sampling w/ UCT
0.0
Reference
Transformer
Reference
Transformer
Reference
Transformer
Max Reward Hit Rate
Norualised Path Length
Normalzed Cumnlative Reward
Discouated Cunalative Gain
Highest Reward
Normalzzd Jang Datinee
Figure 3: Behavior cloning results on the multi-reward tree search problem, where each binary tree of depth 6 has 8 different goals and search step budget. is 50.
Uniform Leaf Sampling
Uniform Path Sampling
Greedy Leaf Sampling
1.0
|
H
Reference
Transformer
Reference
Transformer
Reference
Transformer
Path Sampling w/ Pure Exp.
Path Sampling w/ Greedy
Path Sampling w/ UCT
1.0
0.0
Reference
Trausfomer
Refereace
Transforner
Referentee
Trausformerr
Max-Reward Hit Rate
Normalized Path Length
Normalized Connulstive Reward
Discomted Cumulative Gain
Highestt Reward
Normalized Jump Distance
Figure 4: Behavior cloning results the multi-reward navigation problem, where the size of each problem is 4 x 4, the wall density of each problem is 0.4, and a search step budget is 50.
As shown in Figures 8 and 4 Transformers successfully mimic the performance of the respective reference algorithms. To quantitatively analyze these results, we calculate l2 distance between two vectors of the statistios of the metric values obtained by a reference algorithm and its corresponding Transformer model. We simply employ their means and standard deviations of all six metric values to measure the distance. Figure 5 displays that the results of Transformers are mostly closest to ones of the associated reference algorithms.
Beyond matching the reference algorithm's performance, test whether the Transformer replicates its deci- sion process by measuring the KL divergence between their next-state selection distributions under the same trajectory. This KL divergence is computed online by following the traces of the reference algorithms and comparing the predicted next-state probabilities generated by the reference strategies and the Transformer modols. Eventually, these probabilities are used for calculating the KL divergence DKL.
DKL =
Σ p(xx) log
p(x)
(9)
((x)
10

---

Published in Transactions on Machine Learning Research (03/2026)
Reterence
Referetie
Uniturn Laod
Ginnaly Land
Palls w/ Grondy
Uniforn
Pach wy Pure Exg.
Paatl wy 11TT
KL Divergonos
Dovel
2
2
Dietance
Diel
101
20
(a) Multi-reward tree search
(b) Multi-reward navigation
Step
Figure 5: Comparison of the metric values obtained by reference algo- Figure 6: KL divergences between rithms and Transformers using the results shown Figures 3 and 4 1 the probabilities of reference algo- distance between two vectors regarding the metrics presented. The rithmss and Transformers. The traces shortest and second shortest distances in cach row marked as 1 and of reference algorithms are followed for both.
2, respectively.
where p(z) and q(x) is the next-token probabilities over token of reference algorithm and Transformer model, respectively. Note that the next-state probabilities of the reference algorithms reflect uniform tie- breaking among equally preferred nodes. In particular, to calculate the KL divergence, we should obtain the next-state probabilities of reference algorithms. Uniform leaf sampling assigns equal probabilities to all possible next states, and uniform path sampling uses depth-level equal probabilities to calculate the probabilities of possible next states. Greedy and policy-guided path sampling follow their respective tree traversal rules, while also considering all tied states when selecting among equally preferred options. As result, empirically estimate the probabilities of the reference algorithms by repeating the sampling process 100 times with the same trace, varying only the random seed.
As shown in Figure G the leaf sampling algorithms achieve low KT, divergence with their reference coun- terparts, suggesting successful behavior alignment. In contrast, for path sampling algorithms, they exhibit significantly higher KL divergence, compared the leaf sampling algorithms. This suggests that they are unlikely fully replicate the reference algorithms"' behavior despite matching the performance on all of our metries; the KL divergence results do not agrec with the resulting performance reported in Figure This discrepancy arises because of the following rationales: (i) the training traces for path-based methods do not explicitly encode the tree-policy path for each step; (ii) the randomness included in the process of particular path-based methods, such as the uniform path sampling, pure exploration policy-guided path sampling, and UCT policy-guided path sampling, makes KL divergence higher. Regarding the first rationale, the Trans- former must. infer the tree-based policy and select the next state a single step, without intermediate tree traversal-related outputs. Due this, it is a naturally challenging problem under our tokenization scheme. In addition, this aligns with our thcoretical analysis in Section 4 where the trace format as de- scribed in Section C.1 explicitly requires the Transformer to output the tree traversal policy before selecting the next state. Regarding the second rationale, the order of some subsequent states may be interchangeable, since their ordering does not significantly affect the eventual outcome. This rationale is supported by two observations: (i) pure exploration policy-guided path sampling shows periodic KL divergence patterns with recurring low-divergence regions; (ii) greedy policy-guided path sampling yields lower KT. divergence than other policy-guided path sampling mothods.
We carry out generalization analysis the numbers of test steps, test wall densities, the numbers of test goal states, and test tree depths in Figures 7 that our behavior-cloned Transformers present evidence of generalization to novel tasks. In particular, our Transformer models generalize well within a local range, whereas their generalization performance degrades over longer ranges. We conjecture that this behavior stems from practical limitations such as short context length and limited model capacity. It noteworthy that we test unseen tasks that are harder than training tasks, similar to the recent work by Lee et al. (2025). For example, later test steps are more difficult than earlier ones, as the search space expands with each step.
19 and 20 respectively. These empirical results show
11

---

Published in Transactions on Machine Learning Research (03/2026)
Uniform Leaf
Reference
Transformee
Uniform Path
Transformer
0.5
Greedy Leaf
Relerence
Transforuce
Path w/ Pure Exp.
Referemce
Trunsforar
1.0
0.5
0.0
8
2
Path w/ Greedy
Reference
Transformee
Path w/ UCT
Reference
Transformer
0.5
0.0
2
Test Gool States
Test Good States
#Tee Goal States
# Teed Goal States
Max-Reward Hit Rate
Normalised Path Lengih:
Normalzzed Cummlative Reward
Discouated Cunalative Gain
Highest Reward
Normalize Jang Datance
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 unseen. A gray dashed line indicates the setting used in training.
Table 1: Comparison of Qwen3-8B and fine-tuned Qwen3-8B models, with and without UCT policy-guided sampling. FT stands for fine-tuned.
Qwen3-8B
FT Qwen3-8B UCT Policy
Performance Metric
Qwen3-8B
+ UCT Policy
FT Qwen3-8B
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 Normalized Path Length
0.153 (0.00 (4)04))
0.188 (0.052)
0.262 (0.050)
0.285 (0.017)
0,262 (0.069)
0.286 (0.071)
0.449 (0.075)
0.497 (0.076)
Finding 2:
The Transformers trained from scratch are capable of approximating diverse search
algorithms via behavior cloning and exhibiting the possibility of generalizing to unseen tasks.
More results and discussion Transformers' behavior eloning abilities be found in Section
6 Targeted Fine- Tuning for Enhancing Search Capabilities of LLMs
There is room for improving the search abilities of existing LLMs by utilizing targeted fine-tuning. To tackle this issue, train pretrained LLM, i.e., Qwen3-8B (Qwen Team 2025), on the search trajectories of prompts and their associated responses for the Academic Paper Search problem, where this problem requires an agent to navigale from a start instance to a larget by interacting with the corresponding environments. We choose this problem because its state expansion and evaluation are more structured than in more generic reasoning tasks such as GSM8K (Cobbe et al.) (021) and HotpotQA (Yang al., 2018). The Academic Paper Search problem provides fixed of state expansions, since each paper has a predetermined set of child papers, and supports semantic-based state evaluation, since the distance between papers can be
12

---

Published in Transactions on Machine Learning Research (03/2026)
Uniform Leaf
Thansformer
Uniform Path
Reference
Transformee
Seere
0.5
20
Greedy Leaf
Referesane
Truasocner
Path w/ Pure Exp.
Reference
Transfornce
1.0
0.5
0.0
20
-10
60
80
1000
120
40
80
100
20
400
600
100
200
00
80
100
Path w/ Greedy
Reference
Transformer
Path w// UCT
Referencee
Transformee
Store
0.5
0.0
60
20
90
100
40
00
100
60
Pest Strpe
TTss Stepps
Test Steps
#TTs Stepe
Max-Rewaed Rate
Normalised Path Length:
Normalzzed Cummuatave Reward
Discouated Cunalative Gain
Highest Reward
Normalized Jang Datance
Figure 8: Generalization analysis over the numbers of test steps on the multi-reward navigation problem of size 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.
defined in a straightforward manner. Consequently, this setting aligns well with our problem formulation of unknown tree search with bandit foodback.
The trajectories generated by Gemini-2.5-Flash and policy-guided path sampling with UCT are used for training the pretrained LLM. To verify the performance of fine-tuned LLM, we compare the fine-tuned LLM two methods: (i) pretrained LLM; (ii) the pretrained LLM with specialized scarch algorithm. Specifically, the pretrained and fine-tuned LLMs perform all three sleps of the elfective problemn-ollving process, i.e., space expansion, next-state selection, and state evaluation, whereas the pretrained LLM com- bined with a search algorithm carries out only the two steps, space expansion and state evaluation, with the external algorithm handling next-state selection. The missing details of this experiment can be found in Section b.
As shown in Table D our fine-tuned Qwen3-8B significantly superior the two baseline methods in this problem. These results imply that the search abilities of LLMs are not fully exploited and the additional training explicitly designed for search can successfully unlock these abilities. Importantly, it aligns with our findings discussed in Sections and 5 which demonstrate that the Transformer models, i.e., the underly- ing architectures of LLMs, are capable of implementing various search algorithms both theoretically and empirically.
Finding 3: Targeted fino-tuning enhances the search capabilities of LLMs compared to both LLM--nll method and an LLM with an external algorithm.
Despite our successful experimental results presented in this section, the LLM used in these scttings may introduce unknown biases into the trained model. This issue can potentially be mitigated by using the search trajectories obtained from multiple LLMss. In addition, fine-tuned Qwen3-8B with the specialized
13

---

Published in Transactions on Machine Learning Research (03/2026)
search algorithm is the best compared to the other methods. It is natural outcome since performance through an LLM with an external search algorithm can be considered as ground-truth performance. By optimizing fine-tuned model with a search algorithm, the model's performance could be further improved; this exploration is left for future work.
7 Discussion and Conclusion
To understand whether LLMs truly perform structured search or not, we have tackled the search problem itself, where space extensions and bandit feedback are externally given. Then, show both thcorctically and empirically, that Transformers these theoretical and empirical results, we empirically demonstrate that the trained Transformers showcase evidence of generalization to unseen problem conditions. Finally, the additional fine-tuning of an existing LLM allows us to boost its search abilitics.
approximalely implement diverse search algorithms. In addition to
Limitations and Future Directions Since multi-turn interactions between an agent and an environment are assumed, the process of determining a final state or an answer tends to be inefficient. Tf this overall search process can be internalized within LLMs (Deng et al. 2024 Hao et al., 2024), it may lead to a more efficient alternative to the current approach. Furthermore, to scale up our experiments, we should investigate the search capabilities of Transformers on significantly longer sequences within larger problem instances, which would require larger model sizes and greater computational resources (Kaplan et al., 2020). The search behavior of Transformers in those settings may differ from the current observations. Although this work empirically demonstrates the optimization and generalization of Transformers' parameters for implementing search algorithms, it does not address them theoretically; future work can extend it to establish their theoretical learnability and generalizability by following existing work (Zhang et al. 2024 While our empirical analysis demonstrates the effectiveness of Transformer models in terms of generalization, their generalization abilities still remain limited. Therefore, as future work, one possible direction is to explore a more extensive analysis of generalization, similar to existing studies (Lee ct al.) 2025 Golowich et al. 2025 Cai et al. 2025). On the other hand, we anticipate that RL can enhance the performance of Transformer- based search algorithms with respect to given evaluation metric following recent literature in RL for LLMs DeepSSck-AT, 2025 Shenfeld et al., 2025). Our framework for enhancing the search capabilities of LLMs has the potential to address real-world problems such as planning- based prompting on trees by leveraging tools introduced in recent literature (Xiong el al., 2025 Lu et al.j 20255] Minegishi et al. 2025); we leave this for future work. These applications will enable us to investigate the search performance of LLMs from practical standpoint. Finally, discussed in Section [ while targeted fine-tuning helps internalize the search capabilities of LLMs, it is still hard to reach the performance through exact search algorithms.
Following previous work on tool use (Yao et al. 2023b (chiik et al.] 2023). it implies that the use of external search algs may be a more practical strategy, even if this approach does not internalize the search capabilities within the LLM.
Broader Impact Statement
This work is not directly related to any unethical or harmful tasks since it focuses on the understanding of Transformer models and potentially the improvement of LLMs. However, similar to many advances in learning algorithms, there is a risk that improved search capabilities could be applied to harmful tasks, such optimizing for undesired behaviors in interactive systems. Therefore, we should mind this risk and ensure alignment with ethical standards.
Acknowledgments
This work was supported by National Science Foundation (NSF) award DMS-2023239, NSF CAREER award CCF-2339978, the Amazon Research Award, the Google Cloud Research Credits Program, and a grant from FuriosaAI. In addition, it used resources through CIS251382 from the Advanced Cyberinfrastructure Co- ordination Ecosystem: Services & Support (ACCESS) program, which is supported by NSF awards OAC- 2138259, OAC-2138286, OAC-2138307, OAC-2137603, and OAC-2138296, and resources from the University
14

---

Published in Transactions on Machine Learning Research (03/2026)
of Pittsburgh Center for Research Computing and Data SEIID:SCR_022333 supported by National In- stitutes of Health (NTH) award S100D028483 and NSF award OAC-2117681. We also acknowledge partial support from the National Research Foundation of Korea (NRF) grant funded by the Ministry of Science and ICT (MSIT), Korea [BS-2024-008453511, RS-2024-0044800033 the ICT Challenge and Advanced Net- work of HRD (ICAN) Support Program [RS-2023-0025993] ASS-005-02288048) supervised by the Institute for Information & Communications Technology Planning Evaluation (TTTP); and the Yonsei University Research Fund (2025-22-002).
References
M. Abdin, Aneja, H. Behl, S. Bubeck, R. Eldan, S. Gunasekar, M. Harrison, R. J. Hewett, M. Javaheripi, P. Kauffmann, J. R. Lee, T. Lee, Y. Li, W. Liu, C. C. T. Mendes, A. Nguyen, E. Price, G. de Rosa, O. Saarikivi, A. Salim, S. Shah, X. Wang, R. Ward, Y. Wu, D. Yu, C. Zhang, and Y. Zhang. Phi-4 technical report. arXiv preprint ar Xrx::2442. 08905, 2024.
T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry,
A. Askell, Agarwal, A. Herbert- Voss, G. Krueger, Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, Sutskever, and D. Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NewrIPS), volume 33, pp. 1877-1901, Virtual, 2020.
S. Bubeck, R. Munos, G. Stoltz, and C. Szepesvári. r-armed bandits. Journal of Machine Learning Research, 22(5):1655-169) 2011.
Z. Cai, N. Lee, A. Schwarzschild, S. Oymak, and D. Papailiopoulos. Extrapolation by association: Length generalization transfer in transformers. arXiv preprint ar arXim:2006.00 09251, 2025.
L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and L. Mordatch. Decision Transformer: Reinforcement lcarning via sequence modcling. In Advances in Newral Information Processing Systems (NeurIPS), volume 34, pp. 15084 15097, Virtual, 2021.
K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and J. Schulman. Training verifiers to solve math word problems. arXie preprint arXiv:2110.144: 14168, 2021.
DeepSeck-AI DeepSeck-Ri. Incentivizing reasoning capability in LLMs via reinforcement learning. arXiv preprint ar Xiv:22500 12948, 2025.
Y. Deng, Y. Choi, and S. Shieber. From explicit CoT to implicit CoT: Learning to internalize CoT step by step. arXio preprint 2iv: 14838, 2024.
N. Dziri, X. Lu, M. Sclar, X. L. Li, L. Jiang, B. Y. Lin, P. West, C. Bhagavatula, R. Le Bras, J. D. Hwang, S. Sanyal, S. Welleck, X. Ren, A. Ettinger, Z. Harchaoui, and Y. Choi. Faith and fate: Limits of Transformers on compositionality. In Advances in Neural Information Processing Systems (NeurIPS), volume 36, pp. 00293-70332, New Orleans, Louisiana, USA, 2023.
Gemini Team, Google. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507 062261, 2025.
N. Golowich, S. Jelassi, D. Brandfonbrener, S. M. Kakade, and E. Malach. The role of sparsity for length generalization LLMs. Tn Proceedings of the International Conference on Machine Learning (TCML), pp. 19809 19840, Vancouver, British Columbia, Canada, 2025.
M. Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Compulational Linguislics, 8:156 171, 2020.
T. Han, Z. Wang, C. Fang, S. Zhao, S. Ma, and Z. Chen. Token- bukte-budget-awarrr LLM reasoning. In Findings of the Association for Computational Lúnguisties: ACT, pp. 24842 24855, Vienna, Austria, 2025.
15

---

Published in Transactions on Machine Learning Research (03/2026)
S. Hao, Y. Gu, H. Ma, J. J. Hong, Z. Wang, D. Z. Wang, and Z. Hu. Reasoning with language model is planning with world model. In Proceedings of the Conference on Empirical Melhods in Natural Language. Processing, pp. 8154 8173, Singapore, Singapore, 2023.
S. Hao, S. Sukhbaatar, D. Su, X. Li, Z. Hu, J. Weston, and Y. Tian. Training large language models to reason in & contimuous latent space. arXie preprinl arXiv:24 16.669,, 2024.
C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. Kerkwijk, M. Brett, A. Hal- dane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weekesser, H. Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy. Nature, 585:357-362, 2020.
E. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. LoRA: Low-rank adaptation of large language models. In Proceedings of the International Conference on Learning Representations (ICLR), Virtual, 2022.
K. Järvelin and J. Kekäläinen. Cumulated gain-based ovaluation of IR techniques. ACM Transactions Information Systems, 20(4):422 446, 2002.
S. Kanbhampati, K. Valmeekam, L. Guan, M. Verma, K. Stechly, S. Bhambri, L. Saldyt, and A. Murthy. Position: LLMs can't plan, but can help planning LLM-Modulo frameworks. In Proceedings of the International Conference on Machine Learning (ICML), pp. 22895 22907, Vienna, Austria, 2024.
J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. arXiv preprint (XX:::000100866 2020.
M. Khalifa, R. Agarwal, L. Logeswaran, J. Kim, H. Peng, M. Lee, H. Lee, and L. Wang. Process reward models that think. arXiv preprint arXie:2504.16828 2025.
L. Kocsis and C. Szepesvári. Bandit based Monte-Carlo planning. Proceedings of the European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML-PKDD), pp. 282-222222222 223, Berlin, Germany, 2006.
M. Laskin, L. Wang, J. Oh, E. Parisotto, S. Spencer, R. Steigerwald, DJ Strouse, S. S. Hansen, A. Filos, F. Brooks, M. Gazeau, H. Sahni, S. Singh, and V. Mnih. In-context reinforcement learning with algorithm distillation. In Proceedings of the International Conference on Learning Representations (ICLR), Kigali, Rwanda, 2023.
N. Lee, K. Sreenivasan, J. D. Lee, K. Lee, and D. Papailiopoulos. Teaching arithmetic to small Transformers. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 25001 -20054,4 Vienna, Austria, 2024.
N. Lee, Z. Cai, A. Schwarzschild, K. Lee, and D. Papailiopoulos. Self-improving Transformers overcome easy- to-hard and length generalization challenges. In Proceedings of the International Conference on Machine Learning (ICML), pp. 32930 32964, Vancouver, British Columbia, Canada, 2025.
L. Lehnert, S. Sukhbaatar, D. Su, Q. Zheng, McVay, M. Rabbat, and Y. Tian. Beyond A Better planning with Transformers via scarch dynamics bootstrapping. In Proceedings of the Conference on Language Modeling (COLM), Philadelphia, Pennsylvania, USA, 2024.
H. Lightman, V. Kosaraju, Y. Burda, H. Edwards, B. Baker, T. Lee, J. Leike, J. Schulman, I. Sutskever, and K. Cobbe. Let's verify step by step. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 39578 39601, Vienna, Austria, 2024.
Llama Team, AI Meta. The Llama herd of models. arXiv preprint arXiv:2407.2178. 2024.
Loshchilov and F. Hutter. Decoupled weight decay regularization. In Proceedings of the International Conference Learning Represendations (ICLR), Vancouver, British Columbia, Canada, 2018
16

---

Published in Transactions on Machine Learning Research (03/2026)
J. Lu, Z. Xu, and M. Kankanhalli. Reasoning LLMs are wandering solution explorers. arXiv preprint ar WXXi:2505.20296 2025.
G. Minegishi, H. Furuta, T. Kojima, Y. Iwasawa, and Y. Matsuo. Topology of reasoning: Understanding large reasoning models through reasoning graph properties. arXiv preprint arXi::2566.077 05744, 2025.
V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller. Playing Atari with deep reinforcement learning. arviv preprint 1312 2013.
N. Nolte, O. Kitouni, A. Williams, M. Rabbat, and M. Ibrahim. Transformers can navigate mazes with multi-step prediction. arXiv preprint arXiv:2412.05117, 2024.
OpenAL. GPT-4 technical report. arXiv preprint or arX1::3033008777 2023.
OpenAL Introducing GPT-4.1 in the API, hatpss//ppena..comnnnd////--oo-p 2025a.
OpenAl. Introducing GPT-5.
2025b.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An imperative style, high-performance deep learning library. In Advances in Newral Information Processing Systems (NeurIPS), volume 32, pp. 8026-8008 8037, Vancouver, British Columbia, Canada, 2019.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapcau, M. Brucher, M. Perrot, and É. Duchesnay. Scikit-learn Machine learning in Python. Journal of Machine Learning Research, 12:2825-2888 2011.
J. Pérez, J. Marinković, and P. Barceló. On the turing completeness of modern neural network architec- tures. In Proceedings of the International Conference on Learning Representations (ICLR), New Orleans, Louisiana, USA, 2019.
J. Priem, H. Piwowar, and R. Orr. OpenAlex: A fully-open index of scholarly works, authors, venues, institutions, and concepts. arXiv preprint arXin:225..00 018333, 2022.
Qwen Team. Qwen3 technical report. arXiv preprint arXiv:2505.090 093888 2025.
S. J. Russell and P. Norvig. Artificial Intelligence: Modern Approach. Pearson Education, edition, 2010.
T. Schick, J. Dwivedi- Yu, R. Dessi, R. Raileanu, M. Lomeli, E. Hambro, L. Zettlemoyer, N. Cancedda, and T. Scialom. Toolformer: Language models can teach themselves to use tools. In Advances in Newral Information Processing Systems (NeurIPS), volume 36, pp. 68539 68551, New Orleans, Louisiana, USA, 2023.
J. Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Has- sabis, T. Graepel, T. Lillicrap, and D. Silver. Mastering Atari, Go, chess and shogi by planning with learned model. Nature, 588(7839) 600-609, 2020.
I. Shenfeld, J. Pari, and P. Agrawal. RL's razor: Why online reinforcement learning forgets less. arXiv preprint arXiv:2509.04. 04259, 2025.
D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. van den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, S. Dieleman, D. Grewe, J. Nham, N. Kalchbrenner, I. Sutskever, T. Lil- licrap, M. Leach, K. Kavukcuoglu, T. Graepel, and D. Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484 489, 2016.
C. Snell, J. Loe, K. Xu, and A. Kumar. Scaling LLM test-time compute optimally can be more effec- tive than scaling parameters for reasoning. In Proceedings of the International Conference on Learning Singapore, Singapore, 2025.
Representations (ICLR), pp.
17

---

Published in Transactions on Machine Learning Research (03/2026)
D. Su, S. Sukhbaatar, M. Rabbat, Y. Tian, and Q. Zheng. Dualformer: Controllable fast and slow thinking by learning with randomized reasoning traces. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 95080 95117, Singapore, Singapore, 2025.
J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu. RoFormer: Enhanced Transformer with rotary position embedding. Neurocomputingg 568:127063, 2024.
Z. Sun, I.. Yu, Y. Shen, W. Liu, Y. Yang, S. Welleck, and C. Gan. Easy-to-hard generalization: Scalable alignment beyond human supervision. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, pp. 51118-51168 Vancouver, British Columbia, Canada, 2024.
R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018
H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302.1999 2023.
K. Valmeekam, M. Marquez, A. Olmo, S. Sreedharan, and S. Kambhampati. PlanBench: An extensible benchmark for evaluating large language models on planning and reasoning about change. Tn Advances in Newral Information Processing Systems (NeurIPS), volume 36, pp. 38975 38987, New Orleans, Louisiana, USA, 2023a. Datasets and Benchmarks Track.
K. Valmeekam, M. Marquez, S. Sreedharan, and S. Kambhampati. On the planning abilities of large lan- guage models: a critical investigation. In Advances in Neural Information Processing Systems (NeurIPS), volume 36, pp. 75993-76005, New Orlcans, Louisiana, USA, 2023b.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, Jones, A. N. Gomez, Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), volume 30, pp.
Long Beach, California, USA, 2017.
P. Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Pe- terson, W. Weckesser, J. Bright, S. J. van der Walt, M. Brett, J. Wilson, K. J. Millman, N. Mayorov, A. R. J. Nelson, E. Jones, R. Kern, E. Larson, C. J. Carey, i. Polat, Y. Feng, E. W. Moore, J. VanderPlas, D. Laxalde, J. Perktold, R. Cimrman, 1. Henriksen, E. A. Quintero, C. R. Harris, A. M. Archibald, A. H. Ribeiro, F. Pedregosa, P. van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: fundamental algorithms for scientific computing Python. Nature Methods, 77:261-272. 2020.
G. Weiss, Y. Goldberg, and E. Yahav. Thinking like Transformers. In Proceedings of the International Conference on Machine Learning (ICML), pp. 110S0 11090, Virtual, 2021.
Z. Xiong, Y. Cai, Z. Li, and Y. Wang. Mapping the minds of LLMs: A graph- based analysis of reasoning LLM. Xiu preprint or 2rii:22005. 112005..18800 13890, 2025.
7. Yang, P. Qi, S. Zhang, Y. Bengio, W. W. Cohen, R. Salakhutdinov, and C. D. Manning. HotpotQA: A dataset for diverse, explainable multi-hop question answering. In Proceediiggs of the Conference Empirical Methods in Natural Language Processing (EMNLP), Brussels, Belgium, 2018.
S. Yao, D. Yu, J. Zhao, 1. Shafran, T. L. Griffiths, Y. Cao, and K. Narasimhan. Tree of Thoughts: Deliberate problem solving with large language models. In Advances in Neural Information Processing Systems (NeurIPS), volume 36, pp. 11809-11822, New Orleans, Louisiana, USA, 2023a.
S. Yao, J. Zhao, D. Yu, N. Du, I. Shafran, K. Narasimhan, and Y. Cao. ReAct: Synergizing reasoning and acting in language models. In Proceedings of the International Conference on Learning Representations (ICLR), Kigali, Rwanda, 2023b.
D. Yarotsky. Error bounds for approximations with deep ReLU networks, Neural Networks, 94:103 114, 2017.
18

---

Published in Transactions on Machine Learning Research (03/2026)
C. Yun, S. Bhojanapalli, A. S. Rawat, S. J. Reddi, and S. Kumar. Are Transformers universal approxi- mators of sequenceeto-ssequenee functions? In Proceedings of the International Conference on Learning Representations (ICLR), Virtual, 2020.
Y. Zeng, S. Zhang, W. Kang, S. Wu, L. Zou, Y. Fan, H. Kim, 7. Lin, J. Kim, H. I. Koo, D. Papailiopoulos, and K. Lec. ReJump: A tree-jump representation for analyzing and improving LLM reasoning. arXiv preprinl ar Xrx::2552 00831, 2025.
Zhang, S. Frei, and L. Bartlett. Trained Transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55 55, 2024.
A. Zhou, K. Yan, M. reasoning, acting, and planning in language models. In Proceedings of the International Conference on Machine Learning (ICML), pp. 02138-62100 Vienna, Austria, 2024.
H. Wang, and Y.-.. Wang. Language agent tree search unifies
19

---

Published in Transactions on Machine Learning Research (03/2026)
A
Construction of Tree-Structurees Spaces for Multi-Reward Navigation
U1
4
91
8
U2
U3
U1 U2
U4
V
91
Figure 9: Construction of tree-structured space from particular instance of multi-reward navigation. and 91 are a start node and a goal node, respectively. In addition, vi, v2, 03 and v4 are visitable nodes.
Figure presents how tree-structured search space is constructed from a particular instance of multi-reward navigation. Deeper nodes are omitted for clarity.
20

---

Published in Transactions on Machine Learning Research (03/2026)
B Details and Pseudocode of Reference Search Algorithms
In this section, we provide the pscudocode of reference search algorithms.
Algorithm Fully-Exploree Sub-Tree Exclusion
Input:
The current state 8, unvisited child states F, a hidden search tree T= (S,
Output:
Obtain all visitable child states of s, denoted as C. if C 0 then
return True
else
5: return False
6: end
Algorithm 1 verifies if all possible child states of state all explored. This algorithm is denoted E(s F, where F is unvisited child states and T is hidden search tree. If E(s F, is True, scarch algorithm should exclude this sub-tree of which the root is from further search. Importantly, with this process, underlying tree structures are not given to search algorithms, and we only provide unvisited states to the algorithms.
Without loss of generality, we can define a modified successor function N which only provides immediate child states that have unvisited states:
(*)s)= (s' (ss) -E(s' F,T)).
(10)
where logical not operation.
B.1 Uniform Leaf Sampling
Algorithm Uniform Leaf Sampling
Input: root state 8p, a value estimation function ((s,, a step budget T, a hidden scarch tree (8, Output: search trajectory [[50, up (so)), (81, N(si)),.... (,,,)),,,,,,,,,, (ST,UT, [(s)))]- Initialize trajectory 7 [[so. vo, (ss)))
Update unvisited child states Fo.
3: for ( 1,2, do
4: Choose the next state 8, from Ft-1 uniformly at random.
5: Evaluate 8, to obtain its value using (s:).
Update with (st, Pt. ((s)).
Update unvisited child states Fr.
end for
Algorithm shows the pseudocode of uniform leaf sampling. This algorithm can be considered the randomized algorithm of DFS.
21

---

Published in Transactions on Machine Learning Research (03/2026)
B.2 Greedy Leaf Sampling
Algorithm 3 Greedy Leaf Sampling
Input: root state 8o, a value estimation function (s), a step budget T, a hidden search tree T = (S,N) Output: search trajectory T [(50, vo, (s))), (81, (.,,)....,, (8,, ((rr))
Initialize trajectory т [(so. vo. (s))).
Update unvisited child states Fo- for t 1.2...., do
Choose state & from the parent states of F,-1 which has the highest reward estimated value. 5: Sample the next state 8, from child states [ss)) uniformly.
Evaluate 8 to obtain its value using (st).
Update with (st, vt: (s:)..
Update unvisited child stales Ft. end for
Algorithm 3 demonstrates the pseudocode of greedy leaf sampling. considers the reward values of the parent states of unvisited states.
B.3 Uniform Path Sampling
Algorithm Uniform Path Sampling
Input: A root state So, a value estimation function (s), a step budget T, a hidden search (8, [[so, to; N(s)), (81, U1, (N)s)), (Sr)),,,.,,,,)... N(r)))
Output:
A scarch trajectory
1: Initialize trajectory r = [(so, vo, [ss)))
2: Update unvisited child states Fo.
3: for t = 1, 1.2....,7 do
4: Assign a state as sy = while of $ F do Choose one of (*()(( uniformly random, ensuring that fully-explored sub-trees are considered. Update of using the sampled state.
end while
Assign the next state as = sy Evaluate 8, to obtain its value vs using (st).
10:
11:
Update 7 with (st, vt. (s:)).
12:
Update unvisited child states Fe.
13:
end for
Al Algorithm presents a search algorithm of uniform path sampling. contrast to uniform leaf sampling, this algorithm can viewed as a randomized variant of BFS. The definition of N shown in (10).
22

---

Published in Transactions on Machine Learning Research (03/2026)
B.4 Path Sampling with Pure Exploration, Greedy, or UCT Policy
Algorithm 5 Path Sampling with a Specific Tree Traversal Policy
Input:
root state So, a value estimation function (s), a tree traversal policy P, a stop budget T, a hidden search tree NS.
Output: search trajectory T [(so. to, (so)), (s1, un, (s:)), (7, Initialize trajectory 7 [(so= to. [ss)))
Update unvisited child states Fo.
for 4= ........77 do
4: Assign state as si = 50-
5: while 5/ $ F do Determine unvisited immediate child states F ( P.-1)-
6:
if |7-1>> then
Choose one F uniformly at random.
else
10: Choose one of (+)) with specific tree traversal policy T; see Algorithm ensuring that
fully-explored sub-trees not considered. end if
11:
12:
Update 5y using the chosen state. end while
13:
14:
Assign the next state as 8, = sf
Evaluate 8, to obtain its value is using (st).
15:
16: Update with (st, Ut. (s.)).
17:
Update unvisited child states Ft.
18:
end for
Algorithm 5 shows search strategy of path sampling with a specific tree traversal policy such as pure exploration, greody, or UCT policy. No is defined in (10).
Algorithm 6 Tree Traversal Policies
Input: A tree traversal policy P, the current state modified successor function N
Output: A selected next state s".
1: Pure Exploration Policy
arg min,/en* count(s'), where count returns the number of visit counts.
Greedy Policy
arg maxxeN=(() value(s'), where value returns the summation of all values of the child nodes of s'.
UCT Policy
6: s* rameter.
arg max (aX)) valoc"") /')occ(n(ss') + (22/oo((2comm())
where is balancing hyperpa-
Algorithm 6 presents the details of tree traversal policics. The pure exploration policy only considers the number of visit. counts and chooses the least visited state the next stale. The greedy policy takes into account the values of child states in order to choose the next state. The UCT policy (Kocsis & Szepesvári, 2006) balances exploration and exploitation considering both the number of visits and the estimated values. Across all experiments, we use 0.1. See this reference (Sutton Barto 2018 for more details.
23

---

Published in Transactions on Machine Learning Research (03/2026)
Trace Formats C
This section presents the trace formats used for thcorctical constructions, Transformers trained from scratch, and pretrained existing LLMs.
C.1 Trace Format for Theoretical Constructions
83
8g
$10
$11
$12
$13
% un # (5a) ? % #N( % 44 # (saa) 82 % 02 # N(22)
Figure 10: Visualization of our tokenization scheme. Same colors indicate state transitions.
In our thcoretical analysis, represent each scarch trace sequence of discrete tokens. Figure 10 shows the tokenization scheme for leaf sampling algorithms.
C.2 Trace Format for Transformers Trained from Scratch
To simplify both training and inference of Transformer models trained from scratch, we provide the full set of unvisited states Fi-1 at step t. Additionally, each unvisited state, denoted as S,-1,1, St-1,2,... ......... St F1 is accompanied by its O-based index, denoted as i1, i2,..., & F, 1/- As result, the Transformer models should predict these indices when predicting next states. It implies that the task becomes easier than the task of predicting states directly. The example this trace format as follows:
V # in Sc-1,1 i2 S i S.-1I ? St Ve # in St,1 i2 S 11.2111 St,
Tokens for single step
where V, is estimated value of S, at step t. In this setting, discretize Vi to two decimal places and represent it using 101 tokens corresponding to the values 0.00, 0.01, 0.02, 0.99, 1.00.
24

---

Published in Transactions on Machine Learning Research (03/2026)
start_os _iteration rod0>iod1 r0d0>i1d1 2 r0d0>i2d1 selected_ child
tteetteee_ee__ chward
0.00 start_ of iiteratiiinnn 0 r0d0>10d1 r0d0>11d1 r0d0>12d1>11d2 rod0~12d1>11d2
r0d0>1241>1122 selected, seaected_child_eeee aal_aantddd_eaarrreeeeeeeeeeeeee
0.00 start_of_iteratioo 0 rod0>i1d1
r040>1241>>1022 rod0>12d1>11d2 rod0>1241>1222 4 r040>1041>1002 5 ro40>10d1>11d2 6
r0d0>10d1>112d2 polected_child and seed__hee__eeee then_reward 0.00
0 r0d0>i1d1
r0d0>12d1>>1042 r040-12d1>i1d2 r040>12a1>1142 4 r0d0>1041>>1022 5 r0d0>10d1>11d2 6
r080-110d1-1222-0063 7 ro00-10d1::222>111. ro00-1001>>222>-128 selected ceelected_chiid_aant shen_reward
0.40
coart_of_iteratiio rod0>i1d1 r040>12d1>10022 rod0>i2a1>11d2 r0d0>12d1>12d2
4 ro40>1041>11d2 r000011001>1222>1063 6 ro00>1041>1140>>116 r000~1001--220>1228 8
7000>1001>1022>1066 9 ro00>10a111142>11111
r0d0>1001111442112aa
1.00 start_of_ ittirtettiion 0 rod0>11d1 rod0>12d1>1022 2 ro40>12t1>i1d2 3 r0d0>i2d1>i242
start_of_iteration
#0d0>10d1>x112 #060:1001::22211063 r000>1001>122>1113 r060>10d11:262>11288
+040-1001-:002>-118 9 r0000~0011>104>>22228 selected_chhie child and Sthen_rewark 0.00
Figure 11: Trace example of multi-reward tree search problems.
25

---

Published in Transactions on Machine Learning Research (03/2026)
start_of_ of__teaaiiin x0y0>x0y1 1 x0y0>x1y0 selected, child and then_ 0 0.00
start_ serrt_of_iteratiii x0y0>xly0 x0y0>x0y1xx0y² x0y0xx01>xx0y0 x0y0>x0y1xx1y1
selected_child_ aaapaidddnndd_heeee then_reward
rewardd 1 0.10 start_of offiieratiion x0y0>xly0 x0y0xxy1>xx0y
x0y0>x0y1xxly1 x0y0>x0y1>x022x0013 4 x0y0xx011x0022x0y1 selected_ child aad_thhn_rewwrre
0.10 start_of_ o x0y0>x1y0 x0y0>x0y1xx0y0 2 x0y0>x0y1xxly1 3 x0y0xxy1>x0y22x01
x000=x0111xx7222xx3>xxxxx #000x0y1xx0/2x0xxxxx2 selected_ childdhandddddddeeeeeeee and and__the__reeard 4 0.
start_ of it_iitraaiion x0y0>xly0 x0y0>x0112x070 2 x0y0x001>xxyy1 x0y0>x0y1x0022x001 4
x0y0xx0x1xx02220033>xxx22
x0000011x00222003xxxxxxxxxx
secected_chilldaaee then_
taa___the_____eeeeeeeeeeeee reward 4 0.00 chart_of_iteration x0y0>xly0 x0y0>xy11xx0y0
x0y0>x01>xxy11 xx00xx011x0022x0011 4 =ggpx0011xx022xxxxxxxxx2222
20x0=x00)xx022xx33xxxxxxxxx 200Dxx011xxx22xxxxxxxxxxxx2 7 xxy0000y2xxxxy22>xxxxxxxx2xxXx
selected_ chileddd__hldd_eeeeedee_d and_ neward 0.00
Figure 12: Trace example of multi-reward navigation problems.
Figures 11 and show the trace examples of multi-reward Iree search and navigation problems. These tokens are split by whitespaces, which implies that a single word corresponds to a single token. For the multi-reward navigation problem, to reduce the number of state tokens, we additionally split the state tokens with ">."
C.3 Trace Format for Pretrained Large Language Models
For the cxisting pretrained LLM, we first initialize it with the system prompt shown Figure of the trajectory (st; Ue, (s:)) is encoded into the following pair of messages:
13
Each step
(a) An assistant message
SLAECTED_STATER
(b) A uSer message
SRODBACK_SCORAA for STATE 8:: VF
WEILD_STATES for STATE 8: 8t,0 Sr,1
We then let the LLM generate the next assistant message and parse out the newly selected state.
26

---

On each subsequent turn, the user will send information about the most recently selected state Published in Transactions on Machine Learning Research (03/2026)
You Tree-Search Assistant. Your mission is locate, within fixed number of iterations, the leaf node with the highest reward. You'l1 follow a strict, turn-by-turn protocol with the
### Initialization
When the user's first message is 'START', immediately reply: State the root.
SELECTED SSATE:
#** Turn Format
N:
1. SEEDBACK_SCORE for STATE: <score>
This score is a stochastic estimate of the **average reward** the leaf nodes reachable from
this state.
2. CHILD_STARE for STATE N: <c1> <c2>
<ck>
If there
IDs after CHILD_STATES then N is leaf.
**Important:**
The FEEDBACK_SONE SCORE is a stochastic estimate of the average reward of the leaf nodes reachable from this state.
You only learn state's score when you select it. You not know the scores its children or any other unselected states until you select them.
### 3. Your Selection Rule
At each turn, pick exactly one unvisited state from any the child-state lists you've received so far (including lists from previous turns). Never select a leaf node that you have already visited. Reply with only:
SELECTED_STATE, X where X is the ID of the new state you're choosing.
### 4. Objective
Multiple leaf nodes have non-zert rewards; your goal is find one the highest reward.
Use all received FEEDBACK SCORES and CHILD STATES data to guide your choices.
Aim to identify the higaedt-reward leaf node in few selections possible.
### 5. Strict Output Format
Only output lines of the form SELECTED_STAR X. Do not echo input, add commentary, output anything else.
Begin when the says START.
Figure 13: System prompt used to drive the Tree-Search Assistant.
27

---

Published in Transactions on Machine Learning Research (03/2026)
C.4 Trace Format for Academic Paper Search Problem
In experiments the Academic Paper Search problem,
pretrained LLMs with the system prompts At each iteration, the LLM selects the next paper.
shown in Figures 14 and 15
For all methods, we evaluate the entire path as a single sequence and ask LLMs to assign one suitability score belween 0.00 and 1.00 using the prompt Figure 14 For the method without external search, we instead ask the LLM select the single most suitable candidate paper for reaching the target paper, using the prompt in Figure 15
You are an AI assistant specialized in finding a navigable academic paper from a given start paper (start_title to target paper (target_title).
1. Input:
PATH: The previously selected sequence of papers.
Example: PATH: "Paper A" -> "Paper B" -> "Paper D" ->
2. Task:
Evaluate the entire PATH whole.
Assign one suitability score from 0.00 to 1.00 indicating how appropriate the current PATH is for reaching title)
3. Dutput:
On the last line only, produce a single "score" and its value with two decimal places.
Example: 0.76
No extra text or explanation.
Figure 14: System prompt used to score the path in the Academic Paper Search problem.
You are an AI assistant specialized in finding a navigable academic paper from a start paper to target paper.
1. Input:
PATH: The previously selected sequence of papers.
Example: PATH: "Paper A" -> "Paper B" -> "Paper D" ->
(OORE_LIST, list of the scores for the corresponding papers in PATH.
Example: SCORE_LISS LIST: 0.43 0.22 0.83
CANDIDATE_LIST: 1-based indexed list of candidate papers.
Example: CANDIDATE LIST: 1. "Paper A", 2. "Paper B", 3. "Paper C"
2. Task:
Balance exploration and exploitation when selecting the next paper from CANDIDATE_
Select the single most suitable paper for reaching the target paper.
Do not estimate any candidate the target itself.
3. Dutput:
On the last line only, output exactly one integer: the index of the chosen paper from CANDIDATE_LIST.
Example:
No extra text explanation.
Figure 15: System prompt used to select the candidate in the Academic Paper Search problem.
28

---

Published in Transactions on Machine Learning Research (03/2026)
D Theoretical Constructions of Transformers
In this section, we provide the details of our thcoretical results missing in Section
D.1 Theoretical Transformer Model
This section defines a simplified version of the Transformer architecture used in this work. We make similar assumptions as in (Hahn 2020; WWiss et al., 2021p. We omit layer normalization and replace the fully connected layer with any arbitrary token-wise function (since MLPs are universal approximators). We also replace the conventional softmax attention mechanism with hard attention.
Let X E Rdx= be the imput matrix representing a sequence of n token embeddings, each of dimension d. A single-head Transformer block consists of residual self-attention mechanism followed by position-wise feed-forward function f:
Transformer Block
First, the self-a attention mechanism updates the input X:
Attn = X VX Mardmax(XTK") QX) E R
Here, Q.K.V Rdxd are the learnable query, key, and value weight matrices, respectively. The term X KTQX calculates the attention weight matrix, where entry (i,5) represents the attention from token j to token and we define the hardmax operator as follows:
1 1 = max 2R
hardmax(2);
Σ3 1 1 max 2k
This operator assigns uniform probability distribution over the index (or indices) corresponding to the maximum value(s) in and zero otherwise. Next, a foed-forward function f is applied independently to each token's representation (cach column vector):
Block(X) Attn(X) + [A))tt())
Rd
where:
Attn(X).g denotes the j-th column (token representation) after the attention step;
f Rd - Rd is any arbitrary piecewise continuous function. The output of f for each token is added residually to the output of the attention layer.
Transformer Network A Transformer network is formed by stacking l such blocks. Given an initial input sequence and fixed positional emboddings P1, P2..... , Rd, we first incorporate positional information:
x( X + [pI. ....P....PPll... .P]]
Then, we iteratively apply the blocks:
(k) = Block (ool(xxee) (((k-))), for k 1....,e.
A final linear layer (the "uncmbedding" matrix) U E Roxd is applied to the final representation of the last token in the sequence, x:()) in, where is the vocabulary size:
logits = U E R.,
The predicted next token is the one corresponding the highest logit value:
next token arg mlogits,,
i.e., we do greedy decoding. We break ties uniformly random.
29

---

Published in Transactions on Machine Learning Research (03/2026)
Discussion on Model Assumptions The omission of layer normalization follows established practice in theoretical Transformer constructions including universal approximation proofs (Yun et al. 2020) and constructing programming languages that can be compiled into Transformers (Weiss et al. 2021
Our hard attention assumption is mild that can be approximated arbitrarily closely by softmax at- tention through logit scaling. It is also a standard assumption used in previous Transformer complexity analyses (Pérez al. 2019).
Our requirement for feed-forward networks to represent arbitrary piccewise-contimmous functions aligns with classical universal approxination theorems. Specifically, ReLU network have been shown to be able to approximate any function in Sobolev spaces (Yarotsky 2017 with tight bounds on the parameters needed to achieve arbitrary closeness.
Lastly, we note, these assumptions match those in the RASP programming model (Weiss et al. 2021 which similarly removes normalization layers and assumes full expressivity of foed-forward layers.
D.2 Proof of Theorem 1
Embedding Construction We embed each token into vector whose coordinates ("registers") store interpretable features. Table 2 lists the registers, their meaning, and the initial values for each token type. The id; register is a collection of registers, each corresponding to one state token S, The last three registers (isVisited, inhValue, and pos) are initialized zero and will be updated by the Transformer. Because, we require a register for each state, follows that embedding dimension is size @(T
Notation We write 6 for the standard basis vector with an 1 in the coordinate corresponding to register and elsewhere. For registers and b, define
Ma-ob = eb e
which is the matrix that copies the value from register into register b.
Table 2: Initial embeddings. Each row is register (coordinate), and each column gives its initial value for the token types V. (value), "# ? 6n separators, and state tokens St.
Register
Va
#
?
%
S
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
id,
0
0
1
A:]
bias
1
1
1
Dynamically updated registers:
isVisited
0
0
0
0
inhValue
0
0
0
0
pos
0
0
0
0
We also add a positional embodding by sctting
X
X₂ + Egan:
so that the pos register of the token at position stores its index in the trace.
Layer Marking Visited States We want each state token to know whether follows a ? (i.c., was selected by the Transformer for selection) if it was generated by the environment indicate the children states.
30

---

Published in Transactions on Machine Learning Research (03/2026)
Choose
Q M1ia-11 K M187-+11 V -ar-staiistted
Here bias queries uniformly, K gives score 1 to separators (tokens with is#7=1), and V copies the separator's #? (+1 for #, -1 for ?) into the isVisited register of the querying token. Since a state immodiately after sees one more -1 than +1, its isVisited becomes uegalive; after sees equal counts and remains zero. Thus isVisited
exactly for visited states.
We then apply ford-forward that scales the isValue register of value tokens by their index in the trace:
(1) = (VieValue Upos) CisValue:
which will be used in subsequent layers that state tokens attend to only the closest preceding value token.
Layer Propagating Inherited Rewards Next we want each frontier state (states that are children of visited states but are themselves not yet visited) to inherit the rollout estimate value of its parent state. Parent values live in the most recent value token Va before the state's position. Set
Q Moias-11 K Misvalue
M
Each state token attends (via hard attention) to the immediately preceding value token (due to the feedfor- ward in the previous layer), and the V matrix writes that value a into its inhValue register.
We then use a piecewise feed-forward on the id; registers:
(y) {
yintvalues ILVivviatte
otherwise,
with all other coordinates set to zero. Non-state tokens and visited state tokens get a negative score, and each frontier state S, has its inherited reward stored in id,
Layer 3: Selecting the Maximum Finally, we collect all frontier scores and choose the largest:
Q = Moise K Miestate-11:
Σ Mia,-iid
At the embedding corresponding to the last ? in the trace, the query attends to all preceding state tokens (isState=1) and V sums each state St's id; value into the querying token's id, We set. (() = as we do not need the feedforward in this layer. Now assuming tokens are indexed so that the embedding coordinate id; corresponds exactly to the vocabulary index of state token S,, we choose the unembedding matrix U be zero everywhere except on the submatrix mapping the id, coordinates to their corresponding token logits, where it is the identity. In this setup, the id, registers directly become the logits for cach state token, greedy decoding picks the state with the highest inherited reward.
Remark (Uniform-Leaf Sampling) If instead all V. embeddings are set to the same constant vector, then every frontier state receives identical score in Layer 2, yielding uniform sampling over leaves by the same construction.
D.3 Proof of Theorem
For this theorem, we analyze the conventional MCTS (Sutton & Barto 2018 which is a widely-used form, where the selected state s* is determined by taking the arg min over N(s), rather than (s) as in Algorithm 0 In our empirical experiments, we used N (s) due to the shallow nature of the search trees. However, for theoretical clarity and consistency with classical MCTS, the Transformer construction presented here uses N(s).
31

---

Published in Transactions on Machine Learning Research (03/2026)
Embedding Construction Tokens again embedded into interpretable "registers." Table 3 defines the initial static register values; all dynamic registers are initialized to zero. Registers of type X, denote a set [x,] TB with one register per state token. The oid; registers include two additional coordinates used specifically for the and tokens.
As in the previous construction, we apply a positional embedding via x()) + x; + t Cpos: so that cach token's 8 pos register holds its position in the trace.
Table 3: Initial embeddings for the MCTS- TETTS-Tranforme:: static registers encode each token's fixed features, while dynamic registers are initialized to zero and updated by the network conditioned on the trace.
Register
Va
#
?
[bos]
%
S
value
0
0
0
0
0
#?
+1
0
0
0
0
isValue
0
0
0
0
is#?
0
0
is>
O
0
0
0
0
is?
O
1
0
0
0
0
isBos
0
0
0
0
11111
1
0
0
is#Bos
0
1
0
0
0
0
1000
isState
0
0
0
000
1
id;
0
0
1
=
A:
bias
Dynamically updated registers:
is? opos
isValue pos
isState pos
0
closest?pos
0
0
0
0
parentpos
0
0
0
isVisited
0
0
wasVisited
vid,
cid,
pid,
psid,
0
nsid,
0
oid,
0
0
0 0
iter
pos
0
For any coordinates where we do not explicitly define how it is updated by a feedforward function, the output is defined be zero.
Layer 1: Iteration Counter and Positional Precomputation This layer precomputes positional dependent registers. The attention mechanism defined by the following:
Q Moias->1,
K M
V MisBos Mitrr
32

---

Published in Transactions on Machine Learning Research (03/2026)
This attention is used to compute the number of iterations that have passed which will be stored in the iter register. We apply the feed-forward function f();;
(1) 1/2=?Upom
f()) (1) seState VisStatel/pos
f(ll (1) LeValue-pos VieValueUpoos
(1) ster (L/Viser) Vivers
if Viter else 0,
which processes the iter register to hold the correct value and preproccess several other positionally depen- dent registers for later computation.
Layer 2: Compute closest?pos This layer will do computation so that when isState is 1, closest?pos will store the position of the closest preceding token in the trace. Otherwise, it is set to 0. The attention matrices are given by the following:
Q Mbise-1:
Mist-oa--11
Mpoa-
Through the hard attention, each token attends to the preceding token with the largest is? pos value (i.e., the value of the closest preceding ? token). Its pos value is copied to the querying token's closest?pos register. The subsequent food-forward function f(2) is:
((2)
clocest?pos VclosestPpos VieState
VelosestPpos
which zeros out closest?pos if the current token does not correspond to a state.
Layer 3: Compute Per-Iteration Reward and Count Statistics This layer computes the reward and tree policy visitation count statistics for each iteration, which will be stored the respective Va token at the end of its iteration. The attention is:
Mbias->1,
K Melosest7poa-11;
V = Σ (Maa + Mıa,
For a Vo token, this attention summs statistics from state tokens visited by the tree policy within the current iteration, identified as the tokens with the largest value in closest?P. The id; values from these sy. tokens are copied and summed into the V, token's vid, and cid, registers. The feed-forward function f(3), applicd Va tokens, then sets:
f (3) (y)via, Ivalue [ 0] Dvid.
(3) (y)cia, = [[ves, > 0] - Yvid.
storing in cid, when state S, is visited this iteration (used to compute visitation counts).
Layer 4: Aggregate Statistics (Backpropagation) This layer aggregates reward and count statistics from all iterations. These accumulated statistics will be stored the final V, token of the trace. We set:
Q Motas>11:
K MisValue+11
V = Σ (Myid + Meid.-ridd.)
33

---

Published in Transactions on Machine Learning Research (03/2026)
The token intended for final aggregation attends all previous Va tokens (where isValue= 1). It sums up the vid; (values) and cid; (counts) from these iteration-specific Va tokens into its own registers. We then use the feedforward layer to unormalize by the number of iterations (since the attention will average them);
(y)cia. = Meid; Miter Meid.
( (v)via Divi; Miver 1/vid,
Layer 5: Identify State Tokens Selected by Transformer This layer identifies which state tokens in the trace corresponds to those generated by the Transformer during the tree policy trace. We use attention:
Q Mbias-1;
K M1
V Mra-s188electedd
Each state token sums the ?# values from preceding relevant tokens. For states immediately following a #, this sum in isSelected will be O. It will be nonzero otherwise. The feed-forward ((3) uses this fact to compute the register:
( = ViaState 1 [ViSelleceed A 40] 1/ /isselccted
This sets isSelected to 1 for state tokens selected by the Transformer and for the state token generated by the environment to indicate the child nodes.
Layer Identify Parent Position for Neighbors This layer serves as preprocessing to identify the parent node of each neighbor token (i.e., those state tokens immediately succeeding the token), by first storing the parent's position in the parentpos register of the parent token itself. No attention is needed, we set all the attention weights to 0. The feed-forward f(() is:
( (1) paraatt ViaState V/sSelacted 1/paa-
This operation stores the current token's position into its own parentP register if it is state token and corresponds to a state visited by the tree policy in some iteration.
Layer 7: Propagate Parent id to Its Children States Children state tokens use the information about the parent's position to retrieve the parent's id. The attention is:
Q Moiss-11
K Mparentpos-11,
V
Σ Madd-pid,
Each neighbor state token attends to the immediately preceding non-neighbor state (which will be it's parent state). It then copies the parent token's id, into its own pid, registers. No feed-forward f(T) is needed, set it to 0.
Layer 8: Store Previously Selected State id in Embedding of > Token The token stores the id of the state that was just selected by the tree policy. We set:
Q Motaa-11,
Misstatappa--li
V
Σ Mid.
The > token attends to the immediately preceding state token and copies that state's id, into its psid,
34

---

Published in Transactions on Machine Learning Research (03/2026)
Layer 9: Collect Children States of Previously Selected State in Embedding of > Token The token collects the id of all children of the previously selected state. The attention is:
Q
Σ
Mpsid
K
Σ
Mpid
V =
Σ Mia -nsidd
The token uses its psid, registers as the query. State tokens use their pid, registers as keys. If key and query match, the neighbor's id; is copied/summeet into the token's nsid, register. No feed-forward ((9) is needed, so we set it to
Layer 10: Final UCT Computation for > Token At the token, the next state to select is calculated using UCT scores, First, it gathers aggregated statistics using the following attention weights:
Q Mbias-11:
K = pos--lue-pos-t1.
Σ (Mesd +
The token attends the token holding globally aggregated statistics (computed layer 4) into own registers. The feed-forward function f(10) then does the UCT computation:
Veid. + c (loloo()) Holde Upsid,) yosa. Via> = 1, Ihoid, and Yeid, + 0, if Vis> 1, Unsid, and Yeid, otherwise.
(1)0
0c,
0,
Layer 11: Generate Start State Sy After ? This layer ensures the Transformer generates Sy (the root node) when it sees the initial token. The attention matrices are
Q 0,
O, V Mis?-roidn
We set (11) to the function in this layer.
Layer 12: Generate or After State Selection This layer determines whether to generate or % after a state token has been chosen by the tree policy, Specifically is generated if the tree policy has not reached new undiscovered node. The % token is generated otherwise. For this layer, set the attention matrices to the following:
Σ Msa
K
Σ Mad--ri
V Missalectrd-wassSleeetet
Intuitively for the current state token in the tree policy path, it will attend to all previous state tokens that represent the same states. It aggregates the isSelected values into the wasSelected register. Thus wasSelected > 0 if and only if the state was previously visited in another iteration. Thus we let the feed forward layer be
(y)osa, ViaStata (VosViaatedd > 0], (( ((12)
VinStaat UoasWistted 0).
35

---

Published in Transactions on Machine Learning Research (03/2026)
Final Unembedding Assuming without loss of generality that oid; corresponds to the i'th state token, it again sulfices to let U be the matrix that is zero everywhere except on the submatrix mapping the oid, to the corresponding token. The token with the highest logit value is chosen by greedy decoding.
Note on on Different Tree Policies For path sampling with a greedy policy and uniform path sampling, Layer 10 modified. For path sampling with a greedy policy, the feed-forward f(10) at the token computes:
(( (y)osa, = Yvid, Yier Unsid.
For uniform path sampling, f())) would set:
f(10 (y)o1d. = Blor Unsid,
36

---

Published in Transactions on Machine Learning Research (03/2026)
E Details of Empirical Analysis on Transformers' Search Capabilities
In this work, we use NumPy (Harris et al. 2020 SciPy (Virtanen et al. 2020). PyTorch (Paskke et al. 2019), Scikit-learn (Pedregosa el al. 2011) and other scientific packages in Python. Moreover, as shown in the main article, we employ a variety of I.I.M APIs and checkpoints such as OpenAT GPT, Google Gemini, and Alibaba Qwen.
To run our experiments, we utilize commercial Intel and AMD CPUs such as AMD EPYC 9374F and Intel Xeon Platinum 8352Y, and NVIDIA GPUs such NVIDIA L40S.
E.1 Experimental Details of Pretrained LLMs on Multi-Reware Tree Search Problem
To test the performance of pretrained LLM's on the Multi-Reward Tree Search problem, we use the prompting method as in Section CK3 We use tree-depth of 6 and branching factor of 2. From among all leaf nodes, we uniformly sample 8 distinct leafs to serve "goal" nodes. Each selected goal leaf is then randomly assigned unique reward from the set {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 1.0}. Tn total, generate 200 independent trees according to this procedure, resulting in 200 distinct multi-reward tree search instances.
Model Evaluation We evaluate the following models using the prompting strategy described above: GPT- 4.1-mini and GPT-4.11 (OpenAl 2025a); Gemini 2.5-Flash and Gemini 2.5-Flash (Thinking) (Gemini Team, Each model is tested for 50
Google 2025); and Qwen3-8B and Qwen3-8B (Thinking) (Qwen Team 2025 steps per tree instance.
Model Configuration The "Thinking" variants refer to configurations in which explicit reasoning modes are activated. For Gemini 2.5-Flash (Thinking), set the reasoning_effe effortt parameter to medium. For Qwen3-8B, which by default engages in "Thinking" (i.c., explicit reasoning chains), we create non-thinking variant by appending /tt__tiiinnnnn think to the beginning of the system prompt to suppress intermediate reasoning steps. Lastly, we use greedy decoding for all variants.
37

---

Published in Transactions on Machine Learning Research (03/2026)
E.2 Experimental Details of Fine- Tini-Tnnng for Enhancing Search Capabilities of LLMs
We evaluate the search abilities of pretrained and fine-tuned LLMs on the Academic Paper Search problem. Given a start paper, the goal of this problem is to find target paper where the child states of particular paper are defined by the papers of the co-authors of that particular paper. We randomly choose pairs of start and target papers as shown in Table We employ OpenAlex (Priem et al., 2022) for designing this experiment. The prompts used generate the search trajectories these experiments are shown in Section C44
We fine-tune the Qwen3-8B model (Qwen Team 2025) using Low-Rank Adaptation (LoRA) (Hu et al. 2022 parameter--fficient approach that updates only small set of parameters. The maximum sequence length is set to 10,240, and training is run for one epoch. We use the paged AdamW (32-bit) optimizer (Loshchilov & Hutter 2018) for memory efficiency, with a learning rate x 10 10- a cosine learning-rate schedule, 3% warmup, and weight decay 0,001, LoRA is configured with rank 16, scaling factor 32, and dropout 0.05, and is applied to the main attention projections and MLP projections.
Table Start-targer paper pairs in the Academic Paper Search problem.
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
W26445342909
W4403312580
W2360871073
W2290528378
W2184245987
72075179071
W2413728222
W1497510206
W193222316
W1928336576
W4294761907
W2378250669
W2167634255
W70597564
W201963696
W2413184738
W2994021882
W2111031504
W153886177 733
W31 25578107
W42054161
W199931 1920
W1648404904
W957051818
W1496033586
W1043586152
W1003374900
W4225958650
W4322489422
W1983902283
W2800480410
W657940841
W4387063953
W2124237366
W19861 39416
W2564522209
W2246-96633
W2312519939
W22631 84888
W3038306494
W3041988068
W2367761479
W2935041779
W2272635734
W3203890190
W2363737347
W2368854302
W204856840
W21 36736708823
W2330377032
W2045968954
W43120791099
W3082198414
W22909479339
W3182725060
W2053134422
W4404914457
W639808865
W658822377
44322827068
W21064420077
W2316765899
W 44350875192
W42424471155
W4205383455
W2415588074
W301 3098164
W1525894204
W1538920755
W3173752466
W2952700754
W270839233
W385335894
W20087489400
W2140973411
W2090948737
W2080692806
W4392218955
W4391238054
W20094400780
W2505367746
W46357476
477314777899
W2491 706897
W4232504235
W1461752455
W154001 1986
229339278
W2962025900
W2372889019
W2355250576
W57577176h
W427966303
W2372242901
W2998013669
4216908-4099
W4291372000
W4410927769
W4403703781
W2394974804
W4256391018
W3209857023
W2001 19911
W2369940005
W2356021002
W2393662227
W2358750724
W3153204805
W3153990973
W2134325687
W2337799504
38

---

Published in Transactions on Machine Learning Research (03/2026)
F Additional Empirical Results
In this section, show additional empirical results on the multi-reward tree search and multi-reward navigalion problems.
F.1 Additional Empirical Results on Multi- Tree Search
In addition to the experiments shown in Figure we present additional results on multi-reward tree search problems with more diverse settings.
Unilorn Leaf Sumpling
Uniforn Path Sampling
Greedy Leaf Sampling
I bill
0.0
Reference
Transformer
Reference
Transformer Path Sampling w/ Gredy
Reference
Trannformer
Path Sampling w/ Pure Exp.
Path Sampling w/ UCT
0.5
Reference
Transformer
Reference
Transformer
Reference
Trrnnformer
Maxx-Rewaed Hit Rate
Normalized Path Length
Normalzed Comulstave Revard
Discouated Cunulative Glan
Highest Reward
Normalized Jump Distance
Figure 16: Behavior cloning results on the multi-reward tree search problem, where each binary tree of depth has 8 different goals and search step budget is 50.
Figure shows the experimental results that test deeper multi-reward trees, which are depth 8. These results follow our observation that Transformers successfully mimic the behaviors of reference search algorithms. Since this setting harder than the setting of Figure the search performance of each algorithm is relatively poor.
Uniform Leaf Sampling
Uniform Path Sampling
Greedy Leaf Sampling
1.0
Score
Reference
Trausorner
Reference
Transfomer
Reference
Transformer
Path Sampling w/ Pure Exp.
Path Sampling w/ Greody
Path Sampling w/ UCT
0.0
Reference
Transformer
Reference
Transformer
Refermnce
Transformer
Max. Rewad Hit Rate
Normalisrd Path Length
Normulizel Complative Rewsari
Discountdd Currlaativ Gain
Highest Reward
Normalizel Jump Distance
Figure 17: Behavior cloning results the multi-reward tree search problem, where each quaternary tree of depth 4 has 8 different goals and search step budget is 50.
39

---

Published in Transactions on Machine Learning Research (03/2026)
Similarly, we conduct experiments wider multi-reward trees in Figure 17 These trees have 4 child states for each parent state. Transformers successfully show similar performance of the respective search strategies, following our expectation.
Uniloru Loaf Sampling
Unifona Path Sumpling
Greedy Leaf Sampling
0.5
Reference
Transformer
Reference
Transformer
Reference
Transformer
Puth Sampling w/ Pure Exp.
Path Sampling w/ Greddy
Path Sampling w/ UCT
Seore
0.5
Reference
Transformer
Reference
Transformer
Reference
Transformer
Max-Reward Rate
Normalized Path Length
Normalzzed Cummlatave Reward
Discouated Cunulative Gain
Highest Reward
Normalzzed Jamp Datance
Figure 18: Behavior cloning results on the multi-reward tree search problem, where each binary tree of depth has 4 different goals and search step budget is 50.
We test a problem with a fewer number of goal sates. As shown in Figure 18 these results follow the findings observed the previous experiments.
40

---

Published in Transactions on Machine Learning Research (03/2026)
F.2
Additional Empirical Results Generalization Analysis
Uniform Leaf
Reference
Transoomee
Uniform Path
Reference
Trasformet
0.5
0.3
0.4
0.1
Greedy Leaf
Refermmon
Transformer
Path w/ Pure Exp.
Reference
Transforme
1.0
0.5
0.0
0.1
0.2
0.3
0.4
0.1
0.2
0.3
0.1
0.1
0.2
0.3
0.1
0.1
0.2
0.3
0.4
Path w/ Greedy
Referemoc
Transformee
Path w/ UCT
Reforence
Transformet
0.1
Test Wall Density
Test Wall Density
Test Wall Density
Test Wall Density
Rax.Rrward Hit Rate
Normalized Path Length
Normalined Comolative Roward
Cunulativn Gain
Highest Reward
Normalizet Jura
Figure 19: Generalization analysis over test wall densities on the multi-reward navigation problem of size
x 4 with step budget of 50. The wall density smaller than 0.4 is unseen in a training phase. A gray dashed line indicates the sctting used in training.
41

---

Published in Transactions on Machine Learning Research (03/2026)
Uniform Leaf
Reference
Transformee
Uniform Path
Reference
Thansformer
Greedy Leaf
Relerence
Transforuce
Path w/ Pure Exp.
Reference
Transforar
1.0
0.5
0.0
6
Path w/ Greedy
Reference
Transformee
Path w/ UCT
Reference
Transfrmmr
1.0
0.5
0.0
7
# Test Tree Degti
#Tes Thee Depth
#Tost Tree Dopth
#Tost Tree Depth
Max Rlwed Hit Rate
Normalised Path Lengih
Normalzzed Commlative Reward
Discouated Cunalative Gain
Highest Reward
Normalzed Jang Datance
Generalization analysis over test tree depths on the binary tree search problem with different goal states and a step budget of 50. The test tree depth larger than unseen training. A gray dashed line indicates the setting used in training.
Figure 20:
Figures 19 and 20 demonstrate generalization analysis over test wall densities, the numbers of test goal states, and test tree depths. In Figure 19 test cases with lower wall densities are more challenging than those with higher densities, as they tend to produce deeper and wider trees. Notably, according to Figure 20 generalizing to unseen tree depths is more challenging than other generalization settings. This can be seen as a challenge analogous to length generalization, similar to what is discussed in the previous work (Lee el al.
2024 2025
G Discussion on Safeguards
To promote reliability and reproducibility in our experiments, we aim to establish a set of safeguards focused on improving traceability and robustness. Our approach includes, in part, logging key inputs, outputs, and intermediate states, as well as verifying that the model behaves consistently under fixed conditions. While these components have been partially implemented in current work, consider them part of broader, ongoing effort toward responsible model development. So far, have not observed any misbehavior in experiments.
H Declaration of LLM Usage
In addition to being the primary focus of our study, LLMs were used throughout our research process. We leveraged them to generate and refine system prompts for all experiments involving pretrained LLM models, assisted in writing the code used to run these experiments, and revised the manuscript itself. Specifically, we employed LLM platforms developed by OpenAl, Google Gemini, and Ollama.
42


Reviews:
[{"review_id": "Cf0hFifwnM", "reviewer": "Reviewer_ML1j", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The work analyzes whether transformers can implement tree-search strategies on two controlled benchmarks which implement unknown tree search with bandit feedback. The environment reveals the next plausible nodes given the current state and a noisy value estimate for each. The model then selects the next node from this list. The paper first shows explicit constructions showing constant-depth Transformers can exactly implement several selection policies. Then, the work shows behavior cloning results where small Transformers trained from scratch can match reference algorithms’s aggregate search metrics on two synthetic environments (tree-search and navigation) and generalize to some shifted conditions (more steps, deeper trees). The work also shows that fine-tuning improves performance. Overall, it studies selection policies under uncertainty, and it surfaces a methodological point that I really liked: matching metrics does not necessarily mean the model learned the underlying algorithm (kl divergence of next-state selection distributions indicate very different behavior).", "claims_and_evidence_yn": "Yes", "claims_explanation": "I worry that the paper’s framing overreaches the evidence provided.  However, if claims are adjusted it’s a nice analysis of the selection problem.\n\nW1 The benchmark removes two of the most difficult components of search in LLM settings: generating candidates and evaluating the value of partial solutions. By providing these externally, the task is effectively reduced to supervised ranking or classification over an enumerated set. While the model performs selection -- claiming it has internalized the full search process seems to me like a strong claim.\n\nW2 There is a tension between the paper’s core claim that Transformers do implement search algorithms and its own experimental results. The KL divergence analysis shows that for path-based strategies, the model produces action distributions that differ distinctively from the underlying algorithms. This undermines the claim that the model is implementing the target algorithm; it is simply achieving a similar outcome, likely through different heuristics.\n\nW3 While the theoretical guarantees are technically correct, they feel weakly connected to the empirical results in this paper. The proofs rely on several idealizations, essentially saying that a Transformer can encode a lookup-table implementation of these algorithms. This does not offer much predictive power regarding how the models behave in the experiments, nor does it provide a way for the trained models to utilize the theoretical constructions to implement the underlying algorithm.", "audience_yn": "Yes", "audience_explanation": "Knowing whether LLMs can be enhanced using external search algorithms or whether they can internalize them is an interesting problem. The paper itself has several nice aspects:\n\nThe paper is methodologically strong in its decomposition of the search problem. By externalizing expansion and bandit feedback, the benchmark focuses directly on the selection problem and enables controlled comparisons across policies. It shows transformers have capacity to implement the search algorithms in these settings.\n\nThe experimental evaluation is also extensive, covering multiple metrics across both tree and maze environments.\n\nI particularly appreciated the nuance in the KL divergence analysis. The work shows that path-based policies can match aggregate performance metrics while diverging significantly in their next-action distributions (high KL divergence). This is interesting evidence showing that models do not successfully internalize the specific algorithms but achieve similar outcomes.", "requested_changes": "Address the above weaknesses.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "1vr15B9WMa", "reviewer": "Reviewer_N5Wq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates whether Transformer architectures can internalize search algorithms to solve problems where the search space is not known a priori but revealed incrementally (unknown tree search with bandit feedback). The authors propose a simplified framework where tree extensions and feedback signals are externally specified, allowing for controlled evaluation of search capabilities. Theoretical analysis shows that Transformers are expressive enough to implement various search strategies and empirically they can learn them via behavior cloning. But off-the-shelf LLMs do not currently exploit this capacity well; targeted fine-tuning on search trajectories can substantially enhance their search behavior.\n\nPros: \n- The paper introduces a novel \"unknown tree search with bandit feedback\" set up isolating the model's search selection ability\n- Theorems 1 & 2 provide well-demonstration that Transformers with constant depth and specific embedding dimensions can exactly implement leaf-based and path-based search policies.\n- Provide comprehensive empirical analysis and generalization experiments for training Transformers from scratch to clone various search algorithms, and generalize to unseen conditions.\n\nCons:\n- The theoretical results are based on hard attention, the practical connection to actual multi-head softmax-based Transformers is unknown (which is acknowledged). \n- The experiments are primarily on synthetic problems and academic paper search, which may not fully capture the complexity of real-world reasoning tasks. Generalization to longer sequences and larger problem instances remains challenging (which is acknowledged).", "claims_and_evidence_yn": "Yes", "claims_explanation": "The claims are mostly well-supported\n- Theoretical Claims: Theorems 1 & 2 provide explicit constructions and proofs to show that Transformers with constant depth and specific embedding dimensions can exactly implement leaf-based and path-based search policies.\n- Empirical Claims: Comprehensive experiments are conducted, with clear metrics and comparisons to reference algorithms. Figures 7, 8, 19, and 20, show the testing on the models on unseen tree depths, step counts, and wall densities to support the generalization claim. But the claim of \"generalizing to unseen conditions\" is somewhat overstated given that generalization to deeper trees shows significant degradation (Figure 20). The improvement in LLM performance after targeted fine-tuning is also evaluated in Table 1, with comparisons to both pretrained LLMs and LLMs combined with external algorithms.", "audience_yn": "Yes", "audience_explanation": "Understanding whether the Transformer architecture itself can internalize search loops is a core problem of LLMs’ ability. It will be interesting to broader community and research related to thinking, reasoning-focused, search-based inference, and LLM-as-planners, agentic behavior.", "requested_changes": "- Explicitly discuss the use of hardmax vs softmax, also whether multi-head attention or layer norm might help or hinder such constructions. Also it would be interesting to ablate the sensitivity of components of the Transformer architecture to the search abilities. \n- Explicitly discuss tokenization assumptions.\n- More explanation on the path-based KL divergence results. \n- Provide at least one real world experiment could strengthen the paper.\n- The fine-tuning experiments rely on trajectories generated by proprietary models like Gemini. This can potentially inject unknown biases and behaviors into the trained models. A short comment on this would be appropriate.", "broader_impact_concerns": "No major ethical concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "gykscRznit", "reviewer": "Reviewer_SX3A", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates whether Transformers can internalize and execute structured search algorithms without reliance on external search modules. The authors propose a framework called \"Unknown Tree Search with Bandit Feedback\" to isolate the model's search policy from the environment dynamics. \n\nThe contributions are threefold:\n\n- Theoretical Analysis: The authors provide constructive proofs demonstrating that Transformers of constant depth are theoretically expressive enough to implement various leaf-based and path-based search strategies.\n\n- Empirical Verification (Behavior Cloning): Through experiments on synthetic tree and navigation tasks, the paper shows that Transformers trained from scratch can clone the behavior of reference algorithms, replicating both performance outcomes and behavioral patterns (e.g., jump distances).\n\n- Targeted Fine-Tuning: The authors demonstrate that fine-tuning a pre-trained LLM (Qwen3-8B) on search trajectories significantly improves its performance on a realistic \"Academic Paper Search\" task compared to standard prompting baselines.\n\n\n**Strengths**\n\n- Clear Motivation and Accessible Narrative: The paper is well-written and easy to follow. The motivation is compelling: moving beyond simply \"using\" LLMs for reasoning to understanding if they can internalize the search process itself. The progression from theoretical proofs to synthetic experiments and finally to a realistic downstream task creates a logical and cohesive narrative structure.\n\n- Effective Use of Synthetic Environments to Isolate Search Logic: A key strength is the introduction of the \"Unknown Tree Search with Bandit Feedback\" framework utilizing controlled synthetic environments. By externalizing state expansion and using abstract states, the authors effectively decouple the \"Search Policy\" (decided by the Agent) from the environment dynamics. This design is crucial for isolating the model's pure algorithmic ability to balance exploration and exploitation, ensuring that the evaluation is not confounded by the model's vast pre-trained background knowledge or semantic priors.\n\n- Comprehensive Experimental Design on Synthetic Data: Chapter 5 provides a thorough empirical analysis. The authors go beyond simple performance metrics and include behavioral metrics such as KL divergence and normalized jump distance. This multi-faceted analysis convincingly demonstrates that the model is not just achieving similar results but is actually mimicking the process of the reference algorithms. Furthermore, the inclusion of generalization tests (varying wall densities, step budgets, etc.) adds robustness to the claims of behavior cloning.\n\n\n**Weaknesses**\n\n- Concerns Regarding Synthetic Data Scale and Generalization: My primary concern lies in the complexity of the synthetic environments. The experiments primarily rely on small-scale settings, such as binary trees of depth 6 or 4x4 grids. While the paper claims generalization capabilities, the results in Appendix Figure 20 paint a different picture. There is a significant performance drop when the tree depth increases from 6 (seen) to 8 (unseen). This sharp degradation questions the core claim that the Transformer has truly \"internalized\" the recursive search algorithm. It suggests the model might be relying on shallow pattern matching or local heuristics that fail to scale, rather than learning a generalizable search procedure.\n\n- Lack of Detail and Motivation in Targeted Fine-Tuning (Chapter 6): The transition to the \"Academic Paper Search\" task in Chapter 6 feels somewhat disconnected from the rigorous synthetic analysis. It lacks crucial experimental details regarding the fine-tuning process. For instance, the size of the training dataset, the number of epochs are not clearly detailed in the main text. The motivation for this specific task setup needs strengthening. It is unclear why this specific domain was chosen as the primary validator for \"real-world\" performance over more standard reasoning benchmarks.\n\n- Reliability of LLM-Generated Bandit Feedback: The real-world scenarios rely on the assumption that \"we can prompt the LLM to directly estimate state values... These realizations serve as noisy bandit feedback signals\". This introduces a critical unverified variable: How accurate are these LLM-generated values? If the \"bandit feedback\" is highly inaccurate or inconsistent, the search agent is effectively navigating with a broken compass. The paper focuses on the search policy but takes the reward signal for granted. A high search capability on top of unreliable value estimation does not yield a good solver. The paper should evaluate the correlation between the LLM's estimated values and the ground truth (or task success) to validate this setup", "claims_and_evidence_yn": "Yes", "claims_explanation": "the claims made in the submission are supported by accurate, convincing, and clear evidence, particularly regarding the feasibility of Transformers internalizing search algorithms. The experimental design in Chapter 5 is careful, and the results provide robust support for the authors' central hypothesis.\n\n- Isolation of Variables: By using synthetic environments rather than natural language tasks, the authors effectively eliminate confounding factors such as the model's pre-training knowledge or semantic hallucinations. This ensures that any observed performance is due to the learned search policy .\n\n- Comprehensive Metrics: The authors go beyond simple success rates. They employ process-oriented metrics such as Normalized Jump Distance and KL Divergence. This design choice is crucial because it verifies that the model is not just achieving the same result by chance, but is actually replicating the behavioral patterns of the reference algorithms.\n\n- Successful Behavior Cloning: The experiments demonstrate that Transformers trained from scratch can match the performance of specific reference algorithms with high precision. The low KL divergence for leaf-based strategies further confirms that the models can approximate the target probability distributions accurately.\n\n- Evidence of Internalization: The analysis of Jump Distance provides convincing evidence that the models have internalized the structural traversal logic of the algorithms, distinct from random exploration.\n\n- Clear Evidence of Generalization **(with Limitations)** The authors provide clear evidence that the learned behaviors are not merely memorized but generalized:", "audience_yn": "Yes", "audience_explanation": "This is a highly relevant direction as it attempts to rigorously analyze the planning and reasoning limits of LLMs by isolating their search policy from their world knowledge using the \"Unknown Tree Search with Bandit Feedback\" framework. The paper provides a comprehensive analysis by combining theoretical proofs of the Transformer's expressivity to represent these algorithms with empirical evidence that they can be trained to clone complex search behaviors and generalize to unseen conditions.", "requested_changes": "- Contextualize Generalization Limits (Critical): While the paper claims the models can generalize, the results in Appendix Figure 20  show a clear performance degradation when the tree depth increases from 6 to 8. This suggests that the \"internalized search\" might be limited to a finite horizon or specific pattern matching rather than a fully recursive algorithm.\n\n  Request: Please explicitly discuss this degradation in the main text (Section 5 or Limitations). Analyze why the performance drops—is it due to context length limitations, attention dispersion, or a breakdown of the learned heuristic? The claims regarding generalization should be calibrated to reflect these empirical bounds (e.g., \"generalizes within a local range\" rather than broad generalization).\n\n- Provide Missing Details for Targeted Fine-Tuning (Critical) Section 6 demonstrates that fine-tuning Qwen3-8B improves search capabilities, but lacks critical reproducibility details.\n\n  Request: Please add a detailed Appendix describing the experimental setup for the \"Academic Paper Search\" fine-tuning.\n\n- Validate LLM-Generated Bandit Feedback (Strengthening) The framework relies on the premise that LLMs can provide useful \"noisy bandit feedback signals\". However, LLMs are known to struggle with self-evaluation and scalar reward estimation. If the feedback is too noisy, the search agent is effectively navigating blindly.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)