Taming Transformer Without Using Learning Rate Warmup
Authors: Xianbiao Qi
Yelin He
Jiaquan Ye
Chun-Guang Li
Bojia Zi
Xili Dai
Qin Zou
Rong Xiao
Venue: ICLR 2025 Poster (2025)
Keywords: Transformer
Training Dynamics
Model Crash
Primary area: optimization
TL;DR: We revisiting the training dynamics of Transformer to tame its training process without using learning rate warmup.

Abstract: Scaling Transformer to a large scale without using some technical tricks  such as learning rate warump and an obviously lower learning rate, is an extremely challenging task, and is increasingly gaining more attention. In this paper, we provide a theoretical analysis for the process of training Transformer and reveal a key problem behind model crash phenomenon in the training process, termed *spectral energy concentration* of ${W_q}^{\top} W_k$, which is the reason for a malignant entropy collapse, where ${W_q}$ and $W_k$ are the projection matrices for the query and the key in Transformer, respectively. 
To remedy this problem, motivated by *Weyl's Inequality*, we present a novel optimization strategy, \ie, making the weight updating in successive steps steady---if the ratio $\frac{\sigma_{1}(\nabla W_t)}{\sigma_{1}(W_{t-1})}$ is larger than a threshold, we will automatically bound the learning rate to a weighted multiple of $\frac{\sigma_{1}(W_{t-1})}{\sigma_{1}(\nabla W_t)}$, where $\nabla W_t$ is the updating quantity in step $t$. Such an optimization strategy can prevent spectral energy concentration to only a few directions, and thus can avoid malignant entropy collapse which will trigger the model crash. We conduct extensive experiments using ViT, Swin-Transformer and GPT, showing that our optimization strategy can effectively and stably train these (Transformer) models without using learning rate warmup.

