Transformers Struggle to Learn to Search
Authors: Abulhair Saparov
Srushti Ajay Pawar
Shreyas Pimpalgaonkar
Nitish Joshi
Richard Yuanzhe Pang
Vishakh Padmakumar
Mehran Kazemi
Najoung Kim
He He
Venue: ICLR 2025 Poster (2025)
Keywords: search
reasoning
transformers
scaling laws
mechanistic interpretability
circuit analysis
Primary area: interpretability and explainable AI

Abstract: Search is an ability foundational in many important tasks, and recent studies have shown that large language models (LLMs) struggle to perform search robustly. It is unknown whether this inability is due to a lack of data, insufficient model parameters, or fundamental limitations of the transformer architecture. In this work, we use the foundational graph connectivity problem as a testbed to generate effectively limitless high-coverage data to train small transformers and test whether they can learn to perform search. We find that, when given the right training distribution, the transformer is able to learn to search.

We analyze the algorithm that the transformer has learned through a novel mechanistic interpretability technique that enables us to extract the computation graph from the trained model. We find that for each vertex in the input graph, transformers compute the set of vertices reachable from that vertex. Each layer then progressively expands these sets, allowing the model to search over a number of vertices exponential in the number of layers.

However, we find that as the input graph size increases, the transformer has greater difficulty in learning the task. This difficulty is not resolved even as the number of parameters is increased, suggesting that increasing model scale will not lead to robust search abilities. We also find that performing search in-context (i.e., chain-of-thought) does not resolve this inability to learn to search on larger graphs.

Published as a conference paper at ICLR 2025
TRANSFORMERS STRUGGLE TO LEARN TO SEARCH
Abulhair Saparoy*: x.1
Srushti Pawar Shreyas Pimpalgaonkar
Nitish Joshi² Seyed Mehran Kazemi
Richard Yuanzhe Pang Najoung Kim*.4
Vishakh Padmakumar"
He He
Purdue University, NNew York University, "Google, 'Boston University
asaparov@purdue. edu
ABSTRACT
Search is an ability foundational in many important tasks, and recent studies have shown that large language models (LLMs) struggle to perform search robustly. It is unknown whether this inability is due to a lack of data, insufficient model pa- rameters, or fundamental limitations of the transformer architecture. In this work, we use the foundational graph connectivity problem as a testbed to generate ef- fectively limitless high-coverage data to train small transformers and test whether they can learn to perform search. We find that, when given the right training dis- tribution, the transformer is able learn to search.
We analyze the algorithm that the transformer has learned through novel mech- anistic interpretability technique that enables us to extract the computation graph from the trained model. We find that transformers perform search at every vertex in parallel: For each vertex in the input graph, transformers compute the set of vertices reachable from that vertex. Each layer then progressively expands these sets, allowing the model to search over a number of vertices exponential Mayerr
However, we find that as the input graph size increases, the transformer has greater difficulty in learning the task. This difficulty is not resolved even as the number of parameters increased, suggesting that increasing model scale will not lead robust search abilities. We also find that performing search in-context (i.e., chain- of-thought) does not resolve this inability learn to search larger graphs.
INTRODUCTION
The ability to search is fundamental in many important tasks, such as reasoning (Yao et al., 2024; Kazemi et al., 2023; Hao et al., 2023), planning (Stein & Koller, 2023; Valmeekam et al., 2022), and navigation (Ding et al., 2024). Recent work has demonstrated that transformer- models (LLMs) struggle with proof search (Saparov & He, 2022; Valmeekam et al., 2022; Kamb- hampati et al., 2024). It is unknown whether this shortcoming is due to lack of data, insufficient model parameters, or a fundamental limitation of the transformer architecture. In any case, as the scale of L.LMs continues increase, it is yet unclear whether future L.MMs, equipped with more data, parameters, and compute, will be able to perform search and planning. Chain-of-thought and similar prompting techniques (Wei et al., 2022b; Nye et al., 2021) have enabled L.LMs to decompose the search task into the repeated task of predicting the next step along the path to the goal. However, even in this setting, in the worst case, in order to avoid making a "wrong turn," the model must perform the search within its forward pass to determine the correct next step. And LLMs have been observed producing errors or hallucinations after taking such wrong turn (Saparov & He, 2022).
large language
We aim to shed light on this question by training small transformer models on a simple yet foun- dational search task: Given a directed acyclic graph (DAG), a start vertex, and a goal vertex, find the next vertex along a path from the start to the goal vertex. This task is the backbone of many reasoning problems as is equivalent to proof search in a simplified logic which is a subset of al- most any logic: The model must solve this task if there is any chance to generalize to more complex search and reasoning tasks. We demonstrate experimentally that transformers can indeed be taught to search, when given the right training distribution. The training distribution must be carefully con- structed so as to preclude the usefulness of shortcuts or hcuristics that would otherwise prevent the transformer from learning robust and generalizable search algorithm. By automatically generating such examples, we provide the transformer with effectively limitless and idealized training data, with which estimate "upper bound" on the transformer's 's ability to learn to search.
"quul contribution. MMuch of the work was completed while author was New York University.

---

Published as a conference paper at ICLR 2025
When the model does learn the constructed training distribution (i.e., reaches 100% training accu- racy), it is able to correctly perform search in almost any graph. We aim to determine the algorithm that the model has learned to solve the search task, and to measure the extent to which the model uses such an algorithm or relies on heuristics. We develop a mechanistic interpretability analysis to study the learned algorithm. We find that transformers perform search simultaneously on all vertices of the input graph, where for each vertex, the transformer stores the set of vertices reachable within certain number of steps. Each layer successively extends the sets of reachable vertices, thereby allowing the model to search over number of vertices exponential in the number of layers.
However, we find that as the input graph size increases, the transformer has increasing difficulty in learning the training distribution. We find that increasing model scale does not alleviate this difficulty, suggesting that simply increasing the size of the transformer will not lead to the robust acquisition of searching and planning abilities.
We also consider modified versions of the search task where the model is permitted to output inter- mediate tokens (akin chain-of of-thought prompting: Kojima et al., 2022; Wei et al., 2022b). More specifically, we test depth-firs first search and (zero-shot) selection -eneetiinnifeeence prompting (Creswell et al., 2023). We find that while it is easier to teach the model solve this task, requiring constant num- ber of layers, the model still struggles on larger input graphs.
Thus, our results suggest that future transformer- based models will not solve the search task with standard training, and alternative training approaches may be necessary."
RELATED WORK
Search abilities of transformers. A number of studies have explored the search capabilities of transformers and LLMs. Benchmarks have shown that LLMs can perform some graph reasoning tasks and related tasks such as repeated function composition and the k-hop induction head task (Fan et al., 2024; Fu et al., 2024; Sanford et al., 2024a;b), but they are limited to small graphs, relative to graphs we consider. Ruoss et al. (2024); Gandhi et al. (2024); Shah et al. (2024) find that a transformer can learn to approximate or simulate search algorithm, but with performance gap, and they do not test whether this gap is lessened by increasing model scale or training. Wang et al. (2023); Bachmann & Nagarajan (2024) show that LLMs can do some graph reasoning but are fooled by spurious correlations. Similarly Zhang et al. (2023) find that transformers are unable to learn to perform proof search since they strongly prefer heuristics. In this work, we also show that transformers are highly sensitive to the training distribution, but if extra care is taken to re- move heuristics, they are able to learn to search. Zhang et al. (2024) show that LLMs struggle on real-world graph reasoning tasks. Borazjanizadeh et al. (2024) find that LLMs have difficulty on diverse search problems. However, this is in contrast with work the theoretical expressiveness of transformers. Merrill & Sabharwal (2024) show that with thain-of-thought, transformers can sim- ulate any Turing machine. However, their results do not indicate whether it is possible to train transformer to perform any task. In fact, we find that even if the transformer is permitted to generate intermediate tokens, it is challenging to learn to search on larger graphs. Sanford et al. (2024a:b) show that transformers need a logarithmic number of layers to perform the graph connectivity task, which is supported by our identification of the algorithm that transformers acquire during training. Mechanistic interpretability. There is a large amount of work that seek an algorithmic understand- ing of transformers trained on various tasks. Hou et al. (2023); Kim et al. (2024) look for evidence of specific circuiss'aloorthhm in the transformer's 's activations and attention patterns. In our approach, we do not require the algorithm be known priori. Rather, we reconstruct the computation graph from the model activations and attention patterns. We make heavy use of activation patching (Vig et al., 2020b; Geiger et al., 2021; Heimersheim & Nanda, 2024). Brinkmann et al. (2024); Kim et al. (2024); Stolfo et al. (2023) apply mechanistic interpretability analysis to better understand transformer behavior in reasoning, and Yang et al. (2024); Jenner et al. (2024) present evidence that LLMS perform shallow searches during the forward pass. Ivanitskiy et al. (2023) train small transformer models on a maze search task and find internal representations of the maze map.
supprooellly-upporteed Scaling laws. Scaling laws empirically supported hypotheses about the long-term behavior of machine learning models on a task, as a function of the model size, the amount of data, and compute (Kaplan et al., 2020; Henighan et al., 2020; Hoffmann al., 2022). Caballero al. (2023) explore scaling laws on a wide multitude of tasks, but not including search, reasoning, or planning. Wei
All cnde for generating data, training and evaluation is open-source and freely available gi titubuu. - com/ aawaaparv//eerriniai to search.
2

---

Published as a conference paper at ICLR 2025
start vertex goal vertex
Graph search example:
Model input:
graph edges
Start
8
goal
Label: 3
Equivalent proof search example:
If Alice has fur, then Alice sheds. If Alice is mammal, Alice is warm-blooded, If Alice is warm-blooded then Alice shiver. IF Alice mammal, then Alice has fur. Alice bird, Alice warm-blooded. Given Alice is mammal, prove Alice can shiver.
Label: Alice is warm-blooded.
FIGURE I: (top left) Example of a search example on directed acyclic graph and (top right) the corre- sponding transformer input and output label. (bottom) An equivalent proof search problem in implicational propositional logic, rendered in natural language.
et al. (2022a) posit that there may exist phase transitions scale increases, where certain abilities "emerge." " However, Schaeffer et al. (2023) find that this emergent behavior can be explained by careful selection of metrics. Du al. (2024) argue that training loss is better measure of model ability, independent of scale.
3 SEARCH IN DIRECTED ACYCLIC GRAPHS
In order to test whether transformers can learn to perform search, need way produce large number of search problems on which the model can be trained and evaluated. We do so by generating search problems directed acyclic graphs (DAGs). Each example consists of DAG, a start vertex, and goal vertex (which is guaranteed to be reachable from the start vertex). The model's task is to output the next vertex the path from the start to the goal vertex. We characterize the difficulty of an example by its lookahead: the number of steps that the model must search from the start vertex to find the goal. More precisely, for any example graph, let P be the path from the start to the goal vertex, and S; be the paths from the start vertex that are otherwise disjoint with P. then the lookahead L is min max, [S.]]. An example graph, along with the corresponding transformer input, is shown in Figure 1. In this graph, the lookahead is 2.
We experiment with three distributions of DAG search problems: (1) two simple distributions where pay no special attention to heuristics, which we call the "naïve" and "star" distributions, and (2) a more carefully constructed distribution where we take care to prevent the model from exploiting heuristics to solve the task, called the "balanced distribution."
random graph distribution (Erdös & Rényi, 1959) to generate directed acyclic graphs: We arrange a set of vertices in linear order from left to right (i.e., topological order) and randomly sample edges between pairs of vertices. To ensure there are cycles, all edges are oriented to the right. To reduce the density of the graphs, we limit the in-degree of any vertex 4, since the lookahead is typically very small in dense graphs. See Section A. for details on the generative process.
Naïve distribution. We adapt the Erdös-
Balanced distribution. While easy to describe and implement, the naïve distribution strongly tends generate graphs where the lookahead is very small (typically L = 2: see Figure 8). fact, it becomes exponentially less likely to generate examples with greater lookaheads. In order to both train and evaluate the model's ability search multiple steps find the goal vertex, we need an efficient way to generate examples where the model is required to search for additional steps in order to find the goal. In addition, to prevent the model from relying on heuristies to perform search, we must take care to ensure that heuristics are not useful to solve the search problems in the training distribution. Thus, we design the balanced distribution to specifically generate graphs with lookahead parameter T. (see Section A. for details), which we use to produce training data where the lookahead is uniformly distributed.
Star distribution. We experiment with an additional distribution over graphs where the vertices are arranged in shar-shape (see Fig 1 in Bachmann & Nagarajan, 2024). The vertex at the center is the start vertex, and there are "spokes" that radiate outwards from the center, where each spoke is a linear chain of I vertices. The goal vertex is the end of one of the spokes.
3

---

