On the Role of Attention Masks and LayerNorm in Transformers
Authors: Xinyi Wu
Amir Ajorlou
Yifei Wang
Stefanie Jegelka
Ali Jadbabaie
Venue: NeurIPS 2024 poster (2024)
Keywords: attention mechanism
transformers
layer normalization
deep learning theory
dynamical systems
Primary area: learning_theory
TL;DR: We rigorously show that sparse or local masked attention can mitigate rank collapse of tokens, while LayerNorm can prevent it in transformers, both enhancing model expressivity.

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.

On the Role of Attention Masks and LayerNorm in Transformers
Xinyi Wu
Amir Ajorlou' 'MIT LIDS
Yifei Wang²
Stefanie Jegelka
Ali Jadbabaie
MIT CSAIL
STU Munich
aainyisw,ajoorlco wwr.roa,,iiiiiiiio..... wwfei_i, jadbaabbi jadbabai)@mitt. edu
Abstract
Self-attention the key mechanism of transformers, which are the essential build- ing 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 compo- nents in transformers that may alleviate the rank collapse issue. In this paper, provide a general analysis of rank collapse under self- attentonnnion, taking into account the effects of attention masks and layer normalization (LayerNorrm) In particular, we find that although pure masked attention still suffers from exponential collapse rank subspace, sparse local masked attention 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 rank one subspace still happens exponentially. However, through construction of nontrivial counterex- amples, then establish that with proper choice of value matrices, general class of sequences may not converge to rank one subspace, and the self-attention dynamics with LayerNorm can simultancously possess 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 much more expressive, versatile nonlinear dynamical system than what was originally thought.
Introduction
The celebrated attention mechanism has proven highly effective the architecture of transform- [35), which serve the key building block of modern foundation models including large language models. From 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 [5]
Transformers are known exhibit the rank collapse phenomenor refers to the observation that increasing the number of self-attention layers leads to homogencous token representations. To gain insight into the rank collapse issue and better understand the effects of self-attention multi-layer models, it is essential study the long-term behavior of tokens
33 37 which
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
This is also referred the oversmoothing
the token uniformity problem.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

