Spike No More: Stabilizing the Pre-training of Large Language Models
Authors: Sho Takase
Shun Kiyono
Sosuke Kobayashi
Jun Suzuki
Venue: COLM 2025 (2025)
Keywords: stable training
llm
pre-training
TL;DR: Theoretical analysis to prevent loss spikes during LLM pre-training

Abstract: Loss spikes often occur during pre-training of large language models.
The spikes degrade the performance of large language models and sometimes ruin the pre-training.
Since the pre-training needs a vast computational budget, we should avoid such spikes.
Based on the assumption that the loss spike is caused by the sudden growth of the gradient norm, we explore factors to keep the gradient norm small through an analysis of the spectral norms of the Jacobian matrices for the sub-layers.
Our findings suggest that stabilizing the pre-training process requires two conditions: small sub-layers and large shortcut.
We conduct various experiments to empirically verify our theoretical analyses.
Experimental results demonstrate that methods satisfying the conditions effectively prevent loss spikes during pre-training.

Published as a conference paper at COLM 2025
Spike No More:
Stabilizing the Pre-training of Large Language Models
Sho Takase & Shun Kiyono
Sosuke Kobayashi & Jun Suzuki Tohoku University
SB Intuitions
{sho. takase, shun. hwn.kkiiooon@siinuttttiii co.
soskēpreferred. jun. hmm.suzukk@toooku.ac.ppp
Abstract
Loss spikes often occur during pre-training large language models. The spikes degrade the performance of large language models and sometimes ruin the pre-training. Since the pre-training needs a vast computational budget, we should avoid such spikes. Based the assumption that the loss spike is caused by the sudden growth of the gradient norm, we ex- plore factors to keep the gradient norm small through an analysis of the spectral norms of the Jacobian matrices for the sub-layers. Our findings suggest that stabilizing the pre-training process requires two conditions: small sub-layers and large shortcut. We conduct various experiments to empirically verify our theoretical analyses. Experimental results demon- strate that methods satisfying the conditions effectively prevent loss spikes during pre-training.
1 Introduction
Large language models (LLMs) have been fundamental assets for various applica- tions (Brown et al., 2020; Chowdhery et al., 2022; Touvron et al., 2023). Increasing the num- ber of parameters in (neural) language models and the number of training data usually leads to better LLMs (Kaplan et al., 2020). Consequently, pre-training requires a vast budget, and thus, minimizing the risk of failure of the pre-training is a paramount concern.
375
Scaled Embed
3.50
Train
3.25
3.00
277
2:50
0
5000
10000
15000 20000
25000 30000 35000
Train
Despite their widespread use as the founda- tional architecture for LLMs, a comprehen- sive theoretical understanding of Transform- ers (Vaswani et al., 2017) has not yet been achieved. One of the crucial unresolved ques- tions is the reason for the frequent occur- rence of pre-training failures in Transformer- based LLMs due to spikes in loss values (loss spike) that can lead to catastrophic diver- gence (Chowdhery et al., 2022) illustrated in Vanilla in Figure 1. While several empirical strategies have been proposed to mitigate this problem (Chowdhery al., 2022; Le Scao et al., 2022; Zeng et al., 2023), the absence of theoretical justification for these methods casts unclear on their generalizability to other situations, such as varying sizes of model parameters.
Figure 1: Training loss values of Trans- formers, whose dimensions and the number of layers are the same as the 1.7 billion parameters configuration in Narayanan al. (2021). In Vanilla, some spikes occur at the beginning of the training, and its loss value ex- ploded at about 13000 steps.
In this research, we provide theoretical analyses focusing on the loss spike problem during LLM pre-training. We identify the upper bound of the gradient norms for the Transformer- based LLMs through analyses on the spectral norms of the Jacobian matrices for the sub- layers. If the upper bound is large, the gradients may spike suddenly, and assume that

---

Published as a conference paper at COLM 2025
this phenomenon causes the loss spike. Then, we indicate that the upper bound is large in the typical setting, such as the widely used implementation, Megatron-LM (Shoeybi et al., 2020), and thus, the loss spike is likely to occur. In addition, to make the upper bound sufficiently small, we introduce two conditions: (1) initializing the parameters of sub- layers with small value and (2) making the standard deviation of each embedding close to 1. The former condition can be satisfied by the widely used initialization method for LLMs (Shoeybi et al., 2020; Le Scao et al., 2022; Biderman et al., 2023). On the other hand, the latter condition was satisfied in the original Transformer by scaling embed- dings (Vaswani al., 2017), but such scaling is missing from recent implementations. To sum up, through theoretical analyses, we re-evaluate several previous techniques in terms of the stabilization of LLM pre-training.
Building on our theoretical analysis, we further substantiate our claims through a series of empirical experiments, which provides a clear distinction between effective and ineffective methods over different training scenarios. Our results demonstrate that methods satisfy- ing the conditions avoid the occurrence of loss and gradient spikes. contrast, methods that fail to meet these conditions remain susceptible to gradient spikes, even when pre- viously recommended as empirical solutions of the loss spike problem. Furthermore, we demonstrate that method satisfying the conditions enables LLMs to be pre-trained with a comparatively larger learning rate, leading to superior performance outcomes.
2 Preliminary
2.1 Pre-LN Transformer
This paper mainly focuses on the neural architecture used in the series (Radford et al., 2018; 2019; Brown et al., 2020). They use the Prre-LN Transformer (Xiong et al., 2020), which the de facto standard architecture in recent implementations of Transformers because the training with the architecture is more stable than the original Transformer architecture when we stack many layers (Xiong et al., 2020; Liu al., 2020; Takase et al., 2023). Let x e R2 be an input of layer of the Transformer, where d denotes the dimension of the layer. The layer outputs y with the following equations:
y x + FFN(LN(X))),
x' = x + Attn(LN(x)),
(2)
where LN is the layer normalization function'. We call the first terms in Equations (1) and (2), i.e,, x and x', shortcut. In addition, the ford-forward network (FFN) and multi-head self-attention (Attn) are defined as follows?:
FxN = W2 x)),
(3)
Attn(x) = Wo Wo(concat(heady (x),,,,
(4)
(Wor x) [(Wii x) (Wiz X)T,
head;(x) = softmax
(5)
Vhhead
where F is an activation function, concat concatenates input vectors, softmax applies the softmax function to an input vector, and W1 E Rdma xd W2 E Rdxdm, Wa: E R Xi
WKi E R xd Wy: E Rhedd xd and Wo E Rdxdd are parameter matrices, and din and dhead are the internal dimensions of FFN and multi-head self-attention sub-layers, respectively. In addition, we pack the sequence of input vectors into a matrix as X E RáxL, where L is the input sequence length, to compute the self-attention.
We discuss the difference from the architecture using Root Mean Square layer normalization (RMSNorm) (Zhang & Sennrich, 2019) instead of LN in Appendix B.4, and the original Transformer architecture, i.e., Post-LN Transformer in Appendix H. 2To simplify equations, we omit bias terms.
2

---

Published as a conference paper at COLM 2025
2.2 Gradients of Pre-LN Transformers
Let L be the loss function of the N layered Pre-LN Transformer and Jn be the Jacobian matrix of the n-th layer. We can calculate the gradient of & using the relations in Equations (1) and (2) as:
ac dc N-1
ac
N-1 II ==1
dyn dx dx dx,
dyn
dyn dxny
II I ==1
where
Jn
(6)
dx1
dyn
dyN
Zxn
ax dxn
Using the submulttplicatiiity of the spectral norm, i.e., ||AB||2 S (6), we can derive an upper bound of the norm of the gradient of as:
||||||]BBl and Equation
ac
ac
N dy dx
ac
N-1
dy
ax
ax1
dyN
dx dxnn
āyn
ax
dxn
#=1
==1
Thus, we can estimate the upper bound of the gradient norm of L by analyzing the spec- tral norms of the Jacobian matrices for the FFN layer and the self-attention layer, namely,
dyk dxg
12 and dx, 112-
dxy
2.3
Motivation to Suppress the Upper Bound
In our preliminary experiments, when the gradient norms grow suddenly during LLM pre-training, we observe that the loss spike problem is likely to occur. Thus, we assume that we can prevent the loss spike problem by maintaining the gradient norm small. To prevent the growth of the gradient norm, we explore the way to suppress the upper bound described by Equation (7). To suppress the upper bound, we analyze the Jacobian matrices to find a factor to control the upper bound in the following sections, and then, provide two conditions: small sub-layers and large shortcut. We verify our assumption and theoretical analyses through experiments on LLM pre-training.
3 Analyses on Gradients of Sub-Layers
For the theoretical analyses in this section, we employ the following assumption:
Assumption 1.
Let x and x' be the input and intermediate vectors of each layer. Moreover, let W. denote the model parameter matrix in each layer. We assume that x, x', and W, for all layers follow a normal distribution with mean of ie,, μ = 0.
This assumption is valid when we initialize parameters with the normal distribution, the number of heads in Equation (4) and iis identity function. Empirically, the outputs of each sub-layer are close to the normal distribution as illustrated in Appendix D.
3.1 Jacobian Matrix of FFN
Based on Equation (1),
dyn dz
llz in Equation (7) can be rewritten as:
āy Ox
a(x' FFN(LNN(Z)))
= I+ a(FFN(LN(x)))
(8)
dx'
dx
2
2
We can then derive an upper bound of
dyn 2xc
by applying the subadditivity, i.e., ||A +
Bl]2 < ||
||B||2, and submultiplicativity properties of the spectral norm as follows:
dy
OFFN(LN(x')) aLN(x') aLN(x')
11
πx
dx'
The right-hand side of this inequality indicates that we can estimate the upper bound of dy
x lla by separately computing the spectral norms of Jacobian matrices for FFN and LNN
3

---

Published as a conference paper at COLM 2025
Regarding the FFN part, we assume that the activation function F is identity function to simplify the discussion. Under this assumption, the following equation holds:
)FFN(LN(2')) aLN(x')
(10)
Therefore, we can straightforwardly derive the relation W|W2Will2 < ||Wil22|| W2ll2 from the submultiplicatiiity of the spectral norm. Furthermore, let 01 and 02 be the standard deviations of W1 and W2, respectively. From Assumption 1, the spectral norms of Wi and W2 are obtained by their standard deviations and dimensions (Vershynin, 2018), iee, ||ill2 R 01 (Vd + vddt)) and | x 02(vd vdtn). Finally, we can express an upper bound of the spectral norms of the Jacobian matrices for FFN the following inequality:
aFFN(LN(x')) aL
S <002(Vd vdmt))
(11)
where the right-hand side has the relation (102(Vd + vmm) = Welll
Next, regarding the LN part, the Jacobian matrix of LN can be written as:
aLN(x')
Vd (
I x'x'T | VVVd ( (1- --
(I- x'x'
- 11-221 2 0x
1
(12)
dx
|
The leftmost equation appears in the proof by Xiong et al. (2020). The second equation uses
||xl|22 = V., which can be obtained based on Assumption 1. The last equation is derived from the well-known formula of z (x' /Oz, which converts a normal distribution, x', to the standard normal distribution z, where Ml = 0 in Assumption 1.
We consider the variance (var) of each element in the matrix TzT. Since ž follows X²2 with degree of freedom, and 2iz,(i + j) the multiplication of two independent values following the standard normal distribution, the variances are as follows:
if
var(z,z,) = 2 otherwise
(13)
2z = in LLMs due to d >> 1. Therefore, the spectral norm of Equation (13) indicates that J the Jacobian matrix of LN can be written as:
aLN(x') dx'
1 oF
aLN(x') dx
1 1.
where
(14)
Ox'
Finally, Equation (9) can be rewritten by substituting Equations (11) and (14) as:
dy Ox
0902
< 1 +
Cin
(15)
02
where Cit = (vā + V ) for the simplification.
According to the discussion in Section 2.3 and Equation (15), the standard deviations, 01 and 02 of W2 and W2, respectively, should be sufficiently small, and the standard deviation, Tz', of the shortcut, x', should satisfy 0102 < dz order to keep the upper bound small.
3.2 Jacobian Matrix of Self-Attention
Similar to FFN, we can rewrite
dr' dx 2| in Equation (7) by using Equation (2) as:
dx'
0(x + Attn(LN(x)))
| a(Att(LN(())) dx
(16)
dx
dx
3 AAppendix E discusses the case where we use the ReLU, SIL.U, and SwiGLU as the activation functions, which leads to the same conclusion.
4

---

Published as a conference paper at COLM 2025
We can then derive upper bound of || tiplicativity of the spectral norm, namely:
dn' oz
2l2 by applying the subadditivity and submul-
dx'
(Attn(LN(x)) aLN(x)
xLN(x)
- ++
(17)
dx
2
Therefore, estimate the upper bound of || àr' Jacobian matrices for Attn and LN.
or
ll, we compute the spectral norms of the
Let Z(-) concat(heads (+), ..., head, and let 1 be the Jacobian of the Z(-)*, we can rewrite the spectral norm of the Jacobian matrix of Attn as:
@Attn(LN(x)) xx)
aWoZ(LN(x)) (Z((N(x))) 1/ aZ((N(x))
(18)
aLN(x)
Therefore, we can straightforwardly derive the relation submultiplicatiiity of the spectral norm.
/212 [oollll from the
Let to be the standard deviation of Wo. The relation ||Wollz = (2vd) is derived from Assumption 1. We assign this value Equation (18) and obtain the following inequality:
aAttn(LN(x)) aLN(x)
S <0o(2\0)11211.
(19)
Therefore, we rewrite Equation (17) by substituting Equations (14) and (19) as follows:
Ox'
00 11
CAun
(20)
dx
0+
where CAttn = (2V0)1112112 for the simplification.
Thus, similar to the discussion at the end of Section 3.1, the standard deviation, do, of Wo should be small and the standard deviation, xx, of the shortcut, x, should satisfy do & Cx in order to keep the upper bound small.
Conditions to Avoid Spikes
Based on the discussions in Section 3, we have to pay attention to values of 01, 02 do, and the standard deviation of the shortcut to stabilize the pre-training of LLMs. To make U2, and do small, we have to initialize the corresponding parameters with a small value. Let us consider the actual settings in detail. The widely used initialization method for LLMs (Shoeybi et al., 2020; Le Scao et al., 2022; Biderman et al., 2023), initializes all param-
eters with normal distribution (0, o²) where o = V 2
5a (Nguyen & Salazar, 2019), and
then scales W22 and Wo to small values based on the number of layers: v 2N where N is the number of layers In this situation, 01 02 and to are sufficiently small values.
However, in this situation, the standard deviation of the shortcut is also too small. For example, at shallow layers, the standard deviation is close to V 2 because the embed-
ding matrix is also initialized by (N, o2) where o = = Therefore, to increase the standard deviation of the shortcut, we make the standard deviation of each embedding
We discuss the detail of 1² in Appendix G.
BBiderman et (2023) also scaled W2 and Wo to small values in the initialization, but they used the strategy introduced by Wang & Komatsuzaki (2021) instead of scaling with y 2N- However, its property is the same essentially because they initialize W2 and Wo with o NVa which becomes small based on the number of layers.
5

---

Published as a conference paper at COLM 2025
close to 1 To achieve this, introduce two kinds of modification: "Scaled Embed" and "Embed LN The Scaled Embed scales embeddings with an appropriate value. For example, we multiply embeddings by d, which was used in the original Transformer pa- per (Vaswani et al., 2017)², and then the standard deviations of embeddings become / The Embed LN applies the LN to embeddings. In fact, Le Scao et al. (2022) reported that the Embed LN strategy prevents loss spikes empirically. These two methods are presented as verification examples rather than proposed methods, and alternative approaches could be employed if the conditions are met.
To demonstrate the actual values of the up- per bound described in Equation (15), we take the model with 1.7 billion parameters as an example. In addition to the widely used initialization for LLMs (Vanilla) and the above two modifications: Scaled Em- bed and Embed LN, compare Xavier Init, which initializes all parameters with the Xavier initialization (Glorot & Bengio, 2010), as the situation where we do not scale W2 and Wo based the number of layers. Figure 2 shows the values of Equa- tion (15) for each layer at the beginning of the pre-training. This figure indicates that the methods without suppressing the upper bound, i.e., Xavier Init and Vanilla, rapidly increase the values especially shallow layers. In contrast, Scaled Embed and Embed LN keep small values. In summary, to make the upper bound of the gradient norms small for the stabilization of the LLM pre-training, we have to satisfy two condi- tions: (1) small sub-layers; initializing the parameters of sub-layers with a small value and (2) large shortcut; making the standard deviation of each embedding close to 1.
Xavier Init
Upper boand
Vanilla
Embed LN
10"
Scaled Embed
8
16
24
Layar depth
Figure 2: in Equation (15) for each Transformer layer at the beginning of the LLM pre-training. Be- cause it difficult ues for Ox at all layers, we obtain the empir- ical values by using some inputs, and assign them to Equation (15).
The actual upper bound described
estimate the strict val-
Experiments
We verify the empirical effectiveness of our theoretical analyses. In detail, we demonstrate that controlling the upper bound of the gradient norms also prevents loss and gradient spikes. To assess efficacy in the real situation, we focus on the methods initialized with the widely used method (Shoeybi et al., 2020; Le Scao et al., 2022) in main experiments"
5.1 Datasets
We used C4 (Raffel al., 2020) that consists of clean English texts extracted from Common Crawl as our LLM pre-training corpus. We also used the separated part of C4 as our valida- tion data. We used GPT-2 vocabulary (Radford et al., 2019) that contains Byte Pair Encod- ing (BPE) subword units (Sennrich et al., 2016) as our vocabulary. To evaluate each method,
Based on Equations (15) and (20), the upper bound becomes small the standard deviation of the shortcut increases. However, too large value degrades the performance empirically as described Appendix C.
We can satisfy the condition by initializing embeddings with the normal distribution (O0 o2) where o 1, but do not adopt this strategy this study because we use the same initialization method our experiments.
AAthoough the original Transformer paper introduced this operation, recent implementations ig- nore this. Therefore, I.IMS trained with recent implementations do not satisfy large shortcut.
9hhe Xavier initialization, which does not satisfy the small sub-layers condition, is not widely used for LLMs in recent years. Appendix B.1 shows that the performance of Xavier initialization is worse and fails to avoid spikes.

---

Published as a conference paper at COLM 2025
computed perplexity on WikiText (Merity et al., 2017) and LAMBADA (Paperno et al., 2016) datasets.
5.2 Model Configurations
As described in Section 2, we used the Pre-LN Transformer architecture. We conduct ex- periments with two parameter sizes: 350 million (350M) and 1.7 billion (1.7B). We set the learning rate (Ir) 5.0 x 10-4 Appendix A describes more details on the experimental con- figuration. We compared the following methods. We put v before the method name if the method satisfies both conditions to suppress the upper bound.
This is the standard configuration for LLM pre-training Since this configuration does not suppress the upper bound of the gradient norms, the loss spike is likely to occur.
Vanilla
Embed Detach Zeng et al. (2023) used the shrink embedding gradient tech- nique (Ding al., 2021) to stabilize their LLM pre-training, This method shrinks gradients
on the embedding layer by detaching a part of embeddings from the computational graph as follows:
Embed + yEmmed + (1 - (1 Detac)))EEEbebeee
(21)
where y is hyper-parameter and Detach detaches an input from the computational graph. We assign 0.1 to Y as in Zeng et al. (2023). Zeng al. (2023) indicated that this method empirically prevents the loss spike. However, this method does not satisfy the condition on large shortcut, and thus, we show that this method does not completely solve the loss spike.
Embed LN Dettmers et al. (2022) and Le Scao al. (2022) reported that applying the LN to the embedding layer stabilizes their LLM pre-training, As described in Section 4, this method satisfies the conditions to control the upper bound of the gradient norms.
/ Scaled Embed This method multiplies embeddings by Vd. As described in Section 4, this method satisfies the requirements to control the upper bound of the gradient norms.
5.3 Main Results
Figure 3 shows the loss values of each method in validation data. Figure 4 shows the gra- dient norms of each method. These figures indicate that Vanilla and Embed Detach faced loss and gradient spikes. In contrast, Embed LN and Scaled Embed did not face spikes. These results correspond to our theoretical analyses described in Sections 3 and 4. Thus, only methods that make the upper bound of the gradient norm small have successfully avoided spikes in LLM pre-training.
In comparison between 350M and 1.7B parame- ters, spikes occurred more frequently in 1.7B pa- rameters. Because we initialize embeddings with
Model
Wiki'eet LAMBADA I
350M parameters
Vanilla
30.03
24.73
(0, where o = V the standard devia- tions of embeddings become small as d gets larger in Vanilla and Embed Detach. This means that the upper bounds described by Equations (15) and (20) become large as d gets larger because dx and of are nearly equal to the standard deviation of an input embedding in shallow layers. Therefore, we increase without any technique to control the upper bound of the gradient norms, model becomes more unstable. This result corresponds to the previous study reports (Le Scao et al., 2022;
Embed Detach
30.69
26.93
Embed LN
29.85
25.03
Scaled Embed
29.86
24.37
1.7B parameters
Vanilla Embed Detach
22.58
15.22
22.00
13.88
Embed LN
21.29
13.00
Scaled Embed
21.29
12.53
Table Perplexities of each method.
7

---

Published as a conference paper at COLM 2025
4.00
4.00
Vanilla
3.75
Embed Datach
3.75-
Embed Detach
Embed LN
Embed LN
3.50
Sce od
Valid loss
3.50
Embed
3.25
3.25
3.00
3.00
2.75
2.75
2.50
5000
10000
15000
20000 25000 30000 35000
5000
10000
15000
20000 26000 30000 36000
Train step
Train
(a) 350M parameters.
(b) 1.7B parameters.
Figure 3: Loss curves of each method in validation data.
Vanilla
Vanill
Embed Detach
Embide Detach
Embed
Embed
Grad poa5moomm
Scaled Embed
Crad norm
Scaled Embed
9050
10000
16000
20000
25000
30000
39000
5000
10000
18000
20000
25000
30000
35000
Trahn step
Troin step
(a) 350M parameters.
(b) 1.7B parameters.
Figure 4: Gradient norms of each method during the training.
Chowdhery al., 2022; Zeng al., 2023) that their model became more unstable they increased the number of parameters.
Table 1 shows the perplexities of each method WikiText and LAMBADA. This table shows that Embed LN and Scaled Embed achieved comparable performance, This result implies that methods have no significant difference from each other in their performance each method prevents loss and gradient spikes. In contrast, the perplexities of Vanilla and Embed Detach are worse except for Vanilla with 350M parameters in LAMBADA, and the difference in the performance is larger in large amount of parameters. This result implies that addressing spikes has a more serious influence on the performance as the parameter size gets larger. We conduct experiments on larger model in the following subsection. Moreover, we discuss other settings in Appendix B.
5.4 Results on 13B Parameter Model
We conduct experiments on pre-trainings of 13B parameter models to indicate that our provided conditions can stabilize a model with many more parameters than the ones dis- cussed in Section 5.3. Due to the limitation of our computational budget, focused the comparison between Vanilla and Scaled Embed. We tried two learning rates: 3.0 10-4 and 1.0 10
Appendix describes more details about hyper-parameters.
Figure 5 shows the loss values of each configuration in validation data. As shown in (a) of this figure, the loss value of Vanilla rose from approximately 10000 steps in Ir = 3.0 x 10-1 Then, the gradient of this model became too large to continue its pre-training, In contrast, the loss value of Scaled Embed consistently decreased. This result indicates that Scaled Embed stabilized the pre-training of the model with a large number of parameters.
Table shows the perplexities of each configuration in evaluation data. This table indicates that we can achieve better performance when we use a larger learning rate. In addition, the perplexities of Scaled Embed were comparable to ones of Vanilla when used the small
8

---

Published as a conference paper at COLM 2025
3.5
Scaled Embed
Scaled Embod
3.00
3.00
277
2.75
2.50
2.50
10000
20000
30000
40000
50000
10000
20000
30000
40000
50000
Train
Train
(a) Ir 3.0 10
(b) Ir 1.0 10
Figure Loss values of each method with 13B parameters when we use two learning rates:
= 3.0 x 10 and 1.0 x 10-4
Wikilext
LAMBADA
Model
Ir = 3.0 10
Ir 1.0 x 10-4
Ir 3.0 x 10
= 1.0 x 10
Vanilla
N/A
15.12
N/A
6.50
Scaled Embed
14.47
15.25
5.97
6.53
Table 2: Perplexities of each method with 13B parameters when we use two learning rates: 3.0 10-4 and 1.0 10-4
Model
PIQA
Opee
HellaSwag
WinoGrande 1 ARC-easy ARRC-chaleenge
Ir 3.0 10 4
Vanilla
N/A
Scaled Embed 78.94
39.20
71.03
63.77
60.31
35.49
Ir = 1.0
10
Vanilla 77.80
38.40 39.20
69.10
61.64
57.95 58.92
33.53
Scaled Embed
78.07
68.92
62.59
33.19
Table 3: Performance of each method with 13B parameters on standard tasks.
learning rate: = 1.0 x 10 4, Table 3 shows the performance of each configuration on the standard benchmark datasets: PIQA (Bisk et al., 2020), Open (Mihaylov et al., 2018), HellaSwag (Zellers et al., 2019), WinoGrande (Sakaguchi et al., 2021) and ARC easy and challenge (Clark et al., 2018). This table shows results consistent with the evaluation in terms of perplexity. These results imply that our provided conditions have no considerable risk in pre-training. Thus, we have to satisfy large shortcut in addition to small sub-layers to stabilize the pre-trainings of LLMs.
Related Work
Stability To stabilize trainings of Transformer-based neural language mod- els, there have been various discussions on the architecture (Xiong et al., 2020; Liu al., 2020; Takase et al., 2023; Zeng al., 2023; Zhai et al., 2023), initialization method (Nguyen & Salazar, 2019; Zhang et al., 2019; Huang et al., 2020; Wang et al., 2022), training strategy (Zhang et al., 2022; Li al., 2022), and loss function (Chowdhery et al., 2022; Wortsman et al., 2023).
Xiong et al. (2020) theoretically analyzed gradient scales each part in Transformers, and indicated that the Pre-LN Transformer is more stable than the Post-LN Transformer, that is the original Transformer architecture (Vaswani et al., 2017). Since the Pre-LN Transformer is more stable than the Post-LN Transformer theoretically and empirically, recent studies mainly have used the Pre-LN Transformer to construct an LLM. We also assume using the Pre-LN Transformer in the analysis on the training dynamics in this paper.
To stabilize the LLM pre-training, Le Scao et al. (2022) applied the layer normaliza- tion to the embedding layer. Zeng et al. (2023) used shrink embedding gradient tech-

---

Published as a conference paper at COLM 2025
nique (Ding et al., 2021). Nishida et al. (2024) proposed weight scaling reparameteri- zation (WeSaR) which uses additional parameters to scale parameters of internal layers. In this study, we theoretically proved that the layer normalization to the embedding layer controls the upper bound of the gradient norms sub-layers when we use the widely used initialization method for LLMs (Nguyen & Salazar, 2019; Shoeybi et al., 2020), and thus, it stabilizes the pre-training.
For the initialization methods, recent studies have proposed maximal update parameteri- zation (μIP) and its variants as approaches to transfer hyper-parameters from small mod- els larger ones without incurring the cost of hyper-parameter search Yang Hu (2021); Yang et al. (2024a;b). However, as described in Yang et al. (2024b), such methods may be ineffective when applied to Transformers, as their underlying assumptions do not accu- rately reflect the actual Transformer architectures. In studies focusing on Transformers, Nguyen & Salazar (2019) proposed strategy to initialize parameters of Transformers with small values to stabilize their training. Zhang et al. (2019) and Huang et al. (2020) indi- cated that we can remove layer normalizations in Transformers if we use their proposed initialization methods. Wang et al. (2022) adjusted initial parameter scales based on the number of layers to stabilize the Post-LN Transformer. In this study, we indicated that the widely used initialization method (Shoeybi et al., 2020), which makes parameters small, is necessary to stabilize the LLM pre-training. Moreover, we proved that we can prevent the loss spike problem by making the standard deviation of embeddings close to
Efficiency As shown in Section 5.4 and Appendix B.2, our modification enables the pre-training with relatively larger learning rate, and can achieve better performance. Thus, this study can be regarded as the efficiency of LLM pre-training because our modification can construct a better LLM with a given budget. Strubell et al. (2019) and Schwartz et al. (2019) reported that recent neural methods require substantial computa- tional costs, and thus, they argued that we have to explore a cost-efficient approach. Rajbhandari et al. (2020) proposed ZeRO that reduces memory redundancies during the multi GPU training without increasing communication volume. Dao et al. (2022) focused on GPU memory maad/s and proposed FlashAttention that accelerates the speed of attention mechanisms in Transformers. To reduce the number of computations in the attention mechanism, Shazeer (2019) proposed the multi-query attention that shares one key and value across all of the attention heads in each layer. Takase & Kiyono (2023) ex- plored several parameter sharing strategies, and indicated that parameter sharing across some layers can achieve comparable performance to the vanilla model with small num- ber of parameters. Moreover, several studies have explored a better construction way with a limited budget (Izsak et al., 2021; Takase & Kiyono, 2021). We believe that we can take advantage of their findings to make our LLMs more efficient.
7 Conclusion
This paper explored why large language models (LLMs) sometimes experience loss spikes during pre-training. To provide evidence, we specifically focused on the gradients of sub- layers. We introduced an upper bound for the gradient norms through an analysis of the spectral norms of the Jacobian matrices for the sub-layers. We then theoretically identified two conditions for avoiding loss spikes: small sub-layers and large shortcut. To meet these conditions, we show that using the widely adopted initialization method for LLMSs can make the sub-layer parameters small, and that embedding scaling or incorporating layer normalization into the embedding layer can make the standard deviation of each embed- ding close to 1, resulting in large shortcut. Experimental results indicated that methods satisfying these conditions avoid loss spikes. Furthermore, these methods allow for train- ing with relatively larger learning rate, leading to improved performance. We hope our theoretical analyses and empirical findings will help avoid wasting valuable time and com- putational budgets during LIM construction.
10

---

Published as a conference paper at COLM 2025
Ethics Statement
To stabilize the LLM pre-training, this paper provides theoretical analyses on the spectral norms of the Jacobian matrices for sub-layers to estimate the upper bound of the gradient norm of L. This paper focuses on only the stability of LLM pre-training, and thus, we have to address other issues of LLMs such hallucinations to use the LLM in a real application.
otherissues
Reproducibility Statement
We do not aim to propose a novel method in this paper, but we mainly focus on theoretical analyses the spectral norms of the Jacobian matrices to find the factor to stabilize the pre-training of LLMs. We justify our theoretical analyses through experiments with var- ious situations. To activate our modification,
add only several lines to widely used
implementation, i.e, Megatron-LM Therefore, we believe that it easy to reproduce our experimental results. Moreover, we did not modify the internal architecture of Transform- ers. In particular, Scaled Embed only scales embeddings by constant value, making convenient in terms of portability across various implementations, such as Transformers in HuggingFace"! and VLLM
However, because it is difficult conclude that our provided conditions completely solve the instability during pre- training of LLMs, is better to combine other techniques to stabi- lize the pre-traine pre-training such as an auxiliary loss described by Chowdhery et al. (2022) to make the pre-training more stable in actual pre-training situation.
References
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar Van Der Wal. Pythia: A suite for analyzing large language models across training and scaling. In Proceedings of the 40th International Conference on Machine Learning (ICML), pp. 2397-22922222 2430, 2023.
Yonatan Bisk, Rowan Zellers, Ronan bras, Jianfeng Gao, and Choi Yejin. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 7432-7439, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33 (NeurlPS), pp. 1877-1901, 2020.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Isvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Hen- ryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiri- donov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira,
Joshua
Dhttps///gthnbb.mm/////// MM
11 #httpss//hugggingfacc.o/ docs/ wwwp://tansfoommer//eeeeeeeeeee/ wwwtannswoomeemmmia///eaaaaaaaee manes//aaiiiiiiii/ index PPttps://githbb.com http://iww..omolmmlmmmmmpeeiette
11

---

Published as a conference paper at COLM 2025
Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Bren- nan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. PaLM: Scaling language modeling with pathways, 2022.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018. URL
http://rrivv..rr//////8/./ 05457.
Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Re. Flashattention; Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Informa- fion Processing Systems 35 (NeurIPS), 2022.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block- wise quantization. In Proceedings of the 10th International Conference on Learning Represen- fations (ICLR), 2022.
Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, and Jie Tang. Cogview: Mastering text-to-image generation via transformers. In Advances in Neural Information Processing Systems 34 (NeurIPS), pp. 19822-19835, 2021.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedfor- ward neural networks. In Proceedings of the Thirleenth International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 249-256, 2010.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In Proceedings of the 37th International Confer- ence on Machine Learning (ICML), pp. 4475 4483, 2020.
Peter Izsak, Moshe Berchansky, and Omer Levy. How to train BERT with an academic budget. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 10644-10652, 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020.
Teven Le Scao, Thomas Wang, Daniel Hesslow, Stas Bekman, M Saiful Bari, Stella Bider- man, Hady Flsahar, Niklas Muennighoff, Jason Phang, Ofir Press, Colin Raffel, Victor Sanh, Sheng Shen, Lintang Sutawika, Jaesung Tae, Zheng Xin Yong, Julien Launay, and Iz Beltagy. What language model to train if you have one million GPU hours? In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 765-782, 2022.
Conglong Li, Minjia Zhang, and Yuxiong He. The stability-efficiency dilemma: Investigat- ing sequence length warmup for training gpt models. In Advances in Neural Information Processing Systems 35 (NeurIPS), pp. 26736-26750, 2022.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5747-5763, 2020.
Benjamin Marie, Atsushi Fujita, and Raphael Rubino. Scientific credibility of machine translation research: A meta-evaluation of 769 papers. In Proceedings of the 59th An- nual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (ACL), pp- 7297-7306, 2021.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture Models. In Proceedings of the 5th International Conference on Learning Representa- fions (TCLR), 2017.
12

---

Published as a conference paper at COLM 2025
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 2381 2391,, 2018.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Pat- wary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), pp. 1-15, 2021.
Toan Q. Nguyen and Julian Salazar. Transformers without tears: Improving the normaliza- tion of self-attention. In Proceedings of the 16th International Conference on Spoken Language Translation (IWSLT), 2019.
Kosuke Nishida, Kyosuke Nishida, and Kuniko Saito. Initialization of large language mod- els via reparameterization to mitigate loss spikes. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 226992222222222 22699-22714, 2024.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine trans- lation. In Proceedings of the Third Conference on Machine Translation (WMT), pp. 1-9, 2018.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceed- ings of the 54th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 1525-1534, 2016.
Stephan Peitz, Sarthak Garg, Udhay Nallasamy, and Matthias Paulik. Cross+Self-Attention for transformer models, 2019.
Matt Post. call for clarity in reporting BLEU scores. In Proceedings of the Third Conference on Machine Translation (WMT), pp. 181, 2018
Alec Radford, Karthik Narasimhan, Tim Salimans, and llya Sutskever. Improving language understanding by generative pre-training. 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Lan- guage models unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learn- ing with unified text-to-text transformer. Journal of Machine Learning Research, 21(140): 1-67,2020
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory opti- mizations toward training trillion parameter models, 2020.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: an adversarial winograd schema challenge at scale. Communications of the ACM, 64(9): 99-106, 2021.
Roy Schwartz, Jesse Dodge, Noah A. Smith, and Oren Etzioni. Green AI CoRR, abs/1907.10597, 2019.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 1715-1725, 2016.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. 2019.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models us- ing model parallelism, 2020.
13

---

Published as a conference paper at COLM 2025
Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considera- tions for deep learning in NLP. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 3645-3650, 2019.
Sho Takase and Shun Kiyono. Rethinking perturbations in encoder-decoders for fast train- ing. In Proceedings of the 2021 Conference of the North American Chapler of the Association for Compulational Linguistics: Human Language Technologies (NAACL-HLT), pp. 5767-5780, 2021.
Sho Takase and Shun Kiyono. Lessons on parameter sharing across layers in transformers. In Proceedings of The Fourth Workshop on Simple and Efficient Natural Language Processing (SustaiNLP), Pp. 78-90, 2023.
Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. B2T connection: Serving stability and performance in deep transformers. In Findings of the Association for Compu- tational Linguistics: ACI. 2023, pp. 3078-3095, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schel- ten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, & ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30 (NIPS), pp. 5998-6008. 2017.
Roman Vershynin. High-Dimensional Probability; An Introduction with Applications Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge Uni- versity Press, 2018.
Ben Wang and Aran Komatsuzaki. Gpt-j-6b: A 6 billion parameter autoregressive language model, 2021.
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers to 1,000 layers, 2022.
Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Penning- ton, Jascha Sohl-dickstein, Kelvin Xu, Jaehwon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities, 2023.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Yan Liu. On layer normalization in the trans- former architecture. In Proceedings of the 37th International Conference on Machine Learning (ICML), pp. 10524-10533, 2020.
Greg Yang and Edward J. Hu. Tensor programs VI: Feature learning in infinite-width neural networks. In Proceedings of the 38th International Conference on Machine Learning (ICML), 2021.
Greg Yang, B. Simon, and Jeremy Bernstein. A spectral condition for feature learning, 2024a. URL https: httsss///axvvv.oo//aaa/e 17813.
14

---

Published as a conference paper at COLM 2025
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: Feature learn- ing in infinite depth neural networks. In The Twelfth International Conference on Learning Representations (ICLR), 2024b.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Proceedings of the 57th Amnual Meeting of the Association for Computational Linguistics (ACL), pp. 4791-4800, 2019.
Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Zhiyuan Liu, Peng Zhang, Yuxiao Dong, and Jie Tang. GLMM-300: An open bilingual pre-trained model. In The Fleventh International Conference on Learning Representations (ICLR), 2023.
Shuangfei Zhai, Iatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Josh M. Susskind. Stabilizing transformer training by pre- venting attention entropy collapse. In Proceedings of the 40th International Conference on Machine Learning (ICML), 2023.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems 32 (NeurIPS), 2019.
Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Fixup initialization: Residual learn- ing without normalization. In Proceedings of the 7th International Conference on Learning Representations (ICLR), 2019.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
A Hyper-Parameters
Table 4 shows that hyper-parameters used in experiments LLMs. To make the periments on 13B parameter models close to a realistic situation, we increased the batch size and the number of updates in comparison to ones in experiments on 350M and 1.7B parameter models. In additon,, for Adam 62, most studies have used 0.95 to stabilize their pre-trainings (Brown al., 2020; Zhang et al., 2022; Zeng et al., 2023; Biderman et al., 2023; Touvron et al., 2023), and thus, we also used 0.95 in the experiments on 13B parameter models. As described in Section 5.4, we emphasize that the pre-training of Vanilla is es- sentially unstable even if we use the widely used Adam B2 value, 0.95, which is known as the technique to stabilize the pre-training. Therefore, our provided conditions are also effective for the stabilization in the realistic situation.
B Discussions on Other Configurations
In this section, we conduct experiments on other configurations to describe connections with previous study reports.
B.1 Methods without Small Sub-Layers
Since we applied the widely used initialization method for LLMs in experiments in Section 5, all methods satisfy the condition on the small sub-layers. In this section, we empirically investigate the property of the method that violates the condition. We compare the Trans- former initialized by the Xavier initialization (Glorot & Bengio, 2010) (Xavier Init), and the combination of Xavier Init and Scaled Embed (Xavier Init + Scaled Embed) with Scaled Embed. As described in Table 5, Xavier Init violates both conditions, and Xavier Init + Scaled Embed satisfies the only large shortcut condition. In the same manner as in Section
15

---

Published as a conference paper at COLM 2025
Table 4: Hyper-parameters used in our experiments on the LLM pre-training.
Name
350M
1.7B
13B
Layer num
24
24
40
Hidden dim size
1024
2304
5120
FFN dim size
4096
9216
20480
Attention heads
16
24
40
Dropout rate
0.1
0.1
0.1
Precision
float16
float16
float16
Sequence length
2048
2048
2048
Batch size
528
528
1024
The number of updates
35000
35000
50000
Adam P1
0.9
0.9
0.9
Adam $2
0.999
0.999
0.95
Gradient clipping
1.0
1.0
1.0
Ir decay style
cosine
cosine
cosine
Ir warmup fraction
0.05
0.05
0.05
Weight decay
0.01
0.01
0.01
Method
Small sub-layers Large shortcut
Xavier Init
Vanilla
Embed Detach
Embed LN
Scaled Embed
Xavier Init + Scaled Embed
Table 5: Relations between each method in experiments and two conditions to control the upper bound of gradient norms.
Model
WikiText | LAMBADA І
350M parameters
Xavier Init
33.92
34.72
Xavier Init Scaled Embed
30.50
26.55
Scaled Embed
29.86
24.37
1.7B parameters
Xavier Init
30.10
29.29
Xavier Init Scaled Embed
23.16
15.49
Scaled Embed
21.29
12.53
Table 6: Perplexities of each method.
we trained models of 350M and 1.7B parameters with Ir 5.0 10-4, We also used the hyper-parameters described in Table 4.
Figure 6 shows loss curves validation data for 350M and 1.7B parameters in each method, and Figure 7 shows their gradient norms. These figures show that Xavier Init and Xavier Init + Scaled Embed faced loss and gradient spikes. In particular, the spikes appeared more frequently in Xavier Init, which violates both conditions, in comparison with Xavier Init + Scaled Embed. contrast, Scaled Embed, which satisfies both conditions, avoided the gradient spike and prevented the loss spike problem. These results indicate that we have to satisfy both conditions: small sub-layers and large shortcut to prevent the loss spike problem. Moreover, Table 6 shows perplexities of each configuration in evaluation data. This table indicates that Scaled Embed achieved better performance than Xavier Init and Xavier Init Scaled Embed that faced some spikes.
16

---

Published as a conference paper at COLM 2025
4.0
Xwariist
Novder
38-
Xoarrist- Scated Enhrdd
38-
Xosder Sealed Ermied
3.-
Scoled Enrodd
36-
Scaled Emroed
5.4-
26-
26-
10000
20000
25000
35001
10000
19830
25:00
35/00
Troin stop
Trohn 500p
(a) 350M parameters.
(b) 1.7B parameters.
Figure 6: Loss curves of each method in validation data for the comparison to methods without small sub-layers.
10-
Grat
Gost
2000
10002
15600
20000
25000
20392
1950
10500
150000
20000
25000
35532
Train step
Trah 10ep
(a) 350M parameters.
(b) 1.7B parameters.
Figure 7: Gradient norms of each method during the training for the comparison to meth- ods without small sub-layers.
Vonil
Van
Horda Scaved Uibbe
375-
Suale Embee
7-5-
Frebed
3.5
3.50
3.00
332-
325-
100
275-
2.00
subec 3500
230
15000
25000 33000 soice
2:00
22500
Train mmp
Tmanseep stop
Transup
(a) 350M with Ir = 1.0 10-3 3
(b) 350M with Ir = 5.0 x 10m4,
(c) 350M with
1.0 x 10
4.00
Vandal
355
350
a.53-
300
300
300
250
5950 19000 1000e 25003 25000 20dec 35300
2:0
0
5903 12000 15000 29200
28000 32000 20000
230
5000
10660
19230 20006 25000 20002 1900 Trai-stpp
Traisstop
Tmmh 196p
(d) 1.7B with Ir 1.0 10-3
(e) 1.7B with Ir 5.0 10
(f) 1.7B with lr 1.0 x 10
Figure 8: Loss values of each method with 350M and 1.7B parameters when we vary learning rate.
B.2 Varying Learning Rate
Le Scao al. (2022) reported that the stable method, such as Embed LN, was worse than Vanilla. However, in Section 5.3, the stable methods, Scaled Embed and Embed LN, achieved better performance than Vanilla in the 1.7B parameter configuration. We suppose that the difference in the learning rate causes this gap in findings. In this section, we tried to train Vanilla and Scaled Embed with larger and smaller learning rates: Ir 1.0 x 10-3 and 1.0 10 respectively.
17

---

Published as a conference paper at COLM 2025
WikiText Ir 5.0 x 10
LAMBADA
Model
Ir 1.0 x 10
Ir 1.0 X 10
Ir 1.0 10 5.0 10
Ir 10
350M parameters
Vanilla
29.96 28.09
30.35
34.51
25.12 22.03
24.73
32,49
Scaled Embed
29.86
35.66
24.37
37.14
1.7B parameters
Vanilla
N/A
22.58
23.54
N/A
15.22
16.17
Scaled Embed
20.95
21.29
23.78
12.26
12.53
15.39
Table 7: Perplexities of each method with 350M and 1.7B parameters when we vary learn- ing rate.
Figure 8 shows loss values of each configuration in validation data. As shown in this figure, the larger the learning rate we used, the more frequent the spikes occurred in Vanilla. In particular, in Ir 1.0 x 10 the training of Vanilla with 1.7B parameters failed because its gradient exploded. In contrast, Scaled Embed stabilized the training, and thus, its loss values consistently decreased.
Table 7 shows the perplexities of each configuration in evaluation data. This table indi-
cates that Vanilla with 350M parameters achieved better performance in Ir = 1.0 x 10 that is the situation where its training did not face any spike. This result corresponds to the report of Le Scao et al. (2022). Thus, we suppose that they conducted the comparison with too-small learning rate to stabilize Vanilla. In contrast, the stable methods are more effective in training with a large learning rate, shown in Figure 8 and Table 7. Therefore, Le Scao et al. (2022) used a relatively large learning rate in their experiments, their stable method could achieve better performance.
B.3 Varying Sequence Length
Li et al. (2022) indicated that it is better to train with a short sequence at the early stage to stabilize the LLM pre-training. They justified their method based on the curriculum learning strategy. On the other hand, in this section, we provide the theoretical to their method in terms of the standard deviation of the shortcut.
jusooeticalusiifficatiio justification
As described in Section 2.1, Transformers add the output of each sub-layer to the shortcut. Since the standard deviation of the self-attention layer tends to decrease with the length of an input se-
4.00
L-128
L-258
L-512
quence especially at the early s a long se- quence tends to keep the standard deviation of the shortcut small. Therefore, the long sequence makes the pre-training of Vanilla more unstable.
3.50
L=1024
L-2048
Valid
3.25
3.00
2.75
We conducted experiments with varying the length of the input sequence from 128 2048. To use the same number of tokens to update pa- rameters, we adjusted the batch size. Figure shows loss values of Vanilla with each L config- uration in the validation data. This figure shows that spikes occurred only in the large L, i.e., 1024 and 2048. Moreover, the spikes are more likely to occur at the early stage of the pre-training. There- fore, using a short sequence stabilizes the training at the early stage, reported in Li et al. (2022).
5000
10000
16000 20000
25000 30000 35000
Train
Figure 9: Loss curves of Vanilla with 350M parameters in validation data when we vary the input sequence length. We adjust the batch size to use the same number of tokens for the training of each model.
13see Appendix F for details.
18

---

Published as a conference paper at COLM 2025
4.00
Venile
RMSNorm
3.50
Graid norm
2.0
Vale
3.25
1.5
1,0
277
0.54
10000
15000 2000 26000
30000
36000
0
5000
10000
15000
25000 30000 35000 Train
20000
(a) Loss curves validation data.
(b) Gradient norms.
Figure 10: Loss values and gradient norms of Vanilla and RMSNorm.
B.4 RMSNorm
Some recent LLMs use the RMSNorm (Zhang & Sennrich, 2019) instead of the LN in their Transformers (Touvron al., 2023). We discuss such an architecture this section. the same as LN discussed in Section 3.1, we can obtain the Jacobian matrix of the RMSNorm
with the following equation:
aRMSNorm(x) dx
(22)
Ox
Thus, the upper bound the gradient norm the same in LN if we use RMSNorm.
Figure 10 shows the loss values and gradient norms of the Vanilla configuration in Sec- tion 5.3 and the one using RMSNorms instead LNs ("RMSNorm" in figures) with 350M parameters. We trained them with = 5.0 x 10-4 as in Section 5.3 As shown these fig- ures, RMSNorm faced loss and gradient spikes in similar location to the ones of Vanilla. These empirical results also indicate that the RMSNorms have the same problem as LNs regarding the instability.
C Scaling Embeddings with Larger Value
Equations (15) and (20) indicate that we can stabilize the LLM pre-training by adjusting the standard deviation of the shortcut to large value. In fact, our experimental results show that we can stabilize the LLM pre-training by making the standard deviation of each embedding close 1. To investigate how about larger value, we conducted experiments with making the standard deviation of each embedding close to 5 and 50.
Figure 11 shows loss curves in validation data for 350M and 1.7B parameters in each situ- ation. This figure indicates that although all settings prevented the loss spike problem, the larger standard deviation than 1 degraded the performance. Therefore, it is unnecessary to scale the standard deviation of each embedding with larger value than to prevent the performance degradation.
D Distributions of Sub-Layer Outputs
Figures 12 and 13 shows output distributions of each sub-layer at the beginning of pre- training and after pre-training. These figures indicate that each sub-layer output is close to normal distribution in various configurations both at the beginning of pre-training and after pre-training. Therefore, the assumption in this study, which is that the vector x at each layer follows the normal distribution, is reasonable.
1e tried train them with Ir 1.0 10-3 but RMSNorm exploded.
19

---

Published as a conference paper at COLM 2025
3.75
Scaked Emied
3.75
Scaled Embed
3.50
tiee Scaled Embed
3.50
Scaled Ecaledlmbbed Embed
times Staled
50 times Scalled Embed
3.25
0.25
Vald
3.00
3.00
2.75
2.75
2.50
5000
10000
15000 20000 25000 30000 36000
0
5000
10000
15000 20000
25000 30000 36000 Train
Train stop
(a) 350M parameters.
(b) 1.7B parameters.
Figure 11: Loss curves in validation data when we scale the standard deviation of embed- dings with larger than 1.
(a) self-attention of 350M param- (b) self-attention of 1.7B parame- (c) self-attention of 13B parame- eters.
ters.
ters.
02
da
-11
02
-22 -o
dh
65
(d) FFN of 350M parameters.
(c) FFN of 1.7B parameters.
FFN of 13B parameters.
Figure 12: Output distributions of each sub-layer at the beginning of pre-training.
E Discussion on Various Activation Functions in FFN
E.1 ReLU
We consider the case where we use the ReLU function as F instead of the identity function. Because we assume that parameters and the input vector at each layer follow the normal distribution, the internal layer also follows the normal distribution. Therefore, each ele- ment of the FFN internal layer is negative value with half probability. In this case, can regard that the ReLU function cuts the elements of the internal layer by half. Thus,
xd and W2 E Rdx 4tp in the replace W1 E Rdm xd and W2 E Iddxd with W1 IR fif
discussion Section 3.1 when we use the ReLU function as F.
E.2 SiLU
We consider the case where we use the SiLU function as F. The definition of the SiLU function is as follows:
SiLU(x) x Sigmoid(xx
(23)
20

---

Published as a conference paper at COLM 2025
10
Qvreit
-100
-223
-050
-028
800
022
022
075
100
470
-00
050
0.5
1.00
-100
-073
-000
4025
666
02
0t
6n
1.0
(a) sell-attention of 350M param- (b) self-attention of 1.7B parame- (c) self-attention of 13B parame- ters.
eters.
ters.
Dereay
Dordky
422
050
-000
6.79
9150 1073 -0.5
(d) FFN of 350M parameters.
(e) FFN of 1.7B parameters.
(f) FFN of 13B parameters.
Figure 13: Output distributions of each sub-layer after pre-training,
where Sigmoid is the sigmoid function.
aSilU((x)
Sigmoid(x) + aSigmoid
dx
dx
Sigmoid (x) Sigmoid (x) o (1 Sigmoid(x))
Sigmoid(x) o (1 + x (1 Sigmoid(x))
(24)
Let D = diag
aS aW,x
E d xdmn, Then, we obtain the Jacobian of the FFN as follows:
OFFN(LN(X')) W2DW1 aLN(x")
(25)
Therefore,
)FFN(LN(x*))
< [Walaa|Plla/llll
(26)
aLN(x')
The spectral norm of the diagonal matrix D is the maximum absolute value of its diagonal elements:
dSiLU((;;
|DI| max
(27)
dx;
Moreover, we find that its maximum occurs at x R 2.4 and is approximately 1.1, and thus, ||D||2 < 1.1. Because ||lll (v(V vdm) and ||Wallz 02(Vd vdun), we can express the upper bound as the following inequality:
)FFN(LN((x)) aLN(x')
< 1.1100002(d++ )2
(28)
2
Finally, Equation (9) can be rewritten as:
dy dx
0102
$ + 1.1
Ciin
(29)
2
21

---

Published as a conference paper at COLM 2025
E.3 SwiGLU
We consider the case where we use the SwiGLU function as F. When we use the SwiGLU function, the FFN layer expressed follows:
FFN(x) (1(Swish(Wix) (Vx))
(30)
where V E R xd and follows a normal distribution N (0, Jacobian of the (xN as follows:
Then, we compute the
aFFN(x) (F) aW1x dx
aFFN aVx
+
(31)
dx
dwix
dVx
dx
(FF
@Swish(W1x)
W2
diag(Vxx Vx) diag
(32)
dw
dW1x
(FFN(x) W2 ((((lla((Swis((Wx)))
(33)
aVx
Therefore, we can rewrite Equation (31) as follows:
dFFN(x) = W2 (tiag(Vx) diag
aSwish (Wixs))(W)))) )) W4 + (W(diag((Switt(W;)))) (34)
dx
aw
)) Wi and J2 Wu(ddag(Swish(W::))))
Let J1 W2 (diag(x)) diag
@Swish(Wix)
dW1x
aFFN(x) dx
/1 + 1/2. We can derive the upper bound of aFFN(x) dx
|| as follows:
(FFN
< |1ll2 + 12/222
(35)
dx
For [1112, we have:
@Swish (Wix) |Iill2 aW,
) ||
1hll - [Walla||daag(
diag
(36)
Each element of Vx is sum of d independent random variables with variance oz06 and thus, var(Vx) dofas. Therefore, from the expected maximum of d Gaussian random variables,
Idiag (xiag(Vx))l22 S 0,0V /2 log dim
(37)
The derivation of the Swish function is bounded by 1.1 in the same manner as the SiLU function:
(S1ish(W1x) aW,x
< 1.1
(38)
2
The spectral norms of W1 and W2 can be obtained |[Wll z (v vdmm and ||2ll2 z 02( (d(vd + vdin) as described in Section 3.1. We can obtain the upper bound of |// ll with these equations:
11 < (V Vdm)) 22222/2222 22d log dim
(39)
For [2]2|,, we have:
|| 12ll < [MMila/ldingg(SSsss((())))))))))
(40)
Due Swwsh(W,x)) W, and var(W1x) do we can obtain:
[[ddiag(Siish(WW:)))3355555 0 /2 d log d
(41)
22

---

Published as a conference paper at COLM 2025
Therefore,
1122222 & 01ovo102(VB+ + vdm) 2 d log dim
(42)
Based these equations, we derive the upper bound follows:
(x)
- |1 + 11/2|12
dx
= 2.1 02000102(V / d log dim
0=0v0102Cuwiguu
(43)
where Cswiglu includes 2.1 (vd + vdtn) /2 d log dim for the simplification.
Finally, we consider the actual Transformer layer that includes layer normalization and residual connection:
dy 0x
000102 Cwigll
1
(44)
0x'
We note that a in Equation (43) is equal to in the actual Transformer layer because we apply the layer normalization to the input the FFN layer.
Relation between Input Length and the Standard Deviation of Self-Attention
We explain that the standard deviation of the self-attention layer becomes small as the input length is long. Because we assume that parameters and the input vector at each layer follow the normal distribution, the expectation of each element of (Wo: x)T(Wx is O. Therefore, the expectation after the softmax function is 1 where L is the length of the input sequence. Thus, the long input sequence decreases the standard deviation of the self-attention layer.
(Wk: X)
To simplify, we consider the case where the number of self-attention heads is 1. this case, can obtain the variance of each calculation with the following equation.
var (o(x)) var(Wo)var(x) d var(Wy(x)) var(Wy)varrx))
(45)
(46)
where var represents the variance of the matrix/vector. Thus, the variance of the self- attention layer, var(Attn(x),, is as follows:
var(W var(Wy)var(xx d
var(Attn(x)) var(Wo) Σ
(47)
1,2
(Wa var(x)d
(48)
G Details on Jacobian Matrix of Self-Attention
We can represent concat(head, head, with the summation of the matrix multi- plications follows:
h (ee. = Σ Wea
(49)
where W; E Rhehesd xd whose corresponding element is 1 and the others are Let J be the Jacobian of the (x),,(x), we can represent 1 in Section 3.2 as follows:
Σ J -=1
(50)
23

---

Published as a conference paper at COLM 2025
In addition, the self-attention consists of the interaction among inputs and outputs of each position in the sequence. Thus, we add indices to Jacobians to represent the positions of the input and output. Let xj be the input of the position j, and z be the i-th head of the output position k, and Jí ox, the position j, we can convert Equation (51) into the following Equation:
a= Because 1 can be regarded as the Jacobian of the input of
h
1 = k=1-
(51)
where is the length of the input and output sequences. Therefore, we compute R to obtain 12 in Section 3.2.
We can obtain a head of the output position k, i.e., Zg, as follows
Zk
(52)
(Wo xx) x) V
where Au is the I-th element of the attention vector, softmax Wyx. Therefore, to obtain Jaj. we differentiate 2ạ with respect to xj as:
and By is
azk Σ dxj
aAk dxj
dvj Ari ox,
Jjj
01
(53)
1=1
dul
Wydij
(54)
dxj
( ifl ifl=j
813 10 otherwise
(55)
Thus,
azz
1.
aAu 1i + Aw;W ax,
Σ
(56)
dxj
Here, we assume that the attention vector uniform. In this assumption, since Axj = 1, we can obtain the spectral norm of the second term for Equation (56) as Av = " +
/ / where dy is the standard deviation of Wy. To calculate the first term, we compute aAu
dx,
aAg
aSyi dx,
aSkm
Axi
Akm
(57)
dx,
dxj
m=
aSg
1
(W] + WFWox
(58)
dxj
ddnead
Let Dj be Wxxxx and Eg be WF Woxx. Then,
aSg
(Drōey + Ex51T)
(59)
dx,
V dhead
15 simplify the equations, omit the index i to represent i-th head from the head and parame-
ters.
24

---

Published as a conference paper at COLM 2025
Therefore,
aAu
Agi
(Diotz+ + Exbii))
Σ
Akm (Dmorf + Exōmj)
(60)
dxj
v hdhhad
m1
Agt
Dr
L
b1,Exx AxiEx + dzj
Σ
Akm Dm
(61)
v Váhead
m==1
We assign Equation (61) to the first term of Equation (56) and use the assumption Azj
aA
_
Σ
Agt
51 AcjEx + tkk
Di
Σ
Akm
Dm
(62)
dxj
v hnead
m==
01, 7 + Ezj
Di
Σ
{
Dm
11
(63)
Vndaa
L
1
(bii DeyDi) 1i")
Σ
1
1
Ex
Σ m==
Dma
D1 T
LVdnead
I.
I.
1=1
(64)
L Ex T + Σ
(0ej (be)Droi)) - Ex +
Σ
Dm
Di
(65)
Lydnhaad
1=1
m=1
E T - EgDI" T + dzj Σ 4
(Drii))
Σ M==1
Dm
1j
(66)
Lydhhad
Based on the assumption that parameters and the vector each layer follow the normal distribution, we assume that the mean of D,, Es, and Dy is 0, and thus, we can obtain their norms from their variances. In addition, we assume that the standard deviation of Wo, Wg
and Wv is Then, var(D)) var(Ex) = dddeadot var(Ei Dj) Lddhead and var(o)) =
do2. Thus, ||Ex0]112 || @d || (D, ll2 I| Dm Dm)ofll2
03, LdPahhead Therefore,
AAk dxi
2
1
1k, T
+ ||,,, T| + +0x]
Σ (Dioi
+
[1 Dm 0/
1)
(67)
LV Lhend
1=1
1
203 V d3d2 head + (ōzj + 1)02 v Lachhaad
(68)
LVdneadd
2
dej
+1+
dddhaad
(69)
VI.
VI.
Based on this Equation, we can compute the upper bound of ||g]]|2
dAu T
+ Ax Wy
1aill2
21
dx;
dAu 21
Σ
+ -|AxyyWilk
(70)
dxi
2
1
dej
+ 1 +
2 vT.
03 v /d /dddhead + - [(Vd+ v /dhead)
(71)
25

---

Published as a conference paper at COLM 2025
Moreover, we compute the upper bound of || from Equation (51) follows:
11 2
Fy wi
(72)
M Will2
(73)
1 (1 (1+L+
2L
Lo
0² dd ddeed + I (V + v (dhead)
(74)
VI.
VI.
<v((L+2+
03 dd ddead + +olVd vdhead)
(75)
vL
H
Comparison with Post-LN Transformer
As described in Section 2.1, recent studies use the Pre-LN Transformer architecture to con- struct their LLMs because the architecture more stable. In contrast, some recent studies reported that the Post-LLN Transformer, which is the original architecture, can achieve bet- performance than the Pre-LN if we address the instability issue in the Post-LN, i.e. the vanishing gradient problem (Liu et al., 2020; Takase et al., 2023; Wang et al., 2022). We dis- cuss whether the Pre-LN Transformer entirely underperforms the Post-LN. We conducted experiments machine translation experiments because previous studies mainly focused on them.
We followed the experimental settings in Takase et al. (2023). Table 8 shows the details of hyper-parameters. We used the WMT English-to-Genn -errman training dataset (Vaswani et al., 2017; Ott et al., 2018), and evaluated each model in newstest2010-2016. We used the encoder-decoder architecture proposed by Peitz et al. (2019). To stabilize the Post-LN Transformer, we applied DeepNet (Wang et al., 2022) and B2T connection (Takase et al., 2023). We compared them to Scaled Embed, that is, the Pre-LN Transformer with the stabi- lizing techniques described in this paper.
For used SacreBLEU (Post, 2018) to obtain The signature of SacreBLEU is BLEU+case, mi mixetnnumssfss 1+smooth. 1+smooth.exp-tok. 13atversion. we used two learning rates: Ir = 1.0 x 10-3 and 3.0 x 10-3 For Ir = 1.0 x 10-3, DeepNett and B2T connection outperformed Scaled Embed. Thus, the Post-LN Transformer-based methods achieved better performance than the Pre-IN Transformer-basee method. This result corresponds to reports in previous studies (Liu et al., 2020; Wang et al., 2022; Takase et al., 2023).
Table 9 shows the averaged BLEU scores among newstest2010-2016 the BLEU score calculation, compatible scores (Marie et al., 2021).
As shown in this table,
On the other hand, for = 3.0 x 10-3, Scaled Embed achieved better performance than the others with Ir = 1.0 x 10-3, and the training of the others failed due to the exploding gradients. This result indicates that the Pre-LN Transformer-based method can achieve better performance we use a large learning rate. Therefore, the Pre-LN Transformer (with the stabilizing techniques) is more stable than the Post-LN Transformer-based method, and thus, can achieve better performance when we use a large learning rate that is too large to train the Post-LN Transformers.
26

---

Published as a conference paper at COLM 2025
Table 8: Hyper-parameters used in the comparison with Post-LN Transformer.
Name
Value
Layer num
18
Hidden dim size
512
FFN dim size
2048
Attention heads
Dropout rate
0.5
Precision
float16
Word dropout rate
0.1
Max tokens
7168
Adam Bi
0.9
Adam B2
0.98
Gradient clipping
0.1
Ir decay style
inverse square root
Warmup step
4000
Weight decay
Table Averaged BLEU scores among newstest2010-2016
Model
2010
2011
2012
2013
2014
2015
2016
Average 1
Ir
1.0 10
DeepNet
24.65
22.30
22.87
26.51
27.29
29.77
34.87
26.89
B2T connection
24.46
22.42
22.85
26.51
27.46
29.91
34.65
26.89
Scaled Embed
24.32
22.21
22.40
26.38
26.89
29.98
34.53
26.67
Ir
3.0 x 10
DeepNet
N/A A
N/A
B2T connection
N/A A
N/A A
Scaled Embed
24.52
22.23
22.86
26.54
27.35
29.90
35.16
27


Reviews:
[{"review_id": "wcsYzQylLw", "reviewer": "Reviewer_4pzp", "summary": "This paper tries to understand how to prevent spiking in language model pretraining and explains why two known techniques: Embedding Normalization and Layer-wise parameter initialization scaling, help stablizing the model in training.\n\nThe author proposes to control the upper bound of gradients at initialization because gradient spiking is commonly observed before loss spiking.  They further show that the upper bound of gradient-norm on each layer positively correlates with the weight norm on each layer and negatively correlates with the embedding weight norm. Therefore, the authors propose to make sure the embedding is of norm O(1) and layers has weight norm O(1/L) or O(1/\\sqrt{L}).", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "1. The authors present a very detailed analysis considering the gradient upper bound and perform empirical experiments across different scales to support it.\n\n2. The paper is overall clear and well-written.", "reasons_to_reject": "1. The empirical result seems to have some potential contradiction with the theory. The gradient norm seems to have a lot of spike in Figure 4, it is just that the loss is not spiking. \n\n\n2. While the theoretical analysis is very detailed, the authors didn't argue how this differs from famous previous works: for example, the MuP paper (and its depth variant) [1, 2], which has similar calculations and arrives at a very similar initialization paradigm. \n\n[1] https://arxiv.org/abs/2310.17813\n[2] https://arxiv.org/abs/2310.02244", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JYoMA5tU5k", "reviewer": "Reviewer_2gZh", "summary": "This paper addresses the issue of loss spikes that commonly occur during the pre-training of LLMs. The paper focuses on loss spikes during LLM pre-training, hypothesizing that they are caused by sudden growth in gradient norms. The authors theoretically analyze the spectral norms of the Jacobian matrices for sub-layers in Transformer-based LLMs to identify factors that keep gradient norms small. \nThe paper conducts various experiments to empirically validate its theoretical findings. \nResults show that methods satisfying these conditions effectively prevent loss spikes during pre-training.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "1. Provides theoretical insights into the cause of loss spikes and conditions for stable training.\n2. Offers theoretical backing for empirical methods already used in the field.", "reasons_to_reject": "1. It's unclear why preventing these spikes is crucial.\n2. Some theoretical analyses involve simplifying assumptions, such as the identity function for activation functions in FFNs, even though the appendix explores other activation functions.\n3. Although improving stability, the proposed methods might not always outperform other techniques in terms of final performance, especially at smaller learning rates.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W6QP2ZbAri", "reviewer": "Reviewer_ePfy", "summary": "This paper investigates the factors that maintain a small gradient norm for LLM pretraining by analyzing the spectral norms of Jacobian matrices across sub-layers, under the assumption that loss spikes are caused by a sudden increase in gradient norm. The authors find that stabilizing the pre-training process requires two key conditions: using smaller sub-layers and incorporating larger shortcut connections. They conduct a range of experiments to empirically validate their theoretical insights. Results show that methods satisfying these conditions effectively prevent loss spikes during pre-training.", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "Understanding and addressing loss spiking is a highly important and timely topic, particularly for stabilizing LLM pre-training.", "reasons_to_reject": "1. The proposed solutions and findings are not entirely convincing. For instance, one of the suggested strategies: initializing sub-layer parameters with smaller values, oversimplifies the problem. Prior work has shown that loss spiking results from the interaction of training hyperparameters, data, and model state. In some cases, spiking occurs even when the data and configuration seem correct. Simply using smaller initialization values may not generalize. Moreover, the identification of sub-layers is highly data-dependent, and different data shuffling orders can yield different sub-layer behaviors. Also, since loss spikes typically occur after the warm-up phase, it's unclear whether smaller initializations have any effect at that point. These issues raise serious doubts about the reliability of both the method and the conclusions in this paper.\n2. The experimental setup is insufficient and seems incorrect. The authors rely primarily on the C4 dataset and perform most experiments on 350M and 1.7B models, with only limited testing on a 13B model. However, as noted in the LLaMA paper, significant loss spikes tend to appear at larger scales (e.g., 33B and 65B). For 13B models, such spiking is generally mild. This brings serious doubt on whether the experimental setup in this paper is accurately configured and reflects the loss spiking phenomena that the paper aims to study. As a result, the conclusions in this paper are not only unconvincing but also potentially misleading.", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper analyzes the problem of loss spiking during training LLMs and provide some conditions to bound the gradient norm of the models, connecting to remedies such as scaling the embeddings or applying layer norm that satisfy these conditions. The reviewers appreciated the theoretical analysis and experiments on various scales up to 13B. The discussed remedies have similarities to muP, as discussed by the reviewers. 
Reviewer epfy was concerned that the paper reports spikes on small models (350M, 1.7B) but these are typically stable during training, and attributes this to possibly suboptimal hyperparameter tuning. The results on 13B show that reducing the learning rate also removes the loss spiking phenomenon. The type of spikes in smaller models, which doesn't lead to overall divergence of the training run, seem different from the type in 13B with a larger learning rate. Regardless, the suggested approach seems to resolve both types of spikes and results in similar performance of the final model, and could make the setting of hyperparameters for stable training more robust in training LMs. More extensive experiments on the "spiking regions" of hyperparameters for vanilla vs scaled embed could help clarify how much more robust to hyperparameter settings this is.
Decision: Accept