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
Primary area: interpretability and explainable AI
Keywords: search, reasoning, transformers, scaling laws, mechanistic interpretability, circuit analysis

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.

Decision: accept (poster)
Reviewer scores: 5, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730505668616

thread_type:
official_review

<thread_object>
creator: reviewer_wdgg

<content>
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.

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

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

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696803679

thread_type:
official_review

<thread_object>
creator: reviewer_zlv2

<content>
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.

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. 

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.

weaknesses:
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?

questions:
Please check my questions in the above section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730939891968

thread_type:
official_review

<thread_object>
creator: reviewer_fm1a

<content>
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.
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.

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

questions:
Please refer to the weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731328942649

thread_type:
official_review

<thread_object>
creator: reviewer_nmwf

<content>
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.

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

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?
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732381547484

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Reviewers

comment:
We sincerely thank all reviewers for their thoughtful feedback. We have taken your comments into careful consideration and worked to improve on the paper. We have improved the clarity throughout the paper and tightened many of the claims.

Additionally, we added more complete experimental results, notably in Section 6, where we experiment with in-context (i.e., step-by-step) methods for solving the task, akin to chain-of-thought. We added results on scaling experiments for the depth-first search task (Section 6.1; Figure 14). We also experimented with selection-inference “prompting,” where each step of the search is decomposed into two subtasks: (1) select a previously-visited vertex that has unvisited child vertices, and (2) from the current vertex, predict an unvisited child vertex. We added scaling experiment results for this selection-inference task (Section 6.2; Figures 16 and 17). We found that, while depth-first search and selection-inference is easier to learn than direct prompting, the transformer still struggles on larger input graphs, and we find that increasing model scale does not mitigate this difficulty.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper investigates the ability of large language models (LLMs), specifically transformers, to perform search tasks using graph connectivity problems as a testbed. The paper includes experiments on directed acyclic graphs (DAGs) and introduces novel mechanistic interpretability tools for this problem,  revealing that transformers compute reachable vertex sets layer by layer. All the reviewers appreciated the paper's contribution in terms of methodology and empirical findings, and agree than the paper is likely to be of significant interest to the broader ICLR community. The main weakness noted by the reviewers pertained to presentation and writing. Many of those were resolved by a thorough rebuttal and during the discussion phase. Additionally, reviewer fM1a raised concerns about a potential data leakage, which seems to have been resolved by the authors' response. 

Based on the reviews, rebuttal, and discussion, I recommend accepting the paper. The main reasons for this decision are:
* **Significant Contribution**: The paper addresses a fundamental problem in the field, and adds to our growing understansding of the capabilities and limitations of LLMs. The core findings of the paper (e.g. the mechanism by which the models perform search over graphs) are to the best of mine and the reviewers' knowledge, novel. 
* **Thorough Experimentation**: The experiments are well-designed and provide strong evidence for the paper’s claims.
* **Clear Presentation**: The paper is well-written and clearly presents its findings, making it accessible to a broad audience.

Overall, I think this paper would make a good addition to ICLR. I anticipate it will garner significant interest and will be heavily discussed.