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

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.

## introduction the attention mechanism in transformerbased @xcite 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 @xcite @xcite @xcite it is therefore not entirely clear that this flow of (a) injecting random vectors (b) freeze generation 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. 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, @xcite @xcite 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. foot_0 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 @xmath0. formally, for layer ℓ, we define the attention scores a ℓ n as follows: in 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. ## 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 @xmath1, 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. ## 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. 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. @xcite 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"). early exit methods @xcite @xcite , 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. @xcite demonstrated that patching an operator (e.g., the "→" token) from in-context tasks to another context preserves the operation. @xcite 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 ## 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.