such as the one deployed in the BERT family [1| (4].. The vast majority of popular transformer architectures used nowadays in language models, including the GPT models [6] [0], use the causal attention masks where tokens are only allowed to attend to preceding tokens, or have sparse attention structure (sparse attention) g T0 21
where attention is restricted to local interactions between tokens and their chosen neighbors. This limits the practical applicability of the theoretical results developed in [52 16 27 as they heavily rely on the key assumption that attention is fully bidirectional.
22 32
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- dynamics conducted in Geshkovski et al. [[]]. The analysis. which incorporates LayerNorm, shows the exponential convergence of tokens to common point on the unit sphere. However, this result relies on strong assumption that all the value matrices are the identity matrix. In contrast, for multilayer perceptrons (MLPs), Joudaki al. (33 show that LayerNorm improves the isometry of the representations and can prevent rank collapse in that setting. It thus remains 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.
analysis of self-attention
Hence, 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 musks 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 rigorous analysis of the self-attention dynamics. Notably, unlike some previous works [15 16). which regard self-attention as a continuous-time dynamical system, view self-attention
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, make the following contributions:
We establish that with pure self-attention, the exponential convergence of tokens a common representation holds for broad class of attention masks, accounting for the causal mask and wide class of sparse attention patterns such as the sliding window 44] 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 con- vergence of tokens to a common point on the unit sphere holds for broad class of attention masks. Nonetheless, by constructing nontrivial counterexamples, we prove that the self-attention dynamics with LayerNorm can simultancously 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.
2 Related work
Analysis of self-attention dynamics Understanding the attention mechanism is pivotal step towards understanding the inner workings of transformer- basedd models. From 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. [2] first show that as the number of self-attention layers increases, tokens inevitably suffer from exponential rank collapse, while Wu et al. (37] establish similar result for the attention mechanism in graph neural networks, taking additionally the layer-wise nonlinear

---

activation functions into account. Other works [6 take continuous time approach and study more fine-ge -eee-grrned clustering behaviors the dynamics transformers.
The effect of LayerNorm in transformers LayerNorm is one of the most commonly used nor- malization techniques modern neural networks [ and has become an inherent component of transformers [35]. To better understand its role in transformers, Xiong et al. [3) 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 crucial role in improving the expressivity of the attention layer by making 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 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 detailed discussion) which is at odds with the case of simpler models such as MLPs where LayerNorm is shown be pivotal in addressing 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 the demand for efficiency, particular for long-contex. tasks. For example, sparse attention was populated by Longformer [4] and OpenAI [00] and nowadays popular LLMs like Mistral 7B (2 use sliding window attention by
default. Other popular sparse attention models include, but are not limited to BigBird [42), Recurrent Memory Transformers (RMTs) 77. and Streaming Attention
38) Besides language tasks, sparse
attention is also common in vision transformers 19 25
3 Problem Setup
Notation Let ||| 2| ||- ||e be the 2-norm and Frobenius norm, respectively. We use the shorthand [n] == (I,...,.1) n.). We denote the all-one vector of length N by 1 E IR-. For a matrix M, we denote its i-th row by Mr. and its j-th column by M.g-
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.c. (j,i) EGG), for the attention mechanism it means that token j serves all neighbors of node i is then {k: (k, i) E E(G)).
direct context for token token attends to token j. The set N₁, of
Furthermore, we will be using the following
theoreic terminology:
Definition (Reachability). We say a node is reachable from in directed graph G if and only if there is directed path (u, n1), (M1, ni....... (ny, S v) from u to v. Definition (Strongly Connected). directed graph G is said to be strongly connected if and only if any two distinct of nodes are reachable from each other.
Definition (Center Node), node u from which every node in the directed graph G is reachable is called a center node.
Definition 4 (Quasi-Strongly Connected). A directed graph G is said to be quasi-strongly connected if has at least one center node.
Definition 5 (Radius). The radius of a quasi-strongly connected graph is defined 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.
3.1 (Masked) Attention Mechanism
Given token representations X E Ndd the raw attention score matrix is computed as
R = XWo(XWx) /ddx
where Wo, Wx E Rdxxd are the query and the key matrix, respectively, and ddq is a temperature term control the scale of raw attention scores. To enforce masked attention, create sparse
3

---

attention matrix A E IR NxN based on R. whose sparsity pattern is specified by directed graph G: we normalize Riy among all allowed token attention interactions (k., E Eg) such that
exp(Riz) Aij softmaxg Ray) ΣKEN, exp(Rst) if E ((9), As otherwise.
3.2 LayerNorm
Given token representations X E R RNxd LayerNorm subtracts the mean across different columns in each row and then scales each row to have unit 2-norm. In this work, we consider LayerNorm only perform the scaling operation, which is common assumption in theoretical analyses of the attention mechanism 13 34 Mathematically, let = diag(d1, d2,..., dy) where d: = 1/
for all E [M] then LN(X
DX.
3.3
Self-attention dynamics
For our analysis, we consider single-head (masked) self-attention networks (SANs), where the layerwise update rule can be written
A softmaxgo)) (x Vdok)
x(+1) = w010
(2)
X(+1) = LN(X == D(0 x
where W E RRdxd is the value matrix. For simplicity, throughout the paper, we assume that d d= d' and (()) G. i.e. the attention mask is the same for all the attention layers. Yet the results can be easily generalized the case where masks are time-varying and satisfy similar regularity conditions.
4 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 pt-) RN - Rzo:
(xx) = || - 1xxlee, where xx = XXN..
(4)
This measure is mathemattially equivalent to the measure res(X) = arg min,er || 1x || used in [12], but the form in is easier to work with in the general analysis and more direct to compute. Another advantage of our formulation is that clearly demonstrates that Theorem and Theorem are not dependent the specific choice of pu-):: these results apply any Lipschitz / (-) with Lipschitz constant L such that (xX) = if and only X,, = N Vi.j E [N]. as we can use the formulation directly derive that
(x) = |/' - (1,,, & L|X 1,x IE = Lx(X).
Finally, we adopt the following assumptions in our analysis:
A1 G contains self-loops. i.c. i) E for every token [N]. w
A2 There exist constants C R such that max { [
|
A1 ensures that every token has neighbor that the masked attention computation 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
"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.

---

4.1 Masked Attention
We first analyze the case without LayerNorm and focus on the effect of the attention mask. To ensure boundedness of the token trajectories ( for all 2 0 even without LayerNorm, we further assume that
{|[T (1
//
} is bounded.
A3 The sequence
k=0
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 22. Under A1-A3, if g is quasi-strongly connected graph, then there exists O where for all t 2 0,
1 ce, for all (j,1) E E.
(5)
As a result, a rank collapse of tokens happens exponentially with respect to pl)), i.e. there exists
C > O such that
)(x(1) S C (1-
where r the radius of 9, 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, 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 being 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 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 below.
32
We discuss few interesting implications
Local vs. global attention The exponential rate (1-e")a/r is monotone 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 e, which means that rank collapse is slower with smaller e. One can interpret as how "focused" attention distributed among reachable tokens, as c 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 dok. As larger values of doæ k 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 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-tee sequence-to-sequeme sequenne function universal, yet with a mask G they need at least the function approximation property. This
diameter of G layers to achieve the full implies that masks with smaller diameters (and thus smaller radii, as radius I diameter S 2 radius) 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 crfft-ivvarian property of softmax(-), it is necessary that all the tokens are allowed to attend 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 obtain general result for masked attention.
4.2 Masked Attention with LayerNorm: Rank Collapse
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 common vector can still happen for certain classes of value matrices.
5

---

Theorem 2. Consider the self-attention dynamics with LayerNorm defined in 3 Let G be strongly
connected graph. Assume A1-A2, and that W is orthogonal for all t z 0, and in addition, the initial inpur X(0) satisfies that
(*) N S d, x!0) has full rank.
Then there exist C > O, > O such that Ne 1 and
p(x(1) - C(1 )
Vt 0,
(7)
where r the radius of J. meaning that tokens converge to a common point on gd-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- non-negative matrices and their ergodicity, and can account for time-varying weights and general attention masks, as opposed to fixed WK. Wq. Wv over time and g being complete (which
is the case in (1)). One way to satisfy the condition (*) on the initial input X (0) is to require N S d and to initialize tokens uniformly randomly on g then the condition hold almost surely. This is how the condition is dealt with in (15].
Note that the condition (+) implies that there exists E Sd-1 such that (x:-) o) 0 for all i [N] by either the hyperplane separation theorem Farkas' lemma (see Lemma 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 e0 = min (X (()). indicating the minimal cosine similarity between tokens. the
[]E[[]
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 quasi-strongly connected graph. Under A A1-A2, if W( is orthogonal for all t > 0 and ())0) 0. then there exist C' > 0, >> O such that Ne < and
p(x(1)) I C(1 C(1 -211- 2
Vt V22 0.
(8)
where r is the radius of G, meaning that tokens converge to common point on Sd-1 exponentially.
Full mask vs. causal mask We can refine Corollary by specifying the number of center nodes n in the mask G. then the upper bound for the exponential rate would be (1 (1- -ne²r) 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.
The definition of LayerNorm in B follows from the original transformer paper (35) and nowadays it is referred as post-I.N (39. An alternative use of LayerNorm in many LLMss where LayerNorm comes before self-attention, called pre-I.N and can be written instead as ((++1)= ( (N(X(())(()) ( D(9) x (((()
Post-LN vs. Pre-LN
Note that Theorem and Corollary apply directly to the case of pre-LN with similar proofs.
4.3 Masked Attention with LayerNorm: Counterexample
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 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 can prevent rank collapse for a general class of input sequences rather than specific input sequence. We will then show general result stating that, with LayerNorm and proper choice of value matrices, the self-attention dynamics can have rich set of equilibria with any possible runk 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.

---

No LayerNorm
With LayerNorm
(0. 1)
(0,
Both tokens
First token
Second token
Figure 1: Long-term behavior of tokens in the case of N = 2, = Without LayerNorm (left), both tokens collapse to the same point IR;; whereas with LayerNorm (right), such a collapse would 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).
necessarily happen and token representations
4.3.1 Illustrative Counterexample
For simplicity illustration, we consider N
2, 2, and be the causal mask. Then let
w = w = O, which leads to the attention matrices
1
1
0
vt 0.
1/2
1/2
We further let
w
1 w
Vt 2 0,
1
for w E 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 (0, converges to (0, 1). Then the convergence of the second token as / + do is illustrated in Fig. where
(0,-1). Suppose the first token
2 - B (-a.--11- w
A rigorous proof can be found in Appendix Note that due to the scaling effect of LayerNorm, any scaled version ( ,(()) > 0 W works equivalently here.
A =
Remark 2. For any orthogonal 7 Rdxd Wz Z WZ works equivalently in this example, and the resulting token trajectories follow (
This nontrivial counterexample suggests that with the LayerNorm dynamics in 3 there are proper choices for w((t) w w((") matrices that can prevent tokens from collapsing to rank one subspace, for nonzero measure of input sequences.
4.3.2 Generalization of the Counterexample
We conclude this section with the following statement generalizing the previous illustrative example: with proper choice of weight matrices, tokens under the attention dynamics with LayerNorm can simultancously have equilibria of any rank between one and full. Moreover, for general class of input sequences, tokens provably do not converge
rank one subspace.
Theorem 3. Let G be the causal graph and consider the attention dynamics with LayerNorm defined satisfying A2-A3 such that the corresponding in 3 Then there exists (w(() w((), w(tt) t=0
dynamic has the following properties:
1. For any 1 < < min (N,d), there exist least 2 equilibria of rank k;
For a general class of input sequences ()) in (Sd-13N with measure greater than 0, (1) does not converge to rank subspace Dc.

---

The detailed proof is provided in Appendix [F] causal graph, there an inherent chronological order among tokens, in the convergence analysis of each token we assume all tokens preceding it have already converged.
For the sake of simplicity and observing that in
The above result is in direct contrast to the case without LayerNorm (Theorem I 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 [1]. and suggests that it is an important component of the self-attention dynamics in consformers-adiiine 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 matrices W 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
4.4 Discussion
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 aligned. In the counterexample provided in Section 1. the key insight to prevent the second token from aligning with the first token 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 19 14 I7]. Interestingly, as we show in Appendix F.2 the full rank equilibria described in Theorem 3 align with this observation, tokens lie in a narrow region on I While still remains unclear how exactly such a representation geometry is useful for downstream tasks [1] [7],, empirical studies have found that anisotropy is not necessarily harmful for semantic representations and can assist with tasks like clustering [ 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 rank one subspace. However, 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 specific configuration at equilibrium. This observation seems in line with the fact that MLP modules are crucial for the universal eequence-to-sqquene approximation power of transformers 40 [1 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.
observe in simulation that these
5 Numerical Experiments
In this section, we validate our theoretical findings via numerical experiments. Following [1], 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 found in Appendix G
8

---

The effect of attention masks and LayerNorm We first verify the effect of different attention masks and LayerNorm on rank collapse with respect to p(-). We use BERT [T] as the backbone transformer model and consider five different model variants for controlled experiment: self- attention 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 μu( (()) with standard deviations over the 3000 samples in different architectures. We see that in SANs, μu( ( 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, (+) 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.T
54te (n aed)
SANs Selp teed)
SANG UN (irlstees))
SANa L Iriturlees)
Tra Taaalommrr al
2000
CANs Contrathed)
SANe 24to (eretra nes)
500e LN iomtranedi
SANe LV Stp contrattes
Traeformen ipsinisedi
campata shin alcesitdoosss
sumber byyer
Figure 2: Evolution of (1) (in log-log scale) the number of layers increases. Rank collapse happens exponentially for pure attention, despite different attention masks having different conver- gence rates. However, as soon as we solely add in LayerNorm, p(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.
SANS (Initialized): complete
SANs (initialiied)) causal
SANS (initialized): sliding window
30'
207
102
10
307
10a
4000
100
10"
10"
13
10-2
10-2
11-2
10t9
10
10
10"1
107
10
12"
10
10
SANs (pretrained): complete
SANs (pretrained): causal
SANs (pretrained): sliding window
30'
307
10
19
102
162
4000
100
30F
10a
dor 1/64
10-2
dor
doc
10-1
10-
12*2
dor 4096
10 10
101
10 16
207
15 16F
10
number of layers
Figure Evolution of ((((0)) (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 Ap- pendix G.I. Compared with full transformers where pxX) stays relatively stable, there is clear

---

discrepancy. In that case, LayerNorm emerges as 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 do 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 dok with different attention masks. Fig. shows the average values of y (()()) with standard deviations over the 3000 samples with different masks in pretrained SANs. We observe that while smaller temperature term dok alleviates the (initial) rate 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 transformers. Specifically, to capture more fine- mne-grained details than yu(),, 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 [II] GPT2 [30], T5 31 and ALBERT [24] and the average results over 3000 samples are shown Fig. We see that all models exhibit similar behaviors 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 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.
Transformers (pretrained)
Transformers (pretrained)
Transformers (pretrained)
mln. singular value
abs. cosine similarrity
9.8
66 fuill rank
6.9
101
3.6
BERT
10
6.7
GPT2
T5
3/2
ALBERT
10-1
10
161
number of layers
number of layers
number of layers
Figure 4: Evolution of token geometry as the number of layers increases. We see that tokens are indeed able to maintain full rank, while the same time the representations are anisotropic, meaning that they concentrate in a narrow region, as indicated by the average pairwise absolute cosine similarities.
6 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 dy- namics, 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 simul- taneously, which is the case real transformers as well. would be interesting to study how such co-existence of these two characteristics would help with learning tasks whether the full-rank yet -oo---o-ddmmmiinnllll geometry enables efficient learning and generalization while still letting tokens capture meaningful fine- fine-grained details.
Acknowledgments
X. K.W., A.A., and A.J. would like thank Bernard Chazelle for helpful discussions. X. W., A.A., and A.J. were supported by ONR Award N00014-23-1-2200 and Vannevar Bush fellowship from the
10

---

Office of the Under Secretary of Defense for Research and Engineering (USD(R&E)) Y.W. and S.J. were supported by ONR Award N00014-20-1-2022 (MURI ML-SCOPE), NSF AI Institute TILOS (NSF CCF-2112665),, NSF Award 2134108, and the Alexander von Humboldt Foundation.
References
[1] Mira Ait-Saada and Mohamed Nadif. Is anisotropy truly harmful? case study on text clustering. In ACL, 2023.
[2] Josh Alman and Zhao Song. Fast attention requires bounded entries. In NewrIPS, 2023.
[3] Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. ArXiv, abs/1607.0 2016.
[4] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. ArXiv, abs/2004.05150, 2020.
[5] Shaked Brody, Uni Alon, and Eran Yahav. On the expressivity role of layernorm in transformers' attention. In ACI. 2023.
[6] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In NeurlPS, 2020.
[7] Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Recurrent memory transformer. In NeurlPS, 2022.
[8] Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Scaling transformer to Im tokens and beyond with rmt. ArXiv, abs/2304. 1062, 2023.
[9] Xingyu Cai, Jiaji Huang, Yu-Lan Bian, and Kenneth Ward Church. Isotropy in the contextual embedding space: Clusters and manifolds. In ICLR, 2021.
[10] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv, abs/1904, 10509, 2019.
[11] Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In ACL, 2019.
[12] Yihe Dong, Jean- Baptiste Cordonnier, and Andreas Loukas. Attention not all you need: Pure attention loses rank doubly exponentially with depth. In ICML, 2021.
[13] Kawin Ethayarajh. How contextual contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings. EMNLP, 2019.
[14] Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tie-Yan Liu. Representation degeneration problem in training natural language generation models. In ICLR, 2019.
[15] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. mathematical perspective on transformers. ArXiv, abs/2312.100994 2023.
[16] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. In NeurlPS, 2023.
[17] Nathan Godey, Eric Villemonte de la Clergerie, and Benoit Sagot. Anisotropy is inherent to self-attention in transformers. ArXiv, abs/2401.11113,, 2024.
[18] Darald J. Hartfiel. Nonhomogeneous Matrix Products. 2002.
[19] Ali Hassani, Steven Walton, Jiacheng Li, Shengjia Li, and Humphrey Shi. Neighborhood attention transformer. In CVPR, 2023.

---

[20] Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andy Brock, Samuel L. Smith, and Yee Whye Teh. Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation. In ICLR, 2023.
[21] Dongseong Hwang, Weiran Wang, Zhuoyuan Huo, Khe Chai Sim, and Pedro Moreno Mengibar. Transformerfam: Feedback attention working memory. ArXiv, abs/2404.09) 1977., 2024.
[22] Albert Qiaochu Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, L'elio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril Thomas Wang. Timothée Lacroix, and William El Sayed. Mistral 7b. ArXiv, 2023.
[23] Amir Joudaki, Hadi Daneshmand, and Francis R. Bach. On the impact of activation and normalization in obtaining isometric embeddings at initialization. In NeurlPS, 2023.
[24] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-super learning of language representations. In ICLR, 2020.
[25] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
[26] Timothee Mickus, Atie Grönroos, and Joseph Attich. Isotropy, clusters, and classifiers ArXiv, abs/2402.311111 2024.
[27] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurélien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. In NeurlPS, 2022.
(28) Xuran Pan, Tianzuu Ye, Zhuofan Xia, Shiji Song, and Gao Huang. Slide-transformer: Hierar- chical vision transformer with local self-attention. In CVPR, 2023.
[29] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In NeurlPS, 2019.
[30] Alec Radford, Jeff Wu, Rewon Child. David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[31] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 2020.
[32] Aurko Roy, Mohammad Taghi Saffar, Ashish Vaswani, and David Grangier. Efficient content- based sparse attention with routing transformers. Transactions of the Association for Computa- tional Linguistics, 2020.
[33] Han Shi, Jiahui Gao, Hang Xu, Xiaodan Liang, Zhenguo Li, Lingpeng Kong, Stephen M. S. Lee, and James Tin-Ye Yau Kwok. Revisiting over-smoothing in bert from the perspective of graph. In ICLR, 2022.
[34] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Shaolei Du. Scan and snap: Understand- ing training dynamics and token composition in I-layer transformer. In NeurIPS, 2023.
[35] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NewrIPS, 2017.
[36] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the- art natural language processing. In EMNLP, 2020.
12

---

[37] Xinyi Wu, Amir Ajorlou, Zihui Wu, and Ali Jadbabaie. attention-base based graph neural networks. In NeurlPS, 2023.
Demystifying oversmoothing in
[38] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In ICLR, 2024.
[39] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Yan Liu, On layer normalization the transformer architecture. In ICML, 2020.
[40] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of penuence-to-sequenee functions? ICLR, 2020.
[41] Chulhee Yun, Yin-Wei Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. O(n) connections are expressive enough: Universal approximability of sparse transformers, In NeurIPS, 2020.
[42] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. NeurlPS, 2020.
[43] Biao Zhang and Rico Sennrich. Root mean square layer normalization. NeurIPS, 2019.
A Comment about Section 3.3 in Dong et al. [12]
From the mathematical form of LayerNorm in Eq. 33. since LayerNorm normalizes each row of (() to have 2-norm one, the operator that represents this scaling effect D(O) should be multiplied from the left hand side (scaling each row) rather than from the right hand side (scaling each column). 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 (+) being multiplied from the right hand side and can be merged with w would not work.
B Proof of Theorem 1
B.1 Auxiliary results
Lemma 1. Under A1-A3, there exists >> 0 such that A 2 efor all t 2 (j.t) e
Proof. Writing 2
recursively, we get that the token trajectories
x(+1) (t)) x (
(10)
stay uniformly bounded for all t 2 by A3. Then follows from A2 that there exists C' E such that for all & ' 0.
(x
| sc.
sc.
4:12
(x
| sC.
2:
Hence for all i,j E [N].
-C2 s
S C².
This implies that there exists such that A 2 for all (j,i) E.
13

---

Fix a vector (0) E R and a { J in Ag.e. Let Let = A (0) We further denote that
max iEd
oi == M(0)
min
T3 := m(
Red
Note that M monotone non-increasing in i, whereas m monotone non-decreasing
Lemma 2. Let (t
=pm (1) + (1 p)M for 0 s 1. Then for any T E N.
1+T
1+ -1
A )) M
(t-T) 2;
p
II
nn)
m (1) +
1 1 p
II
(11)
and
(+T-11
/14
II i )) m(
(1-T) 2 p
(n)
M +
1 p
(3))
(12)
n=t
n=
Proof. Given xi = =mm!!l) + (1 - p) ((),
we get that
N
at
(((11)
Σ
100
#=1
S A (pmlt) (1) + (1 - ()) + +(1 31-AA (1-AA)))ar
(1-
pA m + (1 (M )M
Subsequently,
N
(++
A
(1+1) (1+1)
zj
j=1
< A
(t+1)
(pA (9) (t) + (1-pA
(! M() + )
1+1
+++1
(())
1 (1-P II (11 MO0
=
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(m) j in Ag.. and be the radius of G. Then for 2 0 and 0 S p S 1,
M(t+kr) (t+kr) < - -pe") (M
-m( l Nzo-
(13)
Proof. Let in E G be a center token and fix > 0 and 0 S < 1. Without loss of generality, suppose < pm (t) + (1 - (1-p) M()
(14)
From Lemma 2 we get that for all T E
N,
[+T-1
t+T-1
)) M
x to
(L+T)
Sp
A
ioio
mlt) +
p
II
A (a)
igio
S Pe ml (1) + - (1-pe-) (1)
Denote V4 as the set of tokens that are exactly h-p neighbors of io:
Vv = (i V dist(io: j) k}
For any i V1, follows that
(++1)
(1) (1)
1100 x + (1 )MM
S (pm (t) + (1 - (!)
+ (1 m)
(t) (1) + (1 pc) 2)
(15)
14

---

Apply Lemma to 15 we further get that for T E N,
((+T) S (0)m (1 - pE-) M (E
Ti
Then let i2 E V2. For simplicity, we still use ii to denote the intermediate neighbor between in and 22. Similarly to the case of i1, we get that
(1+2)
I A (c+1) (e+1)
((ott)) )arl()
Vis
ian Ti + - 3281
ppm (1) + (1- pc) M
+ (1 -1- (1-c)
pe m() (t) + (1 - pe²) M
(16)
and applying Lemma 2 to (16), it follows that for T N22
Fig S Spe"
((+2)
ppmmm" (t))) + (1 pe ) M
Iterating this process for tokens in Va, -.; V,, where the radius of G. get that for any token E g.
(++)
S pe"m!!) (1) + (1 p<") ML0
meaning that
M(++r) S pmmm! (1) + (1 - -pppp) M1
We thus conclude that
M(+r) - ((+)) S pC mo'm" (0) + (1 - pcp) M - m0
= - (1-pe") (M((() m (1)
(17)
For the other case of (4). i.e. Fig > pm) + (1 -p) argument by bounding m(++) from below through (12). This completes the proof of (13)
is obtained using a symmetric
Corollary 2. There exists C(r!O)) > 0 such that
(0) (9) < CC(1-)) ,220
where r the radius of the graph G.
B.2
Proof of Theorem 1
With Corollary we are ready to prove the main theorem. Recall 10 1
we then get that
x (1+11)
(0) (0) (0) (
...W(())
(t)
V
):
(0) ( (W/'0 ...((()))
w (1) A00 (.)
(18)
where W(()) = W.'... ...W())
This means that for any j [d]. there exists Cj 0 such that
x(((2)) -
x -
A x ( mm (A (0) (0) )m)
Wis (A19) 10) m (A (0) x()0
W
max (1))) ...A(0) (0)
min (A10) ...) x.
[E[N]]
[E[N]
$ <C, Σ Wiy (1 -7t//
(19)
$ C, (1- Vm, e [M]
(20)
15

---

where 19 follows from and 20 follows from A3.
Then by 20
((X()) 1|x
T /NIF
| ((())
11 /N12
1=1
(1) X
x 12
2N
mi
j=1
N
07 e) s c(L-
C (1 E
where C = N 2 Di= and C' == vc here.
C Proof of Theorem 2
C.1 Auxiliary results
Lemma 4. Under Al and A2, there exists
0 such that A 2 for all t 122 0, (j. F.
Proof. Note due to the way the dynamical system is defined in 3 for every t 2 0, every token x!! lies within the unit sphere gd-1 Then it follows from A2 that there exists C E such that for all t 2 0,
(x
w
sc,
2:
(xWI12
w SC.
2:112
Hence for all i.j E [N].
C2 S (X (x )as C².
This implies that there exists > such that A 2 e for all (i,j) E E.
Lemma 5. Under the same assumptions as Theorem 2 there exists C₁, C2 > 0 such that C <
D for all 2 0, i e [N].
Proof. We adopt the following notation:
Notation 1. For X E R.d we use Conv(X) denote the convex hull of {X1, ... XN.).
First, since for
W x always lies in gd-1 D/0 2
Then for all 2 0, each row vector of (A (()))))))); lies in the convex cone generated by x())((() where W²(0 = W(( To see this,
(() (1) w(())
fjo Jo::
10=1
N
(t) ((--1)
Σ
A x((=") wi(l
ijo Jo.jo
jo.ji
11=1
Σ
A
(1) ((- 1) (
1(t-2) (0) (0)
in- 1 A
(00)
(0) (()
16

---

Hence all t 1 0,
(1
Σ
(2)
W())
1/ J.:
jml
where /) 2 0, Vj E [M] and Σラ= aif
2 1.
We then make the following observation:
Lemma 6. Under assumption (+), there exists + 0 such that
(x v) > E (X (0)
meaning that all the points of (0) lie in open hemisphere H.
To see this, observe that since (0) is full rank and N d. 0 & (o)v(X( Then the hyperplane
separation theorem states that there exists hyperplane strictly separating and Conv (x (0) letting
us conclude Lemma 6
Since all the points of (0) lie in open hemisphere H. there exists y 0 such that for all E Conv(X (xnv(X(0)) (0) and W orthogonal,
[ W ||~||2 2 7.
Then
(1(10)) (110)x00)) x(00(1)) )a:112
N
2==1 j==
N
ajj
X
(0)
W
1g
j=1
j=
N
(t)
x(0) V7(1)
j=1
Uij
Jj:
Σド= dij
(1)
2
27.
We conclude that D00 < 1/7, for all 2 0.
Writing the layer-wise update rule 3 recursively, we get the following formulation of (
(1) (o) A x (0) (')"...
We focus the infinite product lim;- D A (.) A We first define family of attention matrices as follows:
Definition Let e
We define Ag. to be the set of row-stochustic matrices satisfying the
following conditions:
1. 6 S Aij $1, (f.1) E F()),
2. Ag = 0, (f.t) $ E.
We also define
Dc,,s = (diag(d) C1 Sew d Sew C2).
(21)
and subsequently.
Mg.c.C..C, = (DA: D E Dc,, AE Ag.ch.
(22)
To study the infinite product of matrices from Mg,e,Cr,0s, introduce some necessary concepts:
17

---

Definition 7. A non-negative matrix M is called row allowable if it has positive entry in each of its rows. Definition 8. Consider a sequence of row allowable matrices {M Define the partial product
p())
We say that the sequence { ergodic if there exists a sequence of positive rank one matrices {5t) Ii such that
P Vi,j E [N].
lim
11
(23)
Sij
Our goal is to show the following key result:
Lemma Let e, C1. C2. 0. Given (D(1) A Mg,e,Ci The sequence of partial products { 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 Jem Let ar and be be the smallest and largest entries in M(), respectively. If Et= b. ergodic.
at = oc, then ]2 is
We cannot directly apply Lemma 8 to ( yet unless in the special case of full bidirectional attention such as the used in BERT i.e. G is undirected, complete graph, a, 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) := ( will eventually become strictly positive for strongly connected G. For lo < (1, we denote
p(hitt)
.(to)
Lemma 9, Under A1, 9 is strongly connected, there exist T E N and C3 C4 such that for all t t2 0.
C3 S p(t+T:)) I C4,, V1 MisijsN.
Proof. Note that by A1 and the strong connectivity of G, there exists T' E N, e.g. the diameter of G, and e, C,, C2 such that for all 0.
(cC,)T < Pij s SC, Visijs
(++T:)
Then we apply Lemma to [P(k+ Solk+IyTTTes 1
J k=o and concludes Lemma
As result, from 23 we see that any p()) can now be written in the following form:
p(t) = = (T) + E(t)
where (1)), (()) are positive vectors such that n (T((i) S(t), Without loss of generality, assume ||000]|2 1.
Then follow from 23) that Lemma 10.
E (t)
lim
= 0.
Vi,j E [N].
1-000 u;
18

---

Proof. Given
Eij
E
11 y
where by
23
since also
10g (1) (8)
lim
0,
We get that
Eft) (2)
lim 8-000
0,
u,
and hence
(t(t)
bj
lim
0.
(-000 u/
Thus the 2-norm of the 7th token of ((+11), || 112
[P x ||u (v(0)) x(0) (x00))
E x =, 1, Vi [N].
(24)
Then by Lemma 10
(()) the
lim +
(ul)) T (1)
+
c
(0)
(i(0))
=
lim +
(1000)] [ 100
1000 ub
lim [c)
Note that
T x (1))
= min ((00))
min
where 0min(-) denotes the smallest singular value. Since we assume (0) has full rank, and (1(1) is orthogonal
lim 60000
(1(1)) T 1x
2 Omin ((x0 > 0.
Then given Vi E [N].
lim (1())))x(100000
lim 10000
Vi [N] (t)
2
as result,
lim 61000
(8)
lim
10) 00, Vi, [N].
(25)
This together with Lemma 10 implies that
lim 10000
E
0, Vi [N]. E [d],
and
lim ((((1)) lim (o((0)) (0) + (0)00011 0.
(26)
1-000
t-000
19

---

C.1.1 Exponential convergence rate
Having established the convergence in 226 we then establish the exponential convergence rate. We define that
(X x/))
min [JE[N]
Similar to the derivation of 25 note that
maxx (1x10)
max ||
2
it follows that
lim
(vlt)) x (x(0))((())00))))
max (((())(((1))) < oo.
Since Vi [M],
lim 400
[1000) x w/
= lim
Vi E [M]
8000 u
result,
lim u
(8) lim
t100 u
(0) O. Vi, ks E [N),
meaning that oo x sd-1 for all i [N]
Then given the convergence established in 26 we get that oo (t)) 1, from which there exists to E N such that (+) > 0, Vt 2 to. follows that for 2 to.
((++)) min iJE[M]
(xttr)
x x(((t)) = (x sott..)
N
1++--1)
(1+r-1)
((-r-1)
x/( x(((tr-l)
((+t),
bi
gictr),ds
1=1
(1tr-1) ((--))
1)
(8)
A
bo- Idl (x (x)))
(t)
(r (N
> N + (1-Ne" 22
Thus
1 - - (1
Writing recursively, we get that for k: 2 0,
1 1 +foo+ kr)
< (11 Ne ((1-
Let C= == f0/r Hence for all lllt f20, 0,
l-a (t) < C(1 - Ne 2
(27)
Since by definition, 1 1- ((t) 2 1-
> |x x 10111 1l3/2, it follows that
N
M(x(1)) = ||x(0 11 x((1 /N
Σ
|X 8.
1 X
N 2N
xx(!)
x
==
$ C111Ne
where C² VCN, completing the proof.
20

---

D Proof of Corollary 1
First, since for Conv(X((), W(O x always lies in St-1, D for all i E [N], 2 0.
For all 2 0, since () 2 and g is qtasilytrongly connected, we have that
g((tt)
min [JE[M]
(x((tr), X3= x((tr))
((+t)
(x((t)) 201
x j(in))
N
N Σ
(++-11)
A ((rr-1) x (pitt-l))
((tr) An
jittr),11
=
4=1
N
A 1++r-1)
X g(ttt-1) 1)
A ( 1)
x (((tr-1))
j(ett)
4=1
Σ
(1+r-1)
((++))kk
Ao- My ((tte)
A
(k [e(N]] (1..)))
2 2r + (1 - )0())
Thus
1- (- S
- (1-
6(0).
Writing recursively, we get that for k 2 0,
1 (Kr)
$ (1
- (1-
2r) - -
Let := == 1/(1 2²)). Hence for all 12 0,
(t)) S SC(1 -27)
(28)
Since by definition, 1 1- 21 (X ( 2 |x
xx"" (1) /12/2, it follows that
m(X(1)) = =|x(0 11 /N
Σ
1x 1 T X
x ( 1/2
1
2N
S )(
where C² VCN, completing the proof.
E
Detailed analysis of the illustrative example in Section 4.3.1
First token
Given that
((41))
(11)
+
w,
x((+I)
1,1
we get that
((++1)
+ wl.
((+1)
x0
and thus
X(1 lim
1,2
1-000 x 00.
21

---

Since the sign of x10 1,1 is invariant, depending on its initial sign, we get that
[0, 1] [0. -1]
if x > 0
lim X 2-+00
if
1 < 0
Second token
Assume that the first token converges to [0, 1]. Then
(1-1)
x
(1)
2,2
2,2
(t-1)
X-()
w
+
x
9
x(())
2,1
X X
x x
Solving
gives X2,1 = - 1/w and
x2
X112
if X2,1 < -1/w
199)
x 2.11
(29)
222
otherwise.
XOFF
2,1
To show the convergence, we utilize the following useful property g()): Lemma 11. g is non-decreasing in X2.2/ X2.1.
Proof. Let y $2.2/X2,1 Given the identity that
2
8²
+1,
X((1)
we get that
0
X2.1
1
2
2y.
Oy
X2,1
Hence
X2,1
X2.1V = -2,22
Oy
which means that
0g + X2,1 20.
Oy
Given X2,2/X211 is invariant at X2,1 -1/w (symmetric points A and B in Fig. together
with Lemmal 11 they imply the invariance of the 4 segments on S1, as shown in Fig. Furthermore,
(29 shows which point each region would converge to. Thus it suffices choose X 2,: be on the dark blue or orange segment.
Proof of Theorem 3
F11 Finding the equilibria
Step Without loss of generality, we set that w w 0, for all 2 0. The system of equations corresponds to the equilibrium condition for the attention dynamics with LayerNorm defined in (3) when the attention mask 9 is causal is as follows:
XI.: = (X) = di Xi. Wv
X2. = Fx(X) = de (1. Wv + Wv)
XN.: FN(X) d~ ( Wy + ... WV)
22

---

(0. 1)
B
Figure 5: Convergence analysis of the second token in d = 2.
where d; i Ez= x3, Wr We define that Br d:/i i for i [N]. Then the above system
equation becomes
X1. WVv
X2.: B2 [s X1; X2.
XN;: = BN ( XN-1,: XN.Wv)
after substitution. Then for E [N- 1].
N4+1: Bit1
X1: + Xit1,WY)
1
Xi:
(30)
33
Let k E [N]. If BB 1, then for E 1),
X Xi; + Xi.
(31)
Then if Xx.: is a generalized eigenvector of rank h corresponding to Wy with eigenvalue 1, then from (31) that (X....., }1.) is the Jordan chain generated by X Xx.:, which known to linearly independent.
Hence we Wy to matrix with generalized eigenvector of rank corresponding to eigenvalue 1. For example,
For k 1, Wv can set as lai
2. For l Wv can be as
0
0
0
1
0
0
0
14
0.
0
0
0
w
3. For & = 3, Wv can be set as
0 1
24 2 0.
0
w
0
0
1
23

---

4. For k = N, Wy can be set as
w
0
0
1
w
0
w
Vt 2 0.
(32)
0
For own purpose, consider & N d. Then to solve for equilibrium defined by the condition in (31. for i E N 1-11 1].
0
w
0
0
0
w
0
0
X1+1;
:
X
0
107
0
w
0
0
which means that
[0. 00-wX;+111, W -X [X X1,2, X,,, .... Xta)
(33)
Then setting X1.: = [0, 10,0..., O.....00 0, 1.. iteratively solving (31). we get that an equilibrium satisfying
(30
as follows:
11, = [0, ..., 0,0, 1]
X2.: 0,0, 0,
X3.: 0,0,
(34)
107
XN. (-1)N-1
w
w
20
16
which is valid if w > 1. Since from 31 there is a free coordinate, fact there a N many such equilibria. Hence there exists (W
w:(0) 122 such that the corresponding attention dynamics with LayerNorm has an equilibrium with rank k = N.
K w:(
show that there are equílibria of any rank between and full co-existing for correspond- ing dynamics given the same set of (W(() K w ( consider w((t to be the case of (32). Then notice that to show there exists equilibrium of & beyond the full rank one as derived in Step 1, notice the following:
Step 2
We can set the first - k + tokens to be [0, .... 0, 1). which still satisfies the equilibrium condition.
This implies that 8 1/(N K 1). Then recall (30):
X Bi+1
X Xx.1.
B
Let By-642 1, we get that
XN-k42. W) [0, 0...
0, k ++
24

---

which yields
N - -k+ 1
(N k + 1)
X 2-
Oa-2,
w²
Continue the process by setting 3, for 2 N k 2, we get rank equilibrium desired (there are 2 such equilibria), long w 2 N.
Step 3
Finally to ensure ( W w/ T satisfies A2-A3:
We see that (W w
1 satisfies A2 by construction;
Since this is a dynamics with LayerNorm, applying w fact equivalent to applying Wy /v/lWvll under the dynamics, where the latter has 2-norm at most and thus using Wy satisfies A3.
F.2
Anisotropy of the full rank equilibria
To show that the full rank equilibria X lie in narrow region, we make use of the notion stable rank:
| 3
SRank(XX
|
We would like to show that for any 8 > O, weights can be chosen such that
1 < SRank(X*) < (1 + 6) + O(1/N)
Since for any matrix M. |M 2 M inequality, consider the equilibrium in Then direct calculation yields that
e. the first inequality trivially holds. To show the second
X 112 Tx(x N.
Since by definition,
|X sup
| u1|2,
||+||==1
and thus any ||X2 such that E RN ||ul| serves lower bound for ||X||2. We consider = 0,..,... 0, T e Idd Then it follows that
(N 1) (1
X B
2
which implies that
|
SRank (X*)
X
N
Hence for any 8 > suffices to choose w > such that
w 2
8
E3 Convergence analysis
The convergence analysis for the case d = 2 can be found in Appendix E Here, we deal with the general where 2.
To show that tokens do not converge to an rank one subspace, suffices to show that given certain conditions on the initial input ( there are invariant space of the tokens that least rank 2. To that end, we analyze the convergence of the first two tokens.
25

---

First token
Given that
x
X
+ w,
x 1,1
x((t))
x
we get that
((++1)
x1
+
wt.
x
X1
1,1
1,1
[ Let X r2w, for r2 R. It follows from above that
1.2
(r2 + t)w.
(35)
X (
Then for all 3 2 d.
(1+1)
X
x19
1,1
((+1))
= w
x(1)
+
x-(1)
(36)
1,1
1,1
setting
X
1,1
r,w
(0)
where ry E I. and substituting from the base case
35
to
36
recursively, we get that
x (t)
1-1 ja
1.1 w'-1
rt + 11-1 1
2
J1
-
J1
...
x( (t)
==
0
+ r2 Σ
Σ Σ
11
(37)
h= =1=0=10
which implies that
((+1) x(()
lim
1,1
(38)
x((+1) xx( 5
Further, for 2 3
x (++1)
wx 1, x(1)
x((+1) x(( (1)
x
(++1) 1-2-1
w +
(39)
x 1
w (2)
x
1 x
x 1
1,1-2
1
Then 38 and 39
yields that
x 1.4
lim 100
x0
which let conclude that if x > 0, the first token would converge [0, 0, ..., 0, 1].
Second token Assume that the first token has converged to [0, ..., 0, 1]. Similar the first token, we can show that for 2 < i < d 1, it holds that
lim
00
(
26

---

and
x
wi-l
ri ri
Σ
ji
ri-3
Σ
Ji+
x21
==
1==0حı=
1-1
+ 12 Σ
Σ
(40)
j--=
fi-2=0
1= f1=0
by setting 2 Xx = -,,0"1 where r; > 0. Then note that
( 1)+1)
2,d
wx 2d-- + (1
2,d
+1
((+1)
2,1
(()
2
(1
2,d
1
2,d-1
(++1)
w. 2,d-1 + X2,
22
X21
X
2,
(f)
X
2
2,3
= h
((()) (1)
ru, ro..... rd-1
2,2
We will show a useful monotone property of l(())
Lemma 12. Fix t, ..., rd-1 0, h is non-decreasing in 22.d-2
Proof. Let y = X2.d/ X2.d-1. Given the identity that
X2 + + X?
3²² 1
2
=² + f(r2....
(1
X²
2,1-1
we get that
X2
2
= 2y X2,0-1
θy
Hence
i
X
X2.d-191 X2dd
Oy
which means that
(++1)
0h
X2,0) 2 0. 2
(1
Oy
(((-1)
x
2,d-1
With this, we will show the following claim proving the invariance of the subspace of X2 S -1/w given some initial conditions:
Lemma 13. Suppose x Offori E (d-22) 2}.
X2
(0)
S -1/w such that
rd-1
ra-2
+ 4t
rd-3 +..+ w
rd-1
1
Td-1
12
+
2
1
( 1 S w2 1, (41)
(ra-1)
Vv
then (1)
2d-11 S -1/ for all 1 0.
Proof. When x(())
-1/w, and given the conditions on r2,...; in
2d-1
x-(1)
2dd
= =
w²
1
(()
2.1
+... +
(() 2
2,d-2
(()
(x)(t)
2
2,d-1
=+ w
f(t, r,,.., Ta-2)
27

---

is well defined, where f(-) is monotone non-increasing in t for fixed r1, r1.....dd- > 0. We denote
v 2 1
Note that if ( S -1/w, we have
X2
E
(1 2
then by Lemma 12 get that
X (t)
2,d
E -h r,, n (0 (2..,,r,
h
t, ..... rd-1
X-
2
and since by (40).
(t)
221
2 1
x((+1) x
2
implies that
((+1) x(t)
(+)(0))
2 g(t) s g(t+1)
X (e+1)
x((1
2
which means that
x((+1) ((+1)
2
- (1+1) ( y(++1)]
(1++)
E
and thus ((()) 2 S -1/w for / O.
This let us conclude that for general class sequences (0), (!) does not converge to a rank one subspace as t +
F.4 Further discussion
Directionality of attention mask G In the proof of Theorem we would like to note that G being the causal graph gives key structure to construct the counterexamples. causal graph essentially introduces asymmetries among token interactions where some tokens can attend to some other tokens not verse versa and we exploit such asymmetries constructing the counterexamples. similar construction thus does not naturally apply to the case where G is undirected, e.g. complete graph.
As a result, 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 as the attention mask? If not, what differentiates the case between undirected graph and a directed graph?
G Experiments
Here we provide more details about numerical experiments presented in Section some additional experimental results. All models were implemented with PyTorch 29 and Transformers library [36].
Compute
We ran all of our experiments on CPUs.
28

---

Licenses
Libraries
Wikipedia: MIT license
tranformers [6] Apache License 2.0
Pretrained models
BERT (bert-base-uncased)
Apache License 2.0
GPT2 (openai-communiyyggn))) 30 MIT license
T5 (google-tS/t--base)
31 Apache License 2.0
ALBERT (albert/albert-aase-22
24 Apache License 2.0
G.1 Deeper models
We provide the results for deeper models for the experiment presented Fig. 2 Here, we conduct same experiment in 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, of them suffer from exponential rank collapse. But as soon as we add in LayerNorm, μu( (X(1) no longer converges to 0, as pl-) does not show any decreasing trend even after 128 layers.
SANs (initiaiized), #layer=128
SANs (initiaizeed, #layer=128
Transformers (initiaiized), #layer=128
10
122
10
13a
16-
(X0)
complete
10-1
causal
siide
104
10-2
13-4
sldde-window-uni
100
13P
number of layers
number of layers
number of layers
Figure 6: Evolution of (9) (in log-log scale) as the number of layers increases in 128 layer
models.
G.2 The effect the temperature term dox in initialized transformers
The effect of the temperature term dox 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 more significant under global attention than under sparser masked attention, and in shallower layers than deeper layers.
SANs (initialized): complete
SANs (initialized): causal
SANs (initialized): sliding window
19
15
131
101
M(X)
dox
10-1
10-
30-
dox
64
13'
20-1
number layers
number layers
number of layers
Figure 7: Evolution of ((x()) (in log-log scale) as the number of layers increases in SANs with different dok at initialization.
G.3 The evolution singular values token representations transformers
In this section, investigate the full set of singular values of token representations (+) in trans- formers. For a clear presentation, we randomly select out of the 3000 sequence samples described in Section 5 and calculate all the singular values of (1). The results in randomly initialized transformers presented in Fig, 8
29

---

We see that while the rank is full for ( there is usually one principal component that is much more dominant than the rest. As have discussed in the main text, this is the anisotropic characteristic of token representations- there is strong non-uniformity in different directions. Nonetheless, it important note that the "small" singular values are fact not negligible here the absolute sense-the minimal singular values are never close to zero, as Fig. (middle) shows.
G.4 Initial token geometry
We verify the initial token geometry in transformers. For clear presentation, we randomly select 3 out of the 3000 sequence samples described in Section | and calculate all the pairwise cosine similarities between tokens. The results in randomly initialized transformers presented in Fig. 2 We see that in particular in BERT and ALBERT, the initial pairwise cosine similarities are all non-negative.
Bniiialized,
Gniiaaizedl,
Geidalizedl
1 singalor value
Snouar
(nttblied)
Iinialeed 56g
Innileed 508
Seguar value
Intol
number of ofaayers
number oflayers
number oflayers
Figure 8: The full set of singular values of (1) at each layer in initialized and pretrained transformers.
BERT firetialioed), sec
ВERT (riitiiized), sep
NERT finklalleed), seg
tequerry
wer
2004
da Onkialkeed
ena
200
knnallez)) sca
T5 (nida 5ea
TS inntalized) See
1200
1800
400
1000
ALBERT (eibatadd),
ALBERi (iniiaized),
ALSERT (inibalaedl,
640
TEQQUSFY
Figure 9: Pairwise cosine similarities between tokens calculated from the initial embeddings in randomly initialized and pretrained transformers, respectively.
30

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We confirm that abstract and introduction accurately reflect the paper's contri- butions 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 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 long is clear that these goals are not attained by the paper.
2. Limitations
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 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 (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
31

---

Answer: [Yes]
Justification: We confirm that the paper provides the full set assumptions and a complete and correct proof for each theoretical result.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We confirm that we fully disclose all the information needed to reproduce the experimental results of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
32

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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/
oliicy for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.ccppblle wwwppppguuuiiieeeee /esii///uuide//oeeuuuiiiiii
(CodeSubiisse wwtaaabmmsssooni ssiiiiolcccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We confirm that 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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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, confi-
dence intervals, statistical significance tests, 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 with given experimental conditions).
33

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: Yes]
Justification:
We provide information on the computer resources in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: Yes]
Justification: We confirm that the research conducted in the paper conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Our work is about machine learning theory with the goal to gain deeper insights into the attention mechanism. We believe there no direct societal impacts of the work that are worth discussing at this time.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
34

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include 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 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 dataset.
35

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release assets in this work.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details the dataset/code/moddl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
36


Reviews:
[{"review_id": "TKF4jKS0SQ", "reviewer": "Reviewer_yhKR", "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.", "questions": "I like this work; however, I have two additional questions:\n\nIf local attention alleviates rank collapse, how do the experimental results compare to those of full attention or causal masked attention?\n\nWhich 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": "Please see above.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. The authors provide detailed analysis to demonstrate that attention masks and layer normalization can help address this issue.\n\n3. The experimental results support the authors' analysis effectively.\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your thoughtful comments and positive assessment of our work. After carefully reviewing your feedback, below we provide answers to the comments you raised.\n\n\n> Q1: 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.\n\nThank you for the comment. While the conclusion might seem intuitive, we would like to point out that formalizing it in theory with rigorous mathematical proofs turns out to be very nontrivial — causal masking has been populated by GPTs for years, yet none of the existing theoretical works on analyzing rank collapse [1,2,3,4] in transformers can accommodate transformers with causal masking or local attention. As we discussed in the paper, all those works require having full attention as a *necessary* assumption to derive their theoretical results on rank collapse. To tackle this technical difficulty, we take a novel graph-theoretic approach by formalizing causal masking and local attention through the lens of directed graphs and make use of both graph theory and discrete dynamical systems theory tools. Moreover, our results establish a formal connection between the rate of rank collapse and the graph structure. We hope this novel analysis technique would be a useful tool for the literature and future research in the community.\n\n> Q2: If local attention alleviates rank collapse, how do the experimental results compare to those of full attention or causal masked attention?\n\nThe experimental results can be found in Figure 1, left column. We choose sliding window attention (SWA) deployed in Longformer and more recently Mistral 7B as the representative for local attention. Compared with both full and causal attention, the convergence is clearly slower, which confirms our theoretical results.\n\n> Q3: 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.\n\nFor pure self-attention networks, local attention would be more effective for mitigating rank collapse. Our theory suggests that the rate of rank collapse is directly affected by the diameter of the directed graph (Theorem 1), which is confirmed with numerical experiments (Figure 1).\n\nRegarding the use of local attention, it is indeed popular in vision tasks. See [5,6,7] for references. However, due to the demand for efficiency in long-context, local attention is also getting popular in language tasks. For example, local and sparse attention was populated by Longformer [8] and OpenAI [9] and nowadays LLMs like Mistral 7B use sliding window attention (SWA) [10]. Other popular sparse attention for language tasks include, but not limited to, BigBird [11], Recurrent Memory Transformers (RMTs) [12,13] and StreamingLLM [14].\n\nWe appreciate your questions and comments very much. Please let us know for any further questions.\n\n----------------\n**References**\n\n[1] Dong et al. Attention is not all you need: pure attention loses rank doubly exponentially with depth. In ICML, 2020.\n\n[2] Noci et al. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. In NeurIPS, 2022.\n\n[3] Geshkovski et al. A mathematical perspective on transformers. 2023.\n\n[4] Geshkovski et al. The emergence of clusters in self-attention dynamics. In NeurIPS, 2023.\n\n[5] Liu et al. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.\n\n[6] Pan et al. Slide-Transformer: Hierarchical Vision Transformer with Local Self-Attention. In CVPR, 2023.\n\n[7] Hassani et al. Neighborhood Attention Transformer. In CVPR, 2023.\n\n[8] Beltagy et al. Longformer: The long-document transformer. 2020.\n\n[9] Child et al. Generating Long Sequences with Sparse Transformers. 2019.\n\n[10] Jiang et al. Mistral 7B. 2023.\n\n[11] Zaheer et al. Bigbird: Transformers for longer sequences. In NeurIPS, 2020.\n\n[12] Bulatov et al. Recurrent Memory Transformer. In NeurIPS, 2022.\n\n[13] Bulatov et al. Scaling Transformer to 1M tokens and beyond with RMT. 2024.\n\n[14] Xiao et al. Efficient Streaming Language Models with Attention Sinks. In ICLR, 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DrsFbwgKLx", "reviewer": "Reviewer_7spk", "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.", "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)?\n2. 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.\n3. 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.\n4. Why does this work only provide exponential convergence rates whereas the original rank collapse paper has doubly exponential?\n\nTypo:\n1. $d_N$ not $d_d$ in line 128.", "limitations": "There is not a clear limitations section but the authors do mention a limitation on line 283.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is mostly well written and clear.\n2. 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. \n3. Moreover, the paper shows a more nuanced analysis for the collapse possibilities when a more complete architecture is considered (in particular with layer normalisation) .\n4. 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).\n\n1. 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.\n2. 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).\n3. 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). \n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your positive assessment and insightful comments, which have helped strengthen our work. Below, we provide individual responses the comments you raised.\n\n**Weaknesses**\n\n**W1** We agree with the reviewer that one should be aware of the effect of skip connections when analyzing rank collapse in transformers. Because we were mainly motivated by analyzing the sole effect of attention masks and LayerNorm, we excluded skip connections in our analysis for a controlled study. Building on our work, future research could include analyzing the effects of skip connections as well as positional encoding, among others.\n\nIn the direction of the comment, we have conducted preliminary experiments measuring the evolution of $\\mu(X)$ in SANs with skip connections alone and with both LayerNorm and skip connections. The results (Fig. 1 & 2 in the rebuttal) offer interesting insights:\n\n- For both initialized and pretrained models, adding pure skip connections indeed prevents $\\mu$ from converging to 0, confirming existing theory. However, it seems to make $\\mu$ unstable, particularly in deeper layers (Fig. 2). Compared with full transformers where $\\mu$ stays relatively stable, there is a clear discrepancy.\n\n- Incorporating LayerNorm seems to alleviate this stability issue:\n    - In initialized models, LayerNorm alone effectively prevents rank collapse without causing $\\mu$ to diverge. \n    - In pretrained models, both LayerNorm and skip connections mitigate rank collapse while LayerNorm is key to maintaining tokens at a scale consistent with full transformers, as we discussed in line 265 of the paper.\n\nThese findings underscore the complex interplay between different components in transformers. LayerNorm emerges as a crucial element in mitigating rank collapse and stabilizing token representations while also counteracting potential negative effects of pure skip connections. We plan to further explore this in theory in the future.\n\n**W2** After carefully checking the code, we found that we accidentally loaded the results for initialized models when plotting for the pretrained models. We feel indebted to the reviewer for spotting this mistake. The results for pretrained models in Fig. 2 & 3 are now reported in Fig. 1 & 3 of the rebuttal, respectively. \n\n- For the effect of the temperature, the error bars are indeed very narrow for initialized models. One hypothesis for the effect to disappear after 5 or 6 layers at initialization is that in deeper layers as rank collapse happens, tokens align with each other and the temperature naturally has much less effect (especially given random $W_K$ and $W_Q$ cannot account for the increased similarities between tokens). We leave the thorough investigation for future work. As for pretrained models, the effect of temperature now can be observed in all layers and it is significant for all masks.\n\n- Due to the space limit, the results for 128 layer initialized models were provided in Appendix G.1. They exhibit the same trend as 12 layer models. We have also included 128 layer versions of Fig. 2 & 3 for initialized models over 3000 samples in the rebuttal.\n\n- For pretrained SANs, we follow Dong et al. (their official implementation: https://github.com/twistedcubic/attention-rank-collapse) by taking the existing pretrained models and only considering the self-attention layers.\n\n**W3** Dong et al. claims that self-attention with LayerNorm *cannot* prevent rank collapse. These counterexamples are sufficient to show that self-attention with LayerNorm *can* prevent it from happening. We do not claim that they are the only possible weights that work. It merely establishes that the set of such desirable designs is not empty.\n\nFig. 1 in the rebuttal shows that just adding LayerNorm to SANs (no any other components such as skip connections) prevents rank collapse in initialized models, while in pretrained models LayerNorm mitigates the issue together with other modules and helps token representations maintain a scale consistent to full transformers. This empirically validates our theory.\n\n**W4** We appreciate this detailed feedback. We would like to clarify a misunderstanding here:\n\n- The counterexample works not because of zero activations in the center node (token 1 in this case). *Even with zero activations in the center node, rank collapse can still occur*: \n\n    - In Section 4.3.1, when token 2 is initialized on the olive segment in Fig. 1 of the paper, rank collapse still happens.\n    - In the case without LayerNorm, suppose we have the causal mask, and $X_{1,:}^{(0)}$ is a vector with a zero in its elements, and all $W_V^{(t)}$ are identity with $W_K^{(t)}$ and $W_Q^{(t)}$ satisfying A2). Then token 1 would stay at its initial value (so there is a deliberate zero activation from the start) and this would be the case where “attention allows other tokens to see the center node's activations but the zero activations are not being transferred to other tokens” --- yet rank collapse would still happen here (Theorem 1).\n\n- *The right intuition here can be better elucidated from a dynamical system perspective.* Without LayerNorm, rank collapse happens as the center node attracts all tokens to align. In the counterexample, the key insight to prevent token 2 from aligning with token 1 (assuming token 1 has converged to $v_1$) is that token 2's updated representation must cancel token 1's attraction, i.e. $X^{(t)}\\_{2,:}W$ needs to have a component negating $v_1$. In the example, $W$ satisfies $Wv_2 = v_1+v_2$ for some $v_2$, so a $-v_2$ component in $X^{(t)}_{2,:}$ can negate token 1's effect ($v_1$). The crucial role of LayerNorm here is to stabilize this process by readjusting token scales to ensure that the cancellation still persists in subsequent updates. This key scaling effect of LayerNorm can also be observed from the newly added experiments, as we discussed in W1.\n\n----\n**Due to the character limit, we answer the reviewer's questions in a comment below.**", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NLw69501RZ", "reviewer": "Reviewer_ivrv", "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:\n\t- 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.\n\t- local attention mask or focus attention can slow down the collapse.\n\t- layernorm plays an important role in self-attention dynamics.\n- This paper provides extensive analysis and numerical experiments on the self attention dynamics, Expressivity and Versatility.", "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": "None", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n- 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.\n- This paper shows an interesting counterexample of collapse and shows that self-attention dynamics can process a rich set of sequences stablely.\n- 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. \n- 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? \n- The numerical experiments are somewhat weak as the results are obtained from only 32 samples.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our paper and providing constructive feedback. In line with the reviewer’s suggestion, we have rerun all the numerical experiments, validating our theoretical findings on 3000 examples. The results are provided in the rebuttal pdf file. In what follows, we provide detailed responses to the rest of the comments raised by the reviewer.\n\n> Q1: Application of the conclusion, I wonder can the findings in this paper guide some design choices in the future transformer models.\n\nThank you for the question. Here are some ideas along our theoretical results:\n\n- The graph-dependent rate of rank collapse established in the paper can be used as a principle for attention mask design. It is, however, important to also note that the design cannot be solely based on this result. One also needs to consider other implications of the mask, including its effects on universal approximation property [1]. We have discussed this remark in in line 178-184 of the manuscript. \n\n- Our result could also further explain the role and importance of LayerNorm in transformers. It clarifies a misconception popularized by a previous work [2] that “layer normalization plays no role” for rank collapse in transformers. We rigorously prove that layer normalization can in principle prevent rank collapse and stabilize token representations in transformers, enabling the model to utilize depth more efficiently.\n\n> Q2: How do the findings connect to (or explain) the strong performances of existing LLMs? It claims that sparse attention can slow down collapse, why do most existing LLMs not use sparse attention?\n\n- Existing implementations of LLMs utilize LayerNorm, which is in line with our results on the role of LayerNorm in preventing rank collapse. This could justify why this module is preserved in the design of pioneering LLMs and improved variants.\n\n- As we discussed in the response to Q1, rank collapse is not the only issue one needs to consider when designing the attention mask. One also needs to consider aspects such as the universal approximation power of masked attention. How to balance the trade-off for more powerful LLMs is a vital research direction which we are currently exploring.\n\n- In fact, while many existing LLMs do not use sparse attention, sparse attention is used in practice and is gaining popularity due to the demand for efficiency. For example, sparse attention was populated by Longformer [3] and OpenAI [4] and nowadays LLMs like Mistral 7B use sliding window attention (SWA) [4]. Other popular sparse attention models include, but not limited to BigBird [5], Recurrent Memory Transformers (RMTs) [6,7] and StreamingLLM [8].\n\n- Besides language tasks, sparse attention is also common in vision transformers. Examples of the use can be found in the following papers [9,10,11].\n\nWe will update our manuscript for an improved literature review on sparse attention.\n\n> Q3: The numerical experiments are somewhat weak as the results are obtained from only 32 samples.\n\nThank you for your constructive suggestion. Our original sample size 32 was chosen according to the experiments in Dong et al. [2]. In line with the reviewer’s suggestion, we validate our theoretical findings on 3000 examples instead. The results are provided in the rebuttal pdf file. They are still similar to the original results; however, as suggested by the reviewer, the larger sample size makes them more reliable.\n\nWe appreciate your questions and comments very much. Please let us know if there are any further questions.\n\n-----------------\n**References**\n\n[1] Yun et al.  O(n) connections are expressive enough: Universal approximability of sparse transformers. In NeurIPS, 2020.\n\n[2] Dong et al. Attention is not all you need: pure attention loses rank doubly exponentially with depth. In ICML, 2020.\n\n[3] Beltagy et al. Longformer: The long-document transformer. 2020.\n\n[4] Child et al. Generating Long Sequences with Sparse Transformers. 2019.\n\n[5] Jiang et al. Mistral 7B. 2023.\n\n[6] Zaheer et al. Bigbird: Transformers for longer sequences. In NeurIPS, 2020.\n\n[7] Bulatov et al. Recurrent Memory Transformer. In NeurIPS, 2022.\n\n[8] Bulatov et al. Scaling Transformer to 1M tokens and beyond with RMT. 2024.\n\n[9] Xiao et al. Efficient Streaming Language Models with Attention Sinks. In ICLR, 2024.\n\n[10] Liu et al. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.\n\n[11] Pan et al. Slide-Transformer: Hierarchical Vision Transformer with Local Self-Attention. In CVPR, 2023.\n\n[12] Hassani et al. Neighborhood Attention Transformer. In CVPR, 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## Response to all reviewers


We would like to thank the reviewers for carefully reading our paper and giving insightful comments and constructive feedback. 
We are glad that our work was recognized as “interesting”, “rigorous” (Reviewer ivrv) and “offering valuable contributions to the field” (Reviewer yhKR), including “a graph-theoretic approach to analyze self-attention” (Reviewer ivrv) and “new insights into the rank collapse phenomenon” (Reviewer 7spk). We are also encouraged that the reviewers found our paper "well-written" (Reviewer 7spk).

We have provided detailed responses to each of the reviews separately. We also include additional numerical results in the rebuttal pdf under the same experimental setup as in our paper. Specifically, we increase the sample size of the experiments to 3000 in all experiments and consider the setting of 128-layer randomly initialized models when measuring the effects of different attention masks, LayerNorm and the temperature term.

The authors.

Meta-review: 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.
Decision: Accept (poster)