[name=Theorem,numberwithin=section]thm
[name=Lemma,numberwithin=section,numberlike=thm]lem
[name=Corollary,numberwithin=section,numberlike=thm]cor
[name=Proposition,numberwithin=section,numberlike=thm]prop
[name=Assumption,numberlike=thm]ass

#  Signal Propagation in Transformers: Theoretical Perspectives and the Role
of Rank Collapse

Lorenzo Noci∗1
<EMAIL_ADDRESS>
&Sotiris Anagnostidis∗1
<EMAIL_ADDRESS>
&Luca Biggio∗1,2
<EMAIL_ADDRESS>
&Antonio Orvieto∗1
<EMAIL_ADDRESS>
&Sidak Pal Singh∗1
<EMAIL_ADDRESS>
&Aurelien Lucchi3
<EMAIL_ADDRESS>
(tests)

###### Abstract

Transformers have achieved remarkable success in several domains, ranging from
natural language processing to computer vision. Nevertheless, it has been
recently shown that stacking self-attention layers — the distinctive
architectural component of Transformers — can result in rank collapse of the
tokens’ representations at initialization. The question of if and how rank
collapse affects training is still largely unanswered, and its investigation
is necessary for a more comprehensive understanding of this architecture. In
this work, we shed new light on the causes and the effects of this phenomenon.
First, we show that rank collapse of the tokens’ representations hinders
training by causing the gradients of the queries and keys to vanish at
initialization. Furthermore, we provide a thorough description of the origin
of rank collapse and discuss how to prevent it via an appropriate depth-
dependent scaling of the residual branches. Finally, our analysis unveils that
specific architectural hyperparameters affect the gradients of queries and
values differently, leading to disproportionate gradient norms. This suggests
an explanation for the widespread use of adaptive methods for Transformers’
optimization.

††1Dept of Computer Science, ETH Zürich, 2Robotics & ML, CSEM SA, Alpnach,
Switzerland, 3Department of Mathematics and Computer Science, University of
Basel

###  1 Introduction

Since its first appearance in Vaswani et al. (2017), the Transformer
architecture has revolutionized the field of Natural Language Processing
(NLP), achieving remarkable success in tasks such as text classification (Yang
et al., 2019), machine translation (Conneau and Lample, 2019), reading
comprehension (Brown et al., 2020) and question answering (Raffel et al.,
2019) among others. Recent efforts have effectively extended its applicability
to computer vision (Dosovitskiy et al., 2020) and other domains (Baevski et
al., 2020; Huang et al., 2018; Biggio et al., 2021; Polu et al., 2022),
further popularizing it outside NLP.

The Transformer operates on inputs comprising a sequence of tokens. At its
core, it relies on stacked attention layers, which compute a measure of
relevance for the whole sequence by assigning token-wise importance weights —
obtained by matrix multiplication of the queries and keys, and finally
normalized with the softmax function. The output of an attention layer is then
a linear combination of the importance weights and the so-called values. Then,
the architecture includes fully-connected sub-layers, residual connections (He
et al., 2016), and layer normalization (LN), as illustrated in Fig. 2.

Figure 1: Evolution of the cosine of the angle between tokens for training
POST-LN Transformers of increasing depth, with the Adam optimizer, for the
IWSLT’14 De-En translation task. Unless adequate residual scaling is used at
initialization, increasing depth leads to an increase in the tokens’ alignment
at initialization, which can inhibit training.

In the absence of residual connections, Dong et al. (2021) proved that at
initialization the rank of the sequence representation collapses doubly
exponentially with depth, and both layer normalization and fully connected
layers can only partially alleviate the speed of degeneracy. Under rank
collapse, the model does not distinguish between representations of different
tokens, which are perfectly aligned in feature space at initialization.
Crucially, the precise implications of rank collapse in Transformers are not
fully understood.

In this paper, we show that a high alignment of the tokens’ representations at
initialization — corresponding to rank collapse in the extreme case of perfect
alignment — affects training by causing vanishingly small gradients of the
queries and keys’ parameter matrices. This problem severely diminishes the
capabilities of the model to learn meaningful attention weights and is further
exacerbated in very deep networks, where the rank deficiency — and hence the
vanishing gradient problem of the queries and keys — affects several layers
(see Fig. 1). In order to shed light on this problem, we take inspiration from
the flourishing literature on signal propagation in random networks and start
our analysis by computing the expected gradients of an attention layer with
respect to the queries, keys, and values, which leads to Theorem 3.1 on the
vanishing gradients for the queries and keys. From here, we pursue two
different directions.

Firstly, we investigate under which conditions rank collapse can be avoided by
studying the evolution of the input sequence in a Transformer at
initialization. Our theory reveals that a depth-dependent scaling of the
residual branches, beyond stabilizing the norm of the activations at
initialization, also approximately preserves the cosine of the angle between
tokens, and hence also stabilizes the rank of the propagating sequence. We
show that this holds even in the infinite-depth limit.

Secondly, we illustrate that there are factors, other than the average tokens’
correlation, that affect differently the gradient norm of the queries and keys
compared to the values. In particular, the propagating sequence’s squared norm
has a linear dependence in the values, while a cubic one in the queries and
keys, justifying the use of layer normalization. We also highlight a different
dependence on the embedding dimension and the length of the input sequence,
implying that the gradient norm of a subset of parameters can potentially be
of different orders of magnitude, as empirically hinted by previous works (Liu
et al., 2020). Our analysis brings to light fundamental issues in the signal
propagation in Transformers, opening the way for new, well-founded and
motivated approaches to improve optimization in these models.

###  2 Background

###### Transformers.

A Transformer architecture consists of $L$ stacked attention blocks, as show
in Fig. 2. Layer normalization is usually applied token-wise either after the
residual connections or to the inputs of the self-attention and position-wise
feed-forward sub-layers, leading to the POST-LN (Vaswani et al., 2017) and
PRE-LN (Wang et al., 2019; Xiong et al., 2020) variants respectively.

Formally, given an input sequence ${\bf X}\in\mathbb{R}^{n\times d_{v}}$, with
$n$ tokens of dimension $d_{v}$, the single-head unmasked scaled dot-product
self-attention111Our analysis also easily generalizes to the case of cross-
attention. is defined as:

${\bf S}^{\ell}:={\bf A}^{\ell}{\bf X}^{\ell}{\bf W}^{V},\text{ where }{\bf
A}^{\ell}=\text{softmax}\left(\frac{1}{\sqrt{d_{k}}}{\bf X}^{\ell}{\bf
W}^{Q}\left({\bf X}^{\ell}{\bf W}^{K}\right)^{\top}\right),$ (1)

where the softmax function is applied independently across each row, and the
superscript $\ell$ indexes the $\ell$-th layer. The matrices ${\bf W}^{Q},{\bf
W}^{K}\in\mathbb{R}^{d_{v}\times d_{k}}$ and ${\bf
W}^{V}\in\mathbb{R}^{d_{v}\times d_{v}}$ are learnable parameters, and each
layer is initialized with an independent set of weights. In the literature,
the matrices ${\bf X}^{\ell}{\bf W}^{Q},{\bf X}^{\ell}{\bf W}^{K},{\bf
X}^{\ell}{\bf W}^{V}$ are referred to as queries, keys and values,
respectively. The complete Transformer block, in the absence of layer
normalization, can be written recursively as:

$\displaystyle{\bf Z}^{\ell}=\alpha_{1}{\bf S}^{\ell}+{\bf X}^{\ell}$ (2)
$\displaystyle{\bf Y}^{\ell}=\sigma({\bf Z}^{\ell}{\bf W}^{F_{1}}){\bf
W}^{F_{2}}$ (3) $\displaystyle{\bf X}^{\ell+1}=\alpha_{2}{\bf Y}^{\ell}+{\bf
Z}^{\ell},$ (4)

where the introduced $\alpha_{1},\alpha_{2}$ parameters indicate the strength
of the residual block, ${\bf W}^{F_{1}}$, ${\bf
W}^{F_{2}}\in\mathbb{R}^{d_{v}\times d_{v}}$ 222In practice, one commonly uses
${\bf W}^{F_{1}}\in\mathbb{R}^{d_{v}\times d_{F}}$, ${\bf
W}^{F_{2}}\in\mathbb{R}^{d_{F}\times d_{v}}$ where $d_{F}=\gamma d_{v}$, with
$\gamma\in\\{2,4,8\\}$. Our results then hold up to a constant factor that
depends on $\gamma$. are matrices of learnable parameters; we set ${\bf
X}^{0}:={\bf X}$, and $\sigma:\mathbb{R}\rightarrow\mathbb{R}$ is an
activation function. In our case, $\sigma$ is the ReLU function, but we relax
this assumption to the linear activation from Section 3.2 on.

Figure 2: A single Transformer block.

At initialization, each weight is sampled independently from a distribution
with zero-mean and variance $\sigma^{2}_{v}=\frac{1}{d_{v}}$ for the values
and feedforward weights333One should explicitly write the layer dependence
${\bf W}^{Q,\ell},{\bf W}^{K,\ell},{\bf W}^{V,\ell},{\bf W}^{F_{1},\ell},{\bf
W}^{F_{2},\ell}$. We at times suppress the $\ell$ index to improve
readability. In case $\sigma$ is the ReLU function, we set ${\bf W}^{F,1}$ to
have variance $\frac{2}{d_{v}}$., and $\sigma^{2}_{k}=\frac{1}{d_{k}}$ for the
queries and keys. This is the standard “Xavier” (Glorot and Bengio, 2010b) or
“He” (He et al., 2015) initialization, commonly used in deep learning.

###### Rank Collapse in Transformers.

Interestingly, Dong et al. (2021) proved that when the residual branches are
omitted, the matrix of the tokens’ representations ${\bf X}^{\ell}$ converges
to a rank-1 matrix in which all the representations are the same and equal to
a vector $\bm{x}\in\mathbb{R}^{d_{v}}$, i.e. ${\bf
X}^{\ell}\to\bm{1}_{n}\bm{x}^{\top}$, where $\bm{1}_{d_{v}}$ is the vector
with all ones in $\mathbb{R}^{d_{v}}$. Note that this is a slightly stronger
notion of a rank-$1$ matrix, as it implies that all the tokens’
representations are both perfectly aligned and have the same norm. Indicating
the inner product with the usual bracket notations
$\langle\cdot,\cdot\rangle$, and the cosine of the angle between two tokens as
$\theta_{k,k^{\prime}}$, perfect alignment happens when $\langle{\bf
X}^{\ell}_{k},{\bf X}^{\ell}_{k^{\prime}}\rangle=\left\lVert{\bf
X}^{\ell}_{k}\right\rVert\left\lVert{\bf
X}^{\ell}_{k}\right\rVert\cos\theta_{k,k^{\prime}}$ with
$\cos\theta_{k,k^{\prime}}=1$ for all $k,k^{\prime}\in[n]$. Note that perfect
alignment together with equal norm between all the tokens implies that all the
representations are the same. One of our main contributions is to provide an
explanation of how rank collapse affects the gradients of a Transformer at
initialization.

###### Vanishing Gradient Problem.

Traditionally considered one of the core issues that prevents successful
training, the vanishing gradient problem has a long and rich history that
dates back to before the popularization of deep learning (Hochreiter, 1991;
Bengio et al., 1994). In its essence, given a loss function
$\mathcal{L}:\mathbb{R}^{n\times d_{v}}\to\mathbb{R}$, vanishing gradients
occur when the norm of the gradient of the loss $\mathcal{L}$ with respect to
the parameters of the network ${\bf W}$ — which we indicate as
$\left\lVert\frac{\partial\mathcal{L}}{\partial{\bf W}}\right\rVert$ — is too
small to provide enough backpropagating signal, thus hindering gradient-based
optimization methods. Despite extensive research toward understanding and
overcoming the problem in disparate contexts (Glorot and Bengio, 2010a; He et
al., 2015; Hanin, 2018; Zhang et al., 2019), a formal explanation of its role
in relatively new architectures such as Transformers is largely missing in the
literature, with a few exceptions (Xiong et al., 2020; Wang et al., 2022;
Huang et al., 2020). In our paper (Section 3.1), we show how vanishing
gradient occurs in conjunction with the rank collapse issue identified by Dong
et al. (2021).

###### Signal Propagation in Random Networks at Initialization.

After addressing the question on the effects of rank collapse, we take a step
back and rigorously analyze its causes by looking at how the properties of the
input sequence ${\bf X}$ are lost/preserved as it propagates through a
randomly initialized Transformer. More specifically, we focus on two aspects
of the propagating sequence: the expected Frobenius norm
$\mathbb{E}\left\lVert{\bf X}^{\ell}\right\rVert^{2}$ and the expected inner
product between different tokens $\mathbb{E}\langle{\bf X}_{k},{\bf
X}_{k}^{\prime}\rangle$, with $k\neq k^{\prime}$. The former is linked to a
number of studies on the initialization of neural networks at the _edge of
chaos_ (Poole et al., 2016; Schoenholz et al., 2016), and vanishing/exploding
gradients (Hanin, 2018). The latter quantity describes how the geometry of the
feature space changes after applying a Transformer block, and is related to
the concept of _dynamical isometry_ (Saxe et al., 2013). To understand the
evolution of the inner product, we analyze the following measure of
correlation (Nachum et al., 2021; Cho and Saul, 2009):

$\rho^{\ell}_{kk^{\prime}}:=\frac{\mathbb{E}\langle{\bf X}^{\ell}_{k},{\bf
X}^{\ell}_{k^{\prime}}\rangle}{\sqrt{\mathbb{E}\left\lVert{\bf
X}_{k}^{\ell}\right\rVert^{2}\mathbb{E}\left\lVert{\bf
X}_{k^{\prime}}^{\ell}\right\rVert^{2}}}.$ (5)

Note that $\rho^{\ell}_{kk^{\prime}}=1$ if and only if the $k$-th and
$k^{\prime}$-th tokens are perfectly aligned ($\cos\theta_{kk^{\prime}}=1$).
We stress that in our case — differently from the aforementioned works —
instead of analyzing the relationship between two different data points, we
study the relationship between tokens of the same sequence.

###  3 Theoretical Results

####  3.1 Vanishing Gradients for Queries and Keys under Rank Collapse

To investigate the problem of vanishing gradients in the attention layers, we
make use of the framework of matrix calculus (Magnus and Neudecker, 2019;
Singh et al., 2021). In particular, we compare the expected Frobenius norm of
the gradient of a self-attention layer with respect to its parameters:
$\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}}\right\|^{2}_{F}$, where here ${\bf W}$ indicates one of the keys, queries
or values weight matrices. Due to the well-known difficulty of computing
expectations of the softmax (Daunizeau, 2017; Shekhovtsov and Flach, 2018),
throughout this manuscript, we make the simplifying assumption that the
softmax output is the uniform distribution at initialization, i.e. the
$n\times n$ matrix containing $\frac{1}{n}$ in each entry. [Uniform attention]
We assume that ${\bf A}^{\ell}=\frac{1}{n}\bm{1}_{n\times n}$, where
$\bm{1}_{n\times n}$ is the matrix with all entries equal to $1$. Crucially,
in Appendix A.5, we formally show that _this assumption holds almost surely_
in the limit $d_{k}\to\infty$. There, we also experimentally show that even in
the more realistic case where $d_{k}=d_{v}\approx 512$, the empirical
simulations provide a surprisingly faithful approximation of the theoretical
insights presented in this paper.

We define the mean token $\bar{\bm{x}}^{\ell}$ through its components
$\bar{\bm{x}}^{\ell}_{i}=\frac{1}{n}\sum_{k=1}^{n}{\bf X}^{\ell}_{ki}$,
$i\in[d_{v}]$. In the following theorem, we compute the expected gradients of
an attention layer at initialization, and set the basis for our following
analysis. We provide the results only for the queries, as the case for the
keys is analogous.

###### Lemma 3.1.

Let ${\bf X}^{\ell}$ be the representations of the input sequence at the
$\ell$-th layer. Under the uniform-attention assumption, we have
$\displaystyle\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}^{V,\ell}}\right\|^{2}_{F}$
$\displaystyle=d_{v}n\mathbb{E}\|\bar{\bm{x}}^{\ell}\|^{2}~{};$ (6)
$\displaystyle\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\|^{2}_{F}$
$\displaystyle=\frac{\sigma^{2}_{v}\sigma^{2}_{k}d_{v}}{n^{2}}\cdot\mathbb{E}\left[\|{\bf
X}^{\ell}\|^{2}_{F}\cdot\|({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}-n\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\|^{2}_{F}\right]~{};$
(7) $\displaystyle\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
X}^{\ell}}\right\|^{2}_{F}$
$\displaystyle\leq\frac{8\sigma^{2}_{q}\sigma^{2}_{k}\sigma^{2}_{v}d_{k}d_{v}}{n}\,\cdot\mathbb{E}\left\lVert({\bf
X}^{\ell})^{\top}{\bf
X}^{\ell}-n\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\right\rVert^{2}_{F}+2d_{v}^{2}\sigma^{2}_{v}\;.$
(8)

We defer the precise study of the scaling of these quantities as a function of
$n$ and $d_{v},d_{k}$, to Section 3.3. At this stage, it is crucial to note
that $\frac{1}{n}({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}-\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}$ is the centered
empirical covariance matrix of the tokens’ representations. It is easy to see
that if ${\bf X}^{\ell}$ is a rank-$1$ matrix, then all the rows of ${\bf
X}^{\ell}$ are proportional to a fixed $d_{v}$-dimensional vector, and the
empirical covariance matrix has all zero entries. Introducing a loss function
$\mathcal{L}:\mathbb{R}^{n\times d_{v}}\to\mathbb{R}$, we make the statement
on vanishing gradients more formal in the following theorem:

[Vanishing gradients under rank collapse] Suppose that the uniform-attention
assumption holds. If additionally ${\bf X}^{\ell}$ for any $l\in[L]$ has
rank-1, and there exists a vector $\bm{x}\in\mathbb{R}^{d}$ such that ${\bf
X}^{\ell}=\bm{1}_{n}\bm{x}^{T}$, then:
$\mathbb{E}\left\lVert\frac{\partial\mathcal{L}}{\partial{\bf
W}^{Q,\ell}}\right\rVert_{F}^{2}=0,\;\;\;\;\mathbb{E}\left\lVert\frac{\partial\mathcal{L}}{\partial{\bf
W}^{K,\ell}}\right\rVert_{F}^{2}=0,$ (9) where the expectation is taken over
the weight matrices. This implies that these quantities are vanishing almost
surely, due to the non-negativeness of the norm.

The proof simply relies on expanding the norm of the gradient of the loss with
the aid of the chain rule and then bounding it by the product of the norms of
each term of the chain. The final result holds with an application of Lemma
3.1, in which the rank-$1$ assumption makes
$\mathbb{E}\left\lVert\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\rVert$ vanish. In light of Theorem 3.1, we can conclude
that the rank collapse issue originally identified in Dong et al. (2021)
corresponds to an initialization in a region of vanishing gradient signal in
the subspace of parameters identified by the queries and keys. How can this
affect training? One may argue that if rank collapse does not happen in the
very first layer, then the corresponding gradients are non-zero, and the rank
of the subsequent layers — affected by rank collapse — can be increased with
the first few steps of gradient descent. In practice, we show empirically in
Fig. 1 that escaping this pathological landscape is harder in deeper nets.

####  3.2 Forward Signal Propagation and the Importance of Scaling the
Residual Branches

We now turn our attention to the study of the influence of skip connections in
transformers. Dong et al. (2021) showed that simply adding skip connections
prevents rank collapse. Somewhat surprisingly, we show that while the claim
holds for any finite depth, the average angle between different tokens quickly
increases with just a few layers, and as $L\to\infty$ a Transformer can still
lose rank unless the residual branches are adequately initialized. As Dong et
al. (2021) showed that layer normalization does not avoid rank collapse, we
omit it in our analysis. Firstly, we introduce two lemmas on the propagation
of inner products (Lemma 3.2) and the norm (Lemma 3.3) of the tokens’
representations.

###### Lemma 3.2 (Propagation of inner products).

Let $C({\bf X}^{\ell})=\sum_{k,k^{\prime}}\langle{\bf X}_{k}^{\ell},{\bf
X}_{k^{\prime}}^{\ell}\rangle$ and ${\bf X}$ the input sequence. Under the
Assumption 3.1 and if $\sigma$ is the linear activation function, we have
that: $\mathbb{E}\left[C({\bf
X}^{L})\right]=(\alpha_{2}^{2}+1)^{L}(\alpha_{1}^{2}+1)^{L}C({\bf X}).$ (10)
hence, under the depth scaling for the residual block parameters
$\alpha_{1}^{2}=\frac{\tilde{\alpha}_{1}}{L},\alpha_{2}^{2}=\frac{\tilde{\alpha}_{2}}{L}$
with $\tilde{\alpha}_{1},\tilde{\alpha}_{2}\in\mathbb{R}$ independent of $L$,
we have that: $\lim_{L\to\infty}\mathbb{E}[C({\bf
X}^{L})]=\text{e}^{\tilde{\alpha}_{1}+\tilde{\alpha}_{2}}C({\bf X}).$ (11)

Note that $C({\bf
X}^{\ell})=n^{2}\left\lVert\bar{\bm{x}}^{\ell}\right\rVert^{2}$. The lemma on
the propagation of the norm is slightly more involved:

###### Lemma 3.3 (Propagation of the norm).

Let ${\bf X}^{L}$ be the representations of the input sequence at the final
layer. Under the assumptions of Lemma 3.2, we have that:
$\mathbb{E}\left\lVert{\bf
X}^{L}\right\rVert_{F}^{2}=n(\alpha_{2}^{2}+1)^{L}\alpha_{1}^{2}\sum_{k=0}^{L-1}(\alpha_{1}^{2}+1)^{k}\left\lVert\bar{\bm{x}}\right\rVert^{2}+(\alpha_{2}^{2}+1)^{L}||{\bf
X}||_{F}^{2},$ (12) hence, under the depth scaling for the residual block
parameters
$\alpha_{1}^{2}=\frac{\tilde{\alpha}_{1}}{L},\alpha_{2}^{2}=\frac{\tilde{\alpha}_{2}}{L}$
with $\tilde{\alpha_{1}},\tilde{\alpha_{2}}\in\mathbb{R}$ independent of $L$,
we have that: $\lim_{L\to\infty}\mathbb{E}\left\lVert{\bf
X}^{L}\right\rVert_{F}^{2}=n\text{e}^{\tilde{\alpha}_{2}}(\text{e}^{\tilde{\alpha}_{1}}-1)\left\lVert\bar{\bm{x}}\right\rVert^{2}+\text{e}^{\tilde{\alpha}_{2}}||{\bf
X}||_{F}^{2}.$ (13)

The proof of Lemma 3.3 consists in expanding $\mathbb{E}\left\lVert{\bf
X}^{L}\right\rVert_{F}^{2}$ according to the defining equations for the
Transformer, and simplifying the expression by using iterated expectations
$\mathbb{E}\left\lVert{\bf
X}^{L}\right\rVert_{F}^{2}=\mathbb{E}[\mathbb{E}[\left\lVert{\bf
X}^{L}\right\rVert_{F}^{2}|{\bf X}^{\ell}]]$ to exploit the conditional
independence between different layers, and then computing the expectations
using the independence assumption on the weights. The expression on the right-
hand side will then depend on ${\bf X}^{\ell}$ only through its norm
$\left\lVert{\bf X}^{\ell}\right\rVert$ and the norm of the mean token
$\left\lVert\bar{\bm{x}}^{\ell}\right\rVert^{2}$. Using Lemma 3.2 then allows
us to unroll the recursion and get the final result. The complete proof,
together with the proof of Lemma 3.2, can be found in Appendix A.3.

The previous Lemma provides theoretical justification that scaling the
residual branches by setting the alpha parameters to be
$\mathcal{O}(1/\sqrt{L})$ allows both the norm of the propagating input and
the inner products between different tokens to be approximately preserved.
Hence, the information contained in the input is not lost, even in the
infinite depth limit.

Figure 3: Effect of the residual scaling to the norm of the gradients of the
network at initialization with respect to some loss. From left to right: (a)
the cosine of the angle between tokens increases with depth. Note how larger
values of $\alpha_{1},\alpha_{2}$ imply a faster token alignment with depth
(Theorem 3.2). Subplots (b) and (c) show the gradients of the queries-keys and
values parameters respectively by increasing depth, compared to the
corresponding norms of the first layer. Gradients for the queries-keys
diminish with depth, while the opposite happens for the values. We use POST-LN
to disentangle the effect of the variance of the input.

###### Residual Scaling Preserves Correlations.

We now prove that without the depth-dependent residual scaling (i.e. with
$\alpha_{1}=\alpha_{2}=1$) the correlation between the tokens quickly
increases, and reaches perfect alignment in the infinite depth limit. More
specifically, our argument shows that in this limit, the correlation between
different tokens $\rho^{\ell}_{k,k^{\prime}}$ as in Eq. (5) converges to 1,
implying rank collapse. Furthermore, we show how setting the residual
parameters $\alpha_{1}$ and $\alpha_{2}$ as dictated by Theorem 3.3, ensures
that the correlation measure is dependent on the input in a non-trivial way
even at infinite depth. To this end, we introduce the average correlation at
layer $\ell$:

$\rho^{\ell}=\frac{1}{n(n-1)}\sum_{k\neq
k^{\prime}}\rho^{\ell}_{kk^{\prime}}.$ (14)

Note that $\rho^{\ell}=1$ if and only if every pair of tokens is perfectly
aligned.

We are now ready to formalize the influence of the $1/\sqrt{L}$-scaling on the
correlation between tokens’ representations by stating Theorem 3.2.

[] Let the input tokens have the same norm, i.e. $\left\lVert{\bf
X}_{k}\right\rVert=\left\lVert\bm{x}\right\rVert\;\forall k\in[n]$ for some
$\bm{x}\in\mathbb{R}^{d_{v}}$. Under the depth scaling for the residual block
parameters
$\alpha_{1}^{2}=\frac{\tilde{\alpha}_{1}}{L},\alpha_{2}^{2}=\frac{\tilde{\alpha}_{2}}{L}$
with $\tilde{\alpha}_{1},\tilde{\alpha}_{2}\in\mathbb{R}$ independent of $L$,
we have that:
$\lim_{L\to\infty}\rho^{\ell}=\frac{n\text{e}^{\tilde{\alpha}_{1}}C({\bf
X})}{(n-1)[(\text{e}^{\tilde{\alpha}_{1}}-1)C({\bf X})+n\left\lVert{\bf
X}\right\rVert_{F}^{2}]}-\frac{1}{n-1}.$ (15) On the other hand, if
$\alpha_{1},\alpha_{2}\neq 0$ are some constants independent of $L$, we have
that: $\lim_{L\to\infty}\rho^{\ell}=1.$ (16)

The proof consists in noting that due to the symmetry of the problem at
initialization, for a fixed layer the expected norm of each token is the same.
Hence, by our definition of $\rho^{\ell}_{kk^{\prime}}$, we can write
$\mathbb{E}\langle{\bf X}^{\ell}_{k},{\bf
X}^{\ell}_{k^{\prime}}\rangle=\rho^{\ell}_{kk^{\prime}}\mathbb{E}\left\lVert\bm{x}^{\ell}\right\rVert^{2}$.
By summing over the $k,k^{\prime}$ indexes, the resulting equation will depend
on $\mathbb{E}[C({\bf X}^{\ell})]$ and $\mathbb{E}\left\lVert{\bf
X}^{\ell}\right\rVert^{2}$, which can be expanded using Lemma 3.2 and 3.3
respectively. The result is then given by solving for $\rho^{\ell}$.

Figure 4: Evolution of Correlation in Transformers with (dashed lines) and
without (solid lines) $1/\sqrt{L}$-scaling for PRE-LN, POST-LN and without
layer normalization (No-LN).

Note that under the $1/\sqrt{L}$-scaling, the correlation term is one if and
only if $C({\bf X})=n\left\lVert{\bf X}\right\rVert^{2}$, which holds in the
degenerate case where all the input tokens are perfectly aligned. In Appendix
A.4, we give precise formulas for the expected correlations at any depth,
showing that $\rho^{\ell}$ reaches values close to one even for relatively
shallow networks when the $1/\sqrt{L}$-scaling is not adopted (see also Fig. 3
(left)). Additionally, in Fig. 4, we empirically show that in the presence of
the $1/\sqrt{L}$-scaling, layer normalization (either PRE or POST) does not
significantly affect the evolution of the correlations. On the other hand,
without the residual scaling, PRE-LN seems to alleviate the rate of increase
of $\rho^{\ell}_{kk^{\prime}}$. It is intriguing that most deep Transformer
models use this configuration (Brown et al., 2020). We provide more extensive
empirical results in Appendix B.

Note that the $1/\sqrt{L}$ scaling for the residual branches has been
previously studied in the context of stabilization of residual networks (see
Section 4), here we extend these results to Transformers and provide new
insights on its role in the context of rank preservation. Finally, note that
by setting $\tilde{\alpha}_{1},\tilde{\alpha}_{2}=0$, we recover the so called
"ReZero" initialization (Bachlechner et al., 2021). In this context, the
$1/\sqrt{L}$ scaling extends this framework as it allows for wider range of
values for $\tilde{\alpha}_{1},\tilde{\alpha}_{2}$ while still guaranteeing
stability. We mention here that extending these results from the linear
activation to the ReLU case is known to be a hard problem, due to the
technical difficulty of propagating the inner products across ReLU layers that
are shared among the tokens (this is the case in the position-wise feed-
forward layers in Transformers). Exact formulas can be found only in the case
of one ReLU layer with Gaussian inputs in Cho and Saul (2009).

####  3.3 Dependence on the Angle between Tokens and the Input Norm

In this section, we drop the superscript $\ell$ as it is obvious from context
and assume for simplicity that $d_{k}=d_{v}$. To gain a better intuition on
the factors that affect the gradients and provide additional insights, we
study the case in which every pair of distinct tokens are zero-mean Gaussian
random variables, correlated in the same way, i.e
$\rho^{\ell}_{ii^{\prime}}=\rho$ for $i\neq i^{\prime}$ or more precisely

$\mathbb{E}\left[{\bf X}_{i,j}{\bf
X}_{i^{\prime},j^{\prime}}\right]=\begin{cases}0&j\neq j^{\prime}\ \
\text{(independent dimensions)}\\\ \sigma^{2}_{x}&i=i^{\prime},j=j^{\prime}\\\
\rho\sigma^{2}_{x}&i\neq i^{\prime},j=j^{\prime}\end{cases}.$

To see that this equation satisfies our definition of the correlation metric,
note that $\mathbb{E}[\left\lVert{\bf X}_{i}\right\rVert^{2}]=d\sigma_{x}^{2}$
and $\mathbb{E}\langle{\bf X}_{i},{\bf
X}_{i^{\prime}}\rangle=d\sigma_{x}^{2}\rho$, for $i\neq i^{\prime}$. Then, the
expected norm of the gradients for the values (Eq. (6)) simplifies to

$\mathbb{E}\left\lVert\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}\right\rVert_{F}^{2}=\sigma_{x}^{2}d^{2}\left(1+\rho(n-1)\right).$
(17)

By making the additional assumption that the norm and the correlation
propagate independently, the respective norm for the queries — and
symmetrically the keys — (Eq. (7)) reduces to:

$\mathbb{E}\left\|\frac{\partial{\bf
S}}{\partial\mathbf{W}^{Q}}\right\|_{F}^{2}=\sigma_{x}^{6}\frac{(n-1)}{n}(1-\rho)^{2}d(n+d).$
(18)

In Appendix A.2 we provide a rigorous proof, that relies on Isserlis theorem
(Isserlis, 1918) to compute higher-order moments. The above expressions reveal
the different dependencies on four main actors, that we inspect separately
here. The gradients of the queries depend via a cubic function on the
_variance of the input, $\sigma_{x}^{2}$_, compared to a linear for the
values. This provides an additional interpretation of the successful use of
layer normalization, as in Xiong et al. (2020), either in the POST-LN or PRE-
LN format, that standardizes the input variance $\sigma_{x}^{2}$ to the value
$1$.

Next, we emphasize the dependence on the _correlation between the tokens_ ,
also illustrated in Fig. 3. Importantly, note how the queries/keys have
opposite monotonic functional dependence with respect to $\rho$ compared to
the values. As revealed by Theorem 3.2 and Fig. 3 (center), inappropriate
scaling of the residual branches can already lead to this phenomenon even in a
relatively shallow network.

Finally, Eq. (17) and (18) reveal a different scaling in terms of the
_embedding size $d$_ and the _sequence length $n$_ due to the self-attention
operation itself. We hope that the identification of the different
dependencies in the gradients of the parameters will inspire a new line of
works aimed at solving some of the difficulties in training Transformers.

####  3.4 Are Adaptive Methods really needed for training Transformers?

Figure 5: Adaptive learning rates computed by Adam in Transformers.

The existence of the discrepancy in the magnitude of the gradients with
respect to the weights ${\bf W}^{Q},{\bf W}^{K}$ and ${\bf W}^{V}$, might
explain the success of adaptive optimization algorithms, as illustrated in
Fig. 5, where we plot the effective learning rate computed by Adam (Kingma and
Ba, 2014) in a toy encoder task (more details in Appendix C). Hence, we embark
on a preliminary exploration to train a Transformer architecture with SGD with
the intent of matching Adam’s performance. Based on our theory, we propose a
simple architectural modification, an inverse temperature scaling
$\tau\in\mathbb{R}$ inside the softmax:

${\bf S}^{\ell}:=\text{softmax}\left(\frac{\tau}{\sqrt{d_{k}}}{\bf
X}^{\ell}{\bf W}^{Q}\left({\bf X}^{\ell}{\bf W}^{K}\right)^{\top}\right){\bf
X}^{\ell}{\bf W}^{V}.$ (19)

A direct consequence of our analysis is that $\tau$ allows controlling the
magnitude of the gradients for the queries and keys’ parameters.

We evaluate our proposal, consisting of residual scaling and the
aforementioned inverse temperature parameters, on the widely used IWSLT14
German-to-English (De-En) benchmark translation task. All details regarding
the experimental setup and the choice of inverse temperature used are provided
in Appendix C. We train a Transformer encoder-decoder of varying depth with
SGD, after removing all normalization layers and adequately initializing the
residual connections. For our training with SGD, we avoid using any learning
rate warm-up, as commonly done for Adam, and instead use a step-scheduler to
decrease the learning rate at 40% and 80% of training. We compare against the
following methods that make use of Adam; POST-LN and PRE-LN refer to the
aforementioned alternatives to apply layer normalization. We also compare
against other successful techniques that rely on specific initializations to
avoid layer normalization, such as ReZero (Bachlechner et al., 2021) and
T-Fixup (Zhang et al., 2019). We report the average BLEU score (Papineni et
al., 2002) across 5 runs in Fig. 7 and Table 7.

Figure 6: BLEU scores by increasing the number of transformers blocks. ‘X’
Transformer blocks implies in total ‘X’ encoder self-attention, ‘X’ decoder
self-attention and ‘X’ decoder cross-attention layers.

Method (6L-Encoder / 6L-Decoder) |  BLEU $\uparrow$
---|---
SGD POST-LN | 31.36
SGD res-scale | 32.79
SGD temperature | 35.69
Adam POST-LN (Vaswani et al., 2017) | 35.39
Adam PRE-LN (Vaswani et al., 2017) | 35.10
ReZero (Bachlechner et al., 2021) | 34.55
T-Fixup Zhang et al. (2019) | 35.59

Figure 7: BLEU scores for the IWSLT14 German-to-English translation task. SGD
res-scale refers to the training of SGD without layer normalization and
initialization of the residual scaling $a_{1}=a_{2}=\frac{1}{\sqrt{L}}$. SGD
temperature additionally employs an inverse temperature inside the softmax.

Our proposed method considerably improves training with SGD, keeping up and in
some cases surpassing any results achieved by the Adam optimizer. We are also
able to train deeper networks without the use of layer normalization. We leave
for future work to further investigate modifications or alternatives to the
self-attention operation.

###  4 Related Work

Our work builds upon the rich literature on forward and backward signal
propagation in random neural networks (Poole et al., 2016; Schoenholz et al.,
2017; Xiao et al., 2018; Pennington et al., 2017; Orvieto et al., 2021; Noci
et al., 2021; Zavatone-Veth and Pehlevan, 2021). The $1/\sqrt{L}$ scaling
scheme has been investigated in the literature for the stabilization of
residual networks (Hanin and Rolnick, 2018; Arpit et al., 2019; Allen-Zhu et
al., 2019; Hayou et al., 2021).

Our work draws inspiration from a series of recent works studying the rank of
the representations of random feed-forward neural networks at initialization
(Daneshmand et al., 2020b, a). In the context of Transformers, Dong et al.
(2021) has recently identified the rank collapse issue object of study of the
present work. Thanks to our analysis of the backward pass, we are able to
demonstrate that rank collapse in Transformer architectures leads to
vanishingly small gradients of queries and keys, thereby preventing effective
training and allowing us to complete the analysis of (Dong et al., 2021).

Among the architectural components in Transformers, layer normalization is,
arguably, one of the most important – and debated – ones (Chen et al., 2018;
Wang et al., 2019; Nguyen et al., 2010; Xiong et al., 2020). In the original
architecture (Vaswani et al., 2017), layer normalization is used to stabilize
the forward pass by reducing the variance of the inputs to the following
sublayer. Our analysis of the forward pass shows that its inclusion is not
strictly necessary for the purpose of controlling the norm of the
representations. For a theoretical analysis of signal propagation in the
presence of layer norm, we refer the reader to Xiong et al. (2020).

Additionally, our theoretical study of the backward pass provides a rigorous
explanation of the empirically observed discrepancy between the magnitude of
the gradients of the queries and the values, which Liu et al. (2020)
hypothesize to be one of the causes of the success of adaptive methods in
training Transformers (Liu et al., 2019; Zhang et al., 2020; Huang et al.,
2020).

Finally, properly rescaled residual connections have been found to be
beneficial for training Transformers by a number of recent research works
(Zhang et al., 2019; Bachlechner et al., 2021; Wang et al., 2022). However,
none of these studies characterize the impact of skip connections on rank
propagation, while our analysis suggests a theoretically-grounded way to
stabilize it.

###  5 Conclusions and Future Work

In this paper, we showed how, at initialization, rank collapse and more
generally high correlation in the tokens, causes vanishing gradients of the
queries and keys of a Transformer architecture. While residual connections
help mitigate rank collapse at finite depth, we showed that they alone cannot
prevent high alignments of the tokens’ representations — unless properly
scaled by a $1/\sqrt{L}$-factor. Finally, we have also discovered counter-
intuitive dependencies on the variance of the input, embedding size, and
sequence length, potentially causing large differences between the gradients
of queries/keys compared to the values’ parameters. Hence, we conclude that
one of the strengths of Transformers lies in their carefully designed
architecture together with an adequate initialization. Finally, we gave
preliminary evidence that one of the factors contributing to the higher
efficacy of Adam compared to SGD in training Transformers arises from the
disproportionate magnitude of gradients as postulated by our theory.
Nonetheless, other factors might further accentuate the difference between
these two algorithms during training, leaving the door open for further
research regarding the benefits of adaptive optimization methods with
Transformers.

### Acknowledgements

We thank our colleague Jonas Kohler who provided insights and expertise that
greatly assisted the research.

### References

  * Allen-Zhu et al. [2019] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In _International Conference on Machine Learning_ , pages 242–252. PMLR, 2019.
  * Arpit et al. [2019] Devansh Arpit, Víctor Campos, and Yoshua Bengio.  How to initialize your network? robust initialization for weightnorm & resnets.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Bachlechner et al. [2021] Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley.  Rezero is all you need: Fast convergence at large depth.  In _Uncertainty in Artificial Intelligence_ , pages 1352–1361. PMLR, 2021.
  * Baevski et al. [2020] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli.  wav2vec 2.0: A framework for self-supervised learning of speech representations.  _Advances in Neural Information Processing Systems_ , 33:12449–12460, 2020.
  * Bengio et al. [1994] Yoshua Bengio, Patrice Simard, and Paolo Frasconi.  Learning long-term dependencies with gradient descent is difficult.  _IEEE transactions on neural networks_ , 5(2):157–166, 1994.
  * Biggio et al. [2021] Luca Biggio, Tommaso Bendinelli, Alexander Neitz, Aurelien Lucchi, and Giambattista Parascandolo.  Neural symbolic regression that scales.  In _International Conference on Machine Learning_ , pages 936–945. PMLR, 2021.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Chen et al. [2018] Mia Xu Chen, Orhan Firat, Ankur Bapna, Melvin Johnson, Wolfgang Macherey, George Foster, Llion Jones, Niki Parmar, Mike Schuster, Zhifeng Chen, et al.  The best of both worlds: Combining recent advances in neural machine translation.  _arXiv preprint arXiv:1804.09849_ , 2018.
  * Cho and Saul [2009] Youngmin Cho and Lawrence Saul.  Kernel methods for deep learning.  _Advances in neural information processing systems_ , 22, 2009.
  * Conneau and Lample [2019] Alexis Conneau and Guillaume Lample.  _Cross-Lingual Language Model Pretraining_.  Curran Associates Inc., Red Hook, NY, USA, 2019.
  * Daneshmand et al. [2020a] Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, and Aurelien Lucchi.  Batch normalization provably avoids rank collapse for randomly initialised deep networks.  _arXiv preprint arXiv:2003.01652_ , 2020a.
  * Daneshmand et al. [2020b] Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, and Aurelien Lucchi.  Batch normalization provably avoids ranks collapse for randomly initialised deep networks.  _Advances in Neural Information Processing Systems_ , 33:18387–18398, 2020b.
  * Daunizeau [2017] Jean Daunizeau.  Semi-analytical approximations to statistical moments of sigmoid and softmax mappings of normal variables.  _arXiv preprint arXiv:1703.00091_ , 2017.
  * 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 _International Conference on Machine Learning_ , pages 2793–2803. PMLR, 2021.
  * Dosovitskiy et al. [2020] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Glorot and Bengio [2010a] Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pages 249–256, 2010a.
  * Glorot and Bengio [2010b] Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pages 249–256. JMLR Workshop and Conference Proceedings, 2010b.
  * Hanin [2018] Boris Hanin.  Which neural net architectures give rise to exploding and vanishing gradients?  _Advances in neural information processing systems_ , 31, 2018.
  * Hanin and Rolnick [2018] Boris Hanin and David Rolnick.  How to start training: The effect of initialization and architecture.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Hayou et al. [2021] Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau.  Stable resnet.  In _International Conference on Artificial Intelligence and Statistics_ , pages 1324–1332. PMLR, 2021.
  * He et al. [2015] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  In _Proceedings of the IEEE international conference on computer vision_ , pages 1026–1034, 2015.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _In IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 770–778, 2016.
  * Hochreiter [1991] Sepp Hochreiter.  Untersuchungen zu dynamischen neuronalen netzen.  _Diploma, Technische Universität München_ , 91(1), 1991.
  * Huang et al. [2018] Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Ian Simon, Curtis Hawthorne, Noam Shazeer, Andrew M Dai, Matthew D Hoffman, Monica Dinculescu, and Douglas Eck.  Music transformer: Generating music with long-term structure.  In _International Conference on Learning Representations_ , 2018.
  * Huang et al. [2020] Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs.  Improving transformer optimization through better initialization.  In _International Conference on Machine Learning_ , pages 4475–4483. PMLR, 2020.
  * Isserlis [1918] Leon Isserlis.  On a formula for the product-moment coefficient of any order of a normal frequency distribution in any number of variables.  _Biometrika_ , 12(1/2):134–139, 1918.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Liu et al. [2019] Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han.  On the variance of the adaptive learning rate and beyond, 2019.  URL https://arxiv.org/abs/1908.03265.
  * Liu et al. [2020] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han.  Understanding the difficulty of training transformers.  _arXiv preprint arXiv:2004.08249_ , 2020.
  * Loshchilov and Hutter [2017] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Magnus and Neudecker [2019] Jan R Magnus and Heinz Neudecker.  _Matrix differential calculus with applications in statistics and econometrics_.  John Wiley & Sons, 2019.
  * Nachum et al. [2021] Ido Nachum, Jan Hązła, Michael Gastpar, and Anatoly Khina.  A johnson–lindenstrauss framework for randomly initialized cnns.  _arXiv preprint arXiv:2111.02155_ , 2021.
  * Nguyen et al. [2010] XuanLong Nguyen, Martin J Wainwright, and Michael I Jordan.  Estimating divergence functionals and the likelihood ratio by convex risk minimization.  _IEEE Transactions on Information Theory_ , 56(11):5847–5861, 2010.
  * Noci et al. [2021] Lorenzo Noci, Gregor Bachmann, Kevin Roth, Sebastian Nowozin, and Thomas Hofmann.  Precise characterization of the prior predictive distribution of deep relu networks.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Orvieto et al. [2021] Antonio Orvieto, Jonas Kohler, Dario Pavllo, Thomas Hofmann, and Aurelien Lucchi.  Vanishing curvature and the power of adaptive methods in randomly initialized deep networks.  _AISTATS 2022 (to appear)_ , 2021.
  * Ott et al. [2019] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli.  fairseq: A fast, extensible toolkit for sequence modeling.  In _Proceedings of NAACL-HLT 2019: Demonstrations_ , 2019.
  * Papineni et al. [2002] Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu.  Bleu: a method for automatic evaluation of machine translation.  In _Proceedings of the 40th annual meeting of the Association for Computational Linguistics_ , pages 311–318, 2002.
  * Pennington et al. [2017] Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli.  Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper/2017/file/d9fc0cdb67638d50f411432d0d41d0ba-Paper.pdf.
  * Polu et al. [2022] Stanislas Polu, Jesse Michael Han, Kunhao Zheng, Mantas Baksys, Igor Babuschkin, and Ilya Sutskever.  Formal mathematics statement curriculum learning.  _arXiv preprint arXiv:2202.01344_ , 2022.
  * Poole et al. [2016] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli.  Exponential expressivity in deep neural networks through transient chaos.  _Advances in neural information processing systems_ , 29, 2016.
  * Raffel et al. [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.  _arXiv preprint arXiv:1910.10683_ , 2019.
  * Saxe et al. [2013] Andrew M Saxe, James L McClelland, and Surya Ganguli.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  _arXiv preprint arXiv:1312.6120_ , 2013.
  * Schoenholz et al. [2016] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein.  Deep information propagation.  _arXiv preprint arXiv:1611.01232_ , 2016.
  * Schoenholz et al. [2017] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein.  Deep information propagation.  In _ICLR_ , 2017.
  * Shekhovtsov and Flach [2018] Alexander Shekhovtsov and Boris Flach.  Feed-forward propagation in probabilistic neural networks with categorical and max layers.  In _International conference on learning representations_ , 2018.
  * Singh et al. [2021] Sidak Pal Singh, Gregor Bachmann, and Thomas Hofmann.  Analytic insights into structure and rank of neural network hessian maps.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=otDgw7LM7Nn.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _arXiv preprint arXiv:1706.03762_ , 2017.
  * Wang et al. [2022] Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei.  Deepnet: Scaling transformers to 1,000 layers.  _arXiv preprint arXiv:2203.00555_ , 2022.
  * Wang et al. [2019] Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao.  Learning deep transformer models for machine translation.  _arXiv preprint arXiv:1906.01787_ , 2019.
  * Xiao et al. [2018] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington.  Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks.  In Jennifer Dy and Andreas Krause, editors, _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pages 5393–5402. PMLR, 10–15 Jul 2018.  URL https://proceedings.mlr.press/v80/xiao18a.html.
  * Xiong et al. [2020] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu.  On layer normalization in the transformer architecture.  In _International Conference on Machine Learning_ , pages 10524–10533. PMLR, 2020.
  * Yang et al. [2019] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le.  Xlnet: Generalized autoregressive pretraining for language understanding.  _Advances in neural information processing systems_ , 32, 2019.
  * Zavatone-Veth and Pehlevan [2021] Jacob Zavatone-Veth and Cengiz Pehlevan.  Exact marginal prior distributions of finite bayesian neural networks.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Zhang et al. [2019] Hongyi Zhang, Yann N Dauphin, and Tengyu Ma.  Fixup initialization: Residual learning without normalization.  _arXiv preprint arXiv:1901.09321_ , 2019.
  * Zhang et al. [2020] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra.  Why are adaptive methods good for attention models?  _Advances in Neural Information Processing Systems_ , 33:15383–15393, 2020.

## Appendix

###  Appendix A Proof of Theorems

Recall the defining equations of a Transformer:

$\displaystyle{\bf Z}^{\ell}=\alpha_{1}{\bf S}^{\ell}+{\bf X}^{\ell}$
$\displaystyle{\bf Y}^{\ell}=\sigma({\bf Z}^{\ell}{\bf W}^{F,1}){\bf W}^{F,2}$
$\displaystyle{\bf X}^{\ell+1}=\alpha_{2}{\bf Y}^{\ell}+{\bf Z}^{\ell},$

where the self attention layers are defined as follows:

${\bf S}^{\ell}:={\bf A}^{\ell}{\bf X}^{\ell}{\bf W}^{V},\text{ where }{\bf
A}^{\ell}=\text{softmax}\left(\frac{1}{\sqrt{d_{k}}}{\bf X}^{\ell}{\bf
W}^{Q}\left({\bf X}^{\ell}{\bf W}^{K}\right)^{\top}\right).$

We remark that when it is clear from the context, we suppress the $\ell$ index
to improve readability.

Initialization: Recall that we initialize our weights with the so called
“Xavier” [Glorot and Bengio, 2010b] or “He” [He et al., 2015] initialization:
each weight is sampled independently from a distribution with zero-mean and
variance $\sigma^{2}_{v}=\frac{1}{d_{v}}$ for the values and feedforward
weights and $\sigma^{2}_{k}=\frac{1}{d_{k}}$ for the queries and keys.

Kronecker Delta: we introduce the Kronecker Delta notation

$\delta_{ab}=\begin{cases}1&a=b\\\ 0&a\neq b\end{cases}$

and, similarly: $\delta_{a\neq b}=1-\delta_{ab}$.

Notation: in this section, we also adopt the following shorthand notation for
the argument of the softmax ${\bf M}:=\frac{{\bf
X}\mathbf{W}^{Q}{\mathbf{W}^{K}}^{\top}{\bf X}^{\top}}{\sqrt{d_{k}}}$. We will
first compute the gradients with respect to values, queries and input (note
that the gradients of the keys have the form as the queries, hence we omit the
derivation). Recall that for a matrix ${\bf X}$, we use ${\bf X}_{k}$ to
indicate its $k$-th row. Finally, we indicate with $\bm{1}_{n\times n}$ the
$n\times n$ matrix with all ones, with $\bm{1}_{n}$ the columns vector with
all ones, and with ${\bf I}_{n}$ the $n$-dimensional identity matrix.

####  A.1 Backward Pass: Proofs of Lemma 3.1 and Theorem 3.1

In this section, we now look at the proofs of Lemma 3.1 and Theorem 3.1. We
will first introduce our notation for the gradients, as well as some useful
properties of the Kronecker product.

#####  A.1.1 Preliminaries

For the gradients, we avoid directly working with tensors by vectorizing the
matrices in a row-wise fashion ($\operatorname{vec}_{r}$) and arranging the
Jacobian in the numerator layout. More formally,

$\dfrac{\partial\mathbf{Y}}{\partial\mathbf{X}}:=\dfrac{\partial\operatorname{vec}_{r}(\mathbf{Y})}{\partial\operatorname{vec}_{r}(\mathbf{X})^{\top}}\,.$

Alongside this, we use the following rule ($\mathbin{\mathop{\otimes}}$ is the
Kronecker product):

$\displaystyle\frac{\partial\mathbf{A}\mathbf{W}\mathbf{B}}{\partial\mathbf{W}}=\mathbf{A}\mathbin{\mathop{\otimes}}\;\mathbf{B}^{\top}\,.$
(20)

For the proof of this rule, we refer to Singh et al. [2021], and to Magnus and
Neudecker [2019] for a complete introduction to matrix calculus. We will also
use the following well-known properties of the Kronecker product.

###### Lemma A.1.

Given the matrices ${\bf A}\in\mathbb{R}^{n\times m}$, ${\bf
B}\in\mathbb{R}^{p\times q}$, ${\bf C}\in\mathbb{R}^{m\times r}$, ${\bf
D}\in\mathbb{R}^{q\times s}$, then the following holds:
$\operatorname{tr}({\bf A}\mathbin{\mathop{\otimes}}{\bf
B})=\operatorname{tr}({\bf A})\operatorname{tr}({\bf B}),$ (21) and $({\bf
A}\mathbin{\mathop{\otimes}}{\bf B})({\bf C}\mathbin{\mathop{\otimes}}{\bf
D})=({\bf A}{\bf C})\mathbin{\mathop{\otimes}}({\bf B}{\bf D}).$ (22)

#####  A.1.2 Proof of Lemma 3.1

In Lemma A.2 and Lemma A.3 we compute the gradients with respect to the
queries, values and ${\bf X}$, respectively. Then we use these results to
prove Lemma 3.1 by computing the expectation of the Frobenius norms.

###### Lemma A.2 (Gradients of Self Attention for parameter matrices).

The gradients of the self attention layer defined in Eq. (1) have the
following form: $\displaystyle\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}=\operatorname{softmax}\left(\frac{{\bf
X}\mathbf{W}^{Q}{\mathbf{W}^{K}}^{\top}{\bf
X}^{\top}}{\sqrt{d_{k}}}\right){\bf X}\,\mathbin{\mathop{\otimes}}{\bf
I}_{d_{v}}\,$ $\displaystyle\frac{\partial{\bf
S}}{\partial\mathbf{W}^{Q}}=\left({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\mathbf{W}^{V}}^{\top}{\bf
X}^{\top}\right)\frac{\partial{\bf A}}{\partial{\bf M}}\left(\frac{{\bf
X}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{K}}{\sqrt{d_{k}}}\right),$
where the gradients of the softmax with respect to its inputs are as follows:
$\frac{\partial{\bf A}}{\partial{\bf
M}}=\operatorname{blockdiag}\Bigg{(}\dfrac{\partial{\bf A}_{i}}{\partial{\bf
M}_{i}^{\top}}\Bigg{)}$ (23) and where $\dfrac{\partial{\bf
A}_{i}}{\partial{\bf M}_{i}^{\top}}=\operatorname{diag}({\bf A}_{i})-{\bf
A}_{i}{\bf A}_{i}^{\top}$ with ${\bf A}_{i}$ being the $i$-th row of ${\bf A}$
in column vector format.
Finally, note that under the uniform-attention assumption, Eq. (23) simplifies
to: $\frac{\partial{\bf A}}{\partial{\bf M}}=\frac{1}{n}{\bf
I}_{n}\mathbin{\mathop{\otimes}}\left({\bf I}_{n}-\frac{1}{n}\bm{1}_{n\times
n}\right).$ (24)

###### Proof.

Let’s start with the simple case of the values’ weights $\mathbf{W}^{V}$.
Using the rule in Eq. (20), it is immediate that:

$\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}=\operatorname{softmax}\left(\frac{{\bf
X}\mathbf{W}^{Q}{\mathbf{W}^{K}}^{\top}{\bf
X}^{\top}}{\sqrt{d_{k}}}\right){\bf X}\,\mathbin{\mathop{\otimes}}{\bf
I}_{d_{v}}={\bf A}{\bf X}\,\mathbin{\mathop{\otimes}}{\bf I}_{d_{v}}\,.$

For the queries, a simple application of the chain rule and then again Eq.
(20) gives:

$\displaystyle\frac{\partial{\bf S}}{\partial\mathbf{W}^{Q}}$
$\displaystyle=\frac{\partial{\bf S}}{\partial{\bf A}}\frac{\partial{\bf
A}}{\partial\mathbf{W}^{Q}}=\frac{\partial{\bf S}}{\partial{\bf
A}}\frac{\partial{\bf A}}{\partial{\bf M}}\frac{\partial{\bf
M}}{\partial\mathbf{W}^{Q}}$ $\displaystyle=\left({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\mathbf{W}^{V}}^{\top}{\bf
X}^{\top}\right)\frac{\partial{\bf A}}{\partial{\bf M}}\left(\frac{{\bf
X}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{K}}{\sqrt{d_{k}}}\right)\,,$

which is the desired results. Finally, for the gradients of the softmax note
that:

$\frac{\partial{{\bf A}_{pq}}}{\partial{\bf
M}_{ij}}=\frac{\partial}{\partial{\bf M}_{ij}}\frac{\exp({\bf
M}_{pq})}{\sum_{k}\exp({\bf M}_{pk})}=\delta_{ip}\delta_{jq}{\bf
A}_{ij}-\delta_{ip}{\bf A}_{iq}{\bf A}_{ij}.$

By writing the above expression in the matrix notation described above, we
obtain the desired result. More specifically, the block diagonal structure is
given from the term $\delta_{ip}$ which stems from the fact that the softmax
is applied row-wise. ∎

###### Lemma A.3 (Gradients of Self Attention with respect to the Embedding
matrix).

The gradients of the self attention layer with respect to the embedding matrix
${\bf X}$ defined in Eq. (1) have the following form
$\displaystyle\frac{\partial{\bf S}}{\partial{\bf X}}$
$\displaystyle=\frac{1}{\sqrt{d_{k}}}({\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}{\bf
X}^{\top})\,\frac{\partial{\bf A}}{\partial{\bf M}}\,\left({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{K}\mathbf{W}^{Q\top}+{\bf
K}_{nn}({\bf I}_{n}\mathbin{\mathop{\otimes}}{\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top})\right)\,+\,{\bf
A}\mathbin{\mathop{\otimes}}{\mathbf{W}^{V}}^{\top},$ (25) where the gradients
of the softmax with respect to its inputs are denoted by $\frac{\partial{\bf
A}}{\partial{\bf M}}$ as before.

###### Proof.

Remember that we defined ${\bf
S}=\operatorname{softmax}({\frac{1}{\sqrt{d_{k}}}}{\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top}{\bf X}^{\top}){\bf X}\mathbf{W}^{V}$.
Alongside with our previous shorthands ${\bf A}$, ${\bf M}$, let us define the
remaining ${\bf X}\mathbf{W}^{V}$ as a matrix ${\bf T}$, so that ${\bf S}={\bf
A}\,{\bf T}$. Both ${\bf A}$ and ${\bf T}$ are functions of ${\bf X}$. So the
matrix differential can be written as:

$\displaystyle\frac{\partial{\bf S}}{\partial{\bf X}}$
$\displaystyle=\frac{\partial{\bf S}}{\partial{\bf A}}\frac{\partial{\bf
A}}{\partial{\bf X}}+\frac{\partial{\bf S}}{\partial{\bf T}}\frac{\partial{\bf
T}}{\partial{\bf X}}$ (26) $\displaystyle=\frac{\partial{\bf S}}{\partial{\bf
A}}\frac{\partial{\bf A}}{\partial{\bf M}}\frac{\partial{\bf M}}{\partial{\bf
X}}+\frac{\partial{\bf S}}{\partial{\bf T}}\frac{\partial{\bf T}}{\partial{\bf
X}}$ (27) $\displaystyle=({\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}{\bf
X}^{\top})\,\frac{\partial{\bf A}}{\partial{\bf M}}\,\frac{\partial{\bf
M}}{\partial{\bf X}}+({\bf A}\mathbin{\mathop{\otimes}}{\bf I}_{d})({\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top})$ (28) $\displaystyle=({\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}{\bf
X}^{\top})\,\frac{\partial{\bf A}}{\partial{\bf M}}\,\frac{\partial{\bf
M}}{\partial{\bf X}}+({\bf A}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top})$
(29)

Next, we use the matrix differential and then the identification theorem of
matrix derivatives to compute the matrix gradient $\frac{\partial{\bf
A}}{\partial{\bf X}}$

$\displaystyle\mathrm{d}{\bf A}$
$\displaystyle={\frac{1}{\sqrt{d_{k}}}}\mathrm{d}({\bf
X})\,\mathbf{W}^{Q}\mathbf{W}^{K\top}{\bf
X}^{\top}+{\frac{1}{\sqrt{d_{k}}}}{\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top}\,\mathrm{d}({\bf X}^{\top}).$

Vectorizing both sides:

$\displaystyle\mathrm{d}\operatorname{vec}_{r}({\bf A})$
$\displaystyle={\frac{1}{\sqrt{d_{k}}}}({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf
X}\mathbf{W}^{K}\mathbf{W}^{Q\top})\mathrm{d}(\operatorname{vec}_{r}({\bf
X}))\,+\,{\frac{1}{\sqrt{d_{k}}}}({\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top}\mathbin{\mathop{\otimes}}{\bf
I}_{n})\,\mathrm{d}(\operatorname{vec}_{r}({\bf X}^{\top}))$
$\displaystyle={\frac{1}{\sqrt{d_{k}}}}({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf
X}\mathbf{W}^{K}\mathbf{W}^{Q\top})\mathrm{d}(\operatorname{vec}_{r}({\bf
X}))\,+\,{\frac{1}{\sqrt{d_{k}}}}({\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top}\mathbin{\mathop{\otimes}}{\bf I}_{n}){\bf
K}_{dn}\,\mathrm{d}(\operatorname{vec}_{r}({\bf X})).$

Recall, for an arbitrary matrix ${\bf B}\in\mathbb{R}^{m\times n}$, the
commutation matrix ${\bf K}_{mn}$ transforms columnwise vectorization into
rowwise vectorization. More precisely,

$\displaystyle{\bf K}_{mn}\operatorname{vec}_{c}({\bf
B})=\operatorname{vec}_{c}({\bf B}^{\top})$

and $\operatorname{vec}_{c}({\bf B})=\operatorname{vec}_{r}({\bf B}^{\top})$.
Therefore, for rowwise vectorization, we have a similar result:

$\displaystyle{\bf K}_{mn}\operatorname{vec}_{r}({\bf B}^{\top})$
$\displaystyle=\operatorname{vec}_{r}({\bf B})$
$\displaystyle\operatorname{vec}_{r}({\bf B}^{\top})$ $\displaystyle={\bf
K}_{nm}\operatorname{vec}_{r}({\bf B}),$

where in the last line we used the fact the commutation is a permutation
matrix, so ${\bf K}_{mn}^{-1}={\bf K}_{mn}^{\top}={\bf K}_{nm}$. Thus, we get
the required matrix derivative as follows:

$\frac{\partial{\bf A}}{\partial{\bf X}}={\frac{1}{\sqrt{d_{k}}}}{\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf
X}\mathbf{W}^{K}\mathbf{W}^{Q\top}+{\frac{1}{\sqrt{d_{k}}}}({\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top}\mathbin{\mathop{\otimes}}{\bf I}_{n}){\bf
K}_{dn}\,.$

Next, we will use a property of commutation matrix to make things simpler
(Theorem 7.9, Magnus and Neudecker [2019]):

$\frac{\partial{\bf A}}{\partial{\bf X}}={\frac{1}{\sqrt{d_{k}}}}{\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf
X}\mathbf{W}^{K}\mathbf{W}^{Q\top}+{\frac{1}{\sqrt{d_{k}}}}{\bf K}_{nn}({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{Q}\mathbf{W}^{K\top}).$

Plugging this into the above Eq. (29), we get:

$\displaystyle\frac{\partial{\bf S}}{\partial{\bf X}}$
$\displaystyle={\frac{1}{\sqrt{d_{k}}}}({\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}{\bf
X}^{\top})\,\frac{\partial{\bf A}}{\partial{\bf M}}\,\left({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{K}\mathbf{W}^{Q\top}+{\bf
K}_{nn}({\bf I}_{n}\mathbin{\mathop{\otimes}}{\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top})\right)\,+\,{\bf
A}\mathbin{\mathop{\otimes}}{\mathbf{W}^{V}}^{\top}.$

As a sanity check, we can calculate if the shapes of the matrices are
consistent. LHS should be a $nd\times nd$ matrix, while the constituent
matrices of the first term on RHS: ${\bf
I}_{n}\mathbin{\mathop{\otimes}}{\mathbf{W}^{V}}^{\top}{\bf
X}^{\top}\in\mathbb{R}^{nd\times n^{2}}$, $\frac{\partial{\bf A}}{\partial{\bf
M}}\in\mathbb{R}^{n^{2}\times n^{2}}$, the additive term next to it is a
$n^{2}\times nd$ matrix, and the second term on RHS is a Kronecker product of
a $n\times n$ and a $d\times d$ matrix. ∎

See 3.1

###### Proof.

Here, we suppress the index $\ell$.

Gradient with respect to the values matrix.

Recall that from Lemma A.2 we have that:

$\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}=\operatorname{softmax}\left(\frac{{\bf
X}\mathbf{W}^{Q}{\mathbf{W}^{K}}^{\top}{\bf
X}^{\top}}{\sqrt{d_{k}}}\right){\bf X}\,\mathbin{\mathop{\otimes}}{\bf
I}_{d_{v}}\overset{\text{Ass.
}\ref{ass:uniform_softmax}}{=}\frac{1}{n}\bm{1}_{n\times n}{\bf
X}\mathbin{\mathop{\otimes}}{\bf I}_{d_{v}}.$

By direct computation:

$\displaystyle\left\lVert\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}\right\rVert_{F}^{2}=\operatorname{tr}\left(\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}\,\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}^{\top}\right)$
$\displaystyle{=}\frac{1}{n^{2}}\operatorname{tr}((\bm{1}_{n\times n}{\bf
X}{\bf X}^{\top}\bm{1}_{n\times n})\mathbin{\mathop{\otimes}}{\bf I}_{d_{v}})$
$\displaystyle{=}\frac{1}{n^{2}}\operatorname{tr}(\bm{1}_{n\times n}{\bf
X}{\bf X}^{\top}\bm{1}_{n\times n})\operatorname{tr}({\bf I}_{d_{v}})$
$\displaystyle=\frac{d}{n^{2}}\operatorname{tr}({\bf X}{\bf
X}^{\top}\bm{1}_{n\times n}\bm{1}_{n\times n})$
$\displaystyle=\frac{d_{v}}{n}\operatorname{tr}({\bf X}{\bf
X}^{\top}\bm{1}_{n\times n})$
$\displaystyle=\frac{d_{v}}{n}\operatorname{tr}({\bf X}{\bf
X}^{\top}\bm{1}_{n}\bm{1}_{n}^{\top})$
$\displaystyle=\frac{d_{v}}{n}\bm{1}_{n}^{\top}{\bf X}{\bf
X}^{\top}\bm{1}_{n}$ $\displaystyle=\frac{d_{v}}{n}\|{\bf
X}^{\top}\bm{1}_{n}\|^{2}=d_{v}n\left\lVert\bm{\bar{x}}\right\rVert^{2}.$

Gradients with respect to the queries/keys matrix.

First, recall the expression for the gradient of the softmax under the
uniform-attention assumption (Eq. (24)):

$\frac{\partial{\bf A}}{\partial{\bf M}}=\frac{1}{n}{\bf
I}_{n}\mathbin{\mathop{\otimes}}\left({\bf I}_{n}-\frac{1}{n}\bm{1}_{n\times
n}\right).$

Hence, we can rewrite the expression of Lemma A.2 for the gradients of the
queries as:

$\displaystyle\frac{\partial{\bf S}}{\partial\mathbf{W}^{Q}}$
$\displaystyle=\left({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\mathbf{W}^{V}}^{\top}{\bf
X}^{\top}\right)\frac{\partial{\bf A}}{\partial{\bf M}}\left(\frac{{\bf
X}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{K}}{\sqrt{d_{k}}}\right)$
$\displaystyle=\frac{1}{\sqrt{d_{k}}n}\left({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\mathbf{W}^{V}}^{\top}{\bf
X}^{\top}\right)\left[{\bf I}_{n}\otimes\left({\bf
I}_{n}-\frac{1}{n}\bm{1}_{n\times n}\right)\right]\left({\bf
X}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{K}\right)$
$\displaystyle=\frac{1}{\sqrt{d_{k}}n}{\bf
X}\mathbin{\mathop{\otimes}}\left[{\mathbf{W}^{V}}^{\top}{\bf
X}^{\top}\left({\bf I}_{n}-\frac{1}{n}\bm{1}_{n\times n}\right){\bf
X}\mathbf{W}^{K}\right],$

where in the last step we have used twice the property of the Kronecker
product in Eq. (22) of Lemma A.1.

Hence,

$\displaystyle\left\|\frac{\partial{\bf
S}}{\partial\mathbf{W}^{Q}}\right\|_{F}^{2}=\operatorname{tr}\left(\frac{\partial{\bf
S}}{\partial\mathbf{W}^{Q}}\frac{\partial{\bf
S}}{\partial\mathbf{W}^{Q}}^{T}\right)\overset{\eqref{eq:trace_kro}}{=}\frac{1}{d_{k}n^{2}}\|{\bf
X}\|^{2}_{F}\cdot\left\|{\mathbf{W}^{V}}^{\top}{\bf X}^{\top}\left({\bf
I}_{n}-\frac{1}{n}\bm{1}_{n\times n}\right){\bf
X}\mathbf{W}^{K}\right\|_{F}^{2},$

where we have used the property on the trace of the Kronecker product (Lemma
A.1, Eq. (21)). Note that if we are conditioning on ${\bf X}$, then we only
have to take the expectation of the last term with respect to the weights
${\bf W}^{K}$ and ${\bf W}^{V}$. Let us call ${\bf L}:={\bf
I}_{n}-\frac{1}{n}\bm{1}_{n\times n}$ for notation simplicity.

Note: for a matrix ${\bf W}\in\mathbb{R}^{d\times d}$ whose entries
$w_{ij}\sim\mathcal{N}(0,\sigma^{2})$, then $\mathbb{E}{\bf W}{\bf
W}^{\top}=d\sigma^{2}\,{\bf I}_{d}$. Thus, exchanging the order of trace and
expectation, we can write:

$\displaystyle\mathbb{E}\|\mathbf{W}^{V\top}{\bf X}^{\top}{\bf L}{\bf
X}\mathbf{W}^{K}\|^{2}_{F}$
$\displaystyle=\mathbb{E}\operatorname{tr}(\mathbf{W}^{V\top}{\bf
X}^{\top}{\bf L}{\bf X}\mathbf{W}^{K}\,\cdot\,\mathbf{W}^{K\top}{\bf
X}^{\top}{\bf L}{\bf X}\mathbf{W}^{V})$ $\displaystyle=\operatorname{tr}({\bf
X}^{\top}{\bf L}{\bf X}\mathbb{E}[\mathbf{W}^{K}\mathbf{W}^{K\top}]{\bf
X}^{\top}{\bf L}{\bf X}\mathbb{E}[\mathbf{W}^{V}\mathbf{W}^{V\top}])$
$\displaystyle=\sigma^{2}_{v}\sigma^{2}_{k}d_{k}d_{v}\operatorname{tr}({\bf
X}^{\top}{\bf L}{\bf X}\,\cdot\,{\bf X}^{\top}{\bf L}{\bf X})$
$\displaystyle=\sigma^{2}_{v}\sigma^{2}_{k}d_{k}d_{v}\|{\bf X}^{\top}{\bf
L}{\bf X}\|^{2}_{F}$
$\displaystyle=\sigma^{2}_{v}\sigma^{2}_{k}d_{k}d_{v}\left\lVert{\bf
X}^{\top}({\bf I}_{n}-\frac{1}{n}\bm{1}_{n}\bm{1}_{n}^{\top}){\bf
X}\right\rVert^{2}_{F}$
$\displaystyle=\sigma^{2}_{v}\sigma^{2}_{k}d_{k}d_{v}\|{\bf X}^{\top}{\bf
X}-n\bar{\bm{x}}\bar{\bm{x}}^{\top}\|^{2}_{F},$

where, $\bar{\bm{x}}=\frac{1}{n}{\bf X}^{\top}\bm{1}_{n}\in\mathbb{R}^{d}$ is
the mean embedding. Multiply this by $\frac{1}{d_{k}n^{2}}\|{\bf X}\|^{2}_{F}$
to get the final answer.

Gradient with respect to the input.

Plugging in the values of $\frac{\partial{\bf A}}{\partial{\bf M}}$ and ${\bf
A}$ under the uniform-attention assumption into Eq. (25) gives rise to the
following:

$\displaystyle\frac{\partial{\bf S}}{\partial{\bf X}}$
$\displaystyle=\frac{1}{n{\sqrt{d_{k}}}}{\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}{\bf X}^{\top}\left({\bf
I}_{n}-\frac{1}{n}\bm{1}_{n}\bm{1}_{n}^{\top}\right){\bf
X}\mathbf{W}^{K}\mathbf{W}^{Q\top}$
$\displaystyle+\frac{1}{n{\sqrt{d_{k}}}}\left[{\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}{\bf X}^{\top}\left({\bf
I}_{n}-\frac{1}{n}\bm{1}_{n}\bm{1}_{n}^{\top}\right)\right]{\bf K}_{nn}({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf X}\mathbf{W}^{Q}\mathbf{W}^{K\top})\,$
$\displaystyle+\frac{1}{n}\bm{1}_{n\times
n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}$

Let’s refer to the matrices on the right-hand side as ${\bf A}_{1},{\bf
A}_{2},{\bf A}_{3}$ respectively. We compute the expected squared Frobenius
norm of these as follows:

For ${\bf A}_{3}$:

$\displaystyle\mathbb{E}[\left\lVert{\bf A}_{3}\right\rVert^{2}_{F}]$
$\displaystyle=\frac{1}{n^{2}}\mathbb{E}[\operatorname{tr}(n\bm{1}_{n\times
n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}\mathbf{W}^{V})]$
$\displaystyle{=}\frac{1}{n}\operatorname{tr}(\bm{1}_{n\times
n})\operatorname{tr}(\mathbb{E}[\mathbf{W}^{V}\mathbf{W}^{V\top}])=d_{v}^{2}\sigma^{2}_{v}.$

Similarly, for ${\bf A}_{1}$:

$\displaystyle\mathbb{E}[\left\lVert{\bf A}_{1}\right\rVert^{2}_{F}]$
$\displaystyle=\frac{\sigma^{2}_{q}\sigma^{2}_{k}\sigma^{2}_{v}d_{k}d_{v}}{n^{2}}\,\operatorname{tr}({\bf
I}_{n})\operatorname{tr}({\bf X}^{\top}{\bf L}{\bf X}\,\cdot\,{\bf
X}^{\top}{\bf L}{\bf X})$ (30)
$\displaystyle=\frac{\sigma^{2}_{q}\sigma^{2}_{k}\sigma^{2}_{v}d_{k}d_{v}}{n}\,\left\lVert{\bf
X}^{\top}{\bf X}-n\bar{\bm{x}}\bar{\bm{x}}^{\top}\right\rVert^{2}_{F}$ (31)
$\displaystyle=\sigma^{2}_{q}\sigma^{2}_{k}\sigma^{2}_{v}d_{k}d_{v}n\,\left\lVert\frac{1}{n}{\bf
X}^{\top}{\bf X}-\bar{\bm{x}}\bar{\bm{x}}^{\top}\right\rVert^{2}_{F}.$ (32)

Finally, for ${\bf A}_{2}$:

$\displaystyle\mathbb{E}[\left\lVert{\bf A}_{2}\right\rVert^{2}_{F}]$ (33)
$\displaystyle=\frac{1}{n^{2}d_{k}}\mathbb{E}\left[\operatorname{tr}(\left[{\bf
I}_{n}\mathbin{\mathop{\otimes}}\mathbf{W}^{V\top}{\bf X}^{\top}{\bf
L}\right]{\bf K}_{nn}({\bf I}_{n}\mathbin{\mathop{\otimes}}{\bf
X}\mathbf{W}^{Q}\mathbf{W}^{K\top}\mathbf{W}^{K}\mathbf{W}^{Q\top}{\bf
X}^{\top}){\bf K}_{nn}\left[{\bf I}_{n}\mathbin{\mathop{\otimes}}{\bf L}{\bf
X}\mathbf{W}^{V}\right])\right]$ (34)
$\displaystyle=\frac{\sigma^{2}_{q}\sigma^{2}_{k}\sigma^{2}_{v}d_{k}d_{v}}{n^{2}}\,\operatorname{tr}(({\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf X}{\bf X}^{\top})[{\bf
I}_{n}\mathbin{\mathop{\otimes}}{\bf L}{\bf X}{\bf X}^{\top}{\bf L}])$ (35)
$\displaystyle=\frac{\sigma^{2}_{q}\sigma^{2}_{k}\sigma^{2}_{v}d_{k}d_{v}}{n}\,\operatorname{tr}({\bf
X}^{\top}{\bf L}{\bf X}\,\cdot\,{\bf X}^{\top}{\bf L}{\bf
X})=\mathbb{E}[\left\lVert{\bf A}_{1}\right\rVert^{2}_{F}],$ (36)

where in the second line we have taken the expectation inside and used the
fact that ${\bf K}_{nn}$, being a commutation matrix, is orthogonal. Then, by
simple properties of Kronecker product and cyclic property of trace, we have
the result, which is the same as that for ${\bf A}_{1}$.

Finally, by the triangle inequality

$\displaystyle\mathbb{E}\left\|\frac{\partial{\bf S}}{\partial{\bf
X}}\right\|^{2}$ $\displaystyle\leq 2\mathbb{E}\|{\bf A}_{1}+{\bf
A}_{2}\|^{2}+2\mathbb{E}\|{\bf A}_{3}\|^{2}$ (37) $\displaystyle\leq
4\mathbb{E}\|{\bf A}_{1}\|^{2}+4\mathbb{E}\|{\bf
A}_{2}\|^{2}+2\mathbb{E}\|{\bf A}_{3}\|^{2}$ (38)
$\displaystyle=8\mathbb{E}\|{\bf A}_{1}\|^{2}+2\mathbb{E}\|{\bf A}_{3}\|^{2}$
(39)
$\displaystyle=\frac{8\sigma^{2}_{q}\sigma^{2}_{k}\sigma^{2}_{v}d_{k}d_{v}}{n}\left\lVert{\bf
X}^{\top}{\bf
X}-n\bar{\bm{x}}\bar{\bm{x}}^{\top}\right\rVert^{2}_{F}+2d_{v}^{2}\sigma_{v}^{2}.$
(40)

This completes the proof. ∎

#####  A.1.3 Proof of Theorem 3.1

See 3.1

Before starting the proof, it is interesting to note that, even though the
gradients of queries and keys vanish in the rank collapse regime (i.e.
$\left\lVert{\bf X}^{\top}{\bf
X}-n\bar{\bm{x}}\bar{\bm{x}}^{\top}\right\rVert=0$), the gradient with respect
to the values and the input does not (see Theorem 3.1). From this simple
remark, we can conclude that, even in the rank collapse regime, information
still propagates in the backward pass. In Section 3.4 (main paper), we show
that even if gradients effectively propagate, the phenomenon studied in this
theorem still greatly affects training.

###### Proof.

By using the chain rule and the fact that for two matrixes ${\bf A},{\bf B}$
we have that $\left\lVert{\bf A}{\bf B}\right\rVert_{F}^{2}\leq\left\lVert{\bf
A}\right\rVert_{F}^{2}\left\lVert{\bf B}\right\rVert_{F}^{2}$, we can upper
bound the gradient as:

$\displaystyle\left\lVert\frac{\partial\mathcal{L}}{\partial{{\bf
W}^{Q,\ell}}}\right\rVert_{F}^{2}$
$\displaystyle\leq\prod_{i=\ell+1}^{L-1}\left\lVert\frac{\partial{\bf
X}^{i+1}}{\partial{\bf
X}^{i}}\right\rVert_{F}^{2}\left\lVert\frac{\partial\mathcal{L}}{\partial{\bf
X}^{L}}\right\rVert_{F}^{2}\left\lVert\frac{\partial{\bf
X}^{\ell+1}}{\partial{\bf W}^{Q,\ell}}\right\rVert_{F}^{2}$
$\displaystyle\leq\prod_{i=\ell+1}^{L-1}\left\lVert\frac{\partial{\bf
X}^{i+1}}{\partial{\bf
X}^{i}}\right\rVert_{F}^{2}\left\lVert\frac{\partial\mathcal{L}}{\partial{\bf
X}^{L}}\right\rVert_{F}^{2}\left\lVert\frac{\partial{\bf
X}^{\ell+1}}{\partial{\bf
Z}^{\ell}}\right\rVert_{F}^{2}\left\lVert\frac{\partial{\bf
Z}^{\ell}}{\partial{\bf W}^{Q,\ell}}\right\rVert_{F}^{2}$
$\displaystyle\leq\prod_{i=\ell+1}^{L-1}\left\lVert\frac{\partial{\bf
X}^{i+1}}{\partial{\bf
X}^{i}}\right\rVert_{F}^{2}\left\lVert\frac{\partial\mathcal{L}}{\partial{\bf
X}^{L}}\right\rVert_{F}^{2}\left\lVert\frac{\partial{\bf
X}^{\ell+1}}{\partial{\bf
Z}^{\ell}}\right\rVert_{F}^{2}\left(\left\lVert\frac{\partial\alpha_{1}{\bf
S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\rVert_{F}^{2}+\underbrace{\left\lVert\frac{\partial{\bf
X}^{\ell}}{\partial{\bf W}^{Q,\ell}}\right\rVert_{F}^{2}}_{=0}\right),$

where we recall that ${\bf Z}^{\ell}=\alpha_{1}{\bf S}^{\ell}+{\bf X}^{\ell}$
and in the last step we have used that ${\bf X}^{\ell}$ does not depend on
${\bf W}^{Q,\ell}$, hence the gradient vanishes. By taking expectation and
using the tower property, we have that:

$\mathbb{E}\left\lVert\frac{\partial\mathcal{L}}{\partial{{\bf
W}^{Q,\ell}}}\right\rVert_{F}^{2}\leq\mathbb{E}\left[\underbrace{\mathbb{E}\left[\prod_{i=\ell+1}^{L-1}\left\lVert\frac{\partial{\bf
X}^{i+1}}{\partial{\bf
X}^{i}}\right\rVert_{F}^{2}\left\lVert\frac{\partial\mathcal{L}}{\partial{\bf
X}^{L}}\right\rVert_{F}^{2}\left\lVert\frac{\partial{\bf
X}^{\ell+1}}{\partial{\bf Z}^{\ell}}\right\rVert_{F}^{2}\right]}_{=:G({\bf
X}^{\ell})}\left\lVert\frac{\partial\alpha_{1}{\bf S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\rVert_{F}^{2}\right],$

where the expectations are taken with respect to ${\bf X}^{\ell}$ for the
outer one and conditioning on ${\bf X}^{\ell}$ for inner one. Indeed, the
first three terms only depend on the network values after ${\bf X}^{\ell}$.
Now, a repeated application of the tower property in $G({\bf X}^{\ell})$,
together with the results on the gradients of Lemma 3.1, easily shows that
$G({\bf X}^{\ell})$ stays bounded under our hypothesis. To see this one can
also simply note that, since the softmax and its derivatives are almost surely
bounded, the boundedness of $G({\bf X}^{\ell})$ is implied by an analogous
statement for a vanilla linear MLP (i.e removing the softmax). In this
setting, the random variable inside the expectation in $G({\bf X}^{\ell})$ is
a finite linear combination of Gaussian products — which has bounded
expectation.

All in all, we have that

$\mathbb{E}\left\lVert\frac{\partial\mathcal{L}}{\partial{{\bf
W}^{Q,\ell}}}\right\rVert_{F}^{2}\leq\mathbb{E}\left[B_{{\bf
X}^{\ell}}\left\lVert\frac{\partial\alpha_{1}{\bf S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\rVert_{F}^{2}\right],$

where $B_{{\bf X}^{\ell}}$ is an almost-surely-bounded function of ${\bf
X}^{\ell}$. Hence, to show that
$\mathbb{E}\left\lVert\frac{\partial\mathcal{L}}{\partial{{\bf
W}^{Q,\ell}}}\right\rVert_{F}^{2}=0$, we now just need to show that:

$\mathbb{E}\left\lVert\frac{\partial\alpha_{1}{\bf S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\rVert_{F}^{2}=0$

under the rank-1 hypothesis for ${\bf X}^{\ell}$. Let ${\bf
X}_{1}^{\ell},\dots\bf X_{n}^{\ell}\in\mathbb{R}^{d_{v}}$ be the
representations for the $n$ tokens. Under the rank-1 assumption, each token
can be written as a multiple of a single vector $\bm{x}\in\mathbb{R}^{d_{v}}$,
and hence there exists $a_{1},\dots,a_{n}\in\mathbb{R}$ such that ${\bf
X}_{1}=a_{1}\bm{x},\dots,{\bf X}_{n}=a_{n}\bm{x}$. From Lemma 3.1, we know
that:

$\mathbb{E}\left\|\frac{\partial{\bf
S}^{\ell}}{\partial\mathbf{W}^{Q}}\right\|^{2}_{F}=\frac{\sigma^{2}_{v}\sigma^{2}_{k}d^{2}}{n^{2}}\cdot\mathbb{E}\left[\|{\bf
X}^{\ell}\|^{2}_{F}\cdot\|({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}-n\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\|^{2}_{F}\right].$

The mean token simplifies to $\bar{\bm{x}}^{l}=\frac{\bm{x}}{n}\sum_{k}a_{k}$
and hence
$\left(\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\right)_{ij}=\frac{1}{n^{2}}(\sum_{k}a_{k})^{2}x_{i}x_{j}$.
Similarly, $\left(({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}\right)_{ij}=\sum_{k}a_{k}^{2}x_{i}x_{j}$. If furthermore all the
coefficients $a_{i}$ are the same (which corresponds to the rank collapse
assumption ${\bf X}^{\ell}=\bm{1}_{n}\bm{x}^{T}$ analyzed here), then it is
easy to see that $\left(({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}\right)_{ij}-n\left(\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\right)_{ij}=0\;\forall
i,j$ and hence $\|({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}-n\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\|^{2}_{F}=0$. ∎

####  A.2 Gradient Analysis of Section 3.3

Throughout this section we assume that between every pair of tokens, the same
dimension is a zero-mean Gaussian random variable with the same correlation,
meaning that

$\mathbb{E}\left[{\bf X}_{i,j}{\bf
X}_{i^{\prime},j^{\prime}}\right]=\begin{cases}0&j\neq j^{\prime}\ \
\text{(independent dimensions)}\\\ \sigma^{2}_{x}&i=i^{\prime},j=j^{\prime}\\\
\rho\sigma^{2}_{x}&i\neq i^{\prime},j=j^{\prime}.\end{cases}$

As we will deal with the computation of 4-th order moments of correlated
Gaussian random variables, we will make use of Isserlis theorem [Isserlis,
1918]:

###### Theorem A.1 (Isserlis).

Let $X_{1},\dots X_{m}$ be $m$ zero-mean Gaussian random variables. Then:
$\mathbb{E}[X_{1}\cdots X_{m}]=\begin{cases}\sum_{p\in
P_{m}^{2}}\prod_{(i,j)\in p}\mathbb{E}[X_{i}X_{j}]&m\text{ even}\\\ 0&m\text{
odd}\end{cases}$ (41) where $P_{m}^{2}$ is the set of all the possible
pairings of the indexes $1,\dots,m$.

In particular, we will only need the 4-th order term, which reads:

$\mathbb{E}[X_{1}X_{2}X_{3}X_{4}]=\mathbb{E}[X_{1}X_{2}]\mathbb{E}[X_{3}X_{4}]+\mathbb{E}[X_{1}X_{3}]\mathbb{E}[X_{2}X_{4}]+\mathbb{E}[X_{1}X_{4}]\mathbb{E}[X_{2}X_{3}].$

Now we can prove Eq. (17) , which we re-state here:

$\mathbb{E}\left\lVert\frac{\partial{\bf
S}}{\partial\mathbf{W}^{V}}\right\rVert_{F}^{2}=\sigma_{x}^{2}d^{2}\left(1+\rho(n-1)\right).$

Also, from Eq. (6) we have that

$\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}^{V,\ell}}\right\|^{2}_{F}=dn\mathbb{E}\|\bar{\bm{x}}^{\ell}\|^{2}.$

Now,

$\mathbb{E}\|\bar{\bm{x}}^{\ell}\|^{2}=\mathbb{E}\left(\sum_{i=1}^{d}(\bar{\bm{x}}^{\ell}_{i})^{2}\right).$

Each $\bar{\bm{x}}^{\ell}_{i}=\frac{1}{n}\sum_{k=1}^{n}{\bf X}^{\ell}_{ki}$ is
equally distributed with mean

$\mathbb{E}[\bar{\bm{x}}^{\ell}_{i}]=\mathbb{E}\left[\frac{1}{n}\sum_{k=1}^{n}{\bf
X}^{\ell}_{ki}\right]=0$

and variance

$\text{Var}[\bar{\bm{x}}^{\ell}_{i}]=\text{Var}\left[\frac{1}{n}\sum_{k=1}^{n}{\bf
X}^{\ell}_{ki}\right]=\frac{1}{n^{2}}\left(n\sigma_{x}^{2}+n(n-1)\rho\sigma_{x}^{2}\right)=\frac{1}{n}\sigma_{x}^{2}(1+\rho(n-1)).$

Finally we get

$\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}^{V,\ell}}\right\|^{2}_{F}=\sigma_{x}^{2}d^{2}(1+\rho(n-1)).$

We know prove Eq. (18), which reads:

$\mathbb{E}\left\|\frac{\partial{\bf
S}}{\partial\mathbf{W}^{Q}}\right\|_{F}^{2}=\sigma_{x}^{6}\frac{(n-1)}{n}(1-\rho)^{2}d(n+d).$

For the queries (and the keys respectively), recall from Eq. (7) that

$\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\|^{2}_{F}=\frac{\sigma^{2}_{v}\sigma^{2}_{k}d^{2}}{dn^{2}}\cdot\mathbb{E}\left[\|{\bf
X}^{\ell}\|^{2}_{F}\cdot\|({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}-n\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\|^{2}_{F}\right].$

To proceed, we drop the superscript $\ell$ and we make the additional
assumption that $\|{\bf X}\|^{2}_{F}$ is uncorrelated from the correlation
magnitude $\|{\bf X}^{\top}{\bf L}{\bf X}\|^{2}_{F}=\|{\bf X}^{\top}{\bf
X}-n\bar{\bm{x}}\bar{\bm{x}}^{\top}\|^{2}_{F}$.

Let us proceed with an expansion:

$\displaystyle\mathbb{E}\left[\|{\bf X}^{\top}{\bf L}{\bf X}\|^{2}_{F}\right]$
$\displaystyle=\sum_{i,j=1}^{d}\mathbb{E}\left[\left(\sum_{a,b=1}^{n}{\bf
X}_{ai}{\bf L}_{ab}{\bf X}_{bj}\right)^{2}\right]$
$\displaystyle=\sum_{i,j=1}^{d}\sum_{a,b,a^{\prime},b^{\prime}=1}^{n}{\bf
L}_{ab}{\bf L}_{a^{\prime}b^{\prime}}\mathbb{E}\left[{\bf X}_{ai}{\bf
X}_{a^{\prime}i}{\bf X}_{bj}{\bf X}_{b^{\prime}j}\right].$

Now, we have 2 cases: if $i\neq j$, which gives $d(d-1)$ equal terms, we need
to compute

$\text{A}:=\sum_{a,b,a^{\prime},b^{\prime}=1}^{n}{\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}\mathbb{E}\left[{\bf X}_{ai}{\bf
X}_{a^{\prime}i}\right]\cdot\mathbb{E}\left[{\bf X}_{bj}{\bf
X}_{b^{\prime}j}\right],$

where $(i,j)$ is any tuple with $i\neq j$ and we used uncorrelation of
different dimensions. Otherwise, we get $d$ each equal to

$\text{B}:=\sum_{a,b,a^{\prime},b^{\prime}=1}^{n}{\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}\mathbb{E}\left[{\bf X}_{ai}{\bf X}_{a^{\prime}i}{\bf
X}_{bi}{\bf X}_{b^{\prime}i}\right],$

where $i$ is any index.

###### Term A.

Note that

$\mathbb{E}\left[{\bf X}_{ai}{\bf
X}_{a^{\prime}i}\right]\cdot\mathbb{E}\left[{\bf X}_{bj}{\bf
X}_{b^{\prime}j}\right]=\begin{cases}\sigma_{x}^{4}&a=a^{\prime},b=b^{\prime},\quad
n^{2}\ \text{terms}\\\ \rho\sigma_{x}^{4}&a=a^{\prime},b\neq b^{\prime},\quad
n^{2}(n-1)\ \text{terms}\\\ \rho\sigma_{x}^{4}&a\neq
a^{\prime},b=b^{\prime},\quad n^{2}(n-1)\ \text{terms}\\\
\rho^{2}\sigma_{x}^{4}&a\neq a^{\prime},b\neq b^{\prime},\quad n^{2}(n-1)^{2}\
\text{terms}\end{cases}.$

So basically $A$ is the sum of 3 terms:

$\displaystyle A_{1}:=\sigma_{x}^{4}\sum_{a,b}{\bf
L}_{ab}^{2}=(n-1)\sigma_{x}^{4}$ $\displaystyle
A_{2}:=2\rho\sigma_{x}^{4}\sum_{a,b}\sum_{b^{\prime}\neq b}{\bf L}_{ab}{\bf
L}_{ab^{\prime}}=-2(n-1)\rho\sigma_{x}^{4}$ $\displaystyle
A_{3}:=\rho^{2}\sigma^{4}_{x}\sum_{a,b}\sum_{a^{\prime}\neq
a}\sum_{b^{\prime}\neq b}{\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}=\rho^{2}\sigma^{4}_{x}(n-1),$

where we leveraged the following direct calculations:

$\displaystyle A_{1}$ $\displaystyle=\sigma_{x}^{4}\sum_{a,b}{\bf L}_{ab}^{2}$
$\displaystyle=\sigma_{x}^{4}\left(n\left(\frac{n-1}{n}\right)^{2}+(n-1)n\frac{1}{n^{2}}\right)$
$\displaystyle=\sigma_{x}^{4}\frac{(n-1)^{2}+(n-1)}{n}$
$\displaystyle=\sigma_{x}^{4}(n-1).$

Next, we compute

$\displaystyle A_{2}$ $\displaystyle=2\rho\sigma_{x}^{4}\sum_{a,b}{\bf
L}_{ab}\sum_{b^{\prime}\neq b}{\bf L}_{ab^{\prime}}$
$\displaystyle=2\rho\sigma_{x}^{4}\left(\sum_{a}{\bf
L}_{a,a}\sum_{b^{\prime}\neq b}{\bf L}_{ab^{\prime}}+\sum_{a}\sum_{b\neq
a}{\bf L}_{ab}\sum_{b^{\prime}\neq b}{\bf L}_{ab^{\prime}}\right)$
$\displaystyle=2\rho\sigma_{x}^{4}\left(\sum_{a}\frac{n-1}{n}\left[-(n-1)\frac{1}{n}\right]+\sum_{a}\sum_{b\neq
a}{\bf L}_{ab}\left[\frac{n-1}{n}-\frac{n-2}{n}\right]\right)$
$\displaystyle=2\rho\sigma_{x}^{4}\left(-\frac{(n-1)^{2}}{n}-\frac{1}{n}(n-1)n\frac{1}{n}\right)$
$\displaystyle=-2\rho\sigma_{x}^{4}(n-1).$

Finally, similar computations also lead to the last term. To follow the
calculations, we invite the reader to draw the matrix ${\bf L}$ and to hide
the columns over which summations are not performed:

$\displaystyle A_{3}$ $\displaystyle=\rho^{2}\sigma^{4}_{x}\sum_{a,b}{\bf
L}_{ab}\sum_{a^{\prime}\neq a}\sum_{b^{\prime}\neq b}{\bf
L}_{a^{\prime}b^{\prime}}$
$\displaystyle=\rho^{2}\sigma^{4}_{x}\left(\sum_{a}{\bf
L}_{aa}\sum_{a^{\prime}\neq a}\sum_{b^{\prime}\neq a}{\bf
L}_{a^{\prime}b^{\prime}}+\sum_{a}\sum_{b\neq a}{\bf
L}_{ab}\sum_{a^{\prime}\neq a}\sum_{b^{\prime}\neq b}{\bf
L}_{a^{\prime}b^{\prime}}\right)$
$\displaystyle=\rho^{2}\sigma^{4}_{x}\left(\sum_{a}{\bf
L}_{aa}(1-\frac{1}{n}+\sum_{a}\sum_{b\neq a}{\bf
L}_{ab}(-1\frac{1}{n})\right)$
$\displaystyle=\rho^{2}\sigma^{4}_{x}\left(n(1-\frac{1}{n})(1-\frac{1}{n})+n(n-1)(-\frac{1}{n})(-\frac{1}{n})\right)$
$\displaystyle=\rho^{2}\sigma^{4}_{x}(n-1).$

All in all, we get:

$\text{A}=(n-1)(1-\rho)^{2}\sigma^{4}_{x}.$

###### Term B.

We make use of Isserlis theorem, stating that:

$\mathbb{E}\left[{\bf X}_{ai}{\bf X}_{a^{\prime}i}{\bf X}_{bi}{\bf
X}_{b^{\prime}i}\right]=\underbrace{\mathbb{E}{\bf X}_{ai}{\bf
X}_{a^{\prime}i}\mathbb{E}{\bf X}_{bi}{\bf
X}_{b^{\prime}i}}_{Q_{1}}+\underbrace{\mathbb{E}{\bf X}_{ai}{\bf
X}_{bi}\mathbb{E}{\bf X}_{a^{\prime}i}{\bf
X}_{b^{\prime}i}}_{Q_{2}}+\underbrace{\mathbb{E}{\bf X}_{ai}{\bf
X}_{b^{\prime}i}\mathbb{E}{\bf X}_{a^{\prime}i}{\bf X}_{bi}}_{Q_{3}}.$

By using our independence assumptions, we get:

$Q_{1}=\sigma_{x}^{4}(\delta_{aa^{\prime}}+\rho\delta_{a\neq
a^{\prime}})(\delta_{bb^{\prime}}+\rho\delta_{b\neq
b^{\prime}})=\sigma_{x}^{4}(\delta_{aa^{\prime}}\delta_{bb^{\prime}}+\rho\delta_{aa^{\prime}}\delta_{b\neq
b^{\prime}}+\rho\delta_{a\neq
a^{\prime}}\delta_{bb^{\prime}}+\rho^{2}\delta_{a\neq a^{\prime}}\delta_{b\neq
b^{\prime}}).$

Similarly for $Q_{2}$ and $Q_{3}$:

$Q_{2}=\sigma_{x}^{4}(\delta_{ab}\delta_{a^{\prime}b^{\prime}}+\rho\delta_{ab}\delta_{a^{\prime}\neq
b^{\prime}}+\rho\delta_{a\neq
b}\delta_{a^{\prime}b^{\prime}}+\rho^{2}\delta_{a\neq b}\delta_{a^{\prime}\neq
b^{\prime}})$

and

$Q_{3}=\sigma_{x}^{4}(\delta_{ab^{\prime}}\delta_{a^{\prime}b}+\rho\delta_{ab^{\prime}}\delta_{a^{\prime}\neq
b}+\rho\delta_{a\neq b^{\prime}}\delta_{a^{\prime}b}+\rho^{2}\delta_{a\neq
b^{\prime}}\delta_{a^{\prime}\neq b}).$

Hence,

$B=\sum_{a,b,a^{\prime},b^{\prime}=1}^{n}{\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}\mathbb{E}\left[{\bf X}_{ai}{\bf X}_{a^{\prime}i}{\bf
X}_{bi}{\bf X}_{b^{\prime}i}\right]=\sum_{a,b,a^{\prime},b^{\prime}=1}^{n}{\bf
L}_{ab}{\bf L}_{a^{\prime}b^{\prime}}(Q_{1}+Q_{2}+Q_{3}).$

Let’s study it term by term. We will also use ${\bf
L}_{ab}=(\delta_{ab}-\frac{1}{n})$, and so ${\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}=(\delta_{ab}\delta_{a^{\prime}b^{\prime}}-\frac{\delta_{ab}}{n}-\frac{\delta_{a^{\prime}b^{\prime}}}{n}+\frac{1}{n^{2}})$.

First term: we have that
$\sigma_{x}^{4}\sum_{aa^{\prime}b^{\prime}b^{\prime}}{\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}Q_{1}$ which is equal to (omitting the constant
$\sigma_{x}^{4}$):

$\displaystyle=\sum_{a,a^{\prime},b,b^{\prime}}(\delta_{ab}\delta_{a^{\prime}b^{\prime}}-\frac{\delta_{ab}}{n}-\frac{\delta_{a^{\prime}b^{\prime}}}{n}+\frac{1}{n^{2}})(\delta_{aa^{\prime}}\delta_{bb^{\prime}}+\rho\delta_{aa^{\prime}}\delta_{b\neq
b^{\prime}}+\rho\delta_{a\neq
a^{\prime}}\delta_{bb^{\prime}}+\rho^{2}\delta_{a\neq a^{\prime}}\delta_{b\neq
b^{\prime}})$
$\displaystyle=\rho^{2}\left(n(n-1)-2(n-1)(n-1)+(n-1)^{2}\right)+\rho(-4(n-1)+2(n-1))+n-2+1$
$\displaystyle=\rho^{2}(n-1)\left(n-2(n-1)+(n-1)\right)-2\rho(n-1)+(n-1)$
$\displaystyle=\rho^{2}(n-1)-2\rho(n-1)+(n-1).$

Second term: we have that
$\sigma_{x}^{4}\sum_{aa^{\prime}b^{\prime}b^{\prime}}{\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}Q_{2}$ which is equal to (omitting the constant
$\sigma_{x}^{4}$):

$\displaystyle=\sum_{a,a^{\prime},b,b^{\prime}}(\delta_{ab}\delta_{a^{\prime}b^{\prime}}-\frac{\delta_{ab}}{n}-\frac{\delta_{a^{\prime}b^{\prime}}}{n}+\frac{1}{n^{2}})(\delta_{ab}\delta_{a^{\prime}b^{\prime}}+\rho\delta_{ab}\delta_{a^{\prime}\neq
b^{\prime}}+\rho\delta_{a\neq
b}\delta_{a^{\prime}b^{\prime}}+\rho^{2}\delta_{a\neq b}\delta_{a^{\prime}\neq
b^{\prime}})$
$\displaystyle=\rho^{2}(n-1)^{2}+\rho(-2n(n-1)+2(n-1))+n^{2}-2n+1$
$\displaystyle=\rho^{2}(n-1)^{2}-2\rho(n-1)^{2}+(n-1)^{2}.$

Third term: we have that
$\sigma_{x}^{4}\sum_{aa^{\prime}b^{\prime}b^{\prime}}{\bf L}_{ab}{\bf
L}_{a^{\prime}b^{\prime}}Q_{3}$ which is equal to (omitting the constant
$\sigma_{x}^{4}$):

$\displaystyle=\sum_{a,a^{\prime},b,b^{\prime}}(\delta_{ab}\delta_{a^{\prime}b^{\prime}}-\frac{\delta_{ab}}{n}-\frac{\delta_{a^{\prime}b^{\prime}}}{n}+\frac{1}{n^{2}})(\delta_{ab^{\prime}}\delta_{a^{\prime}b}+\rho\delta_{ab^{\prime}}\delta_{a^{\prime}\neq
b}+\rho\delta_{a\neq b^{\prime}}\delta_{a^{\prime}b}+\rho^{2}\delta_{a\neq
b^{\prime}}\delta_{a^{\prime}\neq b})$
$\displaystyle=\rho^{2}\left(n(n-1)-2(n-1)(n-1)+(n-1)^{2}\right)+\rho(-4(n-1)+2(n-1))+n-2+1$
$\displaystyle=\rho^{2}(n-1)-2\rho(n-1)+(n-1).$

Summing all the three terms, we get:

$\displaystyle
B=\sigma_{x}^{4}(n-1)\left[\rho^{2}(n+1)-2(n+1)\rho+(n+1)\right]=\sigma_{x}^{4}(n-1)(n+1)(1-\rho)^{2}.$

Plugging in the values of A and B we get:

$\mathbb{E}[\|{\bf X}^{\top}{\bf L}{\bf X}\|^{2}_{F}]=d\cdot B+d(d-1)\cdot
A=\sigma^{4}_{x}(1-\rho)^{2}d(n-1)(n+d),$

and finally assuming Xavier initialization

$\displaystyle\mathbb{E}\left\|\frac{\partial{\bf S}^{\ell}}{\partial{\bf
W}^{Q,\ell}}\right\|^{2}_{F}$
$\displaystyle=\frac{\sigma^{2}_{v}\sigma^{2}_{k}d^{2}}{dn^{2}}\cdot\mathbb{E}\left[\|{\bf
X}^{\ell}\|^{2}_{F}\cdot\|({\bf X}^{\ell})^{\top}{\bf
X}^{\ell}-n\bar{\bm{x}}^{\ell}(\bar{\bm{x}}^{\ell})^{\top}\|^{2}_{F}\right]$
$\displaystyle=\sigma_{x}^{6}\frac{n-1}{n}(1-\rho)^{2}d(n+d).$

####  A.3 Forward Pass: Proofs of Lemma 3.2 and 3.3

First, we characterize the evolution of the correlations between tokens ${\bf
X}_{k},{\bf X}_{k^{\prime}}$ with depth, under the assumptions of Theorem 3.3,
namely uniform-attention assumption, and the adoption of a linear activation.

###### Lemma A.4 (Expectation of Linear Layers).

Let ${\bf D}={\bf X}{\bf W}$, where ${\bf W}\in\mathbb{R}^{d\times d}$ is a
random matrix with i.i.d random entries with variance $\sigma^{2}=\frac{1}{d}$
and ${\bf X}\in\mathbb{R}^{n\times d}$ is a fixed matrix: $\mathbb{E}[{\bf
D}_{kj}{\bf D}_{k^{\prime}j}]=\frac{1}{d}\langle{\bf X}_{k},{\bf
X}_{k^{\prime}}\rangle$

Note that by summing over the indexes, Lemma A.4 implies:

$\displaystyle\mathbb{E}\left\lVert{\bf
D}\right\rVert_{F}^{2}=\mathbb{E}\left\lVert{\bf X}\right\rVert_{F}^{2}$
$\displaystyle\mathbb{E}C({\bf D})=\mathbb{E}C({\bf X}).$

###### Proof.

$\mathbb{E}[{\bf D}_{kj}{\bf D}_{k^{\prime}j}]=\sum_{zz^{\prime}}{\bf
X}_{kz}{\bf X}_{k^{\prime}z^{\prime}}\mathbb{E}[{\bf W}_{zj}{\bf
W}_{z^{\prime}j}]=\sigma^{2}\sum_{z}{\bf X}_{kz}{\bf
X}_{k^{\prime}z}=\frac{1}{d}\langle{\bf X}_{k},{\bf X}_{k^{\prime}}\rangle.$

∎

###### Lemma A.5 (Expectation of skip connection).

Let ${\bf A},{\bf B}\in\mathbb{R}^{p\times q}$. Let ${\bf D}:=\alpha{\bf
A}+{\bf B}$ with $\mathbb{E}[{\bf A}|{\bf B}]=\bm{0}$ and
$\alpha\in\mathbb{R}$. Then: $\mathbb{E}\left[{\bf D}_{ij}{\bf
D}_{i^{\prime}j}\right]=\alpha^{2}\mathbb{E}[{\bf A}_{ij}{\bf
A}_{i^{\prime}j}]+\mathbb{E}[{\bf B}_{ij}{\bf B}_{ij^{\prime}}]$ (42) holds
for all $i,i^{\prime}\in[p],j\in[q]$.

Note that by summing over the indexes, Lemma A.5 implies:

$\displaystyle\mathbb{E}\left\lVert{\bf
D}\right\rVert_{F}^{2}=\alpha^{2}\mathbb{E}\left\lVert{\bf
A}\right\rVert_{F}^{2}+\mathbb{E}\left\lVert{\bf B}\right\rVert_{F}^{2}$
$\displaystyle\mathbb{E}C({\bf D})=\alpha^{2}\mathbb{E}C({\bf
A})+\mathbb{E}C({\bf B}).$

###### Proof.

$\displaystyle\mathbb{E}[{\bf D}_{ij}{\bf D}_{i^{\prime}j}]$
$\displaystyle=\mathbb{E}\left[(\alpha{\bf A}_{ij}+{\bf B}_{ij})(\alpha{\bf
A}_{i^{\prime}j}+{\bf B}_{i^{\prime}j})\right]$
$\displaystyle=\mathbb{E}\left[\alpha^{2}{\bf A}_{ij}{\bf
A}_{i^{\prime}j}+\alpha{\bf A}_{ij}{\bf B}_{i^{\prime}j}+\alpha{\bf
A}_{i^{\prime}j}{\bf B}_{ij}+{\bf B}_{ij}{\bf B}_{i^{\prime}j}\right]$
$\displaystyle=\alpha^{2}\mathbb{E}\left[{\bf A}_{kj}{\bf
A}_{i^{\prime}j}\right]+\mathbb{E}\left[{\bf B}_{ij}{\bf
B}_{i^{\prime},j}\right],$

where using iterated expectations $\alpha\mathbb{E}[{\bf A}_{i^{\prime}j}{\bf
B}_{ij}]=\alpha\mathbb{E}[\mathbb{E}[{\bf A}_{i^{\prime}j}|{\bf B}]{\bf
B}_{ij}]]=0$ and identically $\alpha\mathbb{E}[{\bf A}_{ij}{\bf
B}_{i^{\prime}j}]=0$. ∎

###### Lemma A.6 (Expectation of Attention Layers).

Under the uniform-attention assumption: $\mathbb{E}[{\bf S}_{kj}{\bf
S}_{k^{\prime}j}]=\frac{1}{d_{v}n^{2}}\mathbb{E}C({\bf X}).$

In this case, by summing over the indexes we have that:

$\displaystyle\mathbb{E}\left\lVert{\bf
S}\right\rVert_{F}^{2}=\frac{\mathbb{E}C({\bf X})}{n}$
$\displaystyle\mathbb{E}C({\bf S})=\mathbb{E}C({\bf X}).$

###### Proof.

Note that under the uniform-attention assumption:

${\bf S}_{kj}=\frac{1}{n}\left(\bm{1}_{n\times n}{\bf X}{\bf
W}^{V}\right)_{kj}=\frac{1}{n}\sum_{zi}{\bf X}_{zi}{\bf W}^{V}_{ij}.$

Hence, using the fact that the weights are i.i.d with variance
$\sigma_{v}^{2}=\frac{1}{d_{v}}$:

$\displaystyle\mathbb{E}[{\bf S}_{kj}{\bf
S}_{k^{\prime}j}]=\frac{\sigma_{v}^{2}}{n^{2}}\sum_{z,z^{\prime}}\sum_{i}\mathbb{E}[{\bf
X}_{zi}{\bf
X}_{z^{\prime}i}]=\frac{1}{d_{v}n^{2}}\sum_{k,k^{\prime}}\langle{\bf
X}_{z},{\bf X}_{z^{\prime}}\rangle=\frac{1}{d_{v}n^{2}}\mathbb{E}C({\bf X}).$

∎

See 3.2

###### Proof.

First, note that for the residual blocks we have that
$\mathbb{E}[Y^{\ell}_{kj}|Z^{\ell}_{k^{\prime}j}]=0$ due to the independence
assumption on the feedforward weights, and similarly
$\mathbb{E}[S^{\ell}_{kj}|X^{\ell}_{k^{\prime}j}]=0$. Hence, we can use Lemma
A.5 in both the skip connections of the Transformer architecture. Therefore,
using Lemma A.5 (skip), Lemma A.4 (linear) and Lemma A.6 (attention):

$\displaystyle\mathbb{E}[C({\bf X}^{\ell+1})]$
$\displaystyle{=}\alpha_{2}^{2}\mathbb{E}C({\bf Y}^{\ell})+\mathbb{E}C({\bf
Z}^{\ell})$ $\displaystyle{=}\alpha_{2}^{2}\mathbb{E}C({\bf
Z}^{\ell})+\mathbb{E}C({\bf Z}^{\ell})$
$\displaystyle=(\alpha_{2}^{2}+1)\mathbb{E}C({\bf Z}^{\ell})$
$\displaystyle{=}(\alpha_{2}^{2}+1)\left(\alpha_{1}^{2}\mathbb{E}C({\bf
S}^{\ell})+\mathbb{E}C({\bf X}^{\ell})\right)$
$\displaystyle{=}(\alpha_{2}^{2}+1)(\alpha_{1}^{2}+1)\mathbb{E}[C({\bf
X}^{\ell})]$
$\displaystyle{=}(\alpha_{2}^{2}+1)^{\ell+1}(\alpha_{1}^{2}+1)^{\ell+1}C({\bf
X}),$

where in the last step we have unrolled the recursion until the input layer.

For the limit as $L\to\infty$, simply note that:

$\lim_{L\to\infty}\left(\frac{\tilde{\alpha}_{i}}{L}+1\right)^{L}=\text{e}^{\tilde{\alpha}_{i}},$

with $i\in\\{1,2\\}$. ∎

Now we are ready to re-state and prove Lemma 3.3.

See 3.3

###### Proof.

The proof is in the same spirit as Lemma 3.2 but slightly more involved.
Again, using Lemma A.5 in both the skip connections of the Transfomer
architecture. Therefore, using Lemma A.5 (skip), Lemma A.4 (linear) and Lemma
A.6 (attention):

$\displaystyle\mathbb{E}[||{\bf X}^{\ell+1}||_{F}^{2}]$
$\displaystyle{=}\alpha_{2}^{2}\mathbb{E}||{\bf
Y}^{\ell}||_{F}^{2}+\mathbb{E}||{\bf Z}^{\ell}||_{F}^{2}$
$\displaystyle{=}(\alpha_{2}^{2}+1)\mathbb{E}||{\bf Z}^{\ell}||_{F}^{2}$
$\displaystyle{=}(\alpha_{2}^{2}+1)\left(\alpha_{1}^{2}\mathbb{E}[||{\bf
S}^{\ell}||_{F}^{2}]+\mathbb{E}[||{\bf X}^{\ell}||_{F}^{2}]\right)$
$\displaystyle{=}(\alpha_{2}^{2}+1)\left(\frac{\alpha_{1}^{2}}{n}\mathbb{E}[C({\bf
X}^{\ell})]+\mathbb{E}[||{\bf X}^{\ell}||_{F}^{2}]\right)$
$\displaystyle=(\alpha_{2}^{2}+1)\frac{\alpha_{1}^{2}}{n}\mathbb{E}[C({\bf
X}^{\ell})]+(\alpha_{2}^{2}+1)\mathbb{E}[||{\bf X}^{\ell}||_{F}^{2}]$
$\displaystyle{=}(\alpha_{2}^{2}+1)^{\ell+1}\alpha_{1}^{2}(\alpha_{1}^{2}+1)^{\ell}\frac{C({\bf
X})}{n}+(\alpha_{2}^{2}+1)\mathbb{E}[||{\bf X}^{\ell}||_{F}^{2}]$
$\displaystyle{=}(\alpha_{2}^{2}+1)^{\ell+1}\alpha_{1}^{2}\frac{C({\bf
X})}{n}\sum_{k=0}^{\ell}(\alpha_{1}^{2}+1)^{k}+(\alpha_{2}^{2}+1)^{\ell+1}||{\bf
X}||_{F}^{2},$

where in the second to last step we have used Lemma 3.2 and in the last step
we have unrolled the recursion for $\left\lVert{\bf
X}^{\ell}\right\rVert_{F}^{2}$ until the input layer.

For the second part, we now show that for a network of $L$ layers, the choice
$\alpha_{1}^{2}=\frac{\tilde{\alpha}_{1}}{L}$ and
$\alpha_{2}^{2}=\frac{\tilde{\alpha}_{2}}{L}$ stabilizes the norm of the
activations in the forward pass. Using the product law for the limits, we can
study the converges of $(\alpha_{2}^{2}+1)^{\ell}$ and
$\alpha_{1}\sum_{k=0}^{\ell}(\alpha_{1}^{2}+1)^{k}$ separately.

Let $i\in\\{1,2\\}$. For the latter term we have that:

$\displaystyle\lim_{L\to\infty}\frac{\tilde{\alpha}_{i}}{L}\sum_{l=0}^{L-1}\left(1+\frac{\tilde{\alpha}_{i}}{L}\right)^{\ell}$
$\displaystyle=\lim_{L\to\infty}\frac{\tilde{\alpha}_{i}}{L}\frac{1-\left(1+\frac{\tilde{\alpha}_{i}}{L}\right)^{\ell}}{1-1-\frac{\tilde{\alpha}_{i}}{L}}$
$\displaystyle=\lim_{L\to\infty}-1+\left(1+\frac{\tilde{\alpha}_{i}}{L}\right)^{\ell}$
$\displaystyle=\text{e}^{\tilde{\alpha}_{i}}-1,$

while for the former term we have that
$\lim_{L\to\infty}(\frac{\tilde{\alpha}_{i}}{L}+1)^{\ell}=\text{e}^{\tilde{\alpha}_{i}}$.
Hence, the norm of the representations converges to:

$\lim_{L\to\infty}\mathbb{E}[||{\bf
X}^{\ell}||_{F}^{2}]=\text{e}^{\tilde{\alpha}_{2}}(\text{e}^{\tilde{\alpha}_{1}}-1)\frac{C({\bf
X})}{n}+\text{e}^{\tilde{\alpha}_{2}}||{\bf X}||_{F}^{2}.$

The final results as stated in the theorem hold because of the following:

Remark: note that $C({\bf X})=\sum_{k,k^{\prime}}\sum_{j}{\bf X}_{kj}{\bf
X}_{kj^{\prime}}=\sum_{j}(\sum_{k,k^{\prime}}{\bf X}_{kj}{\bf
X}_{k^{\prime}j})=\sum_{j}(\sum_{k}{\bf
X}_{kj})^{2}=n^{2}\left\lVert\bar{\bm{x}}\right\rVert^{2}$. ∎

####  A.4 Proof of Theorem 3.2: Correlations are Preserved under Residual
Scaling

See 3.2

###### Proof.

Due to the rotational symmetries of the Gaussian random matrices, if the input
tokens have the same norm, then the expected norm at layer $\ell\in[L]$ is
also the same across the token’s representations. Hence, we can write
$\mathbb{E}\left\lVert{\bf
X}^{\ell}\right\rVert_{F}^{2}=n\mathbb{E}\left\lVert\bm{x}^{\ell}\right\rVert^{2}$,
where $\left\lVert\bm{x}^{\ell}\right\rVert^{2}$ is the norm of every token at
layer $\ell$. Furthermore, by definition of our correlation coefficient
$\rho^{l}_{kk^{\prime}}$, we have that $\mathbb{E}\langle{\bf
X}^{\ell}_{k},{\bf
X}^{\ell}_{k^{\prime}}\rangle=\rho^{\ell}_{kk^{\prime}}\mathbb{E}\left\lVert\bm{x}^{\ell}\right\rVert^{2}$.
By summing over the indexes $k,k^{\prime}$, we can expand the relation as:

$\underbrace{\sum_{k,k^{\prime}}\mathbb{E}\langle{\bf X}^{\ell}_{k},{\bf
X}^{\ell}_{k^{\prime}}\rangle}_{\mathbb{E}C({\bf
X})}=\sum_{k,k^{\prime}}\rho^{\ell}_{kk^{\prime}}\mathbb{E}\left\lVert\bm{x}^{\ell}\right\rVert^{2}=(n+\sum_{k\neq
k^{\prime}}\rho^{\ell}_{k,k^{\prime}})\mathbb{E}\left\lVert\bm{x}^{\ell}\right\rVert^{2}=\underbrace{n\mathbb{E}\left\lVert\bm{x}^{\ell}\right\rVert^{2}}_{\mathbb{E}\left\lVert{\bf
X}^{\ell}\right\rVert_{F}^{2}}(1+(n-1)\rho^{\ell}).$

By solving for $\rho^{\ell}$, we have that:

$\rho^{\ell}=\frac{\mathbb{E}C({\bf X}^{\ell})}{(n-1)\mathbb{E}\left\lVert{\bf
X}^{\ell}\right\rVert^{2}}-\frac{1}{n-1}.$

Now we plug in the expressions for $\mathbb{E}C({\bf X}^{\ell})$ and
$\mathbb{E}\left\lVert{\bf X}^{\ell}\right\rVert^{2}$ with the aid of Lemma
3.2 and Lemma 3.3, respectively. Finally, by taking the limits with respect to
$L$, we get the desired result. ∎

####  A.5 Motivation for Assumption 3.1

We motivate here the following assumption, stated in the main paper. This
assumption is crucial to compute expectations involving the softmax function.

See 3.1

###### Theoretical analysis.

We first show that this assumption holds when taking $d_{k}$ to infinity,
keeping $d_{v}$ fixed.

###### Lemma A.7.

Consider initializing each entry of $\mathbf{W}^{Q}\in\mathbb{R}^{d_{v}\times
d_{k}}$ and $\mathbf{W}^{K}\in\mathbb{R}^{d_{v}\times d_{k}}$ independently
with variance $\sigma^{2}_{k}=2/(d_{v}+d_{k})$ — i.e. Glorot initialization
[Glorot and Bengio, 2010b]. Let ${\bf M}=\frac{1}{\sqrt{d_{k}}}{\bf
X}^{\ell}{\bf W}^{Q,\ell}{{\bf W}^{K,\ell}}^{\top}{{\bf X}^{\ell}}^{\top}$;
for any $(i,j)\in[n]\times[n]$ we have $\mathbb{E}[{\bf M}_{i,j}\ |\ {\bf
X}]=0,\qquad\mathbb{E}[{\bf M}_{i,j}^{2}\ |\ {\bf
X}]=\sigma_{k}^{4}\cdot\|{\bf X}_{i,:}\|^{2}\cdot\|{\bf X}_{j,:}\|^{2}.$ (43)
While keeping $d_{v}<\infty$ fixed, taking $d_{k}$ to infinity yields
$\mathbb{E}[{\bf M}_{i,j}^{2}\ |\ {\bf
X}]=\mathcal{O}\left(\frac{1}{d_{k}^{2}}\right).$ (44) In other words, ${\bf
M}$ converges to $\bm{0}_{n\times n}$ in $L^{2}$ as $d_{k}\to\infty$.

###### Proof.

First, note that

$\displaystyle{\bf
M}_{i,j}=\frac{1}{\sqrt{d_{k}}}\sum_{a,c=1}^{d_{v}}\sum_{b=1}^{d_{k}}{\bf
X}_{i,a}\mathbf{W}^{Q}_{a,b}\mathbf{W}^{K}_{c,b}{\bf X}_{j,c}.$

Since $\mathbf{W}^{Q}$ is independent from $\mathbf{W}^{K}$ at initialization,
$\mathbb{E}[{\bf M}_{i,j}\ |\ {\bf X}]=0$. Next, we compute

$\displaystyle\mathbb{E}[{\bf M}_{i,j}^{2}]$
$\displaystyle=\frac{1}{d_{k}}\sum_{a,c,a^{\prime},c^{\prime}=1}^{d_{v}}\sum_{b,b^{\prime}=1}^{d_{k}}{\bf
X}_{i,a}{\bf X}_{i,a^{\prime}}{\bf X}_{j,c}{\bf
X}_{j,c^{\prime}}\mathbb{E}\left[\mathbf{W}^{Q}_{a,b}\mathbf{W}^{Q}_{a^{\prime},b^{\prime}}\mathbf{W}^{K}_{c,b}\mathbf{W}^{K}_{c^{\prime},b^{\prime}}\right]$
$\displaystyle=\frac{1}{d_{k}}\sum_{a,c,a^{\prime},c^{\prime}=1}^{d_{v}}\sum_{b,b^{\prime}=1}^{d_{k}}{\bf
X}_{i,a}{\bf X}_{i,a^{\prime}}{\bf X}_{j,c}{\bf
X}_{j,c^{\prime}}\mathbb{E}\left[\mathbf{W}^{Q}_{a,b}\mathbf{W}^{Q}_{a^{\prime},b^{\prime}}\right]\mathbb{E}\left[\mathbf{W}^{K}_{c,b}\mathbf{W}^{K}_{c^{\prime},b^{\prime}}\right]$
$\displaystyle=\frac{\sigma_{k}^{4}}{d_{k}}\sum_{a,c=1}^{d_{v}}\sum_{b=1}^{d_{k}}{\bf
X}_{i,a}^{2}{\bf X}_{j,c}^{2}$ $\displaystyle=\sigma_{k}^{4}\|{\bf
X}_{i,:}\|^{2}\|{\bf X}_{j,:}\|^{2}.$

This concludes the proof. ∎

The following classical result implies almost sure convergence of the softmax
matrix as $d_{k}\to\infty$.

###### Lemma A.8 (Borel-Cantelli).

Let $(X_{i})$ be a sequence of random variables. If for any $\epsilon>0$

$\sum_{i=0}^{\infty}\mathbb{P}[|X_{i}-X|>\epsilon]<\infty,$

then $X_{i}$ converges to $X$ almost surely444That is,
$\lim_{i\to\infty}X_{i}(\omega)=X(\omega)$ for almost every $\omega\in\Omega$
(i.e. with probability one)..

###### Theorem A.2 (Almost-sure convergence).

Consider initializing each entry of $\mathbf{W}^{Q}\in\mathbb{R}^{d_{v}\times
d_{k}}$ and $\mathbf{W}^{K}\in\mathbb{R}^{d_{v}\times d_{k}}$ independently
with variance $\sigma^{2}_{k}=2/(d_{v}+d_{k})$ — i.e. Glorot initialization
[Glorot and Bengio, 2010b]. Let $d_{v}<\infty$ be fixed, as $d_{k}\to\infty$
we have that, for any ${\bf X}$, ${\bf
A}:=\operatorname{softmax}\left(\frac{1}{\sqrt{d_{k}}}{\bf X}{\bf W}^{Q}{{\bf
W}^{K}}^{\top}{{\bf X}}^{\top}\right)\stackrel{{\scriptstyle
a.s.}}{{\to}}\frac{1}{n}\bm{1}_{n\times n}$ and $\frac{\partial{\bf
A}}{\partial{\bf M}}\stackrel{{\scriptstyle a.s.}}{{\to}}\frac{1}{n}{\bf
I}_{n}\otimes\left({\bf I}_{n}-\frac{1}{n}\bm{1}_{n\times n}\right).$

###### Proof.

Thanks to Lemma A.7 and Markov Inequality, we have fast convergence in
probability: for any fixed ${\bf X}$,

$\mathbb{P}[|{\bf M}_{i,j}|>\epsilon]\leq\frac{\mathbb{E}[{\bf
M}^{2}_{i,j}]}{\epsilon^{2}}\leq\frac{C_{\epsilon}}{d_{k}^{2}}.$

Borel Cantelli then directly yields almost sure convergence of ${\bf M}$ to
$\bm{0}_{n\times n}$ as $d_{k}\to\infty$. Next, note that both ${\bf A}$ and
$\frac{\partial{\bf A}}{\partial{\bf M}}$ are continuous functions of ${\bf
A}$, hence we can apply standard continuity event-per-event. For almost every
$\omega\in\Omega$,

$\lim_{d_{k}\to\infty}{\bf A}({\bf A}(\omega))={\bf
A}\left(\lim_{d_{k}\to\infty}{\bf A}(\omega)\right)={\bf A}(\bm{0}_{n\times
n})=\frac{1}{n}\bm{1}_{n\times n}.$

Hence ${\bf A}\to\frac{1}{n}\bm{1}_{n\times n}$ almost surely. This can also
be seen as a simple application of the continuous mapping theorem. The same
reasoning yields almost sure convergence of

$\frac{\partial{\bf A}}{\partial{\bf
M}}=\operatorname{blockdiag}\Bigg{(}\operatorname{diag}({\bf A}_{i:})-{\bf
A}_{i:}{\bf A}_{i:}^{\top}\Bigg{)},$

to the corresponding limiting quantity. ∎

###### Empirical analysis.

Figure 8: Evolution of $\frac{1}{n^{2}}||{\bf
A}^{\ell}-\frac{1}{n}{\bf{1}}_{n\times n}||_{F}^{2}$ as a function of $d_{k}$
for $d_{v}$ fixed at 100.

We empirically assess the validity of Assumption 3.1 and of its theoretical
justification by performing the following experiments: for a range of
increasing values of $d_{k}$, we compute ${\bf A}$ and we calculate
$\frac{1}{n^{2}}||{\bf A}^{\ell}-\frac{1}{n}{\bf{1}}_{n\times n}||_{F}^{2}$,
i.e. its average (entry-wise) distance from a uniform matrix with entries all
equal to $1/n$. For each value of $d_{k}$, we repeat this calculation 200
times, each time with different random weight matrices. Fig. 8 displays how
the $||{\bf A}-\frac{1}{n}{\bf{1}}_{n\times n}||_{F}^{2}$ averaged over 200
runs, tends to zero with a trend inversely proportional to $d_{k}^{2}$, as
predicted by our theoretical analysis.

###  Appendix B Additional Results

####  B.1 On the Roles of the $1/\sqrt{L}$-Scaling of the Residuals and Layer
Normalization

We present some additional results on the propagation of the norm and the
correlations in Figure 9. In particular, we empirically show that, with an
adequate depth-dependent residual scaling, the norm and the correlation are
stabilized, even for very deep networks. Furthermore, we demonstrate the
propagation of the correlation and the gradient norms for the PRE-LN
configuration in Figure 10. As also hinted in the main text, in Figure 4, the
increase in correlation with depth for PRE-LN is much less wild. This also
results in better stabilized gradients for the queries and keys’ parameters.
We also observe the opposite trend for the gradients of the values, in
relation to the POST-LN case in Figure 3. We speculate that this different
dependence, along with the better preserved correlation, is the main reason
PRE-LN configured Transformers have been shown to scale better with depth. We
plan to investigate this dependence more in future work.

Figure 9: (Left) Propagation of the Frobenius norm of the input sequence;
(Right) Propagation of the average token correlation. Figure 10: Same figure
as 3 but with a PRE-LN architecture. The correlation at depth 0 originates
from the correlations in the randomly initialized tokens’ embeddings and
positional encodings.

####  B.2 Further Empirical Assessment of Assumption 3.1

Here, we empirically test the accuracy and limitations of the uniform-
attention assumption.

For the empirical verification of Assumption 3.1 in the forward pass analysis,
we plot the density of the norm of the representations for only-encoder
Transformers of increasing depth. The results are shown in Fig 11. Note that
when the standard deviation of the input is set to $1/\sqrt{d}$, then the
uniform-attention assumption provide an excellent approximation to the common
Xavier-initialization. On the contrary, we observe a deviation when the
standard deviation of the input is increased. Also, note how as the depth
increases, the distribution becomes more heavy-tailed. This heavy-tailedness
was recently formally shown for standard MLPs with and without ReLU activation
[Noci et al., 2021, Zavatone-Veth and Pehlevan, 2021].

Figure 11: Density plots for $\left\lVert{\bf X}^{\ell}\right\rVert_{F}^{2}$
for Transformers of depths $L$ from $1$ to $10$. The input ${\bf X}$ contains
i.i.d Gaussian entries, simulating an embedding layer. We set
$d:=d_{v}=d_{q}=30$. The empirical mean at $L=10$ is highlighted in a vertical
dashed red line, while the theoretical mean (Lemma 3.3) is a dashed blue line.
The densities are estimated by sampling 1000 times the weights of the network.
(Left): we adopt the uniform-attention. The standard deviation of the input is
set to $1/\sqrt{d}$. (Center): Same, but removing the uniform-attention
assumption. (Right): We remove the uniform-attention assumption, and set the
standard deviation of the input to $1$.

For the verification of the assumption in the backward pass, we additionally
show in Fig. 12 how the norm of the gradients w.r.t queries and keys depends
on the hidden dimension, the sequence length, the input correlation and the
input variance. _Ground-truth_ gradients are calculated with automatic
differentiation, and they are compared with our theoretical results based on
Assumption 3.1. As shown in Fig.12, our theoretical predictions show a very
good agreement with the true gradients. Again, we notice that the smaller the
values of the input standard deviation the tighter the agreement of the theory
with the simulations. Intuitively, a higher input variance causes the argument
of the softmax to have a large range of values. This in turn causes a
deviation from the uniform distribution (i.e. maximum entropy), towards the
distribution of minimum entropy (a Delta Dirac, corresponding to attending to
only one token).

####  B.3 Empirical Verification of the Gradient Analysis of Section 3.3

Finally, in Figures 13 and 14 we show the dependence of the norm of the
gradients for the keys and values based on the parameters of the architecture
and the task-specific parameters. Figure 13 illustrates the true dependence
and Figure 14 the one expected by the theory based on our assumptions. In
short, the main takeaways are the following.

  * •

As the correlation between the tokens increases ($x$-axis in the global plot),
the norm of the gradients of the queries quickly diminishes compared to the
one of the values.

  * •

The dependence on the variance of the input $\sigma_{x}^{2}$ is different
($y$-axis in the global plot), being linear for the values and cubic for the
queries. This highlights the importance of a stabilized forward pass and
provides another explanation regarding the successful use of layer norm in
Transformers.

  * •

The dependence on $n$ ($x$-axis in each subplot) and $d$ ($y$-axis in each
subplot) is more complicated, also being a function of the correlation $\rho$
(compare the first column where $\rho=0$ to the rest).

Figure 12: Empirical comparison of our theoretical findings. We sample, as
aforementioned, the tokens according to a zero-mean Gaussian distribution,
while varying the hidden dimension, sequence length, input correlation and
input variance. Results are averaged over 20 runs. Figure 13: Log ratio of the
norm of the gradients for the queries compared to those of the values for
varying values of embedding dimension, sequence length, cosine of the tokens
angle and standard deviation. Figure 14: Log ratio of the norm of the
gradients of the queries as expected by the theory, compared to those of the
values for varying values of embedding dimension, sequence length, cosine of
the tokens angle and standard deviation. We use Equations (17) and (18).

###  Appendix C Experimental Setup

Here we provide more details regarding the experimental setup.

####  C.1 Toy Example

In Figure 5, we focus on a toy example where the task is to reverse a sequence
of tokens. More specifically, given a sequence of $20$ numbers in the range
$0-9$, we predict the same tokens in the inverted order. We use an embedding
layer of size 16, initializes with variance 1, and sinusoidal positional
encodings to initially embed the input. We use a 5-layer POST-LN Transformer
encoder model, with a single head attention operation and a two-layer feed-
forward layer with a ReLU nonlinearity. We use residual scaling in this case
equal to $\alpha_{1}=\alpha_{2}=1$. We train using Adam with betas parameters
$(0.9,0.999)$, learning rate $0.01$ and weight decay 0.

####  C.2 Translation Task

Introducing an inverse temperature scaling $\tau$ inside the softmax, modifies
the attention operation to

${\bf S}^{\ell}:=\text{softmax}\left(\frac{\tau}{\sqrt{d_{k}}}{\bf
X}^{\ell}{\bf W}^{Q}\left({\bf X}^{\ell}{\bf W}^{K}\right)^{\top}\right){\bf
X}^{\ell}{\bf W}^{V}.$ (45)

Then the gradient of the queries and keys parameters are directly scaled by
this temperature value, following the same proof as for Eq. (7). We choose a
temperature value of $\tau_{\text{final}}=8.5$ to match the gradient norms of
the values and queries as in Equations. (17) and (18). Doing so, we assume a
constant small correlation between tokens (also empirically verified in Fig.
15) and set the sequence length $n$ to the average found in our training
dataset. Due to instabilities in training, we use warm-up on this temperature
value. In short:

$\tau=\tau_{\text{final}}\cdot\text{max}(1,\frac{\text{step}}{\text{steps}_{\text{warmup}}}),$

with ‘$\text{steps}_{\text{warmup}}=1000$’ and ‘step’ the current training
step.

We base our implementation on fairseq [Ott et al., 2019]. For the
hyperparameter configuration, we mostly rely on the extensive search already
done in fairseq [Ott et al., 2019] and Liu et al. [2020]. The final used
parameters are exhibited in Table 1. For the final evaluation, we use the
best-performing model on the left-out validation set. We apply weight decay as
in Loshchilov and Hutter [2017] for both SGD and Adam.

| Hyperparameters | Value
---|---|---
| Max tokens | 4096
Label smoothing | 0.1
clip-norm | 0.0
General Dropout | 0.3
Attention Dropout | 0.1
ReLU Dropout | 0.1
Hidden size | 512
FFN inner hidden size | 2048
Attention Heads | 4
Adam | Learning rate | $7\epsilon^{-4}$
Learning rate scheduler | inverse sqrt
Warm-up updates | 6000
Warm-up init learning rate | 1e-7
Adam $(\beta_{1},\beta_{2})$ | (0.9, 0.98)
Training updates | 100K
Weight decay | 0.0001
SGD | Learning rate | $2\epsilon^{-2}$
Learning rate scheduler | step
Step scheduler $\gamma$ | 0.1
Step scheduler update steps | [100K, 200K]
Training updates | 250K
Weight decay | 0.001
Table 1: Hyperparameters for the IWSLT’14 De-En translation task.

Finally, in Figure 15 we display the evolution of correlations, residual
scaling, and norm of the activations, with depth, for our best trained model.
The residual scaling $\alpha_{1},\alpha_{2}$ are trainable parameters. This
enables them to weight differently the residual branches if deemed necessary.
Although these values increase during training, the correlation between the
tokens does not significantly increase, which as implied by our main results,
allows efficient propagation of the gradients. The norm of the propagated
forward signal tends to slightly increase with depth.

Figure 15: Evolution of the cosine of the angles, the trained residual
$\alpha_{1},\alpha_{2}$ and the activation norm throughout our training.