Transformers Struggle to Learn to Search (2412.04703v2)

Abstract: Search is an ability foundational in many important tasks, and recent studies have shown that 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 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 in $n_{\text{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.

• The paper shows that a balanced training distribution enables transformers to learn a 'path-merging' algorithm for efficient graph search.
• It reveals that as graph size increases, scaling issues arise and simply adding model parameters does not resolve inherent architectural limits.
• Mechanistic interpretability techniques confirm that transformers merge reachable vertex sets, but suboptimal merging leads to errors in extended searches.

Below is a detailed explanation of the research, its motivation, design, and findings.

---

This work investigates whether a transformer—a type of neural network used in LLMs—can learn to perform "search" over a graph. In the experiments, the network is given a graph (which is a set of points connected by arrows), a start point, and a goal. Its task is to decide which neighbor along the path from the start will eventually lead to the goal. Although this may sound simple, it is a basic component of many reasoning and planning problems.

Background and Motivation

• Search as a Fundamental Task:

Searching through a graph is a core ability in many higher-level problems like planning and reasoning. If transformers can learn to search effectively, then they might be able to solve more complex tasks, from solving mazes to proving logical statements.

• Challenges with LLMs:

Recent studies have noted that LLMs sometimes take wrong turns or hallucinate when they try to plan or search. This paper asks whether the error comes from the way the transformer is built, the training data, or simply the size of the model.

Overview of the Approach

• Testbed Using Directed Acyclic Graphs (DAGs):

The authors choose graph problems as they allow for almost unlimited examples which are easy to generate. A directed acyclic graph has no cycles so that every path from the start eventually moves forward toward the goal.

• Various Data Distributions:
  • A “naïve” version that randomly generates graphs.
  • A “star” distribution where vertices radiate from a center.
  • A “balanced” distribution that is carefully designed so the model is forced to learn how to search rather than find shortcuts. In the balanced distribution, the “lookahead” (or number of steps the model must search) is uniformly distributed.
• Training Setup:

The models used are similar in size and design to GPT-2 but with modifications to ease mechanistic interpretation. They are given one-hot token representations so that every token directly indicates its value. Also, special care is taken with attention (how the model pays attention to different parts of the input) by not using a causal mask.

How the Model Learns to Search

• Learning a “Path-Merging” Algorithm:
  • Each vertex’s embedding (a vector that represents its state) stores the set of vertices it can reach.
  • Each transformer layer then updates these sets by merging information from adjacent vertices.
  • For example, if a vertex “A” knows it can reach “B” and “B” knows it can reach “C,” then “A” can record that it may reach “C.” This gradual merging lets the model “search” over many steps in just a few layers.
• Mechanistic Interpretability Studies:

To understand exactly how the model is making its decisions, the researchers use techniques that perturb parts of the transformer and then measure how the output changes. In other words, they “patch” the network’s activations at many places and then reconstruct the computation graph (a map of which operations led to the final output). This analysis gave them evidence that the model was indeed copying and merging the sets of reachable vertices consistently from the first to the final layer.

Key Findings

1. Importance of the Training Distribution:
   • The balanced distribution is crucial for teaching the network to learn proper search.
   • Models trained on simpler distributions struggled when the “lookahead” (number of steps) increased.
2. Scaling Issues:
   • As the input graph grows larger, the transformer has increasing difficulty learning the search task.
   • Simply increasing the number of model parameters did not alleviate the issue, suggesting there may be a fundamental architectural limitation.
3. Chain-of-Thought and In-Context Computation:
   • The researchers also experimented with giving the model the chance to write out intermediate steps—similar to “chain-of-thought” prompting.
   • Even when allowed intermediate steps (like a depth-first search), the model continued to struggle on larger graphs.
4. Mechanistic Insights:
   • The mechanistic interpretability method revealed that the transformer does not always perform an optimal merge of the available information. In other words, it does not always combine the largest sets possible. This non-maximal merging may contribute to errors when the input graphs are larger than those seen during training.

Relevance and Implications

Understanding whether transformers can learn to perform a search in structured problems is important because:

• It informs us about the limits of current architectures. If a transformer struggles with a fundamental algorithmic problem like graph search, then that may explain why larger LLMs still sometimes err in planning and reasoning.
• The techniques developed in the paper help researchers “peek inside” the transformer’s computations. This mechanistic interpretability effort is valuable for improving model architectures and training methods.
• The findings hint that alternative training procedures (perhaps using different curricula or architectural modifications) might be necessary to scale reasoning and planning abilities in large models.

Conclusion

In summary, this research shows that while transformers can learn to search on small graphs if trained carefully with a balanced distribution, they struggle with larger graphs. The paper raises important questions about the fundamental limitations of the transformer architecture regarding planning and search. The mechanistic interpretability method not only confirms that a “path-merging” algorithm is being employed but also highlights areas where the algorithm is suboptimal, especially as the scale increases.

This work is an important step toward understanding the emergent abilities of LLMs and provides guidance on how we might enhance their capacity for structured reasoning in the future.