Published as a conference paper at ICLR 2025
Naiive distr.
0.08
0.81
11.09
1.00
0.90
11.955
0.22
00.33
00.51
01.96
00.33
01445
01:40
(0
0::53
01:12
00.00
0.552
0..2
0.49
Trained on:
Star distr:
0.63
1:00
0.98
0.59
1.00
0.99
100
0.98
0.59
0.599
0.97
1.00
0.996
0.
0
0.99
0.995
0.90
0.52
1.00
Balanced distrs
1.00
1.00
1.00
1.00
0:97
Balanced distr:
0.96
1.00
1.00
1.00
1.00
1.00
0.96
0.50
0.49
(lonkahoud x 12)
Naivo distr
Star dists.
15
1
19
hp
Tested on:
Balanced distribution with lookshead
FIGURE 2: Accuracy of model with a maximum input graph size of 41 vertices trained on 883M examples from the naive distribution, vs star distribution, vs the balanced distribution with lookaheads L S 20 (which is the maximum for the input size), vs the balanced distribution with L < 12 for the last row. All evaluation is on held-out examples.
3.1 EXPERIMENTS
We train transformer models, with the same architecture as GPT-2 (Radford et al., 2019) with ReLU activation. In order to facilitate mechanistic interpretation of the trained model behavior, we use h-hot embeddings for both the token and the absolute positional embedding. Furthermore, the token embedding and positional embeddings are concatenated rather than added to form the transformer input. We use attention head per layer. The feed-forward dimension is the same as the model dimension. Since the edges in our inputs are randomly ordered, it would help for each token to be able to attend to any other token, rather than only the preceding tokens. As such, the model does not use causal mask when computing attention. We train an 6-layer model with hidden dimension 16.
To simulate training on samples of a very large corpus of data, we utilize streaming training. We continually sample batches from the generative process throughout training, instead of sampling batches from a fixed training set. The first few batches are reserved for testing, and subsequent batches are filtered via exact matching to remove any examples that appear in the test set, to ensure that the examples in the test set are unseen. In all our experiments, we use the Sophia optimization algorithm (Liu et al., 2024) with learning rate of 10-5, weight decay of 0.11, and no dropout. We use a batch size of 1024 examples, unless otherwise stated. We minimize the cross-entropy loss during training. Some graphs contain multiple paths from the start the goal vertex. During training, we select one path uniformly at random as the ground truth when computing the loss. During evaluation, we consider the model's ' prediction to be correct if the output vertex lies on any path to the goal.
3.1.1 SENSITIVITY TO TRAINING DISTRIBUTION
We first investigate the effect of the training distribution on the transformer's ability learn the search task. We do so by training one model on the naive distribution, another model on the star distribution (where the number of spokes it and spoke lengths L are uniformly distributed), and a third model on the balanced distribution (where the lookahead L is uniformly distributed), all with input size 128 tokens. Then, we evaluate the accuracy of each model on held-out test sets from the naïve, star, and balanced distributions, for all possible lookaheads L. We observe in Figure 2 that the model trained on the naive distribution is not able to robustly handle graphs with larger lookaheads, showing low accuracy even for observed number of lookaheads (e.g., L = 5). This is due to the fact that the probability of generating graphs with larger lookaheads with the naive distribution decreases exponentially, and so the model is not shown sufficient examples with large lookaheads. The model trained on the star distribution performs reasonably on examples from the balanced distribution but not well on examples from the naive distribution. In contrast, the model trained on the full balanced distribution performs near perfeetly in all test settings, and generalizes to unobserved numbers of lookaheads. This result demonstrates that it is possible to teach transformer to perform search almost perfectly on the space of lookaheads observed during training, when provided with an appropriate training distribution. Furthermore, the model trained on the balanced distribution with lookahead < 12 was able to generalize to lookaheads 13 and 14 but not on any larger lookaheads.
3.1.2 PROOF SEARCH IN NATURAL LANGUAGE
Our earlier experiments were conducted with a symbolic input representation for the graphs in each example. To demonstrate that our findings generalize to inputs expressed in natural language, we re-run our earlier experiment where we train small transformers from scratch, but with a modified input: Each edge (e.g., "vertex connects vertex 2"") is expressed "If Alex is a wumpus, then Alex is a vumpus") where each word and punctuation symbol is a token. The task is correspondingly modified: Given start proposition (e.g., "Alex is a wumpus"), "), find the next step in the proof of a goal proposition. Thus, this modification defines a one-to-one mapping between graph search problems in the symbolic representation and reasoning problems in natural
conditional sentence (e.g.,
4

---

Published as a conference paper at ICLR 2025
Step 1: Identify attention operations in the last layer with greatest contribution to the prediction.
4
1
3
E
3
8
4
3
Q
8
8
Layer
Layer
Layer
Layer
Oulput: 3
Step 2: Identify atlention operations whose perturbations strongly affect the operations from Step the output prediction.
E
3
8
8
Layer
Layer
Layer
Layer
Output: 3
Step 3: Use activation patching to explain each operation from Step 2 and select operations that on an explainable path from the inputs to oulput.
Layer
Layer
Layer
Layer
Output: 3
FIGURE 3: Overview of method reconstruct the computation graph from transformer for specific input.
language representation (in implicational propositional logic). See Figure 1 for an example of this correspondence. We see in Figure 9 that the training behavior in this setting is qualitatively very similar that in the graph search setting. The model has increasing difficulty learning the task as the graph size increases, especially in terms of FLOPs, as evident from the last row of the Figure.
MECHANISTIC INTERPRETATION OF TRANSFORMERS ON GRAPH SEARCH
We observed in the previous section that transformers are sometimes able to learn to search during training, and the resulting model is able to robustly and correctly answer almost any graph search problem in the input space. We aim to better understand the algorithm that the model acquired during training to solve the task, to determine whether and to what extent the model is utilizing a correct algorithm to solve the task, opposed to a heuristic.
4.1 RECONSTRUCTING ALGORITHMS FROM INPUTS
In order to understand how the transformer learns to solve the graph search task, we develop a new method for mechanistically interpreting the model's behavior. Our method involves closely examining the model's behavior for given input example in order to reconstruct computation graph that explains the model's activations, attention patterns, and output prediction. Our method consists of the following steps, as depicted visually in Figure
Compute activations, attention weights, and output logits. For given input example, perform ordinary forward pass to compute the activations, attention weights, and output logits.
II Identify important attention operations in the last layer. For each element of the attention matrix in the last layer corresponding to the last token, perturb the weight by changing it to 0 and recomputing the logit of the model's original prediction. If the resulting decrease the logit of the prediction is greater than threshold parameter a, then this attention operation is marked as important. Similarly perturb each weight by changing it large value² (the largest attention weight in the row of the attention matrix and renormalize). If the resulting decrease in the logit is greater than a, then this operation also marked as important."
We this since we found that at some layers, token will attend strongly every other token except one token, and information actually transferred from exceptional token due to the small attention weight. "Sometimes, this step yields no important operations. In such cases, fallback, the set of attention operations that the largest decrease in the logit of the model's prediction marked important.
5

---

Published as a conference paper at ICLR 2025
Set raachable van- tloss encodied emm bedding for veriex
Set reschable var- Nosa mendted in em- beadingg for vertox
Layer input:
start
Goal
2
Layer input:
Start
goal
Layer input:
Set of reachable vertices encoded in embedding for vertex 3
start
Goal
Layer input:
Set of reachable vertices encoded in embedding for vertex 5
start
Goal
Layer input:
Start
GOal
FIGURE 4: Visualization of the exponential path-merging algorithm, showcasing the layer-by-layer compu- tation of the reachability of vertex 9 from vertex 1. We hypothesize that transformers learn this algorithm search. this algorithm, each token corresponding to vertex stores information about which other vertices are reachable from this vertex (or from which vertices is this vertex reachable). For example, in layer 3, the model knows that vertex 3 is reachable from 1, and that 5 is reachable from 3, and computes that 5 is reachable from 1, shown in the input to layer 4. We posit the model performs this computation for all vertices simultancously.
III. Identify important attention operations in all other layers. For each element of the attention matrix all layers except the last layer, perturb the weight by changing it to 0 or to large value (i.e., setting the weight equal to the largest weight in that and renormalizing). Perform forward pass and inspect the log attention weight each important operation in the last layer. the resulting change in the log attention weight is greater than K2
v'd where d is the model dimension
and Ki is sensitivity parameter, then we mark this attention operation as important. Similarly, the resulting decrease in the logit of the output prediction is larger than o, mark this attention operation as important.
IV. Explain each important attention operation. For each important attention operation, let be the row and i be the column of the corresponding entry in the attention matrix, and be the layer of the operation. We say token i is the source token of this attention operation and token is the target. We use activation patching (Vig et al., 2020a; Zhang & Nanda, 2024) to determine which features of the input are causally significant for this operation. We test perturbations two types input features:
Token perturbations: For each vertex ID v of the input, produce perturbed input where is substituted for v, where v' is a vertex ID that does not appear in the original input. Position perturbations: For each position of the input, produce a perturbed input where
the positional embedding for the 7th token is set to zero.
Suppose we perturb input feature f. We then compute the forward pass the perturbed input while freezing the attention matrices up to the layer of the current attention operation." 4 At layer I, we compute the dot products:
Q,K and Q,K,T where Q is the perturbed query vector corresponding token j. Q is the unperturbed query vector, K, is the perturbed key vector corresponding to token i, and K; is the unperturbed key vector. We compare Q, KT to the original scaled dot product Q If the resulting change in the dot product is greater than vil where K2 is sensitivity parameter, then we say that the
(1)
"We also freeze the ReLU activations in that any value that set zero by ReLU in the original forward pass will also be to the effect the perturbation on the current layer in isolation changes behavior preceding layers.
the perturbed forward pass. The aim of freezing the previous layers is to measure
'e also require the change in dot product to be in the correct direction: If this attention operation has large attention weight, then we require the perturhed dot product to be smaller than the threshold, and vice versa for attention operations with small attention weight.

---

Published as a conference paper at ICLR 2025
embedding of the token at index j contains information about the perturbed feature f. and that this information is used by attention layer to perform this attention operation. Repeat this for all input features and obtain the set of features fi".... ..... f of the target embedding that causally affect this attention operation. Similarly, we compare Y,K, to determine whether information about the perturbed feature f is encoded in the embedding of the token at index i and is significant for this attention operation. Repeat this for all input features and we obtain the set of features fe ff.....ff of the source embedding that causally affect this attention operation. The result of this step is description of each important attention operation, showing why the operation happens. That is, what are the input features (token and positional embeddings) that are encoded in the source and target embeddings that causally affect the attention weight corresponding to this operation.
V. Reconstruct the computation graph/cireuit. Starting from the first layer, let & be the token at position of the input. We say each input vector "explainably contains" information about the token value &g and position k. Next, we consider the attention operations in the first layer. Suppose an attention operation copies source token into target token j. and depends on the source token embedding containing features S ..... ................................ lu and depends on the target token embedding containing
features "".....'"" to perform this operation (as computed in Step IV.). We say this attention op- eration is explainable if the embedding of token explainably contains all features / s and the embedding of token j explainably contains all features .f.....""" f
If the attention operation is explainable, we say the output embedding of the target token j explainably contains the union of the features: fi .......ff fu fi ..... f. We repeat this for every layer, computing all explainable attention operations throughout the model. Pscudocode for this procedure is shown in Algorithm 1. Finally, we filter out attention operations for which there does not exist a path of explainable attention operations to the output prediction (i.e., we can't explain how this operation is useful for the model's 's output on this example). The result is computation tree, where each node corresponds to an embedding vector in some layer in the model, which explainably contains information about a set of input features, and where each edge corresponds to an explainable attention operation.
We apply the above method on a trained model repeatedly for different input examples. The result is set of computation graphs/circuiss, one for each input example, and we can perform further analysis on these circuits to describe the model's 's computation across many inputs. While this method is able to produce fine-grained description of the processing in the transformer, it requires many forward passes Nonetheless, we are able to apply it to our smaller trained models.
4.2 EXPERIMENTS
We perform the above analysis on models trained on the balanced distribution that have achieved near- rear-perfect test accuracy. We hypothesize that the transformer performs search on all vertices simultaneously. where the embedding for each vertex explainably contains information about the set of vertices reachable from the current vertex within a certain number of steps. At each layer, for each vertex, the attention mechanism copies from a source vertex that is at the edge of the current vertex's 's reachable set, computing the union of the reachable sets of both vertices and storing the resulting set in the embedding of the current vertex. Thus, the reachable set can theoretically double in size every layer. In theory, the model may perform the search either in the forward or backwards direction: Rather than storing the set of reachable vertices, it may store the set of vertices from which the current vertex is reachable. A visual depiction of this algorithm is shown in Figure 4.
To test whether the transformer utilizes this algorithm, we perform the analysis described in Sec- tion 4.1 for multiple held-out inputs from both the naive and balanced distributions (on a total of 2000 inputs; 100 for each lookahead). We set = 0.4, Ki 20, and K2 = 10. For each input, we reconstruct and inspect the computation graph of attention operations. We categorize each attention operation into one of the following: (1) path-merge operations,? or (2) copy operations from ver- tices that are specifically reachable from either the start or the goal vertex. If the attention operation does not fall into either category, it is discarded. We say the input is explained by the aeth-merging algorithm if for every vertex along the path from the start to the goal vertex, there exists an unbroken sequence of path-merge operations that ultimately copy from the corresponding token in the first layer into the last token at the last layer.
" "Ln"mt mF' where L is the number of layers, n is the model's input size, is the number of input examples, and the number of perturbed features.
We check for either "token-man matching" or position- position-based path path-merge operations. In loken-matching the attention head selects another token by looking for vertex IDs of overlapping sets of reachable vertices. In mosition-hasedd op, the attention head looks for vertices that one step from vertex in the reachable set.
7

---

Published as a conference paper at ICLR 2025
Proportion of examples explained by path-merging algorithm
Trained on:
Naïve dista.
0.99
1.00
1:00
0.30
0.57
0.55
0.57
0:7h
0.77
0.39
0.52
0.55
0.38
0.663
00:44
0.43
00:45
00:55
00.00
0.07
Balanced distr.
u59
0.99
0.94
0.94
0.99
0.388
0.98
0.19
0:96
0.85
0.99
0.99
0.99
0:08
0106
0.113
0..04
0.994
0,00
0.02
Balanced distr.
(lookahead & 12)
100
0.76
0.990
0.98
0:97
0.97
0.90
0.05
0.012
0.87
0.64
Random model
10
Tested naive dietr.
Tested on balanced distributtio with lookahend
Proportion of "maximal" path merge operations
Trained on:
Naive distr.
0.52
0.18
0.:33
61:49
0::47
0:23
00:25
0.119
0.225
01:15
00:23
0:28
0:32
0:.22
0.11
00.26
0.22
0.20
0.34
0:32
Bolanced diste.
0.72
0.75
0.009
0:33
0.75
0.77
0.77
0.77
0.80
0:81
0:81
0.82
0.54
0.84
0.85
0.80
Balanced distr. (lonkahead 12)
Tosted on
%
%
16
11
nxive distr.
Tested on halanced distribation with loolahead
FIGURE 5: (top) The proportion examples which path merging algorithm identified the com- putation graph, as reconstructed using our mechanistic interpretability analysis. Each cell contains a random held-out sample of 100 examples. We perform our analysis on the same models as in Section 3.1.1 (and Fig- ure 2). A randomly-intialized (untrained) model is shown in the last row as the control. (bottom) The propor- tion of wath-merge operations that "maximal," averaged over 100 random examples. We say path-merge operation is maximal if it merging largest available reachable sets. This contrast with a suboptimal path-merge operation where one or both reachable sets are not the largest available at that layer.
Fraction of converged seeds
1.0
Minimmn loss
0.8
10
10-3
0.2
0.0
10
50
40
Maximum inpul, graph size
Maximum input graph size
FIGURE 6: (left) The minimum test loss after training on 236M examples (in vertices). For each maximum input graph size, we run 14 experiments with different random seeds. models were trained on the balanced distribution. Test loss was evaluated on held-out examples from the maive distribution. We only show results for models that have converged (i.e., it's training accuracy is greater than 0.995). (right) The fraction of models (initialized with different random seeds) that converged maximum input graph size after training on 236M examples.
maximum input graph size
The proportion of examples for which our method identifies the path-merging algorithm is shown in the top part of Figure 5. We observe that our method highly specific, identifying the algorithm in the trained models but not in the random (untrained) model. Our analysis provides more fine- grained view of the transformer's computation, and we are able to count individual path-merge operations and inspect whether they are "maximal" in the sense that they are merging the largest reachable sets that are available at that layer. For example, a path-merge operation that copies the embedding of vertex 2 into that of vertex 1 would be maximal if the embedding for 1 contains the reachable set {1, 2} and the embedding for contains the set {2,3). whereas the operation would be maximal if the embedding for 2 has the set {2) (see the example in Layer 2 of Figure 4). We observe the bottom portion Figure 5 that the proportion of path-merge operations that are maximal is notably less than indicating that the model does not learn maximal path-merge operations. And the model has no reason to do so since it is trained on lookaheads that do not align with a power of and it has more layers than it needs to learn search on the graphs with the largest lookaheads that can fit in its input. This explains why, in Figure 2, the model trained on lookaheads S 12 is not able to fully generalize to larger lookaheads.
8

---

Published as a conference paper at ICLR 2025
100
Training loss
10a
Test loss
10
Non-cthhedcing
Non-cmbeedding
10-2
10
10
10
10
10
10"
10a
Training examples
Training examples
FIGURE 7: Training and test loss number of training examples seen, for models with varying numbers of non-embedding parameters. All models were trained the balanced distribution. Test loss was evaluated on held-out examples from the naive distribution. Test loss is smoothed by averaging over a window 81 data points, where each data point is recorded every 2 262K examples. 4 seeds are shown for each model size.
DOES SCALING HELP?
Even it is possible to train a transformer to search, it is unclear how this ability scales with respect to input graph size and model scale. We investigate scaling behavior by running two sets of experi- ments on small transformer models: (1) training models on increasing graph sizes, and (2) training models with increasing model dimension d. We observe large variance in performance across differ- ent initial random seeds, which has been observed in other tasks (Kim & Linzen, 2020; Zhou et al., 2024). Thus, in each of these experiments, from the balanced distribution, and measure the minimum loss on a held-out test set generated by the naive distribution. We set the batch size 256 examples due to GPU memory limitations.
train the models using multiple seeds examples
In Figure 6, we observe that, when the number of layers is fixed to 8 and the hidden size 16, as the maximum input graph size increased, the likelihood that the model learns the training distribution (i.e., reaches accuracy 2 0.995) becomes vanishingly small. In addition, as the maximum input graph size increases, the minimum test loss over 14 seeds grows at an increasing rate. See Figure 10 in the Appendix for more detailed visualization of the training dynamics versus input graph size.
To determine whether larger models can more easily learn to search on large input graphs, we fix the input graph size to 31 and train models of widely varying sizes. In Figure 7, we see that while larger models are able to more quickly find the local minimum (at loss near 10%), there is no discernible pattern between the size of the model and the amount of training needed to find the global minimum.
We also experiment with decoder-only models with learned token embeddings and rotary positional embeddings (Su et al., 2024), which are multiplied rather than concatenated, and are more predom- inant in contemporary LLMs. But we find that this does not change the model's scaling behavior (see Section A.5).
6 DOES IN-CONTEXT EXPLORATION (I.E., CHAIN-OF-THOUGHT HELP?
Though we have shown that transformers are not able to learn to perform search for larger input graphs, they may be able learn to search if permitted take intermediate steps, akin chain-of- thought prompting. To test this, we repeat our earlier experiments with two "prompting" approaches: (1) depth-first search, and (2) delection-ifferenc.
6.1 DEPTI-FIRST SEARCH
We generate graphs and perform a depth-first search (DFS) from a randomly-selected vertex random goal vertex. From the sequence of visited vertices (i.e., DFS trace), we randomly select a "current" vertex. Each input the model is: (1) the graph, as list of edges, and (2) the sequence of visited vertices up to and including the current vertex. The model's 's task is to predict the next vertex in the DFS trace.
In our earlier search task, the edges of the graph always appeared in the same token position across inputs, since the current and start vertices were identical. However, in this modified task, the se-
Agaii note that there may be multiple correct predictions, since there are typically many correct DFS traces for given graph. Similar to the setup in Section 3.1. we randomly select one of them to be the ground truth label when computing the cross-entropy loss during training. But during evaluation, allow the model to make any prediction that follows valid DFS trace.

---

Published as a conference paper at ICLR 2025
quence of visited vertices can vary in length. In order to avoid complicating the task for the model, we add padding to the input, between the graph and the sequence of visited vertices: The right-most most tokens are reserved for the sequence of visited vertices. All other tokens contain the graph edges, just as in the original search task. See the example in Figure 13. The graphs in the training distribu- tion are sampled such that the backtrack distance is uniformly distributed, analogous to the balanced distribution in our earlier experiments (see Section A.7 for details on the graph distribution).
In the experiments, we first fix the model size, setting the number of layers to 3. and vary the input graph size. As evident in the top row of Figure 14, the model is able to learn the training distribution across all tested graph sizes, suggesting that only a constant number of layers needed to learn the DFS task. However, the model struggles as the graph size increases. Next, we fix the maximum graph size to 35 vertices and instead vary the model size. In the middle and bottom rows of Figure 14, we observe that increasing model scale does not help the transformer to learn this task more easily. Interestingly, we also find that larger models are able to learn the task from fewer training examples. However, the benefit from scale disappears when considering the cost of training larger models: Larger models require many more FLOPSs to learn the task than the smaller models.
6.2 SELECTION-INFERENT
In this setting, each search step is decomposed into two subtasks: (1) given a graph and a list of visited vertices, select a visited vertex that has an unvisited child vertex, and (2) given a selected vertex, predict (i.e., infer) an unvisited child vertex. Starting from the start vertex, if these two subtasks are repeated sufficiently many times, the goal vertex will be found. To construct a graph distribution that is analogous to the balanced distribution, we define two variables: (1) the frontier size F, which is the number of visited vertices that have unvisited children, and (2) the branch count B, which is the number of child vertices of the current vertex (in an inference step). We generate graphs such that the pair (F, B) is uniformly distributed. Each input consists of: (1) the graph, and (2) the sequence of visited edges. 9 See Section A.10 for further details.
In the experiments, we first fix the model size, setting the number of layers to 4, and vary the input graph size. We again note from the top row of Figure 16 that the model struggles as the graph size increases. Next, we fix the maximum graph size to 45 vertices and vary the model size. We note in Figure 17 that increasing model scale does not help to learn the task. Therefore, transformers struggle to learn to perform DFS search and selection-inferee inference larger graphs and additional scaling does not seem to make it casier.
7 CONCLUSION
Through the use of graph connectivity testbed, we found that transformers can learn to search when given the right training distribution. We developed and applied a new mechanistic inter- pretability technique on the trained model to determine the algorithm that the model learned to perform search. We found that the model uses an exponential nuthmeegiigg algorithm, where the embedding of each vertex stores information about the set of reachable vertices from that vertex. As the input graph size increases, the transformer has verrinnceasing difficulty in learning the task, and increasing the scale of the model does not alleviate this difficulty. Lastly, even if the model is permitted to use intermediate steps, they still struggle on larger graphs, regardless of scale.
It is possible that scaling to much larger model sizes may lead to emergent searching ability. Alter- nate training procedures may help transformers to more easily Iearn to search, such as curriculum learning. Alternate architectures may help as well, such as looped transformers. While the path- merging algorithm is able to explain almost all examples for the trained models, there may be other algorithms or heuristics that the model simultaneously utilizes on some examples. Our mechanistic analysis has potential broader applications in reasoning: Some form of the path-merging algorithm may be used by transformers in more general reasoning tasks. In such a case, the representation of each fact would store information about the set of facts provable from the current fact. Our mecha- nistic interpretability tools may be useful other settings, as well, where they may help to uncover the algorithms that transformers learn to solve other tasks. Though additional work is welcome to improve the scalability of our analysis to larger models, our analysis can provide insights on smaller models that can be tested separately in larger models.
Similar to the DFS task, we add padding to ensure the graph edges appear in the same positions across examples.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Emmanouil Antonios Platanios, Jannik Brinkmann, Dan Friedman, Arvid Frydenlund, Gonzalo Conzalez-Pumariega. and Will Merrill for their invaluable insight and discussion. This work was supported by Open Philanthropy, and in part through the NYU IT High Performance Computing resources, services, and staff expertise, and the Rosen Center for Advanced Computing at Purdue University.
REFERENCES
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: : /pwpneeviewww..e/////p/pp./ooewe wttt:/toooo...m//2222220
Béla Bollobás, Christian Borgs, Jennifer T. Chayes, and Oliver Riordan. Directed scale-free graphs. In Proceedings of the Fourleenth Annual ACM-SIAM Symposium Discrete Algorithms, January 12-14, 2003, Baltimore, Maryland, USA, pp. 132-139. ACM/SIAM, 2003. URL http: /dl. 644108 41.8..8... 644133
acm. org/ cewaationnccmiiim
Nasim Borazjanizadeh. Roei Herzig, Trevor Darrell, Rogério Feris, and Leonid Karlinsky. Nav- igating the labyrinth: Evaluating and enhancing Ilms' ability to reason about search prob- lems. ArXiv, abs/2406. 12172,
2024. URL https:/ /api
orgg/
CorpusID: 27005s
Jannik Brinkmann, Abhay Sheshadri, Victor Levoso, Paul Swoboda. and Christian Bartelt. A mech- anistic analysis transformer trained on a symbolic multi-step reasoning task. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Lin- guistics, ACI. 2024, Bangkok, Thailand and virtual meeting, August 11-16, 2024, pp. 4082-4102. Association for Computational Linguistics, 2024. doi: 10. 10.18653/V1/000 1/201/..1000.022... ACC.2222. URL httph::tp///......t. /doi org 111.1...1..11111 186533vv12 /1///221222222///////22 /2224 Findings-acll .242.
Ethan Caballero, Kshitij Gupta, Irina Rish, and David Krueger. Broken neural scaling laws. In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL /ttttts::///ppeeeeee http///ppppeeeeeee//e/ee net ht//foormm?// scegve1
Antonia Creswell, Murray Shanahan, and Irina Higgins. Selection-inferencee Exploiting large lan- guage models for interpretable logical reasoning. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. net wtw.:e/...eeeetttiiee.eeem /tt/:/orouirdPePPee--- 3ww A4.
URL https:://///m. /
Peng Ding, Jiading Fang, Peng Li, Kangrui Wang, Xiaochen Zhou, Mo Yu, Jing Li, Hongyuan Mei, and Matthew Walter. MANGO: benchmark for evaluating mapping and navigation abilities of large language models. In First Conference on Language Modeling, 2024. URL https:
/ //opnreevew.... net wtwt/foott//ieeeeeooeoo BBvormm/d..--
Zhengxiao Du, Aohan Zeng, Yuxiao Dong, and Jie Tang. Understanding emergent abilities of language models from the loss perspective. CoRR, abs/2403.15796, 2024. doi: 10.48550/ARXIV. 2403.15796 URL https:://doooo http://oo.oog//////// http:///...555//xx1...o.
Erdös and A Rényi. On random graphs i. Publicationes Mathematicue Debrecen, 6:290-297, 1959.
Lizhou Fan, Wenyue Hua, Lingyao Li, Haoyang Ling, and Yongfeng Zhang. Nphardeval: Dy- namic benchmark on reasoning ability of large language models via complexity classes. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meet- ing of the Association for Computational Linguistics (Volume 1: Long Pupers), ACL 2024, Bangkok, Thoiland, August 11-16, 2024, pp. 4092-4114 14. Association for Computational Linguis- URL https::/ddo.. www.:ii.org/1...11
tics, 2024. doi: 10. long 225. v1/2024 20c22244.c--200...
Deqing Fu, Ghazal Khalighinejad, Ollie Liu, Bhuwan Dhingra, Dani Yogatama, Robin Jia, and Willie Neiswanger. Isobench: Benchmarking multimodal foundation models isomorphic

---

Published as a conference paper at ICLR 2025
representations. CoRR, abs/2404.01266, 2024. doi: 10.48500ARRXIV2404.01166 URL attps://dbi.org//0 http:///....50////....4....4 01266.
Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, and Noah D. Goodman. Stream of search (sos): Learning to search in language. CoRR, abs/2404.03683, 2024. doi: 10.48550/ARXIV.24044.000 48550/arXiv.200 03683.
URL btttp:://di..rr/
Atticus Geiger, Hanson Lu, Thomas F Icard, and Christopher Potts. Causal abstractions of neural networks. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: : / /pwpneeeeeeeweeee///eeeee net forum? id=RmuXDt: jDhG.
Shibo Hao, Yi Gu, Haodi Ma, Joshua Jíahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu. Reasoning with language model is planning with world model. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pp. 8154-8173. Associ- ation for Computational Linguistics, 2023, doi: 10.18653/V B1A.:02...MLL---PP https: / htt:///oiiioat org/10 18 8633 http/l11/222// em..1. ma ia.iiin50..
URL
Stefan Heimersheim and Neel Nanda. How to use and interpret activation patching. ArXiv, abs/2404. 15255, 2024. URL https: wen.e.ssaaaiicchol....
Org/CorpusID:
26 269302704.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B. Brown, Prafulla Dhariwal, Scott Gray, Chris Hallacy, Benjamin Mann, Alec Rad- ford, Aditya Ramesh, Nick Ryder, Daniel M. Ziegler, John Schulman, Dario Amodei, and Sam McCandlish. Scaling laws for autoregressive generative modeling. CoRR, abs/2010.14701, 2020. /2cs/200.. 14701.
URL https:://// / www.i.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. CoRR, abs/2203. 15556, 2022. doi: 10.48550/ a 220322220... 15556.
ARXIV.2203. 155566. URL httph:ttt//:/
Yifan Hou, Jiaoda Li, Yu Fei, Alessandro Stolfo, Wangchunshu Zhou, Guangtao Zeng, Antoine Bosselut, and Mrinmaya Sachan. Towards a mechanistic interpretation of multi-step reasoning ca- pabilities of language models. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.). Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. EMNLP 2023, Singapore, December 6-10, 2023, pp. 2023. doi: 10. 1/66:3////0223EMMM.PPMPM v1 22022202..22 emn 1p-ma 2.99.
Association for Computational Linguistics, 299. URL. https:://dattt org
/ 1 18653/
Michael I. Ivanitskiy, Alex F. Spies, Tilman Räuker, Guillaume Corlouer, Chris Mathwin, Lucia Quirke, Can Rager, Rusheb Shah, Dan Valentine, Cecilia G. Diniz Behn, Katsumi Inoue, and Samy Wu Fung. Linearly structured world representations in maze-solving transformers. In Mareo Fumero, Emanuele Rodolà, Clémentine Dominé, Francesco Locatello, Karolina Dziugaite, and Mathilde Caron (eds.), Proceedings of UniReps: the First Workshop on Unifying Representa- tions in Neural Models, 15 December 2023, Ernest Morial Convention Center, New Orleans, USA, volume 243 of Proceedings of Machine Learning Research, pp. 133-143. PMLR, 2023. URL https: http:sss/ppoccedeieee http//ppoccees....//e//p mll www.r.press//ee v2///24444.. 4 www.iiisiiiiiii.. htmll.
Erik Jenner, Shreyas Kapur, Vasil Georgiev, Cameron Allen, Scott Emmons, and Stuart Russell. Evidence of learned look-ahead in a chess-playing neural network. CoRR, abs/2406.00877,, 2024. doi: 10.48550/ARXIV.2406000 00877.
URL https: /owii
485.0//55550......../ 2406.
Subbarao Kambhampati, Karthik Valmeekam, Lin Guan, Kaya Stechly, Mudit Verma, Siddhant Bhambri, Lucas Saldyt, and Anil Murthy. Lims can't plan, but can help planning in Ilm-modulo frameworks. CORR, abs/2402.01817, 2024. doi: 10-45500/ARXX122222.11111 URL https:
//doi. org/10.. 48550 2402. 01 817.
12

---

Published as a conference paper at ICLR 2025
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, abs/2001.08361, 2020. URL https: /ttp/vaaivvv/..a///i//eee/aaaa.. org/ abs 2001. 08361.
Mehran Kazemi, Najoung Kim, Deepti Bhatia, Xin Xu, and Deepak Ramachandran. LAMBADA: Backward chaining for automated reasoning in natural language. In Anna Rogers, Jordan Boyd- Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 6547-6568, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.1 10.8866///22222..--1ongggg.. https: a/ www.aaaonnoooooooooooooooooooooe wtgg2023. acl- 361.
URL.
Geonhee Kim, Marco Valentino, and André Freitas. A mechanistic interpretation of syllogistic reasoning in auto-regressive language models. ArXiv, 0bs/2408.08590 2024. URL https:
//api. www.sssssaticcccosseoeo wttp:a/rogg///ooooooeaeaa : 271892176.
Najoung Kim and Tal Linzen. COGS: compositional generalization challenge based on semantic interpretation. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu (eds.), Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pp. 9087-9105. Association for Computational Linguistics, 2020. doi: URL https:: /doi. http//o..oo.....//// 11t.:/...1111/111111 22000.
emnlp-main. 731.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwa- sawa. Large language models are zero-shot reasoners. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, cttppperrrppppppppppppee http:paper_fileeee tttp:/ppppee/2//e/eeee/ //ash//
2022, 2022. URL http: /ttp:://papprsa.eee httppppppppppp.a.p..... cecips.cc/peee 8bb0d291acd oee.eeeeeecccaacfeeeeeeeeeee 1eteee111999c11322 Abstract wow.eetttcoooeeeec.hehee html.
Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- view.net, 2024. URL https:
/ wwtp//eeeeeeeeeeeeeeeeee/ee
net wttp:oommttxxxxeeaaBka...
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Aus- tria, May 7-11, 2024. OpenReview.net, 2024. URL https: http///peeeeeeeeee httts://////peeeeiiiee id=N JNG1Ph.Wh.
net forum?
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Au- gustus Odena. Show your work: Scratchpads for intermediate computation with language models.
CoRR, abs/21 12.01122001114 14, 2021. URL https: /arxivivvvv. http://wwww...o.//////////// abt////s//111/ 2012.001111010 14.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Anian Ruoss, Grégoire Delétang, Sourabh Medapati, Jordi Grau-Moya, Li Kevin Wenliang, Elliot Catt, John Reid, Cannada A. Lewis, Joel Veness, and Tim Genewein. Amortized planning with large-scale transformers: A case study on chess. CoRR, Abs/24002.0449, 2024. doi: 10.48550/ ARXIV.2402.04494, URL https: doi (aci.oore 085111...508888 arxiv. 2402. 04494.
Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Seyed Mehran Kazemi, Jonathan Halcrow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. CoRR, abs/2405.1887 2024a. doi: 10-45550ARXXV22005.15555 URL
https: /oi. /10. 8 arxiv 2405. 18512.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and log- arithmic depth. In Forty-first International Conference Machine Learning, ICML 2024. Vi- enna, Austria, July 2 21-27, 2024. OpenReview.net, 2024b. URL htt.htttss//ppeerrrrree het:/forrmmiid-gc7aabe
13

---

Published as a conference paper at ICLR 2025
Abulhair Saparov and He He. Language models are greedy reasoners: A systematic formal analysis of chain-of-thought. CoRR, abs/2210.01240,, 2022. doi: 1044550/ARX1:210000240 URL https: http://d....rgg11.4555533
Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models mirage? In Alice Oh, Tristan Naumann, Amir Glober- son, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural In- formation Processing Systems 36: Annual Conference on Neural Information Pro- cessing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL
ffrfrsflllssse wttps//apppr////2 2222//222hast adc982664 -tt.aaeeeebettttttttoooooeeeeeee Connttttoofffeeeeeeeee html.
Kulin Shah, Nishanth Dikkala, Xin Wang, and Rina Panigrahy. Causal language modeling can elicit search and reasoning capabilities on logic puzzles. CoRR, abs/2409. 2024. doi: 10.48550/ ARXIV.2409.1050 URL https: http://................... wtt.g//.....5///2.......
Katharina Stein and Alexander Koller. Autoplanbench: Automatically generating benchmarks for Ilm plamners from pddl. arXiv preprint arXiv:2311.09834 2023.
Alessandro Stolfo, Yonatan Belinkov, and Mrinmaya Sachan. A mechanistic interpretation of arith- metic reasoning language models using causal mediation analysis. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.). Proceedings of the 2023 Conference on Empirical Methods in Natu- ral Language Processing, EMNLP 2023, Singapore, December 6- 10, 2023, pp. 7035-7055 035-7052. Asso- ciation for Computational Linguistics, 2023. doi: 10.16633N12023.MMNN https: /oo. www//////00..... 653 v1/.v22.222222/2 emn ponellp-maann 435.
435. URL
Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputingg 568:127063, 2024. doi: N0.0011. 2023. 127063.
URL htt.hhttt:///...ooo
neucom.
Karthik Valmeekam, Alberto Olmo, Sarath Sreedharan, and Subharao Kambhampati. Large lan- guage models still can't plan (a benchmark for LLMs planning and reasoning about change). In NeurIPS 2022 Foundation Models for Decision Making Workshop, 2022. URL https:
/ /popneeviiew net htt//fooummddddeee -OS----OO--XXXL0000
Jesse Vig, Scbastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart M. Shieber. Investigating gender bias in language models using causal media- tion analysis. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Bal- can, and Hssan-TTie Lin (eds.), Advances in Neural Information Processing Systems 33: An- nual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020a. URL https: /ttpe//ee///eeeeeeeeeeeoe nwwrips..i... cc/ paper/202C/ 2eee222277755ee112efffaaaaa00000 Abstract. html.
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart M. Shieber. Causal mediation analysis for interpreting neural nlp: The case of gen- der bias. ArXiv, abs/2004.12265,, 2020b. URL https: //ap:://ppppe. /ttp//pppaaatttte......../.e.a.a iescantcccccoaa....m
Co COrpusID:216553692
Heng Wang, Shangbin Feng. Tianxing He, Zhaoxuan Tan, Xiaochuang Han, and Yulia Tsvetkov. Can language models solve graph problems in natural language? In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: /wwp//eeeeee.eeeeeeeee wet.twwwwwweeeeeeeeeeee.ee.eeeee nwt:/fruut forum? id-UDqihhbqJJV
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language mod- els. Trans. Mach. Learn. Res., 2022, 2022a. URL https:/ /owpnreview...eeooo net tant/forunn forum? 1d=yzkSU5zddD.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoe V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language
14

---

Published as a conference paper at ICLR 2025
models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Ad- vances in Neural Information Processing Systems, 2022b. URL https: http////ppeeeeeeeee/eeeeeee
Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel. Do large language models latently perform multi-hop reasoning? In Anmual Meeting of the Association for Com- putational Linguistics, 2024. URL https: //api. semanti icschollscr org/ Corpusl:: 268032054
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances in Neural Information Processing Systems, 36, 2024.
Fred Zhang and Neel Nanda. Towards best practices of activation patching in language mod- els: Metrics and methods. In The Twelfth International Conference on Learning Representa- tions, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: //openreview.... net/ fort/ffoumm/ um?id=Hf] 7 6u9BC.
Honghua Zhang, Liunian Harold Li, Tao Meng, Kai- Wei Chang, and Guy Van den Broeck. On the paradox of learning to reason from data. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, IJCAI 2023, 19th-25th August 2023, Macao, SAR, China, pp. 3365-3373. ijeai.org. 2023. doi: 10.24966/IIAAL222233555 URL https: /do. .org/ 19.24993/ijcas att.:aa......255..... 2023/375.
Yizhuo Zhang, Heng Wang, Shangbin Feng, Zhaoxuan Tan, Xiaochuang Han, Tianxing He, and Yulia Tsvetkov. Can Ilm graph reasoning generalize beyond pattern memoriza- tion?
ArXiv, abs/2406.
2024. URL https:/ /atp://app... semanticscholarr .rrg/
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang. Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. CoRR, abs/2402.09371, 2024, doi: 10. WWW.AAXXXX.2402.031777 URL https: /woi.di../////////
A APPENDIX
A.1 GRAPH GENERATION DETAILS
A.1.1 NAïVE DISTRIBUTION
To sample graph from this distribution, we first sample the number of vertices
|V| Uniform((3. (3...., Vmnx)), }),
(2)
where VILA is the maximum number of vertices that can fit the input. Then, for each i = -,,,.... 1,..., [V.
we sample number of parent vertices n. parents from Vi...., Vin1
[1
with probability 5
parents
2,3, with probability 8.
Note that this differs from Pedös-Rényi where the number of parents is geometrically distributed. We want to avoid generating overly-dense graphs where the lookahead is too small, and so we choose to sample fewer parents per vertex.
Finally, we sample the parents of V; uniformly without replacement from (Vi,..., V until we have sampled ni parents vertices, or we have sampled all available vertices. We draw an edge from each sampled vertex to V,
After generating the graph, we randomly permute the vertex IDs so that the IDs contain no informa- tion about the graph topology. Observe that this distribution can generate any directed graph with maximum in-degree in topologically-sorted order, and therefore, it can generate any DAG with maximum in-degree
We select the start and goal vertices uniformly random from V. If there is no path from the start to the goal vertex, or if the example does not fit within the model input, we reject the sample and try again.
15

---

Published as a conference paper at ICLR 2025
Example frequency
10-1
20
10
10
10
15
20
Lookahead
FIGURE 8: Histogram lookaheads 10M graphs sampled from the naive distribution (note the log y-axis). The number vertices n 41. Lookaheads are possible, but astronomically unlikely.
A.1.2 BALANCED DISTRIBUTION
Each graph is sampled according to the following procedure: Given a lookahead parameter L, we start by creating linear chain of vertices containing L edges. The first vertex in this chain set as the start vertex, and the last vertex is the goal vertex. We then sample a number of additional "branches":
B ~ Uniform ((1,....,
Vin ()).
(4)
Next, we sample the total number of vertices:
V min (L(B + +1)+++++++0 1)+ u, Vmax):
(5)
where u ~ Uniform (O, 6}). We create B additional chains of vertices: For = 1,..., B, we create chain of vertices where the length of the chain in edges 1, is given by
L, if no additional vertices are available, i.c., Σj 4 L(B i- +1= [V. (6)
11
Uniifrr((L, L 1}}, otherwise.
Each additional branch is added to the start vertex.
While the graphs resulting from the above process will require the model to search at least L steps to find the goal, the graphs still admit heuristics since the start vertex is always the singular source vertex of the graph (i.e., has zero in-degree), and all other vertices have in-degree exactly equal to 1. To prevent such heuristics, we sample additional vertices as follows: We create an "incoming" linear chain with length صنر Uniform .... IV| - Ej= l + 1}. In contrast with the other branches, the start vertex is located at the end of this chain. Finally, to increase the degrees of the vertices in the graph, we create additional vertices until we have a total of [V] vertices. For each new vertex V,, sample a number of child and parent vertices:
children Uniform(0, 10,
(7)
parents
Uniform (n. = -.)...,3).
(8)
where 1xx} is the indicator function whose value is 1 if x is true, and otherwise. We sample children child vertices from Vi...., Vi- without replacement where the probability of sampling mi
Vú is proportional to 2/2 + deg (V,) where deg (v) is the in-degree of (at the time of sampling). Similarly, we sample ni parenis parent vertices from (V,..., .., V.-1) descendantss(v) without replace- ment where the probability of sampling V; proportional 1 + deg+(V,) where deg (v) is the out-degree of v. Note we avoid sampling from the descendants of V/ in order to avoid creating cycles. We chose this sampling scheme in order to produce handful of vertices with high in- or out-degree, and to prevent the model from exploiting heuristic when all vertices have low degree. This distribution an example of scale-free distribution over directed graphs (Bollobás et al., 2003).
As in the naive distribution, after generating the graph, we randomly permute the vertex IDs so that the IDs contain no information about the graph topology. If the resulting graph does not fit within the model input, we reject the sample and try again.
16

---

Published as a conference paper at ICLR 2025
A.2 NATURAL LANGUAGE PROOF SEARCH RESULTS
Figure 9 shows the results of our experiments on the natural language proof search task, as described in Section 3.1.2.
100
1.
Training loss
Test loss
4.x 10-1.
10-1
10-1-
Maximurn Input graph size
Maxtmum Inpur: graph size
19
x 10-1
106
107
10
10
107
10a
Training examples
Training examples
Minimum train loss
109
Minimum test loss
6 10-1
4 10-1.
10-1
3 10-1
Maximurn input graph
Maximum ingor: graphh
10"
2 10
108
107
10
10
107
10
Training examppes
Training examples
Minimum train loss
10
Mininm tost loss
6 10-1_
4xx 10-1
10-1
3 10
Maximurn graph
Maximum graph
10-2
2 10-1
10
10
10
10
101h
10t4
1
1016
Flonting-point operations
Floating-point operations
FIGURE 9: Training and test loss number of training examples seen, for models trained on the natural language proof search task, with varying maximum input graph sizes. All models were trained on balanced distribution. Test loss was evaluated on held-out examples from the naïve distribution. Test loss is smoothed by averaging over a window of 81 data points. where each data point is recorded at every 2 262K examples. In the top row, 5 seeds are shown for each maximum input graph size. In the middle and bottom row, the minimum loss over the seeds shown. In the bottom row, the x-axis rescaled FLOPs.
17

---

Published as a conference paper at ICLR 2025
A.3 ALGORITHM RECONSTRUCTION PSEUDOCODE
Algorithm Pscudocode of the procedure to compute the set of explainable attention operations of a
transforerr M for given input z.
function reconstruct. wtt.eccoopuuttttt. number of layers M
-n..graphh (input example x, transformer M)
initialize Mi, as an empty set for / E (O...., ,.....,) L and all i E {l...., [x|}
initialize & as an empty set for E ....../21 | do add token feature No. add position feature to No,
for 1 E 1,..., , do for i E ..... |2| de
for 1..... | do
11
e represent the attention operation at layer / that copies from source token to target token j
/* this is computed in Step IV. as described in Section 4.1 x/ fī. .f..... f
12
features in token on which e depends
let ff. ...... ff be the features in token j on which depends
if (f...... fa)) C Ni and .... N
/* mark this attention operation as explainable +/ add to E add [L .......... 1 (...... f. No.
15
Ni, now contains the set of features that are explainably contained in the embedding vector of token at layer ./ /. & contains the set of all explainable attention operations */ next, filter the explainable attention operations that not useful for the model's prediction */ lct be an empty stack push (L, onto S initialize 8" as empty set
17
18.
19
while S is empty do
21
pop (1,j) from for attention operation e E at layer whose target is token j do add to E*
22
29
let i be the source token of the attention operation push (I i) onto S
26
return N. g
A.4 ADDITIONAL SCALING RESULTS
Figure 10 provides a visualization of the training dynamics of the transformer when trained with varying maximum input graph sizes. Figure focuses on slice at 236M training examples.
A.5 SCALING DECODER-ONLY MODELS WITH ROTARY POSITIONAL EMBEDDINGS
We repeat the experiments in Section on decoder-only transformers with learned token embed- dings (initialized randomly from a Gaussian distribution) and rotary positional embeddings (RoPF). In our other experiments, the token embedding was concatenated with the positional embedding. In contrast, RoPE positional embeddings are multiplied with the token embedding in these exper- iments. For each maximum input graph size, we set dmotel be the same the corresponding experiment in Section 5 and A.4 for the same maximum input graph size. We observe in Figure 11 that decoder-only models with RoPE similarly struggle learn search on larger graphs. addi- tion, in Figure 12, we see that increasing the model scale does not help the model to learn the task more easily.
A.6 DFS EXAMPLE
Figure 13 shows DFS example, described for the task in Section 6.1.
18

---

Published as a conference paper at ICLR 2025
109
Training loss
109
Test loss
10
Maximum impat graph size
Maximum insut grash size
10-2
10
10
107
107
Training examples
Training examples
Minimum training loss
Minimum test loss
109
109
Maximam Iapat graph dizc
Maximum input gtoph siae
21
31
10-2
21
31
41
21
35
43
10-1
45
10
10
106
107
10
10
107
10a
Training examples
Training examples
FIGURE 10: Training and test loss vs number of training examples seen, for models with trained on varying maximum input graph sizes. All models were trained on the balanced distribution. Test loss was evaluated on held-out examples from the naive distribution. Test loss is smoothed averaging
window of 81 data points, where each data point recorded at every 2 262K examples. In the top row, seeds are shown for each maximum input graph size. the bottom row, the minimum loss over the seeds is shown.
Minimum train loss
Minimum test loss
109
10a
10
Maximum Iaput graph sizc
Maximum insut graph size
10
10
10-3
107
10
107
Training examples
Training examples
FIGURE 1I: Training and test loss number of training examples seen, for decoder-only transformers using rotary positional embeddings. Test loss evaluated held-out examples from the naive distribution. We fix the model size and vary the maximum input graph size. All models were trained on the balanced distribution. Test loss is smoothed by averaging over a window of 81 data points, where each data point is recorded at every 2 262K examples. For each point, we plot the minimum loss over 2 seeds.
A.7 GENERATING DFS EXAMPLES WITH SPECIFIC BACKTRACK DISTANCES
To generate graphs with a specific backtrack distance B, we first sample a DAG from the naive distribution. We then divide the coppllgicalll-sorteee graph into two subgraphs: the first V B vertices form the first subgraph G1 and the last B vertices form the second subgraph Gg. G₁ will contain the goal vertex, and G2 will contain the list of vertices visited so far (i.e., the DFS trace).
Let G1.-1 be the last vertex G1 (in the topological ordering) and let G2 be the first vertex in Gz. Next, select the start vertex If G1,-1 is the only parent vertex of G2,1, sample the start vertex uniformly at random from G1 \ (G1,-1)- If not, then we sample a start vertex uniformly at random from parents(Ga, goal). Then we sample the goal vertex uniformly at random from the set of vertices in G2 that come after s.
{G1,-1) (since need to leave at least vertex in G1 to be the
19

---

Published as a conference paper at ICLR 2025
Non- embedding
Minimum train loss
peramelerr
Miniman test loss
10
Nut-embohhding parsmelers
2.2 1M
0MM
PM
10t5
10
10
Floating-point operations
Floating-pond point operations
Training and test loss vs FLOPs, for decoder-only transformers using rotary positional embed- dings. Test loss computed held-out examples from the naïve distribution. We fix the maximum input graph size vertices and vary the model size. All models were trained loss is smoothed by averaging over a window of 81 data points, where each data point is recorded at every 2 262K examples. For each point, we plot the minimum loss over 2 seeds.
FIGURE 12:
balanced distribution. Test
start vertex goal vertex
Depth-first search example:
Model input (with standard padding):
DFS trace with standard padding
graph edges
Start
goal 6
DFS with random padding
Model input (with random padding):
Label:
FIGURE 13: (left) Example of a depth-first search example on a directed acyclic graph where the model has visited the vertices 8, 4, and 1 so far. (right) The corresponding transformer input and output label. We experiment with two padding methods: (1) standard padding where the DFS trace is left-padded,, and (2) random padding where padding is randomly inserted between vertices the trace.
We have to make sure there exists path from to every vertex in G that comes after s. Iterating the vertices in G from left right, starting with the vertex right after g, if there is path from 8 to that vertex, we add an edge between & and that vertex.
Now that have generated the graph, we next produce the DFS trace: We start with and visit every vertex in G2. The next correct step in the DFS algorithm would be to backtrack from a vertex in G2 any child vertex of that G1. Thus, the backtrack distance of this example is |G2| = B. As with the other graph and DFS example distributions, randomly permute the vertex IDs the last step.
20

---

Published as a conference paper at ICLR 2025
A.8
DFS SCALING
Maximom input gtoph sise
Minimum train loss
Minimum tost loss
10
10a
x 10
Moxtmmmm graph
3
106
107
10a
10
107
10
Training examples
Training examples
3 x 10"
Nos-embeddding
6 10-1
Minimum train loss
Minimum tost loss
2 x 10"
4
10P
3 10-1
cadecadeedicx paromeseers
GM 4M
10
x 10
109
10
107
10
Training examples
Training examples
3.x 10"
embedding peramelers
6 10
Minimum train loss
Minimum tost loss
2 10"
4.x 10
26
109
3 10
Non-conbeddingg
parauusees
4M
5.324
6 1100-
25.334
20.04
1015
1034
1045
1046
102h
10
1045
10a0
Floating point operations
Floating point operations
FIGURE 14: Training and test loss vs number of training examples seen, for models trained on the depth first search task. Test loss computed on held-out examples from the halanced distribution with backtrack 8. In the top row, fix the model size and vary the maximum input graph size. In the middle and bottom rows, fix the maximum input graph size and vary the model size. Note the x-axis in bottom row FLOPs. models were trained on the balanced distribution. Test loss is smoothed by averaging over a window of 81 data points, where each data point is recorded at every 2 = 262K cxamples. For each point, we plot the minimum loss over 15 seeds.
A.9 NON-STANDARD PADDING IN DES
We train 7-layer transformer with input size 128 the DFS task and evaluate on held-out examples with various backtrack distances. The results are shown the top row of Figure 15. We observe that
Standard training
1.00
0.81
0.56
0.46
0.14
0.32
0.33
0.22
0.20
0.21
0.23
0.22
0.22
0.20
0.28
0.15
Trained with random padding
1.00
1.00
0.92
2
14
Tested on examples with backtrack distance
FIGURE 15: Accuracy of model trained the depth-fire first search task with and without random padding, All evaluation is performed on held-out examples. Both models were trained on 746M examples.
21

---

Published as a conference paper at ICLR 2025
the backtrack distance in the DFS search task is analogous to the lookahead distance in the original search task, where the probability of generating examples from the naive distribution with higher backtrack distances is vanishingly low. However, unlike the original search task, we can augment training examples by adding padding between vertices in the sequence of visited vertices to help teach the model to backtrack greater distances. More precisely, starting with the right-most token, we add l padding tokens where k ~ of available padding tokens. We repeat this process for every token, from right to left. We observe in the second row of Figure 15 that when trained on these randomly- paddeddppddddddd inputs, the model more robustly learns to backtrack to greater distances. Unlike in the original search task where we carefully developed the balanced distribution to train the model to search to greater lookaheads, simple post-hoe augmentation of the examples was sufficient successfully train the model to perform DFS search.
where Tavastable is the total number
We note that while random padding may help the model to generalize to larger backtrack distances than those shown in training. they do not help the model to learn to search on larger graphs, as our scaling experiments have shown. In the scaling experiments, the training distribution contains examples of all possible backtrack distances, uniformly distributed.
A0 GENERATING SELECTION-INFERENC EXAMPLES
To generate a selection selection-inference inference example with a given frontier size F and branch count B. we first
sample the graph size [V] (number of vertices) from (niform([2... .2...... Fmaxx F + 1}), where Fmax is the maximum number of edges that fit for the given transformer input size. Next, we arrange [V] vertices from left to right (in topological order) without edges: Vi,---, Vini We select the index of the "current" vertex from Uniform (I....., V - B}). We then sample the indices of the start and goal vertices:
Uniform((1,.... min(c, [V| - F))), g ~ caiffomm(max(cc 1, + IV))).,
(9)
(10)
Add initial edges: Next, we iterate over each vertex in the graph V,, from left to right, and add edges as follows: First sample a number of parent vertices for V; from n. parents ~ Uniform ...., mmin(i - 1, + 1)}) where n is the transformer input size (in tokens). Next, we sample parent vertices from among (Vi,... ...113 one at a time, with probability proportional to the out-degree of each vertex. We add an edge between the selected parent and V/2 and repeat until we have n parcats parent vertices. However, if one of the potential parents is the current vertex Ve, and the number of child vertices of V. is B. we exclude it from the set of potential parents, as we want to ensure the branch count of V. is not larger than B.
Construct the frontier: Next, we sample the frontier vertices, which will be the vertices that have unvisited child vertices. Vg and V. are automatically added to the set of frontier vertices. We sample the remaining vertices from { [V++1..... Viv13 \ {V₂} uniformly at random until we have F frontier vertices. We then perform selection-inferent at random, but we avoid selecting an outgoing edge from any frontier vertex, and we perform the search until no available edges remain. Let & be the list of visited edges (in the order they were visited). Note that there may still exist frontier vertices that have not been reached in E. For each of these frontier vertices Vj: We randomly select an ancestor V, that has been reached in & and replace a random parent of V/2 with V., and add the edge V, - V; into E. However, it is possible that there is no path from V. to V,, in which case no ancestor of V; has been reached in E. In this case, we select a vertex from (V..... Vz-1) that has been reached in E, uniformly at random, and add it as a parent of V, The new edge is added to E. Note that each time an edge added to E, we move into a random valid position.
from Vs, selecting edges to explore uniformly
Ensure each frontier vertex has an unvisited child: At this point, we have guaranteed that every frontier vertex has been reached in E. Next, we have to ensure that every frontier vertex has at least one unvisited child vertex. For each frontier vertex V; without unvisited child vertices, we select new child vertex V, from [Vi+,..... Viv)) that has not been reached in E. Next, we randomly select parent of V, that has not been reached in E and replace it with V. If V, has no such parent, we simply add the edge V42 Vv. Vy.
Make sure the current node has B child vertices: Next, we want to ensure that V. has exactly B child vertices. First, we add frontier vertices children of V:: we select frontier vertex Vy from (Vet1,..., Viv } uniformly at random. If this vertex does not already have an edge from V.,, we add one. We replace the edge in & containing Vy as the target with the new edge V. - Vy. We repeat
22

---

Published as a conference paper at ICLR 2025
further frontier until max(0, 2F + B Emax) frontier vertices are children of Ve (or there non- froffronterr vertex vertices available). Then, we add non- non-frontier vertices as children of Vg: select a non V; from (Ve+1...., Viv uniformly random. We randomly sample parent of V3 that has not been reached in & and replace it with Ve. If V3 has no such parent, we simply add the edge V. Vj- Repeat until Ve has B child vertices.
Make sure the goal vertex is reachable from the start vertex: If Vg is reachable from V.. select a random reachable vertex Vi such that i < and add the edge - Vgg.
Remove some superfluous edges: We next remove number of superfluous edges (i.e., edges that are not in &, are not needed to keep V, and Vg connected, or are not needed to connect frontier vertices to an unvisited child vertex). If there are more than Emas edges, we remove superfluous edges randomly until Emax edges remain. Otherwise, we randomly remove n superfluous edges where n. is the number of edges we have added since adding the initial edges.
Add more edges to E: Next, continue the selection-inferen inferrnnee procedure from carlier to add additional edges to E, taking care that each frontier vertex still has least one unvisited child vertex. We continue until have nf edges, where n' = í with probability proportional i and n E {|E)...., Emmr); }. This step helps to make sure the size of & more uniformly distributed.
As with the other graph and DFS example distributions, we randomly permute the vertex IDs as the last step.
Note that since selection-infernnce consists two subtasks, we have encode the inputs differently. Rather than a list of visited vertices, we encode the list of visited edges. The example in Figure 13 would look like:
Input: E 1 E
Label:
Input: E
Label:
The top example is one of the selection subtask, where the model must predict previously-vsitted vertex with unvisited child vertices. The bottom an example of the inference subtask, where given the vertex e, the model must predict an unvisited child vertex. Interestingly, we find transformers do well on the selection subtask, but fare poorly on the inference subtask when given large input graphs.
23

---

Published as a conference paper at ICLR 2025
A.11
SELECTION-INFERENCE INFERENCE SCALING
x 10
Maximum input gtaph sise
3 10
6 x 10
Training loss
Test loss
2 10
4 10
3 10
2 10
Maximum graph
10a
10
10
10
10a
10
107
10
Training examples
Training examples
FIGURE 16: Training and test loss number of training examples seen, for models trained the selection- inference task. Test loss is computed on held-out examples from the balanced distrihution with frontier size and branch count 4. We fix the model size and vary the maximum input graph size. All models were trained on the balanced distribution. Test loss is smoothed by averaging over a window of 81 data points, where each data point is recorded at every 2 262K examples.
x 10
Not: embeeddnng parametens
3.234
Training loss
3.834
3 10
13.433
Test loss
6 10-1
Non-cabeddiigg
paramestess
3.
x 10
4 10
14.834
40.354
1014
102
10
10
1044
1045
1016
Proating-poitt operations
Floating- costing--pintt point operations
FIGURE 17: Training and test loss vs FLOPs, for models trained on the celection-inference task. Test loss is computed on held-out examples from the balanced distribution with frontier size 4 and branch count 4. We fix the maximum input graph size 45 vertices and vary the model size. All models were trained the balanced distribution. Test loss smoothed by averaging over window of 81 data points, where each data point recorded at every 28 262K examples.
A.12 DOES THE TRANSFORMER SEARCH FORWARDS OR BACKWARDS?
We investigate whether the trained model performs search forwards from the start vertex or back- wards from the goal vertex. To do so, we perform mechanistic interpretability analysis on the trained model, using 100 examples of graphs with lookahead 9. We then restrict our attention to the vertices along the path from the start to the goal vertex, and count the number of path-merge oper- ations between vertices along this path. We perform this analysis for two trained models: (1) the model trained on the balanced distribution with maximum lookahead 20, and (2) the model trained on the balanced distribution with maximum lookahead 12. Figure 18 visualizes the distribution of path-merge operations for the first model, and Figure 19 visualizes the same for the second model.
In the first model, we see a clear pattern in the first 4 layers where each target vertex is copying information from earlier and earlier source vertices (vertices that are progressively closer to the start vertex). Specifically, in layers and 2, the 1th vertex along the path from the start to the goal vertex predominantly copying from the i tth vertex along the path. In layer 3, the 7th vertex is copying from the i 2th vertex. In layer 4, the 5th vertex is copying from the i tth vertex. Thus, this particular model seems to be performing the search backwards from the goal.
However, in the second model (trained with maximum likelihood 12), the pattern is less clear. Layers and 4 suggest the search is performed backwards, similar to the first model, since most path- merge operations lie above the main diagonal. However, layers 2 and 3 suggest the search is being performed forwards, as most of the path-merge operations lie below the main diagonal. We conclude that the search direction may be dependent the random initialization seed, with some trained
24

---

Published as a conference paper at ICLR 2025
models performing search forwards and others performing search backwards. Further experiments
are necessary to obtain a more complete description.
Layer
Layer
Layer 3
Sourre vertox (distance from start)
Souroe vertex (distance from start)
Source vertex (distance from start)
0
0
0
6
6
6
Target vertex (distance from start)
Target (distannc from start)
TTrrett vertex (distance from start)
Layer
Layer
Layer 6
Source vertex (distance from start)
Sourve vertex (distance from start)
Sourre vertex (distance from start)
0
6
6
6
8
8
2
Target (distance from start.)
Target (distance from start)
Target (distance from start)
0.0
0.2
0.4
0.6
0.8
1.0
frequency
FIGURE 18: Visualization of the distribution of path-merge operations at each layer of the model trained on the balanced distribution with maximum lookahead 20. 100 input examples with lookahead were analyzed to obtain the set of path-mo merge operations. We restrict the visualization to path-merge operations between vertices along path from the vertex to the goal vertex. The cell at row and column j depicts the number of path-merge operations where source is th vertex on path the target the jth vertex on the path, divided by the total number of path-merge operations whose target is the j vertex on the path.
25

---

Published as a conference paper at ICLR 2025
Layer
Layer 2
Layer 3
Sonree vertex (distance from start.
Source vertex (distannc from start)
Source vertex (distance from start)
0
0
0
2
2
6
6
6
8
8
Target vertex (distance from start)
Target vertex (distance from start)
Target vertex (distance from start)
Layer
Layer
Layer 6
Source vertex (distance from start)
Sourve vertex dissance from start)
Source vertex (distance from start)
0
6
6
6
Target vertex (distance from start)
Target vertex (distance from start)
Target vertex (distance from start)
0.0
0.2
0.4
0.6
0.8
1.0
frequency
FIGURE 19: Visualization of the distribution of path-merge operations at each layer of the model trained on the balanced distribution with maximum lookahead 12. 100 input examples with lookahead were analyzed to obtain the along path from start vertex to goal yertex. cell at row and column j depicts number of path-merge operations where the source is the vertex on the path and the target is the j" vertex on the path, divided by the total number of path-merge operations whose target is the jth vertex on the path.
path-merge operations. We restrict the visualization to path-merge operations between vertices
26


Reviews:
[{"review_id": "xTociuEEZ7", "reviewer": "Reviewer_NMwf", "summary": "This study investigates whether transformers can learn to perform search by training small models on graph connectivity data. Results show that transformers can learn to search under certain conditions, but struggle with larger graphs, indicating that simply scaling LLMs may not enable robust search. The study introduces a new interpretability method to analyze the model's learned algorithm.", "questions": "- Interpretation of Figure 3: The paper claims that transformers can search under restrictive training distributions, but Figure 3 only seems to show limited generalization to larger lookaheads. Can the authors explain how this supports claims about the training distribution?\n- Pattern and Hypothesis Formation (Line 359): What specific pattern did the authors observe, and how did it lead to the hypothesis about the algorithm the model uses? Could they provide a clear link between the observed pattern and their hypothesis?\n- Explainable Attention Path (Line 337): What exactly is meant by a \"path of explainable attention operations\"? Was this path derived through manual inspection, or was there a specific method used?\n- Quantifying Examples (Line 358): The authors mention using \"a number of input examples\" but do not specify the exact number. Could they provide this detail to strengthen the robustness of their conclusions?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- The study tackles an intriguing and practical research question: understanding the mechanisms behind search capabilities in LLMs. This is not only scientifically interesting but also has meaningful implications for real-world applications.\n- Training a small GPT model on synthetic graph data is a reasonable and well-justified approach to investigate this research question.", "weaknesses": "The logical flow of the paper is weak in several areas. The authors should clarify the connection between their empirical results and the statements made, as well as provide more intuition behind their hypotheses. \nFor example, in line 51, the authors state, \"We demonstrate experimentally that transformers can indeed be taught to search, but only under fairly restrictive conditions on the training distribution.\" However, Figure 3 does not fully support this claim. While it may indicate that the model does not generalize well to a larger number of lookaheads than seen in the training data, it does not substantiate any firm conclusions about the training distribution itself. \nIn line 359, the authors mention, \"We noticed a pattern and formed a hypothesis about the algorithm the model has acquired to solve the search problem.\" However, the pattern observed and its connection to the proposed hypothesis remain unclear and should be elaborated upon.\n\nAdditionally, the proposed method and analysis require clarification. For instance, in line 337, the phrase \"path of explainable attention operations\" is used—was this path inspected manually? And in line 358, the authors mention \"a number of input examples\" without specifying the exact number. Providing this detail would help improve the robustness of their claim.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "cNJKwPguKc", "reviewer": "Reviewer_fM1a", "summary": "This paper aims to explore LLMs' internal mechanisms for graph connectivity problems tasks: given a graph (nodes and connections between nodes), a starting vertex, and a goal vertex, the LLM outputs the next vertex from the starting vertex.\nSpecifically, the paper constructs a training set to train a small decoder-only transformer. It improves the Mechanistic Interpretation visualization method to explore which tokens influence the LLM's output. Based on experimental results, the authors conclude that LLMs must be trained with in-context samples to fully understand the graph search problem.", "questions": "Please refer to the weaknesses part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The topic is interesting and may have influence in the community.", "weaknesses": "1. There are potential data leakage issues in the training and testing datasets constructed by the authors:\nThe authors use a generation method to generate training data online and save the first few generated results as test data. While the authors claim they will remove overlapping samples between training and test data, they don't explain how they compare whether two graphs are identical. If only using string matching, it cannot determine whether two graphs are completely equal. For example, these two graphs: node 1 -> node 2 and node 3 -> node 4 are completely equivalent but cannot be detected through string matching. Therefore, the test set is likely included in the training set. Additionally, given the number of vertices and max number of in-out edges, DAG generation is finite. Thus, the authors' claim about infinite graph generation may be incorrect.\nThe authors trained a simplified model that doesn't correspond to currently widely-used LLMs:\nFirst, since the authors used full attention rather than causal attention, they actually trained an encoder-only model rather than a decoder-only model.\n\n2. Second, the authors only used one-hot embedding for each token and position embedding when training this encoder-only model.\nFinally, for the graph connectivity problem, the authors' trained model only outputs one token, while current LLMs typically have reasoning steps.\n\n3. The authors' improved Mechanistic Interpretation has the following issues:\nThe proposed method requires performing perturbation and forward pass once for each element in every attention map of the LLM, and then forward pass to see the effects on the output logits, which is time consuming.\nIn Line 318, determining the influence of modified tokens on attention through frozen previous layers is not reasonable, as modified tokens also influence previous attention calculation and thus influence the activation for each layer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "ngGUrQRSxk", "reviewer": "Reviewer_zLV2", "summary": "The paper explores the behavior of transformers models when trained on search questions on directed acyclic graphs (DAGs). The authors show the importance of training data distribution for better generalization. Then, they conduct mechanistic understanding of the trained models to discover a progressive message passing algorithm utilized by the model to explore search paths. However, the models struggle to learn from larger graphs. Finally, the authors propose proxy in-context examples that help the model for robust exploration of the graph before solving the search problem. Overall, the paper represents a significant step towards our understanding of the inner mechanisms of transformer models.", "questions": "Please check my questions in the above section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The major strength of the paper lies in its motivation to understand transformer mechanisms in search based tasks. The authors take carefully designed experimental exploration to train transformers on directed acyclic graphs, with careful design discussion on data distribution, and propose a new mechanistic approach to analyze the learned algorithm. The authors discover a message passing algorithm, where the neighborhood information are shared progressively among the vertices, which leads to an exponential path-merging algorithm. \n\nThe authors also touch upon the difficulty of training transformer models on larger graph structures, and propose in-context tasks to help the model explore the graph better. Overall, the paper conducts an in-depth analysis of transformer model training on search problems and will be an important contribution to the academic community.", "weaknesses": "As such, the paper doesn't have many weaknesses. I have a couple of questions regarding the experimental setup.\n\na) **Sequence length in In-context exploration:** As the experiments require training on higher sequence length, how are the samples in training data distribution decided? How many steps in DFS traces are necessary for the model to learn? If the authors had provided same 'K' padding tokens to the experiments in the experiments in section 4, would the models generalize better?\n\nb) **Distribution of path-merge operations:** Are there patterns in the distribution of path-merge operations and copy operations across the layers in the trained model? \n\nc) **Evaluation with density of graphs?:** Do the trained models generalize to extremely sparse graphs? \n- Furthermore, on cases where the graph contains $2$ disconnected components, what will the model output be for start and goal vertices not in the same component? \n\nd) **Values of $\\alpha$, $\\kappa_1$ and $\\kappa_2$ in section 4**: How are these values decided in experiments?\n\ne) **Clarification questions:**\n\n- \"the log attention weight of each important operation in the last layer.\" (line 303) - \nWhat does log attention weight mean? How do you define important operation?\n- \"it requires many forward passes (linear in the number of attention operations and in the number of input examples).\" (line 354) -  Can the authors give details on the number of passes necessary? Furthermore, do the number of necessary passes depend (logarithmically) on the length of the search process for a given input example?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "zzZNZpnrIa", "reviewer": "Reviewer_wdgg", "summary": "This paper explores the ingredients that constitute the search ability in pre-trained language models. The authors introduce a synthetic setting--searching over DAGs represented in the natural language space--and pre-train autoregressive transformers of varying scale. The finding is mixed: transformers can implement search under restrictive data distribution, but face significant challenges with scaled problem sizes.  The authors have explored training strategies that encourage the transformer to generalize better.", "questions": "See weakness.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Combining mechanistic interpretability with the search problem is, to my knowledge, novel, as most prior works have focused on \"classification\"-style tasks that feature a very restricted output space in terms of vocabulary and length. The problem is more challenging in term of complexity, and would presumably require chain-of-thought capabilities to solve effectively. This new setting also prompts the author to introduce a new algorithm for mechanistic analysis, which may be of interest to the interpretability community.\n- I enjoyed the exposition style presentation of the paper, with each section introducing problem setting of growing complexity, as well as experimental findings that sufficiently supports these findings.\n- The authors study nuanced challenges for transformers to generalize on algorithmic tasks in the presence of distribution shift.", "weaknesses": "- My primary concern of this paper stems from the broader implication of the authors findings. Several prior works have found that large-language models can implement certain graph algorithms [1][2], including graph connectivity, and that this type of algorithmic reasoning can be improved with appropriate adaptations of chain-of-thought [3]. It is unclear whether the authors findings contradict, confirm, or offer more nuanced insights to prior works.\n- While the paper does a fine job surveying relevant works in mechanistic interpretability, it is somewhat lacking when situating itself in the LLM planning/search and theoretical expressivity literature. Aside from the aforementioned works, several works have directly studied whether LLM can internalize search (in the form of MCTS) [4] and explore in-context [5]. The lack of a theoretical analysis, or a proper discussion of them make understanding the authors' contribution challenging.\n- While the strategy that strengthens the LLM's search ability in the means of data augmentation is nice, it may not directly translate to practical guidance due to the synthetic nature of the task setup.\n\nOverall, this is an interesting paper in terms of its mechanistic study; but I would encourage the author to situate its theses more broadly with the recent growing body of work in LLM exploration, search, and algorithmic reasoning. \n\n[1] NPHardEval: Dynamic Benchmark on Reasoning Ability of Large Language Models via Complexity Classes, https://arxiv.org/abs/2312.14890\n\n[2] IsoBench: Benchmarking Multimodal Foundation Models on Isomorphic Representations, https://arxiv.org/abs/2404.01266\n\n[3] The Expressive Power of Transformers with Chain of Thought, https://arxiv.org/abs/2310.07923\n\n[4] Amortized Planning with Large-Scale Transformers: A Case Study on Chess, https://arxiv.org/abs/2402.04494\n\n[5] Can large language models explore in-context?, https://arxiv.org/abs/2403.15371", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)