Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs

We thank the reviewer for their insightful comments. We respond to the questions as follows:

For example, it would be helpful to confirm whether deactivating attention heads with attention sinks truly has no impact on the final output.

We agree with this point. Figure 6(b) shows the results of deactivating (zeroing out) attention heads with attention sinks. In the Wikipedia domain, the head exhibits a pure attention sink, and deactivating it has minimal impact on the final output.

The paper’s explanation for why Adam tends to cause residual state peaks is insufficient, as it mainly describes the phenomenon without deeper analysis.

The intuition is that the residual state of the extreme token tends to have a smaller gradient as the training dynamics converge. However, Adam makes small gradients to become constant updates, leading to a linear increase in the norm of the residual states. We will provide more details with simulations in the revised version.

Here is a potentially relevant study [1] where the "broadcasting" mechanism mentioned has a formation process similar to the attention sink.

Thanks for sharing. The active-dormant mechanism is ubiquitous in transformers. We will make more discussions of the related mechanism observed in other works in the revised version.

We thank the reviewer for their insightful comments. We respond to the questions as follows:

Section 3, a specific attention head in a specific layer is shown to exhibit behavior similar to the conjectured mechanism from Section 2. However, I find this unconvincing, as demonstrating a single head is insufficient to establish that the active and dormant phases align with what was observed in Section 2 for the same reasons.

We acknowledge the limitation of demonstrating the attention pattern for only one head. However, brute-force searching for attention patterns across active and dormant phases for all heads is computationally impractical. Instead, the training dynamics of OLMo provide alternative evidence for the active-dormant mechanism. Since the dynamics of attention sinks and value state drains consistently follow the mutual reinforcement mechanism predicted by the theory, it is likely that active-dormant mechanism is the driving force of the extreme-token phenomena in LLMs. We believe that dynamics analysis offers a stronger and more efficient method for demonstrating the active-dormant mechanism across all attention heads.

My second concern is the lack of demonstrated practical benefits from this analysis.

We want to emphasize the practical significance of understanding extreme-token phenomena. Numerous studies have shown that these phenomena can cause issues in quantization, interpretability, and long-context inference [Xiao et al., 2023; Bondarenko et al., 2023], leading to the development of mitigation strategies. Our analysis provides insights into two key practical questions: Are extreme-token phenomena beneficial or harmful to the statistical performance of LLM? We found that these phenomena are neither inherently beneficial nor harmful; they represent a dormant phase of attention heads, and their elimination does not affect the statistical performance of transformers. What causes the extreme-token phenomena? By analyzing the model structure and optimizer, we identify SoftMax and ReLU as two key reasons for the extreme-token phenomena. We agree that validating the effectiveness of mitigation strategies is crucial and represents a promising direction for future research. While we do not explore these strategies fully in this work, we believe our work lays a solid foundation for further investigation into this phenomenon.

What is the motivation for using the BB task and designing the trigger and non-trigger tokens?

Finding the minimal condition to recapitulate extreme-token phenomena is crucial for further understanding of them. The BB model can produce all extreme-token phenomena and bear close resemblance with LLMs, providing a nice framework for further studying the phenomena. The trigger token corresponds to the “non-sink” phase of attention heads, and the “non-trigger” tokens correspond to the “sink” phase. The BB task is designed so that a pre-trained attention needs to switch between “non-sink (active)” and “sink (dormant)” given different inputs.

In Claim 4, could you clarify the case where "adding any value state from previous tokens worsens the prediction"?

In the BB model, the transformer needs to memorize the Bigram transition probabilities of non-trigger tokens, which only require the information from the current token. Including previous tokens in this process introduces noise, making memorization more challenging. Consequently, the attention mechanism remains dormant on non-trigger tokens.

Why does the causal mask cause training dynamics to favor the <s> token?

Take the token on the second position as an example. Due to the causal mask, the <s> is the only token that it can sink its attention weights to. For the third, fourth, and following tokens, they all share <s> as the sink token. Therefore, <s> tends to become the extreme token.

