Library

On the Role of Attention Masks and LayerNorm in Transformers

Self-attention is the key mechanism of transformers, which are the essential building blocks of modern foundation models. Recent studies have shown that pure self-attention suffers from an increasing degree of rank collapse as depth increases, limiting model expressivity and further utilization of model depth. The existing literature on rank collapse, however, has mostly overlooked other critical components in transformers that may alleviate the rank collapse issue. In this paper, we provide a general analysis of rank collapse under self-attention, taking into account the effects of attention masks and layer normalization (LayerNorm). In particular, we find that although pure masked attention still suffers from exponential collapse to a rank one subspace, sparse or local masked attention can provably slow down the collapse rate. In the case of self-attention with LayerNorm, we first show that for certain classes of value matrices, collapse to a rank one subspace still happens exponentially. However, through construction of nontrivial counterexamples, we then establish that with proper choice of value matrices, a general class of sequences may not converge to a rank one subspace, and the self-attention dynamics with LayerNorm can simultaneously possess a rich set of equilibria with any possible rank between one and full. Our result refutes the previous hypothesis that LayerNorm plays no role in the rank collapse of self-attention and suggests that self-attention with LayerNorm constitutes a much more expressive, versatile nonlinear dynamical system than what was originally thought.

cs.LGstat.MLEngineeringComputer ScienceMathematics
View graphSource

Paper

References (41)

01TransformerFAM: Feedback attention is working memoryDongseong Hwang, Weiran Wang, Zhuoyuan Huo et al.2024 · arXiv.org · 19 citations In Library
02Anisotropy Is Inherent to Self-Attention in TransformersNathan Godey, Eric Villemonte de la Clergerie, Benoît Sagot2024 · Conference of the European Chapter of the Association for Computational Linguistics · 69 citations In Library
03A mathematical perspective on TransformersBorjan Geshkovski, Cyril Letrouit, Yury Polyanskiy et al.2023 · Bulletin of the American Mathematical Society · 190 citations In Library
04Efficient Streaming Language Models with Attention SinksGuangxuan Xiao, Yuandong Tian, Beidi Chen et al.2023 · International Conference on Learning Representations · 2.2k citations In Library
05On the impact of activation and normalization in obtaining isometric embeddings at initializationA. Joudaki, Hadi Daneshmand, F. Bach2023 · Neural Information Processing Systems · 16 citations In Library
06Demystifying Oversmoothing in Attention-Based Graph Neural NetworksXinyi Wu, A. Ajorlou, Zihui Wu et al.2023 · Neural Information Processing Systems · 82 citations In Library
07Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer TransformerYuandong Tian, Yiping Wang, Beidi Chen et al.2023 · Neural Information Processing Systems · 121 citations In Library
08The emergence of clusters in self-attention dynamicsBorjan Geshkovski, Cyril Letrouit, Yury Polyanskiy et al.2023 · Neural Information Processing Systems · 144 citations In Library
09On the Expressivity Role of LayerNorm in Transformers' AttentionShaked Brody, Uri Alon, Eran Yahav2023 · Annual Meeting of the Association for Computational Linguistics · 55 citations In Library
10Scaling Transformer to 1M tokens and beyond with RMTA. Bulatov, Yuri Kuratov, M. Burtsev2023 · arXiv.org · 117 citations In Library
11Slide-Transformer: Hierarchical Vision Transformer with Local Self-AttentionXuran Pan, Tianzhu Ye, Zhuofan Xia et al.2023 · Computer Vision and Pattern Recognition · 117 citations In Library
12Fast Attention Requires Bounded EntriesJosh Alman, Zhao Song2023 · Neural Information Processing Systems · 112 citations In Library

Scroll for more · 29 remaining

Similar papers

Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention LayersThiziri Nait Saada, Alireza Naderi, Jared Tanner2024 · arXiv (Cornell University) In Library
Rank, Head-Channel Non-Identifiability, and Symmetry Breaking: A Precise Analysis of Representational Collapse in TransformersG. Cirrincione2026 · 0 citations In Library
Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank CollapseLorenzo Noci, Sotiris Anagnostidis, Luca Biggio et al.2022 · arXiv (Cornell University) · 138 citations In Library
On the Expressivity Role of LayerNorm in Transformers' AttentionShaked Brody, Uri Alon, Eran Yahav2023 · arXiv (Cornell University) · 55 citations In Library
How Smooth Is Attention?Valérie Castin, Pierre Ablin, Gabriel Peyré2023 · arXiv (Cornell University) · 41 citations In Library
Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with DepthYihe Dong, Jean-Baptiste Cordonnier, Andreas Loukas2021 · arXiv (Cornell University) · 569 citations In Library
Noise Accumulation and Rank Collapse in Dense Self-Attention: A Theoretical FrameworkLeonardo Cofone2026 · Zenodo (CERN European Organization for Nuclear Research) In Library
A Primal-Dual Framework for Transformers and Neural NetworksTan M. Nguyen, Tam W. Nguyen, Nhat Ho et al.2024 · arXiv (Cornell University) · 19 citations In Library
Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators?Tokio Kajitsuka, Issei Sato2023 · arXiv (Cornell University) · 41 citations In Library
The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer trainingM. Saponati, Pascal Sager, Pau Vilimelis Aceituno et al.2025 · International Conference on Machine Learning · 9 citations In Library

Peer review

Accept (poster)NeurIPS 2024 · via OpenReview
Reviewer ivrv4/10 · confidence 4/52024-07-05

Summary

- This paper investigates the role of attention masks and layer normalization (LayerNorm) in mitigating the rank collapse issue in transformer models, and gets following conclusions: - A long as there is a token which all other tokens in the sequence can directly or indirectly attend to over a fixed number of layers, exponential rank collapse is guaranteed. - local attention mask or focus attention can slow down the collapse. - layernorm plays an important role in self-attention dynamics. - This paper provides extensive analysis and numerical experiments on the self attention dynamics, Expressivity and Versatility.

Strengths

- The paper provides a rigorous mathematical analysis of the role of attention masks and LayerNorm in transformers, contributing to a deeper understanding of these models' inner workings. - This paper uses a graph-theoretic approach to analyze self-attention, so it is more general and can be extended to sparse attention and causal attention. - This paper shows an interesting counterexample of collapse and shows that self-attention dynamics can process a rich set of sequences stablely. - The paper supports its theoretical findings with numerical experiments.

Weaknesses

- Application of the conclusion, I wonder can the findings in this paper guide some design choices in the future transformer models. - And how the findings connect to (or explain) the strong performances of of existing LLMs? It claims that sparse attention can slow down collapse, why most existing LLMs do not use sparse attention? - The numerical experiments are somewhat weak as the results are obtained from only 32 samples.

Questions

I wonder what is the differences of self-attention dynamics between the causal attention and bidirectional attention? I cannot find clear analysis about this in the paper, since different self-attention methods are generalized with graph connections.

Rating

4

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

None

Reviewer 7spk6/10 · confidence 4/52024-07-07

Summary

This paper theoretically studies the role of attention masks and layer norm in the convergence to the rank collapse degeneracy in Transformers, which are two architectural components that have previously been overlooked in studying the rank collapse phenomenon. The authors first define the problem through the lens of graph theory. They then show that self-attention with general mask structure leads to rank collapse, but that factors such as the “graph diameter” of the mask (related to locality of attention) and the uniformity of attention (related to the temperature). They continue to show that LayerNorm does not prevent rank collapse for general attention masks with orthogonal value matrices, but that counterexamples exist for other choices of value matrices.

Strengths

1. The paper is mostly well written and clear. 2. The paper provides new insights into the rank collapse phenomenon. In particular, the graph theoretic formulations provides general conditions in terms of attention mask (e.g. giving a quasi-strongly connected graph) that lead to rank collapse. 3. Moreover, the paper shows a more nuanced analysis for the collapse possibilities when a more complete architecture is considered (in particular with layer normalisation) . 4. The paper also shows the impact of hyperparams like attention context length or temperature in causing/preventing rank collapse.

Weaknesses

My concerns are presented in order of importance (to me). 1. My main concern is that the theoretical setting considered omits several other important architectural components which are used in practice. In particular skip connections (which are known to be effective in preventing rank collapse properties a lot, e.g. Dong et al 2020, Noci et al 2022 or He et al 2023), or positional encoders. The paper is motivated by the fact that rank collapse papers do not study standard archs with all components included, but the omission of such architectural components (especially skips) also means one could argue the paper does not meet its goal. 2. The experimental verification of the theory seems a bit unconvincing at the moment. For example, the effect of temperature seems to disappear after 5 or 6 layers in the non-sliding window settings, and there do not seem to be error bars in figs 2/3 which makes one wonder if the results are significant (are the error bars just smaller than what is visible?). Moreover, it would be good to consider depths of larger than 10, as modern deep transformers can be up to ~100 layers deep and the effects should be more pronounced at larger depths. Also, there doesn't seem to be a difference between the top row of Figure 2 vs the bottom. Finally, how do you pretrain a SAN (figure 2 bottom left) without skip connections as the rank collapse literature would tell you that such models are not trainable (Noci et al 2022). 3. Somewhat related to my first point, it is not clear in my reading that the results presented concerning the counterexamples with particular value weights provide new insights into the fundamental behaviours of transformers in practice or rather represent theoretical edge cases which are nice to know exist but never occur in practice. Some questions that could address this: do such value weights (or value weights with these properties) appear in trained Transformers, or do these theoretical insights generate new initialisations/parameterisations for value weights to train transformers, given that orthogonal weights will lead to rank collapse even with layernorm (Theorem 2). 4. I would also argue that the intuition for the given counterexamples is not particularly clear at present (though as I say most other parts of the paper are well presented). As I understand it, the mathematical argument is that the "center node" has a neuron that has 0 activation due to the layernorm, and as a result even when attention allows other tokens to see the center node's activations the zero activations are not being transfered to other tokens and this prevents collapse. This can be clearer e.g. in section 4.3.1, but as I say I have reservations about how fundamental/important this insight could be.

Questions

1. For the results without LayerNorm (e.g. Theorem 1 or figure 2 left plots), it seems like the definition of mu(X) doesn't rule out that simply X converges to 0 in Frobenius norm, so that the activations simply go to 0, as opposed to saying something about the different tokens becoming aligned (in terms of cosine similarity going to 1) which is what some previous works have described as rank collapse. Do different attention masks lead to this trivial case (as opposed to the alternative where X is non-zero and all the tokens are identical non-zero)? 2. Does theorem 2/corollary 1 hold for non-orthogonal value weights? For example other random matrices like iid Gaussian/uniform should all give the same cosine similarity properties as orthogonal in the large-width limit. 3. The LayerNorm without centering on line 128 is exactly RMSNorm https://arxiv.org/abs/1910.07467 which is popularly used in LLMs like LLama or Mistral instead of LayerNorm. 4. Why does this work only provide exponential convergence rates whereas the original rank collapse paper has doubly exponential? Typo: 1. $d_N$ not $d_d$ in line 128.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

There is not a clear limitations section but the authors do mention a limitation on line 283.

Authorsrebuttal2024-08-07

Answers to the questions raised by Reviewer 7spk

**Questions** > Q1: The definition of mu(X) doesn't rule out that simply X converges to 0 in Frobenius norm, so that the activations simply go to 0, as opposed to saying something about the different tokens becoming aligned (in terms of cosine similarity going to 1) which is what some previous works have described as rank collapse. Do different attention masks lead to this trivial case (as opposed to the alternative where X is non-zero and all the tokens are identical non-zero)? Thank you for the question. - In the case with LayerNorm, X would not converge to 0 and hence the convergence of $\mu(X)$ to zero implies convergence to the same point on the unit sphere. In this case, it is equivalent to the cosine similarity going to one. - Without LayerNorm, whether or not $X$ goes to zero as $\mu$ goes to zero depends on many factors, including the mask, the initial inputs, and the model parameters. > Q2: Does Theorem 2/Corollary 1 hold for non-orthogonal value weights? For example other random matrices like iid Gaussian/uniform should all give the same cosine similarity properties as orthogonal in the large-width limit. This is a good point. We believe that Theorem 2 and Corollary 1 could be extended to more general classes of matrices, such as families of random matrices, as the reviewer mentioned. We leave this for future work. > Q3: The LayerNorm without centering on line 128 is exactly RMSNorm which is popularly used in LLMs like LLama or Mistral instead of LayerNorm. Thank you for the reference! This is a good call. We were aware of RMSNorm but instead chose the name “LayerNorm” following the convention of previous works such as Geshkovski et al. 2023 and Tian et al. 2023 to avoid confusion of terminology. We will add a footnote for this point. > Q4: Why does this work only provide exponential convergence rates whereas the original rank collapse paper has doubly exponential? This is a good question. Our result here is actually tight because in general for linear systems (which is the special case of $W^{(t)}_K=W^{(t)}_Q=0$, $W_V^{(t)}$ being the identity matrix and no LayerNorm), one can never have a doubly exponential convergence (see Antsaklis and Michel, Section 4.8). For this reason, we believe that the results in Dong et al. 2020 are specific to their setting and would not trivially extend to other settings, including ours. We appreciate your questions and comments very much. Please let us know if there are any further questions. ---------------------------------- **References** Dong et al. Attention is not all you need: pure attention loses rank doubly exponentially with depth. In ICML, 2020. Geshkovski et al. A mathematical perspective on transformers. ArXiv, abs/2312.10794, 2023. Tian et al. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In NeurIPS, 2023. Antsaklis and Michel. A Linear Systems Primer. 2000.

