Title: On the Role of Attention Masks and LayerNorm in Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: attention mechanism

transformers

layer normalization

deep learning theory

dynamical systems
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.

Xinyi Wu 1

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.

1 Introduction

The celebrated attention mechanism has proven highly effective in the architecture of transformers [35], which serve as the key building block of modern foundation models including large language models. From a theoretical perspective, understanding the underlying mechanism of transformers and attention in general has become pivotal for elucidating existing models and paving the way for developing more powerful future models [12, 15].

Transformers are known to exhibit the rank collapse phenomenon 1 [12, 15, 20, 27, 33, 37], which refers to the observation that increasing the number of self-attention layers leads to homogeneous token representations. To gain more insight into the rank collapse issue and better understand the effects of self-attention in multi-layer models, it is essential to study the long-term behavior of tokens under self-attention dynamics [12, 15, 16].

However, many existing studies do not take into account architectural components that are commonly used in practice. For instance, first, the theoretical analysis of long-term self-attention dynamics often assumes that attention is fully bidirectional -that is, all tokens are allowed to attend to all other tokens in the sequence [12, 15, 16, 27], which only applies to certain attention mechanisms

such as the one deployed in the BERT family [11, 24]. The vast majority of popular transformer architectures used nowadays in language models, including the GPT models [6, 30], use the causal attention masks where tokens are only allowed to attend to preceding tokens, or have sparse attention structure (sparse attention) [4, 8, 10, 21, 22, 32, 38, 41, 42], where attention is restricted to local interactions between tokens and their chosen neighbors. This limits the practical applicability of the theoretical results developed in [12, 15, 16, 27], as they heavily rely on the key assumption that attention is fully bidirectional.

Second, layer normalization (LayerNorm) is another inherent component of transformers. Dong et al. [12] put forth a hypothesis that LayerNorm plays no role in preventing rank collapse in self-attention networks. This hypothesis is then partially validated by a continuous-time analysis of self-attention dynamics conducted in Geshkovski et al. [15]. The analysis, which incorporates LayerNorm, shows the exponential convergence of tokens to a common point on the unit sphere. However, this result relies on a strong assumption that all the value matrices are the identity matrix. In contrast, for multilayer perceptrons (MLPs), Joudaki et al. [23] show that LayerNorm improves the isometry of the representations and can prevent rank collapse in that setting. It thus remains a question whether it is truly the case that LayerNorm could not have a similar effect in transformers under more general assumptions due to self-attention.

Hence, in this work, we rigorously study the effect of attention masks and LayerNorm on the token dynamics and rank collapse in transformers. The main questions we address are as follows:

Can attention masks alleviate rank collapse of tokens under self-attention? If so, what type of attention mask would be the most effective?

Can LayerNorm alleviate rank collapse of tokens under self-attention? If so, what long-term behavior of tokens would LayerNorm lead to?

We answer these questions through a rigorous analysis of the self-attention dynamics. Notably, unlike some previous works [15, 16], which regard self-attention as a continuous-time dynamical system, we view self-attention as a discrete-time dynamical system, more closely resembling the architecture used in practice. Furthermore, our analysis extends the results to more general masked attention by making novel use of a graph-theoretic approach and incorporating advanced results on infinite products of inhomogeneous non-negative matrices that may be of independent interest.

In summary, we make the following contributions:

We establish that with pure self-attention, the exponential convergence of tokens to a common representation holds for a broad class of attention masks, accounting for the causal mask and a wide class of sparse attention patterns such as the sliding window [4, 42]. The key property that leads to the exponential rank collapse is a token that serves as a common 'context' for all the other tokens in the sequence to directly or indirectly attend to. Our results also show that local attention can slow down the rate of rank collapse, suggesting its potential advantage over full bidirectional attention from an expressivity perspective at finite depth.

We show that with LayerNorm, when the value matrices are orthogonal, the exponential convergence of tokens to a common point on the unit sphere holds for a broad class of attention masks. Nonetheless, by constructing nontrivial counterexamples, we prove that the self-attention dynamics with LayerNorm can simultaneously have a rich set of equilibria of any possible rank ranging from one to full. Moreover, we rigorously establish that self-attention with LayerNorm, together with proper choice of value matrices, can provably prevent complete collapse of tokens to a rank one subspace for a generic class of input sequences.

2 Related work

Analysis of self-attention dynamics Understanding the attention mechanism is a pivotal step towards understanding the inner workings of transformer-based models. From a dynamical systems perspective, one could abstract the forward pass of the model as tokens undergoing a nonlinear, time-varying dynamics determined by the self-attention mechanism and other parameters of the model. Specifically, Dong et al. [12] first show that as the number of self-attention layers increases, tokens inevitably suffer from exponential rank collapse, while Wu et al. [37] establish a similar result for the attention mechanism in graph neural networks, taking additionally the layer-wise nonlinear

activation functions into account. Other works [15, 16] take a continuous-time approach and study more fine-grained clustering behaviors of the dynamics in transformers.

The effect of LayerNorm in transformers LayerNorm is one of the most commonly used normalization techniques in modern neural networks [3] and has become an inherent component of transformers [35]. To better understand its role in transformers, Xiong et al. [39] study it from an optimization perspective and show that LayerNorm stabilizes the gradients during the backward pass. In addition, Brody et al. [5] find that LayerNorm plays a crucial role in improving the expressivity of the attention layer by making it easier for the model to compute the most frequent token in the input and avoid the problem of 'unselectable' keys. Most relevant to our work, Dong et al. [12] mention the effect of LayerNorm in terms of mitigating the rank collapse issue in transformers. The paper makes a hypothesis that LayerNorm has no effect for token rank collapse. The argument is based on a heuristic observation (see Appendix A for a detailed discussion) which is at odds with the case of simpler models such as MLPs where LayerNorm is shown to be pivotal in addressing a similar rank collapse problem [23]. It thus remains an open question what effect LayerNorm would have on rank collapse in transformers.

Sparse and local Attention While many existing transformer models and LLMs do not use sparse attention, sparse and local attention is gaining popularity due to the demand for efficiency, particular for long-context tasks. For example, sparse attention was populated by Longformer [4] and OpenAI [10] and nowadays popular LLMs like Mistral 7B [22] use sliding window attention by default. Other popular sparse attention models include, but are not limited to BigBird [42], Recurrent Memory Transformers (RMTs) [7], and Streaming Attention [38]. Besides language tasks, sparse attention is also common in vision transformers [19, 25, 28].

3 Problem Setup

Notation Let ∥ · ∥ 2 , ∥ · ∥ F be the 2 -norm and Frobenius norm, respectively. We use the shorthand [ n ] := { 1 , . . . , n } . We denote the all-one vector of length N by 1 ∈ R N . For a matrix M , we denote its i -th row by M i, : and its j -th column by M : ,j .

Throughout the analysis in the paper, we formalize the attention mask to be a directed graph G . Formally, we represent a directed graph with N nodes by G and let E ( G ) be the set of directed edges of G . If there is a directed edge going from j to i in G , i.e. ( j, i ) ∈ E ( G ) , for the attention mechanism it means that token j serves as a direct context for token i or token i attends to token j . The set N i of all neighbors of node i is then { k : ( k, i ) ∈ E ( G ) } .

