Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse
Authors: Lorenzo Noci
Sotiris Anagnostidis
Luca Biggio
Antonio Orvieto
Sidak Pal Singh
Aurelien Lucchi
Venue: NeurIPS 2022 Accept (2022)
Keywords: Theory of Transformers
Transformers
Rank Collapse
Gradient Vanishing

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

Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse
Lorenzo Noci"!
Sotiris Anagnostidis"! sotirios anagnostidisthng wtw.sst...eeeee.e.....oeee ch
lorenzo. noci@int ethz.ch
Luca Biggio*1.2 x1,2
Antonio Orvieto"
Sidak Pal Singh 1,3
luca. big.iiiggiiii ethz. ch
antonio orvieto8inf. ethz, ch
sidak, singh@inf. etzz. ch
Aurelien Lucchi' aurelien. lucchi Batcch@@uibbas.co
Abstract
Transformers have achieved remarkable success in several domains, ranging from natural language processing to computer vision. Nevertheless, it has been recently shown that stacking self-attention layers the distinctive architectural component Transformers can result rank collapse of the tokens' representations at initialization. The question of if and how rank collapse affects training is still largely unanswered, and its investigation is necessary for a more comprehensive understanding of this architecture. In this work, we shed new light on the causes and the effects of this phenomenon. First, we show that rank collapse of the tokens' representations hinders training by causing the gradients of the queries and keys to vanish at initialization. Furthermore, we provide thorough description of the origin of rank collapse and discuss how to prevent it via an appropriate depth- dependent scaling of the residual branches. Finally, our analysis unveils that specific architectural hyperparameters affect the gradients of queries and values differently, leading to disproportionate gradient norms. This suggests an explanation for the widespread use of adaptive methods for Transformers' optimization.
Introduction
Since its first appearance in Vaswani et al. [2017], the Transformer architecture has revolutionized the field of Natural Language Processing (NLP), achieving remarkable success in tasks such as text classification [Yang et al., 2019], machine translation [Conneau and Lample, 2019], reading comprehension [Brown et al., 2020] and question answering [Raffel et al., 2019] among others. Recent efforts have effectively extended its applicability to computer vision [Dosovitskiy et al., 2020] and other domains [Baevski al., 2020, Huang et al., 2018, Biggio al., 2021, Polu al., 2022], further popularizing it outside NLP.
The Transformer operates inputs comprising sequence of tokens. At its core, it relies stacked attention layers, which compute a measure of relevance for the whole sequence by assigning token- wise importance weights obtained by matrix multiplication of the queries and keys, and finally normalized with the softmax function. The output of an attention layer is then linear combination
Dept of Computer Science, ETH Zürich, "Robotics & ML, CSEM SA, Alpnach, Switzerland, "MPI for Intelligent Systems, Tübingen, 'Department of Mathematics and Computer Science, University of Basel
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Res scaling: d=
Residual scaling: a
1
VC
12
Costre angle Unittalizatto)
1.0
1.0
0.8
10
randon modii
0.8
Fardon predutto
0.6
Depth
6.6
Depic3
Cosire angle
Drydtt 12
0.4
0.4
2500 5000
75.00 10000
Depth
Steps
Depth
Steps
Figure 1: Evolution of the cosine of the angle between tokens for training POST-LN Transformers of increasing depth, with the Adam optimizer, for the IWSLT' 14 De-En translation task. Unless adequate residual scaling is used at initialization, increasing depth leads to an increase in the tokens' alignment at initialization, which can inhibit training.
of the importance weights and the so-called values. Then, the architecture includes fully-connecter sub-layers, residual connections [He et al., 2016], and layer normalization (LN), illustrated in Fig.
In the absence of residual connections, Dong et al. [2021] proved that at initialization the rank of the sequence representation collapses doubly exponentially with depth, and both layer normalization and fully connected layers can only partially alleviate the speed of degeneracy. Under rank collapse, the model does not distinguish between representations of different tokens, which are perfectly aligned in feature space at initialization. However, the precise implications of rank collapse in Transformers are not fully understood.
In this paper, we show that a high alignment of the tokens' representations at initialization corresponding rank collapse the extreme case of perfect alignment affects training by causing vanishingly small gradients of the queries and keys' parameter matrices. This problem severely diminishes the capabilities of the model learn meaningful attention weights and is further exacerbated in very deep networks, where the rank deficiency - and hence the vanishing gradient problem of the queries and keys affects several layers (see Fig. 1). In order to shed light on this problem, we take inspiration from the flourishing literature on signal propagation in random networks and start our analysis by computing the expected gradients of attention layer with respect the queries, keys, and values, which leads to Theorem 3.2 on the vanishing gradients for the queries and keys. From here, we pursue two different directions.
Firstly, we investigate under which conditions rank collapse can be avoided by studying the evolution of the input sequence in Transformer at initialization. Our theory reveals that a depth-dependent scaling of the residual branches, beyond stabilizing the norm of the activations at initialization, also approximately preserves the cosine of the angle between tokens, hence stabilizing the rank of the propagating sequence. We show that this holds even in the infinite-dppth limit.
Secondly, we illustrate that there are factors, other than the average tokens' correlation, that affect differently the gradient norm of the queries and keys compared to the values. In particular, the propagating sequence's squared norm has a linear dependence in the values, while cubic one in the queries and keys, justifying the use of layer normalization. We also highlight a different dependence on the embedding dimension and the length of the input sequence, implying that the gradient norm of a subset of parameters can potentially be of different orders of magnitude, as empirically hinted by previous works [Liu et al., 2020]. Our analysis brings 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.
Background
Transformers. A Transformer architecture consists of stacked attention blocks, as show in Fig. Layer normalization is usually applied token-wise either after the residual connections or to the inputs

---

of the self-attention and position-wist wise feed-forward sub-layers, leading to the POST-LN [Vaswani al., 2017] and PRE-LN [Wang et al., 2019, Xiong al., 2020] variants respectively.
Formally, given an input sequence X E Raxd,, with a tokens of dimension du, the single-head unmasked scaled dot-product self-attention' is defined as:
s = A'X' where A softmax (Vde x'we (X'Wk)7)
where the softmax function is applied independently across each row, and the superscript l indexes the E-th layer. The matrices W W E
Rd.xd= and W E Xd. are learnable parameters, and each layer is initialized with an independent set of weights. In the literature, the matrices X'W WV, X X'WV are referred to as queries, keys and values, respectively. The complete Transformer block, in the absence of layer normalization, can be written recursively as:
POST-nom
X
Y
Position-wise
Feed-forward
z' = a1 +X X" Y = olZ' W W
(2)
PRE-norm
(3)
Z
x a2 a +z',
(4)
POST-nom
where the introduced , 02 parameters indicate the strength of the resid- ual block, Fi F E Rddxdd xd. are matrices of learnable parameters;
K- aq
we set X = X, and a IR + R is an activation function. In our case, a the ReLU function, but we relax this assumption to the linear activation from Section 3.2 on.
Self-Attention
At initialization, each weight is sampled independently from distribution with zero-mean and variance o o for the values and feedforward weights', and o 1 for the queries and keys. This is the standard "Xavier" [Glorot and Bengio, 2010] or "He" [He al., 2015] initialization, commonly used in deep learning.
PRE-norm
X
de
Figure 2: A single Trans- former block.
Rank Collapse in Transtormers. Interestingly, Dong et al. [2021] proved that when the residual branches are omitted, the matrix of the tokens' representations X converges to rank-1 matrix in which all the
representations are the same and equal to a vector æ E Rd, i.e. X + 1,, where 1a, is the vector with all ones in Rd,, Note that this is slightly stronger notion of a rank-1 matrix, as it implies that all the tokens' representations are both perfectly aligned and have the same norm. Indicating the inner product with the usual bracket notations (--)). and the cosine of the angle between two tokens { || cos 0x. with cos 0x = 1 for all k, k' E [n]. Note that perfect alignment together with equal norm between all the tokens implies that all the representations are the same. One of our main contributions is to provide an explanation of how rank collapse affects the gradients of a Transformer at initialization.
as &g perfect alignment happens when (X. k. x +
Vanishing Gradient Problem. Traditionally considered one of the core issues that prevents suc- cessful training, the vanishing gradient problem has a long and rich history that dates back to before the popularization of deep learning [Hochreiter, 1991, Bengio al., 1994]. In its essence, given a loss function L Raxd, - IR. 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 [ac small to provide enough backpropagating signal, thus hindering gradient-based optimization methods.
aa - is too
Our analysis also easily generalizes to the case of cross-attention. In practice commonly uses Fi gd Xdp, F E dd where di yd,, with 7 {2, 8. Our results then hold up to constant factor that depends on y
'One should explicitly write the layer dependence w W. W.e wFi.e F2.e. We at times suppress the f index to improve readability. In case is the ReLU function, A. to have variance a.
3

---

Despite extensive research toward understanding and overcoming the problem in disparate contexts [Glorot and Bengio, 2010, He al., 2015, Hanin, 2018, Zhang et al., 2019], a formal explanation of its role in relatively new architectures such as Transformers is largely missing in the literature, with a few exceptions [Xiong et al., 2020, Wang et al., 2022, Huang et al., 2020]. In our paper (Section 3.1), we show how vanishing gradient occurs in conjunction with the rank collapse issue identified by Dong et al. [2021].
Signal Propagation in Random Networks at Initialization. After addressing the question on the effects of rank collapse, we take a step back and rigorously analyze its causes by looking at how the properties of the input sequence X are lost/preserved as it propagates through a randomly initialized Transformer. More specifically, we focus on two aspects of the propagating sequence: the expected
Frobenius norm E X+|| I| and the expected inner product between different tokens (X,,,, X with k + K. 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 [Hanin, 2018]. The latter quantity describes how the geometry of the feature space changes after applying a Transformer block, and is related the concept of dynamical isometry [Saxe et al., 2013]. To understand the evolution of the inner product, we analyze the following measure of correlation [Nachum al., 2021, Cho and Saul, 2009]:
E X E X
PRA
(5)
Note that P = 1 if and only if the k-th and k'tth tokens are perfectly aligned (cos O&s 1). We stress that in our case differently from the aforementioned works - instead of analyzing the relationship between two different data points, we study the relationship between tokens of the same sequence.
3 Theoretical Results
The goal of this section is twofold. 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. In Section 3.2, show that the problem can be mitigated with an appropriate choice of the residual branch parameters an and a2 that inversely scales with the depth of the network L. Under the proposed scaling, the correlations are well behaved even in the infinite depth limit. Finally, in Section 3.3 we analyze the scaling of the gradients with respect to other network's 's parameters, and in 3.4 we draw some connections between our findings and optimization of Transformers.
3.1 Vanishing Gradients for Queries and Keys under Rank Collapse
To investigate the problem of vanishing gradients in the attention layers. make use of the framework of matrix calculus [Magnus and Neudecker, 2019, Singh et al., 2021]. In particular, we compare the expected Frobenius norm of the gradient of self-attention layer with respect to its parameters
as DW where here W indicates one of the keys, queries or values weight matrices. Due to the
well-known difficulty of computing expectations of the softmax [Daunizeau, 2017, Shekhovtsov and Flach, 2018], throughout this manuscript, we make the simplifying assumption that the softmax output the uniform distribution at initialization, i.e. the
matrix containing 1 in each entry. n.
Assumption 3.1 (Uniform attention). We assume that A 1 n 1nxm
where 1n is the matrix with all entries equal to 1. Crucially, in Appendix A.5, we formally show that this assumption holds almost surely in the limit dę + cc. There, we also experimentally show that even in the more realistic case where d, = d, = 512, the empirical simulations provide a surprisingly faithful approximation of the theoretical insights presented in this paper.
We define the mean token æ through its components of = Σr= X E [do] In the following theorem, we compute the expected gradients of an attention layer at initialization, and set the basis for our following analysis. We provide the results only for the queries, as the case for the keys is analogous.

---

Lemma 3.1. Let X be the representations the input sequence at the l-th layer. Under the uniform-attenioon assumption, we have
as
E
d dn
1112
(6)
OWV.
os
2
o
of
E
X [(xx)) X no n
(7)
aWQ
n²
2
o
TE
[(x)) - no (z') 20
(8)
0X
We defer the precise study of the scaling of these quantities as a function of and de dg. to Section 3.3. At this stage, is crucial to note that _ (X)) X x'("") ) is the centered empirical covariance .. matrix of the tokens' representations. It is easy see that if X is rank-1 matrix, then all the rows of X are proportional to a fixed dy di-dimensional vector, and the empirical covariance matrix has all zero entries. Introducing differentiable loss function L Raxd -> I. make the statement on vanishing gradients more formal in the following theorem:
Theorem 3.2 (Vanishing gradients under rank collapse). Suppose that the uniform-attention
assumption holds. If additionally X for any E [L] has rank-1, and there exists vector æ E R" such that X ,, then:
OC
OC
E
=0,
E
0,
(9)
OWQ
aW
where the expectation is taken over the weight matrices. This implies that these quantities vanishing almost surely, due to the non-negativeness of the norm.
The proof simply relies on expanding the norm of the gradient of the loss with the aid of the chain rule and then bounding by the product of the norms of each term the chain. The final result holds with application of Lemma 3.1, in which the rank-1 assumption makes E aw
ase
vanish. The
proof of the results presented in this section deferred to Appendix A.I.
In light of Theorem 3.2, we can conclude that the issue of rank collapse originally identified in Dong et al. [2021] corresponds to an initialization in a region of vamishing gradient signal in the subspace of parameters identified by the queries and keys. How can this affect training? One may argue that if rank collapse does not happen in the very first layer, then the corresponding gradients non-zero and the rank of the subsequent layers affected by rank collapse with the first few steps of gradient descent. In practice, we show empirically in Fig. that escaping this pathological landscape is harder in deeper nets where rank collapse persists across several layers.
be increased
3.2 Forward Signal Propagation and the Importance of Scaling the Residual Branches
We now turn our attention the study of the influence of skip connections in Transformers. Dong et al. [2021] showed that simply adding this architectural trick prevents rank collapse. Somewhat surprisingly, we show that while the claim holds for any finite depth, the average angle between different tokens quickly increases with just a few layers, and as L - oc a Transformer can still lose rank unless the residual branches are adequately initialized. As Dong et al. [2021] showed that layer normalization does not avoid rank collapse, we omit it in our analysis. Firstly, we introduce two lemmas on the propagation of inner products (Lemma 3.2) and the norm (Lemma 3.3) of the tokens' representations.
5

---

Relative
notrasaled cornecioon
acaaled comnectoon
acaled resicial
100
Residual scaling
2100
10
Deptt
Figure 3: Effect of the residual scaling to the norm of the gradients of the network at initialization with respect some loss. From left right: (a) the cosine of the angle between tokens increases with depth. Note how larger values of (11, 022 imply a faster token alignment with depth (Theorem 3.3). Subplots (b) and (c) show the gradients of the queries-keys and values parameters respectively by increasing depth, compared to the corresponding norms of the first layer. Gradients for the queries- keys diminish with depth, while the opposite happens for the values, We POST-LN to disentangle the effect of the variance of the input.
Lemma 3.2 (Propagation of inner products). Ler C(X') = (X,, and X the input sequence. Under the Assumption 3.1 and if is the linear activation function, we have that:
E [C(x')] = (a2 + 1) 1)
(10)
hence, under the depth scaling for the residual block parameters of , with on, 02 independent of L, we have that:
lim L-0
(11)
Note that C(X!) n 1|02 1|2. The lemma on the propagation of the norm is slightly more involved:
Lemma 3.3 (Propagation of the norm). Let X L be the representations of the input sequence at the final layer. Under the assumptions of Lemma 3.2, we have that:
L-1 E x n + 1) Σ 1|*1 + (a + 11 2
(12)
A=
hence, under the depth scaling for the residual block parameters independent of T,, we have that:
== ā1 I
0z 2
āz
I
with 0i, 02 E R
lim E. L
X
( -1) || a |
(13)
The proof of Lemma 3.3 consists in expanding E X
F according to the defining equations for the Transformer, and simplifying the expression by using iterated expectations F. E|X
E|[ X [ F [ exploit the conditional independence between different layers, and then puting the expectations using the independence assumption on the weights. The expression on the right-hand side will then depend on X only through its norm | and the norm of the mean token
|= Using Lemma 3.2 then allows to unroll the recursion and get the final result. The complete proof, together with the proof of Lemma 3.2, can be found in Appendix A.3.
The previous Lemma provides theoretical justification that scaling the residual branches by setting the alpha parameters to be (()/VL) allows both the norm of the propagating input and the inner products between different tokens to approximately preserved. Hence, the information contained in the input is not lost, even in the infinite depth limit.

---

Residual Scaling Preserves Correlations. We now prove that without the depth dependent residual scaling (i.e. with a₁ Q2 = the correlation between the tokens quickly increases, and reaches perfect alignment in the infinite depth limit. More specifically, our argument shows that in this limit, the correlation between different tokens Pl as Eq. (5) converges 1, implying rank collapse. Furthermore, show how setting the residual parameters a and a2 dictated by Theorem 3.3, ensures that the correlation measure is dependent on the input in non-trivial way even at infinite depth. To this end, we introduce the average correlation at layer
Σ PlA KØK
(14)
n(n
Note that p² 1 if and only if every pair of tokens is perfectly aligned. We ready to formalize the influence of the L/VL-scaling the correlation between tokens' representations by stating Theorem 3.3.
Theorem 3.3. Let the input tokens have the same norm, i.e. | |||||
Vk E [n] for some
w Rd., Under the depth scaling for the residual block parameters of ā,, 0g E IR independent of L, we have that:
2, a² z with
Th ((X)
lim p L-000
(15)
(n - 1))(ee (X)(X) + 72 |X
n-1
On the other hand, af, 02 # 0 are some constants independent of L, we have that:
lim p² 1. L
(16)
The proof consists in noting that due to the symmetry of the problem at initialization, for a fixed layer the expected norm of each token is the same. Hence, by definition of Pkx. write
EX k: x k/ Pla E |== By summing over the k, k' indexes, the resulting equation will depend on (X)) and || X+||? which can be expanded using Lemma 3.2 and 3.3 respectively. The result then given by solving for pf.
Note that under the 1/ 1/VI-scaling, the correlation term is one if and only if C(X) = ||X] which holds in the degenerate case where all the input tokens perfectly aligned. In Appendix A.4, we give precise formulas for the expected correlations at any depth, showing that pf reaches values close to one even for relatively shallow net- works when the 1/v L-scaling is not adopted (see also Fig. (left)). Additionally, in Fig. 4, we empirically show that the presence of the 1/VL-scaling (either PRE or POST) does not significantly affect the evo- lution of the correlations. On the other hand, without the residual scaling, PRE-LN seems to alleviate the rate of in- crease P It is intriguing that most deep Transformer models use this configuration [Brown al., 2020]. We provide more extensive empirical results Appendix
3.2
0.3
Correiotion R 8.7
NG-N
PIEAN
POSTUN
layer normalization
scaling
0.4
200
Depth
Note that the 1/VL 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 the context of rank preservation. Finally, note that by setting â1, āz O, we recover the so called "ReZero" initialization [Bachlechner al., 2021]. In this context, the 1/vT. scaling extends this framework as it allows for wider range of values for ã1, 02 while still guaranteeing stability.
Figure 4: Evolution of Correlation in Transformers with (dashed lines) and without (solid lines) 1/v L/VL--caling for PRE-LN, POST-LN and without layer normalization (No-LN).

---

ReLU extension. We mention here that extending these results from the linear activation to the ReLU case is known to be hard problem, due the technical difficulty of propagating the imner products across ReLU layers that are shared among the tokens (this is the case in the position-wise feed-forward layers in Transformers). Exact formulas can be found only in the case of one ReLU layer with Gaussian inputs in Cho and Saul [2009]. 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. 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. [2021] below Equation (2)). Hence, the perfect alignment (a.k.a rank collapse) that affects the linear case affects the ReLU case well (in which case the rank collapses even faster with depth, as we show in Figure 10).
3.3 Dependence on the Angle between Tokens and the Input Norm
In this section, we drop the superscript as it is obvious from context and assume for simplicity that dę = d,- To gain 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 P = p for i + or more precisely
0 oz
+ j' (independent dimensions) i= j + j' j
poz
To see that this equation satisfies our definition of the correlation metric, note that [||| do and (X, Xx) do p, for a i'. Then, the expected norm of the gradients for the values (Eq. (6)) simplifies to
as
E
= o (1 1)). 1))-
(17)
GWV
By making the additional assumption that the norm and the correlation propagate independently, the respective norm for the queries and symmetrically the keys (Eq. (7)) reduces to:
as awQ
(n (1 2)² d).
E
(18)
In Appendix A.2 we provide a rigorous proof, that relies on Isserlis theorem [Isserlis, 1918] to compute higher-order moments. The above expressions reveal the different dependencies on four main actors, that we inspect separately here. The gradients of the queries depend via cubic function on the variance of the input, o2 compared to a linear for the values. This provides an additional interpretation of the successful use of layer normalization, as in Xiong et al. [2020], either in the POST-LN or PRE-LN format, that standardizes the input variance of to the value 1.
Next, we emphasize the dependence on the correlation between the tokens, also illustrated in Fig. 3. Importantly, note how the queries/keys have opposite monotonic functional dependence with respect to compared to the values. As revealed by Theorem 3.3 and Fig. (center), inappropriate scaling of the residual branches can already lead to this phenomenon even in a relatively shallow network.
Finally, Eq. (17) and (18) reveal different scaling in terms of the embedding size d and the sequence length due to the self-attention operation itself. We hope that the identification of the different dependencies in the gradients of the parameters will inspire new line of works aimed at solving some the difficulties in training Transformers.
3.4 Connections to Optimization and Adaptive Methods
The existence of the discrepancy in the magnitude of the gradients with respect to the weights Wº, W and WV, might explain the success of adaptive optimization algorithms, as illustrated in Fig. 6, where we plot the effective learning rate computed by Adam [Kingma and Ba, 2014] in toy encoder task (more details in Appendix C). Notice that the effective learning rate is increasingly larger (with depth) for the queries compared to the values as postulated by our theory this difference is remarkably constant throughout training. Hence, we conjecture that the success of adaptive methods in Transformers' training be partially explained by the need to fix this
and
8

---

Table I: BLEU scores for the IWSLT14 German-te to- English translation task. SGD res-scale refers the train- ing SGD without layer normalization and initialization of the residual scaling 01 0g vE
+. SGD temperature
additionally employs an inverse temperature inside softmax.
8 18
Adat:: PME
Method (6L-Encoder / 6L-Decoder)
BLEU
Aroer POST LB
Adun: Thu
SGD POST-LN
31.36
Adee Re2mm
18
SGD res-scale
32.79
SGD temperature
35.69
Adam POST-LN [Vaswani et al., 2017]
35.39
Figure 5: BLEU scores by increasing the number of
transformers blocks. Transformer blocks implies total 'X' encoder self-attention, 'X' decoder self- attention and 'X" decoder cross-attention layers.
Adam PRE-LN [Vaswami et al., 2017]
35.10
ReZero [Bachlechner et al., 2021]
34.55
T-Fixup Zhang et al. [2019]
35.59
disproportionate gradient's magnitude. To test this hypothesis, propose simple architectural modification, an inverse temperature scaling T E R inside the softmax:
I Ix' (x'W )) x'w V
S = softmax
(19)
direct consequence of our analysis that T al- lows controlling the magnitude of the gradients for the queries and keys' parameters. tion C.2, we detail how one can choose such that the magnitude of the gradients as derived in
Sec-
104
Queries' depth
Adtan's per paramwter aduptiuity
50*
Equation 17 and 18 is approximately matched at initialization.
20²
We evaluate our proposal, consisting of residual scaling and the aforementioned inverse temper- ature parameters, on the widely used IWSLT14 Berrman-to-Englist (De-En) benchmark transla- tion task. All details regarding the experimen- tal setup and the choice of inverse temperature used provided in Appendix C. We train Transformer encoder-decoder of varying depth with stochastic gradient descent (SGD), after re- moving all normalization layers and adequately initializing the residual connections. For training with SGD, avoid using any learning rate warm-up, as commonly done for Adam, and instead use step-scheduler to decrease the learning rate at 40% and 80% of training. We compare against the following methods that make use of Adam; POST-LN and PRE-LN refer to the aforementioned alternatives to apply layer normalization. We also compare against other successful techniques that rely specific initializations avoid layer normalization, such as ReZero [Bachlechner et al., 2021] and T-Fixup [Zhang et al., 2019]. We report the average BLEU score [Papineni et al., 2002] across runs in Fig. and Table 1.
10.-
Values' depth
training steps
Figure Adaptive learning rates computed by Adam in Transformers.
Our proposed method considerably improves training with SGD, keeping up and in some cases surpassing any results achieved by the Adam optimizer. We are also able to train deeper networks without the use of layer normalization. We leave for future work to further investigate modifications alternatives to the self-attention operation.
4 Related Work
Our work builds upon the rich literature on forward and backward signal propagation in random neural networks [Poole et al., 2016, Schoenholz et al., 2017, Xiao et al., 2018, Pennington et al., 2017, Orvieto et al., 2021, Noci et al., 2021, Zavatone- Verhh and Pehlevan, 2021]. The 1/VT. scaling scheme has been investigated in the literature for the stabilization of residual networks [Hanin and Rolnick, 2018, Arpit et al., 2019, Allen-Zhu et al., 2019, Hayou al., 2021].

---

Our work draws inspiration from series of recent works studying the rank of the representations of random feed-forward neural networks at initialization [Daneshmand et al., 2020, 2021]. In the context of Transformers, Dong et al. [2021] has recently identified the rank collapse issue object of study of the present work. Thanks to our analysis of the backward pass, we are able demonstrate that rank collapse in Transformer architectures leads to vanishingly small gradients of queries and keys, thereby preventing effective training and allowing us to complete the analysis of [Dong et al., 2021].
Among the architectural components in Transformers, layer normalization is, arguably, one of the most important and debated ones [Chen et al., 2018, Wang et al., 2019, Nguyen et al., 2010, Xiong al., 2020]. the original architecture [Vaswani et al., 2017], layer normalization used to stabilize the forward pass by reducing the variance of the inputs to the following sublayer. Our analysis of the forward pass shows that its inclusion is not strictly necessary for the purpose of controlling the norm of the representations. For theoretical analysis of signal propagation in the presence of layer norm, we refer the reader to Xiong et al. [2020].
Additionally, our theoretical study of the backward pass provides a rigorous explanation of the empirically observed discrepancy between the magnitude of the gradients of the queries and the values, which Liu et al. [2020] hypothesize to be one of the causes of the success of adaptive methods in training Transformers [Liu et al., 2019, Zhang et al., 2020, Huang al.. 2020].
Finally, properly rescaled residual connections have been found to be beneficial for training Trans- formers by number of recent research works [Zhang et al., 2019, Bachlechner et al., 2021, Wang et al., 2022]. However, none of these studies characterize the impact of skip connections on rank propagation, while analysis suggests
snreerically-rouueeedddd way to stabilize it.
5 Conclusions and Future Work
In this paper, we showed how, at initialization, rank collapse and more generally high correlation in the tokens causes vanishing gradients of the queries and keys of Transformer architecture. While residual connections help mitigate rank collapse at finite depth, we showed that they alone cannot prevent high alignments of the tokens' representations - unless properly scaled by a 1/ 1/VL- factor. Finally, we have also discovered counter-intuitive dependencies on the variance of the input, embedding size, and sequence length, potentially causing large differences between the gradients of queries/keys compared to the values' parameters, Hence, we conclude that one of the strengths of Transformers lies in their carefully designed architecture together with an adequate initialization. Finally, we gave preliminary evidence that one of the factors contributing the higher efficacy of Adam compared to SGD in training Transformers arises from the disproportionate magnitude of gradients as postulated by our theory. Nonetheless, other factors might further accentuate the difference between these two algorithms during training, leaving the door open for further research regarding the benefits of adaptive optimization methods with Transformers.
Acknowledgements
Sidak Pal Singh would like to acknowledge the financial support from Max Planck ETH Center for Learning Systems and the travel support from ELISE (GA no 951847).
References
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
Devansh Arpit, Víctor Campos, and Yoshua Bengio. How to initialize your network? robust initialization for weightnorm & resnets. Advances in Neural Information Processing Systems, 32, 2019.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian Mc Rezero is all you need: Fast convergence at large depth. Uncertainty in Artificial Intelligence pages 1352-1361. PMLR, 2021.
10

---

Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning speech representations. Advances in Neural Information Processing Systems, 01111244-12661 2020.
Yoshua Bengio, Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157-166, 1994.
Luca Biggio, Tommaso Bendinelli, Alexander Neitz, Aurelien Lucchi, and Giambattista Parascandolo. Neural symbolic regression that scales. In International Conference on Machine Learning, pages PMLR, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 33:1877-1901, 2020.
Mia Xu Chen, Orhan Firat, Ankur Bapna, Melvin Johnson, Wolfgang Macherey, George Foster, Llion Jones, Niki Parmar, Mike Schuster, Zhifeng Chen, et al. The best of both worlds: Combining recent advances in neural machine translation. arXiv preprint a8Xiv:1804. 0.949, 2018.
Youngmin Cho and Lawrence Saul. Kernel methods for deep learning, Advances in neural information processing systems, 22, 2009.
Alexis Conneau and Guillaume Lample. Cross-Lingual Language Model Pretraining. Curran Associates Inc., Red Hook, NY, USA, 2019.
Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, and Aurelien Lucchi. Batch normalization provably avoids ranks collapse for randomly initialised deep networks. Advances in Neural Information Processing Systems, 13:18837--18398 2020.
Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes repre- sentations in deep random networks. Advances in Neural Information Processing Systems, 34, 2021.
Jean Daunizeau. Semi-analytical approximations to statistical moments of sigmoid and softmax mappings of normal variables. arXiv preprint arXiv:1703.0009) 2017.
Yihe Dong, Jean-Baptin Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-2803. PMLR, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arX1v:2010.11929. 2020.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249-256 JMLR Workshop and Conference Proceedings, 2010.
Boris Hanin. Which neural net architectures give rise to exploding and vanishing gradients? Advances in neural information processing systems, 31, 2018.
Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. Advances in Neural Information Processing Systems, 31, 2018.
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable resnet. In International Conference on Artificial Intelligence and Statistics, pages 1324-1332. PMLR, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.

---

Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In In IEEE Conference on Computer Vision and Pattern Recognition (CVPR). pages 770-778, 2016.
Sepp Hochreiter. Untersuchungen zu dynamischen neuronalen netzen. Diploma, Technische Univer- sität München, 91(1), 1991.
Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Ian Simon, Curtis Hawthorne, Noam Shazeer, Andrew M Dai, Matthew Hoffman, Monica Dinculescu, and Douglas Eck. Music transformer: Generating music with long-term structure. In International Conference Learning Representations, 2018.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learning. pages 4475-4483. PMLR, 2020.
Leon Isserlis. On formula for the product-moment coefficient of any order of a normal frequency distribution in any number of variables. Biometrika, 1341/2) 1918.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv 6980, 2014.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond, 2019. URL https:
//arxiv. org/ abs/ 03265.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.0822 000-249900. 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
171.5:1711.05101, 2017.
Jan R Magnus and Heinz. Neudecker. Matrix differential calculus with applications in statistics and econometrics. John Wiley& & Sons, 2019.
Ido Nachum, Jan Hązla, Michael Gastpar, and Anatoly Khina. A jonnson--lindenstaass framework for randomly initialized cnns. arXiv preprint arXiv:2111.02155, 2021.
XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11))5547-5886 2010.
Lorenzo Noci, Gregor Bachmann, Kevin Roth, Sebastian Nowozin, and Thomas Hofmann. Precise characterization of the prior predictive distribution of deep relu networks. Advances in Neural Information Processing Systems, 34, 2021.
Antonio Orvieto, Jonas Kohler, Dario Pavllo, Thomas Hofmann, and Aurelien Lucchi. Vanishing curvature and the power of adaptive methods randomly initialized deep networks. AISTATS 2022 (lo appear), 2021.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 3-1-318, 2002.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: ptatttss:/ppreee //ppoceedings.. neurips. http://p///////////////////////.
Paper. Paper.pdf.
12

---

Stanislas Polu, Jesse Michael Han, Kunhao Zheng, Mantas Baksys, Igor Babuschkin, and Ilya Sutskever. Formal mathematics statement curriculum learning. arXiv preprint arXiv:2202.01344, 2022.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in neural information processing systems, 29, 2016.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 10683, 2019.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 2013.
Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein, Deep information propagation. In ICLR, 2017.
Alexander Shekhovtsov and Boris Flach. Feed-forward propagation in probabilistic neural networks with categorical and max layers. In International conference on learning representations, 2018.
Sidak Pal Singh, Gregor Bachmann, and Thomas Hofmann. Analytic insights into structure and rank of neural network hessian maps. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https:
//openreview. www.eee/./oooo fere/fformmmeee daterrmmiddootggeLLNN...
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. arXiv preprint arXiv: (0e.11006...3 03762, 2017.
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers 1,000 layers. arXiv preprint arXiv:2203.00555 00555 2022.
Qiang Wang, Bei Li. Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. arXiv preprint arXiv: arXiv: 1906.01787, 2019.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Penning- ton. Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks. In Jennifer Dy and Andreas Krause, editors, Proceed- ings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 5393-5402. PMLR, 10-15 Jul 2018. URL https: //proceedings... mlr wwt.eem.......eeeeeeeeeee ppwss/v80///se http//iaa11111 html.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533, PMLR, 2020.
Zhilin Yang. Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xinet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
Jacob Zavatone-Vent and Cengiz Pehlevan. Exact marginal prior distributions of finite bayesian neural networks. Advances in Neural Information Processing Systems, 34, 2021.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint arXiv: 1901: 009321 2019.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 13315533--1393, 2020.
13

---

1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 3
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section
(b) Did you include complete proofs of all theoretical results? [Yes] Most of them in the Appendix
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes] The code will be released upon acceptance.
(b) Did you specify all the training details (e.g., data splits, byperparameters, how they were chosen)? [Yes] See Appendix.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes]
4. If you using existing assets (e.g., code, data, models) curating/releasing assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "wYsFoJVamY", "reviewer": "Reviewer_Wg56", "summary": "This paper focuses on studies about the rank collapse phenomenon of transformers regarding its causes and effects. Authors first show that rank collapse has side effect on training by causing vanishing gradients of the queries and keys of a Transformer. Then the causes of this phenomenon is described and a solution is proposed to prevent it via depth dependent scaling of the residual branches. Finally, preliminary evidences are given about high efficacy of Adam compared to SGD in training Transformers.", "questions": "Please refer to the weakness part in the Strengths And Weaknesses section above.", "limitations": "The limitations and assumptions are discussed in Section 3 of the paper. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n+ Studying the rank collapse phenomenon of transformers is an interesting topic. It is important for training depth and large Transformer models, which is the key for popular foundation models. \n\n+ Thorough theoretical analysis are provided to describe the cause of this phenomenon and potential solution to prevent it. Assumptions and limitations are also described clearly. \n\n+ Experimental evidences are also given to validate the theoretical analysis.\n\n\nWeakness:\n\n- Experimental evidences are mainly provided for natural language processing (NLP) area. It is not clear whether the conclusions can be generalized to computer vision tasks, e.g. the most basic image classification. \n\n- The effectiveness of the proposed depth dependent residual branch scaling seems not linked with evaluation metrics. Does this solution help training deeper Transformer models to achieve better performance for NLP and computer vision models? ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Ty2FKaqzbEI", "reviewer": "Reviewer_FzSt", "summary": "This work investigates the rank collapse problem for transformers at the initialization, showing that it results in vanishing gradient and slow training for self-attention parameters. Based on the theoretical study, the work proposed a method to alleviate this issue via adding a temperature parameter to the self-attention. The authors claim that the work opens the way for new, well-founded, and motivating approaches to improve optimization in these models.\n", "questions": "The reviewer’s questions relate to the weakness of this work.\n1. Although the authors justify Assumption 3.1 to some extent. The paper assumes linear activation, which does not hold in practice. How does the theory well connect to the practice? \n2. Some of the argument is in the setting where layers or dimensions are infinite. How do these results connect to the finite sample case?\n3. The major question is about Section 3.4. It is unclear (at least not well explained) from the theory how a temperature can be helpful. And how to choose the size of the temperature parameter?\n", "limitations": "n/a", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strength: the paper is well-presented and easy to follow;\nThe work studies an interesting problem of the rank collapsing phenomenon at initialization.  The work provides rigorous analysis for explaining why this happens under several assumptions. \nThe work utilizes the understanding, showing it can be useful for more efficient training of deep networks.\n\n*Weakness: The assumption is quite strong (Assumption 3.1), and the sigma is linear. These do not hold in general in practice deep networks. Because of the significant simplification, It is unclear about the connection between theory and practice. \nTo deal with the vanishing gradient problem caused by rank collapse, the author proposed to add a temperature scaling to the self-attention. It is very vague how the theory applies such a temperature would help with training. And it is unclear how to choose its size.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "KdBed8Nl7wq", "reviewer": "Reviewer_w8ur", "summary": "This paper studies the rank collapse problem of the tokens' representations at initialization in Transformer. The authors provide theoretical analysis to further understand and deal with this problem. They conclude that the rank collapse problem can cause the gradient vanishing problem during training and propose that applying an appropriate depth dependent scaling of the residual branches can prevent the rank collapse problem. Besides theoretical analysis, the authors also provide some experimental results to demonstrate the effectiveness of their proposed remedy method. ", "questions": "1. In line 31, you mentioned Dong et.al 2021 study the rank collapse problem in the absence of residual connections, I wonder whether it is proper or what the problem is for the requirement? \n2. What is the Depth in Figures 3 and 4? \n", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. This paper studies a novel problem, rank collapse, to further understand the optimization of Transformer. They investigate how the rank collapse hinders the training process of Transformer from the aspect of gradient vanishing, which has not yet been well studied for new architectures, such as Transformer. \n2. The authors propose a remedy method, an appropriate depth dependent scaling of the residual branches, to avoid the rank collapse problem. The experimental results demonstrate the effectiveness of this approach, as shown in Table 5 and Figure 1. \n3. This paper brings a new problem to studying the signal propagation in Transformers. This would motivate other researchers to further study and improve the optimization of Transformer models.\n4. This paper introduces the problems, their analysis, and their solutions clearly. The paper is well written and easy to follow. The authors present their theoretical analysis clearly and easy to be understood. The notations and equations are easy to read. \n\nWeaknesses:\n1. Some theoretical analyses are based on the previous work, Dong et al. 2021. This would make the novelty of this paper a little bit trivial.\n2. Section 3.4 seems to have less relationship with the main topic of this paper. To my understanding, it is redundant.\n3. The authors only provide some results on NLP tasks. I wonder if this problem is still severe and if we can address it in the same way for vision transformers.  ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept