On the Role of Attention Masks and LayerNorm in Transformers

Authors: xinyi  wu
amir  ajorlou
yifei  wang
stefanie  jegelka
ali  jadbabaie
mit  lids
mit  csail
t u munich
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 22:02:38

Abstract: 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.

Main paper: 
 introduction: 
 The celebrated attention mechanism has proven highly effective in the architecture of transformers [35] , which serve as the key building block of modern foundation models including large language models. From a theoretical perspective, understanding the underlying mechanism of transformers and attention in general has become pivotal for elucidating existing models and paving the way for developing more powerful future models [12, 15] .Transformers are known to exhibit the rank collapse phenomenon1 [12, 15, 20, 27, 33, 37] , which refers to the observation that increasing the number of self-attention layers leads to homogeneous token representations. To gain more insight into the rank collapse issue and better understand the effects of self-attention in multi-layer models, it is essential to study the long-term behavior of tokens under self-attention dynamics [12, 15, 16] .However, many existing studies do not take into account architectural components that are commonly used in practice. For instance, first, the theoretical analysis of long-term self-attention dynamics often assumes that attention is fully bidirectional -that is, all tokens are allowed to attend to all other tokens in the sequence [12, 15, 16, 27] , which only applies to certain attention mechanisms such as the one deployed in the BERT family [11, 24] . The vast majority of popular transformer architectures used nowadays in language models, including the GPT models [6, 30] , use the causal attention masks where tokens are only allowed to attend to preceding tokens, or have sparse attention structure (sparse attention) [4, 8, 10, 21, 22, 32, 38, 41, 42] , where attention is restricted to local interactions between tokens and their chosen neighbors. This limits the practical applicability of the theoretical results developed in [12, 15, 16, 27] , as they heavily rely on the key assumption that attention is fully bidirectional.Second, layer normalization (LayerNorm) is another inherent component of transformers. Dong et al. [12] put forth a hypothesis that LayerNorm plays no role in preventing rank collapse in self-attention networks. This hypothesis is then partially validated by a continuous-time analysis of self-attention dynamics conducted in Geshkovski et al. [15] . The analysis, which incorporates LayerNorm, shows the exponential convergence of tokens to a common point on the unit sphere. However, this result relies on a strong assumption that all the value matrices are the identity matrix. In contrast, for multilayer perceptrons (MLPs), Joudaki et al. [23] show that LayerNorm improves the isometry of the representations and can prevent rank collapse in that setting. It thus remains a question whether it is truly the case that LayerNorm could not have a similar effect in transformers under more general assumptions due to self-attention.Hence, in this work, we rigorously study the effect of attention masks and LayerNorm on the token dynamics and rank collapse in transformers. The main questions we address are as follows:Can attention masks alleviate rank collapse of tokens under self-attention? If so, what type of attention mask would be the most effective?Can LayerNorm alleviate rank collapse of tokens under self-attention? If so, what long-term behavior of tokens would LayerNorm lead to?We answer these questions through a rigorous analysis of the self-attention dynamics. Notably, unlike some previous works [15, 16] , which regard self-attention as a continuous-time dynamical system, we view self-attention as a discrete-time dynamical system, more closely resembling the architecture used in practice. Furthermore, our analysis extends the results to more general masked attention by making novel use of a graph-theoretic approach and incorporating advanced results on infinite products of inhomogeneous non-negative matrices that may be of independent interest.In summary, we make the following contributions:• We establish that with pure self-attention, the exponential convergence of tokens to a common representation holds for a broad class of attention masks, accounting for the causal mask and a wide class of sparse attention patterns such as the sliding window [4, 42] . The key property that leads to the exponential rank collapse is a token that serves as a common "context" for all the other tokens in the sequence to directly or indirectly attend to. Our results also show that local attention can slow down the rate of rank collapse, suggesting its potential advantage over full bidirectional attention from an expressivity perspective at finite depth. • We show that with LayerNorm, when the value matrices are orthogonal, the exponential convergence of tokens to a common point on the unit sphere holds for a broad class of attention masks. Nonetheless, by constructing nontrivial counterexamples, we prove that the self-attention dynamics with LayerNorm can simultaneously have a rich set of equilibria of any possible rank ranging from one to full. Moreover, we rigorously establish that self-attention with LayerNorm, together with proper choice of value matrices, can provably prevent complete collapse of tokens to a rank one subspace for a generic class of input sequences. 
related work: 
 Analysis of self-attention dynamics Understanding the attention mechanism is a pivotal step towards understanding the inner workings of transformer-based models. From a dynamical systems perspective, one could abstract the forward pass of the model as tokens undergoing a nonlinear, time-varying dynamics determined by the self-attention mechanism and other parameters of the model. Specifically, Dong et al. [12] first show that as the number of self-attention layers increases, tokens inevitably suffer from exponential rank collapse, while Wu et al. [37] establish a similar result for the attention mechanism in graph neural networks, taking additionally the layer-wise nonlinear activation functions into account. Other works [15, 16] take a continuous-time approach and study more fine-grained clustering behaviors of the dynamics in transformers.The effect of LayerNorm in transformers LayerNorm is one of the most commonly used normalization techniques in modern neural networks [3] and has become an inherent component of transformers [35] . To better understand its role in transformers, Xiong et al. [39] study it from an optimization perspective and show that LayerNorm stabilizes the gradients during the backward pass.In addition, Brody et al. [5] find that LayerNorm plays a crucial role in improving the expressivity of the attention layer by making it easier for the model to compute the most frequent token in the input and avoid the problem of "unselectable" keys. Most relevant to our work, Dong et al. [12] mention the effect of LayerNorm in terms of mitigating the rank collapse issue in transformers. The paper makes a hypothesis that LayerNorm has no effect for token rank collapse. The argument is based on a heuristic observation (see Appendix A for a detailed discussion) which is at odds with the case of simpler models such as MLPs where LayerNorm is shown to be pivotal in addressing a similar rank collapse problem [23] . It thus remains an open question what effect LayerNorm would have on rank collapse in transformers.Sparse and local Attention While many existing transformer models and LLMs do not use sparse attention, sparse and local attention is gaining popularity due to the demand for efficiency, particular for long-context tasks. For example, sparse attention was populated by Longformer [4] and OpenAI [10] and nowadays popular LLMs like Mistral 7B [22] use sliding window attention by default. Other popular sparse attention models include, but are not limited to BigBird [42] , Recurrent Memory Transformers (RMTs) [7] , and Streaming Attention [38] . Besides language tasks, sparse attention is also common in vision transformers [19, 25, 28] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We confirm that we discuss the limitations of the work, their implications and possible directions for future work.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
problem setup: 
 Notation Let ∥ • ∥ 2 , ∥ • ∥ F be the 2-norm and Frobenius norm, respectively. We use the shorthand[n] := {1, . . . , n}. We denote the all-one vector of length N by 1 ∈ R N . For a matrix M , we denote its i-th row by M i,: and its j-th column by M :,j .Throughout the analysis in the paper, we formalize the attention mask to be a directed graph G.Formally, we represent a directed graph with N nodes by G and let E(G) be the set of directed edges of G. If there is a directed edge going from j to i in G, i.e. (j, i) ∈ E(G), for the attention mechanism it means that token j serves as a direct context for token i or token i attends to token j. The set N i of all neighbors of node i is then {k : (k, i) ∈ E(G)}.Furthermore, we will be using the following graph-theoretic terminology: Definition 1 (Reachability). We say a node v is reachable from u in a directed graph G if and only if there is a directed path (u, n 1 ), (n The radius of a quasi-strongly connected graph is defined to be the longest distance from a center node to any node in the graph. If there are multiple center nodes, then it is the smallest value among them.Given token representations X ∈ R N ×d , the raw attention score matrix is computed asR = XW Q (XW K ) ⊤ / d QK ,where W Q , W K ∈ R d×d ′ are the query and the key matrix, respectively, and d QK is a temperature term to control the scale of raw attention scores. To enforce a masked attention, we create a sparse attention matrix A ∈ R N ×N based on R whose sparsity pattern is specified by a directed graph G: we normalize R ij among all allowed token attention interactions (k, i) ∈ E(G) such thatA ij = softmax G (R ij ) = exp(R ij ) k∈Ni exp(R ik ) if (j, i) ∈ E(G), A ij = 0 otherwise.Given token representations X ∈ R N ×d , LayerNorm subtracts the mean across different columns in each row and then scales each row to have a unit 2-norm. In this work, we consider LayerNorm to only perform the scaling operation, which is a common assumption in theoretical analyses of the attention mechanism [15, 34] 2 . Mathematically, letD = diag(d 1 , d 2 , ..., d N ) where d i = 1/∥X i,: ∥ 2 for all i ∈ [N ],then LN(X) = DX .For our analysis, we consider single-head (masked) self-attention networks (SANs), where the layerwise update rule can be written asEQUATIONEQUATIONEQUATIONwhereW (t) V ∈ R d×d ′is the value matrix. For simplicity, throughout the paper, we assume that d = d ′ and G (t) = G, i.e. the attention mask is the same for all the attention layers. Yet the results can be easily generalized to the case where masks are time-varying and satisfy similar regularity conditions.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: Since all the experiments are for the verification purpose of our theory, they are straightforward and can be easily replicated given the instructions detailed in the paper.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
main results: attention with masking and layernorm: 
 To study how token representations evolve under the self-attention dynamics and behave in the long-term, we measure token similarity via µ(•) : R N ×d → R ≥0 :EQUATION)This measure is mathematically equivalent to the measure res(X) = arg min x∈R d ∥X -1x ⊤ ∥ F used in [12] , but the form in ( 4) is easier to work with in the general analysis and more direct to compute. Another advantage of our formulation is that it clearly demonstrates that Theorem 1 and Theorem 2 are not dependent on the specific choice of µ(•): these results apply to any Lipschitz µ ′ (•) with a Lipschitz constant L such that µ(X) = 0 if and only X i,: = X j,: , ∀i, j ∈ [N ], as we can use the formulation to directly derive thatµ ′ (X) = |µ ′ (X) -µ ′ (1 γ X )| ≤ L∥X -1 γ X ∥ F = Lµ(X) .Finally, we adopt the following assumptions in our analysis:A1 G contains self-loops. i.e. (i, i) ∈ E for every token i ∈ [N ] . A2 There exist constants C ∈ R such that max t∈N W (t) Q 2 , W (t) K 2 ≤ C.A1 ensures that every token has a neighbor so that the masked attention computation is well-defined for every token in every layer, while A2 assumes that the key and query weight matrices are bounded, which is key for efficient attention computation in practice [2] .We first analyze the case without LayerNorm and focus on the effect of the attention mask. To ensure boundedness of the token trajectories X (t) for all t ≥ 0 even without LayerNorm, we further assume that A3 The sequencek t=0 W (t) V 2 ∞ k=0 is bounded.Then with general attention masks G, there remains a strong connection between tokens via attention, and the token representations collapse exponentially to rank one. Theorem 1. Consider the self-attention dynamics without LayerNorm defined in (2) . Under A1-A3, if G is a quasi-strongly connected graph, then there exists ϵ > 0 where for all t ≥ 0,EQUATION)As a result, a rank collapse of tokens happens exponentially with respect to µ(•), i.e. there existsEQUATION)where r is the radius of G, meaning that tokens converge to a common vector exponentially.The detailed proof is provided in Appendix B. The above result suggests that with pure self-attention, as 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 of tokens to a common vector is guaranteed. In particular, it generalizes the main result in [12] from G being a complete graph to a much more general class of attention patterns: the attention pattern G only needs to be quasi-strongly connected, meaning that the result applies to general attention masks used in practice including the causal mask used in decoder-only models such as the GPT family [6, 30] , or sparse attention patterns deployed in many efficient transformer models [4, 10, 21, 32, 42] . We discuss a few interesting implications below.Local vs. global attention The exponential rate (1 -ϵ r ) 1/r is monotone in the graph radius r. This means that rank collapse should be slower for graphs with larger radius r. Our result thus indirectly supports the use of local attention patterns [4, 42] , which not only make the attention computation more efficient (what those works are originally motivated by), but also implicitly alleviate the rank collapse issue.Focused vs. uniform attention In addition, the exponential rate is monotone decreasing in ϵ, which means that rank collapse is slower with smaller ϵ. One can interpret ϵ as how "focused" attention is distributed among reachable tokens, as ϵ is maximized when attention happens uniformly among reachable tokens. Besides applying attention masks and restricting the number of reachable tokens, another way to control how focused attention would be is through the temperature term d QK . As larger values of d QK would make the attention allocation among reachable tokens more even, they should make rank collapse happen faster across layers. Trade-off between rank collapse and universal approximating power Finally, for strongly connected graphs, the above result also reveals a trade-off between universal function approximation power and the rate of rank-collapse. Yun et al. [41] show that transformers with strongly connected graph masks are sequence-to-sequence function universal, yet with a mask G they need at least the diameter of G layers to achieve the full sequence-to-sequence function approximation property. This implies that masks with smaller diameters (and thus smaller radii, as radius ≤ diameter ≤ 2 radius) are more efficient in terms of function approximation power, yet they are more prone to rank collapse. Remark 1. Since the analysis in [12] fundamentally relies on the shift-invariant property of softmax(•), it is necessary that all the tokens are allowed to attend to all the other tokens for their proof to work. On the contrary, we leverage a different graph-theoretic approach to exploit the common structure of the attention matrices to obtain a general result for masked attention.So far, we have considered the pure self-attention dynamics without LayerNorm and focused on the role of the attention mask. What happens if we add LayerNorm and consider the attention dynamics defined in (3) instead? In this section, we first present a negative result, showing that exponential collapse of tokens to a common vector can still happen for certain classes of value matrices. Theorem 2. Consider the self-attention dynamics with LayerNorm defined in (3) . Let G be a strongly connected graph. Assume A1-A2, and that W (t)V is orthogonal for all t ≥ 0, and in addition, the initial input X (0) satisfies that ( * ) N ≤ d, X (0) has full rank.Then there exist C > 0, ϵ > 0 such that N ϵ < 1 and µ(X(t) ) ≤ C(1 -N ϵ 2r ) t 2r∀t ≥ 0 , (7) where r is the radius of G, meaning that tokens converge to a common point on S d-1 exponentially.The detailed proof is provided in Appendix C. The result can be seen as a generalized discrete version of Theorem 4.1 in [15] . Notably, our analysis is based purely on advanced linear algebra tools: infinite products of non-negative matrices and their ergodicity, and can account for time-varying weights and general attention masks, as opposed to fixed W K , W Q , W V over time and G being complete (which is the case in [15] ). One way to satisfy the condition ( * ) on the initial input X (0) is to require N ≤ d and to initialize tokens uniformly randomly on S d-1 , then the condition hold almost surely. This is how the condition is dealt with in [15] .Note that the condition ( * ) implies that there exists v ∈ S d-1 such that ⟨X (0) i,: , v⟩ > 0 for all i ∈ [N ] by either the hyperplane separation theorem or Farkas' lemma (see Lemma 6 in Appendix C). If the initial token geometry satisfies a stronger condition than the above, then ( * ) is no longer necessary and Theorem 2 even directly generalizes to quasi-strongly connected graphs G. We define ϕ (t) := mini,j∈[N ] ⟨X (t) i,: , X (t)j,: ⟩, indicating the minimal cosine similarity between tokens. If the cosine similarities are non-negative for all pairs of tokens initially, then the rank collapse happens exponentially, as long as G is quasi-strongly connected. Corollary 1. Consider the self-attention dynamics with LayerNorm defined in (3) . Let G be a quasi-strongly connected graph. Under A1-A2, if W (t)V is orthogonal for all t ≥ 0 and ϕ (0) ≥ 0, then there exist C > 0, ϵ > 0 such that N ϵ < 1 and µ(XEQUATION)where r is the radius of G, meaning that tokens converge to a common point on S d-1 exponentially.Full mask vs. causal mask We can refine Corollary 1 by specifying the number of center nodes n in the mask G, then the upper bound for the exponential rate would be (1 -nϵ 2r ) instead, meaning that the rate of rank collapse can be negatively affected by the number of center nodes in the mask G. In the case of full attention where G is the complete graph, the mask would have N center nodes, matching the bound in Theorem 2. In the case of causal attention where G is the causal graph, the mask would only have one center node, and the upper bound would be looser, suggesting the advantage of the causal mask in mitigating the rate of rank collapse as compared to the full mask.Post-LN vs. Pre-LN The definition of LayerNorm in (3) follows from the original transformer paper [35] and nowadays it is referred as post-LN [39] . An alternative use of LayerNorm in many LLMs, where LayerNorm comes before self-attention, is called pre-LN and can be written instead asEQUATIONNote that Theorem 2 and Corollary 1 apply directly to the case of pre-LN with similar proofs.The main results from the previous sections seem pessimistic: the self-attention dynamics seem doomed to collapse into a rank one subspace in the long run, with or without LayerNorm. In this section, however, we will first construct a nontrivial counterexample with only LayerNorm such that for a general class of input sequences, tokens converge to an equilibrium where rank collapse does not happen. Notice that for a transformer model to be practical, it is important that it can prevent rank collapse for a general class of input sequences rather than a specific input sequence. We will then show a general result stating that, with LayerNorm and proper choice of value matrices, the self-attention dynamics can have a rich set of equilibria with any possible rank between one and full. Moreover, for a general class of input sequences, tokens provably do not converge to a rank one subspace under the resultant dynamics.(0, 1) (0, 1) (0, 1)A B 1First token Second tokenBoth tokens(0, 1) (0, 1) A B 1 1 Figure 1: Long-term behavior of tokens in the case of N = 2, d = 2.Without LayerNorm (left), both tokens collapse to the same point in R 2 ; whereas with LayerNorm (right), such a collapse would not necessarily happen and token representations can maintain full rank in the long term (first token converges either to (0, 1) or (0, -1). Assuming convergence to (0, 1) for the first token, the second token converges to B, if it is initially located within the red segment).For simplicity of illustration, we consider N = 2, d = 2, and G to be the causal mask. Then letW (t) K = W (t)Q = 0, which leads to the attention matricesA (t) = 1 0 1/2 1/2 ∀t ≥ 0 .We further letW (t) V = 1 w 0 1 ∀t ≥ 0 ,for w ∈ R. Without loss of generality, fix w > 1. Then a careful analysis shows that, depending on its initial position, the first token will either converge to (0, 1) or (0, -1). Suppose the first token converges to (0, 1). Then the convergence of the second token as t → ∞ is illustrated in Fig. 1 , whereA = -1 w , 1 -1 w 2 , B = -1 w , -1 -1 w 2 .A rigorous proof can be found in Appendix E. Note that due to the scaling effect of LayerNorm, any scaled versionc (t) W (t) V , c (t) > 0 of W (t) Vworks equivalently here. Remark 2. For any orthogonal Z ∈ R d×d , W Z = Z ⊤ W Z works equivalently in this example, and the resulting token trajectories follow X (t) Z.This nontrivial counterexample suggests that with the LayerNorm dynamics in (3), there are proper choices forEQUATIONV matrices that can prevent tokens from collapsing to a rank one subspace, for a nonzero measure set of input sequences.We conclude this section with the following statement generalizing the previous illustrative example: with a proper choice of weight matrices, tokens under the attention dynamics with LayerNorm can simultaneously have equilibria of any rank between one and full. Moreover, for a general class of input sequences, tokens provably do not converge to a rank one subspace. Theorem 3. Let G be the causal graph and consider the attention dynamics with LayerNorm defined in (3). Then there existsW (t) Q , W (t) K , W (t) V ∞ t=0satisfying A2-A3 such that the corresponding dynamic has the following properties:1. For any 1 ≤ k ≤ min{N, d}, there exist at least 2 k equilibria of rank k; 2. For a general class of input sequences X (0) in (S d-1 ) N with measure greater than 0, X (t) does not converge to a rank one subspace as t → ∞.The detailed proof is provided in Appendix F. For the sake of simplicity and observing that in a causal graph, there is an inherent chronological order among tokens, in the convergence analysis of each token we assume all tokens preceding it have already converged.The above result is in direct contrast to the case without LayerNorm (Theorem 1), where all input sequences eventually result in complete rank collapse to a one-dimensional subspace. It also stands in contrast to the hypothesis that LayerNorm plays no role in mitigating the collapse [12] , and suggests that it is an important component of the self-attention dynamics in transformers-adding LayerNorm fundamentally changes the behavior of the underlying system. Compared to the case without LayerNorm, first, LayerNorm guarantees that the system never diverges, no matter what the value matrices are. Second, and more importantly, attention with LayerNorm leads to a surprisingly more expressive, versatile dynamics than the system without it-composition with different value matricesW (t)V can result in different long-term token behaviors: in some cases, tokens completely collapse to a single point (Theorem 2), while in others tokens can retain higher rank asymptotically (Theorem 3).Next, we discuss three further intriguing findings from our analysis of the self-attention dynamics with LayerNorm presented in the previous section.Scaling effect of LayerNorm is key From a dynamical system perspective, rank collapse happens under pure self-attention because the attraction of the center node causes all the tokens to be aligned. In the counterexample provided in Section 4.3.1, the key insight to prevent the second token from aligning with the first token is that the updated representation of the second token must generate a component canceling out the attraction from the first token, which acts as a repulsion. The crucial role of LayerNorm here is to stabilize this cancellation process by readjusting the scale of tokens to ensure that the cancellation persists in all the updates, which would not be the case if there is no LayerNorm.Anisotropy of token embeddings Analyzing and understanding the token geometry in transformers has long been of interest to the NLP and general ML community. In particular, empirical findings suggest that contextual token embeddings generated by transformers are anisotropic, meaning they tend to concentrate in a narrow region [1, 9, 13, 14, 17] . Interestingly, as we show in Appendix F.2, the full rank equilibria described in Theorem 3 align with this observation, as tokens lie in a narrow region on S d-1 . While it still remains unclear how exactly such a representation geometry is useful for downstream tasks [1, 9, 14, 17] , empirical studies have found that anisotropy is not necessarily harmful for semantic representations and can assist with tasks like clustering [1, 26] .Stability of the equilibria Finally, we would like to note that due to the combinatorial nature of the analysis with increasing dimensions, we do not prove the exact convergence to specific equilibria beyond not converging to a rank one subspace. However, we observe in simulation that these equilibria are indeed stable in certain directions, despite their regions of attraction being relatively small. This suggests that while LayerNorm improves the expressivity of pure self-attention dynamics, the resulting expressive power might still be limited in a way that it would be difficult for a generic input sequence to reach a specific configuration at equilibrium. This observation seems in line with the fact that MLP modules are crucial for the universal sequence-to-sequence approximation power of transformers [40, 41] . To achieve maximal expressivity, and in particular, to be able to move tokens freely around the whole space (such as in the case for transformers with MLPs), nonlinear power from MLPs would be necessary. 
numerical experiments: 
 In this section, we validate our theoretical findings via numerical experiments. Following [12] , we randomly select 3000 samples of 128-token excerpts (based on the BERT tokenizer) from Wikipedia using the Wikipedia API in Python. More experimental details and additional results can be found in Appendix G.The effect of attention masks and LayerNorm We first verify the effect of different attention masks and LayerNorm on rank collapse with respect to µ(•). We use BERT [11] as the backbone transformer model and consider five different model variants for a controlled experiment: selfattention network (SAN) in which the model has self-attention layers; SAN with skip connections; SAN with LayerNorm where in each layer there is self-attention followed by LayerNorm; SAN with both skip connections and LayerNorm where LayerNorm is put after skip connections; and finally the standard transformer with all the components. For attention masks, we select four representative types: the complete graph, the causal graph, the sliding window (tokens can only attend to the token right before and after and themselves) and the uni-directional sliding window (tokens can only attend to the token right before and themselves). Fig. 2 shows the average values of µ(X (t) ) with standard deviations over the 3000 samples in different architectures. We see that in SANs, µ(X (t) ) converges to zero exponentially for all attention masks, yet more local attention masks slow down the convergence rate. Furthermore, in randomly initialized models, right after we add in LayerNorm, µ(X (t) ) no longer converges to zero and can be maintained at a stable level significantly above zero. In pretrained models, LayerNorm helps prevent the issue together with other components such as skip connections and stabilize the representations. Results in 128-layer initialized models also exhibit the same trend, and can be found in Appendix G.1. ) (in log-log scale) as the number of layers increases. Rank collapse happens exponentially for pure attention, despite different attention masks having different convergence rates. However, as soon as we solely add in LayerNorm, µ(X (t) ) no longer converge to zero in randomly initialized models; in pretrained models, LayerNorm helps prevent the issue together with other components and stabilize the representations. ) (in log-log scale) as the number of layers increases. Smaller temperature terms alleviate the rate of rank collapse, and effect is more significant with global attention than with sparser masked attention, and more in shallower layers than deeper layers.The complex interplay between different components in transformers Previous works have shown that skip connections can help combat rank collapse in transformers. Yet in Fig. 2 , skip connections seem to make µ(X) unstable, particularly in deeper layers (for 128-layer results, see Appendix G.1). Compared with full transformers where µ(X) stays relatively stable, there is a clear discrepancy. In that case, LayerNorm emerges as a crucial element in mitigating rank collapse and stabilizing token representations while also counteracting potential negative effects of pure skip connections. This underscores the complex interplay between different components in transformers.The effect of the temperature term Next, we investigate the effect of the temperature term d QK in the attention calculation. Since the amount of focus of the attention is affected by both the attention mask and the temperature term, we investigate the effect of d QK with different attention masks. Fig. 3 shows the average values of µ(X (t) ) with standard deviations over the 3000 samples with different masks in pretrained SANs. We observe that while a smaller temperature term d QK alleviates the (initial) rate of rank collapse in all cases, the effect is more significant with global attention than with sparser masked attention, and more in shallower layers than in deeper layers. The results in initialized models show similar trends and can be found in Appendix G.2.Evolution of token geometry Finally, we study the evolution of token geometry as the number of layers increases in transformers. Specifically, to capture more fine-grained details than µ(•), we measure how the rank, minimal singular value of token representations, and absolute values of pairwise cosine similarities among tokens change as tokens pass through the transformer layers. We select four representative pretrained transformer models: BERT [11] , GPT2 [30] , T5 [31] and ALBERT [24] and the average results over 3000 samples are shown in Fig. 4 . We see that all models exhibit similar behaviors as tokens evolve along the layers: the models can effectively preserve the full rank of token representations, as also evident from the minimal singular values consistently remaining significantly above zero. Yet the absolute cosine similarities among tokens suggest that tokens gradually concentrate in a narrow region. This empirical observation aligns with the characteristics of the equilibrium in our counterexample that a stable long-term geometry for token representations in transformers can retain full rank while being close to a low dimensional geometry at the same time. Figure 4 : Evolution of token geometry as the number of layers increases. We see that tokens are indeed able to maintain full rank, while at the same time the representations are anisotropic, meaning that they concentrate in a narrow region, as indicated by the average pairwise absolute cosine similarities. 
conclusion: 
 The attention mechanism has led to significant empirical progress in building foundation models. From a dynamical system perspective, the time-varying and nonlinear nature of self-attention dynamics, when coupled with attention masks, LayerNorm and other components in transformers, enables expressive and complex behavior. In particular, we show that the choice of attention masks directly affects the token dynamics. It would hence be valuable for future research to investigate how to effectively design attention masks. Our result further suggests that robust token geometries under multi-layer self-attention can exhibit both full-rankness and anisotropic characteristics simultaneously, which is the case in real transformers as well. It would be interesting to study how such co-existence of these two characteristics would help with learning tasks -whether the full-rank yet close-to-low-dimensional geometry enables efficient learning and generalization while still letting tokens capture meaningful fine-grained details. A Comment about Section 3.3 in Dong et al. [12] From the mathematical form of LayerNorm in Eq. ( 3), since LayerNorm normalizes each row of X(t) to have 2-norm one, the operator that represents this scaling effect D (t) should be multiplied from the left hand side (scaling each row) rather than from the right hand side (scaling each column).Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We confirm that we specify all the details for the experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We confirm that we report standard deviations for all the experimental results.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Our work is about machine learning theory with no release of data or models. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We have properly credited all the data and models used in the work. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: In particular, since in Dong et al. [12] , token representation X takes the same formulation that rows represent different tokens while columns represent different features, the same rule applies. Thus, the argument based on D (t) being multiplied from the right hand side and can be merged withW (t) Vwould not work.Lemma 1. Under A1-A3, there exists ϵ > 0 such that A (t) i,j ≥ ϵ for all t ≥ 0, (j, i) ∈ E.Proof. Writing (2) recursively, we get that the token trajectoriesEQUATIONstay uniformly bounded for all t ≥ 0 by A3. Then it follows from A2 that there exists C ∈ R such that for all t ≥ 0,X (t) W (t) Q i,: 2 = X (t) i,: W (t) Q 2 ≤ C , X (t) W (t) K i,: 2 = X (t) i,: W (t) K 2 ≤ C .Hence for all i, j ∈[N ], -C 2 ≤ (X (t) W (t) Q (X (t) W (t) K ) ⊤ ) i,j ≤ C 2 .This implies that there exists ϵ > 0 such that A (t) i,j ≥ ϵ for all (j, i) ∈ E.Fix a vector x (0) ∈ R d and a {A (n) } ∞ n=0 in A G,ϵ . Let x (t) := A (t) ...A (0) x (0) . We further denote thatmax i∈d x (t) i := M (t) , min i∈d x (t) i := m (t) .Note that M (t) is monotone non-increasing in t, whereas m (t) is monotone non-decreasing in t.Lemma 2. Let x (t) i = pm (t) + (1 -p)M (t) for 0 ≤ p ≤ 1.Then for any T ∈ N,EQUATIONEQUATION)Proof. Given x(t) i = pm (t) + (1 -p)M (t), we get thatx (t+1) i = N j=1 A (t) ij x (t) j ≤ A (t)ii (pm(t) + (1 -p)M (t) ) + (1 -A (t) ii )M (t) = pA (t) ii m (t) + (1 -pA (t) ii )M (t) . Subsequently, x (t+2) i = N j=1 A (t+1) ij x (t+1) j ≤ A (t+1) ii (pA (t) ii m (t) + (1 -pA (t) ii )M (t) ) + (1 -A (t+1) ii )M (t) = p t+1 n=t A (s) ii m (t) + 1 -p t+1 n=t A (s) ii M (t) .We obtain (11) by iterating the process. Similarly, (12) can be derived using a symmetric argument as for the upper bound.Lemma 3. Let {A (n) } ∞ n=0 in A G,ϵand r be the radius of G. Then for t ≥ 0 and 0 ≤ p ≤ 1,EQUATIONProof. Let i 0 ∈ G be a center token and fix t ≥ 0 and 0 ≤ p ≤ 1. Without loss of generality, supposeEQUATION)From Lemma 2, we get that for all T ∈ N,x (t+T ) i0 ≤ p t+T -1 s=t A (s) i0i0 m (t) + 1 -p t+T -1 s=t A (s) i0i0 M (t) ≤ pϵ T m (t) + 1 -pϵ T M (t) .Denote V k as the set of tokens that are exactly k-hop neighbors of i 0 :V k := {j ∈ V : dist(i 0 , j) = k} .For anyEQUATION)Apply Lemma 2 to (15), we further get that for T ∈ N,x (t+T ) i1 ≤ pϵ T m (t) + 1 -pϵ T M (t) . Then let i 2 ∈ V 2 .For simplicity, we still use i 1 to denote the intermediate neighbor between i 0 and i 2 . Similarly to the case of i 1 , we get thatEQUATIONand applying Lemma 2 to (16) , it follows that for T∈ N ≥2 x (t+T ) i2 ≤ pϵ T m (t) + 1 -pϵ T M (t) .Iterating this process for tokens in V 3 , ..., V r , where r is the radius of G, we get that for any token i ∈ G,EQUATIONFor the other case of ( 14), i.e. 17) is obtained using a symmetric argument by bounding m (t+r) from below through (12) . This completes the proof of (13).x (t) i0 > pm (t) + (1 -p)M (t) , (Corollary 2. There exists C(x (0) ) > 0 such thatM (t) -m (t) ≤ C (1 -ϵ r ) t/r , ∀t ≥ 0 ,where r is the radius of the graph G.With Corollary 2, we are ready to prove the main theorem. Recall (10) , we then get thatEQUATIONwhereEQUATIONV . This means that for any j ∈ [d], there exists C j > 0 such thatEQUATIONEQUATIONwhere (19) follows from ( 2) and ( 20) follows from A3.Then by (20) , i,j ≥ ϵ for all t ≥ 0, (j, i) ∈ E.µ(X (t) ) = ∥X (t) -11 ⊤ X (t) /N ∥ F = d j=1 ∥X (t+1) :,j -11 ⊤ X (t) :,j /N ∥ 2 2 = 1 2N d j=1 N m=1 N n=1 |X (t) m,j -X (t) n,j | 2 ≤ N 2 d j=1 C 2 j (1 -ϵ r ) 2t/r ≤ C (1 -ϵ r ) 2t/r = C ′ (1 -ϵ r ) t/Proof. Note due to the way the dynamical system is defined in (3), for every t ≥ 0, every token X (t) i,:lies within the unit sphere S d-1 . Then it follows from A2 that there exists C ∈ R such that for all t ≥ 0,X (t) W (t) Q i,: 2 = X (t) i,: W (t) Q 2 ≤ C , X (t) W (t) K i,: 2 = X (t) i,: W (t) K 2 ≤ C .Hence for all i, j ∈[N ], -C 2 ≤ (X (t) W (t) Q (X (t) W (t) K ) ⊤ ) i,j ≤ C 2 .This implies that there exists ϵ > 0 such that A (t) i,j ≥ ϵ for all (i, j) ∈ E. Lemma 5. Under the same assumptions as Theorem 2, there existsC 1 , C 2 > 0 such that C 1 ≤ D (t) i,i ≤ C 2 for all t ≥ 0, i ∈ [N ].Proof. We adopt the following notation: Notation 1. For X ∈ R N ×d , we use Conv(X) to denote the convex hull of {X 1,: , ..., X N,: }.First, since for x ∈ Conv(X (t) ), W (t)⊤ x always lies in S d-1 , D (t) i,i ≥ 1.Then for all t ≥ 0, each row vector of (A (t) X (t) W (t) ) i,: lies in the convex cone generated by X (0) W (t) , where W (t) := Π t i=0 W (i) . To see this,(A (t) X (t) W (t) ) i,: = N j0=1 A (t) ij0 X (t) j0,: W (t) = N j0=1 A (t) ij0 D (t-1) j0,j0   N j1=1 A (t-1) j0,j1 X (t-1)j1,:W (t-1)   W (t) = (j0,...,jt)∈[N ] t+1 A (t) ij0 D (t-1) j0,j0 A (t-1) j0,j1 D (t-2) j1,j1 ...D (0) jt-1,jt-1 A (0) jt-1,jt X (0) jt,: W (t) Hence all t ≥ 0, (A (t) X (t) W (t) ) i,: = N j=1 a (t) ij X (0) j,: W (t) ,where a(t) ij ≥ 0, ∀j ∈ [N ] and N j=1 a (t) ij ≥ 1.We then make the following observation: Lemma 6. Under assumption ( * ), there exists v ̸ = 0 such that⟨X (0) i,: , v⟩ > 0 i ∈ [N ] ,meaning that all the points of X (0) lie in an open hemisphere H.To see this, observe that since X (0) is full rank and N ≤ d, 0 / ∈ Conv(X (0) ). Then the hyperplane separation theorem states that there exists a hyperplane strictly separating 0 and Conv(X (0) ), letting us conclude Lemma 6. Since all the points of X (0) lie in an open hemisphere H, there exists γ > 0 such that for all x ∈ Conv(X (0) ) and W orthogonal,∥W x∥ 2 = ∥x∥ 2 ≥ γ . Then ∥(A (t) X (t) W (t) ) i,: ∥ 2 =   N j=1 a (t) ij   N j=1 a (t) ij X (0) j,: W (t) N j=1 a (t) ij 2 ≥ N j=1 a (t) ij X (0) j,: W (t) N j=1 a (t) ij 2 ≥ γ .We conclude that D (t)i,i ≤ 1/γ, for all t ≥ 0.Writing the layer-wise update rule (3) recursively, we get the following formulation of X (t) : X (t+1) = D (t) A (t) ...D (0) A (0) X (0) W (0) ...W (t) .We focus on the infinite product lim t→∞ D (t) A (t) ...D (0) A (0) . We first define a family of attention matrices as follows: Definition 6. Let ϵ > 0. We define A G,ϵ to be the set of row-stochastic matrices satisfying the following conditions:1. ϵ ≤ A ij ≤ 1, if (j, i) ∈ E(G), 2. A ij = 0, if (j, i) / ∈ E.We also defineEQUATION) and subsequently,EQUATION)To study the infinite product of matrices from M G,ϵ,C1,C2 , we introduce some necessary concepts: Definition 7. A non-negative matrix M is called row allowable if it has a positive entry in each of its rows.Definition 8. Consider a sequence of row allowable matrices {M (t) } ∞ t=0 . Define the partial productP (t) := M (t) ...M (0) .We say that the sequence {P (t) } ∞ t=0 is ergodic if there exists a sequence of positive rank one matricesEQUATION)Our goal is to show the following key result:Lemma 7. Let ϵ, C 1 , C 2 > 0. Given {D (t) A (t) } ∞ t=0 in M G,ϵ,C1,C2. The sequence of partial products{P (t) } ∞ t=0 is ergodic.For that, we will make use of the following result:Lemma 8 (Corollary 5.1 [18] ). Consider a sequence of row allowable matrices {M (t) } ∞ t=0 . Let a t and b t be the smallest and largest entries in M (t) , respectively. If∞ t=0 at bt = ∞, then {P (t) } ∞ t=0 is ergodic.We cannot directly apply Lemma 8 to D (t) A (t) yet unless in the special case of full bidirectional attention such as the one used in BERT i.e. G is an undirected, complete graph, as a t could equal zero due to the sparsity pattern of attention G. In order to make use of the above result, we first need to show that long products of P (t) := D (t) A (t) will eventually become strictly positive for strongly connected G. For t 0 ≤ t 1 , we denote P (t1:t0) = P (t1) . . . P (t0) . Lemma 9. Under A1, if G is strongly connected, there exist T ∈ N and C 3 , C 4 > 0 such that for all t ≥ 0,C 3 ≤ P (t+T :t) i,j ≤ C 4 , ∀1 ≤ i, j ≤ N .Proof. Note that by A1 and the strong connectivity of G, there exists T ∈ N, e.g. the diameter of G, and ϵ, C 1 , C 2 > 0 such that for all t ≥ 0, (ϵC 1 ) T ≤ P(t+T :t) i,j ≤ C T 2 , ∀1 ≤ i, j ≤ N.Then we apply Lemma 8 to {P (k+1)T :T } ∞ k=0 and concludes Lemma 7. As a result, from (23) we see that any P (t) can now be written in the following form:P (t) = u (t) (v (t) ) ⊤ + E (t) ,where u (t) , v (t) are positive vectors such that u (t) (v (t) ) ⊤ = S (t) . Without loss of generality, assume ∥v (t) ∥ 2 = 1.Lemma 10.lim t→∞ E (t) ij u (t) i = 0, ∀i, j ∈ [N ] . Proof. Given v (t) j E (t) ij u (t) i v (t) j ≤ E (t) ij u (t) i v (t) jwhere by (23) , since alsolim t→∞ E (t) ij u (t) i v (t) j = 0 ,We get thatlim t→∞ E (t) ij u (t) i = 0 ,and hencelim t→∞ E (t) ij u (t) i = 0 .Thus the 2-norm of the i th token of X (t+1) , ∥X (t+1) i,:EQUATIONThen by Lemma 10,lim t→∞ (v (t) ) ⊤ X (0) W (t) + E (t) i,: u (t) i X (0) W (t) 2 = lim t→∞ (v (t) ) ⊤ X (0) W (t) 2 = lim t→∞ 1 u (t) i . Note that min ∥v∥2=1 v ⊤ X (0) W (t) 2 = σ min ((X (0) W (t) ) ⊤ ) ,where σ min (•) denotes the smallest singular value. Since we assume X (0) has full rank, andW (t) is orthogonal lim t→∞ (v (t) ) ⊤ X (0) W (t) 2 ≥ σ min ((X (0) W (t) ) ⊤ ) > 0 . Then given ∀i ∈ [N ], lim t→∞ (v (t) ) ⊤ X (0) W (t) 2 = lim t→∞ 1 u (t) i , ∀i ∈ [N ] ,as a result,EQUATION)This together with Lemma 10 implies thatlim t→∞ E (t) ij = 0, ∀i ∈ [N ], j ∈ [d] ,andEQUATIONHaving established the convergence in (26) , we then establish the exponential convergence rate.We define that α (t) := mini,j∈[N ] ⟨X (t) i,: , X (t) j,: ⟩.Similar to the derivation of ( 25), note thatmax ∥v∥2=1 v ⊤ X (0) W (t) 2 = σ max ((X (0) W (t) ) ⊤ ) , it follows that lim t→∞ (v (t) ) ⊤ X (0) W (t) 2 ≤ σ max ((X (0) W (t) ) ⊤ ) < ∞ . Since ∀i ∈ [N ], lim t→∞ (v (t) ) ⊤ X (0) W (t) 2 = lim t→∞ 1 u (t) i , ∀i ∈ [N ]as a result,lim t→∞ u (t) i = lim t→∞ u (t) k > 0, ∀i, k ∈ [N ] , meaning that lim t→∞ X (t) i,: ∈ S d-1 for all i ∈ [N ] .Then given the convergence established in ( 26), we get that lim t→∞ α (t) = 1 , from which there exists t 0 ∈ N such that α (t) > 0, ∀t ≥ t 0 . It follows that for t ≥ t 0 ,ϕ (t+r) = min i,j∈[N ] ⟨X (t+r) i,:, X (t+r) j,: kr , XEQUATION(t) lr ⟩ ≥ N ϵ 2r + (1 -N ϵ 2r )ϕ (t) . Thus 1 -ϕ (t+r) ≤ (1 -N ϵ 2r )(1 -ϕ (t) ).Writing recursively, we get that for k ≥ 0,EQUATION)Since by definition, 1 -ϕ (t) ≥ 1 -⟨X (t) i,: , X (t) j,: ⟩ ≥ ∥X (t) i,: -X (t) j,: ∥ 2 2 /2, it follows that µ(X (t) ) = ∥X (t) -11 ⊤ X (t) /N ∥ F = N i=1 ∥X (t) i,: -1 ⊤ X (t) /N ∥ 2 2 = 1 2N N i=1 N j=1 ∥X (t) i,: -X (t) j,: ∥ 2 2 ≤ C ′ (1 -N ϵ 2r ) t 2r ,where C ′ = √ CN , completing the proof.First, since for x ∈ Conv(X (t) ), W (t)⊤ x always lies inS d-1 , D (t) i,i ≥ 1, for all i ∈ [N ], t ≥ 0.For all t ≥ 0, since ϕ (0) ≥ 0 and G is quasi-strongly connected, we have thatϕ (t+r) = min i,j∈[N ] ⟨X (t+r) i,:, X (t+r) j,: EQUATIONj (t+r) ,l1 ...A (t) lr-1,lr ⟨X (t) kr , X (t) lr ⟩ ≥ ϵ 2r + (1 -ϵ 2r )ϕ (t) . Thus 1 -ϕ (t+r) ≤ (1 -ϵ 2r )(1 -ϕ (t) ).Writing recursively, we get that for k ≥ 0,EQUATIONSince by definition, 1 -ϕ (t) ≥ 1 -⟨X (t) i,: , X (t) j,: ⟩ ≥ ∥X (t) i,: -X (t) j,: ∥ 2 2 /2, it follows that µ(X (t) ) = ∥X (t) -11 ⊤ X (t) /N ∥ F = N i=1 ∥X (t) i,: -1 ⊤ X (t) /N ∥ 2 2 = 1 2N N i=1 N j=1 ∥X (t) i,: -X (t) j,: ∥ 2 2 ≤ C ′ (1 -ϵ 2r ) t 2r ,whereC ′ = √ CN , completing the proof.E Detailed analysis of the illustrative example in Section 4.3.1First token Given thatX (t+1) 1,2 X (t+1) 1,1 = X (t) 1,2 X (t) 1,1 + w , we get that X (t+1) 1,2 X (t+1) 1,1 = X (0) 1,2 X (0) 1,1 + wt .and thuslim t→∞ X (t) 1,2 X (t) 1,1 = ∞ .Since the sign ofX (t)1,1 is invariant, depending on its initial sign, we get thatlim t→∞ X (t) 1,: = [0, 1] if X (0) 1,1 > 0 [0, -1] if X (0) 1,1 < 0 .Second token Assume that the first token converges to [0, 1]. ThenEQUATION)To show the convergence, we utilize the following useful property of g(•): Lemma 11. g is non-decreasing in X 2,2 /X 2,1 .Proof. Let y = X 2,2 /X 2,1 . Given the identity that 1X (t) 2,1 2 = y 2 + 1 , we get that 2 ∂ 1 X2,1 ∂y 1 X 2,1 = 2y . Hence ∂ 1 X2,1 ∂y = X 2,1 y = X 2,2 , which means that ∂g ∂y = 1 + X 2,1 ≥ 0 .Given X 2,2 /X 2,1 is invariant at X 2,1 = -1/w (symmetric points A and B in Fig. 5 ), together with Lemma 11, they imply the invariance of the 4 segments on S 1 , as shown in Fig. 5 . Furthermore, (29) shows which point each region would converge to. Thus it suffices to choose X2,: to be on the dark blue or orange segment.Step 1 Without loss of generality, we set that W(t) K = W (t) Q = 0,for all t ≥ 0. The system of equations corresponds to the equilibrium condition for the attention dynamics with LayerNorm defined in (3) when the attention mask G is causal is as follows: whereEQUATIONEQUATIONEQUATIONThen if X K,: is a generalized eigenvector of rank k corresponding to W V with eigenvalue 1, then from (31) that {X K,: , ..., X 1,: } is the Jordan chain generated by X K,: , which is known to be linearly independent.Hence we set W V to be a matrix with a generalized eigenvector of rank k corresponding to eigenvalue 1. For example, 1. For k = 1, W V can be set as I d ;2. For k = 2, W V can be set asW (t) V =         1 0 0 ... 0 0 0 1 0 ... 0 0 . . . . . . . . . ... . . . . . . 0 0 0 ... 0 0 0 0 0 ... 1 w 0 0 0 ... 0 1         ∀t ≥ 0 .3. For k = 3, W V can be set asW (t) V =         1 0 0 ... 0 0 0 1 0 ... 0 0 . . . . . . . . . ... . . . . . . 0 0 0 ... w 0 0 0 0 ... 1 w 0 0 0 ... 0 1         ∀t ≥ 0 .4. For k = N , W V can be set asEQUATIONFor our own purpose, we consider k = N = d. Then to solve for equilibrium defined by the condition in (31) ,EQUATIONThen setting X 1,: = [0, 0, ..., 0, 1], iteratively solving (31) , we get that an equilibrium satisfying ( 30) is as follows:                                               X 1,: = [0, 0, ..., 0, 0, 1] X 2,: = 0, 0, ..., 0, -1 w , -1 -1 w 2 X 3,: = 0, 0, ..., 1 w 2 , 1 w 2 -1 w 4 , 1 -1 w 2. . .X N,: = (-1) N -1             1 w N -1 1 w 2(N -2) - 1 w 2(N -1)...EQUATIONwhich is valid if w > 1. Since from (31) , there is a free coordinate, in fact there a 2 N many such equilibria. Hence there exists {WEQUATIONEQUATIONV } ∞ t=0 such that the corresponding attention dynamics with LayerNorm has an equilibrium with rank k = N .Step 2 To show that there are equilibria of any rank between 1 and full co-existing for a corresponding dynamics given the same set of {W(t) K , W (t) Q , W (t) V } ∞ t=0 , consider W (t)V to be the case of (32) . Then notice that to show there exists equilibrium of k beyond the full rank one as derived in Step 1, notice the following:We can set the first N -k + 1 tokens to be [0, ..., 0, 1], which still satisfies the equilibrium condition. This implies that β N -k+1 = 1/(N -K + 1). Then recall (30):X i+1,: = β i+1 1 β i X i,: + X i+1,: W V Let β N -k+2 = 1, we get that X N -k+2,: (I -W ) = [0, 0, ..., 0, N -k + 1] ,which yieldsX N -k+2,: = 0 d-2 , - N -k + 1 w , ± 1 - (N -k + 1) 2 w 2 .Continue the process by setting β i = 1 for i ≥ N -k + 2, we get a rank k equilibrium as desired (there are 2 k such equilibria), as long as w ≥ N .Step 3 Finally to ensure {WEQUATIONEQUATIONV } ∞ t=0 satisfies A2-A3:EQUATIONQ } ∞ t=0 satisfies A2 by construction;• Since this is a dynamics with LayerNorm, applying WV is in fact equivalent to applying ŴV = W V /∥W V ∥ 2 under the dynamics, where the latter has 2-norm at most 1 and thus using ŴV satisfies A3.To show that the full rank equilibria X * lie in a narrow region, we make use of the notion stable rank:SRank(X) = ∥X∥ 2 F ∥X∥ 2 2 .We would like to show that for any δ > 0, weights can be chosen such that1 ≤ SRank(X * ) ≤ (1 + δ) + O(1/N ) .Since for any matrix M , ∥M ∥ 2 F ≥ ∥M ∥ 2 2 , the first inequality trivially holds. To show the second inequality, consider the equilibrium in (34) . Then direct calculation yields that∥X * ∥ 2 F = Tr(X * (X * ) ⊤ ) = N . Since by definition, ∥X * ∥ 2 = sup ∥v∥2=1 ∥X * v∥ 2 ,and thus any ∥X * v∥ 2 such that v ∈ R N , ∥v∥ 2 = 1 serves as a lower bound for ∥X∥ 2 .We consider v = [0, 0, ..., 0, 1] ⊤ ∈ R d . Then it follows that∥X * ∥ 2 2 ≥ 1 + (N -1) 1 - 1 w 2 , which implies that SRank(X * ) = ∥X * ∥ 2 F ∥X * ∥ 2 2 ≤ N N -N -1 w 2Hence for any δ > 0, it suffices to choose w ≥ 1 such thatw ≥ 1 δ + 1 .The convergence analysis for the case d = 2 can be found in Appendix E. Here, we deal with the general case where d > 2.To show that tokens do not converge to an rank one subspace, it suffices to show that given certain conditions on the initial input X (0) , there are invariant space of the tokens that is at least rank 2.To that end, we analyze the convergence of the first two tokens.First token Given that X+ w , we get thatX (t+1) 1,2 X (t+1) 1,1 = X (0) 1,2 X (0) 1,1 + wt . Let X (0) 1,2 X (0) 1,1= r 2 w, for r 2 ∈ R. It follows from above thatEQUATIONThen for allEQUATIONsettingX (0) 1,i X (0) 1,1 = r i w i-1 ,where r i ∈ R, and substituting from the base case ( 35) to ( 36) recursively, we get thatX (t) 1,i X (t) 1,1 = w i-1 r i + r i-1 t + r i-2 t-1 j1=0 j 1 + r i-3 t-1 j2=0 j2 j1=0 j 1 + ...+ + r 2 t-1 ji-3=0... which implies thatEQUATIONFurther, for i ≥ 3EQUATIONThen ( 38) and (39) yields thatlim t→∞ X (t) 1,i X (t) 1,i-1 = ∞ , which let us conclude that if X (0)1,1 > 0, the first token would converge to [0, 0, ..., 0, 1].Second token Assume that the first token has converged to [0, 0, ..., 0, 1]. Similar to the first token, we can show that for2 ≤ i ≤ d -1, it holds that lim t→∞ X (t) 2,i X (t) 2,i-1 = ∞ ,is well defined, where f (•) is monotone non-increasing in t for fixed r 1 , ..., r d-2 > 0. We denote and since by (40) ,ȳ(t) = w 2 -1 -f (t,EQUATIONand thus X This let us conclude that for a general class of sequences X (0) , X (t) does not converge to a rank one subspace as t → ∞.Directionality of attention mask G In the proof of Theorem 3, we would like to note that G being the causal graph gives a key structure to construct the counterexamples. A causal graph essentially introduces asymmetries among token interactions where some tokens can attend to some other tokens but not verse versa and we exploit such asymmetries in constructing the counterexamples. A similar construction thus does not naturally apply to the case where G is undirected, e.g. a complete graph.As a result, we leave the following interesting question for future research about whether the directionality of graphs empowers the attention mechanism differently. We conjecture a negative answer based on preliminary simulations.Does a similar equilibrium exist if one uses an undirected graphs G as the attention mask? If not, what differentiates the case between an undirected graph and a directed graph?Here we provide more details about numerical experiments presented in Section 5 and some additional experimental results. All models were implemented with PyTorch [29] and Transformers library [36] .Compute We ran all of our experiments on CPUs.-Wikipedia: MIT license tranformers [36] : Apache License 2.0• Pretrained models -BERT (bert-base-uncased) [11] : Apache License 2.0 -GPT2 (openai-community/gpt2) [30] : MIT license -T5 (google-t5/t5-base) [31] : Apache License 2.0 -ALBERT (albert/albert-base-v2) [24] : Apache License 2.0We provide the results for deeper models for the experiment presented in Fig. 2 . Here, we conduct the same experiment in a 128 layer randomly initialized BERT with 12 heads and 768 hidden dimension using different attention masks. Similarly, in SANs, while different masks have different convergence rates, all of them suffer from exponential rank collapse. But as soon as we add in LayerNorm, µ(X (t) ) no longer converges to 0, as µ(•) does not show any decreasing trend even after 128 layers. The effect of the temperature term d QK in initialized transformers is shown in Fig. 7 . Similar to the pretrained models, we see that smaller temperature terms alleviate the rate of rank collapse, and effect is more significant under global attention than under sparser masked attention, and in shallower layers than deeper layers. In this section, we investigate the full set of singular values of token representations X (t) in transformers. For a clear presentation, we randomly select 3 out of the 3000 sequence samples described in Section 5 and calculate all the singular values of X (t) . The results in randomly initialized transformers are presented in Fig. 8 .We see that while the rank is full for X (t) , there is usually one principal component that is much more dominant than the rest. As we have discussed in the main text, this is the anisotropic characteristic of token representations-there is a strong non-uniformity in different directions. Nonetheless, it is important to note that the "small" singular values are in fact not negligible here in the absolute sense-the minimal singular values are never close to zero, as Fig. 4 (middle) shows.We verify the initial token geometry in transformers. For a clear presentation, we randomly select 3 out of the 3000 sequence samples described in Section 5 and calculate all the pairwise cosine similarities between tokens. The results in randomly initialized transformers are presented in Fig. 9 . We see that in particular in BERT and ALBERT, the initial pairwise cosine similarities are all non-negative. Justification: We confirm that abstract and introduction accurately reflect the paper's contributions and scope.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.This is also referred to as the oversmoothing[15,33,37] or the token uniformity[12,27] problem. 38th Conference on Neural Information Processing Systems (NeurIPS 2024).This definition of LayerNorm, precisely concides with the RMSNorm[43] that is widely deployed in mainstream LLMs. We choose the terminology following the convention of literature. 


Decision: This paper studies the rank collapse problem of attention, focusing on two previously ignored aspects: attention masks and layer normalization. Being a theoretical paper, the authors did a good job providing rigorous analyses to these two important design choices which I believe could help further the community's understanding of attention and Transformers. There are concerns raised, eg on the practical implications of this work -- I believe these are nice to have but shouldn't be the blocker for accepting a theoretical work like this. I'm leaning towards accepting this work and encourage the authors to update the paper with materials presented in the rebuttal period.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: Please see above.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: There is not a clear limitations section but the authors do mention a limitation on line 283.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: None

Comments: title: Thank you
Thank you
Response to the comments
Thank you
Answers to the additional questions
Additional Questions
Answers to the questions raised by Reviewer 7spk
comment: 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!
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.
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!
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.
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.
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.
**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.