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: Abstract
Revised: Abstract

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

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

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

Original: In this work, we shed new light on the causes and the effects of this phenomenon.
Revised: In this work, we shed new light on the causes and 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 keys to vanish at initialization.
Revised: 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.

Original: Furthermore, we provide a thorough description of the 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 of the 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 specific architectural hyperparameters affect the gradients of queries and values differently, leading to disproportionate gradient norms.
Revised: Finally, our analysis unveils that specific architectural hyperparameters affect the gradients of queries and values differently, leading to disproportionate gradient norms.

Original: This suggests an explanation for the widespread use of adaptive methods for Transformers’ optimization.
Revised: This suggests an explanation for the widespread 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 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.
Revised: [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.

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., 2022], further popularizing it outside NLP.
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., 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 stacked attention 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 finally normalized with the softmax function.
Revised: At its core, it relies on stacked attention 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 finally normalized with the softmax function.

Original: The output of an attention layer is then a linear combination
Revised: The output of an attention layer is then a linear combination

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

Original: C
Revised: C

Original: o s i n e a n g l e ( i n i t i a li z a t i o n )
Revised: o s i n e a n g l e ( i n i t i a li z a t i o n )

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

Original: s i n
Revised: s i n

Original: [Figure]
Revised: [Figure]

Original: l
Revised: l

Original: [Figure]
Revised: [Figure]

Original: ( i n i t i
Revised: ( i n i t i

Original: [Figure]
Revised: [Figure]

Original: z
Revised: z

Original: [Figure]
Revised: [Figure]

Original: t
Revised: t

Original: i
Revised: i

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

Original: of the importance weights and the so-called values .
Revised: of the importance weights and the so-called values .

Original: Then, the architecture includes fully-connected sub-layers, residual connections [He et al., 2016], and layer normalization (LN), as illustrated in
Revised: Then, the architecture includes fully-connected sub-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 the sequence representation collapses doubly exponentially with depth, and both layer normalization and fully connected layers can only partially alleviate the speed of degeneracy.
Revised: [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.

Original: Under rank collapse , the model does not distinguish between representations of different tokens, which are perfectly aligned in feature space at initialization.
Revised: Under rank collapse , the model does not distinguish between representations of different tokens, which are perfectly aligned in feature space at initialization.

Original: However, the precise implications of rank collapse in Transformers are not fully understood.
Revised: However, the precise implications of rank collapse in Transformers are not fully understood.

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 by causing 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 by causing vanishingly small gradients of the queries and keys’ parameter matrices.

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

Original: 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.2 on the vanishing gradients for the queries and keys.
Revised: 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.2 on the vanishing gradients for the queries and keys.

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 evolution of the input sequence in a Transformer at initialization.
Revised: Firstly, we investigate under which conditions rank collapse can be avoided by studying the evolution of the input sequence in a Transformer at initialization.

Original: 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, hence stabilizing the rank of the propagating sequence.
Revised: 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, hence stabilizing the rank of the 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 affect differently 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 affect differently the gradient norm of the queries and keys compared to the values.

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

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

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

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 inputs of the self-attention and position-wise feed-forward sub-layers, leading to the POST-LN [Vaswani et al., 2017] and PRE-LN
Revised: 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

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-head unmasked 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-head unmasked 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: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: where the softmax function is applied independently across each row, and the superscript ℓ indexes the ℓ -th layer.
Revised: where the softmax function is applied independently across each row, and the superscript ℓ indexes the ℓ -th layer.

Original: The matrices W Q , W K ∈ R
Revised: The matrices W Q , W K ∈ R

Original: d v × d k and W V ∈ R
Revised: d v × d k and W V ∈ R

Original: d v × d v are learnable parameters, and each layer is initialized with an independent set of weights.
Revised: d v × d v are learnable parameters, and each layer is initialized with an independent set of weights.

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

Original: The complete Transformer block, in the absence of layer normalization, can be written recursively as:
Revised: The complete 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 ∈ 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 ∈ R d v × d v 2 are matrices of learnable parameters; we set X 0 :

Original: =
Revised: =

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

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

Original: At initialization, each weight is sampled independently from a distribution with zero-mean and variance σ 2 v = 1 d v for the values and feedforward weights 3 , and σ 2 k = 1 d k for the queries and keys.
Revised: At initialization, each weight is sampled independently from a distribution with zero-mean and variance σ 2 v = 1 d v for the values and feedforward weights 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, 2010] or “He”
Revised: [Glorot and Bengio, 2010] or “He”

Original: [He et al., 2015] initialization, commonly used in deep learning.
Revised: [He et al., 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 the tokens’ representations X ℓ converges to a rank-1 matrix in which all the representations are the same and equal to a vector x ∈ R d v , i.e. X ℓ → 1 n x ⊤ , where 1 d v is the vector with all ones in R d v .
Revised: [2021] proved that when the residual branches are omitted, the matrix of the tokens’ representations X ℓ converges to a rank-1 matrix in which all the representations are the same and equal to a vector x ∈ R d v , i.e. X ℓ → 1 n x ⊤ , where 1 d v is the vector with all ones in R d v .

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

Original: Indicating the inner product with the usual bracket notations ⟨· , ·⟩ , and the cosine of the angle between two tokens as θ k,k ′ , perfect alignment happens 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 alignment happens when ⟨ X ℓk , X ℓk ′ ⟩

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

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

Original: X ℓk (cid:13) (cid:13) cos θ k,k ′ with cos θ k,k ′ =
Revised: X ℓk (cid:13) (cid:13) cos θ k,k ′ with cos θ k,k ′ =

Original: 1 for all k, k ′ ∈ [ n ] .
Revised: 1 for all k, k ′ ∈ [ n ] .

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

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

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

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

Original: In its essence, given a loss function L : R n × d v → R , vanishing gradients occur when the norm of the gradient of the loss L with respect to the parameters of the network W — which we indicate as (cid:13)(cid:13) ∂ L ∂ W (cid:13)(cid:13) — is too small to provide enough backpropagating signal, thus hindering gradient-based optimization methods
Revised: In its essence, given a loss function L : R n × d v → R , vanishing gradients occur when the norm of the gradient of the loss L with respect to the parameters of the network W — which we indicate as (cid:13)(cid:13) ∂ L ∂ W (cid:13)(cid:13) — is too small to provide enough backpropagating signal, thus hindering gradient-based optimization methods

Original: One should explicitly write the layer dependence W Q,ℓ , W K,
Revised: One should explicitly write the layer dependence W Q,ℓ , W K,

Original: W V
Revised: W V

Original: W F W F
Revised: W F W F

Original: We at times suppress readability σ is
Revised: We at times suppress readability σ is

Original: ReLU function we set W F
Revised: ReLU function we set W F

Original: 1 to have variance 2 d v .
Revised: 1 to have variance 2 d v .

Original: Despite extensive research toward understanding and overcoming the problem in disparate contexts [Glorot and Bengio, 2010, 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].
Revised: Despite extensive research toward understanding and overcoming the problem in disparate contexts [Glorot and Bengio, 2010, 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].

Original: In our paper (Section 3.1), we show how vanishing gradient occurs in conjunction with the rank collapse issue identified by
Revised: In our paper (Section 3.1), we show how vanishing gradient occurs in conjunction with the rank collapse issue identified by

Original: Dong et al.
Revised: 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 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 X are lost/preserved as it propagates through a randomly initialized
Revised: 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 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 ⟩ , with k ̸ = k ′ .
Revised: X ℓ (cid:13)(cid:13) 2 and the expected inner product between different tokens E ⟨ X k , X ′ k ⟩ , with k ̸ = k ′ .

Original: 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., 2017], and vanishing/exploding gradients
Revised: 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., 2017], and vanishing/exploding gradients

Original: [Hanin, 2018].
Revised: [Hanin, 2018].

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

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

Original:  Theoretical Results
Revised:  Theoretical Results

Original: The goal of this section is twofold.
Revised: The goal of this section is twofold.

Original: In Lemma 3.1 and Theorem 3.2, we provide an explanation of the possible cause of vanishingly small gradients for queries and keys at initialization, namely the high correlations between the tokens representations in X ℓ as the depth increases.
Revised: In Lemma 3.1 and Theorem 3.2, we provide an explanation of the possible cause of vanishingly small gradients for queries and keys at initialization, namely the high correlations between the tokens representations in X ℓ as the depth increases.

Original: In Section 3.2, we show that the problem can be mitigated with an appropriate choice of the residual branch parameters α 1 and α 2 that inversely scales with the depth of the network L .
Revised: In Section 3.2, we show that the problem can be mitigated with an appropriate choice of the residual branch parameters α 1 and α 2 that inversely scales with the depth of the network L .

Original: Under the proposed scaling, the correlations are well behaved even in the infinite depth limit.
Revised: Under the proposed scaling, the correlations are well behaved even in the infinite depth limit.

Original: Finally, in Section 3.3 we analyze the scaling of the gradients with respect to other network’s parameters, and in 3.4 we draw some connections between our findings and optimization of Transformers.
Revised: Finally, in Section 3.3 we analyze the scaling of the gradients with respect to other network’s parameters, and in 3.4 we draw some connections between our findings and optimization of Transformers.

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 framework of 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 framework of matrix calculus [Magnus and Neudecker, 2019, Singh et al., 2021].

Original: In particular, we compare the expected Frobenius norm of the gradient of a self-attention layer with respect to its parameters
Revised: In particular, we compare the 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 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 × n matrix containing 1 n in each entry.
Revised: 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 × 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 , where 1 n × n is the matrix with all entries equal to 1 .

Original: Crucially, in Appendix A.5, we formally show that 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 the limit d k → ∞ .

Original: There, we also experimentally show that even in the more realistic case where d k = d v ≈ 512 , the empirical simulations provide a surprisingly faithful approximation of the theoretical insights presented in this paper.
Revised: There, we also experimentally show that even in the more realistic case where d k = d v ≈ 512 , the empirical simulations provide a surprisingly faithful approximation of the theoretical 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 following theorem, we compute the expected gradients of an attention layer at initialization, and set the basis for our following analysis.
Revised: In the following theorem, we compute the expected gradients of an attention layer at initialization, and set the basis for our following analysis.

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

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]

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 covariance matrix of the tokens’ representations.
Revised: At this stage, it is crucial to note that 1 n ( X ℓ ) ⊤ X ℓ − ¯ x ℓ (¯ x ℓ ) ⊤ is the centered empirical covariance matrix of the tokens’ representations.

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

Original: Introducing a differentiable loss function L : R n × d v → R , we make the statement on vanishing gradients more formal in the following theorem:
Revised: Introducing a differentiable loss function L : R n × d v → R , we make the statement on vanishing gradients 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]
Revised: [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 chain rule 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 chain rule and then bounding it by the product of the norms of each term of the chain.

Original: The final result holds with 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 holds with 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: The proof of the results presented in this section is deferred to Appendix A.1.
Revised: The proof of the results presented in this section is deferred to Appendix A.1.

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

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

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

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

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

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

Original: 1 that escaping this pathological landscape is harder in deeper nets where rank collapse persists across several layers.
Revised: 1 that escaping this pathological landscape is harder in deeper nets where rank collapse persists across several layers.

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.
Revised: Dong et al.

Original: [2021] showed that simply adding this architectural trick prevents rank collapse.
Revised: [2021] showed that simply adding this architectural trick prevents rank collapse.

Original: 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 → ∞ a Transformer can still lose rank unless the residual branches are adequately initialized.
Revised: 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 → ∞ a Transformer can still lose rank unless the residual branches are adequately initialized.

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

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

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

Original: [Figure]
Revised: [Figure]

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: (10) hence, under the depth scaling for the residual block parameters α 21 = ˜
Revised: (10) 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: 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 equations for 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 equations for 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: 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 the right-hand side will then depend on X ℓ only through its norm (cid:13)(cid:13)
Revised: The expression on the right-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 complete proof, together with the proof of Lemma 3.2, can be found in Appendix A.3.
Revised: The complete proof, 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 setting the alpha parameters to be O (1 / √ L ) allows both the norm of the propagating input and the inner products between different tokens to be approximately preserved.
Revised: The previous Lemma provides theoretical justification that scaling the residual branches by setting the alpha parameters to be O (1 / √ L ) allows both the norm of the propagating input and the inner products between different tokens to be approximately preserved.

Original: Hence, the information contained in the input is not lost, even in the infinite depth limit.
Revised: Hence, the information contained in 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 residual scaling (i.e. with α 

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

Original: = α 2 = 1 ) the correlation between the tokens quickly increases, and reaches perfect alignment in the infinite depth limit.
Revised: = α 2 = 1 ) the correlation between the tokens quickly increases, and reaches perfect 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 infinite depth.
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 infinite depth.

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 to formalize the influence of the 1 / √ L -scaling on the correlation between tokens’ representations by stating Theorem 3.3.
Revised: We are now ready to formalize the influence of the 1 / √ L -scaling on the correlation between tokens’ representations by stating 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 fixed layer 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 fixed layer 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: [Equation]
Revised: [Equation]

Original: By summing over the k, k ′ indexes, the resulting equation will depend on E [ C ( X ℓ )] and E (cid:13)(cid:13)
Revised: By summing over the k, k ′ indexes, the resulting equation will depend on E [ C ( X ℓ )] and E (cid:13)(cid:13)

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

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

Original: [Figure]
Revised: [Figure]

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

Original: =
Revised: =

Original: n ∥ X ∥ 2 , which holds in the degenerate case where all the input tokens are perfectly aligned.
Revised: n ∥ X ∥ 2 , which holds in the degenerate case where all the input tokens are perfectly aligned.

Original: In Appendix A.4, we give precise formulas for the 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.
Revised: In Appendix A.4, we give precise formulas for the 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.

Original: 3 (left)).
Revised: 3 (left)).

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

Original: 4, we empirically show that in 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 show that in the presence of the 1 / √ L -scaling, layer normalization (either PRE or POST) does not significantly affect the evolution of the correlations.

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

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

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

Original: Note that the 1 / √ 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.
Revised: Note that the 1 / √ 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.

Original: Finally, note that by setting ˜ α 1 , ˜ α 2 = 0 , we recover the so called "ReZero" initialization [Bachlechner et al., 2021].
Revised: Finally, note that by 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 for wider range of values for ˜ α 1 , ˜
Revised: In this context, the 1 / √ L scaling extends this framework as it allows for wider range of values for ˜ α 1 , ˜

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

Original: ReLU extension.
Revised: ReLU extension.

Original: We mention here that extending these results from the linear activation to the
Revised: We mention here that extending these results from the linear activation to the

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

Original: Exact formulas can be found only in the case of one ReLU layer 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 analyzed here, the linear activation function provides a bound on the correlation with respect to the ReLU case.
Revised: However, in the context of rank collapse analyzed here, 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 ReLU case (for instance, see the contraction argument in Nachum et al.
Revised: In fact, correlations are exactly preserved in expectation in the linear case, but increase in the ReLU case (for instance, see the contraction argument in Nachum et al.

Original: [2021] below Equation (2)).
Revised: [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).
Revised: 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 that d k = d v .
Revised: In this section, we drop the superscript ℓ as it is obvious from context and assume for simplicity that d k = d v .

Original: 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 ρ ℓii ′ =
Revised: 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 ρ ℓii ′ =

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

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

Original: [Equation]
Revised: [Equation]

Original: To see that this equation satisfies our definition of the correlation metric, note that E [ ∥ X i ∥ 2 ] = dσ 2 x and 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 x and 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. (6)) simplifies to
Revised: Then, the expected norm of the gradients for the values (Eq. (6)) simplifies to

Original: [Equation]
Revised: [Equation]

Original: By making the additional assumption that the norm and the correlation propagate independently, the respective norm for the queries — and symmetrically the keys — (Eq.
Revised: By making the additional assumption that the norm and the correlation propagate independently, the respective 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] to compute higher-order moments.
Revised: In Appendix A.2 we provide a rigorous proof, that relies on Isserlis theorem [Isserlis, 1918] to compute higher-order moments.

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

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

Original: This provides an additional interpretation of the successful use of layer normalization, as in Xiong et al.
Revised: This provides an additional interpretation 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 respect to ρ compared to the values.
Revised: Importantly, note how the queries/keys have opposite monotonic functional dependence with respect to ρ compared to the values.

Original: As revealed by Theorem 3.3 and Fig.
Revised: As revealed by Theorem 3.3 and Fig.

Original: 3 (center), inappropriate scaling of the residual branches can already lead to this phenomenon even in a relatively shallow network.
Revised: 3 (center), inappropriate scaling of the 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 sequence length 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 sequence length n due to the self-attention operation itself.

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

Original: 3.4 Connections to Optimization and Adaptive Methods
Revised: 3.4 Connections to Optimization and Adaptive Methods

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 success of adaptive optimization algorithms, as illustrated in Fig.
Revised: W Q , W K and W V , might explain the success of 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: Notice that the effective learning rate is increasingly larger (with depth) for the queries compared to the values — as postulated by our theory — and this difference is remarkably constant throughout training.
Revised: Notice that the effective learning rate is increasingly larger (with depth) for the queries compared to the values — as postulated by our theory — and this difference is remarkably constant throughout training.

Original: Hence, we conjecture that the success of adaptive methods in Transformers’ training can be partially explained by the need to fix this
Revised: Hence, we conjecture that the success of adaptive methods in Transformers’ training can be partially explained by the need to fix this

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: disproportionate gradient’s magnitude.
Revised: disproportionate gradient’s magnitude.

Original: To test this hypothesis we propose a simple architectural modification, an inverse temperature scaling τ ∈ R inside the softmax:
Revised: To test this hypothesis we propose a simple architectural modification, an inverse temperature scaling τ ∈ R inside the softmax:

Original: [Equation]
Revised: [Equation]

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

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

Original: In Section C.2, we detail how one can choose τ such that the magnitude of the gradients as derived in
Revised: In Section C.2, we detail how one can choose τ such that the magnitude of the gradients as derived in

Original: Equation 17 and 18 is approximately matched at initialization.
Revised: Equation 17 and 18 is approximately matched at initialization.

Original: We evaluate our proposal, consisting of residual scaling and the aforementioned inverse temperature parameters, on the widely used IWSLT
Revised: We evaluate our proposal, consisting of residual scaling and the aforementioned inverse temperature parameters, on the widely used IWSLT

Original: German-to-English (De-En) benchmark translation task.
Revised: German-to-English (De-En) benchmark translation task.

Original: All details regarding the experimental setup and the choice of inverse temperature used are provided in Appendix C. We train a
Revised: All details regarding the experimental setup and the choice of inverse temperature used are provided in Appendix C. We train a

Original: Transformer encoder-decoder of varying depth with stochastic gradient descent (SGD), after removing all normalization layers and adequately initializing the residual connections.
Revised: Transformer encoder-decoder of varying depth with stochastic gradient descent (SGD), after removing all normalization layers and adequately initializing the residual connections.

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

Original: We compare against the following methods that make use of Adam;
Revised: We compare against the following methods that make use of Adam;

Original: POST-LN and PRE-LN refer to the aforementioned alternatives to apply layer normalization.
Revised: POST-LN and PRE-LN refer to the aforementioned alternatives to apply layer normalization.

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

Original: We report the average BLEU score [Papineni et al., 2002] across 5 runs in Fig.
Revised: We report the average BLEU score [Papineni et al., 2002] across 5 runs 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 cases surpassing any results achieved by the Adam optimizer.
Revised: Our proposed method considerably improves training with SGD, keeping up and in some cases surpassing any results achieved by the Adam optimizer.

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

Original: We leave for future work to further investigate modifications or alternatives to the self-attention operation.
Revised: We leave for future work to further investigate modifications or 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 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].
Revised: 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].

Original: The 1 / √ L scaling scheme has been investigated in the literature for the stabilization of residual networks [Hanin and
Revised: The 1 / √ L scaling scheme has 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 representations of 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 of random feed-forward neural networks at initialization [Daneshmand et al., 2020, 2021].

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

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

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

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

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

Original: Our analysis of the forward pass shows that its inclusion is not strictly necessary for the purpose of controlling the norm of the representations.
Revised: Our analysis of the forward pass shows that its inclusion is not strictly necessary for the purpose of controlling the norm of 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 the empirically observed discrepancy between the magnitude of the gradients of the queries and the values, which Liu et al.
Revised: 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.

Original: [2020] hypothesize to be one of the causes of the success of adaptive methods in training Transformers
Revised: [2020] hypothesize to be one of the causes of the success of adaptive methods in training Transformers

Original: [Liu et al., 2019, Zhang et al., 2020, Huang et al., 2020].
Revised: [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, Wang et 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, Wang et al., 2022].

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

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 correlation in 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 correlation in 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 alone cannot 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 alone cannot 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 gradients of queries/keys compared to the values’ parameters.

Original: Hence, we conclude that one of the strengths of Transformers lies in their carefully designed architecture together with an adequate initialization.
Revised: Hence, we conclude that one of the strengths of 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 efficacy of Adam compared to SGD in training Transformers arises from the disproportionate magnitude of gradients as postulated by our theory.
Revised: 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.

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

Original: Acknowledgements
Revised: Acknowledgements

Original: Sidak Pal Singh would like to acknowledge the financial support from Max Planck ETH Center for
Revised: Sidak Pal Singh would like to acknowledge the financial support from Max Planck ETH Center for

Original: Learning Systems and the travel support from ELISE (GA no 951847).A A.1 Backward [Table]
Revised: Learning Systems and the travel support from ELISE (GA no 951847).

Original: : Proofs of Lemma 3.1 and Theorem 3.2 . . . . . . . . . . . . .

Original: A.2 Gradient Analysis of Section 3.3 . . . . . . . . . . . . . . . . . . . . . . . . .

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

Original: A.4 Proof of Theorem 3.3: Correlations are Preserved under Residual Scaling . . . .

Original: A.5 Motivation for Assumption 3.1 . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: B Additional Results

Original: B.1 On the Roles of

Original: [Table]

Original: Scaling of the

Original: [Table]

Original: and Layer Normalization . .

Original: [Table]

Original: [Table]

Original: Further Empirical Assessment of

Original: [Table]

Original: [Table]

Original: [Table]

Original: Empirical Verification of the

Original: Analysis of

Original: C

Original: [Table]

Original: Setup

Original: C.1 Toy Example . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: C.2 Tempering the softmax . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: C.3 Translation Task . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: A Proof of Theorems

Original: Recall the defining equations of a Transformer:

Original: [Equation]

Original: where the self attention layers are defined as follows:

Original: [Equation]

Original: We remark that when it is clear from the context, we suppress the ℓ index to improve readability.

Original: Initialization : Recall that we initialize our weights with the so called “Xavier”

Original: [Glorot and Bengio, 2010] or “He”

Original: [He et al., 2015] initialization: each weight is sampled independently from a distribution with zero-mean and variance σ 2 v = 1 d v for the values and feedforward weights and σ 2 k = 1 d k for the queries and keys.

Original: Kronecker Delta : we introduce the Kronecker Delta notation

Original: [Equation]

Original: and, similarly: δ a ̸ = b = 1 − δ ab .

Original: Notation : in this section, we also adopt the following shorthand notation for the argument of the softmax M :=

Original: [Equation]

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

Original: Recall that for a matrix X , we use X k to indicate its k -th row.

Original: Finally, we indicate with 1 n × n the n × n matrix with all ones, with 1 n the columns vector with all ones, and with I n the n -dimensional identity matrix.

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

Original: In this section, we now look at the proofs of Lemma 3.1 and Theorem 3.2.

Original: We will first introduce our notation for the gradients, as well as some useful properties of the Kronecker product.

Original: A.1.1 Preliminaries

Original: For the gradients, we avoid directly working with tensors by vectorizing the matrices in a row-wise fashion ( vec r ) and arranging the Jacobian in the numerator layout.

Original: More formally,

Original: [Equation]

Original: Alongside this, we use the following rule ( ⊗ is the Kronecker product):

Original: [Equation]

Original: For the proof of this rule, we refer to Singh et al.

Original: [2021], and to Magnus and Neudecker [2019]

Original: for a complete introduction to matrix calculus.

Original: We will also use the following well-known properties of the Kronecker product.

Original: Lemma A.1.

Original: Given the matrices A ∈ R n × m , B ∈ R p × q , C ∈ R m × r , D ∈ R q × s , then the following holds:

Original: [Equation]

Original: and

Original: [Equation]

Original: A.1.2 Proof of Lemma 3.

Original: In Lemma A.2 and Lemma A.3 we compute the gradients with respect to the queries, values and X , respectively.

Original: Then we use these results to prove Lemma 3.1 by computing the expectation of the

Original: Frobenius norms.

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

Original: The gradients of the self attention layer defined in Eq.

Original: (1) have the following form:

Original: [Equation]

Original: where the gradients of the softmax with respect to its inputs are as follows:

Original: [Equation]

Original: and where ∂ A i ∂ M ⊤ i = diag( A i ) − A i

Original: A ⊤ i with A i being the i -th row of A in column vector format.

Original: Finally, note that under the uniform-attention assumption, Eq.

Original: (23) simplifies to:

Original: [Equation]

Original: Proof.

Original: Let’s start with the simple case of the values’ weights W V .

Original: Using the rule in Eq.

Original: (20), it is immediate that:

Original: [Equation]

Original: For the queries, a simple application of the chain rule and then again Eq.

Original: (20) gives:

Original: [Equation]

Original: which is the desired results.

Original: Finally, for the gradients of the softmax note that:

Original: [Equation]

Original: By writing the above expression in the matrix notation described above, we obtain the desired result.

Original: More specifically, the block diagonal structure is given from the term δ ip which stems from the fact that the softmax is applied row-wise.

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

Original: The gradients of the self attention layer with respect to the embedding matrix X defined in Eq.

Original: (1) have the following form

Original: [Equation]

Original: where the gradients of the softmax with respect to its inputs are denoted by ∂ A ∂ M as before.

Original: Proof.

Original: Remember that we defined S = softmax( 1 √ d k XW Q W K ⊤ X ⊤ ) XW V .

Original: Alongside with our previous shorthands A , M , let us define the remaining XW V as a matrix T , so that S =

Original: A T .

Original: Both

Original: A and T are functions of X .

Original: So the matrix differential can be written as:

Original: [Equation]

Original: Next, we use the matrix differential and then the identification theorem of matrix derivatives to compute the matrix gradient ∂ A ∂ X

Original: [Equation]

Original: Vectorizing both sides:

Original: [Equation]

Original: Recall, for an arbitrary matrix B ∈ R m × n

Original: [Equation]

Original: the commutation matrix K mn transforms columnwise vectorization into rowwise vectorization.

Original: More precisely,

Original: K mn vec

Original: [Equation]

Original: and vec c ( B ) = vec r ( B ⊤ ) .

Original: Therefore, for rowwise vectorization

Original: [Equation]

Original: we have a similar result:

Original: [Equation]

Original: where in the last line we used the fact the commutation is a permutation matrix, so K − 1 mn = K ⊤ mn =

Original: K nm .

Original: Thus, we get the required matrix derivative as follows:

Original: [Equation]

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

Original: [Equation]

Original: Plugging this into the above Eq.

Original: (29), we get:

Original: [Equation]

Original: As a sanity check, we can calculate if the shapes of the matrices are consistent.

Original: LHS should be a nd × nd matrix, while the constituent matrices of the first term on RHS: I n ⊗ W V ⊤ X ⊤ ∈ R nd × n 2 , ∂ A ∂ M ∈ R n 2 × n 2 , the additive term next to it is a n 2 × nd matrix, and the second term on RHS is a

Original: Kronecker product of a n × n and a d × d matrix.

Original: Lemma 3.1.

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

Original: Under the uniform-attention assumption, we have

Original: [Equation]

Original: Proof.

Original: Here, we suppress the index ℓ .

Original: Gradient with respect to the values matrix.

Original: Recall that from Lemma A.2 we have that:

Original: [Equation]

Original: By direct computation:

Original: [Equation]

Original: Gradients with respect to the queries/keys matrix.

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

Original: [Equation]

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

Original: [Equation]

Original: where in the last step we have used twice the property of the Kronecker product in Eq.

Original: (22) of

Original: Lemma A.1.

Original: Hence,

Original: [Equation]

Original: where we have used the property on the trace of the Kronecker product (Lemma A.1, Eq.

Original: Note that if we are conditioning on X , then we only have to take the expectation of the last term with respect to the weights W K and W V .

Original: Let us call L := I n − 1 n 1 n × n for notation simplicity.

Original: Note: for a matrix W ∈ R

Original: d × d whose entries w ij ∼ N (0 , σ 2 ) , then E WW ⊤ =

Original: dσ 2 I d .

Original: Thus, exchanging the order of trace and expectation, we can write:

Original: [Equation]

Original: where, ¯ x = 1 n X ⊤ 1 n ∈ R d is the mean embedding.

Original: Multiply this by

Original: [Equation]

Original: d k n 2 ∥ X ∥ 2 F to get the final answer.

Original: Gradient with respect to the input.

Original: Plugging in the values of ∂ A ∂ M and A under the uniform-attention assumption into Eq.

Original: (25) gives rise to the following:

Original: [Equation]

Original: Let’s refer to the matrices on the right-hand side as A 1 , A 2 , A 3 respectively.

Original: We compute the expected squared Frobenius norm of these as follows:

Original: For A 3 :

Original: [Equation]

Original: Similarly, for A 1 :

Original: [Equation]

Original: Finally, for A 2 :

Original: [Equation]

Original: where in the second line we have taken the expectation inside and used the fact that K nn , being a commutation matrix, is orthogonal.

Original: Then, by simple properties of Kronecker product and cyclic property of trace, we have the result, which is the same as that for A 1 .

Original: Finally, by the triangle inequality

Original: [Equation]

Original: This completes the proof.

Original: A.1.3 Proof of Theorem 3.

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

Original: Suppose that the uniform-attention assumption holds.

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

Original: [Equation]

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

Original: 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. (cid:13) (cid:13) X ⊤ X − n ¯ x ¯ x ⊤ (cid:13) (cid:13) = 0 ), the gradient with respect to the values and the input does not (see Theorem 3.1).

Original: From this simple remark, we can conclude that, even in the rank collapse regime, information still propagates in the backward pass.

Original: In Section 3.4 (main paper), we show that even if gradients effectively propagate, the phenomenon studied in this theorem still greatly affects training.

Original: Proof.

Original: By using the chain rule and the fact that for two matrixes A , B we have that ∥ AB ∥ 2 F ≤ ∥

Original: A ∥ 2 F ∥ B ∥ 2 F , we can upper bound the gradient as:

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Table]

Original:  

Original: [Table]

Original: where we recall that Z ℓ = α 1 S ℓ + X ℓ and in the last step we have used that X ℓ does not depend on

Original: W Q

Original: [Table]

Original: ℓ , hence the gradient vanishes.

Original: By taking expectation and using the tower property

Original: [Table]

Original: we have that:

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: where the expectations are taken with respect to X ℓ for the outer one and conditioning on X ℓ for inner one.

Original: Indeed, the first three terms only depend on the network values after X ℓ .

Original: Now, a repeated application of the tower property in G ( X ℓ ) , together with the results on the gradients of Lemma 3.1, easily shows that G ( X ℓ ) stays bounded under our hypothesis.

Original: To see this one can also simply note that, since the softmax and its derivatives are almost surely bounded, the boundedness of G ( X ℓ ) is implied by an analogous statement for a vanilla linear MLP (i.e removing the softmax).

Original: In this setting, the random variable inside the expectation in G ( X ℓ ) is a finite linear combination of Gaussian products — which has bounded expectation.

Original: All in all, we have that

Original: [Equation]

Original: , where B X ℓ is an almost-surely-bounded function of X ℓ .

Original: Hence, to show that E (cid:13)(cid:13)

Original: ∂ L ∂ W Q,ℓ

Original: [Equation]

Original: 2 F = 0 , we now just need to show that:

Original: [Equation]

Original: under the rank-1 hypothesis for X ℓ .

Original: Let X ℓ 1 , . . .

Original: X ℓn ∈

Original: R

Original: d v be the representations for the n tokens.

Original: Under the rank-1 assumption, each token can be written as a multiple of a single vector x ∈ R d v , and hence there exists a 1 , . . .

Original: , a n ∈ R such that X 1 = a 1 x , . . .

Original: X n = a n x .

Original: From Lemma 3.1, we know that:

Original: [Equation]

Original: The mean token simplifies to ¯ x l = x n (cid:80) k a k and hence (cid:0) ¯ x ℓ (¯ x ℓ ) ⊤ (cid:1) ij = 1 n 2 ( (cid:80)

Original: k

Original: a k ) 2 x i x j .

Original: Similarly, (cid:0) ( X ℓ ) ⊤ X ℓ (cid:1) ij =

Original: (cid:80) k

Original: a 2 k x i x j .

Original: [-If furthermore all the coefficients a i are the same (which corresponds to the rank collapse assumption X ℓ = 
-]
Revised: [-If furthermore all the coefficients a i are the same (which corresponds to the rank collapse assumption X ℓ = 
-]

Original: n x T analyzed here), then it is easy to see that (cid:0) ( X ℓ ) ⊤ X ℓ (cid:1) ij − n (cid:0)

Original: A.2 Gradient Analysis of Section 3.

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

Original: [Equation]

Original: 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]:

Original: Theorem A.1 (Isserlis) .

Original: Let X 1 , . . .

Original: X m be m zero-mean Gaussian random variables.

Original: Then:

Original: [Equation]

Original: where P 2 m is the set of all the possible pairings of the indexes 1 , . .

Original: , m .

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

Original: [Equation]

Original: Now we can prove Eq.

Original: (17) , which we re-state

Original: here:

Original: [Equation]

Original: Also, from Eq.

Original: (6) we have that

Original: [Equation]

Original: Now,

Original: [Equation]

Original: Each ¯ x ℓi = 1 n (cid:80) nk =1 X ℓki is equally distributed with mean

Original: [Equation]

Original: and variance

Original: [Equation]

Original: Finally we get

Original: [Equation]

Original: We know prove Eq.

Original: (18), which reads:

Original: [Equation]

Original: For the queries (and the keys respectively), recall from Eq.

Original: (7) that

Original: [Equation]

Original: To proceed, we drop the superscript ℓ and we make the additional assumption that ∥ X ∥ 2 F is uncorrelated from the correlation magnitude ∥ X ⊤ LX ∥

Original: [Equation]

Original: x ¯ x ⊤ ∥ 2 F .

Original: Let us proceed with an expansion:

Original: [Equation]

Original: Now, we have 2 cases: if i ̸ = j , which gives d ( d − 1) equal terms, we need to compute

Original: [Equation]

Original: where ( i, j ) is any tuple with i ̸

Original: =

Original: j and we used uncorrelation of different dimensions.

Original: Otherwise, we get d each equal to

Original: [Equation]

Original: where i is any index.

Original: Term A. Note that

Original: E [ X ai X a ′

Original: i ] · E [ X bj X

Original: b ′ j ] =   σ 4 x a = a ′ , b = b ′ , n 2 terms ρσ 4 x a = a ′ , b ̸ = b ′ , n 2 ( n − 1) terms ρσ 4 x a ̸ = a ′ , b = b ′ , n 2 ( n − 1) terms ρ 2 σ 4 x a ̸ = a ′ , b ̸ = b ′ , n 2 ( n − 1) 2 terms .

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

Original: [Equation]

Original: where we leveraged the following direct calculations:

Original: [Equation]

Original: Next, we compute

Original: [Equation]

Original: Finally, similar computations also lead to the last term.

Original: To follow the calculations, we invite the reader to draw the matrix L and to hide the columns over which summations are not performed:

Original: [Equation]

Original: All in all, we get:

Original: [Equation]

Original: Term B. We make use of Isserlis theorem, stating that:

Original: [Equation]

Original: By using our independence assumptions

Original: [Equation]

Original: we get:

Original: [Equation]

Original: Similarly for Q 2 and Q 3 :

Original: [Equation]

Original: and

Original: Q

Original: [Equation]

Original: Hence,

Original: [Equation]

Original: =

Original: [Equation]

Original: Let’s study it term by term.

Original: We will also use L ab =

Original: ( δ ab − 1 n ) , and so L ab L a ′ b ′ =

Original: ( δ ab δ a ′ b ′ − δ ab n − δ a ′ b ′ n + 1 n 2 ) .

Original: First term : we have that σ 4 x

Original: [Equation]

Original: ab

Original: [Equation]

Original: 1 which is equal to (omitting the constant σ 4 x

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: n −

Original: [Equation]

Original: Second term : we have that σ 4 x (cid:80) aa

Original: [Equation]

Original: b

Original: [Equation]

Original: b

Original: [Equation]

Original: L ab L a ′ b ′ Q 2 which is equal to (omitting the constant σ 4 x ):

Original: [Equation]

Original: Third term: we have that σ 4 x (cid:80) aa ′

Original: b

Original: [Equation]

Original: b

Original: [Equation]

Original: L ab L a

Original: [Equation]

Original: b ′ Q 3 which is equal to (omitting the constant σ 4 x ): = (cid:88)

Original: [Equation]

Original: [Equation]

Original: b

Original: [Equation]

Original: b ′

Original: [Equation]

Original: ab

Original: [Equation]

Original: (cid:0) n ( n

Original: [Equation]

Original: [Equation]

Original: ( n

Original: [Equation]

Original: 1)( n

Original: [Equation]

Original: [Equation]

Original: ( n

Original: [Equation]

Original: 1) 2 (cid:1)

Original: + ρ ( −

Original: [Equation]

Original: ( n − 1)

Original: [Equation]

Original: ( n − 1))

Original: + n −

Original: [Equation]

Original: = ρ 2 ( n − 1) − 2 ρ ( n − 1) + ( n − 1)

Original: [Equation]

Original: Summing all the three terms

Original: [Equation]

Original: we get:

Original: [Equation]

Original: Plugging in the values of A and B we get:

Original: [Equation]

Original: and finally assuming Xavier initialization

Original: [Equation]

Original: A.3 Forward Pass: Proofs of Lemma 3.2 and 3.

Original: First

Original: [Equation]

Original: we characterize the evolution of the correlations between tokens X k , X k ′ with depth, under the assumptions of Theorem 3.3, namely uniform-attention assumption

Original: [Equation]

Original: and the adoption of a linear activation.

Original: Lemma A.4 (Expectation of Linear Layers) .

Original: Let D = XW , where

Original: W ∈ R d × d is a random matrix with i.i.d random entries with variance σ 2 = 1 d and X ∈ R n × d is a fixed matrix:

Original: [Equation]

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

Original: [Equation]

Original: Proof.

Original: [Equation]

Original: Lemma A.5 (Expectation of skip connection) .

Original: Let A , B ∈ R p × q .

Original: Let D := α A + B with E [ A | B ] = 0 and α ∈ R .

Original: Then:

Original: [Equation]

Original: holds for all i, i ′ ∈

Original: [ p ] , j ∈ [ q ] .

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

Original: [Equation]

Original: Proof.

Original: [Equation]

Original: where using iterated expectations α E

Original: [ A i ′ j B ij ] =

Original: α E [ E [ A i ′

Original: j | B ] B ij ]]

Original: = 0 and identically α E

Original: [ A ij B i ′ j ] = 0 .

Original: Lemma A.6 (Expectation of Attention Layers) .

Original: Under the uniform-attention assumption:

Original: [Equation]

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

Original: [Equation]

Original: Proof.

Original: Note that under the uniform-attention assumption:

Original: [Equation]

Original: Hence, using the fact that the weights are i.i.d with variance σ 2 v = 1 d v :

Original: E [

Original: [Equation]

Original: Lemma 3.2 (Propagation of inner products) .

Original: Let C ( X ℓ )

Original: =

Original: (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:

Original: [Equation]

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

Original: α 1 L , α 22 = ˜

Original: α 2 L with ˜

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

Original: [Equation]

Original: Proof.

Original: First, note that for the residual blocks we have that E

Original: [ Y ℓkj |

Original: Z ℓk ′

Original: [-j ] = 
-]
Revised: [-j ] = 
-]

Original: due to the independence assumption on the feedforward weights, and similarly E [ S ℓkj |

Original: X ℓk ′ j ] = 0 .

Original: Hence, we can use Lemma

Original: A.5 in both the skip connections of the Transformer architecture.

Original: Therefore, using Lemma A.5 (skip),

Original: Lemma A.4 (linear) and Lemma A.6 (attention):

Original: [Equation]

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

Original: For the limit as L → ∞ , simply note that:

Original: [Equation]

Original: with i ∈ { 1 , 2 } .

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

Original: Lemma 3.3 (Propagation of the norm) .

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

Original: [Equation]

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

Original: α 1 L , α 22 = ˜

Original: α 2 L with ˜

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

Original: [Equation]

Original: Proof.

Original: The proof is in the same spirit as Lemma 3.2 but slightly more involved.

Original: Again, using Lemma

Original: A.5 in both the skip connections of the Transfomer architecture.

Original: Therefore, using Lemma A.5 (skip),

Original: Lemma A.4 (linear) and Lemma A.6 (attention):

Original: [Equation]

Original: where in the second to last step we have used Lemma 3.2 and in the last step we have unrolled the recursion for (cid:13)(cid:13) X ℓ (cid:13)(cid:13) 2 F until the input layer.

Original: For the second part, we now show that for a network of L layers, the choice α 21 = ˜

Original: α 1 L and α 22 =

Original: ˜

Original: α 2 L stabilizes the norm of the activations in the forward pass.

Original: [-Using the product law for the limits, we can study the converges of ( α 22 + 1) ℓ and α 1 (cid:80) ℓk =
-]
Revised: [-Using the product law for the limits, we can study the converges of ( α 22 + 1) ℓ and α 1 (cid:80) ℓk =
-]

Original: ( α 21 + 1) k separately.

Original: Let i ∈ { 1 , 2 } .

Original: For the latter term we have that:

Original: [Equation]

Original: while for the former term we have that lim L →∞ ( ˜ α i L + 1) ℓ = e ˜ α

Original: i .

Original: Hence, the norm of the representations converges to:

Original: [Equation]

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

Original: Remark : note that C ( X ) =

Original: (cid:80) k,k ′ (cid:80)

Original: j X kj X kj ′

Original: = (cid:80) j

Original: ( (cid:80) k

Original: [Equation]

Original: k ′ X kj X k ′ j )

Original: = (cid:80) j

Original: ( (cid:80) k

Original: X kj ) 2 = n 2 ∥ ¯ x ∥ 2 .

Original: A.4 Proof of Theorem 3.3: Correlations are Preserved under Residual Scaling

Original: Theorem 3.3.

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

Original: X k ∥ =

Original: ∥ 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: = ˜

Original: α 1 L , α 22 = ˜

Original: α 2 L with ˜

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

Original: [Equation]

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

Original: [Equation]

Original: Proof.

Original: Due to the rotational symmetries of the Gaussian random matrices, if the input tokens have the same norm, then the expected norm at layer ℓ ∈ [ L ] is also the same across the token’s representations.

Original: Hence, we can write E (cid:13)(cid:13) X ℓ (cid:13)(cid:13) 2 F = n E (cid:13)(cid:13) x ℓ (cid:13)(cid:13) 2 , where (cid:13)(cid:13) x ℓ (cid:13)(cid:13) 2 is the norm of every token at layer ℓ .

Original: Furthermore, by definition of our correlation coefficient ρ lkk ′ , we have that E ⟨ X ℓk , X ℓk ′ ⟩ = ρ ℓ kk ′ E

Original: [Equation]

Original: x ℓ (cid:13) (cid:13) 2 .

Original: By summing over the indexes k, k ′ , we can expand the relation as:

Original: [Equation]

Original: By solving for ρ ℓ , we have that:

Original: [Equation]

Original: Now we plug in the expressions for E C ( X ℓ ) and E (cid:13)(cid:13) X ℓ (cid:13)(cid:13) 2 with the aid of Lemma 3.2 and Lemma 3.3, respectively.

Original: Finally, by taking the limits with respect to L , we get the desired result.

Original: A.5 Motivation for Assumption 3.

Original: We motivate here the following assumption, stated in the main paper.

Original: This assumption is crucial to compute expectations involving the softmax function.

Original: Assumption 3.1 (Uniform attention) .

Original: We assume that A ℓ = 1 n 1 n × n ,

Original: Theoretical analysis.

Original: We first show that this assumption holds when taking d k to infinity, keeping d v fixed.

Original: Lemma A.7.

Original: Consider initializing each entry of W Q ∈ R d v × d k and W K ∈ R

Original: d v × d k independently with variance σ 2 k = 2 / ( d v + d k ) — i.e. Glorot initialization [Glorot and Bengio, 2010].

Original: Let

Original: M = 1 √ d k X ℓ W Q,ℓ W K,ℓ ⊤ X ℓ ⊤ ; for any ( i, j ) ∈ [ n ] × [ n ] we have

Original: [Equation]

Original: While keeping d v < ∞ fixed, taking d k to infinity yields

Original: [Equation]

Original: In other words, M converges to 0 n × n in L 2 as d k → ∞ .

Original: Proof.

Original: First, note that

Original: [Equation]

Original: Since W Q is independent from W K at initialization, E [ M i,j | X ] = 0 .

Original: Next, we compute

Original: [Equation]

Original: This concludes the proof.

Original: The following classical result implies almost sure convergence of the softmax matrix as d k → ∞ .

Original: Lemma A.8 (Borel-Cantelli) .

Original: Let ( X i ) be a sequence of random variables.

Original: If for any ϵ >∞ (cid:88) i =

Original: [Equation]

Original: then X i converges to X almost surely 4 .

Original: Theorem A.2 (Almost-sure convergence) .

Original: Consider initializing each entry of W Q ∈ R d v × d k and

Original: W K ∈ R d v × d k independently with variance σ 2 k = 2 / ( d v + d k ) — i.e. Glorot initialization [Glorot and Bengio, 2010].

Original: Let d v < ∞ be fixed, as d k → ∞ we have that, for any X ,

Original: [Equation]

Original: and

Original: [Equation]

Original: Proof.

Original: Thanks to Lemma A.7 and Markov Inequality, we have fast convergence in probability: for any fixed X ,

Original: [Equation]

Original: Borel Cantelli then directly yields almost sure convergence of M to 0 n × n as d k → ∞ .

Original: Next, note that both A and ∂ A ∂ M are continuous functions of A , hence we can apply standard continuity event-per-event.

Original: For almost every ω ∈ Ω ,

Original: [Equation]

Original: Hence A → 1 n 1 n × n almost surely.

Original: This can also be seen as a simple application of the continuous mapping theorem.

Original: The same reasoning yields almost sure convergence of

Original: [Equation]

Original: to the corresponding limiting quantity.

Original: [Figure]

Original: Empirical analysis.

Original: 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 A and we calculate 1 n 2 || A ℓ − 1 n 1 n × n || 2 F , i.e. its average (entry-wise) distance from a uniform matrix with entries all equal to 1 /n .

Original: For each value of d k , we repeat this calculation 200 times, each time with different random weight matrices.

Original: Fig.

Original: 7 displays how the ||

Original: A − 1 n 1 n × n || 2 F averaged over 200 runs, tends to zero with a trend inversely proportional to d 2 k , as predicted by our theoretical analysis.

Original: B Additional Results B.1 On the Roles of the 1 / √ L -Scaling of the Residuals and Layer Normalization

Original: We present some additional results on the propagation of the norm and the correlations in Figure 8.

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

Original: Furthermore, we demonstrate the propagation of the correlation and the gradient norms for the PRE-LN configuration in Figure 9.

Original: As also hinted in the main text, in Figure 4, the increase in correlation with depth for PRE-LN is much less wild.

Original: This also results in better stabilized gradients for the queries and keys’ parameters.

Original: We also observe the opposite trend for the gradients of the values, in relation to the POST-LN case in

Original: Figure 3.

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

Original: We plan to investigate this dependence more in future work.

Original: We provide the detailed experimental setup for Figures 3 and 9 in Table 2.

Original: Finally, as hinted in the main text, replacing the linear activation function with ReLU leads to an even faster increase in correlations (for instance, see the contraction argument in Nachum et al.

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

Original: We empirically evaluate that, by comparing the propagation of correlation (as in Figure 4) also for a network with ReLU activation functions.

Original: Results are shown in Figure 10.

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: B.2 Further Empirical Assessment of Assumption 3.

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

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

Original: The results are shown in Fig 11.

Original: Note that when the standard deviation of the input is set to 1 / √ d , then the uniform-attention assumption provide an excellent approximation to the common Xavier-initialization.

Original: On the contrary, we observe a deviation when the standard deviation of the input is increased.

Original: Also, note how as the depth increases, the distribution becomes more heavy-tailed.

Original: This heavy-tailedness

Original: [Figure]

Original: Linear

Original: [Figure]

Original: was recently formally shown for standard MLPs with and without ReLU activation [Noci et al.

Original: Zavatone-Veth and Pehlevan, 2021].

Original: For the verification of the assumption in the backward pass we additionally show in Fig.

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

Original: Ground-truth gradients are calculated with automatic differentiation, and they are compared with our theoretical results based on Assumption 3.1.

Original: As shown in Fig.12, our theoretical predictions show a very good agreement with the true gradients.

Original: Again, we notice that the smaller the values of the input standard deviation the tighter the agreement of the theory with the simulations.

Original: Intuitively, a higher input variance causes the argument of the softmax to have a large range of values.

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

Original: B.3 Empirical Verification of the Gradient Analysis of Section 3.

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

Original: Figureillustrates the true dependence and Figure 14 the one expected by the theory based on our assumptions.

Original: In short, the main takeaways are the following.

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

Original: • The dependence on the variance of the input σ 2 x is different ( y -axis in the global plot), being linear for the values and cubic for the queries.

Original: This highlights the importance of a stabilized forward pass and provides another explanation regarding the successful use of layer norm in Transformers.

Original: • 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 ρ (compare the first column where ρ = 0 to the rest).

Original: [Figure]

Original: C Experimental Setup

Original: Here we provide more details regarding the experimental setup

Original: [Figure]

Original: I n p u t s t d 0 .

Original: Cosine angle 0.

Original: [-I n p u t s t d 0 .
-]
Revised: [-I n p u t s t d 0 .
-]

Original: [-I n p u t s t d 0 .
-]
Revised: [-I n p u t s t d 0 .
-]

Original: [-I n p u t s t d 0 .
-]
Revised: [-I n p u t s t d 0 .
-]

Original: [-I n p u t s t d 1 .
-]
Revised: [-I n p u t s t d 1 .
-]

Original: [-I n p u t s t d 1 .
-]
Revised: [-I n p u t s t d 1 .
-]

Original: [Figure]

Original: I n p u t s t d 1 .

Original: [Figure]

Original: I n p u t s t

Original: [Figure]

Original: Cosine

Original: [Figure]

Original: [Figure]

Original: Cosine

Original: [Figure]

Original: [Figure]

Original: angle 0.

Original: [Figure]

Original: [Figure]

Original: d i

Original: [Figure]

Original: s i z

Original: [Figure]

Original: d v

Original: C.1 Toy Example

Original: In Figure 6, we focus on a toy example where the task is to reverse a sequence of tokens.

Original: More specifically, given a sequence of 20 numbers in the range 0 − 9 , we predict the same tokens in the inverted order.

Original: We use an embedding layer of size 16, initializes with variance 1, and sinusoidal positional encodings to initially embed the input.

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

Original: We use residual scaling in this case equal to α 1 = α 2 = 1 .

Original: We train using Adam with betas parameters (0 . 9 , 0 . 999) , learning rate 0 .

Original: 01 and weight decay 0.

Original: C.2 Tempering the softmax

Original: The theory devised in Section 3.3 postulates that a different magnitude between the gradients of the queries/keys and the values should likely be observed.

Original: Here, we propose a simple remedy that consists in introducing an inverse temperature scaling τ inside the softmax that modifies the attention

Original: I n p u t s t d 0 .

Original: Cosine angle 0.

Original: [-I n p u t s t d 0 .
-]
Revised: [-I n p u t s t d 0 .
-]

Original: [-I n p u t s t d 0 .
-]
Revised: [-I n p u t s t d 0 .
-]

Original: [-I n p u t s t d 0 .
-]
Revised: [-I n p u t s t d 0 .
-]

Original: [-I n p u t s t d 1 .
-]
Revised: [-I n p u t s t d 1 .
-]

Original: [-I n p u t s t d 1 .
-]
Revised: [-I n p u t s t d 1 .
-]

Original: [Figure]

Original: I n p u t s t d 1 .

Original: [Figure]

Original: n p

Original: [Figure]

Original: t

Original: [Figure]

Original: Cosine

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: d d i n g s i z e d v operation to

Original: [Equation]

Original: By computing the gradients of the queries/keys as in 3.3, Equation 18 becomes:

Original: [Equation]

Original: Hence, the norm of gradients of the queries/keys scales linearly with τ .

Original: On the contrary

Original: [Equation]

Original: the temperature scaling inside the Softmax will not affect the gradients with respect to the values due to the fact that the Softmax normalizes the activations (see Lemma A.1 in Wang et al.

Original: Hence τ can be heuristically chosen such that the magnitude of the gradients approximately matches:

Original: [Equation]

Original: We stress that this requires a constant correlation ρ .

Original: In practice, this can be estimated as the mean correlation across all pairs of tokens (as we do in the computation of the correlations in Figure 1).

Original: Furthermore, both ρ and the variance σ 2 x change across layers as our analysis in Section 3.2 predicts.

Original: Hence in practice a different temperature per layer should be adopted.

Original: Finally, note that in practice both ρ and σ 2 x change during training, and is hard to study their the dynamics under SGD.

Original: We leave the time evolution of ρ and τ as an exciting future direction.

Original: Also, the value of n is set to be the average number of tokens per sentence.

Original: In this work, we set τ to a fixed value according to our analysis at initialization.

Original: C.3 Translation Task

Original: We now describe the experimental setup regarding the translation task on the IWSLT’14 De-En dataset.

Original: Using the ideas detailed in the previous section, we choose a temperature value of τ final = 8 .to match the gradient norms of the values and queries as in Equations.

Original: (17) and (18).

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

Original: Due to instabilities in training, we use warm-up on this temperature value.

Original: In short:

Original: [Equation]

Original: with ‘steps warmup = 1000 ’ and ‘step’ the current training step.

Original: We base our implementation on fairseq [Ott et al., 2019].

Original: For the hyperparameter configuration, we mostly rely on the extensive search already done in fairseq [Ott et al., 2019] and Liu et al.

Original: The final used parameters are exhibited in Table 3.

Original: For the final evaluation, we use the best-performing model on the left-out validation set.

Original: We apply weight decay as in Loshchilov and Hutter [2017] for both SGD and Adam.

Original: [Table]

Original: Finally, in Figure 15 we display the evolution of correlations, residual scaling

Original: [Table]

Original: and norm of the activations, with depth, for our best trained model.

Original: The residual scaling α 1 , α 2 are trainable parameters.

Original: This enables them to weight differently the residual branches if deemed necessary.

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

Original: The norm of the propagated forward signal tends to slightly increase with depth.

Original: [Figure]