#  On the Role of Attention Masks
and LayerNorm in Transformers

Xinyi Wu1 Amir Ajorlou1 Yifei Wang2 Stefanie Jegelka2,3 Ali Jadbabaie1
1MIT LIDS 2MIT CSAIL 3TU Munich
<EMAIL_ADDRESS>
###### 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, 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 [29], 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 [10, 13].

Transformers are known to exhibit the _rank collapse phenomenon_ 111This is
sometimes also referred to as the _oversmoothing_ [27, 13, 31] or the _token
uniformity_ [10, 22] problem. [10, 22, 13, 27, 31, 17], 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 [10, 13, 14].

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 [10, 22, 14, 13], which only applies to
certain attention mechanisms such as the one deployed in the BERT family [9,
20]. The vast majority of popular transformer architectures used nowadays in
practice, including the GPT models [24, 6], use the causal attention masks
where tokens are only allowed to attend to preceding tokens, or have sparse
attention structure (sparse transformers) [34, 35, 4, 8, 26, 18], where
attention is restricted to local interactions between tokens and their chosen
neighbors. This limits the practical applicability of the theoretical results
developed in [10, 22, 14, 13], 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. [10] 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. [13]. 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. [19] 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 [14, 13], 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, 35]. 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. [10]
first show that as the number of self-attention layers increases, tokens
inevitably suffer from exponential rank collapse, while Wu et al. [31]
establish a similar result for the attention mechanism in graph neural
networks, taking additionally the layer-wise nonlinear activation functions
into account. Other works [13, 14] 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 [29].
To better understand its role in transformers, Xiong et al. [32] 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. [10] 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 [19]. It thus remains an open
question what effect LayerNorm would have on rank collapse in transformers.

##  3 Problem Setup

##### Notation

Let $\|\cdot\|_{2}$, $\|\cdot\|_{F}$ be the $2$-norm and Frobenius norm,
respectively. We use the shorthand $[n]:=\\{1,\ldots,n\\}$. We denote the all-
one vector of length $N$ by $\mathbf{1}\in\mathbb{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 $\mathcal{G}$. Formally, we represent a directed graph with $N$
nodes by $\mathcal{G}$ and let $E(\mathcal{G})$ be the set of directed edges
of $\mathcal{G}$. If there is a directed edge going from $j$ to $i$ in
$\mathcal{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$, then $(j,i)\in
E(\mathcal{G})$. The set $\mathcal{N}_{i}$ of all neighbors of node $i$ is
then $\\{k|(k,i)\in E(\mathcal{G})\\}$.

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

######  Definition 1 (Reachability).

We say a node $v$ is _reachable_ from $u$ in a directed graph $\mathcal{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 $\mathcal{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 $\mathcal{G}$ is
reachable is called a _center node_.

######  Definition 4 (Quasi-Strongly Connected).

A directed graph $\mathcal{G}$ is said to be _quasi-strongly connected_ if
$\mathcal{G}$ has at least one center node.

######  Definition 5 (Diameter).

The diameter 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 longest distance among them.

###  3.1 (Masked) Attention Mechanism

Given token representations $X\in\mathbb{R}^{N\times d}$, the raw attention
score matrix is computed as

$R=XW_{Q}(XW_{K})^{\top}/\sqrt{d_{QK}}\,,$

where $W_{Q},W_{K}\in\mathbb{R}^{d\times d^{\prime}}$ are the query and the
key matrix, respectively, and $\sqrt{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\in\mathbb{R}^{N\times N}$ based on $R$ whose
sparsity pattern is specified by a directed graph $\mathcal{G}$: we normalize
$R_{ij}$ among all allowed token attention interactions $(k,i)\in
E(\mathcal{G})$ such that

$A_{ij}={\mathrm{softmax}}_{\mathcal{G}}(R_{ij})=\frac{\exp(R_{ij})}{\sum_{k\in\mathcal{N}_{i}}\exp(R_{ik})}\,\;\,\text{if
$(j,i)\in E(\mathcal{G})$},\;\;\;A_{ij}=0\,\text{ otherwise}.$

###  3.2 LayerNorm

Given token representations $X\in\mathbb{R}^{N\times 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 [13, 28]. Mathematically, let
$D=\operatorname{diag}(d_{1},d_{2},...,d_{d})$ where $d_{i}=1/\|X_{i,:}\|_{2}$
for all $i\in[N]$, then $\operatorname{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

$\displaystyle
A^{(t)}=\mathrm{softmax}_{\mathcal{G}^{(t)}}\left(X^{(t)}W^{(t)}_{Q}(X^{(t)}W^{(t)}_{K})^{\top}/\sqrt{d_{QK}}\right)$
(1) $\displaystyle\tilde{X}^{(t+1)}=A^{(t)}X^{(t)}W_{V}^{(t)}$ (2)
$\displaystyle
X^{(t+1)}=\operatorname{LN}(\tilde{X}^{(t)})\vcentcolon=D^{(t)}A^{(t)}X^{(t)}W_{V}^{(t)}\,.$
(3)

where $W_{V}^{(t)}\in\mathbb{R}^{d\times d^{\prime}}$ is the value matrix. For
simplicity, throughout the paper, we assume that $d=d^{\prime}$ and
$\mathcal{G}^{(t)}=\mathcal{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
$\mu(\cdot):\mathbb{R}^{N\times d}\to\mathbb{R}_{\geq 0}$:

$\mu(X)\vcentcolon=\|X-\mathbf{1}\gamma_{X}\|_{F},\text{ where
}\gamma_{X}=\mathbf{1}^{\top}X/N\,.$ (4)

This measure is mathematically equivalent to the measure
$\operatorname{res}(X)=\operatorname*{arg\,min}_{x\in\mathbb{R}^{d}}\|X-\mathbf{1}x^{\top}\|_{F}$
used in [10], 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 $\mu(\cdot)$: these results apply to any Lipschitz
$\mu^{\prime}(\cdot)$ with a Lipschitz constant $L$ such that $\mu(X)=0$ if
and only $X_{i,:}=X_{j,:},\forall i,j\in[N]$, as we can use the formulation to
directly derive that

$\mu^{\prime}(X)=|\mu^{\prime}(X)-\mu^{\prime}(\mathbf{1}_{\gamma_{X}})|\leq
L\|X-\mathbf{1}_{\gamma_{X}}\|_{F}=L\mu(X)\,.$

Finally, we adopt the following assumptions in our analysis:

  1. A1

$\mathcal{G}$ contains self-loops. i.e. $(i,i)\in E$ for every token $i\in[N]$
.

  2. A2

There exist constants $C\in\mathbb{R}$ such that
$\underset{t\in\mathbb{N}}{\max}\left\\{\left\|W_{Q}^{(t)}\right\|_{2},\left\|W_{K}^{(t)}\right\|_{2}\right\\}\leq
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\geq 0$ even without LayerNorm, we further assume that

  1. A3

The sequence
$\left\\{\left\|\prod_{t=0}^{k}W_{V}^{(t)}\right\|_{2}\right\\}_{k=0}^{\infty}$
is bounded.

Then with general attention masks $\mathcal{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 $\mathcal{G}$ is a quasi-strongly connected graph, then there exists
$\epsilon>0$ where for all $t\geq 0$,

$A^{(t)}_{i,j}\geq\epsilon,\quad\text{ for all }(j,i)\in E.$ (5)

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

$\mu(X^{(t)})\leq C\left(1-\epsilon^{r}\right)^{t/r}\,,$ (6)

where $r$ is the diameter of $\mathcal{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 [10] from $\mathcal{G}$ being
a complete graph to a much more general class of attention patterns: the
attention pattern $\mathcal{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
[24, 6], or sparse attention patterns deployed in many efficient transformer
models [4, 35, 26, 8, 18]. We discuss a few interesting implications below.

##### Local vs. global attention

The exponential rate $(1-\epsilon^{r})^{1/r}$ is monotone in the graph
diameter $r$. This means that rank collapse should be slower for graphs with
larger diameter $r$. Our result thus indirectly supports the use of local
attention patterns [35, 4], 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 $\epsilon$, which
means that rank collapse is slower with smaller $\epsilon$. One can interpret
$\epsilon$ as how "focused" attention is distributed among reachable tokens,
as $\epsilon$ 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. [34] show that transformers with strongly connected graph
masks are sequence-to-sequence function universal, yet with a mask
$\mathcal{G}$ they need at least the diameter of $\mathcal{G}$ layers to
achieve the full sequence-to-sequence function approximation property. This
implies that masks with smaller diameters are more efficient in terms of
function approximation power, yet they are more prone to rank collapse.

######  Remark 1.

Since the analysis in [10] fundamentally relies on the shift-invariant
property of $\mathrm{softmax}(\cdot)$, 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
$\mathcal{G}$ be a strongly connected graph. Assume A1-A2, and that
$W_{V}^{(t)}$ is orthogonal for all $t\geq 0$, and in addition, the initial
input $X^{(0)}$ satisfies the following conditions:

  * (i)

There exists $v\in\mathbb{S}^{d-1}$ such that $\langle
X^{(0)}_{i,:},v\rangle>0$ for all $i\in[N]$;

  * (ii)

$N\leq d$, $X^{(0)}$ has full rank.

Then there exist $C>0$, $\epsilon>0$ such that $N\epsilon<1$ and

$\mu(X^{(t)})\leq C(1-N\epsilon^{2r})^{\frac{t}{2r}}\qquad\forall t\geq 0\,,$
(7)

where $r$ is the diameter of $\mathcal{G}$, meaning that tokens converge to a
common point on $\mathbb{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 [13]. 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 $\mathcal{G}$ being complete (which is the case in [13]). One
way to satisfy the two conditions on the initial input $X^{(0)}$ is to require
$N\leq d$ and to initialize tokens uniformly randomly on $\mathbb{S}^{d-1}$,
then the two conditions hold almost surely. This is how the conditions are
dealt with in [13].

If the initial token geometry satisfies a stronger condition than (i), then
(ii) is no longer necessary and Theorem 2 even directly generalizes to quasi-
strongly connected graphs $\mathcal{G}$. We define
$\phi^{(t)}\vcentcolon=\underset{i,j\in[N]}{\min}\langle
X^{(t)}_{i,:},X^{(t)}_{j,:}\rangle$, 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
$\mathcal{G}$ is quasi-strongly connected.

######  Corollary 1.

Consider the self-attention dynamics with LayerNorm defined in (3). Let
$\mathcal{G}$ be a quasi-strongly connected graph. Under A1-A2, if
$W_{V}^{(t)}$ is orthogonal for all $t\geq 0$ and $\phi^{(0)}\geq 0$, then
there exist $C>0$, $\epsilon>0$ such that $N\epsilon<1$ and

$\mu(X^{(t)})\leq C(1-N\epsilon^{2r})^{\frac{t}{2r}},\qquad\forall t\geq 0\,.$
(8)

where $r$ is the diameter of $\mathcal{G}$, meaning that tokens converge to a
common point on $\mathbb{S}^{d-1}$ exponentially.

While the requirement in the above result on the initial token geometry might
seem restrictive, in practice, such an initial token geometry is indeed
satisfied in some commonly used transformer models such as BERT [9] or ALBERT
[20] after the initial embedding layer. See the empirical measurements in
Section G.4.

###  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.

####  4.3.1 Illustrative Counterexample

For simplicity of illustration, we consider $N=2,d=2$, and $\mathcal{G}$ to be
the causal mask. Then let $W_{K}^{(t)}=W_{Q}^{(t)}=\mathbf{0}$, which leads to
the attention matrices

$A^{(t)}=\begin{bmatrix}1&0\\\ 1/2&1/2\end{bmatrix}\,\quad\forall t\geq 0\,.$

We further let

$W_{V}^{(t)}=\begin{bmatrix}1&w\\\ 0&1\end{bmatrix}\,\quad\forall t\geq 0\,,$

for $w\in\mathbb{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\to\infty$ is
illustrated in Fig. 1, where
$A=\left(-\frac{1}{w},\sqrt{1-\frac{1}{w^{2}}}\right)$,
$B=\left(-\frac{1}{w},-\sqrt{1-\frac{1}{w^{2}}}\right)$. A rigorous proof can
be found in Appendix E. Note that due to the scaling effect of LayerNorm, any
scaled version $c^{(t)}W_{V}^{(t)},c^{(t)}>0$ of $W_{V}^{(t)}$ works
equivalently here.

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

Figure 1: Long-term behavior of tokens in the case of $N=2,d=2$. Without
LayerNorm (left), both tokens collapse to the same point in $\mathbb{R}^{2}$;
whereas with LayerNorm (right), such a collapse would not necessarily happen
and token representations can maintain full rank in the long term (first token
converges either to $(0,1)$ or $(0,-1)$. Assuming convergence to $(0,1)$ for
the first token, the second token converges to $B$, if it is initially located
within the red segment).

####  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 $\mathcal{G}$ be the causal graph and consider the attention dynamics with
LayerNorm defined in (3). Then there exists
$\left\\{W_{Q}^{(t)},W_{K}^{(t)},W_{V}^{(t)}\right\\}_{t=0}^{\infty}$
satisfying A2-A3 such that the corresponding dynamic has the following
properties:

  1. 1.

For any $1\leq k\leq\min\\{N,d\\}$, there exist at least $2^{k}$ equilibria of
rank $k$;

  2. 2.

For a general class of input sequences $X^{(0)}$ in $(\mathbb{S}^{d-1})^{N}$
with measure greater than $0$, $X^{(t)}$ does not converge to a rank one
subspace as $t\to\infty$.

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 [10], 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_{V}^{(t)}$ 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 two further intriguing findings from our analysis of the
self-attention dynamics with LayerNorm presented in the previous section.

##### 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 [7, 15,
11, 12, 1]. Interestingly, as we show in Section F.2, the full rank equilibria
described in Theorem 3 align with this observation, as tokens lie in a narrow
region on $\mathbb{S}^{d-1}$. While it still remains unclear how exactly such
a representation geometry is useful for downstream tasks [7, 12, 15, 1],
empirical studies have found that anisotropy is not necessarily harmful for
semantic representations and can assist with tasks like clustering [1, 21].

##### 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 [33, 34]. 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 [10], we randomly select 32 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 $\mu(\cdot)$. We use BERT [9] as the backbone
transformer model and consider three different model variants for a controlled
experiment: self-attention network (SAN) in which the model has self-attention
layers, SAN with LayerNorm where in each layer there is self-attention
followed by LayerNorm, 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 $\mu(X^{(t)})$ with standard
deviations over the $32$ samples in different architectures. We see that as
our theory predicts, in SANs, $\mu(X^{(t)})$ converges to zero exponentially
for all attention masks, yet more local attention masks slow down the
convergence rate. Furthermore, right after we add in LayerNorm, $\mu(X^{(t)})$
no longer converges to zero and can be maintained significantly above zero,
which holds for both initialized and pretrained models. Results in $128$-layer
initialized models also exhibit the same trend, and can be found in Section
G.1.

Figure 2: Evolution of $\mu(X^{(t)})$ (in log-log scale) as the number of
layers increases. Rank collapse happens exponentially for pure attention,
despite different attention masks having different convergence rates, as our
theory predicts. However, as soon as we solely add in LayerNorm,
$\mu(X^{(t)})$ no longer converge to zero in both randomly initialized and
pretrained models. Figure 3: Evolution of $\mu(X^{(t)})$ (in log-log scale) as
the number of layers increases. Smaller temperature terms alleviate the rate
of rank collapse, and effect is more significant with global attention than
with sparser masked attention, and more in shallower layers than deeper
layers.

##### The 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
$\mu(X^{(t)})$ with standard deviations over the $32$ 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 Section 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 $\mu(\cdot)$, 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 [9], GPT2 [24], T5 [25] and
ALBERT [20] and the average results over $32$ samples are shown in Fig. 4. We
see that all models exhibit similar behaviors as tokens evolve along the
layers: the models can effectively preserve the full rank of token
representations, as also evident from the minimal singular values consistently
remaining significantly above zero. Yet the absolute cosine similarities among
tokens suggest that tokens gradually concentrate in a narrow region. This
empirical observation aligns with the characteristics of the equilibrium in
our counterexample that a stable long-term geometry for token representations
in transformers can retain full rank while being close to a low dimensional
geometry at the same time.

Figure 4: Evolution of token geometry as the number of layers increases. We
see that tokens are indeed able to maintain full rank, while at the same time
the representations are anisotropic, meaning that they concentrate in a narrow
region, as indicated by the average pairwise absolute cosine similarities.
This is in agreement with our theoretical results about LayerNorm.

##  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
and LayerNorm, 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. This research has been supported in part 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)).

## References

  * Ait-Saada and Nadif [2023] Mira Ait-Saada and Mohamed Nadif.  Is anisotropy truly harmful? a case study on text clustering.  In _ACL_ , 2023.
  * Alman and Song [2023] Josh Alman and Zhao Song.  Fast attention requires bounded entries.  In _NeurIPS_ , 2023.
  * Ba et al. [2016] Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton.  Layer normalization.  _ArXiv_ , abs/1607.06450, 2016.
  * Beltagy et al. [2020] Iz Beltagy, Matthew E. Peters, and Arman Cohan.  Longformer: The long-document transformer.  _ArXiv_ , abs/2004.05150, 2020.
  * Brody et al. [2023] Shaked Brody, Uri Alon, and Eran Yahav.  On the expressivity role of layernorm in transformers’ attention.  In _ACL_ , 2023.
  * Brown et al. [2020] 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.
  * Cai et al. [2021] Xingyu Cai, Jiaji Huang, Yu-Lan Bian, and Kenneth Ward Church.  Isotropy in the contextual embedding space: Clusters and manifolds.  In _ICLR_ , 2021.
  * Child et al. [2019] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers.  _ArXiv_ , abs/1904.10509, 2019.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In _ACL_ , 2019.
  * Dong et al. [2021] 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.
  * Ethayarajh [2019] Kawin Ethayarajh.  How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings.  In _EMNLP_ , 2019.
  * Gao et al. [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.
  * Geshkovski et al. [2023a] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet.  A mathematical perspective on transformers.  _ArXiv_ , abs/2312.10794, 2023a.
  * Geshkovski et al. [2023b] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet.  The emergence of clusters in self-attention dynamics.  In _NeurIPS_ , 2023b.
  * Godey et al. [2024] Nathan Godey, Eric Villemonte de la Clergerie, and Benoit Sagot.  Anisotropy is inherent to self-attention in transformers.  _ArXiv_ , abs/2401.12143, 2024.
  * Hartfiel [2002] Darald J. Hartfiel.  _Nonhomogeneous Matrix Products_.  2002\.
  * He et al. [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.
  * Hwang et al. [2024] Dongseong Hwang, Weiran Wang, Zhuoyuan Huo, Khe Chai Sim, and Pedro Moreno Mengibar.  Transformerfam: Feedback attention is working memory.  _ArXiv_ , abs/2404.09173, 2024.
  * Joudaki et al. [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.
  * Lan et al. [2020] 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.
  * Mickus et al. [2024] Timothee Mickus, Stig-Arne Grönroos, and Joseph Attieh.  Isotropy, clusters, and classifiers.  _ArXiv_ , abs/2402.03191, 2024.
  * Noci et al. [2022] 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.
  * Paszke et al. [2019] 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.
  * Radford et al. [2019] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners.  2019\.
  * Raffel et al. [2020] 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.
  * Roy et al. [2020] Aurko Roy, Mohammad Taghi Saffar, Ashish Vaswani, and David Grangier.  Efficient content-based sparse attention with routing transformers.  _Transactions of the Association for Computational Linguistics_ , 2020.
  * Shi et al. [2022] 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.
  * Tian et al. [2023] 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.
  * Vaswani et al. [2017] 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.
  * Wolf et al. [2020] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush.  Transformers: State-of-the-art natural language processing.  In _EMNLP_ , 2020.
  * Wu et al. [2023] Xinyi Wu, Amir Ajorlou, Zihui Wu, and Ali Jadbabaie.  Demystifying oversmoothing in attention-based graph neural networks.  In _NeurIPS_ , 2023.
  * Xiong et al. [2020] 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.
  * Yun et al. [2020a] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar.  Are transformers universal approximators of sequence-to-sequence functions?  In _ICLR_ , 2020a.
  * Yun et al. [2020b] 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_ , 2020b.
  * Zaheer et al. [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.

##  Appendix A Comment about Section 3.3 in Dong et al. [10]

From the mathematical form of LayerNorm in Eq. 3, since LayerNorm normalizes
each row of $\tilde{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. [10], 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.

##  Appendix B Proof of Theorem 1

###  B.1 Auxiliary results

######  Lemma 1.

Under A1-A3, there exists $\epsilon>0$ such that $A^{(t)}_{i,j}\geq\epsilon$
for all $t\geq 0$, $(j,i)\in E$.

###### Proof.

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

$X^{(t+1)}=A^{(t)}...A^{(0)}X^{(0)}W_{V}^{(0)}...W_{V}^{(t)}\,,$ (9)

stay uniformly bounded for all $t\geq 0$ by A3. Then it follows from A2 that
there exists $C\in\mathbb{R}$ such that for all $t\geq 0$,

$\left\|\left(X^{(t)}W_{Q}^{(t)}\right)_{i,:}\right\|_{2}=\left\|X^{(t)}_{i,:}W_{Q}^{(t)}\right\|_{2}\leq
C\,,$
$\left\|\left(X^{(t)}W_{K}^{(t)}\right)_{i,:}\right\|_{2}=\left\|X^{(t)}_{i,:}W_{K}^{(t)}\right\|_{2}\leq
C\,.$

Hence for all $i,j\in[N]$,

$-C^{2}\leq(X^{(t)}W^{(t)}_{Q}(X^{(t)}W^{(t)}_{K})^{\top})_{i,j}\leq C^{2}\,.$

This implies that there exists $\epsilon>0$ such that
$A^{(t)}_{i,j}\geq\epsilon$ for all $(j,i)\in E$. ∎

Fix a vector $x^{(0)}\in\mathbb{R}^{d}$ and a $\\{A^{(n)}\\}_{n=0}^{\infty}$
in $\mathcal{A}_{\mathcal{G},\epsilon}$. Let
$x^{(t)}\vcentcolon=A^{(t)}...A^{(0)}x^{(0)}$. We further denote that

$\underset{i\in d}{\max}~{}x_{i}^{(t)}\vcentcolon=M^{(t)},\qquad\underset{i\in
d}{\min}~{}x_{i}^{(t)}\vcentcolon=m^{(t)}\,.$

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

######  Lemma 2.

Let $x_{i}^{(t)}=pm^{(t)}+(1-p)M^{(t)}$ for $0\leq p\leq 1$. Then for any
$T\in\mathbb{N}$,

$x_{i}^{(t+T)}\leq
p\left(\prod_{n=t}^{t+T-1}A_{ii}^{(n)}\right)m^{(t)}+\left(1-p\left(\prod_{n=t}^{t+T-1}A_{ii}^{(n)}\right)\right)M^{(t)}\,,$
(10)

and

$x_{i}^{(t+T)}\geq
p\left(\prod_{n=t}^{t+T-1}A_{ii}^{(n)}\right)M^{(t)}+\left(1-p\left(\prod_{n=t}^{t+T-1}A_{ii}^{(n)}\right)\right)m^{(t)}\,.$
(11)

###### Proof.

Given $x_{i}^{(t)}=pm^{(t)}+(1-p)M^{(t)}$, we get that

$\displaystyle x_{i}^{(t+1)}$
$\displaystyle=\sum_{j=1}^{N}A_{ij}^{(t)}x_{j}^{(t)}$ $\displaystyle\leq
A_{ii}^{(t)}(pm^{(t)}+(1-p)M^{(t)})+(1-A_{ii}^{(t)})M^{(t)}$
$\displaystyle=pA_{ii}^{(t)}m^{(t)}+(1-pA_{ii}^{(t)})M^{(t)}\,.$

Subsequently,

$\displaystyle x_{i}^{(t+2)}$
$\displaystyle=\sum_{j=1}^{N}A_{ij}^{(t+1)}x_{j}^{(t+1)}$ $\displaystyle\leq
A_{ii}^{(t+1)}(pA_{ii}^{(t)}m^{(t)}+(1-pA_{ii}^{(t)})M^{(t)})+(1-A_{ii}^{(t+1)})M^{(t)}$
$\displaystyle=p\left(\prod_{n=t}^{t+1}A_{ii}^{(s)}\right)m^{(t)}+\left(1-p\left(\prod_{n=t}^{t+1}A_{ii}^{(s)}\right)\right)M^{(t)}\,.$

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

######  Lemma 3.

Let $\\{A^{(n)}\\}_{n=0}^{\infty}$ in $\mathcal{A}_{\mathcal{G},\epsilon}$ and
$r$ be the diameter of $\mathcal{G}$. Then for $t\geq 0$ and $0\leq p\leq 1$,

$M^{(t+kr)}-m^{(t+kr)}\leq\left(1-p\epsilon^{r}\right)^{k}(M^{(t)}-m^{(t)})\,,\quad
k\in\mathbb{N}_{\geq 0}.$ (12)

###### Proof.

Let $i_{0}\in\mathcal{G}$ be a center token and fix $t\geq 0$ and $0\leq p\leq
1$. Without loss of generality, suppose

$x_{i_{0}}^{(t)}\leq pm^{(t)}+(1-p)M^{(t)}\,.$ (13)

From Eq. 11, we get that for all $T\in\mathbb{N}$,

$\displaystyle x_{i_{0}}^{(t+T)}$ $\displaystyle\leq
p\left(\prod_{s=t}^{t+T-1}A_{i_{0}i_{0}}^{(s)}\right)m^{(t)}+\left(1-p\left(\prod_{s=t}^{t+T-1}A_{i_{0}i_{0}}^{(s)}\right)\right)M^{(t)}$
$\displaystyle\leq
p\epsilon^{T}m^{(t)}+\left(1-p\epsilon^{T}\right)M^{(t)}\,.$

Denote $\mathcal{V}_{k}$ as the set of tokens that are exactly $k$-hop
neighbors of $i_{0}$:

$\mathcal{V}_{k}\vcentcolon=\\{j\in\mathcal{V}:\operatorname{dist}(i_{0},j)=k\\}\,.$

For any $i_{1}\in\mathcal{V}_{1}$, it follows that

$\displaystyle x_{i_{1}}^{(t+1)}$ $\displaystyle\leq
A_{i_{1}i_{0}}^{(t)}x_{i_{0}}^{(t)}+(1-A_{i_{1}i_{0}}^{(t)})M^{(t)}$
$\displaystyle\leq\epsilon\left(pm^{(t)}+(1-p)M^{(t)}\right)+(1-\epsilon)M^{(t)}$
$\displaystyle=p\epsilon m^{(t)}+\left(1-p\epsilon\right)M^{(t)}\,.$ (14)

Apply Eq. 11 to (14), we further get that for $T\in\mathbb{N}$,

$\displaystyle x_{i_{1}}^{(t+T)}\leq
p\epsilon^{T}m^{(t)}+\left(1-p\epsilon^{T}\right)M^{(t)}\,.$

Then let $i_{2}\in\mathcal{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

$\displaystyle x_{i_{2}}^{(t+2)}$ $\displaystyle\leq
A_{i_{2}i_{1}}^{(t+1)}x_{i_{1}}^{(t+1)}+(1-A_{i_{2}i_{1}}^{(t+1)})M^{(t)}$
$\displaystyle=\epsilon\left(p\epsilon
m^{(t)}+\left(1-p\epsilon\right)M^{(t)}\right)+(1-\epsilon)M^{(t)}$
$\displaystyle=p\epsilon^{2}m^{(t)}+\left(1-p\epsilon^{2}\right)M^{(t)}\,,$
(15)

and applying Eq. 11 to (15), it follows that for $T\in\mathbb{N}_{\geq 2}$

$\displaystyle x_{i_{2}}^{(t+T)}$ $\displaystyle\leq
p\epsilon^{T}m^{(t)}+\left(1-p\epsilon^{T}\right)M^{(t)}\,.$

Iterating this process for tokens in $\mathcal{V}_{3},...,\mathcal{V}_{r}$,
where $r$ is the diameter of $\mathcal{G}$, we get that for any token
$i\in\mathcal{G}$,

$\displaystyle x_{i}^{(t+r)}\leq
p\epsilon^{r}m^{(t)}+\left(1-p\epsilon^{r}\right)M^{(t)}\,,$

meaning that

$\displaystyle M^{(t+r)}\leq
p\epsilon^{r}m^{(t)}+\left(1-p\epsilon^{r}\right)M^{(t)}\,.$

We thus conclude that

$\displaystyle M^{(t+r)}-m^{(t+r)}$ $\displaystyle\leq
p\epsilon^{r}m^{(t)}+\left(1-p\epsilon^{r}\right)M^{(t)}-m^{(t)}$
$\displaystyle=\left(1-p\epsilon^{r}\right)\left(M^{(t)}-m^{(t)}\right)$ (16)

For the other case of (13), i.e. $x_{i_{0}}^{(t)}>pm^{(t)}+(1-p)M^{(t)}$, (16)
is obtained using a symmetric argument by bounding $m^{(t+r)}$ from below
through (11). This completes the proof of (12). ∎

######  Corollary 2.

There exists $C(x^{(0)})>0$ such that

$\displaystyle M^{(t)}-m^{(t)}\leq
C\left(1-\epsilon^{r}\right)^{t/r}\,,\forall t\geq 0\,,$

where $r$ is the diameter of the graph $\mathcal{G}$.

###  B.2 Proof of Theorem 1

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

$\displaystyle X^{(t+1)}_{:,j}$
$\displaystyle=(A^{(t)}...A^{(0)}X^{(0)}W_{V}^{(0)}...W_{V}^{(t)})_{:,j}$
$\displaystyle=A^{(t)}...A^{(0)}X^{(0)}\left(W_{V}^{(0)}...W_{V}^{(t)}\right)_{:,j}$
$\displaystyle=\sum_{i=1}^{d}{\tilde{W}^{(t)}}_{ij}A^{(t)}...A^{(0)}X^{(0)}_{:,i}$
(17)

where $\tilde{W}^{(t)}\vcentcolon=W_{V}^{(0)}...W_{V}^{(t)}$. This means that
for any $j\in[d]$, there exists $C_{j}>0$ such that

$\displaystyle\left|X^{(t+1)}_{m,j}-X^{(t+1)}_{n,j}\right|$
$\displaystyle=\left|\sum_{i=1}^{d}\tilde{W}_{ij}((A^{(t)}...A^{(0)}X^{(0)}_{:,i})_{m}-(A^{(t)}...A^{(0)}X^{(0)}_{:,i})_{n})\right|$
$\displaystyle\leq\sum_{i=1}^{d}\left|\tilde{W}_{ij}\right|\left|(A^{(t)}...A^{(0)}X^{(0)}_{:,i})_{m}-(A^{(t)}...A^{(0)}X^{(0)}_{:,i})_{n}\right|$
$\displaystyle\leq\sum_{i=1}^{d}\left|\tilde{W}_{ij}\right|\left|\underset{l\in[N]}{\max}(A^{(t)}...A^{(0)}X^{(0)}_{:,i})_{l}-\underset{l\in[N]}{\min}(A^{(t)}...A^{(0)}X^{(0)}_{:,i})_{l}\right|$
$\displaystyle\leq
C_{j}\sum_{i=1}^{d}\left|\tilde{W}_{ij}\right|\left(1-\epsilon^{r}\right)^{t/r}$
(18) $\displaystyle\leq C_{j}\left(1-\epsilon^{r}\right)^{t/r}\quad\forall
m,n\in[N]\,,$ (19)

where (18) follows from (2) and (19) follows from A3.

Then by (19),

$\displaystyle\mu(X^{(t)})$
$\displaystyle=\|X^{(t)}-\mathbf{1}\mathbf{1}^{\top}X^{(t)}/N\|_{F}=\sqrt{\sum_{j=1}^{d}\|X^{(t+1)}_{:,j}-\mathbf{1}\mathbf{1}^{\top}X_{:,j}^{(t)}/N\|_{2}^{2}}$
$\displaystyle=\sqrt{\frac{1}{2N}\sum_{j=1}^{d}\sum_{m=1}^{N}\sum_{n=1}^{N}|X^{(t)}_{m,j}-X^{(t)}_{n,j}|^{2}}$
$\displaystyle\leq\sqrt{\frac{N}{2}\sum_{j=1}^{d}C_{j}^{2}\left(1-\epsilon^{r}\right)^{2t/r}}\leq\sqrt{C\left(1-\epsilon^{r}\right)^{2t/r}}$
$\displaystyle=C^{\prime}\left(1-\epsilon^{r}\right)^{t/r}\,,$

where $C\vcentcolon=\frac{N}{2}\sum_{j=1}^{d}C_{j}^{2}$ and
$C^{\prime}\vcentcolon=\sqrt{C}$ here.

##  Appendix C Proof of Theorem 2

###  C.1 Auxiliary results

######  Lemma 4.

Under A1 and A2, there exists $\epsilon>0$ such that
$A^{(t)}_{i,j}\geq\epsilon$ for all $t\geq 0$, $(j,i)\in E$.

###### Proof.

Note due to the way the dynamical system is defined in (3), for every $t\geq
0$, every token $X^{(t)}_{i,:}$ lies within the unit sphere
$\mathbb{S}^{d-1}$. Then it follows from A2 that there exists $C\in\mathbb{R}$
such that for all $t\geq 0$,

$\left\|\left(X^{(t)}W_{Q}^{(t)}\right)_{i,:}\right\|_{2}=\left\|X^{(t)}_{i,:}W_{Q}^{(t)}\right\|_{2}\leq
C\,,$
$\left\|\left(X^{(t)}W_{K}^{(t)}\right)_{i,:}\right\|_{2}=\left\|X^{(t)}_{i,:}W_{K}^{(t)}\right\|_{2}\leq
C\,.$

Hence for all $i,j\in[N]$,

$-C^{2}\leq(X^{(t)}W^{(t)}_{Q}(X^{(t)}W^{(t)}_{K})^{\top})_{i,j}\leq C^{2}\,.$

This implies that there exists $\epsilon>0$ such that
$A^{(t)}_{i,j}\geq\epsilon$ for all $(i,j)\in E$. ∎

######  Lemma 5.

Under the same assumptions as Theorem 2, there exists $C_{1},C_{2}>0$ such
that $C_{1}\leq D^{(t)}_{i,i}\leq C_{2}$ for all $t\geq 0$, $i\in[N]$.

###### Proof.

We adopt the following notation:

######  Notation 1.

For $X\in\mathbb{R}^{N\times d}$, we use $\operatorname{Conv}(X)$ to denote
the convex hull of $\\{X_{1,:},...,X_{N,:}\\}$.

First, since for $x\in\operatorname{Conv}(X^{(t)})$, ${(W^{(t)})}^{\top}x$
always lies in $\mathbb{S}^{d-1}$, $D^{(t)}_{i,i}\geq 1$.

Then for all $t\geq 0$, each row vector of $(A^{(t)}X^{(t)}W^{(t)})_{i,:}$
lies in the convex cone generated by $X^{(0)}\tilde{W}^{(t)}$, where
$\tilde{W}^{(t)}\vcentcolon=\Pi_{i=0}^{t}W^{(i)}$ . To see this,

$\displaystyle(A^{(t)}X^{(t)}W^{(t)})_{i,:}$
$\displaystyle=\sum_{j_{0}=1}^{N}A^{(t)}_{ij_{0}}X^{(t)}_{j_{0},:}W^{(t)}$
$\displaystyle=\sum_{j_{0}=1}^{N}A^{(t)}_{ij_{0}}D_{j_{0},j_{0}}^{(t-1)}\left(\sum_{j_{1}=1}^{N}A^{(t-1)}_{j_{0},j_{1}}X^{(t-1)}_{j_{1},:}W^{(t-1)}\right)W^{(t)}$
$\displaystyle=\sum_{(j_{0},...,j_{t})\in[N]^{t+1}}A^{(t)}_{ij_{0}}D_{j_{0},j_{0}}^{(t-1)}A^{(t-1)}_{j_{0},j_{1}}D^{(t-2)}_{j_{1},j_{1}}...D^{(0)}_{j_{t-1},j_{t-1}}A^{(0)}_{j_{t-1},j_{t}}X^{(0)}_{j_{t},:}\tilde{W}^{(t)}$

Hence all $t\geq 0$,

$(A^{(t)}X^{(t)}W^{(t)})_{i,:}=\sum_{j=1}^{N}a^{(t)}_{ij}X_{j,:}^{(0)}\tilde{W}^{(t)}\,,$

where $a^{(t)}_{ij}\geq 0,\forall j\in[N]$ and $\sum_{j=1}^{N}a^{(t)}_{ij}\geq
1$.

Given by assumption, all the points of $X^{(0)}$ lie in an open hemisphere
$\mathcal{H}$. Since all the points of $X^{(0)}$ lie in an open hemisphere
$\mathcal{H}$, there exists $\gamma>0$ such that for all
$x\in\operatorname{Conv}(X^{(0)})$ and $W$ orthogonal,

$\|Wx\|_{2}=\|x\|_{2}\geq\gamma\,.$

Then

$\displaystyle\|(A^{(t)}X^{(t)}W^{(t)})_{i,:}\|_{2}$
$\displaystyle=\left\|\left(\sum_{j=1}^{N}a^{(t)}_{ij}\right)\frac{\sum_{j=1}^{N}a^{(t)}_{ij}X_{j,:}^{(0)}\tilde{W}^{(t)}}{\sum_{j=1}^{N}a^{(t)}_{ij}}\right\|_{2}$
$\displaystyle\geq\left\|\frac{\sum_{j=1}^{N}a^{(t)}_{ij}X_{j,:}^{(0)}\tilde{W}^{(t)}}{\sum_{j=1}^{N}a^{(t)}_{ij}}\right\|_{2}$
$\displaystyle\geq\gamma\,.$

We conclude that $D^{(t)}_{i,i}\leq 1/\gamma$, for all $t\geq 0$.

∎

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

$X^{(t+1)}=D^{(t)}A^{(t)}...D^{(0)}A^{(0)}X^{(0)}W^{(0)}...W^{(t)}\,.$

We focus on the infinite product
$\lim_{t\to\infty}D^{(t)}A^{(t)}...D^{(0)}A^{(0)}$. We first define a family
of attention matrices as follows:

######  Definition 6.

Let $\epsilon>0$. We define $\mathcal{A}_{\mathcal{G},\epsilon}$ to be the set
of row-stochastic matrices satisfying the following conditions:

  1. 1.

$\epsilon\leq A_{ij}\leq 1,\,\text{if }(j,i)\in E(\mathcal{G})$,

  2. 2.

$A_{ij}=0,\,\text{if }(j,i)\notin E$.

We also define

$\mathcal{D}_{C_{1},C_{2}}=\\{\operatorname{diag}(\mathbf{d}):C_{1}\operatorname{\leq_{ew}}\mathbf{d}\operatorname{\leq_{ew}}C_{2}\\}\,.$
(20)

and subsequently,

$\mathcal{M}_{\mathcal{G},\epsilon,C_{1},C_{2}}=\\{DA:D\in\mathcal{D}_{C_{1},C_{2}},A\in\mathcal{A}_{\mathcal{G},\epsilon}\\}\,.$
(21)

To study the infinite product of matrices from
$\mathcal{M}_{\mathcal{G},\epsilon,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}^{\infty}$.
Define the partial product

$P^{(t)}\vcentcolon=M^{(t)}...M^{(0)}\,.$

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

$\lim_{t\to\infty}\frac{P_{ij}^{(t)}}{S_{ij}^{(t)}}=1\,,\quad\forall
i,j\in[N]\,.$ (22)

Our goal is to show the following key result:

######  Lemma 6.

Let $\epsilon,C_{1},C_{2}>0$. Given $\\{D^{(t)}A^{(t)}\\}_{t=0}^{\infty}$ in
$\mathcal{M}_{\mathcal{G},\epsilon,C_{1},C_{2}}$. The sequence of partial
products $\\{P^{(t)}\\}_{t=0}^{\infty}$ is ergodic.

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

######  Lemma 7 (Corollary 5.1 [16]).

Consider a sequence of row allowable matrices $\\{M^{(t)}\\}_{t=0}^{\infty}$.
Let $a_{t}$ and $b_{t}$ be the smallest and largest entries in $M^{(t)}$,
respectively. If $\sum_{t=0}^{\infty}\frac{a_{t}}{b_{t}}=\infty$, then
$\\{P^{(t)}\\}_{t=0}^{\infty}$ is ergodic.

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

$P^{(t_{1}:t_{0})}=P^{(t_{1})}\ldots P^{(t_{0})}\,.$

######  Lemma 8.

Under A1, if $\mathcal{G}$ is strongly connected, there exist $T\in\mathbb{N}$
and $C_{3},C_{4}>0$ such that for all $t\geq 0$,

$C_{3}\leq P^{(t+T:t)}_{i,j}\leq C_{4}\,,\quad\forall 1\leq i,j\leq N\,.$

###### Proof.

Note that by A1 and the strong connectivity of $\mathcal{G}$, there exists
$T\in\mathbb{N}$, e.g. the diameter of $\mathcal{G}$, and
$\epsilon,C_{1},C_{2}>0$ such that for all $t\geq 0$,

$(\epsilon C_{1})^{T}\leq P^{(t+T:t)}_{i,j}\leq C_{2}^{T},\quad\forall 1\leq
i,j\leq N.$

∎

Then we apply Lemma 7 to $\\{P^{(k+1)T:T}\\}_{k=0}^{\infty}$ and concludes
Lemma 6.

As a result, from (22) we see that any $P^{(t)}$ can now be written in the
following form:

$P^{(t)}=u^{(t)}(v^{(t)})^{\top}+E^{(t)}\,,$

where $u^{(t)},v^{(t)}$ are positive vectors such that
$u^{(t)}(v^{(t)})^{\top}=S^{(t)}$. Without loss of generality, assume
$\|v^{(t)}\|_{2}=1$.

Then it follow from (22) that

######  Lemma 9.

$\lim_{t\to\infty}\frac{E_{ij}^{(t)}}{u_{i}^{(t)}}=0,\quad\forall
i,j\in[N]\,.$

###### Proof.

Given

$\left|v_{j}^{(t)}\frac{E_{ij}^{(t)}}{u_{i}^{(t)}v_{j}^{(t)}}\right|\leq\left|\frac{E_{ij}^{(t)}}{u_{i}^{(t)}v_{j}^{(t)}}\right|$

where by (22), since also

$\lim_{t\to\infty}\left|\frac{E_{ij}^{(t)}}{u_{i}^{(t)}v_{j}^{(t)}}\right|=0\,,$

We get that

$\lim_{t\to\infty}\left|\frac{E_{ij}^{(t)}}{u_{i}^{(t)}}\right|=0\,,$

and hence

$\lim_{t\to\infty}\frac{E_{ij}^{(t)}}{u_{i}^{(t)}}=0\,.$

∎

Thus the $2$-norm of the $i^{th}$ token of $X^{(t+1)}$,
$\|X^{(t+1)}_{i,:}\|_{2}$ is

$\|P^{(t)}_{i,:}X^{(0)}\tilde{W}^{(t)}\|_{2}=\|u_{i}^{(t)}(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}+E_{i,:}^{(t)}X^{(0)}\tilde{W}^{(t)}\|_{2}=1\,,\forall
i\in[N].$ (23)

Then by Lemma 9,

$\displaystyle\lim_{t\to\infty}\left\|(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}+\frac{E_{i,:}^{(t)}}{u_{i}^{(t)}}X^{(0)}\tilde{W}^{(t)}\right\|_{2}=\lim_{t\to\infty}\left\|(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}\right\|_{2}=\lim_{t\to\infty}\frac{1}{u_{i}^{(t)}}\,.$

Note that

$\underset{\|v\|_{2}=1}{\min}\left\|v^{\top}X^{(0)}\tilde{W}^{(t)}\right\|_{2}=\sigma_{\min}((X^{(0)}\tilde{W}^{(t)})^{\top})\,,$

where $\sigma_{\min}(\cdot)$ denotes the smallest singular value. Since we
assume $X^{(0)}$ has full rank, and $\tilde{W}^{(t)}$ is orthogonal

$\lim_{t\to\infty}\left\|(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}\right\|_{2}\geq\sigma_{\min}((X^{(0)}\tilde{W}^{(t)})^{\top})>0\,.$

Then given $\forall i\in[N]$,

$\lim_{t\to\infty}\left\|(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}\right\|_{2}=\lim_{t\to\infty}\frac{1}{u_{i}^{(t)}}\,,\forall
i\in[N]\,,$

as a result,

$\lim_{t\to\infty}u_{i}^{(t)}=\lim_{t\to\infty}u_{k}^{(t)}<\infty\,,\forall
i,k\in[N]\,.$ (24)

This together with Lemma 9 implies that

$\lim_{t\to\infty}E_{ij}^{(t)}=0,\forall i\in[N],j\in[d]\,,$

and

$\displaystyle\lim_{t\to\infty}\mu(X^{(t)})=\lim_{t\to\infty}\|Bu^{(t)}(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}+BE^{(t)}X^{(0)}\tilde{W}^{(t)}\|_{F}=0\,.$
(25)

####  C.1.1 Exponential convergence rate

Having established the convergence in (25), we then establish the exponential
convergence rate.

We define that

$\alpha^{(t)}\vcentcolon=\underset{i,j\in[N]}{\min}\langle
X^{(t)}_{i,:},X^{(t)}_{j,:}\rangle.$

Similar to the derivation of (24), note that

$\underset{\|v\|_{2}=1}{\max}\left\|v^{\top}X^{(0)}\tilde{W}^{(t)}\right\|_{2}=\sigma_{\max}((X^{(0)}\tilde{W}^{(t)})^{\top})\,,$

it follows that

$\lim_{t\to\infty}\left\|(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}\right\|_{2}\leq\sigma_{\max}((X^{(0)}\tilde{W}^{(t)})^{\top})<\infty\,.$

Since $\forall i\in[N]$,

$\lim_{t\to\infty}\left\|(v^{(t)})^{\top}X^{(0)}\tilde{W}^{(t)}\right\|_{2}=\lim_{t\to\infty}\frac{1}{u_{i}^{(t)}}\,,\forall
i\in[N]\,$

as a result,

$\lim_{t\to\infty}u_{i}^{(t)}=\lim_{t\to\infty}u_{k}^{(t)}>0,\forall
i,k\in[N]\,,$

meaning that $\lim_{t\to\infty}X_{i,:}^{(t)}\in S^{d-1}$ for all $i\in[N]$ .

Then given the convergence established in (25), we get that
$\lim_{t\to\infty}\alpha^{(t)}=1\,,$ from which there exists
$t_{0}\in\mathbb{N}$ such that $\alpha^{(t)}>0$, $\forall t\geq t_{0}$. It
follows that for $t\geq t_{0}$,

$\displaystyle\phi^{(t+r)}$ $\displaystyle=\underset{i,j\in[N]}{\min}\langle
X^{(t+r)}_{i,:},X^{(t+r)}_{i,:}\rangle=\langle
X^{(t+r)}_{i^{(t+r)},:},X^{(t+r)}_{j^{(t+r)},:}\rangle$
$\displaystyle\geq\langle\sum_{k_{1}=1}^{N}A^{(t+r-1)}_{i^{(t+r)},k_{1}}X^{(t+r-1)}_{k_{1}},\sum_{l_{1}=1}^{N}A^{(t+r-1)}_{j^{(t+r)},l_{1}}X^{(t+r-1)}_{l_{1}}\rangle$
$\displaystyle\geq\sum_{(k_{1},...,k_{r})\in[N]^{d}}\sum_{(l_{1},...,l_{r})\in[N]^{d}}A^{(t+r-1)}_{i^{(t+r)},k_{1}}...A^{(t)}_{k_{r-1},k_{r}}A^{(t+r-1)}_{j^{(t+r)},l_{1}}...A^{(t)}_{l_{r-1},l_{r}}\langle
X^{(t)}_{k_{r}},X^{(t)}_{l_{r}}\rangle$ $\displaystyle\geq
N\epsilon^{2r}+(1-N\epsilon^{2r})\phi^{(t)}\,.$

Thus

$1-\phi^{(t+r)}\leq(1-N\epsilon^{2r})(1-\phi^{(t)}).$

Writing recursively, we get that for $k\geq 0$,

$1-\phi^{(t_{0}+kr)}\leq(1-N\epsilon^{2r})^{k}(1-\phi^{(t_{0})})\,,$

Let $C\vcentcolon=2/(1-N\epsilon^{2r})^{t_{0}/r}$. Hence for all $t\geq 0$,

$1-\alpha^{(t)}\leq C(1-N\epsilon^{2r})^{\frac{t}{r}}\,.$ (26)

Since by definition, $1-\phi^{(t)}\geq 1-\langle
X^{(t)}_{i,:},X^{(t)}_{j,:}\rangle\geq\|X_{i,:}^{(t)}-X^{(t)}_{j,:}\|_{2}^{2}/2$,
it follows that

$\displaystyle\mu(X^{(t)})$
$\displaystyle=\|X^{(t)}-\mathbf{1}\mathbf{1}^{\top}X^{(t)}/N\|_{F}=\sqrt{\sum_{i=1}^{N}\|X^{(t)}_{i,:}-\mathbf{1}^{\top}X^{(t)}/N\|_{2}^{2}}$
$\displaystyle=\sqrt{\frac{1}{2N}\sum_{i=1}^{N}\sum_{j=1}^{N}\|X^{(t)}_{i,:}-X^{(t)}_{j,:}\|_{2}^{2}}$
$\displaystyle\leq C^{\prime}(1-N\epsilon^{2r})^{\frac{t}{2r}}\,,$

where $C^{\prime}=\sqrt{CN}$, completing the proof.

##  Appendix D Proof of Corollary 1

First, since for $x\in\operatorname{Conv}(X^{(t)})$, ${(W^{(t)})}^{\top}x$
always lies in $\mathbb{S}^{d-1}$, $D^{(t)}_{i,i}\geq 1$, for all
$i\in[N],t\geq 0$.

For all $t\geq 0$, since $\phi^{(0)}\geq 0$ and $\mathcal{G}$ is quasi-
strongly connected, we have that

$\displaystyle\phi^{(t+r)}$ $\displaystyle=\underset{i,j\in[N]}{\min}\langle
X^{(t+r)}_{i,:},X^{(t+r)}_{i,:}\rangle=\langle
X^{(t+r)}_{i^{(t+r)},:},X^{(t+r)}_{j^{(t+r)},:}\rangle$
$\displaystyle\geq\langle\sum_{k_{1}=1}^{N}A^{(t+r-1)}_{i^{(t+r)},k_{1}}X^{(t+r-1)}_{k_{1}}W^{(t+r-1)},\sum_{l_{1}=1}^{N}A^{(t+r-1)}_{j^{(t+r)},l_{1}}X^{(t+r-1)}_{l_{1}}W^{(t+r-1)}\rangle$
$\displaystyle=\langle\sum_{k_{1}=1}^{N}A^{(t+r-1)}_{i^{(t+r)},k_{1}}X^{(t+r-1)}_{k_{1}},\sum_{l_{1}=1}^{N}A^{(t+r-1)}_{j^{(t+r)},l_{1}}X^{(t+r-1)}_{l_{1}}\rangle$
$\displaystyle\geq\sum_{(k_{1},...,k_{r})\in[N]^{d}}\sum_{(l_{1},...,l_{r})\in[N]^{d}}A^{(t+r-1)}_{i^{(t+r)},k_{1}}...A^{(t)}_{k_{r-1},k_{r}}A^{(t+r-1)}_{j^{(t+r)},l_{1}}...A^{(t)}_{l_{r-1},l_{r}}\langle
X^{(t)}_{k_{r}},X^{(t)}_{l_{r}}\rangle$ $\displaystyle\geq
N\epsilon^{2r}+(1-N\epsilon^{2r})\phi^{(t)}\,.$

Thus

$1-\phi^{(t+r)}\leq(1-N\epsilon^{2r})(1-\phi^{(t)}).$

Writing recursively, we get that for $k\geq 0$,

$1-\phi^{(kr)}\leq(1-N\epsilon^{2r})^{k}(1-\phi^{(0)})\,,$

Let $C\vcentcolon=1/(1-N\epsilon^{2r})$. Hence for all $t\geq 0$,

$1-\phi^{(t)}\leq C(1-N\epsilon^{2r})^{\frac{t}{r}}\,.$ (27)

Since by definition, $1-\phi^{(t)}\geq 1-\langle
X^{(t)}_{i,:},X^{(t)}_{j,:}\rangle\geq\|X_{i,:}^{(t)}-X^{(t)}_{j,:}\|_{2}^{2}/2$,
it follows that

$\displaystyle\mu(X^{(t)})$
$\displaystyle=\|X^{(t)}-\mathbf{1}\mathbf{1}^{\top}X^{(t)}/N\|_{F}=\sqrt{\sum_{i=1}^{N}\|X^{(t)}_{i,:}-\mathbf{1}^{\top}X^{(t)}/N\|_{2}^{2}}$
$\displaystyle=\sqrt{\frac{1}{2N}\sum_{i=1}^{N}\sum_{j=1}^{N}\|X^{(t)}_{i,:}-X^{(t)}_{j,:}\|_{2}^{2}}$
$\displaystyle\leq C^{\prime}(1-N\epsilon^{2r})^{\frac{t}{2r}}\,,$

where $C^{\prime}=\sqrt{CN}$, completing the proof.

##  Appendix E Detailed analysis of the illustrative example in Section 4.3.1

##### First token

Given that

$\frac{X_{1,2}^{(t+1)}}{X_{1,1}^{(t+1)}}=\frac{X_{1,2}^{(t)}}{X_{1,1}^{(t)}}+w\,,$

we get that

$\frac{X_{1,2}^{(t+1)}}{X_{1,1}^{(t+1)}}=\frac{X_{1,2}^{(0)}}{X_{1,1}^{(0)}}+wt\,.$

and thus

$\lim_{t\to\infty}\frac{X^{(t)}_{1,2}}{X^{(t)}_{1,1}}=\infty\,.$

Since the sign of $X_{1,1}^{(t)}$ is invariant, depending on its initial sign,
we get that

$\lim_{t\to\infty}X^{(t)}_{1,:}=\begin{cases}[0,1]\quad\text{ if
}X^{(0)}_{1,1}>0\\\ [0,-1]\quad\text{ if }X^{(0)}_{1,1}<0\end{cases}\,.$

##### Second token

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

$\frac{X_{2,2}^{(t+1)}}{X_{2,1}^{(t+1)}}=\frac{X_{2,2}^{(t)}}{X_{2,1}^{(t)}}+w+\frac{1}{X_{2,1}^{(t)}}\vcentcolon=g\left(\frac{X_{2,2}^{(t)}}{X_{2,1}^{(t)}}\right)\,,$

Solving
$\frac{X_{2,2}^{(t+1)}}{X_{2,1}^{(t+1)}}=\frac{X_{2,2}^{(t)}}{X_{2,1}^{(t)}}$
gives $X_{2,1}=-1/w$ and

$\begin{cases}\frac{X_{2,2}^{(t+1)}}{X_{2,1}^{(t+1)}}>\frac{X_{2,2}^{(t)}}{X_{2,1}^{(t)}}\quad\text{
if }X_{2,1}<-1/w\\\
\frac{X_{2,2}^{(t+1)}}{X_{2,1}^{(t+1)}}<\frac{X_{2,2}^{(t)}}{X_{2,1}^{(t)}}\quad\text{
otherwise}\,.\end{cases}$ (28)

To show the convergence, we utilize the following useful property of
$g(\cdot)$:

######  Lemma 10.

$g$ is non-decreasing in $X_{2,2}/X_{2,1}$.

###### Proof.

Let $y=X_{2,2}/X_{2,1}$. Given the identity that

$\left(\frac{1}{X_{2,1}^{(t)}}\right)^{2}=y^{2}+1\,,$

we get that

$2\frac{\partial\left(\frac{1}{X_{2,1}}\right)}{\partial
y}\frac{1}{X_{2,1}}=2y\,.$

Hence

$\frac{\partial\left(\frac{1}{X_{2,1}}\right)}{\partial
y}=X_{2,1}y=X_{2,2}\,,$

which means that

$\frac{\partial g}{\partial y}=1+X_{2,1}\geq 0\,.$

∎

Given $X_{2,2}/X_{2,1}$ is invariant at $X_{2,1}=-1/w$ (symmetric points $A$
and $B$ in Fig. 5), together with Lemma 10, they imply the invariance of the 4
segments on $\mathbb{S}^{1}$, as shown in Fig. 5. Furthermore, (28) 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.

Figure 5: Convergence analysis of the second token in $d=2$.

##  Appendix F Proof of Theorem 3

###  F.1 Finding the equilibria

##### Step 1

Without loss of generality, we set that $W_{K}^{(t)}=W_{Q}^{(t)}=\mathbf{0}$,
for all $t\geq 0$. The system of equations corresponds to the equilibrium
condition for the attention dynamics with LayerNorm defined in (3) when the
attention mask $\mathcal{G}$ is causal is as follows:

$\begin{cases}X_{1,:}=F_{1}(X)=d_{1}X_{1,:}W_{V}\\\
X_{2,:}=F_{2}(X)=d_{2}\left(\frac{1}{2}X_{1,:}W_{V}+\frac{1}{2}X_{2,:}W_{V}\right)\\\
\vdots\\\
X_{N,:}=F_{N}(X)=d_{N}\left(\frac{1}{N}X_{1,:}W_{V}+...+\frac{1}{N}X_{N,:}W_{V}\right)\end{cases}$

where $d_{i}=i/\|\sum_{j=1}^{i}X_{j,:}W_{V}\|_{2}$. We define that
$\beta_{i}=d_{i}/i$ for $i\in[N]$. Then the above system equation becomes

$\begin{cases}X_{1,:}=\beta_{1}X_{1,:}W_{V}\\\
X_{2,:}=\beta_{2}\left(\frac{1}{\beta_{1}}X_{1,:}+X_{2,:}W_{V}\right)\\\
\vdots\\\
X_{N,:}=\beta_{N}\left(\frac{1}{\beta_{N-1}}X_{N-1,:}+X_{N,:}W_{V}\right)\end{cases}$

after substitution. Then for $i\in[N-1]$,

$X_{i+1,:}=\beta_{i+1}\left(\frac{1}{\beta_{i}}X_{i,:}+X_{i+1,:}W_{V}\right)\,,$
(29)

Let $k\in[N]$. If $\beta_{1}=...=\beta_{k}=1$, then for $i\in[k-1]$,

$X_{i+1,:}=X_{i,:}+X_{i+1,:}W_{V}\,,$ (30)

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

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

  1. 1.

For $k=1$, $W_{V}$ can be set as $I_{d}$;

  2. 2.

For $k=2$, $W_{V}$ can be set as

$W_{V}^{(t)}=\begin{bmatrix}1&0&0&...&0&0\\\ 0&1&0&...&0&0\\\
\vdots&\vdots&\vdots&...&\vdots&\vdots\\\ 0&0&0&...&0&0\\\ 0&0&0&...&1&w\\\
0&0&0&...&0&1\\\ \end{bmatrix}\,\quad\forall t\geq 0\,.$

  3. 3.

For $k=3$, $W_{V}$ can be set as

$W_{V}^{(t)}=\begin{bmatrix}1&0&0&...&0&0\\\ 0&1&0&...&0&0\\\
\vdots&\vdots&\vdots&...&\vdots&\vdots\\\ 0&0&0&...&w&0\\\ 0&0&0&...&1&w\\\
0&0&0&...&0&1\\\ \end{bmatrix}\,\quad\forall t\geq 0\,.$

  4. 4.

For $k=N$, $W_{V}$ can be set as

$W_{V}^{(t)}=\begin{bmatrix}1&w&0&...&0&0\\\ 0&1&w&...&0&0\\\
\vdots&\vdots&\vdots&...&\vdots&\vdots\\\ 0&0&0&...&w&0\\\ 0&0&0&...&1&w\\\
0&0&0&...&0&1\\\ \end{bmatrix}\,\quad\forall t\geq 0\,.$ (31)

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

$X_{i+1,:}\begin{bmatrix}0&-w&0&...&0&0\\\ 0&0&-w&...&0&0\\\
\vdots&\vdots&\vdots&...&\vdots&\vdots\\\ 0&0&0&...&-w&0\\\ 0&0&0&...&0&-w\\\
0&0&0&...&0&0\\\ \end{bmatrix}=X_{i,:}\,,$

which means that

$[0,-wX_{i+1,1},-wX_{i+1,2},...,-wX_{i+1,d-1}]=[X_{i,1},X_{i,2},X_{i,3},...,X_{i,d}]\,.$
(32)

Then setting $X_{1,:}=[0,0,...,0,1]$, iteratively solving (30), we get that an
equilibrium satisfying (29) is as follows:

$\begin{cases}X_{1,:}={[0,0,...,0,0,1]}\\\
X_{2,:}={\left[0,0,...,0,-\frac{1}{w},-\sqrt{1-\frac{1}{w^{2}}}\right]}\\\
X_{3,:}=\left[0,0,...,\frac{1}{w^{2}},\sqrt{\frac{1}{w^{2}}-\frac{1}{w^{4}}},\sqrt{1-\frac{1}{w^{2}}}\right]\\\
\vdots\\\ X_{N,:}=(-1)^{N-1}\begin{bmatrix}\frac{1}{w^{N-1}}\\\
\sqrt{\frac{1}{w^{2(N-2)}}-\frac{1}{w^{2(N-1)}}}\\\ ...\\\
\sqrt{\frac{1}{w^{4}}-\frac{1}{w^{6}}}\\\
\sqrt{\frac{1}{w^{2}}-\frac{1}{w^{4}}}\\\
\sqrt{1-\frac{1}{w^{2}}}\,.\end{bmatrix}^{\top}\,,\end{cases}$ (33)

which is valid if $w>1$. Since from (30), there is a free coordinate, in fact
there a $2^{N}$ many such equilibria. Hence there exists
$\\{W_{K}^{(t)},W_{Q}^{(t)},W_{V}^{(t)}\\}_{t=0}^{\infty}$ 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_{K}^{(t)},W_{Q}^{(t)},W_{V}^{(t)}\\}_{t=0}^{\infty}$, consider
$W_{V}^{(t)}$ to be the case of (31). 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 $\beta_{N-k+1}=1/(N-K+1)$. Then
recall (29):

$X_{i+1,:}=\beta_{i+1}\left(\frac{1}{\beta_{i}}X_{i,:}+X_{i+1,:}W_{V}\right)\,$

Let $\beta_{N-k+2}=1$, we get that

$X_{N-k+2,:}(I-W)=[0,0,...,0,N-k+1]\,,$

which yields

$X_{N-k+2,:}=\left[\mathbf{0}_{d-2},-\frac{N-k+1}{w},\pm\sqrt{1-\frac{(N-k+1)^{2}}{w^{2}}}\right]\,.$

Continue the process by setting $\beta_{i}=1$ for $i\geq N-k+2$, we get a rank
$k$ equilibrium as desired (there are $2^{k}$ such equilibria), as long as
$w\geq N$.

##### Step 3

Finally to ensure $\\{W_{K}^{(t)},W_{Q}^{(t)},W_{V}^{(t)}\\}_{t=0}^{\infty}$
satisfies A2-A3:

  * •

We see that $\\{W_{K}^{(t)},W_{Q}^{(t)}\\}_{t=0}^{\infty}$ satisfies A2 by
construction;

  * •

Since this is a dynamics with LayerNorm, applying $W_{V}^{(t)}$ is in fact
equivalent to applying $\hat{W}_{V}=W_{V}/\|W_{V}\|_{2}$ under the dynamics,
where the latter has $2$-norm at most $1$ and thus using $\hat{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:

$\operatorname{SRank}(X)=\frac{\|X\|^{2}_{F}}{\|X\|^{2}_{2}}\,.$

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

$1\leq\operatorname{SRank}(X^{*})\leq(1+\delta)+O(1/N)\,.$

Since for any matrix $M$, $\|M\|^{2}_{F}\geq\|M\|_{2}^{2}$, the first
inequality trivially holds. To show the second inequality, consider the
equilibrium in $\eqref{eq: equilibrium_sol}$. Then direct calculation yields
that

$\|X^{*}\|_{F}^{2}=\operatorname{Tr}(X^{*}(X^{*})^{\top})=N\,.$

Since by definition,

$\|X^{*}\|_{2}=\underset{\|v\|_{2}=1}{\sup}\|X^{*}v\|_{2}\,,$

and thus any $\|X^{*}v\|_{2}$ such that $v\in\mathbb{R}^{N},\|v\|_{2}=1$
serves as a lower bound for $\|X\|_{2}$.

We consider $v=[0,0,...,0,1]^{\top}\in\mathbb{R}^{d}$. Then it follows that

$\displaystyle\|X^{*}\|^{2}_{2}\geq 1+(N-1)\left(1-\frac{1}{w^{2}}\right)\,,$

which implies that

$\displaystyle\operatorname{SRank}(X^{*})=$
$\displaystyle\frac{\|X^{*}\|^{2}_{F}}{\|X^{*}\|_{2}^{2}}\leq\frac{N}{N-\frac{N-1}{w^{2}}}$

Hence for any $\delta>0$, it suffices to choose $w\geq 1$ such that

$w\geq\sqrt{\frac{1}{\delta}+1}\,.$

###  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

$\frac{X_{1,2}^{(t+1)}}{X_{1,1}^{(t+1)}}=\frac{X_{1,2}^{(t)}}{X_{1,1}^{(t)}}+w\,,$

we get that

$\frac{X_{1,2}^{(t+1)}}{X_{1,1}^{(t+1)}}=\frac{X_{1,2}^{(0)}}{X_{1,1}^{(0)}}+wt\,.$

Let $\frac{X_{1,2}^{(0)}}{X_{1,1}^{(0)}}=r_{2}w$, for $r_{2}\in\mathbb{R}$. It
follows from above that

$\frac{X_{1,2}^{(t)}}{X_{1,1}^{(t)}}=(r_{2}+t)w\,.$ (34)

Then for all $3\leq i\leq d$,

$\frac{X_{1,i}^{(t+1)}}{X_{1,1}^{(t+1)}}=w\frac{X_{1,i-1}^{(t)}}{X_{1,1}^{(t)}}+\frac{X_{1,i}^{(t)}}{X_{1,1}^{(t)}}\,,$
(35)

setting

$\frac{X_{1,i}^{(0)}}{X_{1,1}^{(0)}}=r_{i}w^{i-1}\,,$

where $r_{i}\in\mathbb{R}$, and substituting from the base case (34) to (35)
recursively, we get that

$\displaystyle\frac{X_{1,i}^{(t)}}{X_{1,1}^{(t)}}=w^{i-1}\Bigg{(}r_{i}$
$\displaystyle+r_{i-1}t+r_{i-2}\sum_{j_{1}=0}^{t-1}j_{1}+r_{i-3}\sum_{j_{2}=0}^{t-1}\sum_{j_{1}=0}^{j_{2}}j_{1}+...+$
$\displaystyle+r_{2}\sum_{j_{i-3}=0}^{t-1}...\sum_{j_{2}=0}^{j_{3}}\sum_{j_{1}=0}^{j_{2}}j_{1}...+\sum_{j_{i-2}=0}^{t-1}...\sum_{j_{2}=0}^{j_{3}}\sum_{j_{1}=0}^{j_{2}}j_{1}\Bigg{)}\,,$
(36)

which implies that

$\lim_{t\to\infty}\frac{X_{1,1}^{(t+1)}}{X_{1,i}^{(t+1)}}\frac{X_{1,i}^{(t)}}{X_{1,1}^{(t)}}=1\,.$
(37)

Further, for $i\geq 3$

$\displaystyle\frac{X_{1,i}^{(t+1)}}{X_{1,i-1}^{(t+1)}}$
$\displaystyle=\frac{wX_{1,i-1}^{(t)}+X_{1,i}^{(t)}}{wX_{1,i-2}^{(t)}+X_{1,i-1}^{(t)}}=\frac{X_{1,1}^{(t+1)}}{X_{1,i-1}^{(t+1)}}\frac{X_{1,i-1}^{(t)}}{X_{1,1}^{(t)}}\left(w+\frac{X_{1,i}^{(t)}}{X_{1,i-1}^{(t)}}\right)\,.$
(38)

Then (37) and (38) yields that

$\lim_{t\to\infty}\frac{X_{1,i}^{(t)}}{X_{1,i-1}^{(t)}}=\infty\,,$

which let us conclude that if $X_{1,1}^{(0)}>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\leq i\leq d-1$, it holds that

$\lim_{t\to\infty}\frac{X_{2,i}^{(t)}}{X_{2,i-1}^{(t)}}=\infty\,,$

and

$\displaystyle\frac{X_{2,i}^{(t)}}{X_{2,1}^{(t)}}=w^{i-1}\Bigg{(}r_{i}$
$\displaystyle+r_{i-1}t+r_{i-2}\sum_{j_{1}=0}^{t-1}j_{1}+r_{i-3}\sum_{j_{2}=0}^{t-1}\sum_{j_{1}=0}^{j_{2}}j_{1}+...+$
$\displaystyle+r_{2}\sum_{j_{i-3}=0}^{t-1}...\sum_{j_{2}=0}^{j_{3}}\sum_{j_{1}=0}^{j_{2}}j_{1}...+\sum_{j_{i-2}=0}^{t-1}...\sum_{j_{2}=0}^{j_{3}}\sum_{j_{1}=0}^{j_{2}}j_{1}\Bigg{)}\,,$
(39)

by setting $X_{2,i}^{(0)}/X_{2,1}^{(0)}=r_{i}w^{i-1}$, where $r_{i}>0$. Then
note that

$\displaystyle\frac{X_{2,d}^{(t+1)}}{X_{2,d-1}^{(t+1)}}$
$\displaystyle=\frac{wX_{2,d-1}^{(t)}+X_{2,d}^{(t)}+1}{wX_{2,d-1}^{(t)}+X_{2,d}^{(t)}}=\frac{X_{2,1}^{(t+1)}}{X_{2,d-1}^{(t+1)}}\frac{X_{2,d-1}^{(t)}}{X_{2,1}^{(t)}}\left(w+\frac{X_{2,d}^{(t)}}{X_{2,d-1}^{(t)}}+\frac{1}{X_{2,d-1}^{(t)}}\right)$
$\displaystyle\vcentcolon=h\left(\frac{X_{2,3}^{(t)}}{X_{2,2}^{(t)}},t,r_{0},...,r_{d-1}\right)\,.$

We will show a useful monotone property of $h(\cdot)$.

######  Lemma 11.

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

$\left(\frac{1}{X_{2,d-1}^{(t)}}\right)^{2}=y^{2}+1+\frac{X_{2,1}^{2}+...+X_{2,d-2}^{2}}{X_{2,d-1}^{2}}\vcentcolon=y^{2}+1+f(r_{2},...,r_{d})\,,$

we get that

$2\frac{\partial\left(\frac{1}{X_{2,d-1}}\right)}{\partial
y}\frac{1}{X_{2,d-1}}=2y\,.$

Hence

$\frac{\partial\left(\frac{1}{X_{2,d-1}}\right)}{\partial
y}=X_{2,d-1}y=X_{2,d}\,,$

which means that

$\frac{\partial h}{\partial
y}=\frac{X_{2,1}^{(t+1)}}{X_{2,d-1}^{(t+1)}}\frac{X_{2,d-1}^{(t)}}{X_{2,1}^{(t)}}(1+X_{2,d})\geq
0\,.$

∎

With this, we will show the following claim proving the invariance of the
subspace of $X_{2,d-1}\leq-1/w$ given some initial conditions:

######  Lemma 12.

Suppose $X_{2,i}^{(0)}<0$ for $i\in[d-2]$, $X_{2,d-1}^{(0)}\leq-1/w$ such that

$\frac{1}{w^{2}}\left(\frac{r_{d-2}}{r_{d-1}}\right)^{2}+\frac{1}{w^{4}}\left(\frac{r_{d-3}}{r_{d-1}}\right)^{2}+...+\frac{1}{w^{2(d-3)}}\left(\frac{r_{2}}{r_{d-1}}\right)^{2}+\frac{1}{w^{2(d-2)}}\left(\frac{1}{r_{d-1}}\right)^{2}\leq\sqrt{w^{2}-1}\,,$
(40)

then $X_{2,d-1}^{(t)}\leq-1/w$ for all $t\geq 0$.

###### Proof.

When $X_{2,d-1}^{(t)}=-1/w$, and given the conditions on $r_{2},...,r_{d-1}$
in (40),

$\displaystyle\frac{X_{2,d}^{(t)}}{X_{2,d-1}^{(t)}}$
$\displaystyle=\pm\sqrt{w^{2}-1-\frac{\left(X_{2,1}^{(t)}\right)^{2}+...+\left(X^{(t)}_{2,d-2}\right)^{2}}{\left(X^{(t)}_{2,d-1}\right)^{2}}}$
$\displaystyle=\pm\sqrt{w^{2}-1-f(t,r_{1},...,r_{d-2})}$

is well defined, where $f(\cdot)$ is monotone non-increasing in $t$ for fixed
$r_{1},...,r_{d-2}>0$. We denote

$\bar{y}^{(t)}=\sqrt{w^{2}-1-f(t,r_{1},...,r_{d-2})}\,.$

Note that if $X_{2,d-1}^{(t)}\leq-1/w$, we have

$\frac{X_{2,d}^{(t)}}{X_{2,d-1}^{(t)}}\in\left[-\bar{y}^{(t)},\bar{y}^{(t)}\right]$

then by Lemma 11 we get that

$h\left(\frac{X_{2,d}^{(t)}}{X_{2,d-1}^{(t)}},t,r_{2},...,r_{d-1}\right)\in\left[-h\left(\bar{y}^{(t)},t,r_{2},...,r_{d}\right),h\left(\bar{y}^{(t)},t,r_{2},...,r_{d}\right)\right]\,$

and since by (39),

$\frac{X_{2,1}^{(t+1)}}{X_{2,d-1}^{(t+1)}}\frac{X_{2,d-1}^{(t)}}{X_{2,1}^{(t)}}\leq
1\,,$

it implies that

$\displaystyle\left|h\left(\pm\bar{y}^{(t)}\right)\right|$
$\displaystyle=\frac{X_{2,1}^{(t+1)}}{X_{2,d-1}^{(t+1)}}\frac{X_{2,d-1}^{(t)}}{X_{2,1}^{(t)}}\bar{y}^{(t)}\leq\bar{y}^{(t+1)}\,,$

which means that

$\frac{X_{2,d-1}^{(t+1)}}{X_{2,d}^{(t+1)}}\in\left[-\bar{y}^{(t+1)},\bar{y}^{(t+1)}\right]$

and thus $X_{2,d-1}^{(t)}\leq-1/w$ for $t\geq 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\to\infty$.

###  F.4 Further discussion

##### Directionality of attention mask $\mathcal{G}$

In the proof of Theorem 3, we would like to note that $\mathcal{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 $\mathcal{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
> $\mathcal{G}$ as the attention mask? If not, what differentiates the case
> between an undirected graph and a directed graph?

##  Appendix 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 [23] and Transformers library [30].

##### Compute

We ran all of our experiments on CPUs.

##### Licenses

  * •

Libraries

    * –

Wikipedia: MIT license

    * –

tranformers [30]: Apache License 2.0

  * •

Pretrained models

    * –

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

    * –

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

    * –

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

    * –

ALBERT (albert/albert-base-v2) [20]: 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, $\mu(X^{(t)})$ no longer converges to $0$, as $\mu(\cdot)$
does not show any decreasing trend even after $128$ layers.

Figure 6: Evolution of $\mu(X^{(t)})$ (in log-log scale) as the number of
layers increases in $128$ layer models.

###  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.

Figure 7: Evolution of $\mu(X^{(t)})$ (in log-log scale) as the number of
layers increases in SANs with different $d_{QK}$ at initialization.

###  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 $32$ sequence samples described in Section 5
and calculate all the singular values of $X^{(t)}$. The results in randomly
initialized and pretrained transformers are presented in Fig. 8.

We see that in both pretrained and initialized transformers, 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 $32$ sequence samples
described in Section 5 and calculate all the pairwise cosine similarities
between tokens. The results in randomly initialized and pretrained
transformers are presented in Fig. 9. We see that in particular in BERT and
ALBERT, the initial pairwise cosine similarities are all non-negative.

Figure 8: The full set of singular values of $X^{(t)}$ at each layer in
initialized and pretrained transformers.

Figure 9: Pairwise cosine similarities between tokens calculated from the
initial embeddings in randomly initialized and pretrained transformers,
respectively.