Furthermore, we will be using the following graph-theoretic terminology:

Definition 1 (Reachability) . We say a node v is reachable from u in a directed graph G if and only if there is a directed path ( u, n 1 ) , ( n 1 , n 2 ) , ..., ( n k , v ) from u to v .

Definition 2 (Strongly Connected) . A directed graph G is said to be strongly connected if and only if any two distinct of nodes are reachable from each other.

Definition 3 (Center Node) . A node v 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 G has at least one center node.

Definition 5 (Radius) . The radius of a quasi-strongly connected graph is defined to be the longest distance from a center node to any node in the graph. If there are multiple center nodes, then it is the smallest value among them.

3.1 (Masked) Attention Mechanism

Given token representations X ∈ R N × d , the raw attention score matrix is computed as

where W Q , W K ∈ R d × d ′ are the query and the key matrix, respectively, and √ d QK is a temperature term to control the scale of raw attention scores. To enforce a masked attention, we create a sparse

attention matrix A ∈ R N × N based on R whose sparsity pattern is specified by a directed graph G : we normalize R ij among all allowed token attention interactions ( k, i ) ∈ E ( G ) such that

3.2 LayerNorm

Given token representations X ∈ R N × d , LayerNorm subtracts the mean across different columns in each row and then scales each row to have a unit 2 -norm. In this work, we consider LayerNorm to only perform the scaling operation, which is a common assumption in theoretical analyses of the attention mechanism [15, 34] 2 . Mathematically, let D = diag( d 1 , d 2 , ..., d N ) where d i = 1 / ∥ X i, : ∥ 2 for all i ∈ [ N ] , 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 as

where W ( t ) V ∈ R d × d ′ is the value matrix. For simplicity, throughout the paper, we assume that d = d ′ and G ( t ) = G , i.e. the attention mask is the same for all the attention layers. Yet the results can be easily generalized to the case where masks are time-varying and satisfy similar regularity conditions.

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 µ ( · ) : R N × d → R ≥ 0 :

This measure is mathematically equivalent to the measure res( X ) = arg min x ∈ R d ∥ X -1 x ⊤ ∥ F used in [12], but the form in (4) is easier to work with in the general analysis and more direct to compute. Another advantage of our formulation is that it clearly demonstrates that Theorem 1 and Theorem 2 are not dependent on the specific choice of µ ( · ) : these results apply to any Lipschitz µ ′ ( · ) with a Lipschitz constant L such that µ ( X ) = 0 if and only X i, : = X j, : , ∀ i, j ∈ [ N ] , as we can use the formulation to directly derive that

Finally, we adopt the following assumptions in our analysis:

A1 G contains self-loops. i.e. ( i, i ) ∈ E for every token i ∈ [ N ] .

A2 There exist constants C ∈ R such that max t ∈ N {∥ ∥ ∥ W ( t ) Q ∥ ∥ ∥ 2 , ∥ ∥ ∥ W ( t ) K ∥ ∥ ∥ 2 } ≤ C .

A1 ensures that every token has a neighbor so that the masked attention computation is well-defined for every token in every layer, while A2 assumes that the key and query weight matrices are bounded, which is key for efficient attention computation in practice [2].

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 X ( t ) for all t ≥ 0 even without LayerNorm, we further assume that

Then with general attention masks G , there remains a strong connection between tokens via attention, and the token representations collapse exponentially to rank one.

Theorem 1. Consider the self-attention dynamics without LayerNorm defined in (2) . Under A1 -A3 , if G is a quasi-strongly connected graph, then there exists ϵ > 0 where for all t ≥ 0 ,

As a result, a rank collapse of tokens happens exponentially with respect to µ ( · ) , i.e. there exists C > 0 such that

where r is the radius of G , meaning that tokens converge to a common vector exponentially.

The detailed proof is provided in Appendix B. The above result suggests that with pure self-attention, as long as there is a token which all other tokens in the sequence can directly or indirectly attend to over a fixed number of layers, exponential rank collapse of tokens to a common vector is guaranteed. In particular, it generalizes the main result in [12] from G being a complete graph to a much more general class of attention patterns: the attention pattern G only needs to be quasi-strongly connected, meaning that the result applies to general attention masks used in practice including the causal mask used in decoder-only models such as the GPT family [6, 30], or sparse attention patterns deployed in many efficient transformer models [4, 10, 21, 32, 42]. We discuss a few interesting implications below.

Local vs. global attention The exponential rate (1 -ϵ r ) 1 /r is monotone in the graph radius r . This means that rank collapse should be slower for graphs with larger radius r . Our result thus indirectly supports the use of local attention patterns [4, 42], which not only make the attention computation more efficient (what those works are originally motivated by), but also implicitly alleviate the rank collapse issue.

Focused vs. uniform attention In addition, the exponential rate is monotone decreasing in ϵ , which means that rank collapse is slower with smaller ϵ . One can interpret ϵ as how "focused" attention is distributed among reachable tokens, as ϵ is maximized when attention happens uniformly among reachable tokens. Besides applying attention masks and restricting the number of reachable tokens, another way to control how focused attention would be is through the temperature term d QK . As larger values of d QK would make the attention allocation among reachable tokens more even, they should make rank collapse happen faster across layers.

Trade-off between rank collapse and universal approximating power Finally, for strongly connected graphs, the above result also reveals a trade-off between universal function approximation power and the rate of rank-collapse. Yun et al. [41] show that transformers with strongly connected graph masks are sequence-to-sequence function universal, yet with a mask G they need at least the diameter of G layers to achieve the full sequence-to-sequence function approximation property. This implies that masks with smaller diameters (and thus smaller radii, as radius ≤ diameter ≤ 2 radius) are more efficient in terms of function approximation power, yet they are more prone to rank collapse.

Remark 1. Since the analysis in [12] fundamentally relies on the shift-invariant property of softmax( · ) , it is necessary that all the tokens are allowed to attend to all the other tokens for their proof to work. On the contrary, we leverage a different graph-theoretic approach to exploit the common structure of the attention matrices to obtain a general result for masked attention.

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 a common vector can still happen for certain classes of value matrices.

Theorem 2. Consider the self-attention dynamics with LayerNorm defined in (3) . Let G be a strongly connected graph. Assume A1 -A2 , and that W ( t ) V is orthogonal for all t ≥ 0 , and in addition, the initial input X (0) satisfies that

Then there exist C > 0 , ϵ > 0 such that Nϵ < 1 and

where r is the radius of G , meaning that tokens converge to a common point on S d -1 exponentially.

The detailed proof is provided in Appendix C. The result can be seen as a generalized discrete version of Theorem 4.1 in [15]. Notably, our analysis is based purely on advanced linear algebra tools: infinite products of non-negative matrices and their ergodicity, and can account for time-varying weights and general attention masks, as opposed to fixed W K , W Q , W V over time and G being complete (which is the case in [15]). One way to satisfy the condition ( ∗ ) on the initial input X (0) is to require N ≤ d and to initialize tokens uniformly randomly on S d -1 , then the condition hold almost surely. This is how the condition is dealt with in [15].

Note that the condition ( ∗ ) implies that there exists v ∈ S d -1 such that ⟨ X (0) i, : , v ⟩ > 0 for all i ∈ [ N ] by either the hyperplane separation theorem or Farkas' lemma (see Lemma 6 in Appendix C). If the initial token geometry satisfies a stronger condition than the above, then ( ∗ ) is no longer necessary and Theorem 2 even directly generalizes to quasi-strongly connected graphs G . We define ϕ ( t ) := min i,j ∈ [ N ] ⟨ X ( t ) i, : , X ( t ) j, : ⟩ , indicating the minimal cosine similarity between tokens. If the cosine similarities are non-negative for all pairs of tokens initially, then the rank collapse happens exponentially, as long as G is quasi-strongly connected.

Corollary 1. Consider the self-attention dynamics with LayerNorm defined in (3) . Let G be a quasi-strongly connected graph. Under A1 -A2 , if W ( t ) V is orthogonal for all t ≥ 0 and ϕ (0) ≥ 0 , then there exist C > 0 , ϵ > 0 such that Nϵ < 1 and

where r is the radius of G , meaning that tokens converge to a common point on S d -1 exponentially.

Full mask vs. causal mask We can refine Corollary 1 by specifying the number of center nodes n in the mask G , then the upper bound for the exponential rate would be (1 -nϵ 2 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.

Post-LN vs. Pre-LN The definition of LayerNorm in (3) follows from the original transformer paper [35] and nowadays it is referred as post-LN [39]. An alternative use of LayerNorm in many LLMs, where LayerNorm comes before self-attention, is called pre-LN and can be written instead as

Note that Theorem 2 and Corollary 1 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 a nontrivial counterexample with only LayerNorm such that for a general class of input sequences, tokens converge to an equilibrium where rank collapse does not happen. Notice that for a transformer model to be practical, it is important that it can prevent rank collapse for a general class of input sequences rather than a specific input sequence. We will then show a general result stating that, with LayerNorm and proper choice of value matrices, the self-attention dynamics can have a rich set of equilibria with any possible rank between one and full. Moreover, for a general class of input sequences, tokens provably do not converge to a rank one subspace under the resultant dynamics.

We further let

(0

,

-

1)

(0

,

1)

(0

,

-

1)

With LayerNorm

First token

Second token

No LayerNorm

Both tokens

(0

,

1)

A

B

4.3.1 Illustrative Counterexample

1

B For simplicity of illustration, we consider N = 2 , d = 2 , and G to be the causal mask. Then let W ( t ) K = W ( t ) Q = 0 , which leads to the attention matrices

for w ∈ R . Without loss of generality, fix w > 1 . Then a careful analysis shows that, depending on its initial position, the first token will either converge to (0 , 1) or (0 , -1) . Suppose the first token converges to (0 , 1) . Then the convergence of the second token as t →∞ is illustrated in Fig. 1, where

Note that due to the scaling effect of LayerNorm, any scaled version c ( t ) W ( t ) V , c ( t ) > 0 of W ( t ) V works equivalently here.

Remark 2. For any orthogonal Z ∈ R d × d , W Z = Z ⊤ WZ works equivalently in this example, and the resulting token trajectories follow X ( t ) Z .

This nontrivial counterexample suggests that with the LayerNorm dynamics in (3), there are proper choices for W ( t ) K , W ( t ) Q , W ( t ) V matrices that can prevent tokens from collapsing to a rank one subspace, for a nonzero measure set of input sequences.

4.3.2 Generalization of the Counterexample

We conclude this section with the following statement generalizing the previous illustrative example: with a proper choice of weight matrices, tokens under the attention dynamics with LayerNorm can simultaneously have equilibria of any rank between one and full. Moreover, for a general class of input sequences, tokens provably do not converge to a rank one subspace.

Theorem 3. Let G be the causal graph and consider the attention dynamics with LayerNorm defined in (3) . Then there exists { W ( t ) Q , W ( t ) K , W ( t ) V } ∞ t =0 satisfying A2 -A3 such that the corresponding dynamic has the following properties:

For any 1 ≤ k ≤ min { N,d } , there exist at least 2 k equilibria of rank k ;

1

For a general class of input sequences X (0) in ( S d -1 ) N with measure greater than 0 , X ( t ) does not converge to a rank one subspace as t →∞ .

1

The detailed proof is provided in Appendix F. For the sake of simplicity and observing that in a causal graph, there is an inherent chronological order among tokens, in the convergence analysis of each token we assume all tokens preceding it have already converged.

The above result is in direct contrast to the case without LayerNorm (Theorem 1), where all input sequences eventually result in complete rank collapse to a one-dimensional subspace. It also stands in contrast to the hypothesis that LayerNorm plays no role in mitigating the collapse [12], and suggests that it is an important component of the self-attention dynamics in transformers-adding LayerNorm fundamentally changes the behavior of the underlying system. Compared to the case without LayerNorm, first, LayerNorm guarantees that the system never diverges, no matter what the value matrices are. Second, and more importantly, attention with LayerNorm leads to a surprisingly more expressive, versatile dynamics than the system without it-composition with different value matrices W ( t ) V can result in different long-term token behaviors: in some cases, tokens completely collapse to a single point (Theorem 2), while in others tokens can retain higher rank asymptotically (Theorem 3).

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 be aligned. In the counterexample provided in Section 4.3.1, the key insight to prevent the second token from aligning with the first token is that the updated representation of the second token must generate a component canceling out the attraction from the first token, which acts as a repulsion. The crucial role of LayerNorm here is to stabilize this cancellation process by readjusting the scale of tokens to ensure that the cancellation persists in all the updates, which would not be the case if there is no LayerNorm.

Anisotropy of token embeddings Analyzing and understanding the token geometry in transformers has long been of interest to the NLP and general ML community. In particular, empirical findings suggest that contextual token embeddings generated by transformers are anisotropic, meaning they tend to concentrate in a narrow region [1, 9, 13, 14, 17]. Interestingly, as we show in Appendix F.2, the full rank equilibria described in Theorem 3 align with this observation, as tokens lie in a narrow region on S d -1 . While it still remains unclear how exactly such a representation geometry is useful for downstream tasks [1, 9, 14, 17], empirical studies have found that anisotropy is not necessarily harmful for semantic representations and can assist with tasks like clustering [1, 26].

Stability of the equilibria Finally, we would like to note that due to the combinatorial nature of the analysis with increasing dimensions, we do not prove the exact convergence to specific equilibria beyond not converging to a rank one subspace. However, we observe in simulation that these equilibria are indeed stable in certain directions, despite their regions of attraction being relatively small. This suggests that while LayerNorm improves the expressivity of pure self-attention dynamics, the resulting expressive power might still be limited in a way that it would be difficult for a generic input sequence to reach a specific configuration at equilibrium. This observation seems in line with the fact that MLP modules are crucial for the universal sequence-to-sequence approximation power of transformers [40, 41]. To achieve maximal expressivity, and in particular, to be able to move tokens freely around the whole space (such as in the case for transformers with MLPs), nonlinear power from MLPs would be necessary.

5 Numerical Experiments

In this section, we validate our theoretical findings via numerical experiments. Following [12], we randomly select 3000 samples of 128-token excerpts (based on the BERT tokenizer) from Wikipedia using the Wikipedia API in Python. More experimental details and additional results can be found in Appendix G.

The effect of attention masks and LayerNorm We first verify the effect of different attention masks and LayerNorm on rank collapse with respect to µ ( · ) . We use BERT [11] as the backbone transformer model and consider five different model variants for a controlled experiment: selfattention network (SAN) in which the model has self-attention layers; SAN with skip connections; SAN with LayerNorm where in each layer there is self-attention followed by LayerNorm; SAN with both skip connections and LayerNorm where LayerNorm is put after skip connections; and finally the standard transformer with all the components. For attention masks, we select four representative types: the complete graph, the causal graph, the sliding window (tokens can only attend to the token right before and after and themselves) and the uni-directional sliding window (tokens can only attend to the token right before and themselves). Fig. 2 shows the average values of µ ( X ( t ) ) with standard deviations over the 3000 samples in different architectures. We see that in SANs, µ ( X ( t ) ) converges to zero exponentially for all attention masks, yet more local attention masks slow down the convergence rate. Furthermore, in randomly initialized models, right after we add in LayerNorm, µ ( X ( t ) ) no longer converges to zero and can be maintained at a stable level significantly above zero. In pretrained models, LayerNorm helps prevent the issue together with other components such as skip connections and stabilize the representations. Results in 128 -layer initialized models also exhibit the same trend, and can be found in Appendix G.1.

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

(X)

SANs (initialized)

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs + LN (initialized)

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs + LN + Skip (initialized)

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

Transformers (initialized)

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

(X)

SANs (pretrained)

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

10 0

10 1

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs + LN (pretrained)

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs + LN + Skip (pretrained)

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

Transformers (pretrained)

complete

causal

slide-window

slide-window-uni

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

(X)

SANs (initialized): complete

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs (initialized): causal

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs (initialized): sliding window

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

(X)

SANs (pretrained): complete

10 0

10 1

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs (pretrained): causal

10 0

10 1

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs (pretrained): sliding window

dQK

= 1/64

dQK

= 1

dQK

= 64

dQK

= 4096

The complex interplay between different components in transformers Previous works have shown that skip connections can help combat rank collapse in transformers. Yet in Fig. 2, skip connections seem to make µ ( X ) unstable, particularly in deeper layers (for 128-layer results, see Appendix G.1). Compared with full transformers where µ ( X ) stays relatively stable, there is a clear

discrepancy. In that case, LayerNorm emerges as a crucial element in mitigating rank collapse and stabilizing token representations while also counteracting potential negative effects of pure skip connections. This underscores the complex interplay between different components in transformers.

The effect of the temperature term Next, we investigate the effect of the temperature term d QK in the attention calculation. Since the amount of focus of the attention is affected by both the attention mask and the temperature term, we investigate the effect of d QK with different attention masks. Fig. 3 shows the average values of µ ( X ( t ) ) with standard deviations over the 3000 samples with different masks in pretrained SANs. We observe that while a smaller temperature term d QK alleviates the (initial) rate of rank collapse in all cases, the effect is more significant with global attention than with sparser masked attention, and more in shallower layers than in deeper layers. The results in initialized models show similar trends and can be found in Appendix G.2.

Evolution of token geometry Finally, we study the evolution of token geometry as the number of layers increases in transformers. Specifically, to capture more fine-grained details than µ ( · ) , we measure how the rank, minimal singular value of token representations, and absolute values of pairwise cosine similarities among tokens change as tokens pass through the transformer layers. We select four representative pretrained transformer models: BERT [11], GPT2 [30], T5 [31] and ALBERT [24] and the average results over 3000 samples are shown in Fig. 4. We see that all models exhibit similar behaviors as tokens evolve along the layers: the models can effectively preserve the full rank of token representations, as also evident from the minimal singular values consistently remaining significantly above zero. Yet the absolute cosine similarities among tokens suggest that tokens gradually concentrate in a narrow region. This empirical observation aligns with the characteristics of the equilibrium in our counterexample that a stable long-term geometry for token representations in transformers can retain full rank while being close to a low dimensional geometry at the same time.

0

2

4

6

8

10

12

number of layers

0.5

0.6

0.7

0.8

0.9

1.0

% full rank

Transformers (pretrained)

BERT

GPT2

T5

ALBERT

10 0

10 1

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

min. singular value

Transformers (pretrained)

0

2

4

6

8

10

12

number of layers

0.2

0.4

0.6

0.8

abs. cosine similarity

Transformers (pretrained)

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 dynamics, when coupled with attention masks, LayerNorm and other components in transformers, enables expressive and complex behavior. In particular, we show that the choice of attention masks directly affects the token dynamics. It would hence be valuable for future research to investigate how to effectively design attention masks. Our result further suggests that robust token geometries under multi-layer self-attention can exhibit both full-rankness and anisotropic characteristics simultaneously, which is the case in real transformers as well. It would be interesting to study how such co-existence of these two characteristics would help with learning tasks - whether the full-rank yet close-to-low-dimensional geometry enables efficient learning and generalization while still letting tokens capture meaningful fine-grained details.

Acknowledgments

X.W., A.A., and A.J. would like to thank Bernard Chazelle for helpful discussions. X.W., A.A., and A.J. were supported by ONR Award N00014-23-1-2299 and a Vannevar Bush fellowship from the

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-2023 (MURI ML-SCOPE), NSF AI Institute TILOS (NSF CCF-2112665), NSF Award 2134108, and the Alexander von Humboldt Foundation.

References

Mira Ait-Saada and Mohamed Nadif. Is anisotropy truly harmful? a case study on text clustering. In ACL , 2023.

Josh Alman and Zhao Song. Fast attention requires bounded entries. In NeurIPS , 2023.

Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. ArXiv , abs/1607.06450, 2016.

Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. ArXiv , abs/2004.05150, 2020.

Shaked Brody, Uri Alon, and Eran Yahav. On the expressivity role of layernorm in transformers' attention. In ACL , 2023.

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 are few-shot learners. In NeurIPS , 2020.

Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Recurrent memory transformer. In NeurIPS , 2022.

Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Scaling transformer to 1m tokens and beyond with rmt. ArXiv , abs/2304.11062, 2023.

Xingyu Cai, Jiaji Huang, Yu-Lan Bian, and Kenneth Ward Church. Isotropy in the contextual embedding space: Clusters and manifolds. In ICLR , 2021.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv , abs/1904.10509, 2019.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In ACL , 2019.

Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In ICML , 2021.

Kawin Ethayarajh. How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings. In EMNLP , 2019.

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.

Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers. ArXiv , abs/2312.10794, 2023.

Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. In NeurIPS , 2023.

Nathan Godey, Eric Villemonte de la Clergerie, and Benoit Sagot. Anisotropy is inherent to self-attention in transformers. ArXiv , abs/2401.12143, 2024.

Darald J. Hartfiel. Nonhomogeneous Matrix Products . 2002.

Ali Hassani, Steven Walton, Jiacheng Li, Shengjia Li, and Humphrey Shi. Neighborhood attention transformer. In CVPR , 2023.

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.

Dongseong Hwang, Weiran Wang, Zhuoyuan Huo, Khe Chai Sim, and Pedro Moreno Mengibar. Transformerfam: Feedback attention is working memory. ArXiv , abs/2404.09173, 2024.

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.

Amir Joudaki, Hadi Daneshmand, and Francis R. Bach. On the impact of activation and normalization in obtaining isometric embeddings at initialization. In NeurIPS , 2023.

Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In ICLR , 2020.

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.

Timothee Mickus, Stig-Arne Grönroos, and Joseph Attieh. Isotropy, clusters, and classifiers. ArXiv , abs/2402.03191, 2024.

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 NeurIPS , 2022.

Xuran Pan, Tianzhu Ye, Zhuofan Xia, Shiji Song, and Gao Huang. Slide-transformer: Hierarchical vision transformer with local self-attention. In CVPR , 2023.

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 DeVito, 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 NeurIPS , 2019.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

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 a unified text-to-text transformer. Journal of Machine Learning Research , 2020.

Aurko Roy, Mohammad Taghi Saffar, Ashish Vaswani, and David Grangier. Efficient contentbased sparse attention with routing transformers. Transactions of the Association for Computational Linguistics , 2020.

Han Shi, Jiahui Gao, Hang Xu, Xiaodan Liang, Zhenguo Li, Lingpeng Kong, Stephen M. S. Lee, and James Tin-Yau Kwok. Revisiting over-smoothing in bert from the perspective of graph. In ICLR , 2022.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Shaolei Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In NeurIPS , 2023.

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 NeurIPS , 2017.

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-theart natural language processing. In EMNLP , 2020.

Xinyi Wu, Amir Ajorlou, Zihui Wu, and Ali Jadbabaie. Demystifying oversmoothing in attention-based graph neural networks. In NeurIPS , 2023.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In ICLR , 2024.

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 in the transformer architecture. In ICML , 2020.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In ICLR , 2020.

Chulhee Yun, Yin-Wen 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.

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. In NeurIPS , 2020.

Biao Zhang and Rico Sennrich. Root mean square layer normalization. In NeurIPS , 2019.

A Comment about Section 3.3 in Dong et al. [12]

From the mathematical form of LayerNorm in Eq. (3), since LayerNorm normalizes each row of ˜ X ( t ) to have 2-norm one, the operator that represents this scaling effect D ( t ) should be multiplied from the left hand side (scaling each row) rather than from the right hand side (scaling each column). In particular, since in Dong et al. [12], token representation X takes the same formulation that rows represent different tokens while columns represent different features, the same rule applies. Thus, the argument based on D ( t ) being multiplied from the right hand side and can be merged with W ( t ) V would not work.

B Proof of Theorem 1

B.1 Auxiliary results

Lemma 1. Under A1 -A3 , there exists ϵ > 0 such that A ( t ) i,j ≥ ϵ for all t ≥ 0 , ( j, i ) ∈ E .

Proof. Writing (2) recursively, we get that the token trajectories

stay uniformly bounded for all t ≥ 0 by A3 . Then it follows from A2 that there exists C ∈ R such that for all t ≥ 0 ,

This implies that there exists ϵ > 0 such that A ( t ) i,j ≥ ϵ for all ( j, i ) ∈ E .

Hence for all i, j ∈ [ N ] ,

Fix a vector x (0) ∈ R d and a { A ( n ) } ∞ n =0 in A G ,ϵ . Let x ( t ) := A ( t ) ...A (0) x (0) . We further denote that

Note that M ( t ) is monotone non-increasing in t , whereas m ( t ) is monotone non-decreasing in t .

Lemma 2. Let x ( t ) i = pm ( t ) +(1 -p ) M ( t ) for 0 ≤ p ≤ 1 . Then for any T ∈ N ,

and

Proof. Given x ( t ) i = pm ( t ) +(1 -p ) M ( t ) , we get that

Subsequently,

We obtain (11) by iterating the process. Similarly, (12) can be derived using a symmetric argument as for the upper bound.

Lemma 3. Let { A ( n ) } ∞ n =0 in A G ,ϵ and r be the radius of G . Then for t ≥ 0 and 0 ≤ p ≤ 1 ,

Proof. Let i 0 ∈ G be a center token and fix t ≥ 0 and 0 ≤ p ≤ 1 . Without loss of generality, suppose

From Lemma 2, we get that for all T ∈ N ,

Denote V k as the set of tokens that are exactly k -hop neighbors of i 0 :

For any i 1 ∈ V 1 , it follows that

Apply Lemma 2 to (15), we further get that for T ∈ N ,

Then let i 2 ∈ V 2 . For simplicity, we still use i 1 to denote the intermediate neighbor between i 0 and i 2 . Similarly to the case of i 1 , we get that

and applying Lemma 2 to (16), it follows that for T ∈ N ≥ 2

Iterating this process for tokens in V 3 , ..., V r , where r is the radius of G , we get that for any token i ∈ G ,

meaning that

We thus conclude that

For the other case of (14), i.e. x ( t ) i 0 > pm ( t ) + (1 -p ) M ( t ) , (17) is obtained using a symmetric argument by bounding m ( t + r ) from below through (12). This completes the proof of (13).

Corollary 2. There exists C ( x (0) ) > 0 such that

where r is the radius of the graph G .

B.2 Proof of Theorem 1

With Corollary 2, we are ready to prove the main theorem. Recall (10), we then get that

where ˜ (0) ( t )

where (19) follows from (2) and (20) follows from A3 . Then by (20),

where C := N 2 ∑ d j =1 C 2 j and C ′ := √ C here.

C Proof of Theorem 2

C.1 Auxiliary results

Lemma 4. Under A1 and A2 , there exists ϵ > 0 such that A ( t ) i,j ≥ ϵ for all t ≥ 0 , ( j, i ) ∈ E .

Proof. Note due to the way the dynamical system is defined in (3), for every t ≥ 0 , every token X ( t ) i, : lies within the unit sphere S d -1 . Then it follows from A2 that there exists C ∈ R such that for all t ≥ 0 ,

Hence for all i, j ∈ [ N ] ,

This implies that there exists ϵ > 0 such that A ( t ) i,j ≥ ϵ for all ( i, j ) ∈ E .

Lemma 5. Under the same assumptions as Theorem 2, there exists C 1 , C 2 > 0 such that C 1 ≤ D ( t ) i,i ≤ C 2 for all t ≥ 0 , i ∈ [ N ] .

Proof. We adopt the following notation:

Notation 1. For X ∈ R N × d , we use Conv( X ) to denote the convex hull of { X 1 , : , ..., X N, : } .

First, since for x ∈ Conv( X ( t ) ) , W ( t ) ⊤ x always lies in S d -1 , D ( t ) i,i ≥ 1 .

Then for all t ≥ 0 , each row vector of ( A ( t ) X ( t ) W ( t ) ) i, : lies in the convex cone generated by X (0) ˜ W ( t ) , where ˜ W ( t ) := Π t i =0 W ( i ) . To see this,

Hence all t ≥ 0 ,

Then

where a ( t ) ij ≥ 0 , ∀ j ∈ [ N ] and ∑ N j =1 a ( t ) ij ≥ 1 .

We then make the following observation:

̸

Lemma 6. Under assumption ( ∗ ), there exists v = 0 such that

meaning that all the points of X (0) lie in an open hemisphere H .

To see this, observe that since X (0) is full rank and N ≤ d , 0 / ∈ Conv( X (0) ) . Then the hyperplane separation theorem states that there exists a hyperplane strictly separating 0 and Conv( X (0) ) , letting us conclude Lemma 6.

Since all the points of X (0) lie in an open hemisphere H , there exists γ > 0 such that for all x ∈ Conv( X (0) ) and W orthogonal,

≥ γ .

We conclude that D ( t ) i,i ≤ 1 /γ , for all t ≥ 0 .

Writing the layer-wise update rule (3) recursively, we get the following formulation of X ( t ) :

We focus on the infinite product lim t →∞ D ( t ) A ( t ) ...D (0) A (0) . We first define a family of attention matrices as follows:

Definition 6. Let ϵ > 0 . We define A G ,ϵ to be the set of row-stochastic matrices satisfying the following conditions:

We also define

and subsequently,

To study the infinite product of matrices from M G ,ϵ,C 1 ,C 2 , we introduce some necessary concepts:

Definition 7. A non-negative matrix M is called row allowable if it has a positive entry in each of its rows.

Definition 8. Consider a sequence of row allowable matrices { M ( t ) } ∞ t =0 . Define the partial product

We say that the sequence { P ( t ) } ∞ t =0 is ergodic if there exists a sequence of positive rank one matrices { S ( t ) } ∞ t =0 such that

Our goal is to show the following key result:

Lemma 7. Let ϵ, C 1 , C 2 > 0 . Given { D ( t ) A ( t ) } ∞ t =0 in M G ,ϵ,C 1 ,C 2 . The sequence of partial products { P ( t ) } ∞ t =0 is ergodic.

For that, we will make use of the following result:

Lemma 8 (Corollary 5.1 [18]) . Consider a sequence of row allowable matrices { M ( t ) } ∞ t =0 . Let a t and b t be the smallest and largest entries in M ( t ) , respectively. If ∑ ∞ t =0 a t b t = ∞ , then { P ( t ) } ∞ t =0 is ergodic.

We cannot directly apply Lemma 8 to D ( t ) A ( t ) yet unless in the special case of full bidirectional attention such as the one used in BERT i.e. G is an undirected, complete graph, as a t could equal zero due to the sparsity pattern of attention G . In order to make use of the above result, we first need to show that long products of P ( t ) := D ( t ) A ( t ) will eventually become strictly positive for strongly connected G . For t 0 ≤ t 1 , we denote

Lemma 9. Under A1 , if G is strongly connected, there exist T ∈ N and C 3 , C 4 > 0 such that for all t ≥ 0 ,

Proof. Note that by A1 and the strong connectivity of G , there exists T ∈ N , e.g. the diameter of G , and ϵ, C 1 , C 2 > 0 such that for all t ≥ 0 ,

Then we apply Lemma 8 to { P ( k +1) T : T } ∞ k =0 and concludes Lemma 7.

As a result, from (23) we see that any P ( t ) can now be written in the following form:

where u ( t ) , v ( t ) are positive vectors such that u ( t ) ( v ( t ) ) ⊤ = S ( t ) . Without loss of generality, assume ∥ v ( t ) ∥ 2 = 1 .

Then it follow from (23) that

Lemma 10.

Proof. Given where by (23), since also

We get that and hence

and

Thus the 2 -norm of the i th token of X ( t +1) , ∥ X ( t +1) i, : ∥ 2 is

Note that

where σ min ( · ) denotes the smallest singular value. Since we assume X (0) has full rank, and ˜ W ( t ) is orthogonal

Then given ∀ i ∈ [ N ] , as a result,

This together with Lemma 10 implies that

C.1.1 Exponential convergence rate

Having established the convergence in (26), we then establish the exponential convergence rate. We define that

Similar to the derivation of (25), note that

meaning that lim t →∞ X ( t ) i, : ∈ S d -1 for all i ∈ [ N ] .

Then given the convergence established in (26), we get that lim t →∞ α ( t ) = 1 , from which there exists t 0 ∈ N such that α ( t ) > 0 , ∀ t ≥ t 0 . It follows that for t ≥ t 0 ,

Thus

Writing recursively, we get that for k ≥ 0 ,

Let C := 2 / (1 -Nϵ 2 r ) t 0 /r . Hence for all t ≥ 0 ,

Since by definition, 1 -ϕ ( t ) ≥ 1 -⟨ X ( t ) i, : , X ( t ) j, : ⟩ ≥ ∥ X ( t ) i, : -X ( t ) j, : ∥ 2 2 / 2 , it follows that

where C ′ = √ CN , completing the proof.

it follows that

Since ∀ i ∈ [ N ] , as a result,

D Proof of Corollary 1

First, since for x ∈ Conv( X ( t ) ) , W ( t ) ⊤ x always lies in S d -1 , D ( t ) i,i ≥ 1 , for all i ∈ [ N ] , t ≥ 0 . For all t ≥ 0 , since ϕ (0) ≥ 0 and G is quasi-strongly connected, we have that

Thus

Writing recursively, we get that for k ≥ 0 ,

Let C := 1 / (1 -ϵ 2 r ) . Hence for all t ≥ 0 ,

Since by definition, 1 -ϕ ( t ) ≥ 1 -⟨ X ( t ) i, : , X ( t ) j, : ⟩ ≥ ∥ X ( t ) i, : -X ( t ) j, : ∥ 2 2 / 2 , it follows that

where C ′ = √ CN , completing the proof.

E Detailed analysis of the illustrative example in Section 4.3.1

First token Given that we get that

and thus

Since the sign of X ( t ) 1 , 1 is invariant, depending on its initial sign, we get that

.

Second token Assume that the first token converges to [0 , 1] . Then

Solving X ( t +1) 2 , 2 X ( t +1) = X ( t ) 2 , 2 X ( t ) gives X 2 , 1 = -1 /w and

To show the convergence, we utilize the following useful property of g ( · ) :

Lemma 11. g is non-decreasing in X 2 , 2 /X 2 , 1 .

Proof. Let y = X 2 , 2 /X 2 , 1 . Given the identity that we get that

Hence which means that

Given X 2 , 2 /X 2 , 1 is invariant at X 2 , 1 = -1 /w (symmetric points A and B in Fig. 5), together with Lemma 11, they imply the invariance of the 4 segments on S 1 , as shown in Fig. 5. Furthermore, (29) shows which point each region would converge to. Thus it suffices to choose X (0) 2 , : to be on the dark blue or orange segment.

F Proof of Theorem 3

F.1 Finding the equilibria

Step 1 Without loss of generality, we set that W ( t ) K = W ( t ) Q = 0 , for all t ≥ 0 . The system of equations corresponds to the equilibrium condition for the attention dynamics with LayerNorm defined in (3) when the attention mask G is causal is as follows:

(0

,

-

1)

(0

,

1)

A

B

where d i = i/ ∥ ∑ i j =1 X j, : W V ∥ 2 . We define that β i = d i /i for i ∈ [ N ] . Then the above system equation becomes

after substitution. Then for i ∈ [ N -1] ,

Let k ∈ [ N ] . If β 1 = ... = β k = 1 , then for i ∈ [ k -1] ,

Then if X K, : is a generalized eigenvector of rank k corresponding to W V with eigenvalue 1 , then from (31) that { X K, : , ..., X 1 , : } is the Jordan chain generated by X K, : , which is known to be linearly independent.

Hence we set W V to be a matrix with a generalized eigenvector of rank k corresponding to eigenvalue 1 . For example,

For k = 1 , W V can be set as I d ;

For k = 2 , W V can be set as

For k = 3 , W V can be set as

1

For k = N , W V can be set as

For our own purpose, we consider k = N = d . Then to solve for equilibrium defined by the condition in (31), for i ∈ [ N -1] ,

which means that

Then setting X 1 , : = [0 , 0 , ..., 0 , 1] , iteratively solving (31), we get that an equilibrium satisfying (30) is as follows:

which is valid if w > 1 . Since from (31), there is a free coordinate, in fact there a 2 N many such equilibria. Hence there exists { W ( t ) K , W ( t ) Q , W ( t ) V } ∞ t =0 such that the corresponding attention dynamics with LayerNorm has an equilibrium with rank k = N .

Step 2 To show that there are equilibria of any rank between 1 and full co-existing for a corresponding dynamics given the same set of { W ( t ) K , W ( t ) Q , W ( t ) V } ∞ t =0 , consider W ( t ) V to be the case of (32). Then notice that to show there exists equilibrium of k beyond the full rank one as derived in Step 1, notice the following:

We can set the first N -k +1 tokens to be [0 , ..., 0 , 1] , which still satisfies the equilibrium condition. This implies that β N -k +1 = 1 / ( N -K +1) . Then recall (30):

Let β N -k +2 = 1 , we get that

which yields

Continue the process by setting β i = 1 for i ≥ N -k +2 , we get a rank k equilibrium as desired (there are 2 k such equilibria), as long as w ≥ N .

Step 3 Finally to ensure { W ( t ) K , W ( t ) Q , W ( t ) V } ∞ t =0 satisfies A2 -A3 :

We see that { W ( t ) K , W ( t ) Q } ∞ t =0 satisfies A2 by construction;

Since this is a dynamics with LayerNorm, applying W ( t ) V is in fact equivalent to applying ˆ W V = W V / ∥ W V ∥ 2 under the dynamics, where the latter has 2 -norm at most 1 and thus using ˆ W V satisfies A3 .

F.2 Anisotropy of the full rank equilibria

To show that the full rank equilibria X ∗ lie in a narrow region, we make use of the notion stable rank:

We would like to show that for any δ > 0 , weights can be chosen such that

Since for any matrix M , ∥ M ∥ 2 F ≥ ∥ M ∥ 2 2 , the first inequality trivially holds. To show the second inequality, consider the equilibrium in (34). Then direct calculation yields that

and thus any ∥ X ∗ v ∥ 2 such that v ∈ R N , ∥ v ∥ 2 = 1 serves as a lower bound for ∥ X ∥ 2 .

We consider v = [0 , 0 , ..., 0 , 1] ⊤ ∈ R d . Then it follows that

Hence for any δ > 0 , it suffices to choose w ≥ 1 such that

Since by definition, which implies that

F.3 Convergence analysis

The convergence analysis for the case d = 2 can be found in Appendix E. Here, we deal with the general case where d > 2 .

To show that tokens do not converge to an rank one subspace, it suffices to show that given certain conditions on the initial input X (0) , there are invariant space of the tokens that is at least rank 2 . To that end, we analyze the convergence of the first two tokens.

First token Given that we get that

Then for all 3 ≤ i ≤ d , setting

Let X (0) 1 , 2 X (0) 1 , 1 = r 2 w , for r 2 ∈ R . It follows from above that

where r i ∈ R , and substituting from the base case (35) to (36) recursively, we get that

which implies that

Further, for i ≥ 3

Then (38) and (39) yields that

which let us conclude that if X (0) 1 , 1 > 0 , the first token would converge to [0 , 0 , ..., 0 , 1] .

Second token Assume that the first token has converged to [0 , 0 , ..., 0 , 1] . Similar to the first token, we can show that for 2 ≤ i ≤ d -1 , it holds that

and

by setting X (0) 2 ,i /X (0) 2 , 1 = r i w i -1 , where r i > 0 . Then note that

We will show a useful monotone property of h ( · ) .

Lemma 12. Fix t, r 2 , ..., r d -1 > 0 , h is non-decreasing in X 2 ,d /X 2 ,d -1 .

Proof. Let y = X 2 ,d /X 2 ,d -1 . Given the identity that

we get that

Hence which means that

With this, we will show the following claim proving the invariance of the subspace of X 2 ,d -1 ≤ -1 /w given some initial conditions:

Proof. When X ( t ) 2 ,d -1 = -1 /w , and given the conditions on r 2 , ..., r d -1 in (41),

is well defined, where f ( · ) is monotone non-increasing in t for fixed r 1 , ..., r d -2 > 0 . We denote

Note that if X ( t ) 2 ,d -1 ≤ -1 /w , we have

then by Lemma 12 we get that

and since by (40), it implies that

which means that

and thus X ( t ) 2 ,d -1 ≤ -1 /w for t ≥ 0 .

This let us conclude that for a general class of sequences X (0) , X ( t ) does not converge to a rank one subspace as t →∞ .

F.4 Further discussion

Directionality of attention mask G In the proof of Theorem 3, we would like to note that G being the causal graph gives a key structure to construct the counterexamples. A causal graph essentially introduces asymmetries among token interactions where some tokens can attend to some other tokens but not verse versa and we exploit such asymmetries in constructing the counterexamples. A similar construction thus does not naturally apply to the case where G is undirected, e.g. a complete graph.

As a result, we leave the following interesting question for future research about whether the directionality of graphs empowers the attention mechanism differently. We conjecture a negative answer based on preliminary simulations.

Does a similar equilibrium exist if one uses an undirected graphs G as the attention mask? If not, what differentiates the case between an undirected graph and a directed graph?

G Experiments

Here we provide more details about numerical experiments presented in Section 5 and some additional experimental results. All models were implemented with PyTorch [29] and Transformers library [36].

Compute We ran all of our experiments on CPUs.

Licenses

Libraries

-Wikipedia: MIT license

-tranformers [36]: Apache License 2.0

Pretrained models

-BERT (bert-base-uncased) [11]: Apache License 2.0

-GPT2 (openai-community/gpt2) [30]: MIT license

-T5 (google-t5/t5-base) [31]: Apache License 2.0

-ALBERT (albert/albert-base-v2) [24]: Apache License 2.0

G.1 Deeper models

We provide the results for deeper models for the experiment presented in Fig. 2. Here, we conduct the same experiment in a 128 layer randomly initialized BERT with 12 heads and 768 hidden dimension using different attention masks. Similarly, in SANs, while different masks have different convergence rates, all of them suffer from exponential rank collapse. But as soon as we add in LayerNorm, µ ( X ( t ) ) no longer converges to 0 , as µ ( · ) does not show any decreasing trend even after 128 layers.

10 0

10 1

10 2

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

(X)

SANs (initialized), #layer=128

10 0

10 1

10 2

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs + LN (initialized), #layer=128

10 0

10 1

10 2

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

Transformers (initialized), #layer=128

complete

causal

slide-window

slide-window-uni

G.2 The effect of the temperature term d QK in initialized transformers

The effect of the temperature term d QK in initialized transformers is shown in Fig. 7. Similar to the pretrained models, we see that smaller temperature terms alleviate the rate of rank collapse, and effect is more significant under global attention than under sparser masked attention, and in shallower layers than deeper layers.

10 0

10 1

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

(X)

SANs (initialized): complete

10 0

10 1

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs (initialized): causal

10 0

10 1

number of layers

10

3

10

2

10

1

10 0

10 1

10 2

10 3

SANs (initialized): sliding window

dQK

= 1/64

dQK

= 1

dQK

= 64

G.3 The evolution of singular values of token representations in transformers

In this section, we investigate the full set of singular values of token representations X ( t ) in transformers. For a clear presentation, we randomly select 3 out of the 3000 sequence samples described in Section 5 and calculate all the singular values of X ( t ) . The results in randomly initialized transformers are presented in Fig. 8.

We see that while the rank is full for X ( t ) , there is usually one principal component that is much more dominant than the rest. As we have discussed in the main text, this is the anisotropic characteristic of token representations- there is a strong non-uniformity in different directions. Nonetheless, it is important to note that the 'small" singular values are in fact not negligible here in the absolute sense-the minimal singular values are never close to zero, as Fig. 4 (middle) shows.

G.4 Initial token geometry

We verify the initial token geometry in transformers. For a clear presentation, we randomly select 3 out of the 3000 sequence samples described in Section 5 and calculate all the pairwise cosine similarities between tokens. The results in randomly initialized transformers are presented in Fig. 9. We see that in particular in BERT and ALBERT, the initial pairwise cosine similarities are all non-negative.

0

1

2

3

4

5

6

7

8

9

10

11

12

0

50

100

150

200

singular value

BERT (initialized), seq 1

0

1

2

3

4

5

6

7

8

9

10

11

12

0

50

100

150

200

BERT (initialized), seq 2

0

1

2

3

4

5

6

7

8

9

10

11

12

0

50

100

150

200

BERT (initialized), seq 3

0

1

2

3

4

5

6

7

8

9

10

11

12

0

10

20

30

40

50

60

70

singular value

GPT2 (initialized), seq 1

0

1

2

3

4

5

6

7

8

9

10

11

12

0

10

20

30

40

50

60

70

GPT2 (initialized), seq 2

0

1

2

3

4

5

6

7

8

9

10

11

12

0

10

20

30

40

50

60

GPT2 (initialized), seq 3

0

1

2

3

4

5

6

7

8

9

10

11

12

0

200

400

600

800

singular value

T5 (initialized), seq 1

0

1

2

3

4

5

6

7

8

9

10

11

12

0

200

400

600

800

T5 (initialized), seq 2

0

1

2

3

4

5

6

7

8

9

10

11

12

0

200

400

600

800

T5 (initialized), seq 3

0

1

2

3

4

5

6

7

8

9

10

11

12

number of layers

0

100

200

300

400

500

600

700

singular value

ALBERT (initialized), seq 1

0

1

2

3

4

5

6

7

8

9

10

11

12

number of layers

0

200

400

600

ALBERT (initialized), seq 2

0

1

2

3

4

5

6

7

8

9

10

11

12

number of layers

0

200

400

600

800

ALBERT (initialized), seq 3

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0

1000

2000

3000

4000

5000

6000

7000

frequency

BERT (initialized), seq 1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0

1000

2000

3000

4000

5000

6000

7000

BERT (initialized), seq 2

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0

1000

2000

3000

4000

5000

6000

7000

BERT (initialized), seq 3

0.0

0.2

0.4

0.6

0.8

1.0

0

2000

4000

6000

8000

10000

frequency

GPT2 (initialized), seq 1

0.2

0.0

0.2

0.4

0.6

0.8

1.0

0

2000

4000

6000

8000

10000

GPT2 (initialized), seq 2

0.0

0.2

0.4

0.6

0.8

1.0

0

2000

4000

6000

8000

10000

GPT2 (initialized), seq 3

0.2

0.0

0.2

0.4

0.6

0.8

1.0

0

2000

4000

6000

8000

10000

12000

14000

frequency

T5 (initialized), seq 1

0.2

0.0

0.2

0.4

0.6

0.8

1.0

0

2000

4000

6000

8000

10000

12000

14000

T5 (initialized), seq 2

0.2

0.0

0.2

0.4

0.6

0.8

1.0

0

2000

4000

6000

8000

10000

12000

14000

T5 (initialized), seq 3

0

1000

2000

3000

4000

5000

6000

frequency

ALBERT (initialized), seq 1

0

1000

2000

3000

4000

5000

6000

ALBERT (initialized), seq 2

0

1000

2000

3000

4000

5000

6000

ALBERT (initialized), seq 3

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

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 contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

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 no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We confirm that the paper provides the full set of 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 proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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) If the contribution is primarily a 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 be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Since all the experiments are for the verification purpose of our theory, they are straightforward and can be easily replicated given the instructions detailed in the paper.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We confirm that we specify all the details for the experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We confirm that we report standard deviations for all the experimental results.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to 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.

It is 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 a 96% CI, if the hypothesis of Normality of 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 are 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 or 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 or 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 https://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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]

Justification: Our work is about machine learning theory with the goal to gain deeper insights into the attention mechanism. We believe there are no direct societal impacts of the work that are worth discussing at this time.

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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work is about machine learning theory with no release of data or models. Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have properly credited all the data and models used in the work.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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 new assets in this work.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of 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. Crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to 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 an equivalent approval/review based on 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 on 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nMy 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThere is not a clear limitations section but the authors do mention a limitation on line 283.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

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

STRENGTHS:
1.I like this paper for its strong motivation and very interesting insights. I appreciate the authors for addressing the rank collapse issue in Transformers, a topic often overlooked in the community.

2. The authors provide detailed analysis to demonstrate that attention masks and layer normalization can help address this issue.

3. The experimental results support the authors' analysis effectively.

4. Figure 1 is particularly helpful in understanding the illustration of the effectiveness of layer normalization.

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

QUESTIONS:
I like this work; however, I have two additional questions:

If local attention alleviates rank collapse, how do the experimental results compare to those of full attention or causal masked attention?

Which approach is more effective in addressing the rank collapse issue: local attention or causal attention? Are there any strong insights into this? I assume local attention may not be generally applicable in language tasks but is more suitable for vision tasks.

LIMITATIONS:
Please see above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-15