Talking Heads: Understanding Inter-layer Communication in Transformer Language Models

Summary

This paper investigates the interaction between attention heads at different layers in a transformer. They primarily study the “inhibition-mover subcircuit”, a previously identified attention head interaction from circuit analysis work [1,2]. They show the interaction between heads can be characterized by a low-rank subspace, and show it is helpful in reducing prompt sensitivity on a simple item-recall task. [1] Wang, et al. Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small. 2023. (https://openreview.net/forum?id=NpsVSN6o4ul) [2] Merullo, et al. Circuit Component Reuse Across Tasks in Transformer Language Models. 2024 (https://openreview.net/forum?id=fpoAYV6Wsk)

Strengths

- The work is well-aware of existing related literature, and the questions posed are interesting. - The work is well-motivated - current LLMs are often not robust to prompt variations, and understanding and mitigating the reasons for this is of value. - The experiments demonstrate that the identified low-rank subspace can be used to effectively improve model performance on the item-recall task they introduce.

Weaknesses

- The study is limited to fairly small transformers (GPT2-small, and Pythia-160m in appendix), and it’s hard to know whether the results will generalize to more complicated tasks or methods will scale to larger models. Additionally, It would also be helpful to know whether larger models also struggle with the item-recall task (and whether they have similar inhibition-mover subcircuits). - The composition score presented in Equation 1 should be presented in context (i.e. seems to have come from Elhage, et al. [3]? Is there earlier use of this score elsewhere?). Besides substituting individual SVD components in for QK or OV matrices, are there any other changes compared to how [3] measure the composition score? - It is unclear to me whether the failure mode of GPT2 small on the item-recall task is due to (a) not correctly identifying duplicated objects, or (b) not suppressing correctly identified duplicated objects when going to copy. This is because Figure 4 shows the attention pattern of mover heads can be influenced by either duplicate token head channels or inhibition head channels. Is this because the duplicate token channel influences the inhibition signal? (i.e. its effect on the mover head is mediated by the inhibition head?) - There are some portions of the paper that present results without experimental details. One example is on Lines 228-231: "On OpenWebText ... we find inhibition heads are primarily active in lists and settings where repetition...". Providing details about the experiments run to validate this claim would help strengthen the argument made. In addition, what does it mean for an attention head to be "active" on text or not? Is this measured by its attention pattern or something else? [3] Elhage, et al. A Mathematical Framework for Transformer Circuits. 2021 (https://transformer-circuits.pub/2021/framework/index.html) ___ There are a few places with incomplete sentences or minor typos. I've listed a few I noticed below: - Line 174-175 - “... test how this affects.”, (incomplete sentence) - Line 187 - “communicatoin” channels - Line 214 - “non” - Line 276 - “repersent” - Line 290 - “featuers” - Line 506 - "These results are in", (incomplete sentence)

Questions

- When you say a head is “highly active” on a prompt or task [e.g. Line 229, 255], what is meant by that? Is this measured by its attention pattern or something else? When would an attention head be "non-active"? - When using a subspace to steer the inhibition score, the singular values needed to flip attention from IO to S (or vice versa) seem rather large. What are the typical ranges of the singular values you see in your decomposition of the weight matrices you’re investigating? - As one of the other well-studied/well-known subcircuits, have you done any analysis to understand whether the induction subcircuit is also dominated by a low-rank subspace? Or is this specific to the inhibition-mover circuit you study? - In section 5, how does trying to learn optimal singular vector weightings either via gradient descent or regression compare to doing a grid search over 3-d points?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors have addressed limitations of their work, which are reasonable.

Summary

The paper investigates the communication between attention heads across different layers in Transformer-based language models. First, it establishes that a previous composition metric, which has been shown to be useful in toy settings, is noisy in larger models when tested on the Indirect Object Identification (IOI) task with a known circuit. The authors then propose a modification using Singular Value Decomposition (SVD), rewriting the original matrix as the sum of outer products of the left and right singular vectors, scaled by the corresponding singular values. This new metrics reveals inter-layer communication more clearly and provides a less noisy signal. This is causally verified by zeroing out a single singular value on the IOI task which reduces the inhibition score notably. The authors then study these specific communication channels in detail, showing how interventions in these channels can affect the model's downstream behaviour. Interestingly, they demonstrate that this mechanism is independent of the specific context, functioning as a general pointer over lists. Finally, the authors use these insights to study the seemingly arbitrary sensitivity of language models to the order of items in lists using an artificial laundry list task. They show that as the list length grows, the internal structure in the communication partly collapses and makes it hard for the model to select the correct item. This can in part be addresses using a custom intervention, again demonstrating the causality of this phenomenon on the model behaviour.

Strengths

- This paper addresses a fundamental question in interpretability of transformer-based language models, namely how attention heads selectively communicate with each other. Although it was known from prior work that attention heads compose with heads in other layers, to the best of my knowledge, this is the first in-depth study of inter-layer communication in large language models. - The authors use a variety of techniques to systematically test and causally verify their hypotheses; I also find the use weight-based decompositions instead of activation-based techniques intriguing, as it enables static analysis of neural networks without the need for input data. - The findings are connected to previous observations on language model robustness, explaining a seemingly arbitrary sensitivity of language models to the order of items in lists.

Weaknesses

- The proposed composition score does not work for models using relative positional embeddings like RoPE, which are standard in state-of-the-art language models such as Gemma [1]. This makes the composition score not directly transferable. - The results could be significantly strengthened if the technique was used to conduct an unsupervised search for head compositions within the model, rather than focusing solely on the IOI task with a known circuit. Minor issues: - Typo in line 11: “via analysis *of* their weight matrices” - Typo in line 176: “the OV [or] matrix” - Typo in line 187: “communicatoin” - Overlapping image in Fig. 5, covering part of the “10" [1] G. Team et al., ‘Gemma: Open Models Based on Gemini Research and Technology’, arXiv [cs.CL]. 2024.

Questions

1. Could you elaborate on the intervention used to improve model performance on the laundry list task? Specifically, what operations do you perform to "set the model components in a certain area of the 3D space"? 2. Can the proposed composition score be used to identify compositions of attention heads without prior hypotheses?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

I believe the limitations are properly addressed, although the most important limitation - that the composition score is not directly applicable to models with relative positional embeddings - is somewhat hidden in the appendix.

Summary

This paper explores the routes and mechanisms of information transfer between heads of transformer large language model. They hypothesize the presence of low-ranking communication channels between attention heads from different layers within residual connections. The authors propose a method based on Singular Value Decomposition of weight matrices to detect those channels in a pretrained model; this method doesn’t require any additional data. They show that this method can uncover intricate and interpretable intrinsic structures of a transformer, which in turn can be manipulated to targetly adjust the model’s performance. The authors illustrate it by significantly improving the performance of GPT-2 on the list item recall task.

Strengths

•) Very interesting and novel idea. •) Good analysis; promising results that can facilitate further progress in interpretability and understanding of the innerworkings of large language models. •) The article is well-written, and the text is easy to follow. The figures are nice and illustrative.

Weaknesses

1) The proposed method of Communication Channel Interventions improves model performance on a synthetic task (Laundry List); however, it is not clarified how such modification will affect model's performance on other tasks. 2) The contextualization in the research field (Section 6, Related Work) is short and covers rather little information on previous works about information passage within transformer models and attention mechanism. 3) The limitations of the method do not allow for a straightforward implementation on newer models with relative positional embeddings. 4) [Minor] The possible practical applications of the achieved results are outlined vaguely.

Questions

1) You identify several types of attention heads: mover, duplicator, and inhibitor, but could a single head exhibit traits of multiple types? Also, is this classification not exhaustive (i.e., there are heads that do not fall into any of the aforementioned categories)? 2) If there are two interventions and each of them is beneficial (i.e., increases the model's performance) on its own, how often is their composition (i.e., applying them simultaneously) beneficial? Is it possible to somehow predict such pairs of interventions without directly computing the quality of each pair? 3) How many heads were affected by the interventions in your experiments in Section 5? Do you have information on how the increase in performance depends on the number of modified heads?

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

Authors clearly outline the limitations of their work in the dedicated section (in Appendix B).

Summary

Building upon prior work in Transformer circuits (Elhage et al., 2021; Wang et al., 2022), the paper identifies novel, low-rank communication channels between attention heads across layers via the Composition Score (a generalized cosine between the read-out weights of a lower-level layer and the read-in weights of an upper-level layer) between _low-rank_ factorizations of the weight matrices (obtained by SVD). The low-rank decomposition is critical, resolving the known issue of the vanilla Composition Score being very noisy. Generally focusing on the role of inhibition heads, which are known to prevent the copying of certain tokens in the prompt, the paper identifies a low-rank communication channel for inhibition and validates its role on a synthetic "Laundry List" task. The task requires recalling an item from a list of items specified in the prompt. It is shown that, when the identified (rank-1) inhibition channel is zeroed out (intervention), the model effectively stops passing the "inhibition signal." The role of the channel is validated for indirect object identification (IOI) and token indexing tasks in the GPT2-small model.

Strengths

- Using the SVD as a remedy for fixing the noisy Composition Score is an intuitive and seemingly effective strategy. - The paper’s main contributions, including the method and the new synthetic dataset, are reasonably well-motivated. - The empirical results make sense (for the most part) and support the claim of low-rank communication channels for inhibition heads. It is interesting to see that the inhibition channels are content-independent (just pointing at names in the IOI task) or that it affects the accuracy on the laundry list task.

Weaknesses

- Overall, while the paper is well-motivated and includes potentially interesting results, its presentation significantly hinders the reader from understanding the main takeaways of the paper. - First, a key piece of related work is the composition analysis of two-layer, attention-only models by Elhage et al. (2021), but the review of the relevant background work is simply too terse and unorganized. Terms like inhibition heads, inhibitor-mover subcircuit, value composition, and the QK/OV circuits (e.g., which weight matrices are we talking about exactly? why are they “low-rank”, especially when it’s not as low-rank as the subspaces we’re discussing here?), are not formally defined. Even things like why the composition score looks the way it is (why is it a matrix multiplication?) or the basic shapes of the weight matrices ($d$ by $d_h$ or $d_h$ by $d$?) involved are not written out in the “background” section. Without these in place, I think the paper is borderline incomprehensible to anyone outside the mech interp community. - Second, all the overclaims in the paper really hurt the overall message. The paper’s main thesis could simply have been something like “using the SVD in the Composition Score allows us to identify low-rank communication channels, as showcased by the identification of inhibition head channels,” which would be clear and interesting. But instead, the paper keeps trying to convey a broader, unsubstantiated message that their methods and experiments somehow are more novel and general than they should be. - “Understanding Inter-Layer Communication”: Exactly what is the scope of this inter-layer communication? Is it primarily/always from layer $l-1$ to $l$, or something more? - ”… using subspaces we call communication channels” (Intro): this terminology already existed in Elhage et al. (2021), but this phrasing makes it sound like the current paper came up with it. - “model components like attention heads” (what other components does this paper cover?) - “Three types of composition”: It appears that most of the focus in the paper, including the dataset itself, is on inhibition heads. - “Composition Score”: It should be clear to the reader that the score was first developed by Elhage et al. (2021), especially in Section 3.1. More generally, the intro should do a better job of clearly disambiguating what was introduced already in the literature and what this paper is adding to it. - Lastly, the writing is simply not well polished, and there are many typos and unnatural sentence structures that make it even more difficult to understand the overall paper. Just to name some of them: - Abstract, line 3 & line 11 - Intro, line 34 (\citep) - Page 4, line 114 - Page 4, line 128 (d \times d) - Figure 2 caption, second sentence - Page 5, lines 170—172 (broken parallelism) - Figure 3 caption, second sentence - Page 9, line 276 There’s also a related issue of overusing informal analogies unnecessarily (attention heads “talking” to each other). - Aside from the presentation, the biggest confusion for me is: generally, what implications do we have about inter-layer communication in Transformers from these findings? Besides inhibitions, what compositions do we expect to be represented as low-rank subspaces according to the paper’s approach? How far across layers can the channels be? Do we have any sense of what the rank tells us about the feature? The discussion section should be expanded substantially to at least mention or address some of these points. - In terms of the proposed Laundry List task, a possible caveat is that models more recent (and larger) than GPT2-Small may actually do a good enough job of solving them. I am sympathetic to the fact that it is challenging to redo many of the analyses on larger models, but I think it is worth, at least, to test the accuracy (Fig. 1, left) on more recent, open-weight models, to see how much the motivation in the introduction is still relevant. - The related work on SVD seems terse, knowing that SVD is arguably one of the most basic linear algebra tools in ML (any method based on PCA would also be relevant, for example). - In the intervention experiments, the choices of the exact subspace (the number of components) and the intervention method (e.g., scaling diagonally for 2D) appear a bit heuristic and need better justification.

Questions

- Are there other baseline methods to be considered beyond the vanilla and SVD-based Composition Scores for identifying communication channels? - Page 6, lines 217--218: How exactly is the 2D duplicate token head found? - Page 7, line 227: In what sense are these results “strong”? Is there a baseline? - Page 7, lines 228--232: any quantitative results for this claim? - Page 8, line 266: why is it top-3 this time? how should someone replicating your experiments choose this number? - I think the talking heads analogy makes sense, but it’s one that was already used before by [Shazeer et al. (2020)](https://arxiv.org/abs/2003.02436) in a not-so-unrelated context (the paper is fully referenced in the circuits paper, for example). I would suggest changing the title or at least adding a footnote of disambiguation from that paper.

Rating

5

Confidence

4

Soundness

2

Presentation

1

Contribution

3

Limitations

The paper states that it does not make claims about the different types of compositions across attention heads in different layers. Aside from that, the limitations are not well addressed; see weaknesses above.

Summary

This paper primarily employs low-rank communication channels to elucidate how internal layers within the Transformer model transmit information. Initially, the study utilizes existing research to identify duplicate heads, inhibition heads, and mover heads. The objective is then to explore the interactions among these heads. The findings indicate that directly incorporating these elements into the computation of the Composition Score does not yield significant results due to the complexity of the signals involved. Consequently, the paper proposes the application of Singular Value Decomposition (SVD) on the weight matrix, with subsequent sorting by variance, to identify the principal read/write subspace. This approach identifies such communication channels solely through the model's weight file. Furthermore, in certain recall tasks, such as the Indirect Object Identification (IOI) and Laundry List Task, this paper leverages the low-rank characteristics of communication channels to enhance task accuracy.

Strengths

1. This article provides a detailed background to facilitate readers' understanding of prior research. 2. The article presents a hypothesis, verifies it through experiments, and conducts further in-depth research, thereby demonstrating a complete scientific research process. 3. To validate the proposed hypothesis, the article also develops a Laundry List Task for specific experimental verification.

Weaknesses

1. As part of its contributions, this article presents a method for identifying nearly complete IOI circuit signals in GPT-2 Small directly through the weight file. However, the instructions provided in lines [148-152] are too brief, making the specific method difficult to understand. 2. It is well known that the matrix of the attention head is sparse, which naturally suggests using Singular Value Decomposition (SVD) to compress the original matrix. Although the primary aim of this article is to identify subspaces, there is no fundamental difference between the two approaches. Numerous related works, such as LoRA, have applied SVD to the original model weights. Moreover, the conclusion that inter-layer communication is low-rank seems evident due to the inherent low rank of the attention head weight matrix.

Questions

1. The primary experiments in this paper were conducted on the IOI and Laundry List Task. Notably, the Laundry List Task has significant limitations in excluding repeated words. Inhibiting repeated words can only be considered a part of the recall task. Is it possible that recall tasks are the ones where repeated words appear most frequently? Some tasks should enhance the inhibitory effect on repeated words, while others should reduce it. Therefore, the experimental tasks in this paper are insufficient. 2. In Figure 2, the Inhibition Score is presented, but it is not introduced until Section 3.3.1, which may confuse readers. The use of \( V \) in Equation 2 may confuse readers with the \( V \) matrix in the Transformer model. There is a spelling error in the caption of Figure 3: "inhiibt" should be "inhibit." 3. The model-editing method proposed in this article involves performing SVD decomposition on a specific head (the inhibition head) and retaining the main subspace components. This approach improves performance on the Laundry List Task, indicating an enhanced inhibitory effect on repeated words. The intervention method in this article primarily demonstrates that other signals unrelated to inhibitory ability exist within the inhibition head, and that dimension reduction can enhance inhibitory capacity. It should be noted that identifying the inhibition head is not a contribution of this article; rather, it relies on the conclusions of previous work. This article merely reinforces that the inhibition head functions as an inhibitory mechanism and that its combination with the duplicate head improves the inhibition of repeated words, a conclusion already established in prior research. Additionally, this article does not propose a new mechanism for inter-layer communication in the Transformer model.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The author has not fully addressed the limitations of their research. While the proposed method improves the model's inhibitory effect on repeated words, it primarily relies on existing findings and does not introduce a new mechanism for inter-layer communication in the Transformer model. Additionally, the experimental tasks used in the study, such as the Laundry List Task, have notable limitations, enhance the inhibitory effect on repeated words, while others maybe should reduce it.

Thank you for your reply

I have read your response, as well as the other reviews. Thank you for the additional information regarding results on larger models and your viewpoint on how the simple recall task may relate to more general language modeling capabilities. I do think the other reviewers have brought up some valid concerns, (e.g. limited related work discussion, generality of the findings), but I think the authors have done an adequate job of trying to address most of them. I think the authors’ proposed changes and clarifications (along with cleaning up of the typos/ and general presentation) will strengthen the paper. As it stands, I think the paper would be an interesting contribution to the conference and would like to keep my score.