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

Original: Author(
Revised: Author(

Original: Abstract
Revised: Abstract

Original: Transformers have achieved remarkable success in several domains, ranging fromnatural language processing to computer vision.
Revised: Transformers have achieved remarkable success in several domains, ranging fromnatural language processing to computer vision.

Original: Nevertheless, it has been recentlyshown that stacking self-attention layers — the distinctive architectural componentof Transformers — can result in rank collapse of the tokens’ representations atinitialization.
Revised: Nevertheless, it has been recentlyshown that stacking self-attention layers — the distinctive architectural componentof Transformers — can result in rank collapse of the tokens’ representations atinitialization.

Original: The question of if and how rank collapse affects training is stilllargely unanswered, and its investigation is necessary for a more comprehensiveunderstanding of this architecture.
Revised: The question of if and how rank collapse affects training is stilllargely unanswered, and its investigation is necessary for a more comprehensiveunderstanding of this architecture.

Original: In this work, we shed new light on the causesand the effects of this phenomenon.
Revised: In this work, we shed new light on the causesand the effects of this phenomenon.

Original: First, we show that rank collapse of the tokens’representations hinders training by causing the gradients of the queries and keysto vanish at initialization.
Revised: First, we show that rank collapse of the tokens’representations hinders training by causing the gradients of the queries and keysto vanish at initialization.

Original: Furthermore, we provide a thorough description ofthe origin of rank collapse and discuss how to prevent it via an appropriate depthdependent scaling of the residual branches.
Revised: Furthermore, we provide a thorough description ofthe origin of rank collapse and discuss how to prevent it via an appropriate depthdependent scaling of the residual branches.

Original: Finally, our analysis unveils that specificarchitectural hyperparameters affect the gradients of queries and values differently,leading to disproportionate gradient norms.
Revised: Finally, our analysis unveils that specificarchitectural hyperparameters affect the gradients of queries and values differently,leading to disproportionate gradient norms.

Original: This suggests an explanation for thewidespread use of adaptive methods for Transformers’ optimization.
Revised: This suggests an explanation for thewidespread use of adaptive methods for Transformers’ optimization.

Original:  Introduction
Revised:  Introduction

Original: Since its first appearance in Vaswani et al.
Revised: Since its first appearance in Vaswani et al.

Original: [2017] the Transformer architecture has revolutionizedthe field of Natural Language Processing (NLP) achieving remarkable success in tasks such astext classification [Yang et al., 2019], machine translation [Conneau and Lample, 2019], readingcomprehension [Brown et al., 2020] and question answering [Raffel et al., 2019] among others.
Revised: [2017] the Transformer architecture has revolutionizedthe field of Natural Language Processing (NLP) achieving remarkable success in tasks such astext classification [Yang et al., 2019], machine translation [Conneau and Lample, 2019], readingcomprehension [Brown et al., 2020] and question answering [Raffel et al., 2019] among others.

Original: 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.
Revised: 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.

Original: , 2022],further popularizing it outside NLP.
Revised: , 2022],further popularizing it outside NLP.

Original: The Transformer operates on inputs comprising a sequence of tokens.
Revised: The Transformer operates on inputs comprising a sequence of tokens.

Original: At its core, it relies on stackedattention layers, which compute a measure of relevance for the whole sequence by assigning tokenwise importance weights — obtained by matrix multiplication of the queries and keys , and finallynormalized with the softmax function.
Revised: At its core, it relies on stackedattention layers, which compute a measure of relevance for the whole sequence by assigning tokenwise importance weights — obtained by matrix multiplication of the queries and keys , and finallynormalized with the softmax function.

Original: The output of an attention layer is then a linear combinationof the importance weights and the so-called values .
Revised: The output of an attention layer is then a linear combinationof the importance weights and the so-called values .

Original: Then, the architecture includes fully-connectedsub-layers, residual connections [He et al., 2016], and layer normalization (LN), as illustrated in
Revised: Then, the architecture includes fully-connectedsub-layers, residual connections [He et al., 2016], and layer normalization (LN), as illustrated in

Original: Fig.
Revised: Fig.

Original: In the absence of residual connections, Dong et al.
Revised: In the absence of residual connections, Dong et al.

Original: [2021] proved that at initialization the rank of thesequence representation collapses doubly exponentially with depth, and both layer normalization andfully connected layers can only partially alleviate the speed of degeneracy.
Revised: [2021] proved that at initialization the rank of thesequence representation collapses doubly exponentially with depth, and both layer normalization andfully connected layers can only partially alleviate the speed of degeneracy.

Original: Under rank collapse , themodel does not distinguish between representations of different tokens, which are perfectly alignedin feature space at initialization.
Revised: Under rank collapse , themodel does not distinguish between representations of different tokens, which are perfectly alignedin feature space at initialization.

Original: Crucially, the precise implications of rank collapse in Transformersare not fully understood.
Revised: Crucially, the precise implications of rank collapse in Transformersare not fully understood.

Original: R e s i
Revised: R e s i

Original: d u a l s c a li n g : 1 = 2 =
Revised: d u a l s c a li n g : 1 = 2 =

Original: Depth:
Revised: Depth:

Original: Depth: Depth:
Revised: Depth: Depth:

Original: Depth:48 random prediction
Revised: Depth:48 random prediction

Original: Training loss
Revised: Training loss

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: R e s i
Revised: R e s i

Original: d u a l s c a li n g : 1 = 2 = 1 L
Revised: d u a l s c a li n g : 1 = 2 = 1 L

Original: [Figure]
Revised: [Figure]

Original: T o k e n s a v e r a g e c o s i n e a n g l e
Revised: T o k e n s a v e r a g e c o s i n e a n g l e

Original: [Figure]
Revised: [Figure]

Original: 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 bycausing vanishingly small gradients of the queries and keys’ parameter matrices.
Revised: 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 bycausing vanishingly small gradients of the queries and keys’ parameter matrices.

Original: This problemseverely diminishes the capabilities of the model to learn meaningful attention weights and is furtherexacerbated in very deep networks, where the rank deficiency — and hence the vanishing gradientproblem of the queries and keys — affects several layers (see Fig. 1).
Revised: This problemseverely diminishes the capabilities of the model to learn meaningful attention weights and is furtherexacerbated in very deep networks, where the rank deficiency — and hence the vanishing gradientproblem of the queries and keys — affects several layers (see Fig. 1).

Original: In order to shed light on thisproblem, we take inspiration from the flourishing literature on signal propagation in random networksand start our analysis by computing the expected gradients of an attention layer with respect to thequeries, keys, and values, which leads to Theorem 3.2 on the vanishing gradients for the queries andkeys.
Revised: In order to shed light on thisproblem, we take inspiration from the flourishing literature on signal propagation in random networksand start our analysis by computing the expected gradients of an attention layer with respect to thequeries, keys, and values, which leads to Theorem 3.2 on the vanishing gradients for the queries andkeys.

Original: From here, we pursue two different directions.
Revised: From here, we pursue two different directions.

Original: Firstly, we investigate under which conditions rank collapse can be avoided by studying the evolutionof the input sequence in a Transformer at initialization.
Revised: Firstly, we investigate under which conditions rank collapse can be avoided by studying the evolutionof the input sequence in a Transformer at initialization.

Original: Our theory reveals that a depth-dependentscaling of the residual branches, beyond stabilizing the norm of the activations at initialization, alsoapproximately preserves the cosine of the angle between tokens, and hence also stabilizes the rank ofthe propagating sequence.
Revised: Our theory reveals that a depth-dependentscaling of the residual branches, beyond stabilizing the norm of the activations at initialization, alsoapproximately preserves the cosine of the angle between tokens, and hence also stabilizes the rank ofthe propagating sequence.

Original: We show that this holds even in the infinite-depth limit.
Revised: We show that this holds even in the infinite-depth limit.

Original: Secondly, we illustrate that there are factors, other than the average tokens’ correlation, that affectdifferently the gradient norm of the queries and keys compared to the values.
Revised: Secondly, we illustrate that there are factors, other than the average tokens’ correlation, that affectdifferently the gradient norm of the queries and keys compared to the values.

Original: In particular, thepropagating sequence’s squared norm has a linear dependence in the values, while a cubic one in thequeries and keys, justifying the use of layer normalization.
Revised: In particular, thepropagating sequence’s squared norm has a linear dependence in the values, while a cubic one in thequeries and keys, justifying the use of layer normalization.

Original: We also highlight a different dependenceon the embedding dimension and the length of the input sequence, implying that the gradient normof a subset of parameters can potentially be of different orders of magnitude, as empirically hintedby previous works [Liu et al., 2020].
Revised: We also highlight a different dependenceon the embedding dimension and the length of the input sequence, implying that the gradient normof a subset of parameters can potentially be of different orders of magnitude, as empirically hintedby previous works [Liu et al., 2020].

Original: Our analysis brings to light fundamental issues in the signalpropagation in Transformers, opening the way for new, well-founded and motivated approaches toimprove optimization in these models.
Revised: Our analysis brings to light fundamental issues in the signalpropagation in Transformers, opening the way for new, well-founded and motivated approaches toimprove optimization in these models.

Original:  Background
Revised:  Background

Original: Transformers.
Revised: Transformers.

Original: A Transformer architecture consists of L stacked attention blocks, as show in Fig.
Revised: A Transformer architecture consists of L stacked attention blocks, as show in Fig.

Original: Layer normalization is usually applied token-wise either after the residual connections or to the inputsof the self-attention and position-wise feed-forward sub-layers, leading to the POST-LN [Vaswaniet al., 2017] and
Revised: Layer normalization is usually applied token-wise either after the residual connections or to the inputsof the self-attention and position-wise feed-forward sub-layers, leading to the POST-LN [Vaswaniet al., 2017] and

Original: PRE-LN
Revised: PRE-LN

Original: [Wang et al., 2019, Xiong et al., 2020] variants respectively.
Revised: [Wang et al., 2019, Xiong et al., 2020] variants respectively.

Original: Formally, given an input sequence X ∈ R n × d v , with n tokens of dimension d v , the single-headunmasked scaled dot-product self-attention 1 is defined as:
Revised: Formally, given an input sequence X ∈ R n × d v , with n tokens of dimension d v , the single-headunmasked scaled dot-product self-attention 1 is defined as:

Original: [Equation]
Revised: [Equation]

Original: where A
Revised: where A

Original: [Equation]
Revised: [Equation]

Original: softmax
Revised: softmax

Original: [Equation]
Revised: [Equation]

Original: where the softmax function is applied independently across and each layer is initialized with an independent set of weights.
Revised: where the softmax function is applied independently across and each layer is initialized with an independent set of weights.

Original: In the literature, the matrices
Revised: In the literature, the matrices

Original: X ℓ W Q , X ℓ W K , X ℓ W V are referred to as queries, keys and values, respectively.
Revised: X ℓ W Q , X ℓ W K , X ℓ W V are referred to as queries, keys and values, respectively.

Original: The complete
Revised: The complete

Original: Transformer block, in the absence of layer normalization, can be written recursively as:
Revised: Transformer block, in the absence of layer normalization, can be written recursively as:

Original: [Equation]
Revised: [Equation]

Original: where the introduced α 1 , α 2 parameters indicate the strength of the residual block, W F 1 , W F 2 ∈ 73 R d v × d v 2 are matrices of learnable parameters; we set X 0 :
Revised: where the introduced α 1 , α 2 parameters indicate the strength of the residual block, W F 1 , W F 2 ∈ 73 R d v × d v 2 are matrices of learnable parameters; we set X 0 :

Original: =
Revised: =

Original: X , and σ : R → R is an activationfunction.
Revised: X , and σ : R → R is an activationfunction.

Original: In our case, σ is the ReLU function, but we relax this assumption to the linear activationfrom Section 3.2 on.
Revised: In our case, σ is the ReLU function, but we relax this assumption to the linear activationfrom Section 3.2 on.

Original: [Figure]
Revised: [Figure]

Original: At initialization, each weight is sampled independently from a distributionwith zero-mean and variance σ 2 v = 1 d v for the values and feedforwardweights 3 , and σ 2 k = 1 d k for the queries and keys.
Revised: At initialization, each weight is sampled independently from a distributionwith zero-mean and variance σ 2 v = 1 d v for the values and feedforwardweights 3 , and σ 2 k = 1 d k for the queries and keys.

Original: This is the standard“Xavier”
Revised: This is the standard“Xavier”

Original: [Glorot and Bengio
Revised: [Glorot and Bengio

Original: 2010] or “He” 
Revised: 2010b] or “He” [

Original: [He et al
Revised: He et al

Original:  commonly used 
Revised: 2015] initialization commonly used in deep learning.

Original: Rank Collapse in Transformers.
Revised: Rank Collapse in Transformers.

Original: Interestingly, Dong et al.
Revised: Interestingly, Dong et al.

Original: [2021]proved that when the residual branches are omitted, the matrix of thetokens’ representations X ℓ converges to a rank-1 matrix in which allthe representations are the same and equal to a vector x ∈ R d v , i.e.
Revised: [2021]proved that when the residual branches are omitted, the matrix of thetokens’ representations X ℓ converges to a rank-1 matrix in which allthe representations are the same and equal to a vector x ∈ R d v , i.e.

Original: X ℓ → 1 n x ⊤ , where 1 d v is the vector with all ones in R d v .
Revised: X ℓ → 1 n x ⊤ , where 1 d v is the vector with all ones in R d v .

Original: Note thatthis is a slightly stronger notion of a rank- 1 matrix, as it implies that allthe tokens’ representations are both perfectly aligned and have the samenorm.
Revised: Note thatthis is a slightly stronger notion of a rank- 1 matrix, as it implies that allthe tokens’ representations are both perfectly aligned and have the samenorm.

Original: Indicating the inner product with the usual bracket notations ⟨· , ·⟩ ,and the cosine of the angle between two tokens as θ k,k ′ , perfect alignmenthappens when ⟨ X ℓk , X ℓk ′ ⟩ =
Revised: Indicating the inner product with the usual bracket notations ⟨· , ·⟩ ,and the cosine of the angle between two tokens as θ k,k ′ , perfect alignmenthappens when ⟨ X ℓk , X ℓk ′ ⟩ =

Original: (cid:13)(cid:13)
Revised: (cid:13)(cid:13)

Original: X ℓk (cid:13)(cid:13)
Revised: X ℓk (cid:13)(cid:13)

Original: (cid:13)(cid:13)
Revised: (cid:13)(cid:13)

Original: X ℓk (cid:13)(cid:13)
Revised: X ℓk (cid:13)(cid:13)

Original: cos θ k,k ′ with cos θ k,k ′ = 1 forall k, k ′ ∈ [ n ] .
Revised: cos θ k,k ′ with cos θ k,k ′ = 1 forall k, k ′ ∈ [ n ] .

Original: Note that perfect alignment together with equal normbetween all the tokens implies that all the representations are the same.
Revised: Note that perfect alignment together with equal normbetween all the tokens implies that all the representations are the same.

Original: One of our main contributions is to provide an explanation of how rankcollapse affects the gradients of a Transformer at initialization.
Revised: One of our main contributions is to provide an explanation of how rankcollapse affects the gradients of a Transformer at initialization.

Original: Vanishing Gradient Problem.
Revised: Vanishing Gradient Problem.

Original: Traditionally considered one of the coreissues that prevents successful training, the vanishing gradient problemhas a long and rich history that dates back to before the popularizationof deep learning [Hochreiter, 1991, Bengio et al., 1994].
Revised: Traditionally considered one of the coreissues that prevents successful training, the vanishing gradient problemhas a long and rich history that dates back to before the popularizationof deep learning [Hochreiter, 1991, Bengio et al., 1994].

Original: In its essence,given a loss function L : R n × d v → R , vanishing gradients occur whenthe norm of the gradient of the loss L with respect to the parameters of the network W — which weindicate as (cid:13) (cid:13) ∂
Revised: In its essence,given a loss function L : R n × d v → R , vanishing gradients occur whenthe norm of the gradient of the loss L with respect to the parameters of the network W — which weindicate as (cid:13) (cid:13) ∂

Original: L ∂ W (cid:13) (cid:13) — is too small to provide enough backpropagating signal, thus hindering gradientbased optimization methods.
Revised: L ∂ W (cid:13) (cid:13) — is too small to provide enough backpropagating signal, thus hindering gradientbased optimization methods.

Original: Despite extensive research toward understanding and overcoming theproblem in disparate contexts [Glorot and Bengio, 2010, He et al., 2015, Hanin, 2018, Zhang et al.
Revised: Despite extensive research toward understanding and overcoming theproblem in disparate contexts [Glorot and Bengio, 2010a, He et al., 2015, Hanin, 2018, Zhang et al.

Original: ,2019], a formal explanation of its role in relatively new architectures such as Transformers is largelymissing in the literature, with a few exceptions [Xiong et al., 2020, Wang et al., 2022, Huang et al.
Revised: ,2019], a formal explanation of its role in relatively new architectures such as Transformers is largelymissing in the literature, with a few exceptions [Xiong et al., 2020, Wang et al., 2022, Huang et al.

Original: In our paper (Section 3.1), we show how vanishing gradient occurs in conjunction with therank collapse issue identified by Dong et al.
Revised: In our paper (Section 3.1), we show how vanishing gradient occurs in conjunction with therank collapse issue identified by Dong et al.

Original: Signal Propagation in Random Networks at Initialization.
Revised: Signal Propagation in Random Networks at Initialization.

Original: After addressing the question on theeffects of rank collapse, we take a step back and rigorously analyze its causes by looking at how theproperties of the input sequence X are lost/preserved as it propagates through a randomly initialized
Revised: After addressing the question on theeffects of rank collapse, we take a step back and rigorously analyze its causes by looking at how the properties of the input sequence X are lost/preserved as it propagates through a randomly initialized

Original: Transformer.
Revised: Transformer.

Original: More specifically, we focus on two aspects of the propagating sequence: the expected
Revised: More specifically, we focus on two aspects of the propagating sequence: the expected

Original: Frobenius norm E (cid:13)(cid:13)
Revised: Frobenius norm E (cid:13)(cid:13)

Original: X ℓ (cid:13)(cid:13) 2 and the expected inner product between different tokens E ⟨ X k , X ′ k ⟩ , withk ̸ = k ′ .
Revised: X ℓ (cid:13)(cid:13) 2 and the expected inner product between different tokens E ⟨ X k , X ′ k ⟩ , withk ̸ = k ′ .

Original: The former is linked to a number of studies on the initialization of neural networks atthe edge of chaos [Poole et al., 2016, Schoenholz et al., 2017], and vanishing/exploding gradients[Hanin, 2018].
Revised: The former is linked to a number of studies on the initialization of neural networks atthe edge of chaos [Poole et al., 2016, Schoenholz et al., 2016], and vanishing/exploding gradients[Hanin, 2018].

Original: The latter quantity describes how the geometry of the feature space changes afterapplying a Transformer block, and is related to the concept of dynamical isometry [Saxe et al., 2013].
Revised: The latter quantity describes how the geometry of the feature space changes afterapplying a Transformer block, and is related to the concept of dynamical isometry [Saxe et al., 2013].

Original: To understand the evolution of the inner product, we analyze the following measure of correlation[Nachum et al., 2021, Cho and Saul, 2009]:
Revised: To understand the evolution of the inner product, we analyze the following measure of correlation[Nachum et al., 2021, Cho and Saul, 2009]:

Original: [Equation]
Revised: [Equation]

Original: Note that ρ ℓkk ′
Revised: Note that ρ ℓkk ′

Original: = 

Revised: =

Original: if and only if the k -th and k ′ -th
Revised: if and only if the k -th and k ′ -th

Original: tokens are perfectly aligned ( cos θ kk ′ = 1 ).
Revised: tokens are perfectly aligned ( cos θ kk ′ = 1 ).

Original: Westress that in our case — differently from the aforementioned works — instead of analyzing therelationship between two different data points, we study the relationship between tokens of the samesequence.
Revised: Westress that in our case — differently from the aforementioned works — instead of analyzing therelationship between two different data points, we study the relationship between tokens of the samesequence.

Original:  Theoretical Results
Revised:  Theoretical Results

Original: 3.1 Vanishing Gradients for Queries and Keys under Rank Collapse
Revised: 3.1 Vanishing Gradients for Queries and Keys under Rank Collapse

Original: To investigate the problem of vanishing gradients in the attention layers, we make use of the frameworkof matrix calculus [Magnus and Neudecker, 2019, Singh et al., 2021].
Revised: To investigate the problem of vanishing gradients in the attention layers, we make use of the frameworkof matrix calculus [Magnus and Neudecker, 2019, Singh et al., 2021].

Original: In particular, we comparethe expected Frobenius norm of the gradient of a self-attention layer with respect to its parameters:
Revised: In particular, we comparethe expected Frobenius norm of the gradient of a self-attention layer with respect to its parameters:

Original: E (cid:13)(cid:13)(cid:13) ∂ S ℓ ∂ W (cid:13)(cid:13)(cid:13) 2 F , where here W indicates one of the keys, queries or values weight matrices.
Revised: E (cid:13)(cid:13)(cid:13) ∂ S ℓ ∂ W (cid:13)(cid:13)(cid:13) 2 F , where here W indicates one of the keys, queries or values weight matrices.

Original: Due to thewell-known difficulty of computing expectations of the softmax [Daunizeau, 2017, Shekhovtsovand Flach, 2018], throughout this manuscript, we make the simplifying assumption that the softmaxoutput is the uniform distribution at initialization, i.e. the n × n matrix containing 1 n in each entry.
Revised: Due to thewell-known difficulty of computing expectations of the softmax [Daunizeau, 2017, Shekhovtsovand Flach, 2018], throughout this manuscript, we make the simplifying assumption that the softmaxoutput is the uniform distribution at initialization, i.e. the n × n matrix containing 1 n in each entry.

Original: Assumption 3.1 (Uniform attention) .
Revised: Assumption 3.1 (Uniform attention) .

Original: We assume that A ℓ = 1 n 1 n × n,where 1 n × n is the matrix with all entries equal to 1 .
Revised: We assume that A ℓ = 1 n 1 n × n .

Original: Crucially, in Appendix A.5, we formally showthat this assumption holds almost surely in the limit d k → ∞ .
Revised: Crucially, in Appendix A.5, we formally show that this assumption holds almost surely in thelimit d k → ∞ .

Original: There, we also experimentally showthat even in the more realistic case where d k = d v ≈ 512 , the empirical simulations provide asurprisingly faithful approximation of the theoretical insights presented in this paper.
Revised: There, we also experimentally show that even in the more realistic case whered k = d v ≈ 512 , the empirical simulations provide a surprisingly faithful approximation of thetheoretical insights presented in this paper.

Original: We define the mean token ¯ x ℓ through its components ¯ x ℓi = 1 n (cid:80) nk =1 X ℓki , i ∈
Revised: We define the mean token ¯ x ℓ through its components ¯ x ℓi = 1 n (cid:80) nk =1 X ℓki , i ∈

Original: [ d v ] .
Revised: [ d v ] .

Original: In the followingtheorem, we compute the expected gradients of an attention layer at initialization, and set the basisfor our following analysis.
Revised: In the followingtheorem, we compute the expected gradients of an attention layer at initialization, and set the basisfor our following analysis.

Original: We provide the results only for the queries, as the case for the keys isanalogous.
Revised: We provide the results only for the queries, as the case for the keys isanalogous.

Original: Lemma 3.1.
Revised: Lemma 3.1.

Original: Let X ℓ be the representations of the input sequence at the ℓ -th layer.
Revised: Let X ℓ be the representations of the input sequence at the ℓ -th layer.

Original: Under the uniform-attention assumption, we have
Revised: Under the uniform-attention assumption, we have

Original: [Equation]
Revised: [Equation]

Revised: ℓ

Revised: [Equation]

Original: 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.
Revised: 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.

Original: At this stage, it is crucial to note that 1 n ( X ℓ ) ⊤ X ℓ − ¯ x ℓ (¯ x ℓ ) ⊤ is the centered empirical covariancematrix of the tokens’ representations.
Revised: At this stage, it is crucial to note that 1 n ( X ℓ ) ⊤ X ℓ − ¯ x ℓ (¯ x ℓ ) ⊤ is the centered empirical covariancematrix of the tokens’ representations.

Original: It is easy to see that if X ℓ is a rank- 1 matrix, then all the rowsof X ℓ are proportional to a fixed d v -dimensional vector, and the empirical covariance matrix hasall zero entries.
Revised: It is easy to see that if X ℓ is a rank- 1 matrix, then all the rowsof X ℓ are proportional to a fixed d v -dimensional vector, and the empirical covariance matrix hasall zero entries.

Original: Introducing a loss function L : R n × d v → R , we make the statement on vanishinggradients more formal in the following theorem:
Revised: Introducing a loss function L : R n × d v → R , we make the statement on vanishinggradients more formal in the following theorem:

Original: Theorem 3.2 (Vanishing gradients under rank collapse) .
Revised: Theorem 3.2 (Vanishing gradients under rank collapse) .

Original: Suppose that the uniform-attention assumption holds.
Revised: Suppose that the uniform-attention assumption holds.

Original: If additionally X ℓ for any l ∈
Revised: If additionally X ℓ for any l ∈

Original: [ L ] has rank-1, and there exists a vector x ∈ R d such that X ℓ = 1 n x T , then:
Revised: [ L ] has rank-1, and there exists a vector x ∈ R d such that X ℓ = 1 n x T , then:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where the expectation is taken over the weight matrices.
Revised: where the expectation is taken over the weight matrices.

Original: This implies that these quantities are vanishing almost surely, due to the non-negativeness of the norm.
Revised: This implies that these quantities are vanishing almost surely, due to the non-negativeness of the norm.

Original: The proof simply relies on expanding the norm of the gradient of the loss with the aid of the chainrule and then bounding it by the product of the norms of each term of the chain.
Revised: The proof simply relies on expanding the norm of the gradient of the loss with the aid of the chainrule and then bounding it by the product of the norms of each term of the chain.

Original: The final result holdswith an application of Lemma 3.1, in which the rank- 1 assumption makes E (cid:13)(cid:13)(cid:13) ∂ S ℓ ∂ W Q,ℓ
Revised: The final result holdswith an application of Lemma 3.1, in which the rank- 1 assumption makes E (cid:13)(cid:13)(cid:13) ∂ S ℓ ∂ W Q,ℓ

Original: (cid:13)(cid:13)(cid:13) vanish.
Revised: (cid:13)(cid:13)(cid:13) vanish.

Original: In light of Theorem 3.2, we can conclude that the rank collapse issue originally identified in Donget al.
Revised: In light of Theorem 3.2, we can conclude that the rank collapse issue originally identified in Donget al.

Original: [2021] corresponds to an initialization in a region of vanishing gradient signal in the subspaceof parameters identified by the queries and keys.
Revised: [2021] corresponds to an initialization in a region of vanishing gradient signal in the subspaceof parameters identified by the queries and keys.

Original: How can this affect training?
Revised: How can this affect training?

Original: One may argue that ifrank 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 thefirst few steps of gradient descent.
Revised: One may argue that ifrank 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 thefirst few steps of gradient descent.

Original: In practice, we show empirically in Fig.
Revised: In practice, we show empirically in Fig.

Original: 1 that escaping thispathological landscape is harder in deeper nets.
Revised: 1 that escaping thispathological landscape is harder in deeper nets.

Original: 3.2 Forward Signal Propagation and the Importance of Scaling the Residual Branches
Revised: 3.2 Forward Signal Propagation and the Importance of Scaling the Residual Branches

Original: We now turn our attention to the study of the influence of skip connections in transformers.
Revised: We now turn our attention to the study of the influence of skip connections in transformers.

Original: Dong et al.[2021] showed that simply adding skip connections prevents rank collapse.
Revised: Dong et al.[2021] showed that simply adding skip connections prevents rank collapse.

Original: Somewhat surprisingly,we show that while the claim holds for any finite depth, the average angle between different tokensquickly increases with just a few layers, and as L → ∞ a Transformer can still lose rank unless theresidual branches are adequately initialized.
Revised: Somewhat surprisingly,we show that while the claim holds for any finite depth, the average angle between different tokensquickly increases with just a few layers, and as L → ∞ a Transformer can still lose rank unless theresidual branches are adequately initialized.

Original: As Dong et al.
Revised: As Dong et al.

Original: [2021] showed that layer normalizationdoes not avoid rank collapse, we omit it in our analysis.
Revised: [2021] showed that layer normalizationdoes not avoid rank collapse, we omit it in our analysis.

Original: Firstly, we introduce two lemmas on thepropagation of inner products (Lemma 3.2) and the norm (Lemma 3.3) of the tokens’ representations.
Revised: Firstly, we introduce two lemmas on thepropagation of inner products (Lemma 3.2) and the norm (Lemma 3.3) of the tokens’ representations.

Original: Lemma 3.2 (Propagation of inner products) .
Revised: Lemma 3.2 (Propagation of inner products) .

Original: Let C ( X ℓ )
Revised: Let C ( X ℓ )

Original: =
Revised: =

Original: (cid:80) k,k ′ ⟨ X ℓk , X ℓk ′ ⟩ and X the input sequence.
Revised: (cid:80) k,k ′ ⟨ X ℓk , X ℓk ′ ⟩ and X the input sequence.

Original: Under the Assumption 3.1 and if σ is the linear activation function, we have that:
Revised: Under the Assumption 3.1 and if σ is the linear activation function, we have that:

Original: [Equation]
Revised: [Equation]

Original: hence, under the depth scaling for the residual block parameters α 21 = ˜
Revised: hence, under the depth scaling for the residual block parameters α 21 = ˜

Original: α 1 L , α 22 = ˜
Revised: α 1 L , α 22 = ˜

Original: α 2 L with ˜ α
Revised: α 2 L with ˜ α

Original: [Equation]
Revised: [Equation]

Original: ˜ α 2 ∈ R independent of L , we have that:
Revised: ˜ α 2 ∈ R independent of L , we have that:

Original: [Equation]
Revised: [Equation]

Original: Note that C ( X ℓ ) = n 2 (cid:13)(cid:13)¯ x ℓ (cid:13)(cid:13) 2 .
Revised: Note that C ( X ℓ ) = n 2 (cid:13)(cid:13) x ℓ (cid:13)(cid:13) 2 .

Original: The lemma on the propagation of the norm is slightly more involved:
Revised: The lemma on the propagation of the norm is slightly more involved:

Original: Lemma 3.3 (Propagation of the norm) .
Revised: Lemma 3.3 (Propagation of the norm) .

Original: Let X L be the representations of the input sequence at the final layer.
Revised: Let X L be the representations of the input sequence at the final layer.

Original: Under the assumptions of Lemma 3.2, we have that:
Revised: Under the assumptions of Lemma 3.2, we have that:

Original: [Equation]
Revised: [Equation]

Original: hence, under the depth scaling for the residual block parameters α 21 = ˜
Revised: hence, under the depth scaling for the residual block parameters α 21 = ˜

Original: α 1 L , α 22 = ˜
Revised: α 1 L , α 22 = ˜

Original: α 2 L with ˜
Revised: α 2 L with ˜

Original: α 1 , ˜ α 2 ∈ R independent of L , we have that:
Revised: α 1 , ˜ α 2 ∈ R independent of L , we have that:

Original: [Equation]
Revised: [Equation]

Original: The proof of Lemma 3.3 consists in expanding E (cid:13)(cid:13)
Revised: The proof of Lemma 3.3 consists in expanding E (cid:13)(cid:13)

Original: X L (cid:13)(cid:13) 2 F according to the defining equationsfor the Transformer, and simplifying the expression by using iterated expectations E (cid:13)(cid:13)
Revised: X L (cid:13)(cid:13) 2 F according to the defining equationsfor the Transformer, and simplifying the expression by using iterated expectations E (cid:13)(cid:13)

Original: X L (cid:13)(cid:13) 2 F =
Revised: X L (cid:13)(cid:13) 2 F =

Original: [Figure]
Revised: [Figure]

Original: E [ E [ (cid:13)(cid:13)
Revised: E [ E [ (cid:13)(cid:13)

Original: X L (cid:13)(cid:13) 2 F | X ℓ ]] to exploit the conditional independence between different layers, and then computing the expectations using the independence assumption on the weights.
Revised: X L (cid:13)(cid:13) 2 F | X ℓ ]] to exploit the conditional independence between different layers, and then computing the expectations using the independence assumption on the weights.

Original: The expression on theright-hand side will then depend on X ℓ only through its norm (cid:13)(cid:13)
Revised: The expression on theright-hand side will then depend on X ℓ only through its norm (cid:13)(cid:13)

Original: X ℓ (cid:13)(cid:13) and the norm of the mean token(cid:13)(cid:13) ¯ x ℓ (cid:13)(cid:13) 2 .
Revised: X ℓ (cid:13)(cid:13) and the norm of the mean token(cid:13)(cid:13) ¯ x ℓ (cid:13)(cid:13) 2 .

Original: Using Lemma 3.2 then allows us to unroll the recursion and get the final result.
Revised: Using Lemma 3.2 then allows us to unroll the recursion and get the final result.

Original: The completeproof, together with the proof of Lemma 3.2, can be found in Appendix A.3.
Revised: The completeproof, together with the proof of Lemma 3.2, can be found in Appendix A.3.

Original: The previous Lemma provides theoretical justification that scaling the residual branches by settingthe alpha parameters to be O (1 / √ L ) allows both the norm of the propagating input and the innerproducts between different tokens to be approximately preserved.
Revised: The previous Lemma provides theoretical justification that scaling the residual branches by settingthe alpha parameters to be O (1 / √ L ) allows both the norm of the propagating input and the innerproducts between different tokens to be approximately preserved.

Original: Hence, the information containedin the input is not lost, even in the infinite depth limit.
Revised: Hence, the information containedin the input is not lost, even in the infinite depth limit.

Original: Residual Scaling Preserves Correlations.
Revised: Residual Scaling Preserves Correlations.

Original: We now prove that without the depth-dependent residualscaling (i.e. with α 

Revised: We now prove that without the depth-dependent residualscaling (i.e. with α

Original: = α 2 = 1 ) the correlation between the tokens quickly increases, and reachesperfect alignment in the infinite depth limit.
Revised: = α 2 = 1 ) the correlation between the tokens quickly increases, and reachesperfect alignment in the infinite depth limit.

Original: More specifically, our argument shows that in this limit,the correlation between different tokens ρ ℓk,k ′ as in Eq.
Revised: More specifically, our argument shows that in this limit,the correlation between different tokens ρ ℓk,k ′ as in Eq.

Original: (5) converges to 1, implying rank collapse.
Revised: (5) converges to 1, implying rank collapse.

Original: Furthermore, we show how setting the residual parameters α 1 and α 2 as dictated by Theorem 3.3,ensures that the correlation measure is dependent on the input in a non-trivial way even at infinitedepth.
Revised: Furthermore, we show how setting the residual parameters α 1 and α 2 as dictated by Theorem 3.3,ensures that the correlation measure is dependent on the input in a non-trivial way even at infinitedepth.

Original: To this end, we introduce the average correlation at layer ℓ :
Revised: To this end, we introduce the average correlation at layer ℓ :

Original: [Equation]
Revised: [Equation]

Original: Note that ρ ℓ = 1 if and only if every pair of tokens is perfectly aligned.
Revised: Note that ρ ℓ = 1 if and only if every pair of tokens is perfectly aligned.

Original: We are now ready toformalize the influence of the 1 / √ L -scaling on the correlation between tokens’ representations bystating Theorem 3.3.
Revised: We are now ready toformalize the influence of the 1 / √ L -scaling on the correlation between tokens’ representations bystating Theorem 3.3.

Original: Theorem 3.3.
Revised: Theorem 3.3.

Original: Let the input tokens have the same norm, i.e. ∥
Revised: Let the input tokens have the same norm, i.e. ∥

Original: X k ∥ =
Revised: X k ∥ =

Original: ∥ x ∥ ∀ k ∈ [ n ] for some x ∈ R d v .
Revised: ∥ x ∥ ∀ k ∈ [ n ] for some x ∈ R d v .

Original: Under the depth scaling for the residual block parameters α 

Revised: Under the depth scaling for the residual block parameters α

Original: = ˜
Revised: = ˜

Original: α 1 L , α 22 = ˜
Revised: α 1 L , α 22 = ˜

Original: α 2 L with ˜
Revised: α 2 L with ˜

Original: α 1 , ˜ α 2 ∈ R independent of L , we have that:
Revised: α 1 , ˜ α 2 ∈ R independent of L , we have that:

Original: [Equation]
Revised: [Equation]

Original: On the other hand, if α 1 , α 2 ̸ = 0 are some constants independent of L , we have that:
Revised: On the other hand, if α 1 , α 2 ̸ = 0 are some constants independent of L , we have that:

Original: [Equation]
Revised: [Equation]

Original: The proof consists in noting that due to the symmetry of the problem at initialization, for a fixedlayer the expected norm of each token is the same.
Revised: The proof consists in noting that due to the symmetry of the problem at initialization, for a fixedlayer the expected norm of each token is the same.

Original: Hence, by our definition of ρ ℓkk ′ , we can write
Revised: Hence, by our definition of ρ ℓkk ′ , we can write

Original: E ⟨ X ℓk , X ℓk ′ ⟩ = ρ ℓkk ′ E (cid:13)(cid:13) x ℓ (cid:13)(cid:13) 2 .
Revised: E ⟨ X ℓk , X ℓk ′ ⟩ = ρ ℓkk ′ E (cid:13)(cid:13) x ℓ (cid:13)(cid:13) 2 .

Original: By summing over the k, k ′ indexes, the resulting equation will dependon E
Revised: By summing over the k, k ′ indexes, the resulting equation will dependon E

Original: [ C ( X ℓ )] and E (cid:13)(cid:13) X ℓ (cid:13)(cid:13) 2 , which can be expanded using Lemma 3.2 and 3.3 respectively.
Revised: [ C ( X ℓ )] and E (cid:13)(cid:13) X ℓ (cid:13)(cid:13) 2 , which can be expanded using Lemma 3.2 and 3.3 respectively.

Original: Theresult is then given by solving for ρ ℓ .
Revised: Theresult is then given by solving for ρ ℓ .

Original: [Figure]
Revised: [Figure]

Original: Note that under the 1 / √ L -scaling, the correlation term isone if and only if C ( X )
Revised: Note that under the 1 / √ L -scaling, the correlation term isone if and only if C ( X )

Original: =
Revised: =

Original: n ∥ X ∥ 2 , which holds in thedegenerate case where all the input tokens are perfectlyaligned.
Revised: n ∥ X ∥ 2 , which holds in thedegenerate case where all the input tokens are perfectlyaligned.

Original: In Appendix A.4, we give precise formulas forthe expected correlations at any depth, showing that ρ ℓreaches values close to one even for relatively shallow networks when the 1 / √ L -scaling is not adopted (see also Fig.3 (left)).
Revised: In Appendix A.4, we give precise formulas forthe expected correlations at any depth, showing that ρ ℓreaches values close to one even for relatively shallownetworks when the 1 / √ L -scaling is not adopted (see also Fig. 3(left)).

Original: Additionally, in Fig.
Revised: Additionally, in Fig.

Original: 4, we empirically show thatin the presence of the 1 / √ L -scaling, layer normalization(either PRE or POST) does not significantly affect the evolution of the correlations.
Revised: 4, we empirically showthat in the presence of the 1 / √ L -scaling, layer normalization (either PRE or POST) does not significantly affect theevolution of the correlations.

Original: On the other hand, without theresidual scaling, PRE-LN seems to alleviate the rate of increase of ρ ℓ kk ′ .
Revised: On the other hand, withoutthe residual scaling, PRE-LN seems to alleviate the rateof increase of ρ ℓkk ′ .

Original: It is intriguing that most deep Transformermodels use this configuration [Brown et al., 2020].

Original: Weprovide more extensive empirical results in Appendix B.
Revised: We provide more extensive empiricalresults in Appendix B.

Original: Note that the 1 / √ L scaling for the residual branches hasbeen previously studied in the context of stabilization ofresidual networks (see Section 4), here we extend theseresults to Transformers and provide new insights on itsrole in the context of rank preservation.
Revised: Note that the 1 / √ L scaling for the residual branches hasbeen previously studied in the context of stabilization ofresidual networks (see Section 4), here we extend theseresults to Transformers and provide new insights on itsrole in the context of rank preservation.

Original: Finally, note thatby setting ˜ α 1 , ˜ α 2 = 0 , we recover the so called "ReZero"initialization [Bachlechner et al., 2021].
Revised: Finally, note thatby setting ˜ α 1 , ˜ α 2 = 0 , we recover the so called "ReZero"initialization [Bachlechner et al., 2021].

Original: In this context,the 1 / √ L scaling extends this framework as it allows forwider range of values for ˜ α 1 , ˜
Revised: In this context,the 1 / √ L scaling extends this framework as it allows forwider range of values for ˜ α 1 , ˜

Original: α 2 while still guaranteeing stability.
Revised: α 2 while still guaranteeingstability.

Original: ReLU extension. 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 innerproducts across ReLU layers that are shared among the tokens (this is the case in the position-wisefeed-forward layers in Transformers).
Revised: We mention here that extending these results from the linear activation to the ReLU case isknown 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-forwardlayers in Transformers).

Original: Exact formulas can be found only in the case of one ReLUlayer with Gaussian inputs in Cho and Saul [2009].
Revised: Exact formulas can be found only in the case of one ReLU layer with Gaussian inputs in Cho and Saul [2009].

Original: However, in the context of rank collapse analyzedhere, the linear activation function provides a bound on the correlation with respect to the ReLU case.

Original: In fact, correlations are exactly preserved in expectation in the linear case, but increase in the ReLUcase (for instance, see the contraction argument in Nachum et al.

Original: [2021] below Equation (2)).

Original: Hence,the perfect alignment (a.k.a rank collapse) that affects the linear case affects the ReLU case as well(in which case the rank collapses even faster with depth, as we show in Figure 10).

Original: 3.3 Dependence on the Angle between Tokens and the Input Norm
Revised: 3.3 Dependence on the Angle between Tokens and the Input Norm

Original: In this section, we drop the superscript ℓ as it is obvious from context and assume for simplicity thatd k = d v .
Revised: In this section, we drop the superscript ℓ as it is obvious from context and assume for simplicity thatd k = d v .

Original: To gain a better intuition on the factors that affect the gradients and provide additionalinsights, we study the case in which every pair of distinct tokens are zero-mean Gaussian randomvariables, correlated in the same way, i.e ρ ℓii ′ =
Revised: To gain a better intuition on the factors that affect the gradients and provide additionalinsights, we study the case in which every pair of distinct tokens are zero-mean Gaussian randomvariables, correlated in the same way, i.e ρ ℓii ′ =

Original: ρ for i ̸
Revised: ρ for i ̸

Original: = i ′ or more precisely
Revised: = i ′ or more precisely

Original: [Equation] dimensions
Revised: E [ X i,j X i ′ ,j ′ ] =    j ̸ = j ′ (independent dimensions) σ 2 x i = i ′ , j = j ′ ρσ 2 x i ̸

Original: [Equation]
Revised: = i ′ , j = j ′ .

Original: To see that this equation satisfies our definition of the correlation metric, note that E [ ∥ X i ∥ 2 ] = dσ 2 xand E ⟨ X i , X i ′ ⟩ = dσ 2 x ρ , for i ̸
Revised: To see that this equation satisfies our definition of the correlation metric, note that E [ ∥ X i ∥ 2 ] = dσ 2 xand E ⟨ X i , X i ′ ⟩ = dσ 2 x ρ , for i ̸

Original: = i ′ .
Revised: = i ′ .

Original: Then, the expected norm of the gradients for the values (Eq.
Revised: Then, the expected norm of the gradients for the values (Eq.

Original: (6))simplifies to
Revised: (6))simplifies to

Original: [Equation][Figure] [Table] [Table]
Revised: [Equation]

Original: inverse

Original: [Table] [Table]
Revised: [Equation]

Original: By making the additional assumption that the norm and the correlation propagate independently, therespective norm for the queries — and symmetrically the keys — (Eq.
Revised: By making the additional assumption that the norm and the correlation propagate independently, therespective norm for the queries — and symmetrically the keys — (Eq.

Original: (7)) reduces to:
Revised: (7)) reduces to:

Original: [Equation]
Revised: [Equation]

Original: In Appendix A.2 we provide a rigorous proof, that relies on Isserlis theorem [Isserlis, 1918] tocompute higher-order moments.
Revised: In Appendix A.2 we provide a rigorous proof, that relies on Isserlis theorem [Isserlis, 1918] tocompute higher-order moments.

Original: The above expressions reveal the different dependencies on fourmain actors, that we inspect separately here.
Revised: The above expressions reveal the different dependencies on four [Figure] [Table] main actors, that we inspect separately here.

Original: The gradients of the queries depend via a cubic functionon the variance of the input, σ 2 x , compared to a linear for the values.
Revised: The gradients of the queries depend via a cubic functionon the variance of the input, σ 2 x , compared to a linear for the values.

Original: This provides an additionalinterpretation of the successful use of layer normalization, as in Xiong et al.
Revised: This provides an additionalinterpretation of the successful use of layer normalization, as in Xiong et al.

Original: [2020], either in the
Revised: [2020], either in the

Original: POST-LN or PRE-LN format, that standardizes the input variance σ 2 x to the value 1 .
Revised: POST-LN or PRE-LN format, that standardizes the input variance σ 2 x to the value 1 .

Original: Next, we emphasize the dependence on the correlation between the tokens , also illustrated in Fig.
Revised: Next, we emphasize the dependence on the correlation between the tokens , also illustrated in Fig.

Original: Importantly, note how the queries/keys have opposite monotonic functional dependence with respectto ρ compared to the values.
Revised: Importantly, note how the queries/keys have opposite monotonic functional dependence with respectto ρ compared to the values.

Original: As revealed by Theorem 3.3 and Fig. 3 (center), inappropriate scaling ofthe residual branches can already lead to this phenomenon even in a relatively shallow network.
Revised: As revealed by Theorem 3.3 and Fig. 3(center), inappropriate scaling ofthe residual branches can already lead to this phenomenon even in a relatively shallow network.

Original: Finally, Eq.
Revised: Finally, Eq.

Original: (17) and (18) reveal a different scaling in terms of the embedding size d and the sequencelength n due to the self-attention operation itself.
Revised: (17) and (18) reveal a different scaling in terms of the embedding size d and the sequencelength n due to the self-attention operation itself.

Original: We hope that the identification of the differentdependencies in the gradients of the parameters will inspire a new line of works aimed at solvingsome of the difficulties in training Transformers.
Revised: We hope that the identification of the differentdependencies in the gradients of the parameters will inspire a new line of works aimed at solvingsome of the difficulties in training Transformers.

Original: 3.4 Are Adaptive Methods really needed for training Transformers?
Revised: 3.4 Are Adaptive Methods really needed for training Transformers?

Original: [Figure]
Revised: [Figure]

Original: The existence of the discrepancy in the magnitude of the gradients with respect to the weights
Revised: The existence of the discrepancy in the magnitude of the gradients with respect to the weights

Original: W Q , W K and W V , might explain the successof adaptive optimization algorithms, as illustrated in Fig.
Revised: W Q , W K and W V , might explain the successof adaptive optimization algorithms, as illustrated in Fig.

Original: 6, where we plot the effective learning rate computed by Adam [Kingma and Ba,2014] in a toy encoder task (more details in Appendix C).
Revised: 6, where we plot the effective learning rate computed by Adam [Kingma and Ba,2014] in a toy encoder task (more details in Appendix C).

Original: Hence, we embark on a preliminaryexploration to train a Transformer architecturewith SGD with the intent of matching Adam’sperformance.
Revised: Hence, we embark on a preliminaryexploration to train a Transformer architecturewith SGD with the intent of matching Adam’sperformance.

Original: Based on our theory, we proposea simple architectural modification, an inversetemperature scaling τ ∈ R inside the softmax:
Revised: Based on our theory, we proposea simple architectural modification, an inversetemperature scaling τ ∈ R inside the softmax:

Original: [Equation]
Revised: [Equation]

Original: A direct consequence of our analysis is that τ allows controlling the magnitude of the gradients for thequeries and keys’ parameters.
Revised: A direct consequence of our analysis is that τ allows controlling the magnitude of the gradients forthe queries and keys’ parameters.

Original: In Section C.2, we detail how one can choose τ such that the magnitudeof the gradients as derived in Equation 17 and 18 is approximately matched at initialization.

Original: We evaluate our proposal, consisting of residual scaling and the aforementioned inverse temperatureparameters, on the widely used IWSLT14 German-to-English (De-En) benchmark translation task.
Revised: We evaluate our proposal, consisting of residual scaling and the aforementioned inverse temperatureparameters, on the widely used IWSLT14 German-to-English (De-En) benchmark translation task.

Original: All details regarding the experimental setup and the choice of inverse temperature used are providedin Appendix C.
Revised: All details regarding the experimental setup and the choice of inverse temperature used are providedin Appendix C.

Original: We train a Transformer encoder-decoder of varying depth with SGD, after removingall normalization layers and adequately initializing the residual connections.
Revised: We train a Transformer encoder-decoder of varying depth with SGD, after removingall normalization layers and adequately initializing the residual connections.

Original: For our training with
Revised: For our training with

Original: SGD, we avoid using any learning rate warm-up, as commonly done for Adam, and instead use astep-scheduler to decrease the learning rate at 40% and 80% of training.
Revised: SGD, we avoid using any learning rate warm-up, as commonly done for Adam, and instead use astep-scheduler to decrease the learning rate at 40% and 80% of training.

Original: We compare against thefollowing methods that make use of Adam; POST-LN and PRE-LN refer to the aforementionedalternatives to apply layer normalization.
Revised: We compare against thefollowing methods that make use of Adam; POST-LN and PRE-LN refer to the aforementioned alternatives to apply layer normalization.

Original: We also compare against other successful techniques thatrely on specific initializations to avoid layer normalization, such as ReZero [Bachlechner et al.,
Revised: We also compare against other successful techniques thatrely on specific initializations to avoid layer normalization, such as ReZero [Bachlechner et al.,

Original: 2021]and T-Fixup [Zhang et al., 2019].
Revised: 2021]and T-Fixup [Zhang et al., 2019].

Original: We report the average BLEU score [Papineni et al., 2002] across 5runs in Fig.
Revised: We report the average BLEU score [Papineni et al., 2002] across 5runs in Fig.

Original: 5 and Table 1.
Revised: 5 and Table 1.

Original: Our proposed method considerably improves training with SGD, keeping up and in some casessurpassing any results achieved by the Adam optimizer.
Revised: Our proposed method considerably improves training with SGD, keeping up and in some casessurpassing any results achieved by the Adam optimizer.

Original: We are also able to train deeper networkswithout the use of layer normalization.
Revised: We are also able to train deeper networkswithout the use of layer normalization.

Original: We leave for future work to further investigate modificationsor alternatives to the self-attention operation.
Revised: We leave for future work to further investigate modificationsor alternatives to the self-attention operation.

Original:  Related Work
Revised:  Related Work

Original: Our work builds upon the rich literature on forward and backward signal propagation in randomneural networks [Poole et al., 2016, Schoenholz et al., 2017, Xiao et al., 2018, Pennington et al.
Revised: Our work builds upon the rich literature on forward and backward signal propagation in randomneural networks [Poole et al., 2016, Schoenholz et al., 2017, Xiao et al., 2018, Pennington et al.

Original: ,2017, Orvieto et al., 2021, Noci et al., 2021, Zavatone-Veth and Pehlevan, 2021].
Revised: ,2017, Orvieto et al., 2021, Noci et al., 2021, Zavatone-Veth and Pehlevan, 2021].

Original: The 1 / √ L scalingscheme has been investigated in the literature for the stabilization of residual networks [Hanin and
Revised: The 1 / √ L scalingscheme have been investigated in the literature for the stabilization of residual networks [Hanin and

Original: Rolnick, 2018, Arpit et al., 2019, Allen-Zhu et al., 2019, Hayou et al., 2021].
Revised: Rolnick, 2018, Arpit et al., 2019, Allen-Zhu et al., 2019, Hayou et al., 2021].

Original: Our work draws inspiration from a series of recent works studying the rank of the representationsof random feed-forward neural networks at initialization [Daneshmand et al., 2020, 2021].
Revised: Our work draws inspiration from a series of recent works studying the rank of the representations ofrandom feed-forward neural networks at initialization [Daneshmand et al., 2020b,a].

Original: In thecontext of Transformers, Dong et al.
Revised: In the contextof Transformers, Dong et al.

Original: [2021] has recently identified the rank collapse issue objectofstudy of  the present work.
Revised: [2021] has recently identified the rank collapse issue object of studyof the present work.

Original: Thanks to our analysis of the backward pass, we are able to demonstratethat rank collapse in Transformer architectures leads to vanishingly small gradients of queries andkeys, thereby preventing effective training and allowing us to complete the analysis of [Dong et al.,2021].
Revised: Thanks to our analysis of the backward pass, we are able to demonstrate thatrank 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].

Original: Among the architectural components in Transformers, layer normalization is, arguably, one of themost important – and debated – ones [Chen et al., 2018, Wang et al., 2019, Nguyen et al., 2010, Xionget al., 2020].
Revised: Among the architectural components in Transformers, layer normalization is, arguably, one of themost important – and debated – ones [Chen et al., 2018, Wang et al., 2019, Nguyen et al., 2010, Xionget al., 2020].

Original: In the original architecture [Vaswani et al., 2017], layer normalization is used to stabilizethe forward pass by reducing the variance of the inputs to the following sublayer.
Revised: In the original architecture [Vaswani et al., 2017], layer normalization is used to stabilizethe forward pass by reducing the variance of the inputs to the following sublayer.

Original: Our analysis of theforward pass shows that its inclusion is not strictly necessary for the purpose of controlling the normof the representations.
Revised: Our analysis of theforward pass shows that its inclusion is not strictly necessary for the purpose of controlling the normof the representations.

Original: For a theoretical analysis of signal propagation in the presence of layer norm,we refer the reader to Xiong et al.
Revised: For a theoretical analysis of signal propagation in the presence of layer norm,we refer the reader to Xiong et al.

Original: Additionally, our theoretical study of the backward pass provides a rigorous explanation of theempirically observed discrepancy between the magnitude of the gradients of the queries and thevalues, which Liu et al.
Revised: Additionally, our theoretical study of the backward pass provides a rigorous explanation of theempirically observed discrepancy between the magnitude of the gradients of the queries and thevalues, which Liu et al.

Original: [2020] hypothesize to be one of the causes of the success of adaptive methodsin training Transformers [Liu et al., 2019, Zhang et al., 2020, Huang et al., 2020].
Revised: [2020] hypothesize to be one of the causes of the success of adaptive methodsin training Transformers [Liu et al., 2019, Zhang et al., 2020, Huang et al., 2020].

Original: 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, Wanget al., 2022].
Revised: 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, Wanget al., 2022].

Original: However, none of these studies characterize the impact of skip connections on rankpropagation, while our analysis suggests a theoretically-grounded way to stabilize it.
Revised: However, none of these studies characterize the impact of skip connections on rankpropagation, while our analysis suggests a theoretically-grounded way to stabilize it.

Original:  Conclusions and Future Work
Revised:  Conclusions and Future Work

Original: In this paper, we showed how, at initialization, rank collapse and more generally high correlationin the tokens, causes vanishing gradients of the queries and keys of a Transformer architecture.
Revised: In this paper, we showed how, at initialization, rank collapse and more generally high correlationin the tokens, causes vanishing gradients of the queries and keys of a Transformer architecture.

Original: While residual connections help mitigate rank collapse at finite depth, we showed that they alonecannot prevent high alignments of the tokens’ representations — unless properly scaled by a 1 / √ Lfactor.
Revised: While residual connections help mitigate rank collapse at finite depth, we showed that they alonecannot prevent high alignments of the tokens’ representations — unless properly scaled by a 1 / √ Lfactor.

Original: 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.
Revised: 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 gradientsof queries/keys compared to the values’ parameters.

Original: Hence, we conclude that one of the strengthsof Transformers lies in their carefully designed architecture together with an adequate initialization.
Revised: Hence, we conclude that one of the strengthsof Transformers lies in their carefully designed architecture together with an adequate initialization.

Original: Finally, we gave preliminary evidence that one of the factors contributing to the higher efficacyof Adam compared to SGD in training Transformers arises from the disproportionate magnitudeof gradients as postulated by our theory.
Revised: Finally, we gave preliminary evidence that one of the factors contributing to the higher efficacyof Adam compared to SGD in training Transformers arises from the disproportionate magnitudeof gradients as postulated by our theory.

Original: Nonetheless, other factors might further accentuate thedifference between these two algorithms during training, leaving the door open for further researchregarding the benefits of adaptive optimization methods with Transformers
Revised: Nonetheless, other factors might further accentuate thedifference between these two algorithms during training, leaving the door open for further researchregarding the benefits of adaptive optimization methods with Transformers.