Attend First, Consolidate Later: On the Importance of Attention in Different LLM Layers

Authors: amit ben artzy
roy  schwartz
Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-15 08:38:23

Abstract: In decoder-based LLMs, the representation of a given layer serves two  purposes: as input to the next layer during the computation of the current token; and as input to the attention mechanism of future tokens.  
In this work, we show that the importance of the latter role might be overestimated. 
To show that, we start by manipulating the representations of previous tokens; e.g. by replacing the hidden states at some layer $k$ with random vectors.
Our experimenting with four LLMs and four tasks show that this operation often leads to small to negligible drop in performance. Importantly, this happens if the manipulation occurs in the top part of the model---$k$ is in the final 30--50\% of the layers. In contrast, doing the same manipulation in earlier layers might lead to chance level performance.
We continue by switching the hidden state of certain tokens with hidden states of other tokens from another prompt; e.g., replacing the word "Italy" with "France" in ``What is the capital of Italy?''. We find that when applying this switch in the top 1/3 of the model, the model ignores it (answering "Rome"). However if we apply it before, the model conforms to the switch ("Paris").

Our results hint at a two stage process in transformer-based LLMs: the first part gathers input from previous tokens, while the second mainly processes that information internally.

Main paper: 
 introduction: 
 The attention mechanism in transformerbased (Vaswani et al., 2017) LLMs allows information to flow from the hidden representation of one token to another. While this process is the same for all model layers, previous work has shown that different layers capture different types of information (Niu et al., 2022; Geva et al., 2020 Geva et al., , 2022;; Press et al., 2019; van Aken et al., 2019) It is therefore not entirely clear that this flow of Figure 1 : To evaluate the role of previous hidden states as input to the attention mechanism, we devise two setups: (a) we replace the hidden state at layer k with a random vector, and use it as input to layer k + 1, which continues processing as normal; (b) starting from a given layer k + 1, the hidden representations of previous tokens are frozen, and the attention mechanism attends to their hidden states at layer k.information is equally important for all layers. Can we find a distinction between layers that aggregate information from previous tokens, and those that process this information internally?To better understand these dynamics, we apply various manipulations to the hidden states of all tokens barring the current one, and evaluate their impact on the model's performance over various tasks. We consider several different manipulations, e.g., replacing the hidden state at layer k with random vectors; and replacing the hidden states of all upper layers (ℓ > k) with those of the last unmanipulated layer (k). We note that none of the manipulations in this work involves further training or fine-tuning.We experiment with four LLMs (Llama2-7B, Touvron et al., 2023; Mistral-7B, Jiang et al., 2023; Yi-6B, Ai et al., 2024; and Llemma-7B, Azerbayev et al., 2023) across four tasks. Our results show that transformers are surprisingly robust to manipulations of their previous tokens. Freezing up to 50% of the layers results in some cases in no loss in performance across multiple tasks. Moreover, replacing up to 30% of the top layers with random vectors also results in little to no decrease. Importantly, we identify a distinct point where LLMs become robust to these manipulations: applying them at that point or later leads to little to moderate drop in performance, while doing it earlier leads to a drastic drop in performance.To further study this phenomenon, we consider a third manipulation: switching the hidden states of certain tokens with a hidden state computed based on a separate prompt. E.g., in factual question answering tasks ("What is the capital of Italy?"), we replace the hidden state of the token "Italy" with that of the token "France" from another prompt. Our results are striking: in accordance with our previous results, doing this manipulation at the top 1/3 of the model leads to no change in prediction. However, doing it earlier leads to the generation of the output corresponding to the change ("Paris" instead of "Rome").Finally, we consider dropping the attention mechanism altogether, by skipping the attention block in all layers starting a given layer k. As before, we find in some cases a high variance in how important attention mechanism is across layers: doing this at the bottom layers leads to severe performance degradation, while doing it at higher layers results in smaller drops, and even matches baseline performance in some tasks.Our results shed light on the way information is processed in transformer LLMs. In particular, they suggest a two-phase process: in the first, the model extracts information from previous tokens. In this phase any change to their hidden representation leads to substantial degradation in performance. In the second phase, information is processed internally, and the representation of previous tokens matters less. They also have potential implications for making transformer LLMs more efficient, by allowing both to skip upper attention layers, and accordingly, reducing the memory load of caching these computations. We publicly released our code.1 2 Manipulated LLM Generation Decoder-only transformers consist of a series of transformer blocks. Each block contains an attention block and a feed-forward block.2 Formally, to generate token n + 1, we process the n'th token by attending to all previous tokens i ≤ n. Formally, for layer ℓ, we define the attention scores A ℓ n as follows:A ℓ n = sof tmax( q ℓ n • K ℓ n √ d ) • V ℓ n where q ℓ n = W ℓ q •x ℓ-1 n ; K ℓ n = W ℓ k •X ℓ-1 1,...,n ; V ℓ n = W ℓ v •X ℓ-1 1,...,nIn this setup, W ℓ q/k/v are weight matrices, d is the hidden size dimension, x ℓ-1 n is the representation of the current token in the previous layer, and X ℓ-1 1,...,n is a matrix of the hidden representation of all tokens from the previous layer.We highlight two important properties of transformers. First, the K ℓ n and V ℓ n matrices are the only components in the transformer block that observe the previous tokens in the document. Second, all transformer layers are defined exactly the same, as described above. In this work we suggest that perhaps this uniformity across layers needs to be reconsidered.We aim to ask how sensitive a model is to observing the exact history tokens, or in other wordshow much will observing manipulated versions of them impact it. Below we describe the manipulations we employ. We stress that all the manipulations in this work operate on the pre-trained model, and do not include any training or fine-tuning. See Fig. 1 for visualization of the different approaches. 
noise: 
 First, we ask whether the content of the hidden state even matters. To do so, we replace the hidden states at layer k of all previous tokens (X k 1,...,n-1 ) with random vectors. The next layer (k + 1) gets these random vectors as input, and the following layers continue as normal. We use two policies for introducing noise to the history of tokens, both ensuring the noisy hidden states have the same norm as the original hidden states: shuffle, where we take the original hidden state and shuffle its indexes in a random permutation; and random where we randomize a vector of variance 1 and mean 0 then re-scale it such that the norm would be the same as the original hidden state.If the model is successful in this setup at some layer, we may conclude that it has already gained the majority of the relevant information from previous layers, and in practice it is not making excessive use of this information in higher layers.We next turn to address the question of whether deep processing of the previous tokens is needed. To do so, we freeze the model at layer k and copy the hidden state at that layer to all subsequent layers. 3 Formally, for ℓ > k, we set X ℓ 1,...,n-1 = X k 1,...,n-1 . If this manipulation would result in large performance degradation, we may conclude that subsequent processing in higher layers is important. Alternatively, a minor to no drop in performance would indicate that perhaps the processing up to layer k is sufficient. 
datasets: 
 We aim to understand the flow of information in various test cases. For this purpose, we consider four benchmarks described below: two that we curate, which allow us to perform nuanced, meticulous manipulations; and two other benchmarks for standard tasks-question answering and summarization.Capitals We devise a simple fact-extraction QA dataset, which consists of 194 country-capital pairs. The dataset is in the format of "What is the capital of X?". To align the model to return the capital only, rather than a full sentence such as "The capital of X is Y", we use a 1-shot setting. We report exact match accuracy.We compile a second dataset, consisting of simple math exercises of addition and subtraction of single digit numbers. We consider two variants: 2-term (i.e., the subtraction/addition of two numbers, e.g., "1 + 2 ="), and 3-term (e.g., "1+2+3 ="). In both cases, we only consider cases where the answer is also a single (non-negative) digit. In 3-term, we also verify that each mid-step can be represented as a single token. This results in 110 2-term instances, and 1210 3-term instances.This dataset has a few desired properties. First, it has a clear answer, which facilitates evaluation. Second, perhaps surprisingly, it is not trivial-the math-tailored LLM we experiment with for this task (Llemma) only reaches ≈80% on 2-term and ≈50% on 3-term. Third, it allows us to easily increase the level of difficulty of the problem (by considering 2-term vs. 3-term). We report exact match accuracy.SQuAD We also consider the Stanford Question Answering Dataset (SQuAD; Rajpurkar et al., 2016) , a dataset consisting of question-paragraph pairs, where one of the sentences in the paragraph contains the answer to the corresponding question. The task is to correctly output the segment that contains the answer. We sample 1,000 instances from the SQuAD test set and report exact match.CNN/Daily Mail This dataset (Hermann et al., 2015) contains news articles from CNN and the Daily Mail. The task is to generate a summary of these articles. We sample 100 instances from the CNN/Daily Mail test set and report averaged rouge-1 and rouge-l scores (Lin, 2004; Papineni et al., 2002) .For the textual tasks (Capitals, SQuAD, and CNN/Daily Mail), we experiment with three opensource, decoder-only models, each containing 32 layers: Llama2-7B (Touvron et al., 2023) , Mistral-7B (Jiang et al., 2023) , and Yi-6B (Ai et al., 2024) .For the math exercises dataset, we observe that these models perform strikingly low, so we experiment with Llemma-7B (Azerbayev et al., 2023) , a Code Llama (Rozière et al., 2023) based model finetuned for math. 
results: 
 Capitals Figure 2 shows the results of our manipulations on the capitals dataset. We first note that all LLMs are surprisingly robust to the different manipulations. When freezing the top ≈50% of the model, all models reach similar performance as the baseline (unmanipulated model). For the noise manipulations, we observe the same trend, though at later layers: For both manipulations, the model matches the baseline performance if applied at the top ≈30% of the model.We also note that, interestingly, in almost all cases we observe a critical layer k ′ , for which the model performs almost at chance level if manipulated in layers i ≤ k ′ , while substantially improving if applied afterwards. While this point varies between models and manipulation types, e.g., from layer 8 (Llama2-7B, freeze) to 25 (Yi, shuffle), the phenomenon in general is quite robust.Our results hint that LLMs exhibit a two-phase processing: the first part gathers information from previous tokens. At this part the content of previous hidden states is highly important. In contrast, at the second part, the model mostly consolidates this information, and is far less sensitive to such manipulations.We next consider the math exercises dataset (Fig. 3 ). Again, we observe that freezing the top 1/2 of the model results in a similar performance to the baseline in both setups (2-term and 3-term). However, here the shuffle and random manipulations perform similar to the baseline only when applied at the top 10% of the model. Here we also observe that we do not see a clear transition point from chance level performance to baseline-level, but rather a more steady increase. Interestingly, the trends in the 2-term and 3-term settings are similar; despite the fact that the 2-term problems are substantially easier to the model.We now turn to consider common NLP tasks, and start with SQuAD (Fig. 4 ). We first observe for two of the three LLMs, applying the freeze manipulation leads to the same trend as before: comparable (or even better!) results as the baseline. This happens starting layer 20 (Llama2-7B) or 25 (Yi). In contrast, for Mistral, the manipulated model only matches the performance of the full model after 30 (of 32) layers.Considering the two noise manipulations, we observe that for Llama2-7B, both variants also reach the baseline performance after 25-27 layers. However, the other two models never fully reach it. Nonetheless, we note that for these models, we clearly observe the transition point observed in the capitals experiments: Between 15-23 layers, model performance is at chance level, and afterwards it dramatically improves. These results further support the two-phase setup.CNN / Daily Mail Finally, we consider the CNN/ Daily Mail dataset (Fig. 5 ) Yi-6B matches baseline performance at top layers across all manipulations. As in SQuAD, using the freeze manipulation, Llama2-7B performs similar or even slightly better than the baseline if applied in the final ≈30% of the layers. Results for Mistral-7B with the freeze manipulations are close, though clearly inferior to the baseline. In contrast, the two noise manipulations lead to substantially lower scores in Llama2-7B and Mistral-7B.Discussion Our results demonstrate a few interesting trends. First, we observe that in almost all cases, models are robust to freezing, in some cases as early as after 50% of the layers. We also note that in some cases (capitals, SQuAD with Llama2-7B) LLMs are robust to adding noise, but in general this does lead to noticeable performance degradation. Nonetheless, we still observe a rather consistent phenomenon with these manipulation, which shows that applying them too early leads to chancelevel performance, and at a certain layer, results suddenly improve dramatically (albeit not reaching the baseline performance).Our results suggest two-step phase in the processing of LLMs: a first step that gathers information from previous tokens, and a second that consolidates it. We next turn to further explore this hypothesis, by presenting two additional manipulations-replacing the hidden representation with that of another token from a different prompt; and skipping the attention mechanism altogether. 
injecting information from a different prompt: 
 To further test the two-phase hypothesis, we study the impact of "injecting" new information to the model, in the form of replacing the hidden representation of a given token with that of another token from a different prompt. 4 We experimented with the capitals dataset. For example, given the question "What is the capital of Italy?", we replace the hidden states corresponding to the word "Italy" in layer ℓ with the hidden states corresponding to "France" at layer ℓ from another prompt. We randomized pairs of countries which are represented with the same number of tokens.Our results are shown in Fig. 6 . For each model, we identify a clear transition point: before it, the model answer conforms to the patched value (e.g., "Paris" in the example above), and afterwards the model is unimpacted by the manipulation, returning the original answer ("Rome"). These results further illustrate the two phases we observed in previous experiments.6 Is Attention Needed at Top Layers?Our results so far indicate that the role of previous tokens is far more important in the bottom layers of the model than in top ones. A question that arises now is whether the attention mechanism is even needed in those top layers.To study this question, we experiment with skipping the attention block in those layers, and only applying the feed-forward sub-layer. 5 We find that in three of the four datasets (capitals, Fig. 2 ; SQuAD, Fig. 4 ; and CNN / Daily Mail, Fig. 5 ), the effect of this process is similar to that of the shuffle manipulation. I.e., in some (though not all) cases the models are surprisingly robust to this process.In contrast, and perhaps surprisingly, we find that skipping the attention block in the math exercises dataset leads to chance-level performance in all cases (Fig. 3 ). This might indicate the nature of this problem, where each and every token is critical to give the final answer, forces the model to use the attention mechanism all the way through. 
related work: 
 To better understand the inner-working transformers, previous work has explored the roles of the different transformer layers. Niu et al. (2022) Figure 6 : The effect of injecting information from a different prompt (E.g., replacing "Italy" with "France" in "What is the capital of Italy?"). All models exhibit a clear two-phase behavior: in the first part, the injection changes the output to be the modified answer (e.g., "Paris"). In the second, the model is unaffected by the injection (answering "Rome").found that linguistics features can be extracted from hidden representations. Geva et al. (2020 Geva et al. ( , 2022) ) demonstrated that lower layers are associated with shallow patterns while higher layers are associated with semantic ones. Press et al. (2019) and Mandava et al. (2020) have trained from scratch transformers with different sub-layer ordering and found some variants that outperform the default ordering. Related to our work, they observed that orederings that with more attention layers at the bottom half and more feed-forward at the top tended to perform better, hinting that the attention is more important at the bottom layers.Early exit methods (Schwartz et al., 2020; Xin et al., 2020; Schuster et al., 2022; Elhoushi et al., 2024) , which speed up LLMs by processing only the bottom part of the model, also provide evidence that the top layers of the model have already gained relevant information from previous tokens.Prior work tried to better understand the information encapsulated in hidden representations. Hendel et al. (2023) demonstrated that patching an operator (e.g., the "→" token) from in-context tasks to another context preserves the operation. Ghandeharioun et al. (2024) showed that patching can be seen as a generalization of various prior interpretability methods and demonstrated how this method can be used in other cases, e.g., feature extraction. Both of these works aimed to study the hidden representations in transformers and the features they encode. In contrast, we propose to patch different vectors into some context to learn more about the flow of information between the tokens in context.Previous work also experimented with manipulating LLMs. Meng et al. (2022) corrupted hidden states to identify and edit specific neurons responsible to specific outputs. The concurrent work of Lad et al. (2024) investigated swapping layers to better understand the process of token generating, and hypothesized about stages of inference. Gromov et al. (2024) explored deleting layers for pruning purposes. We, however, utilized different manipulations to better understand the flow of information in LLMs. 
conclusion: 
 We investigated the role of the attention mechanism across a range of layers. We applied various manipulations over the hidden states of previous tokens, and showed that their impact is far less pronounced when applied to the top 30-50% of the model. Moreover, we switched the hidden states of specific tokens with hidden states of other tokens from another prompt. We found that there is a distinct point, at the top 1/3 of the model, where before it the model conforms to the switch, and afterwards it ignores it, answering the original question. Finally, we experimented with dropping the attention component altogether starting from a given layer. We found again, that in some cases (though not all), doing this at the top 30% of the model leads to a small effect, while much larger earlier.Our results shed light on the inner workings of transformer LLMs, by hinting at a two-phase setup of their text generation process: first, they aggregate information from previous tokens, and then they decipher the meaning and generate new token. Our work could further be potentially extended to reduce LLMs costs: First by skipping the attention component in upper layer, and second by alleviating the need to cache their output for future generations. 
 
 Appendix: github.com/schwartz-lab-NLP/ Attend-First-Consolidate-LaterWe omit normalization and residuals for brevity.This is common practice in early-exit setups(Schuster et al., 2022), where some tokens require deeper processing than the previous ones that performed an early-exit.This process is often called "patching"(Hendel et al., 2023;Ghandeharioun et al., 2024).Following preliminary experiments, we also skip the normalization prior to the attention. 


Decision: While the reviewers all agree that the paper presents an exciting setup with interesting findings, a large amount of missing related work is pointed out that the authors should include in a revised version of the paper. Since the topic that is tackled by the paper has been covered in earlier work as well, it is highly important the authors make clear how their work differentiates from these papers.

Reviews: review: paper_summary: paper_summary: The paper aims to shed light on which components of large decoder Transformer models are necessary for next-token prediction.
It does so by 1) freezing or replacing hidden states with noise (by shuffling/generating random ones) of all previous tokens in certain layers or 2) skipping the attention mechanism in a certain layer.
These methods are applied to multiple models applied to two question answering tasks, a summarization task, and a simple arithmetic task.

For the majority of the tasks and models investigated, the drop in performance after these edits is small to negligible in the final ~30% of layers, indicating that the models have acquired all task relevant information from earlier tokens in early layers, and do not need to pay as much attention to these historical tokens in the later layers to predict the correct next token.
All models are more robust to freezing hidden states than to randomizing hidden states. Sometimes, freezing the hidden states even leads to better predictions than letting the model run in full. 
For the mathematical arithmetic task, the model is not as robust to randomization of the hidden states, and not robust at all to skipping the attention mechanism. 

Generally, the steep increase in accuracy on natural language tasks after a certain layer suggests that these decoder Transformer models operate in two phases to solve these tasks: first acquiring relevant information from previous tokens, and then after that layer, consolidating this information internally. There is no such steep increase in accuracy for the mathematical task.
strength_weakness: strength_weakness: It was not entirely clear to me whether the "freeze" approach is actually equivalent to early exiting (This seems to be implied in line 388 when comparing to Schuster 2022) or freezing *only* the hidden states, but not the attention mechanism in the later layers (Which seems to be implied in section L165).
I think this is an important distinction to make, and might be interesting to compare: if early exiting, which is more widely studied, works just as well -- what is there to gain by only skipping the attention components / only skipping the hidden state components of later layers? In other words: how much 'internal processing' or consolidation is necessary? The paper currently does not answer this question.


The paper's engagement with other related interpretability work is somewhat limited:
- There is literature on early exiting/early decoding and interpretability, such as the LogitLens (nostalgebraist 2020) and papers using LogitLens to interpret intermediate layers (e.g. Merullo et al. 2023, which also experiments with a capital city prediction task throughout the layers, Halawi et al. 2023, which also notes that intermediate layers can achieve better accuracy in some cases).
- There is also previous work on intervening on early hidden states by replacing them with those of another input (eg. 'interchange interventions' for encoder Transformers, Geiger et al. 2021).
- The only future work suggestion listed is for improving efficiency, but what could these findings mean for future *interpretability* work?

Because the methods of replacing/freezing states are quite general, and because of the two weaknesses listed above, the paper and its findings may not be novel: a similar paper with more interaction with previous work would be preferable. My assessment of the paper is therefore somewhere between ambivalent and worthy. Still, I have rounded my overall assessment score up to a 4, since the paper is clear, well written, and leads to some interesting general questions.
questions: questions: Suggestions/questions:
- As mentioned above, it would be interesting to see the difference between immediate early exiting and freezing only the hidden states but not the attention. 
- The model seems to be able to recover from replacing the hidden states past a certain layer. Is it a possibility that after that layer, the model is able to ignore noise and recover the previous hidden states from its 'residual stream'? In this case, the model may still be needing previous tokens in those later layers. Would replacing/randomizing the hidden states of all layers l > k for some k (similar to the freezing technique) be a viable experiment?
- What could specifically be happening around layer 28 of Mistral?

Minor suggestions/clarity:
- Section 5: It would be good to mention whether this experiment is performed on all pairs of capitals.
- L060: 'and two additional tasks we compile', it would be clearer to simply list the tasks or at least what the tasks are roughly about.
paper_summary: paper_summary: The paper studies the importance of layers when modifying internal representations of the LLMs. The study uses four LLMs (Llama2-7B, Mistral-7B, Yi-6B, and Llemma-7B) across four tasks: a custom capitals dataset, math exercises, SQuAD, and CNN/Daily Mail summarization. By using two types of perturbations to the previous tokens -replacing the original representations by random vectors, and freezing the intermediate representations by copying the same hidden state as input to all subsequent layers- the authors demonstrate models are highly robust to these changes in the top 30-50% of the layers, suggesting a two-phase processing stage, a first step that gathers information from previous tokens, and a second that consolidates it. To further test this hypothesis, two more experiments are proposed, performing activation patching on the residual stream, and skipping attention in in upper layers, finding supporting evidence.
strength_weakness: strength_weakness: - The different stages in the forward-pass, and the apparently ineffectiveness of later layers have been proposed in previous work (Vedang et al., 204, Gromov et al., 2024, Elhousi et al., 2024), making the findings in this paper less novel. However, here the focus is on the context utilization.

- I missed a more detailed analysis into what the final layers are doing, and in which cases they are relevant.


[Gromov et al., 2024, The Unreasonable Ineffectiveness of the Deeper Layers](https://arxiv.org/abs/2403.17887)

[Vedang et al., 2024, The Remarkable Robustness of LLMs: Stages of Inference?](https://arxiv.org/abs/2406.19384)

[Elhousi et al., 2024, LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding](https://arxiv.org/abs/2404.16710)
questions: questions: - Missing citations

Stages of inference:

[Vedang et al., 2024, The Remarkable Robustness of LLMs: Stages of Inference?](https://arxiv.org/abs/2406.19384)

Referring to activation pathing:

[Geiger et al., 2020, Neural Natural Language Inference Models Partially Embed Theories of Lexical Entailment and Negation](https://arxiv.org/abs/2004.14623)

[Vig et al., 2020, Investigating Gender Bias in Language Models Using Causal Mediation Analysis](https://proceedings.neurips.cc/paper/2020/hash/92650b2e92217715fe312e6fa7b90d82-Abstract.html)

Typos

- Figured are referenced in the text just with their number, instead of "Figure X".
- Figure 2, caption: "... observe that the all" -> "... observe that all"
- line 278 "of even better" -> "or even better"
paper_summary: paper_summary: In this paper, the input to the attention mechanism at different layers is manipulated: previous hidden states are replaced with random vectors, "frozen" hidden representations from previous layers, or representations of counterfactual tokens. In experiments across 3 textual tasks and with 3 models, as well as a math task with two models, authors find that models are relatively robust to this noise in higher layers, while performing at chance level if corruption happens in earlier layers. The paper concludes that the results are evidence of a two-stage process in Transformers, with a first phase where information is gathered from previous tokens, and a second phase where representations are mainly refined internally. These findings could motivate pruning of attention heads in higher layers.
strength_weakness: strength_weakness: - There are two relevant strands of research that are not discussed in related work, and also reduce the novelty of findings.
  - One is research on pruning attention heads (e.g. Michel et al. 2019, Voita et al. 2019), which has already shown that many Transformer attention heads have little relevance on a prediction. This work argues that attention is more important in earlier layers than later ones (mirroring Press et al. 2019).
  - previous interpretability work, for example the Causal Mediation Analysis as used in Meng et al. (2022), works with a similar principle (applying corruption at some point and measuring the effect), and has been used to arrive at similar conclusions (specific hidden states can be decisive for the prediction, and we can localize and edit these decisive states to manipulate the prediction).
- The analysis in this paper is relatively coarse-grained, especially in relation to work mentioned in the previous point. For example, this paper corrupts/skips all attention in top layers, and concludes that attention in top layers is less important and may be skippable, but this conclusion seems to broad. For example, Mistral-7B is sensitive to corruptions in layer 27, and Llemma is sensitive to corruption in all layers. I'd suggest moving to a more fine-grained analysis, but this will likely make the work more similar to Meng et al. (2022).
questions: questions: citations mentioned above:

@inproceedings{NEURIPS2019_2c601ad9,
 author = {Michel, Paul and Levy, Omer and Neubig, Graham},
 booktitle = {Advances in Neural Information Processing Systems},
 editor = {H. Wallach and H. Larochelle and A. Beygelzimer and F. d\textquotesingle Alch\'{e}-Buc and E. Fox and R. Garnett},
 pages = {},
 publisher = {Curran Associates, Inc.},
 title = {Are Sixteen Heads Really Better than One?},
 url = {https://proceedings.neurips.cc/paper_files/paper/2019/file/2c601ad9d2ff9bc8b282670cdd54f69f-Paper.pdf},
 volume = {32},
 year = {2019}
}

@inproceedings{voita-etal-2019-analyzing,
    title = "Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned",
    author = "Voita, Elena  and
      Talbot, David  and
      Moiseev, Fedor  and
      Sennrich, Rico  and
      Titov, Ivan",
    editor = "Korhonen, Anna  and
      Traum, David  and
      M{\`a}rquez, Llu{\'\i}s",
    booktitle = "Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics",
    month = jul,
    year = "2019",
    address = "Florence, Italy",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/P19-1580",
    doi = "10.18653/v1/P19-1580",
    pages = "5797--5808",
}



@article{meng2022locating,
  title={Locating and Editing Factual Associations in {GPT}},
  author={Kevin Meng and David Bau and Alex Andonian and Yonatan Belinkov},
  journal={Advances in Neural Information Processing Systems},
  volume={36},
  year={2022},
  note={arXiv:2202.05262}
}