Mechanism and Emergence of Stacked Attention Heads in Multi-Layer Transformers

Authors: Tiberiu Mușat
Primary area: interpretability and explainable AI
Keywords: mechanistic interpretability, large language models, transformers, emergent abilities, curriculum learning, reasoning

Abstract:
In this paper, I introduce the retrieval problem, a simple yet common reasoning task that can be solved only by transformers with a minimum number of layers, which grows logarithmically with the input size. I empirically show that large language models can solve the task under different prompting formulations without any fine-tuning. To understand how transformers solve the retrieval problem, I train several transformers on a minimal formulation. Successful learning occurs only under the presence of an implicit curriculum. I uncover the learned mechanisms by studying the attention maps in the trained transformers. I also study the training process, uncovering that attention heads always emerge in a specific sequence guided by the implicit curriculum.

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

Reviews and discussion:
<thread>
creation date:
1730610689029

thread_type:
official_review

<thread_object>
creator: reviewer_yct3

<content>
summary:
This paper introduces the retrieval problem, a fundamental task that challenges transformers to retrieve information from multiple positions within an input sequence. The authors demonstrate that solving the retrieval problem requires a certain depth of transformer layers.
 This paper explores how large language models and transformers handle retrieval tasks. They introduce two tasks: the retrieval problem and a variant called the conditional retrieval problem. It shows that large language models can solve retrieval tasks without fine-tuning by leveraging complex mechanisms involving multiple attention heads, which emerge in a specific sequence when trained with an implicit curriculum.

strengths:
● The paper provides insights into how LLMs perform retrieval using attention heads.
● Introducing these tasks gives a clear way to study transformers' retrieval abilities.
● The study highlights the role of learning curriculum in the development of retrieval mechanisms.

weaknesses:
1. **Lack of Experimental Validation for Theoretical Claims**: Theoretical claims like Theorem 1 lack empirical support, making it hard to verify their practical impact.
2. **Insufficient Details on Experimental Setup**:: The paper lacks detailed explanations for key experimental settings. The implicit curriculum (IC) formulation, which performs better than non-IC, is not clearly defined, and Section 8’s description of manually reverse-engineering circuits lacks detail, making the experiments difficult to reproduce.

questions:
1. Assumption 1, which requires shared positional or token information for attention between positions, is overly restrictive and conflicts with the flexibility of Transformer models. Transformers are designed to allow any position to attend to any other, enabling them to learn complex relationships without pre-shared information. This assumption limits the model's ability to capture long-range dependencies and generalize to tasks where connections are context-driven rather than based on shared information. I recommend reconsidering this assumption to better align with the strengths of the Transformer architecture.
2. Theorem 1  would benefit from experimental validation. Testing whether the target embedding truly appears in the residual stream only when \( t \geq \log_3(2D) \) would strengthen the claim and confirm its practical relevance. I recommend adding experiments to verify this behavior across different settings.
3. The specific details of the implicit curriculum formulation remain vague, and the reasons for its superior performance are not clearly explained. Additionally, it would be helpful to clarify why the non-IC formulation, which aligns more closely with traditional problem setups, fails to perform well.
4. In Section 8,  "manually reverse-engineering the circuits learned by the transformers" , the specific details of this process are not fully explained. Providing more information on the methodology used to reverse-engineer these circuits, as well as any criteria or steps involved, would enhance the clarity and reproducibility of this analysis.
5. The experimental results need to be validated under more complex settings, such as with longer retrieval steps.
We will increase the score based on the answer to the question.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626288939

thread_type:
official_review

<thread_object>
creator: reviewer_cntj

<content>
summary:
This paper presents a retrieval task to illustrate the source of the reasoning power of hierarchical transformer models. With the retrieval task, the author takes various structures to verify the connection between successful learning and the presence of an implicit curriculum. In my opinion, this work is helpful in understanding the ability of large language models. As descributed in this paper,  a transformer cannot solve the retrieval problem with a specific difficulty unless it has the minimum number of necessary layers.

strengths:
This is a good joy for understanding the ability of LLM, especially the emergent ability of models.


1.  a novel idea to study the ability of the reasoning ability of LLM.

2. The finding is exciting and fits with human intuition. 

3. The visualization of attention can give describution on how LLM takes retrieval reasoning tasks.

weaknesses:
1. Can you provide a more complex example? 

2. In my opinion, I hope I can see a general framework that can unify more tasks with your retrieval task.

3.  There are a lot of chapters in the article, and I can't quite understand the relationship between different chapters.

questions:
Can you describe the definition of the implicit curriculum in detail? In your experiments, I can see the connection between the number of layers and the model's performance, as shown in Figure 4.

The emergent ability of LLM is more depend on depth or width?  This paper mainly analyzes the effect of model depth on emergence ability.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730640736268

thread_type:
official_review

<thread_object>
creator: reviewer_mjk6

<content>
summary:
The paper proposes a novel retrival problem to conduct an empirical but theoretically grounded anaylsis of mechanisms leading to reasoning in transformers. 
Expanding on the induction problem the authors generalize the task to an arbitrary number of inductions steps and introcude a conditional variant. 
The authors demonstrate that current SOTA LLMs struggle with the retrival problem and their performance approaches random guessing with increased problem complexity. 
Experiments on a minimal problem formulation indicate the importance of implicit curricula and highlight a correlation between the difficulty of the task and the number of layers.
Through manual inspection and subsequent experiments the authors identify learned circuits in a subset of attention heads which emerge one head at a time during training.

strengths:
- The paper makes a strong contribution to the field of mechanistic interpretability and furthering our understanding of the transformer architecture and training behavior. Especially the finding on the sequential emergence of attention heads for reasoning circuits is valuable. 
- The **problem statement** at hand is presented nicely, and the paper follows a logical progression building up to the final insights.
- The **structure and flow of the experiments** are sensible, starting with higher-level analysis and ending in more focused experiments. 
- Furthermore, the experiments were conducted rigorously.

weaknesses:
**Readability:** 
- Section 5 "THEORETICAL ANALYSIS OF INFORMATION FLOW" is quite **hard to follow** and requires some time to understand, especially with limited prior knowledge. Concrete, examples to what "E", "F", ... and so forth might mean in the training and evaluation context could help the reader to grasp the theoretical analysis quicker. Augmenting this section with examples, e.g. from  Sec. 4, would make it more accessible to a larger audience. 

**Missing clarity:** 
- Large portions of the paper's analysis are **based on a strong assumption**. Key aspects are *Assumption 1* along with simplification of the transformer architecture. The authors should provide further justification for the reasonability of these assumptions and why they would generalize to LLMs. For example, an elaboration on the "shared token information" assumption in real-world scenarios would benefit the paper.  
- The **real-world example from the discussion section could be used earlier** to better motivate the theoretical analysis and problem statement.

**Generalization:**
- Similarly, the paper would benefit from an analysis on a real-world question-answer dataset, for example. 

**Method:**
- Lastly, the **method of _manually_ reverse engineering** the circuits is only mentioned but **not described in detail**. This severely affects the potential replicability of the paper and makes independent assessments of the methodology challenging.


**Post-rebuttal edit**
The rebuttal addresses the main weaknesses of the paper. I adjusted my score accordingly.

questions:
1. Can the authors elaborate on why Assumption 1 should hold for LLMs, as well as how their 
2. From section 4 and Figure 2, I take that all variations of the problem formulations (F1 to F4) include 4 choices, given the random guessing probability of 25%. However,  this does not seem to be the case for F5 "Relatives". Is there a specific reason for that?
3. Do the authors have an intuition of how their results might be generalized to real-world data? It would be valuable to asses if reasoning circuits follow similar emergence patterns during training. However, I assume that the manual identification of circuits is not scalable. 
4. What methods were used to "reverse engineer" the attention head circuits?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732240538548

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We would like to thank all reviewers for their insightful feedback!

We have updated our paper to address the two most important criticisms:

1. We have **weakened assumption 1** and improved its justification
2. We have added a detailed description of the **reverse-engineering procedure**

We changed assumption 1 from stating that attention requires a "shared token embedding or positional encoding" to stating that it requires a "shared piece of information" (any information). This assumption is motivated by the fact that a position can only attend to another position if their key and query vectors align. Constructing aligned key-query pairs is only possible if the two positions already share some information. More precisely, the shared information must be located in the row spaces of the query and key matrices of the attending and attended positions, respectively.

To maintain the soundness of the argument, we have introduced the old assumption as a lemma dervied from the new assuption. Its proof is based on the fact that the token values and pair positions are assigned randomly and independently.

Regarding the reverse-engineering procedure, we have added a comprehensive description in Appendix D. We have decided to place it in the appendix (rather than the main paper) since the procedure is not particularly innovative or interesting. In fact, we believe that a great number of different procedures could have been used to arrive at the same circuits. Nonetheless, including the procedure is important for reproducibility.

We would like to thank the reviewers for highlighting these issues and we hope that you will appreciate these improvements!

If the paper is accepted, we will fix the remaining issues in the presentation and clarity of the paper, as highlighter by the reviewers. We are also happy to include additional experiments. Any suggestions in this regard are highly appreciated!
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper analyzes a retrieval problem, that involves extracting information from multiple positions within an input sequence. The authors show that solving this problem necessitates a specific depth of Transformer layers. The study investigates how large language models and transformers approach retrieval tasks, introducing two variants: the retrieval problem and the conditional retrieval problem. It reveals that large language models can tackle retrieval tasks without fine-tuning by utilizing sophisticated mechanisms involving multiple attention heads, which naturally emerge in a specific order when trained using an implicit curriculum. There is also theoretical analysis for this. 
The paper tackles an interesting topic in the LLM interpretability domain and demonstrates results not only for toy LLMs, but also sota LLMs like Gemini etc. The reviewers particularly find the experimental studies, the clear results and the interesting findings of the role of an implicit learning curriculum as strong points of the paper. As weaknesses, there were some initial questions about insufficient experimental details and lack of too strong assumptions, the hard-to-follow theoretical section 5 and lack of results on real-world Q&A datasets.  Given the clear idea and execution and the additional clarifications provided in the rebuttal phase, the AC recommends acceptance.