Published as a conference paper at ICLR 2025
TAMING TRANSFORMER WITHOUT USING LEARNING RATE WARMUP
Xianbiao Qi', Yelin He', Jiaquan Ye', Chun-Guang Li Bojja Zi', Xili Dai', Qin Zou", Rong Xiao I1
Intelliftusonn Inc. 2BPT $CUHK *HKUST (GZ) "WHU
ABSTRACT
Scaling Transformer large scale without using some technical tricks such learning rate warump and using an obviously lower learning rate is an ex- tremely challenging task, and increasingly gaining more attention. In this paper, we provide theoretical analysis for the process of training Transformer and reveal the rationale behind the model crash phenomenon in the training process, termed spectral energy concentration of W Wk. which is the reason for malignant entropy collapse, where W and We are the projection ma- trices for the query and the key in Transformer, respectively. To remedy this problem, motivated by Weyl's Inequality, present novel optimization strat- egy, i.e, making the weight updating in successive steps smooth the ratio a
is larger than a threshold, we will automatically bound the learning WW
(W,)) rate to weighted multiple of W, where VW, is the updating quantity in step Such an optimization strategy prevent spectral energy concentration to only a few directions, and thus can avoid malignant entropy collapse which will trigger the model crash. We conduct extensive experiments using ViT, Swin- Transformer and GPT, showing that our optimization strategy can effectively and stably train these Transformers without using learning rate warmup.
Nothing in life is to be feared. It is only to be understood."
Marre Curie
INTRODUCTION
Transformer Vaswani et al. 2017 has revolutionized various domains of artificial intelligence, including natural language processing [Rafford et al. 2018 2019 Brown et al. 2020 Chowdh-
ery al. 2023 Touvron al. 2023 Dubey et al. 2024 and computer vision Dosovitskly al. 2020 Liu et al.
2021) and many more applications (Radford et al. 2021 Ramesh et al. 2021 Peebles & Xie 2023), owning to their ability capture long-range dependencies through self- attention mechanisms. However, despite their widespread application and empirical success, training deep Transformer models remains quite challenging. Practitioners frequently encounter variant issues, such as gradient explosion (Qi et al. 2023b), rank collapse (Dong et al. 2021 en- tropy collapse (Zhai al. 2023) and general training instability (Kim et 2021 especially during the initial training stage.
Qi et 2023b),
To address these challenges, researchers have proposed various modifications to the original Transformer architecture, including altering the placement of Layer Normalization (Wang et al., 2019 Xiong et al. 2020) (e.g. pre-I.N vs. post-LN schemes), carefully conditioning the residual connections (Bachlechner et al. 2021), and QKNorm (Henry et al.) 2020 Dehghani et al.) 2023) for the sell-attention module. Similarly, DeepNet [Wang et al. 2022 introduces a new normal- ization function to modify the residual connection in Transtormer. ReZero (Bachlechner et al. 2021 introduces a learnable residual scalar parameter for the residual shortcut, and requires ini- tiating it to 0 at the start stage of training. More recent approaches (Kim et al. 2021 Qi et al., 2023a have focused on examining and enforcing Lipschitz continuity properties of Transformer
Corresponding authors

---

Published as a conference paper at ICLR 2025
components, which can provide insights into the network behavior and the training stability. Al- though there are a few works (Bachlechner et al. 2021 Qiet al. 2023a) that can avoid using learn- ing rate warmup to train Transformer successfully, all of them require significant modifications of the network structure.
Learning rate warmup [Loshhhlovv & Hutter 2016
seems be must- have technology for stan-
dard optimizers [Robbins & Monro 1951 Duchi et al. 2011 Kingma & Ba 2014
Loshchilov &
Hutter 2019 in some popular large Transformer models (Radford et al. 2018 2019 Brown al.,
2020 Chowdhery et al. 2023 Touvron et al. 2023). Without the learning rate warmup stage, the
Transformer training will prone diverge.
Although it is usual to train a Transformer by modifying the network structure as mentioned above using the learning rate warmup, two natural and interesting questions remain:
1. What are the training dynamics of Transformer model when its training or succeeds?
2. Can we successfully tame Transformer without changing the network structure or without using learning rate warmup?
This paper aims to answer these questions. To answer the first question, we examine the training processes of three types of Transformers, by visualizing the changing trajectories along with the training process of 15 (or 13) quantities about the parameters, activations, and attention maps. By doing so, we observe that the model crash is accompanied by a weird phenomenon that the entropy of the attention map almost and the spectral norm of Wa WW increases to a very large value. By conducting mathematical analysis for the Transformer training, we identify that the Spectral Energy Concentration (SEC) of Wa Ww is the key problem leading to the model crash. To answer the second question, motivated by Weyl' Inequality, we present a novel opti- mization strategy, Le., making weight updating smooth, and verify empirically that our optimiza- tion strategy can prevent spectral energy concentration and thus achieving stable convergence in training.
Paper Contributions. The contributions of the paper are highlighted as follows.
We visualize the training dynamics of Transformers that train successfully unsuccessfully and summarize two important observations from unsuccessful training that: a) the rank of the attention map matrix tends very low and the entropy of attention probability matrix tends 0; and b) (W, Wx) increases rapidly to very large value.
We present theoretical analysis for the Transformer training, finding that the Jacobian matrix dvec P)
X @ XT, where P X Wg W.X. It implies that the gradient of W, W
dvec(W, a W₁)
is largely dominated by the rank of X @ XT.
We reveal that SEC of W T W; makes the attention map matrix to be sparse yet low-rank and is the reason leading model crash.
Motivated by the Weyl's inequality, we introduce novel strategy to address the problem of SEC of W,, Wg by controlling the rapid growth of singular values, and verify that our strategy leads to stable training process.
2 PRELIMINARIES
Matrix norm. Given a matrix W, its fn-nrrm is defined as: I| supa Tel When p the induced matrix norm is the spectral norm, which is defined the largest singular value of W and is also expressed as the square root of the largest eigenvalue of the Gram matrix
WW The spectral norm of matrix W can be calculated as: IW = max || ll2
V Amax W W) = σ (W), where o (W) denotes the largest singular value of matrix W and Amax (WT W) denotes the largest eigenvalue of W W, and s=-1 denotes unit sphere in R"
Power iteration to compute matrix spectral norm. The power iteration algorithm starts with Wæ, for k 0....,K K- 1. At every iteration, a is multiplied by matrix W and normalized, After K iterations,
vector 2o of unit normm The entire iteration process is as follows: æx+1
iW
2

---

Published as a conference paper at ICLR 2025
| 22 is used as the estimated spectral norm. Usually, it takes 3 to 5 iterations converge, and thus the computation cost cheap.
Adam Optimizer. Adam Kingma Ba 2014 is currently the most widely used optimizer for training neural networks, owing to its efficiency and effectiveness. Adam can be simply defined V P2 V-1 +0 (1, W, = W M, 2 v /V,+ e, where G2 is the gradient at step f, G; is the bementtwise square of G,, and o denotes element-wise division, a; is the learning rate step 1, and $I. are the first-order and the second-ords mentum factors, respectively.
as: M, = M, (1--)++11))+:+
3 TAMING TRANSFORMER REQUIRES TO REVISIT THE TRAINING DYNAMICS
To begin with, we first give some basic notions in a Transformer, which includes an attention module, an FFN module and two normalization modules that used before the attention mod- ule and the FFN module. For the attention module, we usually use a multi-head attention which allows the model to jointly attend to information from different representations with different heads. Here, for the convenience, without losing generality, we only use single-head attention. To be precise, we define each of them follows:
X W T W
Attn(; W WE. W. Wo
W. W,W, softmax
Vdg
FFN(;; WT, W2)
W2ReLU(W11))
y and y = 1
LN(x)
70z+B, where
11
x.
std(y)
where X E Rdxx, Wg E Rdyxd W E Rdgxd, W, Rdoxd, W. Raxd Wi RAxdd W₂ Rdx4d Y Rd and B E Rd. Note that only in a single- head definition, W. can be put before
W;; otherwise, it should be after concatenation operator. For the convenience discussion, we define:
= softmax(
and P=X W.X,
dg
where A is called the attention map and vdg
P
usually called the logit.
3.1 VISUALIZATION: WHIT HAPPENS WHEN TRANSFORMER TRAINING FAILS SUCCEEDS
Visualizations are commonly used as an effective means to help us examine whether the neural network's training succeeds or fails. In particular, what happens when the training of Trans- former tends to crash? And what happens when the training succeeds?
One of the most important aspects of understanding the training of neural networks is the obser- vation of changes in parameters and activations. Since the parameters or activations and their gradients are matrices or vectors, the norm is the best way observe these quantities. In this paper, for a Transformer training, we summarize the following 15 terms to watch:
a (Wa), (WW), (W,,,, (WW,,, T (WW, (W),, ( We).
aL
(1)
σ (W, lml2.
1. 1222. 132l. o (W, WT), ||
ox
where (1 and 71 denote the parameters the first LayerNorm, and 32 and y2 denote the pa- rameters of the second LayerNorm. When RMSNorm IZhang & Sennrich 2019 is used, there are only 13 terms that will be analyzed since does not have 3. the weight matrix, we use the spectral norm. For vector, we use its l norm.
To ensure that the phenomena we observed can generalize well, we visualized them on both VIT (Dosovitskiy al. 2020) and GPT (Radford al. 2018). VIT is pure encoder architectures whereas GPT is pure decoder architecture.
3

---

Published as a conference paper at ICLR 2025
Ma
(a) l1 12
(b) 1/1 l2
(c) 01 (Wa)
my Fy
Bio
(ool(We)
(e) 0 (W T W&)
(1) (W,)
5i9
(g) To (Wo)
(h) llr2llz
(i) 1,3 12
(W))
(k) 0o (W2)
0 (W2W1)
400
(m) (7 (W.W.) W.)
(n) ||æ|||
(o) ||
de
lla
FIGURE 1: Training dynamics of a failure VIT. This figure shows how 15 items of quantities as de- fined Equation change during the training period. Please pay more attention subfigures (a)-(e).
4

---

Published as a conference paper at ICLR 2025
In this section, due to the space limitation, we will only visualize ViT-base, and put more visualization results into the Appendix H For the VIT implementation, use Timm Wightman (019), in which "timm" library provides rich model architectures of many pre-trained image models in PyTorch. For the GPT implementation, we use nanoGPT, which uses LayerNorm without a bias term, and thus only watch 13 terms, rather than 15 terms in VIT.
10h
10h
17.
176
128
(a) Block (successfuu.
(b) Block (successfil).
(c) Block 11 (successful).
30
12-
199
159
156
206
230
2l
105
129
150
13
10
(d) Block (crashed).
(e) Block (crashed).
(f) Block (crashed).
To achieve successful VIT training, we use a long learning rate warmup. For instance, we use 60 epochs of warmup and the whole training process takes 150 epochs. To
FIGURE 2: Visualization of the dynamics process of attention map in different training steps for a successful and a crashed ViT-base model, respectively. Please click the images to play the flash. Best viewed with Acrobat Reader.
obtain the dynamics of failure training of VIT, do not warmup.
Figure I visualizes failure training process of a ViT-base model. The model includes 12 blocks with index from 0 to 11. In Figure
visualize the weight matrices in blocks 20,2,4, 6, 8, 10, 11), where the index of the last block is 11. For the features and the gradients 0L. 2 we hook the input features that enter into the corresponding blocks. Figure 8 in the [ppendix[] visuulzzesss a successful training process of VIT-L model. Meanwhile, in Figure 2 visualize the dynamic processes attention maps during the training period about a successful case and a failure case of the ViT-base model, respectively. We visualize the attention map of the same image at different steps.
We observe the following phenomena from Figures and
As shown in Figure 1 at the beginning, the maximum singular value (W, We) gradually increases, and at a certain point, the maximum singular value suddenly and rapidly increases very large value (e.g. around 200,000), where the loss divergence emerges. However, for a successful training process, d' (W, Ww) gradually increases medium value shown in FigureBBand then vibrates around that value.
As shown in Figure 2 in failure training process of Transformer, the attention maps gradu- ally become sparse and low-rank. Finally, the entropy of the attention map is 0; whereas in successful training process, the attention map too sparse and of medium rank.
The normalization layers exhibit a huge difference: 7/1 and (1 in successful VIT training pro- cess are very smooth, but ll llz and || llz suddenly increase dramatically in an unsuccessful case.
The ranges of the activation and the gradients are very large in a crashed model, and the gradi- ents in different blocks vary much larger than that in successful model.
We summarize that the successful training and the unsuccessful training of a Trans- former exhibit significant differences among their a (W We), their normalization parameters y and B, and their activations and gradients àl.
Remark.
5

---

Published as a conference paper at ICLR 2025
3.2 THEORETICAL ANALYSIS: MATRIX CALCULUS OF TRANSFORMER
To understand the training dynamics of Transformer, we should investigate the process of back- propagation RRmelhart et al. 1986 LeCun al. 2002 1989 1998). In the attention mechanism, however, the input and the output are both matrices, we cannot directly use vector calculus. Instead, we need to use Vectorization (Graham 2018 Petersen al. 2008) and Kronecker Prod- uct] (Graham 2018 Petersen et al. 2008). In matrix calculus, the vectorization of matrix is linear transformation that converts a matrix into vector. Specifically, the vectorization of ma- trix M E RR"*", denoted as vec(M), is a column vector, obtained by an ordered stacking of the b c
columns of the matrix M, i.e,, vec(M) Rm" For example, for a matrix M d f]
the vectorization of M is vec(M)) la d b fl" For the attention module of Transformer, we have the following proposition about the Jacobian matrix of the output with respect to the input X and the parameters.
Proposition (Matrix Calculus for solf-Attentio) Wg q W X, where X E Rd Wa Rddad W E Rdyxd, according to vectorization
Let x
and matrix calculus, we have the following derivations:
Ovec(P) X ② T
dvec(P)
(XT Ovec(X)
W TWTTwo Wg + Wi).
(2)
dvec(W, WR))
dvec(P)
Ovec(P) X @ (W
(W (WxX) OX
(3)
)ve
dvec(Wr)
where & denotes the Kronecker product, I E R denotes an identity matrix with shape x n, K is the commutation matrix, which depends on the dimensions X. Since Rdxn then we know K Rdddxdd The commutation matrix K has the property that vec(X K vec(X) for any matrix X.
In Appendices A and B we supply some elementary information for the vectorization and the Kronecker product and the derivations of the Jacobian matrix for single-head attention.
We have the following observations from Proposition
We have Pvec(P) T & X in Equation 2 and we know about the Kronecker product @ve q We)
that rank(X X rx) which implies that if the rank of X is very low, then the rank of Dvec(P) will also be very low. Note that X being low rank means that the features
@ve Wi
across different timestep highly correlated coherent. all æ; in X collapses to single point, then X @ X will only have large singular value, and the rest are 0.
The Jacobian matrix dver( in Equation 2 is in direct proportion X and Wa T We. If the dveclX]
at will more likely to be spectral norm di (W, Wx) is very large, it implies that the gradient a.X magnified lot.
Equation 3 suggests that changes in the query weights Wa a are related to both the input X and the key representation W, W.X. gquation[Jsuggest that changes in the key weights Wi are related both the input X and the query representation W.
All these relationships are interconnected, with changes in one variable potentially affecting the others. For instance, if W; increases fast, then according Equation dvec(P) will more
)vec(W,
likely be very large. In this way, Wa will likely increase very fast.
3.3 RATIONALE IN MODEL CRASH: SPECTRAL ENERGY CONCENTRATION
Before we reveal the rationale in model crash, let us first discuss a bit on the entropy collapse.
Two Entropy Collapse Modes. In experiments, observe two types of attention entropy collapse modes. Note that when attention collapse happens, the attention map tends to
https://en.co http:////aaaiiiacaaooaaaaa////tt htt.://aaaiiatmmtttmmaattaaaoo (mathematics)

---

Published as a conference paper at ICLR 2025
sparse matrix (f.e., there few dominate nonzero coefficients in attention map), and thus the entropy of the attention map is vanishing. To be more specific, when the attention map is sparse but not low-rank, call it benign collapse, whereas if the attention map is sparse and simultaneously low-rank] we call it a malignant collapse. When benign col- lapse occurs, the attention map is shown in the right panel of Figure that, there is almost an identity matrix. In this way, the diagonal elements are almost 1, and the non-diagonal elements have values around 0. Unfortunately, when malignant collapse happens, the tention map becomes sparse and simultaneously low-rank matrix as shown in the middle panel of Figure g Furthermore, we observe that the distribution of the spectral energy of W Wx for the benign collapse is relatively uniform; whereas the spectral energy of the atten- tion matrix for the malignant collapse tends to concentrate on few dominate singular values.
By analyzing the ma- trix Wy'We in the be- nign collapse when happens in the exper- iments, we find that has the following prop- erty: Wg W, is usually non-symmetric posi- tive quasi-definite square matrix (see Appendix [ for details), and the self-attention layer de- generates into linear projection layer because Y= W, A WX W.X. We give an intuitive analysis in Appendix
(6) Normad attention
(b) Malignant cullapne
(c) Beniign collapre
FIGURE Three modes of attention maps. The left panel shows normal attention map. The middle panel shows a classical at- tention map when model crash occurs, for which the entropy almost 0. The right panel shows an attention map from normal model training but its entropy almost 0.
When the malignant collapse happens, the model will usually crash. We identify that the ratio- nale behind the model crash is phenomenon called spectral energy concentration (SEC). Be- fore we present our theorem about SEC, let us first introduce an index to quantify it. Recall that Wy E Rdg and WA Rdyxd xd, where dy d. We have that W, W e Raxd, but its rank less than or equal to dq. Precisely, we define a SEC index as follows:
of T Wt)
SECldq, 3)
Σ (W, Ww)
where dg is the head dimension and s dg. For instance, if we have dg = 64 and s= 4, and if at this time, SEC(64, 4) 99%, could say the spectral energy of W T W; highly concentrates only four dominant singular values.
To be precise, we have the following theorem for the reason cause malignant collapse. Theorem (Malignant Entropy Collapse) Let = T X and A = softmax
W = W Wx e dddd Suppose that the following
ddg
two conditions are simultaneously satisfied:
X a low-rank matrix;
W is a lou-rank matrix with only a few dominant singular values (e.g. the singular values are greater than Co ydg where Co 1 is a constant).
Then, the attention map A will be sparse and simultaneously low-rank in high probability.
When the malignant entropy collapse happens, the training process will crash. We provide proof for Theorem I in Appendix
For sparse and simultaneously low-rank matrix, we refer the reader recent textbook
WWiight & Ma
2022
7

---

Published as a conference paper at ICLR 2025
According to Theorem we know that the model crash is caused by the spectral energy concen- tration. In Figure we evaluate and compare the curves of SEC(dg, s) three different blocks under successfully trained model and crashed model. In successfully trained model, the spectral energy distributes all directions. However, the spectral energy ofa crashed model only concentrates few directions. As shown in Figure in crashed model, the SEC collapses into less than 10 directions.
Figure reveals how the attention collapse propagates through
San-tiist livegy Canserratte
Spectiid Eiveryy mopcermation
Spe-diial Shangg Concenetaiin
1.0
0.8
55014
SECLA
Boo creas bsaa
Eooi kacarr
(a) Block 0
(b) Block 6
(c) Block 11
FIGURE 4: Comparison of spectral energy concentration index SEC between a successfully trained model and crashed model. The attention maps of three different blocks are shown. The spectral energy distributes in all directions successful training case; whereas the spectral energy only concentrates few directions in crashed model.
x X
X
dvect
Wq
WkI
A X
W W
FIGURE 5: Attribution flow chart of attention collapse.
each term, illustrating the entire attribution chain from the input X' the output X/+1 in the attention module. Note that indicates being low-rank and V means being sparse and simul- taneously low-rank. X is low-rank, then X X T is also low-rank because rank( X)
rank(X) rank(X). According to the gradient computation, we have dvectW, W = & X
dveciP)
thus we know àvec P)
is also low-rank. Meanwhile, it should be noted that the spectral en-
Dvec(W, Wi)
ergy of åvec(P) is over-concentrated, which means that the gradient update will dramatically dvec(W, W2)
change W We. thus W will have large probability be low-rank. In the paper, have proved that being low-rank and the leading singular values of W We are very large will lead to the attention map A over-concentrated (see Appendix for the proof), and become to sparse and simultaneously low-rank matrix. Finally, an cver-concentrated A will lead to X to be lonkr We provide more detailed analysis of Figure Appendix
3.4
OUR SOLUTION: TAMING TRANSFORMER VIA WEYL'S INEQUALITY
The analysis above reveals that spectral energy concentration the key reason leading to un- stable training. One manifestation of spectral energy concentration is the rapid growth of the singular values of of the weight matrices. Thus, our basic idea to prevent malignant collapse is to suppress the fast growth of the singular values. Fortunately, Weyl's inequality provides us simple but effective tool.
Theorem (Weyl's Inequality Singular Values)
2.. be the ordered singular values Let Wi, W2 E Rmnx where m (W)) 2 Wf, and on (W2 (2) ... o (W2) be the ordered singular values of W2. Then we have:
0 dt-j-1(W W2) (W;) (W2).
The proof for Theorem is provided in Appendix H From Theorem 2 it is easy to that o, (W, W2) s o, (W1) + o (W2) Let W, be the weight matrix at time step 6, VW, be the quantity computed from the gradients and their derivations (where vW, be obtained by
8

---

Published as a conference paper at ICLR 2025
Algorithm I AdamW Taming Transformer via Weyl' Inequality without learning rate warmup.
Input: learning rate scheduler ar, weight decay A, and first-order and second-order momentum fI.. P Output: updated weight
1: for I 1.2...., 7 do
2:
G = -V
3:
M, +
(1=
4:
M, M1 (1) Bi). V== (==V((1-Af3)
VW, = MM e Vi+c
VW, the final update quantity, Power iteration compute spectral norm.
81 PowerIter(VWii 01 Powerlter(W;-.)
a, then
7:
If Rule does not meet, adjust the learning rate a,.
01
Th a, 61
9:
end if
10:
if Weight Decay is Yes then
11:
W,- Wi1 W,WW, W-11
12:
13:
W = W,1 - &VVW,
14: end if
15:
end for
SGD (Robbins & Monro 1951), Adagrad (Duchi et al. [011], or Adam Kingma & Ba 2014)), ), is the learning rate at time step Usually, our update equation is W, = W a, VW, According to Weyl's Inequality, we have,
0y (W,)
0 = (W,-1 - a, sowi)oono
(5)
An important observation from Equation is that o (VW) is very large, then W, will be signif- icantly different from W It means that the successive updates of W, at time step t from time step t- 1 would "jump" too much. A "smoother" updating should satisfy the following rule.
Rule (Steady Weights Updating Rule)
Given weights matrix W and the updating quantity VW, , at step 1, with the learning rate a, a steady weights should satisfy the following inequality: [| 1iy---,,WW,lll (1+ 7)1l W, w-1122 1l2,, where
T 0 is small factor.
To meet R what we need is that 0 (W, (W W (1+r)
It is easy to see that:
(W
a, S (VW,),
This inequality tell us that the learning rate should be bounded by ratio of singular values 01 (WW,)) and 01 (VW). Generally, r a small value, egg, 0.004 or 0.005. The intuition behind that if the spectral norm of VW, significantly larger than that of W,, then the model is po- tentially undergoing rapid changes. In such cases, a large learning rate could lead to training (W o,(V then we truncate a; to (V
instability. Therefore, our strategy is that if > Ø1 (W,1)
Since our base optimizer is AdamW [Loshchilov & Hutter 2019 and our algorithm is motivated by Weyl's Inequality, we term our algorithm as 4
For clarity, summarize our 2 in Algorithm where lines 6-9 highlight our improve- ments over the base optimizer, the other codes are same as AdamW. According to line in Al- gorithm (V,)) and (W,-1) ,1) are computed via fast power iteration method. In practice, we set the maximum iterations power iteration to 3. Actually, we find that two iterations are enough to estimate the spectral norm of matrices. According to Equation if a, (_(W,1) W,
then the learning rate a, will be truncated (W(Wei) (v(VW, the algorithm will adjust a, and use the default learning rate set by the learning rate schedule.
EXPERIMENTS
Compared to some previous works (Bachlechner et al. 2021 Wang et al. 2019: Xiong et al. 2020 Wang et al. 2022 Qi al. 2023b) that focuse improving the training stability of Transformer,

---

Published as a conference paper at ICLR 2025
TABLE 1: Quantitative comparison of AdamW and AdamW with and without learning rate
warmup. AdamW demonstrates very competitive performance compared to AdamW.
Method
VIT (Acc. 1)
GPT (Loss !)
Swin-Transformee (Acc. 1)
Configurations
VTT-B
ViT-L
GPT-S
Swin-S
Swin-B B
Parameters
86M
307M
125M
50M
88M
AdamW (with warmup)
80.22
81.65
2.848
83.02
83.48
A (no warmup)
80.58
81.82
2.840
83.14
83.44
our Ad does not need to adjust the network structure and we do not use learning rate warmup. For ViT, Tranffrranformer and GPT, we will use warmup of 60 epochs, 20 epochs, and 2000 steps, respectively. In AdamW2, we directly use a cosine learning rate schedule and decay the learning rate from maximum minimum.
We conduct experi- ments three pop- ular Transformers, iee. VIT (Dosovitskiy et al. 2020 GPT- RRafford al. 2019 and Swin- Transformer (Liu et al. 2021), where
3.2
CIty yea
3.55
210
2.40
28
2.30
14
Can 2013
2.0
100
Spack
263
135
200
VIT and Swin-
(a) Tran-Trrasformer
(b) GPT
Transformer are
pure encoder archi- tectures and GPT is a pure causal decoder. Note that do conduct any adjustments to the networks and directly the original implementation. Our experiments include image classification on ImageNet (Deng) et al. 2009) and large language model on OpenWebText (Gokaslan & Cohen) dataset. We list some training configurations in Appendix N The quantitative results are shown in Table I Our baseline model the corresponding Transformer using a learning rate warmup; whereas baseline models without using learning rate warmup will crash. Adam demonstrates a very competitive performance compared the baseline method. These experimental results verify that our understanding of the training dynamics of the Transformer is rational.
FIGURE 6: Ablation study of r in 2 using Swin S and GPT-S.
also conduct an ablation study of the choice of T in GPT and Sran-Trransformer. The results are shown in Figure 6 We can see that the performance of AdamW2 varies slightly for different values of r, but overall, our approach robust for different choices of T. The curves basically overlap in the later epochs because steady updating rule is never broken in the later epochs.
CONCLUSION
In this paper, we revisited the training dynamics of the Transformers by visualizing the spectral norm of weight matrices, the activations and the attention map, presented a theoretical analysis for the Transformer training and identified two modes of attention entropy collapse, i.e., the be- nign collapse and the malignant collapse, which the malignant collapse accompanies model W, T Wr is the reason
crash. Moreover, we revealed that the spectral energy concentration behind the model crash, which causes the attention map to be sparse and simultaneoush, low- rank. Furthermore, we proposed a steady updating rule to resolve the problem of spectral energy concentration W W. by controlling the rapid growth of singular values, which can prevent the fast spectral energy concentration to few directions and thus avoid the malignant entropy collapse. We conducted extensive experiments to verify the proposed strategy with VIT, Swin Transformer, and GPT, and demonstrated that the proposed strategy could effectively and stably train a model without using any learning rate warmup.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
In this paper, we aim to provide novel approach to train transformers without learning rate warmup. Our work does not involve any human subjects, and we have carefully ensured that poses no potential risks or harms. Additionally, there are no conflicts of interest, sponsorship concerns, or issues related to discrimination, bias, or fairness associated with this study, We have taken steps to address privacy and security concerns, and all data used comply with legal and ethical standards. Our work fully adheres to research integrity principles, and no ethical concerns have arisen during the course of this study.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our work, we provide all the details to reproduce the experi- ments. Theoretical proofs of the claims made in this paper, and detailed experimental settings and configurations are provided in the Appendices.
ACKNOWLEDGMENTS
Chun-Guun Guan-Guaggga was partially supported by the National Natural Science Foundation of China under Grant 61876022. Qin Zou was partially funded by the National Natural Science Foundation of China under Grant 62171324.
REFERENCES
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. Uncertainty in Artificial Intelligence, pp. 1352-1361 PMIRR, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
Aaron Defazio, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, et al. The road less scheduled. arXiv preprint arXiv:2405. 150.222......,,, 2024.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers 22 billion parameters. In International Conference Machine Learning. pp. 7480-7512. PMLR, 2023.
Ji Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255. leee, 2009.
Yihe Dong, Jean- nan-Baptiste Cordonnier, and Andreas Loukas. Attention not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference Machine Learning, pp. 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. In International Conference Learning Representations, 2020.
11

---

Published as a conference paper at ICLR 2025
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus.
Gene Golub and Charles Van Loan, Matrix computations. JHU press, 2013.
Alexander Graham. Kronecker products and matrix calculus with applications. Courier Dover Publications, 2018.
Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang. Flatten transformer: Vi- sion transformer using focused linear attention. In Proceedings of the IEEF/CVF international conference on computer vision, pp. 5961-5971, 2023.
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4246-4253 2020.
Roger A Horn and Charles R Johnson. Topics in matrix analysis, 1991. Cambridge University Presss, Cambridge, 37:39, 1991.
Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge university press, 2012.
Andrej Karpathy. NanoGPT. htttps://gittubb.com///aaaaaeeee
2022.
Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant self- aeff-attention. In International Conference on Machine Learning, pp. 5562-5571. PMLR, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6 6980, 2014.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay bal- ances learning across neural networks. arXiv preprint arXiv:235. 17212, 2023.
Atli Kosson, Bettina Messmer, and Martin laggi. Analyzing & reducing the need for learning rate warmup in gpt training. Advances in Neural Information Processing Systems, 2024.
Yann LeCun, Bernhard Boser, John Denker, Donnie Henderson, Richard Howard, Wayne Hub- bard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1 14):5511-51, 1989.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient based learning applied to document recognition. Proceedings of the IEEE, 96211):2278-2324 1998.
Yann LeCun, Léon Bottou, Genevieve Orr, and Maus-Rober Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-50. Springer, 2002.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the diffi- culty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5747-5763, 2020.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEFE/CVF International Conference on Computer Vision, pp. 10012-10022, 2021.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts, arXiv preprint arXiv:1608 009803983,...000 2016.
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. In International Conference on Learning Representations, 2019.
12

---

Published as a conference paper at ICLR 2025
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role rank col- lapse. Advances in Neural Information Processing Systems, 35:27198-27211,220 2022.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEFFICVP International Conference on Computer Vision, pp- 4195-4205, 2023.
Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical Univer- sity of Denmark, 7(15):510, 2008.
Xianbiao Qi, Jianan Wang, Yihao Chen, Yukai Shi, and Lei Zhang. Lipsformer: Introducing lips- chitz continuity to vision transformers. In The Fleventh International Conference on Learning Representations, 2023a.
Xianbiao Qi, Jianan Wang, and Lei Zhang. Understanding optimization of deep learning via ja- cobian matrix and lipschitz constant. arXiv preprint arXiv:2306. 09338, 2023b.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language un- derstanding by generative pre- rre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, (8):9, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to eoxt-t--image generation. In International conference chine learning, pp. 8821-8831. Pmlr, 2021.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathe- matical statistics, pp. 40. 4000000007000 1951.
David E Rumelhart, Geoffrey E Hinton, and Ronald Williams. Learning repressntttessssn by back- propagating errors. nature, 323(6089) 1986.
Leslie N Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter confer- ence on applications of computer vision (WACV), pp. 464-472. IEEE, 2017.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding train- ing dynamics and token composition in I-layer transformer. Advances in Neural Information Processing Systems, 36:71911-71940 2023a.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demys- tifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535 2023b.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09220 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Roman Vershynin. dineeniimenion probability: An introduction with applications in data sci- ence, volume 47. Cambridge university press, 2018.
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deep- net: Scaling transformers 1,000 layers. arXiv preprint arXiv:2203.0 00555, 2022.
13

---

Published as a conference paper at ICLR 2025
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. arXiv preprint arXiv:1906.01777, 2019.
Ross Wightman. Pytorch image models. pytorch
httpes//ggitaab..cmm/vvvvgtgtttt
model s 2019.
John Wright and Yi Ma. High-dimensiones data analysis with low-dimensionalloolee Principles, computation, and applications. Cambridge University Press, 2022.
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 architec- ture. In International Conference on Machine Learning, pp. 10524-10533. PMLR, 2020.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind. Stabilizing transformer training by preventing tention entropy collapse. In International Conference on Machine Learning, pp. 40770-40803. PMLR, 2023.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances Neural Infor- mation Processing Systems, 32, 2019.
14

---

Published as a conference paper at ICLR 2025
A KRONECKER PRODUCT AND VECTORIZATION
Kronecker product Graham 2018 Petersen al. 2008), also called as matrix direct product,
an operation defined on two matrices of arbitrary size. The specific definition is as follows.
Definition 1 (Kronecker Product)
Let A be an p matrix and B an m q matrix. The mn pq matrix
01, B
B
41.p B
02 B
022
02.p B
A B=
An,1 B
Bn
Qn.p
is called the Kronecker product of and B. is also called the direct product or the tensor product.
1
3
2
4
6
1
2
11
2
3 5
3
5
10
For instance, if A
and B
then A
3 4
3
4
3
6
9
4
8
12
9
12
15
12
16
20
Some basic properties of the Kronecker product include:
A (B C) (A & B) * C,
AB (A B) + (A&C), (A+ (A=B)oC=(AeC))(BBC))
(A B)² A @ B
For matrix A, the rank of A & A can be computed
rank @ A) rank(A) rank(A).
It means if the rank of the matrix A is small, then the rank of A a A will also be very small.
In mathematics, Vectorization (Graham 2018 Petersen et 2008 usually used together with the Kronecker product to express matrix multiplication as linear transformation on matrices. After vectorization, we can calculate the Jacobian matrix the matrix product more conve- niently. A property of vectorization for the matrix product is defined below. Proposition 2 (Property of Vectorization for Matrix Product) Let mxn Rink C Ex1 then we have
vec(ABC) (CT A)vec(B).
Proof. Let C₁ be the i-th row of C. Then we have:
vec(ABC) Σ
Σ
bij vecla,C))
Σ
bij C @
Σ [C] A)b;
(CT A) vec(B).
Furthermore, we have the following properties:
vec(AB) (T 8 A)vec(B) (B) vec(A),
vec(ABC) (C B Im] vec(A)
= =(C² @ A) vec(B)
(Ti AB) vec(C))
15

---

Published as a conference paper at ICLR 2025
where I E R I Rixi, .Im Rm* are all identity matrices.
Together with the Kronecker product, vectorization is an effective tool to compute matrix calcu- lus. We can see the following two examples.
Let P = AB where A E R RMnx, B E R R"*k, we have:
@vec(P) Ave
Qve dvec(B)
= BT @ Im
1 s A.
Let P ABC where A E RMn Be R CE Rkki have,
@ve dvec(A)
Ovec(P) dvec(B)
Ovec(P) dvec(C)
C BT & I.
C. & A,
Ii 8 AB.
Vectorization and Kronecker product provide us a convenient way to analyze the self-attention module. We can compute the Jacobian matrix of the output with respect to the input the weight matrix more conveniently. For more introduction to Kronecker product and vectorization, the readers refer to (Petersen al. 2008 Graham 2018
B DERIVATION OF JACOBIAN MATRIX FOR SINGLE-HEAD SELE-AITENIIO
A single-head eeff-attention can be defined as
Y=W,XA
where softmaxx P. in which P= X W W;X. The matrix is called the attention ma- Vdg
trix and is called the logit, A RuKn X R dxn W, o E Rd,xd Heere, our goal is calculate
v Vig
ovectY) Ovec(X)
In the main body, we have derived dvec( ovec( Here, let us calculate the matrix calculus of A with respect P using Kronecker products and vectorization. We rewrite
softmax( P
it A exp(
P
yd (1 exp( P,)), )), where 1, denotes n-dimensional vector of 1's in R".
vāg
Note that A is obtained by conducting softmax operation in each column individually.
First, let us define two intermediate variables:
P = BB.
exp(
C= exp(
dd
dd
In this way, we represent the attention matrix A as = BeC.
Then, we can vectorize the equation A = B e C' as follows:
vec( vec B e C*) vec BB @ vec(C).
According to the chain rule, we have
θvec(A) òvec(A) dvec(B)
dvec(A) dvec(C))
@vec(P) dvec(B) Ōvec(P)
Ōvec(C) dvec(P)
16

---

Published as a conference paper at ICLR 2025
Let us calculate each individual term. We have Ōvec(A)
dvec(B) = 1 n² e diag(vec(C)),
Ôvec(B) diag(vec(B))
åvec(PP
vdg
Ōvec(A) diag(vec(B) oee(C))
ôvec(C)
Ōvec(C)
(1, @ 1, 1 Idiagtvvev)))
åvec(P)
dd
where In is an identity matrix of and is the Kronecker product, 1n denotes an n²- dimensional vector of I's R"²
By substituting the above four terms into the chain rule, have
àvec(A) [1, e diag(vec(C))) diag(vec(B)) diag(vec(B) @ (vec(C) vec(C))) 1,1, Övec(P)
1 diag(vec(B))
d
diag(vec(A)) diag(vec(B) @ (vec(C) o vec (1, @ 1, 1 diag(vec(())
Vd,
blockdiag(diag(A.1) A, diagLA..)) A.
Vdg
For the simplicity, we denote
J blockdiag(diag((.)) A, diag...,,) A..
When and P are vectorized into vectors, we use a and to denote them, respectively. Then see that
dvec(a) diag(a) aa
Ōvec(p)
dg
If approaches unit vector e, then the Jabobian matrix dvec(a) will tend to 0.
dvecipl
In Section 3.2 we have the following Jacobian matrix
àvec P))
(X T W Wa W Tw))
WW).
@vec(X)
By vectorization Y = W, X A, we have
àvec(Y)) (A @ W,)ovec(X) (Ln @ W, X))vec(A)
Therefore, according to the product rule and chain rule, we can denote the Jacobian matrix Y with respect to X as follows:
òvec(YY)
= (A @ W.) + (I, @ W, X) dvec(A)
ovec(X)
Ōvec(X)
Ōvec(A) Ōvec(P)
= (A @ W,, (1, @ W,
Ōvec(P) Ōvec(X)
Bringing in all these terms, we get the following formula:
òvec(Y) (AT o w.)+((n W.X)
((x WT T WW Kunk W Two). m
dvec(X)
dd
17

---

Published as a conference paper at ICLR 2025
Let us analyze Equation 7 If malignant entropy mode happens, J will approach 0 because Vdg
each in A will be unit vector From the perspective of the forward process, the features Y will collapse several directions. From the perspective of the backward process, J will
Va
XvectX)) will be sparse and simultaneously low-rank matrix. Through A W,, most of the positions in X will get zero gradient, and only very few columns will obtain some large noisy gradients. In a malignant entropy mode, the learned feature is invalid and useless. Similarly, benign entropy mode happens, the attention map A will approach an identity ma- āg attention module degenerates to linear layer.
become O, and dver(y)
trix I and J = when A = Therefore, we have dvecly) = @ Wy). In this way, a self- 3vec(X]
C PROOF OF BENIGN ENTROPY COLLAPSE
Recall that softmax( P.) where P X W W.X. Here, let W W Ww and W Vdg
Raxd In this way, we have that A = softmax
X XWX, ). We know that rank(W)) s dg and dg < d. Vdg
To prove A will always collapse to identity matrix when W is non-symmetric positive quasi- definite square matrix, is equivalent to prove E [ Wæ; > E [x Wæyl for any i # j. It will be very hard prove it mathematically W is form of non-symmetric positive quasi- definite square matrix. Therefore, let us make some simplification assumptions. Assume W real symmetric positive semi-definite square matrix and its trace is in direct proportion to iid the dimension dq. and any æ; is high-dimension random vector and each element in Xt,j N (0, 1).
We break our proof into two sub- sub-problems.
Proposition (Expectation a; W a;;)
Let W be real symmetric positive semi-definite matrix, and any be high-dimennionall ran- dom vector. Then, have xx) W Wx,]] trace(W).
Proof. Let W be real symmetric positive semi-definite, thus it can be decomposed into W UEUT, In this way, we have
W
=E
UEU x1
= E
Σz|
(let z = U xx)
Σ 0 /=1
(E a diagonal matrix)
dg
0 (0+1)
(by independence, mean 0)
= Σ o; trace(W).
For real symmetric positive semi memi-definite, all its singular values are larger or equal to 0. Thus, we know Σ o, 0 considering W not matrix of all zeros.
Proposition 4 (Expectation of Waj j for j)
Let W be real symmetric positive semi-definite square matrix, and any 2; is high-dimension random vector, Etzj [x; W ar =0.
18

---

Published as a conference paper at ICLR 2025
Proof. Let W be real symmetric positive semi-definite. Thus, it can be decomposed into W U.UT. In this way, we have
Eiz [2)
m, Way
W
= E (aj
UUU az1]
= =E Ev]
(let z == U x;, and v =U-[] xz)
dg
= E
Σ
DijziVj
(2 is diagonal matrix. and v are independent)
1=1
According to Proposition 3 and Proposition any i # j. Considering that W is usually high- dimensional matrix and some of its singular val- ues significantly larger than 0. Thus, after the softmax operation, A will always collapse to identity matrix. In this way, the self-attention module degenerates into linear projection mod- ule. The model fitting ability will decline, but model training will not crash. Our proof is based on matrix computations (Golub & Van Loan 2013 and high-dimensional probability |Vershynin, 2018
can have that E(x;; Wæ;
[2) > >E(x; Wæyl for
D PROOF OF MALIGNANT ENTROPY COLLAPSE
Proof. Step 1. To prove the sparsity of A, we must show that the number of non-zero elements in each column of A is small.
Now, let's consider the properties of the matrix P >=x Since and W are low-rank matrices, the matrix P will also be low-rank. Specifically, the rank of P is bounded by the rank of X and W, which is much smaller than the dimension. In particular, P has only small number of significant singular values. This implies that the entries in P are concentrated in a lower- dimensional subspace.
When we apply the softmax function to the rows of P, the function concentrates most of the probability mass a few components of each column. This is because the softmax function is sharply peaked around the largest values in each row. The smaller values each row contribute less to the sum in the denominator of the softmax function, and therefore, their corresponding entries in A will be small.
Thus, in each column of A, only small number of entries will be non-zero with high probability, and the rest will be close to zero. This establishes the sparsity of A.
Step 2. To prove the low-rankness of A, we turn to proving that A is approximately low-rank.
As noted earlier, the matrix P X low-rank. Specifically, the rank of P bounded by the ranks of X and W, which are both small. Therefore, P has only few dominant singular val- ues. Assume that the softmax function does not significantly change the rank of the matrix. The rank of A is controlled by the rank of P, as the softmax operation only introduces nonlinearities that do not increase the rank.
Thus, since P has small number of dominant singular values, A, formed by the softmax of P, will also have a small number of significant singular values. This implies that A is approximately low-rank.
In conclusion, under the assumptions that X and W are low-rank matrices with sufficiently large singular values, the matrix A formed by applying the softmax function to P will exhibit the properties of both sparsity and low-rankness with high probability.
19

---

Published as a conference paper at ICLR 2025
E PROOF OF WEYL'S INEQUALITY ON SINGULAR VALUES
Our derivation depends on Horn & Johnson 1991 more background information.
2012)
Readers can refer to these material for
Proof. Before we prove Weyl's Inequality on singular values, let us review Courant- min- max principle that is important for analyzing the singular values of matrix.
Theorem 3 (Courant-Fischee
Min-max Principle for Singular Values) Let W E R mx be a matrix where m 2 n. W has ordered singular values 01 22022 2.. 00 2 0.
Then, for i = 1,2,..., we have
(W(W)=
max R
min [
min SER" dim(S)=nn-++
max I æ! EES' 1=|=1
WES |=|=1
dim
where the maximum is taken over all i-dimensional subbscess S ofR", and the minimum is taken over all unit vectors æ in S.
Let Wi = UE VT and W₂ = UE2V V be singular value decompositions of W1 and W2 with unitary matrix V [U,...., Unl. V [U u,l where vg, u, R and unitary matrix U [u,...,, uml, = [u)..... Um, where uj E
Let and be positive integers with i. s and i+ n 1. Let Span( Unl and S2 = Sp j...... Dnl: notice that dim(S)) = n n-1+1 and dim(S2) = n -- j+1. then we have
Let k = dim(S) S2),
dim(S) n2) dim(S,) + dim(()) dim(S + S2) (m ()-
2 (n (n- i +
(1+ 11
Because of the bounds assumed for and j. Thus, the subspace S₁ n2 has positive dimension k, fn i+ and we have
O++j-1(W1 W2) m (W1 + W2)
min
[I(W W2)æ112
ScR"
dim |=1|2=
max IIW, + W2)2122 nS |=]|=]
max SES [=l2=1
max W
||2all2
aESinse 1
max I| max = 2ESI |=
|
As a special case, the second part of the theorem follows directly from the general result of part (a). Specifically, for == 1, we have:
o (W o (WW)))
This completes the proof.
20

---

Published as a conference paper at ICLR 2025
F
RELATED WORKS
Training Dynamics of Transformer. Previous works have delved into understanding the training dynamics of Transformers from two different perspectives: a high-level perspective and a low- level perspective. From a high-level perspective, Scan&Snap (Tian et al. 2023a] unveiled com- plex phenomena, particularly in single-layer architectures, relating native bias. These studies linked sparse attention patterns to token co-occurrence frequencies and observed two-stage behaviors in attention logits. JoMA (Tian et al.) 2023b further improved upon previous models by incorporating residual connections and MLP nonlinearity, analyzing joint training of MLP and self-attention layers, and offering qualitative explanations for multi- layer Transformer dynamics. From a low-level perspective, two critical challenges in Transformer training have been identified: rank collapse (Dong et al. 2021 Noci et al. 2022), where attention output converges to rank 1 matrix, potentially causing vanishing gradients; and entropy col- lapse (Zhai et al. 2023), which denotes pathologically low attention entropy, corresponding to highly concentrated attention scores. In this work, we analyze and prove two different entropy collapse modes and identify the key reason for model failure is spectral energy concentration. Fi- nally, we introduce a simple but effective solution address this problem.
frequency and discrimi-
Training Stability of Transformer. ReZero [Bachlechner et al. (021)) introduces a simple yet ef- fective mechanism for improving training stability. The key innovation lies in initializing residual connections to zero, which allows networks to learn identity mappings more easily. Admin (Liu et al. 2020) introduces a new network initialization strategy tailored for Transformers to make the network train stable. DeepNorm (Wang et al. 2022) extends the concept of normalization to accommodate increasingly deeper networks. By dynamically adjusting normalization parame- ters, DeepNorm ensures stability even as network depth increases. LipsFormer Qi et al. (2023a) addresses the specific challenge of stability in transformer networks. By introducing a Lipschitz continuity constraint, Lipsformer effectively mitigates the issue of exploding gradients a com- mon problem in deep transformer architectures. This approach ensures that the network's out- put changes smoothly with respect to its input, promoting overall stability. ReZero, Admin, and DeepNorm can all be considered as an approach to control the Lipschitz constant of the net- work in the initial stage. In this work, by revisiting the training dynamics of Transformer, we can achieve a stable training only by modifying the optimizer instead of using learning rate warmup or changing the network structures as LipsFormer QQ al Dehghani et al. (023)..
2023d and QKNorm Henry al. 2020);
must-have tech- nique for ensuring a stable network training, especially in the initial phases of the optimization process. This method involves gradually increasing the learning rate from small value to the desired initial learning rate over a certain number of training steps or epochs. The cosine learn- ing rate scheduler [Loshchilov & Hutter 2016 has gained popularity due its smooth annealing properties. This schedule decreases the learning rate following cosine curve, starting from an initial value and decaying to a minimum value over a set number of epochs or iterations. Cyclic learning rates (Smith) 2017) involve systematically varying the learning rate between boundary values. The learning rate oscillates between a lower and upper bound, either linearly or follow- ing other patterns (e.g., triangular, cosine). The mentooned learning rate schedules require specification of a stopping time step T, Defazio et al. (2024) introduces a Schedule- Free approach that avoids the need for this stopping time by eschewing the use of schedules entirely.
Learning Rate Schedule. Warmup [Loshchloov & Hutter 2016 has emerged
Compared to the up- up-meniio works, the core novel contributions of our work lie on follows.
1. We present theoretical analysis for Transformer training and point out two entropy collapse modes, the benign collapse and the malignant collapse.
2. We reveal that spectral energy concentration (SEC) of W, We is the main reason of model crash.
3. We introduce 2 new optimization strategy motivated by Weyl's Inequality.
We also observe there are two works (Kosson et al. 2023 2024 discussing the needs of learning rate warmup by explicitly controlling the angular updates via Rotational Optimizer Variants and by limiting the Frobenius norm of the update relative to that of the weights. They provided some different perspectives on the necessity of the learning rate warmup.
21

---

Published as a conference paper at ICLR 2025
G SIMULATION OF THREE ATTENTION MODES
We provide a simple simulation code to simulate three attention modes, but it is important to note that the real picture
complicated. In real case, in the benign attention entropy mode, Wa W is non-symmetric positive quasi defasideeinite square matrix instead of symmetric positive definite matrix in our simulation. The code is just to demonstrate the core ideas behind three attention modes.
CODE 1: Simulation of Three Attention Modes.
import torch
import torch.nn
import matplotlib. pyplot plt
import numpy as np
#Randomly generate data and weight matrices d_q, a, num_tokens= 64, 768, 197
Wq forch.randie (d_9, d) Ak forch.randi (d_9, d)
torch. (d, num_tokens) torch.mm (Rq.I, Wk)
13
14
Normal attention mode w1 W
torch. (torch. mm T, W1), attn mapl p.softmax (dim=l)
18
Mallllgaa Malignant attention entropy collapse mode
22
u, S, forch.svdlie
23
s[0:3] torch. ( [3., 2., 1. +8:33]
24
s [[::] 0.0
w2 torch.. mm (torch. mm (u, torch dsiag(s)), v.T)
torch.. (torch, mm (X.T, w2),
27
attn map2 P.softmax (dim=1)
28
Benign attention entropy collapse mode
31
U, s, torch. svd(W)
32
W3 torch. mm (torch, mm (u, torch.diag(s). u.r)
33 p = torch. mm (torch. mm (X.T, W3), X)
attn_mast
P. softmax (dim=1)
37 Plot figures
fig, axs = pttppp.suppppp subplots 3, figsize= (15, 5))
axs [0] 1ollimmsie imshow (attn_mapl pattn_mapl.detaat detach ().nupyy ())
axs [1] .imshww (attn_map2 detach .ummpy ())
41 axs [2] imshow (attn_map3..cmm detach () .numpy )))
42
plt.show ()
22

---

Published as a conference paper at ICLR 2025
H ATTENTION MAP VISUALIZATION OF GPT
Figure visualizes the dynamic process of attention map as the number of training steps in- creases for successful and unsuccessful GPT-Small model. should be noted that the GPT model uses a lower triangular attention mask.
200
300
(a) Block (successfuD.
(b) Block (unsuccessful).
FIGURE 7: Visualization of the dynamic process of attention map the number of training steps increases for a successful and unsuccessful GPT-Small model. Attention map gradually becomes sparse and low-rank along with the training process in failure case. Please click the images to play the flash. Best viewed with Acrobat Reader.
In Figure 7
the attention values in successful case distribute different position, but the tention values in a unsuccessful case will only concentrate into several directions.
I MORE TRAINING DYNAMICS OF VIT AND GPT
Figure visualizes a successful ViT training process. Compared with Figure we find several significant differences as follows.
successful VIT training process, the value of o (W, (W)) increases 16,000, then starts to oscillate smoothly. But for an unsuccessful training, the value suddenly increases to a very large value, around 300,000, it triggers the model crash,
The yı and ß1 in successful VIT training process very smooth, but they change a lot in unsuccessful case,
The fast increase of (W, W2) is accompanied by fast increase Wa and W
We can observe similar phenomenon in Figure and Figure 15 In successful GPT training process, the value of (WW, Wi) increases unsuccessful GPT training, the value increases to 20,000. The difference between the sclae of value between GPT and VIT may be due to the density and sparsity of the supervision signal. In GPT, each token will contribute gradient, but in VIT, only one class label in an image provides a supervision information.
then starts to oscillate smoothly. But for
23

---

Published as a conference paper at ICLR 2025
(a) 1 12
(b) 1 l2
(c (
0677
St9
5xp
oi (We)
(e) σ (Wg T W&)
(f) (W
( on (Wa)
(h) l122
(i) 1,3g
"T
① o (W1)
(k) 0o (2)
01 (W)
Sta
(m) 0 W.)
(n) ||æ|||
(0) llz
FIGURE 8: Training dynamics of successful VIT training.
24

---

Published as a conference paper at ICLR 2025
(a) 11 12
(b) (W
(c) (We)
54g
(d) oi (Wq W2.
(e) 0 (W
(f) o1 (W
1r2 ll
(h) (W1)
To (W2)
(i) σ ( Wn 1
σ (W. W.)
Crrbel
|
(m)
FIGURE 9: Training dynamics of successful GPT training.
25

---

Published as a conference paper at ICLR 2025
dy phool
(a) 11 12
(b) (W
(c) (W₂)
(d) o ( W2
(e) or (W,)
(f) o (Wo)
(g) Ir2 ll2
(h) 0n (W1)
σ (Wo (.
( o (W2)
(i) 1 (W W1
|
(m)
FIGURE 10: Training dynamics of an unsuccessful GPT training.
26

---

Published as a conference paper at ICLR 2025
J EXPERIMENT OF 1B VIT
To further evaluate the effectiveness of our method at a larger scale, we assessed ViT-g with 1B parameters. The ViT-g model architecture consists of 40 layers with hidden dimension 1408, 16 attention heads, and an MLP dimension of 6144. The total parameter count is 1011M, around one billion parameters. We conducted a comparative study between ViT-g with AdamW and wit- with Adamm, where ViT-g with AdamW was evaluated under two settings: with and without learning rate warmup. Our A does not use warmup. The comparison results are presented in Figure 11 and Figure 12
Figure 11 shows that ViT-g with AdamW crashes after only a few training steps when running without warmup. While the use of warmup enables ViT-g complete training, but the loss spikes one time. Our ViT-g with AdamW not only achieves stable training without warmup but also demonstrates better performance.
ImageNet Validation Loss
VIT-g
VIT-9 wo warmup
VIT-9 AdamW?
5.0
Lo30
4.0
3.0
2.0
1.0
140
160
Epoch
FIGURE 11: Comparison of loss curve of 2 and AdamW on ViT-g model.
ImageNet Validation Accuracy
80.0
70.0
60.0
50.0
Accuracy
40.0
20.0
10.0
ViT-g
warmup
0.0
AdamW?
20
60
100
120
140
160
Epoch
FIGURE 12: Comparison of accuracy of 2 and AdamW on ViT-g model.
27

---

Published as a conference paper at ICLR 2025
K EXPERIMENT OF 774M NANOGPT
We also evaluated the effectiveness of our method larger-scale scale language model, termed as nanoGPT-large.. The model architecture consists of 36 layers with a hidden dimension of 1280 and 20 attention heads. The total parameter count is 774M. Our experimental setup strictly fol- lows the nanoGPT configuration, including all learning rate settings. It is important to note that training nanoGPT-large computationally intensive, requiring two weeks to train 600K steps on 16 A800 GPUs. To reduce the training time, we limited our training to 100K steps instead of the full 600K steps. The comparison results are presented in Figure 13 We can see from Figure 13 nanoGPT-large achieves stable training without warmup and obtains a similar validation loss with its counterpart, GPT2-large. This further verifies our understanding to the model crash of Transformer.
OpenWebText Val Loss
3.00
GPT2-L AdamW?
2.80
Loss
2.70
2.60
2.50
60k
80k
Step
FIGURE 13: Comparison of validation loss of 2 and AdamW on nanoGPT-large model.
EXPERIMENT OF FLATTEN-SWIN
Besides VIT. GPT, and Swin-Transformer, we further validated our approach Flatten- Transformer Han et al. (2023). We used Flatten-Swin and we compared the performance of our method and the baseline method training 150 and 300 epochs. Our method could stably train and demonstrate performance comparable to the baseline. This further verified the correctness of our understanding of neural network stability.
28

---

Published as a conference paper at ICLR 2025
ImageNet Accuracy
80.0
70.0
60.0
Accuracy
40.0
30.0
20.0
Flatten-Swim (30Depochs)
Flatten-Swim AdamW/ (300epochs)
10.0
Flatten-Swim (150epochs)
Flatten- Swit AdamW/ (150epochs)
50
100
150
200
250
300
Epoch
FIGURE 14: Evaluation of Flatten-Swin.
29

---

Published as a conference paper at ICLR 2025
M ACTUAL LEARNING RATE CURVE ALONG WITH TRAINING STEPS
We recorded the actual learning rate throughout the training steps, we sample one point every 50 steps. Our initial setting of the learning rate is cosine learning rate scheduler without warmup. a (W -
W > r, then a, will be truncated to 01 (Wr (,) From the figure, we observe that 71 and
72 in RMSNorm only exceed the preset T during the initial training phase and rarely exceed afterwards. For other curves, they somewhat look like curve with learning rate warmup, but can see that different blocks have different learning rates.
We also observe that shallower layers more likely violate the preset value. means the shallower layers are more likely to lead to a greater update of weight matrix and typically require smaller learning rate. Additionally, notice that for the weight matrix W2, it more prone to exceeding the preset value compared to the weight matrix W1.
6488
3012
Fly Mosa
6000
500
Value
2004
Stsp
(a) learning rate a, of Y1
(b) learning rate a of W.. We and W,
Hap essod Pyy ciocc
clocks
Value
6006
Value
3001
5004
0000
303
0000
8000
30
200
15x
150
Step
Step
(c) learning rate of W.
(d) learning rate 0f
0438
2017
2010300
Value
3006
66
Step
(e) learning rate a; of W
(f) learning rate a, of W
FIGURE 15: Actual Learning Rate Curve along with Training steps.
30

---

Published as a conference paper at ICLR 2025
N TRAINING CONFIGURATIONS
Training Configurations. We list the training configurations VIT, GPT, Swin- Transformer and and Flatten-Swin do not use learning
Flatten- Swin in T For ViT, GPT, rate warmup. For GPT, we follow the experimental configurations of nanoGPT (Karpathy) 2022). all parameters For Trin-Trannformer we use the original code provided by Liu et al. (0021). For Flatten-Swin, we use the original code provided by (Han al. 2023
same GPT2 RRadford al. 2019). For VIT, we use Timm [Wightman 2019
TABLE 2: Training configurations for VIT, GPT and Swin Transformer.
(a) Training configurations for VIT.
(b) Training configurations for GPT.
training config
VTT-B/L/g (224*)
training config optimizer
GPT-S/I.
optimizer
A2
AdamW2
(In default)
0.004 or 0.003
0.01
warmup epochs
0
warmup epochs
0
weight init
Truncated Xavier
weight init baseline learning rate weight decay
Xavier
base learning rate
1e-3
0.0006 or 0.00025 0.1
weight decay
0.05/0.1
optimizer momentum batch size
p1 B2 0.9, 0.99
optimizer momentum tokens seen each update max iters batch size
BI, 1= 0.9,0.99
1024
500,000
training epochs
150
600K or 100K
learning rate schedule
cosine decay
480
randaugment
(9,0.5)
sequence length
1024
mixup
0.8
dropout
0.0
cutmix
1.0
bfloat16
True
random erasing
0
gradient clipping
1.0
label smoothing
0.1
stochastic depth
0.1/0.5
gradient clip
None
exp. mov. avg. (EMA)
no
(c) Training configurations for Swin- Transformer.
(d) Training configurations for Flatten-Swin
training config
Swin S/B (2242)
training config
Flatten-Swin (2242)
optimizer
4
optimizer (In default)
AdamW
т (In default)
0.004
0.004
warmup epochs
warmup epochs
training epochs
300
training epochs
150 or 300
others
same as [iu al. (0021
others
same Han al. 2023
O NON-SYMMETRIC POSITIVE QUASI-DEFINIIE SQUARE MATRIX
When we mention a non-symmetric positive quasi- definite square matrix, we mean it has the following three properties,
1. W Wi is not symmetric because generally, W W: # W W.W.
WT We is a square matrix and most of its eigenvalues are larger than 0, and only very few are less than 0.0. So we call it positive quasi-definit matrix.
3. if we assume W W We is positive definite matrix, if for each element in is sampled from standard Gaussian distribution, we can prove
E [ W x E E( Way)
when i * j, see Appendix C the proof.
31

---

Published as a conference paper at ICLR 2025
P DISCUSSION ABOUT RANK COLLAPSE, ENTROPY COLLAPSE AND SPARSE YET RONK ENTROPY MATRIX
Before we start our discussion, let us see three matrices,
0
0
0
0
0
0
0
0
=
B
0
0
0
0
C
0 0 0
0
0
0
0
1
0
can see that low-rank, B sparse but not low-rank, sparse and low-rank.
In previous papers [Dong et al. 2021 Zhai et al. 2023 researchers have analyzed the problem of model crash via rank collapse of activations and entropy collapse of attention map. Dong et al. (Dong et al. 2021] attributes the model crash into rank collapse of the activations, but Zhai al. (023) think is the entropy collapse of the attention map leading to the model crash.
However, based on our analysis, we can find counterexamples for entropy collapse, and mean- while the rank collapse of the activation cannot fully describe the inner reason of the model crash (the weight matrix instead of activations). When the state of C' usually happens, the model crashed,
Rank collapse of the activations cannot reveal the underlying cause that exists in the weight matrix. Weight matrix is the inner key ingredient of the model instead of activations.
B is counterexample of entropy collapse. we observe that in some successful cases, state B occurs. According to the definition of entropy collapse, state B should lead model crash; however, our experiments show that the model remains stable in this state.
Sparse yet low-rank attention matrix is the state the attention map when a model crashs. We believe rank collapse of activations and entropy collapse of attention map describe the state of the model crash precisely. According to our analysis, the Spectral Energy Concentration (SEC) of the W 1 is the inner reason the model crash, and the sparse yet low-rank attention matrix is the phenomena observed on the attention matrix.
enough to
In summary, our paper, via rigid theoritical analysis, our paper reveals the Spectral Energy Con- centration (SEC) of the W We is the inner reason the model crash, and the sparse yet low-rank attention matrix is the phenomena that is observed on the attention matrix.
Q
ILLUSTRATION OF FIGURE 5
Illustration Arrow
According to the property of Kronecker Product, we have
rank(X @ X) rank(a))aankkeee
Since X is low-rank, then X X is also low-rank. In the following, will also prove the singular values of X @ will also strengthen the concertration of spectral energy into some directions with large singular values.
Illustration of Arrow
32

---

Published as a conference paper at ICLR 2025
According to the computation of Jacobian matrix, we have
dve((P) = =x 8 8x
Övec(W, (W))
where P= P W
Illustration of Arrow 3.
Let X E m be matrix with rank min(m, m). Denote the singular values of X 01 02 ..20, > 0 and 0y+1 =... ... mintmm,)) = 0.
Definition. (Singular Values of Kronecker Product)** For matrix X, define A(X) the set of all possible products of its singular values, i.e.,A(X) = (0:0 1 s,jjsr,
Theorem (Singular Values of Kronecker Product) For low-rank matrix X E RRax with rank r, the singular values of @ X are precisely the elements in A(X). More formally, let (uu) be the of singular values of X X. Then (ux] j wherel iij}r))
m is an orthogonal matrix, V E " is an orthogonal matrix, and Σ = diag(o I...., Jr. Or.0. Or.0.....0f Then, the Kronecker product
Proof. Consider the singular value decomposition (SVD) of X -UVV, where U m
X can be expanded as: X X (UEV @ (UEV*) (U @ U) X)(V @ V The singular values of U & U and V & V are all 1, they are composed of orthogonal matrices.
Therefore, the singular values of X & X are precisely the elements of Σ e Σ, which are exactly ( sijj
In summary, the Singular Values of Kronecker Product X @ X will also strengthen the concertra- tion of spectral energy into some directions with large singular values. this way, the singular values of dvect will also concentrate to few directions.
dvectW, W
When update W Wi according the following equation:
ac
àvec(P)
We q Wew vec(W, Welold dvec(P) ô @ve q Wx)
where is the loss function, and a is the step size. Since the singular values of also concentrate few directions, the update will lead to the singular values of W We tends to concentrate.
dvec
will
dvect Wa
Illustration of Arrow 4.
Our Theorem 1 in the paper demonstrate that the spectral energy concentration of W W and the associated dominant large singular values will lead to A to be sparse yet low-rank ma- trix.
Illustration of Arrow
Since X = A, where V is the value matrix after projection in attention module and A is the attention matrix, according to linear algebra, we have,
rank(X! min/rank(V),, rank(())
Thus, have that X+1 is also low-rank matrix.
33


Reviews:
[{"review_id": "60oXtfztYp", "reviewer": "Reviewer_QdMW", "summary": "This paper deals with the challenges of scaling Transformers without relying on techniques like learning rate warmup. It provides a theoretical analysis that identifies the concentration of spectral energy in weight matrices as a key factor behind model crashes, leading to entropy collapse. To resolve this issue, the authors propose a novel optimization strategy that smooths weight updates, preventing rapid concentration of spectral energy. Extensive experiments with various Transformer models, including ViT, Swin Transformer, and GPT, demonstrate that this approach effectively trains models.", "questions": "I have a few questions and suggestions as follows.\n\n1. In the abstract, there are mathematical notations that should either be predefined or the abstract should be written more clearly without them.\n\n2. In the Introduction, some notations are defined later in the text, such as \"$\\delta vec(P)$\" mentioned in the second contribution point, which should be defined earlier.\n\n3. The paper claims that spectral energy concentration is a key issue behind model crashes. If attention is computed without using softmax (for example, in Linear attention models like Flatten or Vicinity Vision Transformer), will the same crash occur? Is the problem specifically due to softmax, which might cause sparsity or model instability?\n\n4. The paper states that \"the attention map is a sparse yet low-rank matrix\" (in the third contribution point). Does this imply that \"$A = softmax(\\frac{P}{sqrt{d_q}})$\" is sparse? If so, please provide a reference or proof for this assertion.\n\n5. In the experimental results, please clarify the choice of using only pure encoder architectures (like ViT) or pure decoder architectures (like GPT). Is it possible to utilize an architecture that combines both encoder and decoder components?\n\n6. It would be preferable to provide proper references instead of linking to internet sources (such as Wikipedia) as seen on page 6 and elsewhere.\n\n7. Please include some recent variants of Transformers in the empirical results for comparison.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper demonstrates strong originality, with the necessary mathematical definitions and proofs included as required. It is well-structured and organized clearly. The essential mathematical framework is presented both in the main text and supplementary material. Additionally, experiments have been conducted to verify the theoretical claims made in the paper.", "weaknesses": "This paper has several areas for improvement in terms of writing, particularly regarding the use of mathematical notations in the abstract without proper definitions. A similar issue is evident in the Introduction as well.\n\nThe experiments are inadequate; the paper should incorporate additional baseline models for comparison, as it currently only discusses or modifies three models. It would be advantageous to present empirical results using well-known recent models, such as those with linear complexity in attention, like Linformer or Flatten Transformers.\n\n It would be preferable to provide proper references instead of linking to internet sources (such as Wikipedia) as seen on page 6 and elsewhere.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "Y60MPydpVY", "reviewer": "Reviewer_YrRp", "summary": "The paper analyzes a training instability associated with the self-attention layer of a Transformer by studying its training dynamics. Specifically, the paper identifies a ```Spectral Energy Concentration``` condition where-in the weights of the product of query transpose key matrix shows a rank collapse. The paper proposes to bound the base learning rate per step when the optimizer (AdamW specifically) proposes very large update vectors. The paper shows how the need to use learning rate warmup can be  relaxed for training Transformers. Empirical results are provided to support analysis and proposed algorithm.", "questions": "- Would the proposed optimizer update work at reasonable scales like the ones considered by [Worstman et al.](https://arxiv.org/abs/2309.14322)?\n\n- How different is the instability shown by analyzing weights of Q.T, K product different from logit norm growth shown previously in literature?\n\n- (clarification request) The paper describes the product matrix in line 345 in general as a non-symmetric positive quasi-definite\nsquare matrix. A reference to what these technical terms mean and/or definition in the appendix would help clarify what this means to the reader. The reviewer was unable to come up with a definition of quasi-definite for non-symmetric matrices but this maybe perhaps because I missed looking at an appropriate reference. \n\nThe concerns raised in the review is that the work has limited novelty and also limited empirical support. Therefore, the initial decision for this paper is a reject.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The analysis of the  query (transpose), key weight matrix product  is interesting. Specifically, using the singular values of this product matrix to show collapse is interesting\n- Application of Weyl's inequality to the update equation for the optimizer to stabilize Transformer training is a nice development\n- The biggest contribution of this paper is that the proposed algorithm built on analysis frees up a practitioner to use  learning rate warmup.", "weaknesses": "- The paper contains limited empirical results. Specifically, the experimental setup used to study the algorithm is small and no results are provided with larger models (0.5B, 1B, 3B). This limits the contribution as its not clear how the proposed algorithm will work with larger models (albeit can still be considered ``small'') where training instabilities are more apparent\n\n- Related to above, the QK layer norm paper by Deghani et al. (2023) show that training instabilities occur at 8B for vision transformers (ViTs). So, I wonder whether the new update along with no learning rate warmup will work as advertised\n\n- The paper proposes ``Spectral Energy Concentration (SEC) '' to quantify collapse. Wouldn't a smooth rank measure also provide the similar  information while being easier to interpret? Furthermore, this contribution is not novel as large attention logits being a problem has been observed in literature below by Deghani et al (referred to in the paper) and by [Worstman et al.](https://arxiv.org/abs/2309.14322).\n\n- The argument in Section 3.2 is not very clear. Especially the reader could use some help with understanding why the rank of X.T being small would imply the rank of X.T \\otimes X.T is still small. The equation shows a square relationship. Firming up the arguments here would be very helpful to the reader", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "TItm7DfOW6", "reviewer": "Reviewer_9a7P", "summary": "The paper presents a method to eliminate learning rate warmup for training transformes. It begins by analyzing training dynamics, specifically observing the norm of certain parameters to understand behavior patterns when training is successful versus when it crashes. The authors note that, in the presence of warmup, the parameter norms tend to increase steadily, while they diverge in failed training runs. However, this analysis does not introduce new insights, as it’s known that parameter norms typically blow up when training crashes.\nThe authors argue that warmup can be bypassed by controlling the spectral energy concentration (SEC), which maintains the singular values of the product $W_q^TW_k$​ thus allowing successful training. To achieve this, they introduce a modified version of the AdamW optimizer using a power iteration-based step to estimate the largest singular values of the $W_q^TW_k$ matrix. Experiments with the proposed method show that it can perform similar to warmup.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1 .The authors propose an alternate method for transformer training that does not need learning rate warmup.", "weaknesses": "1. The choice of parameter norms for tracking training dynamics seems arbitrary. It’s already established in literature that parameter norms diverge when training fails, so this observation does not seem novel [1]. The paper could benefit from a deeper analysis or rationale for the specific parameters chosen, or alternatively, from an exploration of novel insights that could provide a more compelling argument.\n2. While the authors claim that the attention maps are sparse and low-rank, the plots are difficult to interpret. Hence it is impossible to make conclusions based on the provided plots. It might help to change the scaling of the plots to help visualization or simply estimate the rank of the attention maps.\n3. My major concern is that in the proposed algorithm that controls the step size to ensure the SEC does not collapse, only ends up modifying the step size $\\alpha$. It seems to me that this is another way of performing learning rate warmup and the analysis further justifies the use of warmup as it may help prevent SEC from collapsing. In order to disentangle these ideas, the authors might need to discuss how the modified $\\alpha_t$ looks over time (does it resemble warmup, which seems to be the case).\nExperiments also suggest that the proposed method is in fact very similar to warmup and might be a different parameterization offering the same benefit as warmup.\n4. Further, while replacing warmup the authors introduce an additional power iteration step which is computationally more expensive and adds to each gradient update step, while learning rate warmup does not require this additional step.\n\n\n[1] Pennington, Jeffrey, Samuel Schoenholz, and Surya Ganguli. \"Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice.\" Advances in neural information processing systems 30 (2017).\n\n[2] Defazio, Aaron, et al. \"When, why and how much? adaptive learning rate scheduling by refinement.\" arXiv preprint arXiv:2310.07831 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "pUtkKY5QEv", "reviewer": "Reviewer_1Y1q", "summary": "This paper visualizes the training dynamics of transformer models in both successful and failing cases. Based on these visualizations, the authors observe that the rapid growth of the maximal singular value of the product of the query and key matrices,  W_q^T W_k , is a critical phenomenon associated with model crashes. They provide insights suggesting that when  W_q^T W_k  becomes concentrated along the directions of the top singular vectors, the attention probability matrix becomes sparse and low-rank. They also show crashed models typically exhibit significant spectral energy concentration. Finally, the authors propose a variant of AdamW that addresses this issue by clipping the learning rate to prevent the spectral norm of all parameters from increasing too quickly.", "questions": "See weakness above", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper provides insightful experiments demonstrating that the rapid growth of the spectral norm of the query-key product matrix (and other parameter matrices) is a crucial factor in model crashes. The study shows that in crashed models, the spectrum of the query-key product matrices often concentrates on the top eigenvalues. This provides novel insight on understanding transformer optimization.", "weaknesses": "- The theory part does not enhance the paper as effectively as it could.\n    - In Theorem 1, the statement is somewhat vague. Terms like “malignant collapse” and “spectral energy concentration” are not formally defined, making the theorem unclear. Additionally, the proof of Theorem 1, which appears in Appendix D, lacks clarity and coherence. For example, the theorem aims to establish a sufficient condition for malignant collapse, yet the proof assumes the low rank of  W  due to malignant collapse. Furthermore, the proof relies heavily on the assumption of Gaussian distribution.\n    - I would personally suggest rewriting Theorem 1 as a highlighted insight or intuition, supported by experimental results, such as Figure 1 and 4.\n- In Figure 5, could you elaborate on how the concentrated and low-rank structure of  $\\frac{\\partial \\text{vec}(P)}{\\partial \\text{vec}(W_q^T W_k)}$  leads  W_q^T W_k  to become low rank? It still requires reshaping the matrix, multiplying the activation gradient, and summing over data from  $\\frac{\\partial \\text{vec}(P)}{\\partial \\text{vec}(W_q^T W_k)}$  to the update of  W_q^T W_k, so this implication is not straightforward.\n- Could you clarify how one would generally set the parameter $\\tau$ for an unseen model, and how $\\tau$ might vary as the model scales up?\n- Additionally, does benign collapse occur in real training tasks?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)