Reviewer 7spk2024-08-08

Additional Questions

Thank you for the rebuttal and clarifications. I have some additional questions/comments in light of the response: 1. Doesn't the SANs+LN at initialisation plot (in Figure 1 and Figure 2 of additional pdf) contradict Theorem 2, which states that even with LN and with orthgonal initialisation you should obtain rank collapse? How are the QKV weights initialised in the plots? The mechanism outlined in W4 for the counterexample, while interesting, seems to rely on specific properties of the value weights/representation X which won't occur at random initialisation. 2. Is it necessary to have certain neurons in the activations to have zero values in order to construct the counterexample? If not, it feels like an unnecessary detail which complicates the intuitive picture. The intuition of W4 should be included in the main paper to help readers. 3. Regarding W1, a large body of existing work has shown that the combination of Pre-Norm skip connections reduces the effect of the residual branch which mitigates signal propagation degeneracies like rank collapse. Perhaps the main citation for this is https://arxiv.org/abs/2002.10444, but see also https://arxiv.org/abs/2010.12859 https://arxiv.org/abs/2003.04887 or https://arxiv.org/abs/2102.06171 to name a few. In the context of transformers, Noci et al 2022 and also https://arxiv.org/abs/2311.01906 have shown this effect too. I would recommend including the new plots and a discussion of these more practical architectures in the paper. 4. I still maintain that the fact that the definition of mu(X) doesn't separate out the cosine similarities going to 1 vs the activation norms going to 0 as problematic for the understanding of the mechanisms of these different architectural components on rank collapse. Even if theoretically it is not possible to show, I would devise separate metrics to isolate these two effects and produce the equivalent of Figures 1 and 2.

Authorsrebuttal2024-08-09

Answers to the additional questions

We thank the reviewer for the quick response and the additional questions. Below, we provide point-to-point answers: **Q1**: Thank you for the comment. - QKV are initialized using the $U(-\sqrt{k}, \sqrt{k})$, with $k$=1/input_dim, which is the default initialization used in transformers like BERT implemented in HuggingFace. - Theorem 2 is established under the assumption that the value matrices are orthogonal. It is important to note that while finite random uniform matrices are orthogonal in expectation, each realization is *not* orthogonal. Hence the phenomenon in experiments does not contradict our theoretical results. - Note that due to the existence of LayerNorm, token trajectories are not a continuous function of model parameters, i.e. the value matrices. To get an idea, notice that if $x_1, x_2$ go to zero, the normalized values $x_1/||x_1||_2$ and $x_2/||x_2||_2$ can have a distance as large as two. As a result, any measure on token trajectories, including $\mu$, is not a continuous function of value matrices. **Q2**: Thank you for the question. As detailed in the response to W4 with concrete examples, having zero activations in the center node does *not* play any role in the mechanism of counterexample. If one replaces $W$ in the example with $W_{Z} = Z^{T}WZ$ where $Z$ is an orthogonal matrix, then the trajectory of $X^{(t)}$ of the new dynamics would be $X_Z^{(t)} = X^{(t)}Z$. Notice that under the new dynamics, the first token is going to converge to $(0,1)Z$, which clearly may not have any zero activation. This is a good point. We will add a remark to include a detailed discussion of the right intuition as the one provided to W4 to avoid confusion and misinterpretation. Thank you for bringing this up. **Q3**: Thank you for the comment and the pointers to the references. We are glad that the reviewer found our newly added experiments meaningful and connected to the literature. We will include the newly added experiments in the updated version of the manuscript and discuss the complex interplay between different architectural components in transformers and the related literature. **Q4**: Thank you for the comment. Our definition of $\mu$ is standard: it is mathematically equivalent to the definition of the measure $\textbf{res}$ used in Dong et al. (the original rank collapse paper), and the definition adopted in more recent e.g. Geshkovski et al. to study convergence of tokens in transformers. This measure checks for whether all tokens converge to the same representation or not. While it does not give information about whether the common representation is zero or not, in both cases the model loses representation power as it can no longer map tokens within the same sequence to different values. We appreciate your questions and comments very much. Please let us know for any further questions.

Reviewer 7spk2024-08-10

Thank you

Thanks for the response. Regarding Q1, it would be good to do the relevant plot at initialisation matching the assumptions of Theorem 2 (e.g. orthogonal weights) just as a verification of the theory. I am surpised that SAN+LN does not converge to rank collapse at initialisation, and I remain unconvinced that the theoretical counterexamples provided explain this because I don't think they should not apply at initialisation. But otherwise thanks again for the clarifications.

Authorsrebuttal2024-08-12

Response to the comments

We thank the reviewer for the comment. - To further address the concern regarding the verification of our Theorem 2, we perform an additional set of experiments of SAN+LN under the exact conditions in Theorem 2 with initialization of the value matrices set to be exactly orthogonal. Here is the table summarizing the mean (std) of $\mu$’s in a 128 layer networks under different masks: | layer | complete | causal | slide-window | slide-window-uni | |-----|--------------------|-------------------|-------------------|-------------------| | 0 | 274.7604 (2.0498) | 274.7832 (2.0682) | 274.6891 (1.9408) | 274.7870 (2.0094) | | 32 | 4.6109e-5 (2.0682) | 47.1223 (16.2438) | 111.6686 (3.4742) | 161.0730 (2.5183) | | 64 | 4.6347e-5 (2e-6) | 67.9591 (12.3382) | 95.4022 (4.8795) | 160.5059 (3.8947) | | 96 | 4.6396e-5 (2e-6) | 82.9471 (11.5331) | 85.6265 (5.5486) | 164.1741 (5.2969) | | 128 | 4.6220e-5 (2e-6) | 92.3288 (11.6265) | 79.9934 (5.6181) | 165.205 (5.4153) | In particular, we observe that while $\mu$ converges to zero for complete masks and shows a clear decreasing trend for slide-window masks (both masks are strongly connected), such converging trends do not hold for non-strongly connected masks (causal and sliding-window-uni, as they are only quasi-strongly connected), verifying our Theorem 2 and showing that (1) the convergence rate depends inversely on the graph diameter (2) the tightness of our result that Theorem 2 only applies for strongly connected masks. We will include these sets of experiments to better illustrate our theoretical results. - Finally, we would like to emphasize again that the goal of the counterexample is to establish that self-attention with layer normalization can prevent rank collapse from happening, as opposed to what Dong et al. claims (“layer normalization plays no role”). We do not claim that they are the only possible weights that can work. It merely establishes that the set of such desirable designs exists. We thank the reviewer once again for the discussion and the constructive feedback!

Reviewer 7spk2024-08-13

Thank you

Thank you for the additional clarifications and empirical verification of Theorem 2. I agree that including these results as well as some of the other suggestions (like the updated figures in the rebuttal pdf, clarifying the intuition for the counterexamples, or discussion of omitted architectural components like skips/positional encodings) will strengthen the paper. I would also mention the goal of the counterexample in order to better distinguish when these theoretical results are expected to hold in practice (and in the empirical verifications), for a cleaner message. I think with these updates the paper would make a nice contribution to NeurIPS 2024. I will keep my score of 6.

Authorsrebuttal2024-08-13

Thank you

We are happy to read this sentiment and will diligently incorporate the updates in the camera-ready version of our paper. We sincerely thank you for your valuable input!

Reviewer yhKR7/10 · confidence 3/52024-07-15

Summary

In this work, the authors investigate the issue of rank collapse in Transformers, providing insights into how attention masks and layer normalization can mitigate this problem. The paper includes extensive analysis, addressing two important questions and offering valuable contributions to the field.

Strengths

1.I like this paper for its strong motivation and very interesting insights. I appreciate the authors for addressing the rank collapse issue in Transformers, a topic often overlooked in the community. 2. The authors provide detailed analysis to demonstrate that attention masks and layer normalization can help address this issue. 3. The experimental results support the authors' analysis effectively. 4. Figure 1 is particularly helpful in understanding the illustration of the effectiveness of layer normalization.

Weaknesses

It seems obvious that causal masking and local attention would help mitigate the issue of rank collapse in Transformer/Attention mechanisms compared to full attention. Maybe no need to demonstrate this a lot.

Questions

I like this work; however, I have two additional questions: If local attention alleviates rank collapse, how do the experimental results compare to those of full attention or causal masked attention? Which approach is more effective in addressing the rank collapse issue: local attention or causal attention? Are there any strong insights into this? I assume local attention may not be generally applicable in language tasks but is more suitable for vision tasks.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Please see above.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC