Spike No More: Stabilizing the Pre-training of Large Language Models

## Abstract

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.

## 1 Introduction

Figure: Figure 1: Training loss values of Transformers, whose dimensions and the number of layers are the same as the 1.7 billion parameters configuration in Narayanan et al. (). In Vanilla, some spikes occur at the beginning of the training, and its loss value exploded at about 13000 steps.
Refer to caption: x1.png

Large language models (LLMs) have been fundamental assets for various applications (Brown et al., ; Chowdhery et al., ; Touvron et al., ).
Increasing the number of parameters in (neural) language models and the number of training data usually leads to better LLMs (Kaplan et al., ).
Consequently, pre-training requires a vast budget, and thus, minimizing the risk of failure of the pre-training is a paramount concern.

Despite their widespread use as the foundational architecture for LLMs, a comprehensive theoretical understanding of Transformers (Vaswani et al., ) has not yet been achieved.
One of the crucial unresolved questions is the reason for the frequent occurrence of pre-training failures in Transformer-based LLMs due to spikes in loss values (loss spike) that can lead to catastrophic divergence (Chowdhery et al., ) as illustrated in Vanilla in Figure 1.
While several empirical strategies have been proposed to mitigate this problem (Chowdhery et al., ; Le Scao et al., ; Zeng et al., ), the absence of theoretical justification for these methods casts unclear on their generalizability to other situations, such as varying sizes of model parameters.

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 we assume that 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., ), 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 a 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., ; Le Scao et al., ; Biderman et al., ).
On the other hand, the latter condition was satisfied in the original Transformer by scaling embeddings (Vaswani et al., ), 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 satisfying the conditions avoid the occurrence of loss and gradient spikes.
In contrast, methods that fail to meet these conditions remain susceptible to gradient spikes, even when previously recommended as empirical solutions of the loss spike problem.
Furthermore, we demonstrate that a 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 GPT series (Radford et al., ; ; Brown et al., ).
They use the Pre-LN Transformer (Xiong et al., ), which is 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., ; Liu et al., ; Takase et al., ).
Let $x\in\mathbb{R}^{d}$ be an input of a layer of the Transformer, where $d$ denotes the dimension of the layer.
The layer outputs $y$ with the following equations:

$$ $\displaystyle y$ $\displaystyle=x^{\prime}+\mathrm{FFN}(\mathrm{LN}(x^{\prime})),$ (1) $\displaystyle x^{\prime}$ $\displaystyle=x+\mathrm{Attn}(\mathrm{LN}(x)),$ (2) $$

where $\mathrm{LN}$ is the layer normalization function(^1^11We discuss the difference from the architecture using Root Mean Square layer normalization (RMSNorm) (Zhang & Sennrich, ) instead of LN in Appendix D, and the original Transformer architecture, i.e., Post-LN Transformer in Appendix J.).
We call the first terms in Equations (1) and (2), i.e., $x$ and $x^{\prime}$, shortcut.
In addition, the feed-forward network ($\mathrm{FFN}$) and multi-head self-attention ($\mathrm{Attn}$) are defined as follows(^2^22To simplify equations, we omit bias terms.):

$$ $\displaystyle\mathrm{FFN}(x)$ $\displaystyle=W_{2}(\mathcal{F}(W_{1}\ x)),$ (3) $\displaystyle\mathrm{Attn}(x)$ $\displaystyle=W_{O}(\mathrm{concat}(\mathrm{head}_{1}(x),...,\mathrm{head}_{h} (x))),$ (4) $\displaystyle\mathrm{head}_{i}(x)$ $\displaystyle=\mathrm{softmax}\left(\frac{(W_{Qi}\ x)^{\mathrm{T}}(W_{Ki}\ X)} {\sqrt{d_{\mathrm{head}}}}\right)(W_{Vi}\ X)^{\mathrm{T}},$ (5) $$

where $\mathcal{F}$ is an activation function, $\mathrm{concat}$ concatenates input vectors, $\mathrm{softmax}$ applies the softmax function to an input vector, and $W_{1}\in\mathbb{R}^{d_{\mathrm{ffn}}\times d}$, $W_{2}\in\mathbb{R}^{d\times d_{\mathrm{ffn}}}$, $W_{Qi}\in\mathbb{R}^{d_{\mathrm{head}}\times d}$, $W_{Ki}\in\mathbb{R}^{d_{\mathrm{head}}\times d}$, $W_{Vi}\in\mathbb{R}^{d_{\mathrm{head}}\times d}$, and $W_{O}\in\mathbb{R}^{d\times d}$ are parameter matrices, and $d_{\mathrm{ffn}}$ and $d_{\mathrm{head}}$ 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\in\mathbb{R}^{d\times L}$, where $L$ is the input sequence length, to compute the self-attention.

### 2.2 Gradients of Pre-LN Transformers

Let $\mathcal{L}$ be the loss function of the $N$ layered Pre-LN Transformer and $J_{n}$ be the Jacobian matrix of the $n$-th layer.
We can calculate the gradient of $\mathcal{L}$ using the relations in Equations (1) and (2) as:

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial x_{1}}=\frac{\partial\mathcal {L}}{\partial y_{N}}\prod_{n=1}^{N-1}J_{n}=\frac{\partial\mathcal{L}}{\partial y _{N}}\prod_{n=1}^{N-1}\left(\frac{\partial y_{n}}{\partial x^{\prime}_{n}} \frac{\partial x^{\prime}_{n}}{\partial x_{n}}\right),\quad\mbox{where}\quad J _{n}=\frac{\partial y_{n}}{\partial x_{n}}$ $\displaystyle=\frac{\partial y_{n}}{\partial x^{\prime}_{n}}\frac{\partial x^{ \prime}_{n}}{\partial x_{n}}.$ (6) $$

Using the submultiplicativity of the spectral norm, i.e., $\|AB\|_{2}\leq\|A\|_{2}\|B\|_{2}$, and Equation (6), we can derive an upper bound of the norm of the gradient of $\mathcal{L}$ as:

$$ $\displaystyle\bigg{\|}\frac{\partial\mathcal{L}}{\partial x_{1}}\bigg{\|}_{2}= \bigg{\|}\frac{\partial\mathcal{L}}{\partial y_{N}}\prod_{n=1}^{N-1}\frac{ \partial y_{n}}{\partial x^{\prime}_{n}}\frac{\partial x^{\prime}_{n}}{ \partial x_{n}}\bigg{\|}_{2}$ $\displaystyle\leq\bigg{\|}\frac{\partial\mathcal{L}}{\partial y_{N}}\bigg{\|}_ {2}\prod_{n=1}^{N-1}\bigg{\|}\frac{\partial y_{n}}{\partial x^{\prime}_{n}} \bigg{\|}_{2}\bigg{\|}\frac{\partial x^{\prime}_{n}}{\partial x_{n}}\bigg{\|}_ {2}.$ (7) $$

Thus, we can estimate the upper bound of the gradient norm of $\mathcal{L}$ by analyzing the spectral norms of the Jacobian matrices for the FFN layer and the self-attention layer, namely, $\|\frac{\partial y_{n}}{\partial x^{\prime}_{n}}\|_{2}$ and $\|\frac{\partial x^{\prime}_{n}}{\partial x_{n}}\|_{2}$.

### 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^{\prime}$ 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^{\prime}$, and $W$ for all layers follow a normal distribution with a mean of 0, i.e., $\mu=0$.

This assumption is valid when we initialize parameters with the normal distribution, the number of heads in Equation (4) is $1$, and $\mathcal{F}$ is an identity function. Empirically, the outputs of each sub-layer are close to the normal distribution as illustrated in Appendix F.

###### Assumption 1 .

### 3.1 Jacobian Matrix of FFN

Based on Equation (1), $\|\frac{\partial y_{n}}{\partial x^{\prime}_{n}}\|_{2}$ in Equation (7) can be rewritten as:

$$ $\displaystyle\bigg{\|}\frac{\partial y}{\partial x^{\prime}}\bigg{\|}_{2}$ $\displaystyle=\bigg{\|}\frac{\partial(x^{\prime}+\mathrm{FFN}(\mathrm{LN}(x^{ \prime})))}{\partial x^{\prime}}\bigg{\|}_{2}=\bigg{\|}I+\frac{\partial( \mathrm{FFN}(\mathrm{LN}(x^{\prime})))}{\partial x^{\prime}}\bigg{\|}_{2}.$ (8) $$

We can then derive an upper bound of $\|\frac{\partial y_{n}}{\partial x^{\prime}_{n}}\|_{2}$ by applying the subadditivity, i.e., $\|A+B\|_{2}\leq\|A\|_{2}+\|B\|_{2}$, and submultiplicativity properties of the spectral norm as follows:

$$ $\displaystyle\bigg{\|}\frac{\partial y}{\partial x^{\prime}}\bigg{\|}_{2}$ $\displaystyle\leq 1+\bigg{\|}\frac{\partial\mathrm{FFN}(\mathrm{LN}(x^{\prime} ))}{\partial\mathrm{LN}(x^{\prime})}\bigg{\|}_{2}\bigg{\|}\frac{\partial \mathrm{LN}(x^{\prime})}{\partial x^{\prime}}\bigg{\|}_{2}.$ (9) $$

The right-hand side of this inequality indicates that we can estimate the upper bound of $\|\frac{\partial y}{\partial x^{\prime}}\|_{2}$ by separately computing the spectral norms of Jacobian matrices for $\mathrm{FFN}$ and $\mathrm{LN}$.

Regarding the $\mathrm{FFN}$ part,
we assume that the activation function $\mathcal{F}$ is an identity function(^3^33Appendix G discusses the case where we use the $\mathrm{ReLU}$ as the activation function, which results in the same conclusion with a slightly different coefficient.)
to simplify the discussion.
Under this assumption, the following equation holds:

$$ $\displaystyle\bigg{\|}\frac{\partial\mathrm{FFN}(\mathrm{LN}(x^{\prime}))}{ \partial\mathrm{LN}(x^{\prime})}\bigg{\|}_{2}=\|W_{1}W_{2}\|_{2}.$ (10) $$

Therefore, we can straightforwardly derive the relation $\|W_{1}W_{2}\|_{2}\leq\|W_{1}\|_{2}\|W_{2}\|_{2}$ from the submultiplicativity of the spectral norm.
Furthermore, let $\sigma_{1}$ and $\sigma_{2}$ be the standard deviations of $W_{1}$ and $W_{2}$, respectively.
From Assumption 1, the spectral norms of $W_{1}$ and $W_{2}$ are obtained by their standard deviations and dimensions (Vershynin, ), i.e., $\|W_{1}\|_{2}\approx\sigma_{1}(\sqrt{d}+\sqrt{d_{\mathrm{ffn}}})$ and $\|W_{2}\|_{2}\approx\sigma_{2}(\sqrt{d}+\sqrt{d_{\mathrm{ffn}}})$.
Finally, we can express an upper bound of the spectral norms of the Jacobian matrices for $\mathrm{FFN}$ as the following inequality:

$$ $\displaystyle\bigg{\|}\frac{\partial\mathrm{FFN}(\mathrm{LN}(x^{\prime}))}{ \partial\mathrm{LN}(x^{\prime})}\bigg{\|}_{2}\leq\sigma_{1}\sigma_{2}(\sqrt{d} +\sqrt{d_{\mathrm{ffn}}})^{2},$ (11) $$

where the right-hand side has the relation $\sigma_{1}\sigma_{2}(\sqrt{d}+\sqrt{d_{\mathrm{ffn}}})^{2}=\|W_{1}\|_{2}\|W_{2
}\|_{2}$.

Next, regarding the $\mathrm{LN}$ part, the Jacobian matrix of LN can be written as:

$$ $\displaystyle\frac{\partial\mathrm{LN}(x^{\prime})}{\partial x^{\prime}}=\frac {\sqrt{d}}{\|x^{\prime}\|_{2}}\bigg{(}I-\frac{x^{\prime}x^{\prime\top}}{\|x^{ \prime}\|_{2}^{2}}\bigg{)}$ $\displaystyle=\frac{\sqrt{d}}{\sigma_{x^{\prime}}\sqrt{d}}\bigg{(}I-\frac{x^{ \prime}x^{\prime\top}}{\sigma_{x^{\prime}}^{2}d}\bigg{)}=\frac{1}{\sigma_{x^{ \prime}}}\bigg{(}I-\frac{zz^{\top}}{d}\bigg{)}.$ (12) $$

The leftmost equation appears in the proof by Xiong et al. ().
The second equation uses $\|x^{\prime}\|_{2}=\sigma_{x^{\prime}}\sqrt{d}$, which can be obtained based on Assumption 1.
The last equation is derived from the well-known formula of $z=(x^{\prime}-\mu_{x^{\prime}})/\sigma_{x^{\prime}}$, which converts a normal distribution, $x^{\prime}$, to the standard normal distribution $z$, where $\mu_{x^{\prime}}=0$ in Assumption 1.

We consider the variance ($\mathrm{var}$) of each element in the matrix $zz^{\top}$.
Since $z_{i}z_{i}$ follows $\mathcal{X}^{2}$ with $1$ degree of freedom, and $z_{i}z_{j}(i\neq j)$ is the multiplication of two independent values following the standard normal distribution, the variances are as follows:

$$ $\displaystyle\mathrm{var}(z_{i}z_{j})=\begin{cases}1\ \ \text{if $i\neq j$}\\ 2\ \ \text{otherwise}\end{cases}.$ (13) $$

Equation (13) indicates that $\frac{zz^{\top}}{d}\approx 0$ in LLMs due to $d\gg 1$.
Therefore, the spectral norm of the Jacobian matrix of LN can be written as:

$$ $\displaystyle\bigg{\|}\frac{\partial\mathrm{LN}(x^{\prime})}{\partial x^{ \prime}}\bigg{\|}=\frac{1}{\sigma_{x^{\prime}}},\quad\mbox{where}\quad\frac{ \partial\mathrm{LN}(x^{\prime})}{\partial x^{\prime}}=\frac{1}{\sigma_{x^{ \prime}}}I.$ (14) $$

Finally, Equation (9) can be rewritten by substituting Equations (11) and (14) as:

$$ $\displaystyle\bigg{\|}\frac{\partial y}{\partial x^{\prime}}\bigg{\|}_{2}$ $\displaystyle\leq 1+\frac{\sigma_{1}\sigma_{2}}{\sigma_{x^{\prime}}}C_{\mathrm {ffn}},$ (15) $$

where $C_{\mathrm{ffn}}=(\sqrt{d}+\sqrt{d_{\mathrm{ffn}}})^{2}$ for the simplification.

According to the discussion in Section 2.3 and Equation (15),
the standard deviations, $\sigma_{1}$ and $\sigma_{2}$, of $W_{1}$ and $W_{2}$, respectively, should be sufficiently small, and the standard deviation, $\sigma_{x^{\prime}}$, of the shortcut, $x^{\prime}$, should satisfy $\sigma_{1}\sigma_{2}\ll\sigma_{x^{\prime}}$ in order to keep the upper bound small.

### 3.2 Jacobian Matrix of Self-Attention

Similar to $\mathrm{FFN}$, we can rewrite $\|\frac{\partial x^{\prime}}{\partial x}\|_{2}$ in Equation (7) by using Equation (2) as:

$$ $\displaystyle\bigg{\|}\frac{\partial x^{\prime}}{\partial x}\bigg{\|}_{2}$ $\displaystyle=\bigg{\|}\frac{\partial(x+\mathrm{Attn}(\mathrm{LN}(x)))}{ \partial x}\bigg{\|}_{2}=\bigg{\|}I+\frac{\partial(\mathrm{Attn}(\mathrm{LN}(x )))}{\partial x}\bigg{\|}_{2}.$ (16) $$

We can then derive an upper bound of $\|\frac{\partial x^{\prime}}{\partial x}\|_{2}$ by applying the subadditivity and submultiplicativity of the spectral norm, namely:

$$ $\displaystyle\bigg{\|}\frac{\partial x^{\prime}}{\partial x}\bigg{\|}_{2}$ $\displaystyle\leq 1+\bigg{\|}\frac{\partial\mathrm{Attn}(\mathrm{LN}(x))}{ \partial\mathrm{LN}(x)}\bigg{\|}_{2}\bigg{\|}\frac{\partial\mathrm{LN}(x)}{ \partial x}\bigg{\|}_{2}.$ (17) $$

Therefore, to estimate the upper bound of $\|\frac{\partial x^{\prime}}{\partial x}\|_{2}$, we compute the spectral norms of the Jacobian matrices for $\mathrm{Attn}$ and $\mathrm{LN}$.

Let $Z(\cdot)=\mathrm{concat}(\mathrm{head}_{1}(\cdot),...,\mathrm{head}_{h}(\cdot)))$ and let $J^{Z}$ be the Jacobian of the $Z(\cdot)$(^4^44We discuss the detail of $J^{Z}$ in Appendix I.), we can rewrite the spectral norm of the Jacobian matrix of Attn as:

$$ $\displaystyle\bigg{\|}\frac{\partial\mathrm{Attn}(\mathrm{LN}(x))}{\partial \mathrm{LN}(x)}\bigg{\|}_{2}=\bigg{\|}\frac{\partial W_{O}Z(\mathrm{LN}(x))}{ \partial Z(\mathrm{LN}(x))}\frac{\partial Z(\mathrm{LN}(x))}{\partial\mathrm{ LN}(x)}\bigg{\|}_{2}=\|W_{O}J^{Z}\|_{2}.$ (18) $$

Therefore, we can straightforwardly derive the relation $\|W_{O}J^{Z}\|_{2}\leq\|W_{O}\|_{2}\|J^{Z}\|_{2}$ from the submultiplicativity of the spectral norm.

Let $\sigma_{O}$ be the standard deviation of $W_{O}$.
The relation $\|W_{O}\|_{2}\approx\sigma_{O}(2\sqrt{d})$ is derived from Assumption 1.
We assign this value to Equation (18) and obtain the following inequality:

$$ $\displaystyle\bigg{\|}\frac{\partial\mathrm{Attn}(\mathrm{LN}(x))}{\partial \mathrm{LN}(x)}\bigg{\|}_{2}\leq\sigma_{O}(2\sqrt{d})\|J^{Z}\|_{2}.$ (19) $$

Therefore, we can rewrite Equation (17) by substituting Equations (14) and (19) as follows:

$$ $\displaystyle\bigg{\|}\frac{\partial x^{\prime}}{\partial x}\bigg{\|}_{2}$ $\displaystyle\leq 1+\frac{\sigma_{O}}{\sigma_{x}}C_{\mathrm{Attn}},$ (20) $$

where $C_{\mathrm{Attn}}=(2\sqrt{d})\|J^{Z}\|_{2}$ for the simplification.

Thus, similar to the discussion at the end of Section 3.1, the standard deviation, $\sigma_{O}$, of $W_{O}$ should be small and the standard deviation, $\sigma_{x}$, of the shortcut, $x$, should satisfy $\sigma_{O}\ll\sigma_{x}$ in order to keep the upper bound small.

## 4 Conditions to Avoid Spikes

Based on the discussions in Section 3, we have to pay attention to values of $\sigma_{1}$, $\sigma_{2}$, $\sigma_{O}$, and the standard deviation of the shortcut to stabilize the pre-training of LLMs.
To make $\sigma_{1}$, $\sigma_{2}$, and $\sigma_{O}$ 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., ; Le Scao et al., ; Biderman et al., ), initializes all parameters with a normal distribution $\mathcal{N}(0,\sigma^{2})$ where $\sigma=\sqrt{\frac{2}{5d}}$ (Nguyen & Salazar, ), and then scales $W_{2}$ and $W_{O}$ to small values based on the number of layers: $\sqrt{\frac{1}{2N}}$ where $N$ is the number of layers(^5^55Biderman et al. () also scaled $W_{2}$ and $W_{O}$ to small values in the initialization, but they used the strategy introduced by Wang & Komatsuzaki () instead of scaling with $\sqrt{\frac{1}{2N}}$. However, its property is the same essentially because they initialize $W_{2}$ and $W_{O}$ with $\sigma=\frac{2}{N\sqrt{d}}$ which becomes small based on the number of layers.).
In this situation, $\sigma_{1}$, $\sigma_{2}$, and $\sigma_{O}$ 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 $\sqrt{\frac{2}{5d}}$ because the embedding matrix is also initialized by $\mathcal{N}(0,\sigma^{2})$ where $\sigma=\sqrt{\frac{2}{5d}}$.
Therefore, to increase the standard deviation of the shortcut, we make the standard deviation of each embedding close to $1$(^6^66Based on Equations (15) and (20), the upper bound becomes small as the standard deviation of the shortcut increases. However, a too large value degrades the performance empirically as described in Appendix E.).
To achieve this, we introduce two kinds of modification: “Scaled Embed” and “Embed LN”(^7^77We can satisfy the condition by initializing embeddings with the normal distribution $\mathcal{N}(0,\sigma^{2})$ where $\sigma=1$, but we do not adopt this strategy in this study because we use the same initialization method in our experiments.).
The Scaled Embed scales embeddings with an appropriate value.
For example, we multiply embeddings by $\sqrt{d}$, which was used in the original Transformer paper (Vaswani et al., )(^8^88Although the original Transformer paper introduced this operation, recent implementations ignore this.), and then the standard deviations of embeddings become $\sqrt{\frac{2}{5}}$.
The Embed LN applies the LN to embeddings.
In fact, Le Scao et al. () 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.

Figure: Figure 2: The actual upper bound described in Equation (15) for each Transformer layer at the beginning of the LLM pre-training. Because it is difficult to estimate the strict values for $\sigma_{x}$ at all layers, we obtain the empirical values by using some inputs, and assign them to Equation (15).
Refer to caption: x2.png

To demonstrate the actual values of the upper 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 Embed and Embed LN, we compare Xavier Init, which initializes all parameters with the Xavier initialization (Glorot & Bengio, ), as the situation where we do not scale $W_{2}$ and $W_{O}$ based on the number of layers.
Figure 2 shows the values of Equation (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 in 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 conditions: (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$.

**Table 1: Relations between each method in experiments and two conditions to control the upper bound of gradient norms. We conduct experiments for Xavier Init and Xavier Init + Scaled Embed in Appendix B.**
| Method | Small sub-layers | Large shortcut |
| --- | --- | --- |
| Xavier Init | - | - |
| Vanilla | $\checkmark$ | - |
| Embed Detach | $\checkmark$ | - |
| Embed LN | $\checkmark$ | $\checkmark$ |
| Scaled Embed | $\checkmark$ | $\checkmark$ |
| Xavier Init + Scaled Embed | - | $\checkmark$ |

## 5 Main 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., ; Le Scao et al., ) in main experiments(^9^99The Xavier initialization, which does not satisfy the small sub-layers condition as shown in Table 1, is not widely used for LLMs in recent years. Appendix B shows that the performance of Xavier initialization is worse and fails to avoid spikes.).

### 5.1 Datasets

We used C4 (Raffel et al., ) that consists of clean English texts extracted from Common Crawl(^10^1010[https://commoncrawl.org/](https://commoncrawl.org/)) as our LLM pre-training corpus.
We also used the separated part of C4 as our validation data.
We used GPT-2 vocabulary (Radford et al., ) that contains Byte Pair Encoding (BPE) subword units (Sennrich et al., ) as our vocabulary.
To evaluate each method, we computed perplexity on WikiText (Merity et al., ) and LAMBADA (Paperno et al., ) datasets.

### 5.2 Model Configurations

As described in Section 2, we used the Pre-LN Transformer architecture.
We set the number of layers $N=24$, and varied $d$ to adjust the total number of parameters to 350 million (350M) and 1.7 billion (1.7B).
We set the learning rate ($\mathrm{lr}$) $5.0\times 10^{-4}$.
Section 6.1 shows experiments in varying the learning rate.
Appendix A describes more details on the experimental configuration.

Figure: (a) 350M parameters.
Refer to caption: x3.png

Figure: (a) 350M parameters.
Refer to caption: x5.png

We compared the following methods.
We put $\checkmark$ before the method name if the method satisfies both conditions to suppress the upper bound.
Moreover, Table 1 summarizes whether each method satisfies each condition.

#### Vanilla

This is the standard configuration for the LLM pre-training.
Since this configuration does not suppress the upper bound of the gradient norms, the loss spike is likely to occur.

#### Embed Detach

Zeng et al. () used the shrink embedding gradient technique (Ding et al., ) 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:

$$ $\displaystyle\mathrm{Embed}\leftarrow\gamma\mathrm{Embed}+(1-\gamma)\mathrm{ Detach}(\mathrm{Embed}),$ (21) $$

where $\gamma$ is a hyper-parameter and $\mathrm{Detach}$ detaches an input from the computational graph.
We assign $0.1$ to $\gamma$ as in Zeng et al. ().
Zeng et al. () 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.

#### ✓ ✓ \checkmark ✓ Embed LN

Dettmers et al. () and Le Scao et al. () 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.

#### ✓ ✓ \checkmark ✓ Scaled Embed

This method multiplies embeddings by $\sqrt{d}$.
As described in Section 4, this method satisfies the requirements to control the upper bound of the gradient norms.

### 5.3 Results

Figure 3 shows the loss values of each method in validation data.
Figure 4 shows the gradient 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 parameters, spikes occurred more frequently in 1.7B parameters.
Because we initialize embeddings with $\mathcal{N}(0,\sigma^{2})$ where $\sigma=\sqrt{\frac{2}{5d}}$, the standard deviations 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 $\sigma_{x}$ and $\sigma_{x}^{\prime}$ are nearly equal to the standard deviation of an input embedding in shallow layers.
Therefore, if we increase $d$ without any technique to control the upper bound of the gradient norms, a model becomes more unstable.
This result corresponds to the previous study reports (Le Scao et al., ; Chowdhery et al., ; Zeng et al., ) that their model became more unstable as they increased the number of parameters.

**Table 2: Perplexities of each method.**
| Model | WikiText $\downarrow$ | LAMBADA $\downarrow$ |
| --- | --- | --- |
| 350M parameters |  |  |
| Vanilla | 30.03 | 24.73 |
| Embed Detach | 30.69 | 26.93 |
| Embed LN | 29.85 | 25.03 |
| Scaled Embed | 29.86 | 24.37 |
| 1.7B parameters |  |  |
| Vanilla | 22.58 | 15.22 |
| Embed Detach | 22.00 | 13.88 |
| Embed LN | 21.29 | 13.00 |
| Scaled Embed | 21.29 | 12.53 |

Table 2 shows the perplexities of each method on 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 if 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 a 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 discuss this matter in more detail in Section 6.1.

## 6 Discussions on Other Configurations

In this section, we conduct experiments on other configurations to describe connections with previous study reports.

### 6.1 Varying Learning Rate

Le Scao et al. () reported that the stable method, such as Embed LN, was worse than Vanilla.
However, in Section 5, 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: $\mathrm{lr}=1.0\times 10^{-3}$ and $1.0\times 10^{-4}$ respectively.

Figure: (a) 350M with $\mathrm{lr}=1.0\times 10^{-3}$.
Refer to caption: x7.png

Figure 5 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 $\mathrm{lr}=1.0\times 10^{-3}$, 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 3 shows the perplexities of each configuration in evaluation data.
This table indicates that Vanilla with 350M parameters achieved better performance in $\mathrm{lr}=1.0\times 10^{-4}$ that is the situation where its training did not face any spike.
This result corresponds to the report of Le Scao et al. ().
Thus, we suppose that they conducted the comparison with a too-small learning rate to stabilize Vanilla.
In contrast, the stable methods are more effective in training with a large learning rate, as shown in Figure 5 and Table 3.
Therefore, if Le Scao et al. () used a relatively large learning rate in their experiments, their stable method could achieve better performance.

**Table 3: Perplexities of each method with 350M and 1.7B parameters when we vary a learning rate.**
|  | WikiText $\downarrow$ | LAMBADA $\downarrow$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Model | $\mathrm{lr}\ 1.0\times 10^{-3}$ | $\mathrm{lr}\ 5.0\times 10^{-4}$ | $\mathrm{lr}\ 1.0\times 10^{-4}$ | $\mathrm{lr}\ 1.0\times 10^{-3}$ | $\mathrm{lr}\ 5.0\times 10^{-4}$ | $\mathrm{lr}\ 1.0\times 10^{-4}$ |
| 350M parameters |  |  |  |  |  |  |
| Vanilla | 29.96 | 30.35 | 34.51 | 25.12 | 24.73 | 32.49 |
| Scaled Embed | 28.09 | 29.86 | 35.66 | 22.03 | 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 |

### 6.2 Varying Sequence Length

Li et al. () 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 justification to their method in terms of the standard deviation of the shortcut.

Figure: Figure 6: 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.
Refer to caption: x13.png

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 sequence especially at the early stage(^11^1111See Appendix H for details.), a long sequence tends to keep the standard deviation of the shortcut small.
Therefore, the long sequence makes the pre-training of Vanilla more unstable.

We conducted experiments with varying the length of the input sequence $L$ from $128$ to $2048$.
To use the same number of tokens to update parameters, we adjusted the batch size.
Figure 6 shows loss values of Vanilla with each $L$ configuration 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.
Therefore, using a short sequence stabilizes the training at the early stage, as reported in Li et al. ().

## 7 Related Work

#### Stability

To stabilize trainings of Transformer-based neural language models, there have been various discussions on the architecture (Xiong et al., ; Liu et al., ; Takase et al., ; Zeng et al., ; Zhai et al., ), initialization method (Nguyen & Salazar, ; Zhang et al., ; Huang et al., ; Wang et al., ), training strategy (Zhang et al., ; Li et al., ), and loss function (Chowdhery et al., ; Wortsman et al., ).
Xiong et al. () theoretically analyzed gradient scales of 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., ).
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. () applied the layer normalization to the embedding layer.
Zeng et al. () used shrink embedding gradient technique (Ding et al., ).
In this study, we theoretically proved that the layer normalization to the embedding layer controls the upper bound of the gradient norms of sub-layers when we use the widely used initialization method for LLMs (Nguyen & Salazar, ; Shoeybi et al., ), and thus, it stabilizes the pre-training.
For the initialization methods, Nguyen & Salazar () proposed a strategy to initialize parameters of Transformers with small values to stabilize their training.
Zhang et al. () and Huang et al. () indicated that we can remove layer normalizations in Transformers if we use their proposed initialization methods.
Wang et al. () 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., ), 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 $1$.

#### Efficiency

As shown in Table 3, our modification enables the pre-training with a relatively larger learning rate, and can achieve better performance.
Thus, this study can be regarded as on the efficiency of LLM pre-training because our modification can construct a better LLM with a given budget.
Strubell et al. () and Schwartz et al. () reported that recent neural methods require substantial computational costs, and thus, they argued that we have to explore a cost-efficient approach.
Rajbhandari et al. () proposed ZeRO that reduces memory redundancies during the multi GPU training without increasing communication volume.
Dao et al. () focused on GPU memory reads/writes, and proposed FlashAttention that accelerates the speed of attention mechanisms in Transformers.
To reduce the number of computations in the attention mechanism, Shazeer () proposed the multi-query attention that shares one key and value across all of the attention heads in each layer.
Takase & Kiyono () explored several parameter sharing strategies, and indicated that parameter sharing across some layers can achieve comparable performance to the vanilla model with a small number of parameters.
Moreover, several studies have explored a better construction way with a limited budget (Izsak et al., ; Takase & Kiyono, ).
We believe that we can take advantage of their findings to make our LLMs more efficient.

## 8 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 LLMs 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 embedding close to $1$, resulting in large shortcut.
Experimental results indicated that methods satisfying these conditions avoid loss spikes.
Furthermore, these methods allow for training with a relatively larger learning rate, leading to improved performance.
We hope our theoretical analyses and empirical findings will help avoid wasting valuable time and computational budgets during LLM construction.

#### Ethics Statement

To stabilize the LLM pre-training, this paper provides theoretical analyses om the spectral norms of the Jacobian matrices for sub-layers to estimate the upper bound of the gradient norm of $\mathcal{L}$.
This paper focuses on only the stability of LLM pre-training, and thus, we have to address other issues of LLMs such as hallucinations to use the LLM in a real application.

#### Reproducibility Statement

We do not aim to propose a novel method in this paper, but we mainly focus on theoretical analyses on 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 various situations.
To activate our modification, we add only several lines to a widely used implementation, i.e., Megatron-LM(^12^1212[https://github.com/NVIDIA/Megatron-LM](https://github.com/NVIDIA/Megatron-LM)).
Therefore, we believe that it is easy to reproduce our experimental results.
However, because it is difficult to conclude that our provided conditions completely solve the instability during pre-training of LLMs, it is better to combine other techniques to stabilize the pre-training such as an auxiliary loss described by Chowdhery et al. () to make the pre-training more stable in an actual pre-training situation.

## Appendix A Hyper-Parameters

Table 4 shows that hyper-parameters used in our experiments on LLMs.
In addition to experiments described in Section 5, this table also indicates the hyper-parameters of the model with 13B parameters that we evaluated in Appendix C.

**Table 4: Hyper-parameters used in our experiments on the LLM pre-training.**
| Name | 350M | 1.7B | 13B |
| --- | --- | --- | --- |
| Precision | float16 | float16 | float16 |
| 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 |
| Sequence length | 2048 | 2048 | 2048 |
| Batch size | 528 | 528 | 1024 |
| The number of updates | 35000 | 35000 | 50000 |
| Adam $\beta_{1}$ | 0.9 | 0.9 | 0.9 |
| Adam $\beta_{2}$ | 0.999 | 0.999 | 0.95 |
| Gradient clipping | 1.0 | 1.0 | 1.0 |
| $lr$ decay style | cosine | cosine | cosine |
| $lr$ warmup fraction | 0.05 | 0.05 | 0.05 |
| Weight decay | 0.01 | 0.01 | 0.01 |

## Appendix B 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 Transformer initialized by the Xavier initialization (Glorot & Bengio, ) (Xavier Init), and the combination of Xavier Init and Scaled Embed (Xavier Init + Scaled Embed) with Scaled Embed.
As described in Table 1, Xavier Init violates both conditions, and Xavier Init + Scaled Embed satisfies the only large shortcut condition.
In the same manner as in Section 5, we trained models of 350M and 1.7B parameters with $\mathrm{lr}=5.0\times 10^{-4}$.
We also used the hyper-parameters described in Table 4.

Figure 7 shows loss curves in validation data for 350M and 1.7B parameters in each method, and Figure 8 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.
In 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 5 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.

Figure: (a) 350M parameters.
Refer to caption: x14.png

Figure: (a) 350M parameters.
Refer to caption: x16.png

**Table 5: Perplexities of each method.**
| Model | WikiText $\downarrow$ | LAMBADA $\downarrow$ |
| --- | --- | --- |
| 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 |

## Appendix C Pre-training of the Model with 13B Parameters

Figure: (a) $\mathrm{lr}=3.0\times 10^{-4}$ .
Refer to caption: x18.png

**Table 6: Perplexities of each method with 13B parameters when we use two learning rates: $\mathrm{lr}=3.0\times 10^{-4}$ and $1.0\times 10^{-4}$.**
|  | WikiText $\downarrow$ | LAMBADA $\downarrow$ |  |  |
| --- | --- | --- | --- | --- |
| Model | $\mathrm{lr}=3.0\times 10^{-4}$ | $\mathrm{lr}=1.0\times 10^{-4}$ | $\mathrm{lr}=3.0\times 10^{-4}$ | $\mathrm{lr}=1.0\times 10^{-4}$ |
| Vanilla | N/A | 15.12 | N/A | 6.50 |
| Scaled Embed | 14.47 | 15.25 | 5.97 | 6.53 |

We conducted pre-trainings of models with 13B parameters to indicate that our modification can stabilize a model with many more parameters than the ones discussed in Section 5.
To make this experiment close to a realistic situation, as shown in Table 4, we increased the batch size and the number of updates, and decreased the Adam $\beta_{2}$.
In particular, for Adam $\beta_{2}$, most studies have used $0.95$ to stabilize their pre-trainings (Brown et al., ; Zhang et al., ; Zeng et al., ; Biderman et al., ; Touvron et al., ), and thus, we also used $0.95$ in this experiment.
We tried two learning rates: $3.0\times 10^{-4}$, which is the same value in Touvron et al. (), and $1.0\times 10^{-4}$.

Figure 9 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 $\mathrm{lr}=3.0\times 10^{-4}$.
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.
We emphasize that the pre-training of Vanilla is essentially unstable even if we use the widely used Adam $\beta_{2}$ value, $0.95$, which is known as the technique to stabilize the pre-training, and our modification is also effective for the stabilization in this realistic situation.

Table 6 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 the same as in Section 6.1.
In addition, the perplexities of Scaled Embed were comparable to ones of Vanilla when we used the small learning rate: $\mathrm{lr}=1.0\times 10^{-4}$.
These results imply that our modification has 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.

## Appendix D RMSNorm

Some recent LLMs use the RMSNorm (Zhang & Sennrich, ) instead of the LN in their Transformers (Touvron et al., ).
We discuss such an architecture in this section.
In the same as LN discussed in Section 3.1, we can obtain the Jacobian matrix of the RMSNorm with the following equation:

$$ $\displaystyle\bigg{\|}\frac{\partial\mathrm{RMSNorm}(x)}{\partial x}\bigg{\|}_ {2}=\frac{1}{\sigma_{x}}I$ (22) $$

Thus, the upper bound of the gradient norm is the same in LN if we use RMSNorm.

Figure 10 shows the loss values and gradient norms of the Vanilla configuration in Section 5 and the one using RMSNorms instead of LNs (“RMSNorm” in figures) with 350M parameters.
We trained them with $\mathrm{lr}=5.0\times 10^{-4}$ as in Section 5(^13^1313We tried to train them with $\mathrm{lr}=1.0\times 10^{-3}$ but RMSNorm exploded.).
As shown in these figures, RMSNorm faced loss and gradient spikes in a similar location to the ones of Vanilla.
These empirical results also indicate that the RMSNorms have the same problem as LNs regarding the instability.

Figure: (a) Loss curves in validation data.
Refer to caption: x20.png

## Appendix E 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 a 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 to $1$.
To investigate how about a 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 situation.
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 a larger value than $1$ to prevent the performance degradation.

Figure: (a) 350M parameters.
Refer to caption: x22.png

## Appendix F Distributions of Sub-Layer Outputs

Figure 12 shows output distributions of each sub-layer for each layer at the initialization.
This figure indicates that each sub-layer output is close to the normal distribution in various configurations.
Therefore, the assumption in this study, which is that the vector $x$ at each layer follows the normal distribution, is reasonable.

Figure: (a) self-attention of 350M parameters.
Refer to caption: x24.png

## Appendix G Discussion on ReLU in FFN

We consider the case where we use the ReLU function as $\mathcal{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 element of the FFN internal layer is a negative value with half probability.
In this case, we can regard that the ReLU function cuts the elements of the internal layer by half.
Thus, we replace $W_{1}\in\mathbb{R}^{d_{\mathrm{ffn}}\times d}$ and $W_{2}\in\mathbb{R}^{d\times d_{\mathrm{ffn}}}$ with $W_{1}\in\mathbb{R}^{\frac{d_{\mathrm{ffn}}}{2}\times d}$ and $W_{2}\in\mathbb{R}^{d\times\frac{d_{\mathrm{ffn}}}{2}}$ in the discussion in Section 3.1 when we use the ReLU function as $\mathcal{F}$.

## Appendix H 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 $(W_{Qi}\ x)^{\mathrm{T}}(W_{Ki}\ X)$ is 0 0.
Therefore, the expectation after the softmax function is $\frac{1}{L}$ where $L$ is the length of the input sequence.
Thus, the long input sequence decreases the standard deviation of the self-attention layer.

To simplify, we consider the case where the number of self-attention heads is $1$.
In this case, we can obtain the variance of each calculation with the following equation.

$$ $\displaystyle\mathrm{var}(W_{O}(x))$ $\displaystyle=\mathrm{var}(W_{O})\mathrm{var}(x)\ d$ (23) $\displaystyle\mathrm{var}(W_{V}(x))$ $\displaystyle=\mathrm{var}(\mathrm{W_{V}})\mathrm{var}(x)\ d$ (24) $$

where $\mathrm{var}$ represents the variance of the matrix/vector.
Thus, the variance of the self-attention layer, $\mathrm{var}(\mathrm{Attn}(x))$, is as follows:

$$ $\displaystyle\mathrm{var}(\mathrm{Attn}(x))$ $\displaystyle=\mathrm{var}(W_{O})\ d\sum^{L}\frac{\mathrm{var}(\mathrm{W_{V}}) \mathrm{var}(x)\ d}{L^{2}}$ (25) $\displaystyle=\frac{\mathrm{var}(W_{O})\mathrm{var}(\mathrm{W_{V}})\mathrm{var }(x)d^{2}}{L}$ (26) $$

## Appendix I Details on Jacobian Matrix of Self-Attention

We can represent $\mathrm{concat}(\mathrm{head}_{1}(x),...,\mathrm{head}_{h}(x))$ with the summation of the matrix multiplications as follows:

$$ $\displaystyle\mathrm{concat}(\mathrm{head}_{1}(x),...,\mathrm{head}_{h}(x))= \sum_{i=1}^{h}\mathrm{head}_{i}W_{i}$ (27) $$

where $W_{i}\in\mathbb{R}^{d_{\mathrm{head}}\times d}$ whose corresponding element is $1$ and the others are 0 0.
Let $J^{i}$ be the Jacobian of the $\mathrm{head}_{i}(x)$, we can represent $J^{Z}$ in Section 3.2 as follows:

$$ $\displaystyle J^{Z}=\sum_{i=1}^{h}J^{i}W_{i}$ (28) $$

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 $x_{j}$ be the input of the position $j$, and $z_{k}^{i}$ be the $i$-th head of the output position $k$, and $J_{kj}^{i}=\frac{\partial z_{k}^{i}}{\partial x_{j}}$.
Because $J^{Z}$ can be regarded as the Jacobian of the input of the position $j$, we can convert Equation (29) into the following Equation:

$$ $\displaystyle J^{Z}=\sum_{k=1}^{L}\sum_{i=1}^{h}J_{kj}^{i}W_{i}$ (29) $$

where $L$ is the length of the input and output sequences.
Therefore, we compute $J_{kj}^{i}$ to obtain $J^{Z}$ in Section 3.2.

We can obtain a head of the output position $k$, i.e., $z_{k}$, as follows(^14^1414To simplify the equations, we omit the index $i$ to represent $i$-th head from the head and parameters.)

$$ $\displaystyle z_{k}=\sum_{l=1}^{L}A_{kl}v_{l}$ (30) $$

where $A_{kl}$ is the $l$-th element of the attention vector, $\mathrm{softmax}\left(\frac{(W_{Q}\ x_{k})^{\mathrm{T}}(W_{K}\ X)}{\sqrt{d_{
\mathrm{head}}}}\right)$ and $v_{l}$ is $W_{V}x_{l}$.
Therefore, to obtain $J_{kj}$, we differentiate $z_{k}$ with respect to $x_{j}$ as:

$$ $\displaystyle J_{kj}=\frac{\partial z_{k}}{\partial x_{j}}=\sum_{l=1}^{L}\left (\frac{\partial A_{kl}}{\partial x_{j}}{v_{l}}^{\mathrm{T}}+A_{kl}\frac{ \partial v_{l}}{\partial x_{j}}\right)$ (31) $$

$$ $\displaystyle\frac{\partial v_{l}}{\partial x_{j}}$ $\displaystyle=W_{V}\delta_{lj}$ (32) $\displaystyle\delta_{lj}$ $\displaystyle=\begin{cases}1\ \ \text{if $l=j$}\\ 0\ \ \text{otherwise}\end{cases}$ (33) $$

Thus,

$$ $\displaystyle\frac{\partial z_{k}}{\partial x_{j}}=\sum_{l=1}^{L}\left(\frac{ \partial A_{kl}}{\partial x_{j}}{v_{l}}^{\mathrm{T}}\right)+A_{kj}W_{V}$ (34) $$

Here, we assume that the attention vector is uniform.
In this assumption, since $A_{kj}=\frac{1}{L}$, we can obtain the spectral norm of the second term for Equation (34) as $\|A_{kj}W_{V}\|_{2}\approx\frac{\sigma_{V}}{L}(\sqrt{d}+\sqrt{d_{\mathrm{head}
}})$, where $\sigma_{V}$ is the standard deviation of $W_{V}$.
To calculate the first term, we compute $\frac{\partial A_{kl}}{\partial x_{j}}$.

$$ $\displaystyle\frac{\partial A_{kl}}{\partial x_{j}}$ $\displaystyle=A_{kl}\left(\frac{\partial S_{kl}}{\partial x_{j}}-\sum_{m=1}^{L }A_{km}\frac{\partial S_{km}}{\partial x_{j}}\right)$ (35) $\displaystyle\frac{\partial S_{kl}}{\partial x_{j}}$ $\displaystyle=\frac{1}{\sqrt{d_{\mathrm{head}}}}\left(W_{Q}^{\mathrm{T}}W_{K}x _{l}\delta_{kj}+W_{K}^{\mathrm{T}}W_{Q}x_{k}\delta_{lj}\right)$ (36) $$

Let $D_{l}$ be $W_{Q}^{\mathrm{T}}W_{K}x_{l}$ and $E_{k}$ be $W_{K}^{\mathrm{T}}W_{Q}x_{k}$. Then,

$$ $\displaystyle\frac{\partial S_{kl}}{\partial x_{j}}$ $\displaystyle=\frac{1}{\sqrt{d_{\mathrm{head}}}}\left(D_{l}\delta_{kj}+E_{k} \delta_{lj}\right)$ (37) $$

Therefore,

$$ $\displaystyle\frac{\partial A_{kl}}{\partial x_{j}}$ $\displaystyle=\frac{A_{kl}}{\sqrt{d_{\mathrm{head}}}}\left((D_{l}\delta_{kj}+E _{k}\delta_{lj})-\sum_{m=1}^{L}A_{km}(D_{m}\delta_{kj}+E_{k}\delta_{mj})\right)$ (38) $\displaystyle=\frac{A_{kl}}{\sqrt{d_{\mathrm{head}}}}\left(\delta_{lj}E_{k}-A_ {kj}E_{k}+\delta_{kj}\left(D_{l}-\sum_{m=1}^{L}A_{km}D_{m}\right)\right)$ (39) $$

We assign Equation (39) to the first term of Equation (34) and use the assumption $A_{kj}=\frac{1}{L}$:

$$ $\displaystyle\sum_{l=1}^{L}\left(\frac{\partial A_{kl}}{\partial x_{j}}{v_{l}} ^{\mathrm{T}}\right)$ $\displaystyle=\sum_{l=1}^{L}\left(\frac{A_{kl}}{\sqrt{d_{\mathrm{head}}}}\left (\delta_{lj}E_{k}-A_{kj}E_{k}+\delta_{kj}\left(D_{l}-\sum_{m=1}^{L}A_{km}D_{m} \right){v_{l}}^{\mathrm{T}}\right)\right)$ (40) $\displaystyle=\frac{1}{L\sqrt{d_{\mathrm{head}}}}\sum_{l=1}^{L}\left(\delta_{ lj}E_{k}-\frac{1}{L}E_{k}+\delta_{kj}\left(D_{l}-\sum_{m=1}^{L}\frac{1}{L}D_{m }\right){v_{l}}^{\mathrm{T}}\right)$ (41) $\displaystyle=\frac{1}{L\sqrt{d_{\mathrm{head}}}}\left(\sum_{l=1}^{L}\left( \left(\delta_{lj}E_{k}+\delta_{kj}D_{l}\right){v_{l}}^{\mathrm{T}}\right)+ \left(\sum_{l=1}^{L}\left(-\frac{1}{L}E_{k}-\sum_{m=1}^{L}\frac{1}{L}D_{m} \right){v_{l}}^{\mathrm{T}}\right)\right)$ (42) $\displaystyle=\frac{1}{L\sqrt{d_{\mathrm{head}}}}\left(E_{k}{v_{j}}^{\mathrm{T }}+\sum_{l=1}^{L}\left(\delta_{kj}D_{l}{v_{l}}^{\mathrm{T}}\right)-\left(E_{k} +\sum_{m=1}^{L}D_{m}\right){v_{l}}^{\mathrm{T}}\right)$ (43) $\displaystyle=\frac{1}{L\sqrt{d_{\mathrm{head}}}}\left(E_{k}{v_{j}}^{\mathrm{T }}-E_{k}{v_{l}}^{\mathrm{T}}+\delta_{kj}\sum_{l=1}^{L}\left(D_{l}{v_{l}}^{ \mathrm{T}}\right)-\left(\sum_{m=1}^{L}D_{m}\right){v_{l}}^{\mathrm{T}}\right)$ (44) $$

Based on the assumption that parameters and the vector at each layer follow the normal distribution, we assume that the mean of $D_{l}$, $E_{k}$, and $v_{l}$ is 0 0, and thus, we can obtain their norms from their variances.
In addition, we assume that the standard deviation of $W_{Q}$, $W_{K}$ and $W_{V}$ is $\sigma$.
Then, $\mathrm{var}(D_{l})=\mathrm{var}(E_{k})=dd_{\mathrm{head}}\sigma^{4}$, $\mathrm{var}(\sum_{l}^{L}D_{l})=Ldd_{\mathrm{head}}\sigma^{4}$, and $\mathrm{var}(v_{l})=d\sigma^{2}$.
Thus, $\|E_{k}v_{j}^{\mathrm{T}}\|_{2}\approx\|E_{k}v_{l}^{\mathrm{T}}\|_{2}\approx
\sigma^{3}\sqrt{d^{3}d_{\mathrm{head}}^{2}}$, $\|\sum_{l=1}^{L}(D_{l}v_{l}^{\mathrm{T}})\|_{2}\approx\|(\sum_{m=1}^{L}D_{m})v
_{l}^{\mathrm{T}}\|_{2}\approx\sigma^{3}\sqrt{Ld^{3}d_{\mathrm{head}}^{2}}$.
Therefore,

$$ $\displaystyle\Bigg{\|}\sum_{l=1}^{L}\left(\frac{\partial A_{kl}}{\partial x_{j }}{v_{l}}^{\mathrm{T}}\right)\Bigg{\|}_{2}$ $\displaystyle\leq\frac{1}{L\sqrt{d_{\mathrm{head}}}}\left(\Big{\|}E_{k}{v_{j}} ^{\mathrm{T}}\Big{\|}_{2}+\Big{\|}E_{k}{v_{l}}^{\mathrm{T}}\Big{\|}_{2}+\delta _{kj}\Bigg{\|}\sum_{l=1}^{L}\left(D_{l}{v_{l}}^{\mathrm{T}}\right)\Bigg{\|}_{2 }+\Bigg{\|}\left(\sum_{m=1}^{L}D_{m}\right){v_{l}}^{\mathrm{T}}\Bigg{\|}_{2}\right)$ (45) $\displaystyle\approx\frac{1}{L\sqrt{d_{\mathrm{head}}}}\left(2\sigma^{3}\sqrt{ d^{3}d_{\mathrm{head}}^{2}}+(\delta_{kj}+1)\sigma^{3}\sqrt{Ld^{3}d_{\mathrm{ head}}^{2}}\right)$ (46) $\displaystyle=\frac{1}{\sqrt{L}}\left(\delta_{kj}+1+\frac{2}{\sqrt{L}}\right) \sigma^{3}\sqrt{d^{3}d_{\mathrm{head}}}$ (47) $$

Based on this Equation, we can compute the upper bound of $\|J_{kj}\|_{2}$:

$$ $\displaystyle\|J_{kj}\|_{2}$ $\displaystyle=\Bigg{\|}\sum_{l=1}^{L}\left(\frac{\partial A_{kl}}{\partial x_{ j}}{v_{l}}^{\mathrm{T}}\right)+A_{kj}W_{V}\Bigg{\|}_{2}$ $\displaystyle\leq\Bigg{\|}\sum_{l=1}^{L}\left(\frac{\partial A_{kl}}{\partial x _{j}}{v_{l}}^{\mathrm{T}}\right)\Bigg{\|}_{2}+\|A_{kj}W_{V}\|_{2}$ (48) $\displaystyle\leq\frac{1}{\sqrt{L}}\left(\delta_{kj}+1+\frac{2}{\sqrt{L}} \right)\sigma^{3}\sqrt{d^{3}d_{\mathrm{head}}}+\frac{\sigma}{L}(\sqrt{d}+\sqrt {d_{\mathrm{head}}})$ (49) $$

Moreover, we can compute the upper bound of $\|J^{Z}\|_{2}$ from Equation (29) as follows:

$$ $\displaystyle\|J^{Z}\|_{2}$ $\displaystyle=\bigg{\|}\sum_{k=1}^{L}\sum_{i=1}^{h}J_{kj}^{i}W_{i}\bigg{\|}_{2}$ (50) $\displaystyle\leq\sum_{k=1}^{L}\sum_{i=1}^{h}\|J_{kj}^{i}\|_{2}\|W_{i}\|_{2}$ (51) $\displaystyle\approx h\left(\frac{1}{\sqrt{L}}\left(1+L+\frac{2L}{\sqrt{L}} \right)\sigma^{3}\sqrt{d^{3}d_{\mathrm{head}}}+\frac{L\sigma}{L}(\sqrt{d}+ \sqrt{d_{\mathrm{head}}})\right)$ (52) $\displaystyle=h\left(\left(\sqrt{L}+2+\frac{1}{\sqrt{L}}\right)\sigma^{3}\sqrt {d^{3}d_{\mathrm{head}}}+\sigma(\sqrt{d}+\sqrt{d_{\mathrm{head}}})\right)$ (53) $$

## Appendix J Comparison with Post-LN Transformer

As described in Section 2.1, recent studies use the Pre-LN Transformer architecture to construct their LLMs because the architecture is more stable.
In contrast, some recent studies reported that the Post-LN Transformer, which is the original architecture, can achieve better performance than the Pre-LN if we address the instability issue in the Post-LN, i.e., the vanishing gradient problem (Liu et al., ; Takase et al., ; Wang et al., ).
We discuss whether the Pre-LN Transformer entirely underperforms the Post-LN.
We conducted experiments on machine translation experiments because previous studies mainly focused on them.

We followed the experimental settings in Takase et al. ().
Table 7 shows the details of hyper-parameters.
We used the WMT English-to-German training dataset (Vaswani et al., ; Ott et al., ), and evaluated each model in newstest2010-2016.
We used the encoder-decoder architecture proposed by Peitz et al. ().
To stabilize the Post-LN Transformer, we applied DeepNet (Wang et al., ) and B2T connection (Takase et al., ).
We compared them to Scaled Embed, that is, the Pre-LN Transformer with the stabilizing techniques described in this paper.

Table 8 shows the averaged BLEU scores among newstest2010-2016.
For the BLEU score calculation, we used SacreBLEU (Post, ) to obtain compatible scores (Marie et al., ).
The signature of SacreBLEU is BLEU+case.mixed+numrefs.1+smooth.exp+tok.13a+version.1.5.0.
As shown in this table, we used two learning rates: $\mathrm{lr}=1.0\times 10^{-3}$ and $3.0\times 10^{-3}$.
For $\mathrm{lr}=1.0\times 10^{-3}$, DeepNet and B2T connection outperformed Scaled Embed.
Thus, the Post-LN Transformer-based methods achieved better performance than the Pre-LN Transformer-based method.
This result corresponds to reports in previous studies (Liu et al., ; Wang et al., ; Takase et al., ).

On the other hand, for $\mathrm{lr}=3.0\times 10^{-3}$, Scaled Embed achieved better performance than the others with $\mathrm{lr}=1.0\times 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 if 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, it can achieve better performance when we use a large learning rate that is too large to train the Post-LN Transformers.

**Table 7: Hyper-parameters used in the comparison with Post-LN Transformer.**
| Name | Value |
| --- | --- |
| Precision | float16 |
| Layer num | 18 |
| Hidden dim size | 512 |
| FFN dim size | 2048 |
| Attention heads | 8 |
| Dropout rate | 0.5 |
| Word dropout rate | 0.1 |
| Max tokens | 7168 |
| Adam $\beta_{1}$ | 0.9 |
| Adam $\beta_{2}$ | 0.98 |
| Gradient clipping | 0.1 |
| $lr$ decay style | inverse square root |
| Warmup step | 4000 |
| Weight decay | 0 |

**Table 8: Averaged BLEU scores among newstest2010-2016.**
| Model | 2010 | 2011 | 2012 | 2013 | 2014 | 2015 | 2016 | Average $\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\mathrm{lr}=1.0\times 10^{-3}$ |  |  |  |  |  |  |  |  |
| 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 |
| $\mathrm{lr}=3.0\times 10^{-3}$ |  |  |  |  |  |  |  |  |
| DeepNet | N/A | N/A |  |  |  |  |  |  |
| B2T connection | N/A | N/A |  |  |  |  |  |  |
| Scaled Embed | 24.52 | 22.23 | 22.86 | 26.54 | 27.35 | 29.90 | 35.16 | 26.94 |

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see the above weaknesses section.
soundness: 2
strengths: * The findings (i) and (ii) from the analysis are well presented, although they have been previously utilized in past studies.
* This work examines various learning hyperparameters.
* This work also presents the results for the 13B model in Table 6 of the Appendix.
* The paper is well-written and easy to understand.
summary: This paper presents a strategy to avoid spikes in loss during the training of LMs by keeping the gradient norm small. To manage the upper limit of gradient norms effectively, the method involves (i) using small initial values for sub-layers, and (ii) maintaining the standard deviation of embeddings around 1.
weaknesses: * Although the theoretical analysis is intriguing, I question the practical value of this work, as most practices described in Section 4 are already in use. Utilizing small values for initialization to ensure stable training is well-known, and both Scaled Embed and Embed LM have been introduced in prior literature. If this work could offer a novel, advanced method for embedding normalization, it might receive more interest from the community.
* The activation function F was assumed to be either an identity function or ReLU, as stated on line 152 of page 3. What would be the results if widely used activation functions in recent LLMs, such as SiLU and SwiGLU, were applied?
* I am curious about how loss spikes impact the performance of downstream tasks on LLM leaderboards, beyond just affecting perplexity. Are these spikes also harmful to the accuracy of downstream tasks?
* I believe it would be beneficial to conduct a theoretical analysis of the relationship between learning rate, loss spikes, and model sizes. This suggestion stems from the observation that the learning rates causing loss spikes differ according to model size.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: Where approaches evaluated on benchmarks post-training?
soundness: 3
strengths: The paper conducts mathematical analysis to demonstrate the requisite terms they later leverage.
Paper is clear and provides actionable results.
summary: The authors provide theoretical support to reduce the likelihood of divergence based on stabilization via small sub-layer initialization and large shortcut values. They then propose two techniques to effectively support this assumption and evaluate accordingly.
weaknesses: The authors only tested on smaller models, it is well established that most instability problems happen with larger models (>100B parameters). It would be beneficial to evaluate the loss curves on larger models or more diverse datasets 

Although the focus of this paper was to stabilize training, they underperform on loss-curves compared to vanilla approaches to disprove Le Scao et al's findings. This is hypothesized to be related to learning-rates - which is demonstrated by looking at a absolute min score. However,  by only exploring LR adjustments on smaller models, it isn't immediately clear that the proposed method is consistently better as sub-optimal lr's are more stable, a lr scheduler could account for differences in the long-term.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: The relationship between spikes and poor performance is not clearly established.

Some loss spikes can be recovered during training, while others cannot.

Are you assuming that spikes are (1) always bad, (2) risks that result in divergence in some cases, or (3) sometimes acceptable and not harmful?

L40: What is catastrophic divergence? Does it mean the spike never goes down again?

What model is used in Figure 1?

What is the difference between spikes in LLMs and other smaller models?

L47: The assumption between spikes and gradient norm is unclear. What did you observe in L121?

L53: What is the standard deviation of embedding means? Does this equate to a large shortcut in the later sections?

L63: Evaluation is not explained.

L83: Conventionally, does shortcut mean residual?

Can you comment on the parallel FF+Attn setting where there is no intermediate vector? It is broadly used in many models, e.g., Pythia, Mesh-Transformers, and PaLM.

clear writing in section 2.1, 2.2

L130: What is the difference between W_* and W?

L134: You should mention the single-head attention and F as the identity function (linear settings) in the abstract and introduction.

L134: Please plot the distribution of x, x', and W, and note the standard deviation and mean of each layer to support your assumption. In Appendix F, there are only plots for initialized models. I could not find plots for W or during training and pretrained models.

L160: What are d and d_ffn? Did you mention them somewhere?

L167: This is an estimation as shown in L159. Please be consistent.

L174: It's unclear what this well-known formula is.

Eq 13: Why is variance the degree of freedom?

Condition 2 of large shortcut comes from assumption 1.

L219, Eq 19: Where does the 2 come from?

L234: The writing in Section 4 could follow the format in Section 5.2 to improve readability.

L250: Both "Scaled Embed" and "Embed LN" are existing approaches, not introduced by this paper.

L291: Why should it be close to 1 and not as large as possible? If so, why is it called 'large shortcut' in the previous sections? Similarly, should the sublayer be as small as possible or close to some value? Why do we scale the embed by some value but not a larger one? What is the optimal scale from your theory?

L309: C4 is too small to be the pretraining data for LLM.

L317: Why not use a common public architecture (e.g., LLama, Mixtral, Gemma) for reproducibility? Is there any common pretraining baseline in the literature (e.g., Dettmers et al. (2022) and Le Scao et al. (2022))?

Figure 3 (a): Lines are overlapped and cannot see what happened. Is there any other metric like the number of spikes other than visualization?

L367: Only in the embed layer?

The training risk mentioned in the abstract is not addressed in the experiment. We can only see some marginal improvement over perplexity. In the abstract, it indicates that 'if we don't take spikes seriously, the whole effort will be in vain due to catastrophic divergence.' However, both vanilla and embed detach do not suffer from catastrophic divergence in the experiment results.

Section 6.1: How does the learning rate affect the gradient norm? How does the learning rate affect the four baseline methods? Did you indicate that stabilization can be achieved by a smaller learning rate? If so, there is no actual need for "Scaled Embed" or "Embed LN"? In Figure 5, we can see that vanilla performs better than Scaled Embed. To achieve better final performance as depicted in Table 3, can we use a small learning rate to travel through the risky early training stage and increase the learning rate later to avoid spiking?

Section 6.2: Good point! I expect a similar discussion in Section 6.1 to explain why a smaller learning rate can reduce spikes. Please remind me if I missed something. Also, this indicates the settings of 'short seq in the early stage and long seq later.' Can we say that 'small lr in the early stage and large one later' is also possible?

L490: Can you use your theory to explain why preLN is more stable than post-LN?

L509: The efficiency and learning rate discussion is a bit too far and not quite relevant.
soundness: 2
strengths: The paper effectively addresses the critical issue of loss spikes during training, providing a detailed analysis of their relationship with gradient norms and embedding means. It evaluates existing approaches like "Scaled Embed" and "Embed LN," discussing their effectiveness in mitigating spikes. Additionally, the paper offers valuable insights into the impact of learning rate adjustments on model stability and compares the behavior of spikes in large language models (LLMs) with smaller models, providing a broader context for its findings.
summary: This paper analyzes gradient norms to explain and reduce the spikes of training loss. However, this paper does not address the differences between LLM and small models, differences between deep and shallow layer, and the relationship between spikes and gradient norm. Most importantly, the risk of training LLM is simplified by evaluating the performance on benchmark datasets while the risk of spending computing resources is not well addressed unfortunately requires lots of improvements.
weaknesses: The paper suffers from an unclear relationship between spikes and poor performance, with insufficient explanations of key terms and assumptions. The evaluation section is not well-explained, and there are inconsistencies in terminology. Additionally, it lacks necessary plots and data to support its assumptions, and some figures are difficult to interpret due to overlapping lines. Reproducibility is a concern as common public architectures are not used, and some discussions are considered irrelevant to the main topic.

Meta-review: The paper analyses the important question of divergence in optimizing the parameters of large neural language models. The authors investigate the high gradient norms as the cause for loss spikes, and discuss two mitigation strategies in terms of having small sub-layers and large shortcuts. 

While the considered problem is of great interest, the paper has a series of shortcomings. The reviewers raised questions on the limited analysis of the work in understanding the effect between spikes and poor model performance. Furthermore, some sections of the paper are not clear enough, in particular in the evaluation section.

Reviewers expressed concerns about the significance of the findings. In particular, there are questions on whether the gains of the proposed remedies occur only when the model diverges. In that context, authors are advised to investigate the benefit of the small sub-layers and large shortcuts also in cases when the optimization hyper-parameters are carefully tuned to avoid a divergence.

Considering the mentioned limitations I recommend rejection, but simultaneously encourage reviewers to consider the provided comments in improving the paper.

Decision: Reject