So, is the cause of the attention sink truly the causal mask, or is it the lack of semantic meaning in certain tokens?

Our intuition is that both of them contribute to the attention sink. In real LLMs, both semantic meaningless tokens (delimiters) and <s> would be attention sinks. We also observe similar phenomena in the BB model. Results are in Appendix F.

References

Xiao, Guangxuan, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).

Bondarenko, Yelysei, Markus Nagel, and Tijmen Blankevoort. "Quantizable transformers: Removing outliers by helping attention heads do nothing." Advances in Neural Information Processing Systems 36 (2023): 75067-75096.

We thank the reviewer for their insightful comments. We respond to the questions as follows:

Rather it uses these two components in parallel. The reason for this choice is not clear and should be better motivated.

The motivation for the parallel structure comes from the interventions in the 1-layer BB model (Figure 3(a)), which show that the attention and MLP layers have separable utilities. Even when trained with a sequential structure, the attention layer predicts the next token only for trigger tokens, having no impact on non-trigger tokens. Conversely, the MLP layer predicts the next token exclusively for non-trigger tokens. As a result, the outputs of the attention and MLP layers belong to different subspaces. This indicates that the trained sequential structure is functionally equivalent to a parallel structure.

In addition, the trained models use pre-layer norm and it is not clear if residual connections are used and how.

We use residual connections and sequential structure for the trained models. We will add more details in the revised version.

In the derivation of the loss(alpha, beta), it is not clear to me how the stable distribution is defined. Could you clarify?

The stable distribution is defined as the marginal distribution of tokens in the BB task, which means the frequency of each token in the sequence generated by the BB task.

Can you provide more intuition on the role of the softmax?

The SoftMax requires the total attention weights from each token to previous tokens to sum up to 1. However, sometimes the previous token information is unnecessary for the next word prediction. In this situation, more attention weights will concentrate on tokens with small value states, leading to the extreme-token phenomena.

The natural question then is whether this component is actually needed in practice and why modern LLMs use it. Could you elaborate on this?

We agree that understanding whether softmax is essential is an important question, one that warrants a comprehensive analysis in future work. -. Some work has shown that a SoftMax transformer may be stabler in the training and has better performance [Shen et al., 2023; Hua et al., 2022].

First, are the authors now using a standard transformer architecture with attention layers interleaved with mlps? Second, do the authors have any intuition on why at least 3 layers are needed?

For the first question: Yes. For the second: Our intuition is that the “OV” structure in the attention plus the MLP of the first layer is essential to the rise of massive residual states. The third layer MLP, however, removes the massiveness of the residual states so that it doesn’t hurt the read-out layer. We do not include them as we lack appropriate experiments to support the intuition. A more fine-grained analysis is an interesting future direction.

In the construction of the token embeddings, the authors are assuming them to have dimension V

Since the output of the transformer should have dimension $V$, we set the hidden dimension $d=V$ to avoid complications. All results can generalize to arbitrary embedding dimension $d$ if we include the $V*d$ read-out layer in the theoretical model.

it is shown that a head can have sinks yet being non-dormant (due to trigger tokens). Is this case contemplated by the theory?

Yes. Our theory only considers non-trigger tokens and focuses on explaining the sinks. The trigger tokens would have all the attention weights concentrated on the previous token.

How did you practically identify the heads exhibiting sink tokens in real LLMs?

One way to identify the heads which exhibit sink tokens for a particular sample is to first identify which tokens have the potential to become sink/extreme tokens in a particular head. We can do this by examining the residual state peak phenomenon: the residual state norms of the (potential) sink tokens are significantly greater (usually > 10x greater) than others. Then, we examine heads which have large mean attention score on the sink tokens, i.e., compute the average attention score $\frac{1}{n-i}\sum_{j = i}^{n}A_{hij}$ for each token $i$ at head $h$ and examine the attention score on the potential sink tokens identified in the previous step. If these scores are large then the head has sink tokens, otherwise it does not. This is a straightforward and intuitive method to find such heads, but we do not claim it is optimal in any way; we leave improving this mechanism to future work.

References

Shen, Kai, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. "A study on relu and softmax in transformer." arXiv preprint arXiv:2302.06461 (2023).

Hua, Weizhe, Zihang Dai, Hanxiao Liu, and Quoc Le. "Transformer quality in linear time." In International conference on machine learning, pp. 9099-9117. PMLR, 2022.

We thank the reviewer for their insightful comments. We respond to the questions as follows:

The motivation behind the mitigation of extreme-token phenomena is not clear.

Thank you for pointing out the need to clarify the motivation for studying extreme-token phenomena and their mitigation. The primary goal of our paper is not to propose or advocate for mitigation strategies but rather to provide a theoretical understanding of why extreme-token phenomena arise in the first place.

While our paper explores potential interventions like replacing SoftMax with ReLU, these modifications were used primarily as tools to examine the mechanisms behind the phenomena, rather than as definitive solutions to improve downstream performance. Consequently, we have not exhaustively explored the trade-offs associated with these changes. We agree that mitigating extreme-token phenomena would be an important area for future research.

This assumption may oversimplify the interactions in real-world LLMs, where dependencies between heads and layers are often more intricate.

We agree that our theory simplifies many of the complex interactions present in real-world LLMs. However, we consider this simplification a strength rather than a limitation. By distilling the mechanism to its core, our theory isolates and highlights the essence of the mutual reinforcement mechanism that underlies the formation of extreme tokens. Developing a theory within an over-complicated system would risk hiding the key insights and limiting its general applicability.

What is particularly compelling about our framework is that the behavior of the simplified model, which demonstrates the mutual reinforcement mechanism, aligns closely with the training dynamics observed in OLMo, a real-world LLM with a more intricate structure. This connection suggests that the simplified theory successfully captures the foundational dynamics relevant to more complex systems. Additionally, the assumption that “zeroing out attention heads for intervention outcomes will not impact token prediction accuracy” has been empirically validated in several attention heads of Llama-2, further supporting the practical relevance of our theoretical framework.

In particular, would we still find attention heads with simple and interpretable active behavior for weaker contextual differences?

We agree that it would be interesting to discover simple and interpretable active-vs-dormant behavior for more subtle contextual differences. The example we have in the paper is to illustrate the point that the theoretical principles carry over to large language models. In particular, we do not guarantee that every domain difference corresponds to a head with similar behavior to the head in the paper. As such, we leave the discovery of such heads for particular interesting contexts to future work.

In section 2.3, it is mentioned that no residual peaks appear in a one-layer transformer(BB model) or even with different combinations (2xTF, MLP+TF+ML, Attn+TF+MLP); it would be interesting to know why it seems to have residual peaks in certain instances.

We have three intuitions:

1. The First Layer Creates the Massive Residual State: The first attention and MLP layer have large outputs on the $\langle s \rangle$ token, while keeping other residual states to be small.

2. The Middle Layers Have Minimal Effect on Massive Residual States: The intermediate layers appear to contribute little to the magnitude of the residual states.

3. The Final Layer MLP Mitigates the Residual State Peak: The third-layer MLP reduces the residual state of the $\langle s\rangle$ token. Without this adjustment, an excessively large on the $\langle s\rangle$ token causes troubles in the read-out layer.

Although we have not included these intuitions in the main analysis due to a lack of rigorous experimental support, we consider a more fine-grained investigation of these ideas an important future direction.

It would be helpful to see the comparison of residual state peaks and attention sink and value state drain phenomena for the 1-layer BB model even when residual state peaks do not appear

The residual state peaks strengthen the attention sink and the value state drain in upper layers. The intervention results (Figure 5(b)) show that the residual state peaks lead to attention sink and value state drain in upper layers. Further comparing the extreme-token phenomena in one-layer and three-layer transformers, we have that:

The results show that residual state peaks strengthen the extreme-token phenomena in three-layer transformers compared with 1-layer transformers.

I am happy with the overall response and have decided to raise my score.