Transformers Struggle to Learn to Search Without In-context Exploration
search;reasoning;transformers;scaling laws;mechanistic interpretability;circuit analysis
Yes
No ethics review needed.
- 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?
- 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?
- 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?
- 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?
- 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.
- Training a small GPT model on synthetic graph data is a reasonable and well-justified approach to investigate this research question.
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.
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. 
For 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. 
In 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.

Additionally, 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.
Yes
No ethics review needed.
Please refer to the weaknesses part.
The topic is interesting and may have influence in the community.
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.
Specifically, 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.
1. There are potential data leakage issues in the training and testing datasets constructed by the authors:
The 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.
The authors trained a simplified model that doesn't correspond to currently widely-used LLMs:
First, since the authors used full attention rather than causal attention, they actually trained an encoder-only model rather than a decoder-only model.

2. Second, the authors only used one-hot embedding for each token and position embedding when training this encoder-only model.
Finally, for the graph connectivity problem, the authors' trained model only outputs one token, while current LLMs typically have reasoning steps.

3. The authors' improved Mechanistic Interpretation has the following issues:
The 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.
In 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.
Yes
No ethics review needed.
Please check my questions in the above section.
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. 

The 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.
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.
As such, the paper doesn't have many weaknesses. I have a couple of questions regarding the experimental setup.

a) **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?

b) **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? 

c) **Evaluation with density of graphs?:** Do the trained models generalize to extremely sparse graphs? 
- 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? 

d) **Values of $\alpha$, $\kappa_1$ and $\kappa_2$ in section 4**: How are these values decided in experiments?

e) **Clarification questions:**

- "the log attention weight of each important operation in the last layer." (line 303) - 
What does log attention weight mean? How do you define important operation?
- "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?
Yes
No ethics review needed.
See weakness.
- 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.
- 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.
- The authors study nuanced challenges for transformers to generalize on algorithmic tasks in the presence of distribution shift.
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.
- 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.
- 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.
- 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.

Overall, 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. 

[1] NPHardEval: Dynamic Benchmark on Reasoning Ability of Large Language Models via Complexity Classes, https://arxiv.org/abs/2312.14890

[2] IsoBench: Benchmarking Multimodal Foundation Models on Isomorphic Representations, https://arxiv.org/abs/2404.01266

[3] The Expressive Power of Transformers with Chain of Thought, https://arxiv.org/abs/2310.07923

[4] Amortized Planning with Large-Scale Transformers: A Case Study on Chess, https://arxiv.org/abs/2402.04494

[5] Can large language models explore in-context?, https://arxiv.org/abs/2403.15371
@inproceedings{
anonymous2024transformers,
title={Transformers Struggle to Learn to Search Without In-context Exploration},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=9cQB1Hwrtw},
note={under review}
}
Search is an ability fundamental 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 graph connectivity 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, under specific conditions on the 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 simply increasing the scale of LLMs will not lead to robust search abilities.

Finally, we show that by loosening the task to allow the model to *explore* the graph *in-context*, allowing the model to visit vertices that do not necessarily lead to the goal and backtrack, the transformer is able to more easily learn to search robustly.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
search reasoning transformers scaling laws mechanistic interpretability circuit analysis
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/e6e91c49837af9fb30d01ee72f5c145dd1162a13.pdf
interpretability and explainable AI
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Transformers Struggle to Learn to Search Without In-context Exploration
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission