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

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.

We suppose that the main concerns of the reviewer are as follows:

• The need to clarify the harms caused by loss spikes to establish the advantages of preventing them.
• A suspicion that we used an uncommon architecture, leading the reviewer to question the practicality of our study.

Regarding the first concern, we demonstrated in our experiments that models without loss spikes achieved better performance. For instance, we extracted the perplexities of the 1.7B parameter models from Table 2:

As shown in the table, Embed LN and Scaled Embed, which satisfy the conditions for preventing the loss spike, outperformed Vanilla and Embed Detach, which violate the conditions.

Furthermore, Figures 1, 5 (d), and 9 (a) show that Vanilla failed during pre-training due to catastrophic divergence. We kindly ask the reviewer to consider these empirical results alongside the training costs described in our response to the summary. If the reviewer still believes that certain discussions are unrelated to the main topic, we would greatly appreciate it if the reviewer could specify the points in question.

Regarding the second concern, we used the Pre-LN Transformer architecture, which is commonly employed in various LLMs, such as the Llama series. Therefore, we believe that our findings are practically useful. In addition, as described in Reproducibility Statement (lines 546--548), we add only several lines to a widely used implementation, i.e., Megatron-LM. Therefore, we also believe that our experimental findings are straightforward to reproduce. We would greatly appreciate it if the reviewer could point out any essential differences between the models.

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

The essential contribution of our paper, the theoretical analysis of stability during pre-training, does not depend on the size of the training data. Additionally, we would appreciate it if the reviewer could consider the computational costs required to train each model. We would like to re-emphasize that constructing each 1.7B parameter model costs approximately $5,120.

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))?

As mentioned in the response to the weaknesses, our focus is on the Pre-LN Transformer architecture, which is a common architecture for LLMs.

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

We consider the number of spikes to be a subjective indicator, as it can be defined in a way that is unfavorable to Vanilla. Furthermore, not only Figure 3 (a) but also Figures 3 (b), 5, and 7 support our theoretical analysis. We would appreciate it if the reviewer confirms these figures.

L367: Only in the embed layer?

Yes, Embed LN applies layer normalization to embeddings, as per our theoretical analysis. We would appreciate it if the reviewer could provide more details regarding this comment for further discussion.

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?

Figures 1, 5 (d), and 9 (a) clearly show that the model violating the conditions we provided, Vanilla, diverged during pre-training. Moreover, we would like to emphasize that while a small learning rate stabilizes pre-training, it significantly degrades performance, as shown in Table 3.

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?

We do not think so. Section 6.1 discusses the relationship between our theoretical analysis and Li et al., 22. Specifically, we provided a theoretical justification for the strategy of using a short sequence to stabilize training and empirically demonstrated stability when varying the input sequence length. However, we did not address whether it is necessary to extend the input sequence length during pre-training, as proposed by Li et al., 22, because the purpose of our study is not to re-examine their work. Furthermore, we are unable to link our findings to the strategy of varying the learning rate during pre-training.

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

No, the difference between Pre-LN and Post-LN is beyond the scope of our study. As stated in lines 490--492, Xiong et al., 2020 proved that the Pre-LN Transformer is more stable than the Post-LN Transformer.

L509: The efficiency and learning rate discussion is a bit too far and not quite relevant.

As described in the response to the weaknesses, our study is useful to construct a better model in the limited budget while avoiding the failure. Therefore, we believe that our study can be regarded as efficient pre-training of LLMs.

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

The line 134 states, "This assumption is valid when we initialize parameters with the normal distribution, the number of heads in Equation (4) is 1, and $F$ is an identity function" describing the conditions under which our assumption holds. However, we did not restrict our analysis to single-head attention or $F$ as the identity function. In Section 3.2, we discussed multi-head attention, and Appendix G presents the case where ReLU is used as $F$.

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.

Figure 12 in Appendix F shows the distributions of $x$ and $x'$, both of which follow a normal distribution. Each $W$ was initialized with a normal distribution as described in Sections 4 and 5, and therefore, they clearly follow a normal distribution without requiring confirmation through plotting.

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

For $d$ and $d_{\mathrm{ffn}}$, we describe them in line 76 and lines 92--93, respectively, as follows: "$d$ denotes the dimension of the layer", and "$d_{\mathrm{ffn}}$ and $d_{\mathrm{head}}$ are the internal dimensions of FFN and multi-head self-attention sub-layers".

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

We would appreciate it if the reviewer found it. We would like to replace = with $\simeq$.

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

As described in lines 174–176, we provided the equation that transforms the variable x’ which follows a normal distribution, into the variable z, which follows the standard normal distribution: $z = (x' - \mu_{x'}) / \sigma_{x'}$.

Eq 13: Why is variance the degree of freedom?

We apologize, but we are still unable to understand the intent of this question. We would appreciate it if the reviewer could provide more details. Since we interpret this comment as a question about the degree of freedom, we would like to clarify. As stated in our paper, $z_i z_i$ follows a $\chi^2$ distribution with 1 degree of freedom, and the variance of such a distribution is equal to 2. This is a property of the $\chi^2(1)$ distribution, and the statement aligns with this property. For the case of $z_i z_j (i \neq j)$, the variance is equal to 1, as it is the product of two independent variables that follow the standard normal distribution.

Condition 2 of large shortcut comes from assumption 1.

Is this sentence intended to be a question? We are unable to understand the reasoning of the reviewer.

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

We applied the description of lines 158--159 to $||W_O||_2$, and thus: $||W_O||_2 \simeq \sigma_O (\sqrt{d} + \sqrt{d}) = \sigma_O (2\sqrt{d})$.

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

We thank the reviewer for their feedback, but we are confused as we are unable to understand the intent. We would appreciate it if the reviewer could provide more details.

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

As mentioned in line 543, "We do not aim to propose a novel method in this paper", we did not claim to propose a novel method. Moreover, we would like to emphasize that Scaled Embed and Embed LN were proposed in previous studies; however, they are not standard practices in the pre-training of LLMs.

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?

As mentioned in the response to the summary, our motivation is to construct a model that performs as well as possible while avoiding the failure of pre-training, as the pre-training of LLMs requires an enormous amount of computational resources. Appendix E demonstrates that setting the standard deviation of the embeddings to a value larger than 1 degrades the performance. Therefore, we do not recommend using a significantly larger value.

Then, we would like to address the questions, but we would greatly appreciate it if the reviewer could provide more details for certain questions to facilitate a fruitful discussion.

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?

We assume (1). In fact, our experimental results indicate that loss spikes negatively impact performance, even if a model recovers from them. For example, in the 1.7B parameter models, Vanilla and Embed Detach, which frequently encounter loss spikes, underperformed compared to Scaled Embed and Embed LN, as shown in the above table.

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

Yes, we used it as the meaning that the spike never goes down again, as illustrated in Vanilla as Figure 1.

What model is used in Figure 1?

As mentioned in the caption of Figure 1, we used 1.7B parameter models to illustrate the figure. The learning rate is $1.0 \times 10^{-3}$. In short, the configuration is identical to that of Figure 5 (d). However, since Figures 1 and 5 (d) show the loss curves on the training data and validation data, respectively, the two figures are not identical.

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

If "other smaller models" refers to Transformer-based decoder-only models, we believe our findings are essentially applicable. However, we would need more details about the "other smaller models" for further discussion.

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

We explained in lines 119--120 that "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". This correlation is demonstrated in the pairs of Figures 3 and 4, as well as Figures 7 and 8.

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

For the former question, we used "embedding" to refer to the vector representation of an input subword. Thus, "the standard deviation of embedding" refers to the standard deviation of this vector representation. For the latter question, yes, we explained in lines 291–292 that "(2) large shortcut; making the standard deviation of each embedding close to 1."

L63: Evaluation is not explained.

We interpret this comment as a critique that the Introduction lacks details of experiments. However, we would appreciate it if the reviewer could confirm Sections 5 and 6, which describe our experiments in detail.

L83: Conventionally, does shortcut mean residual?

No. In the residual block, $x + f(x)$, $x$ refers to the shortcut, and $f(x)$ refers to the residual branch. We kindly ask the reviewer to refer to the footnote 3 of Liu et al., 20 as an example.

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.

Our theoretical conclusion remains valid in the parallel FFN+Attn setting. We can easily derive an upper bound as follows:

$

|| J_{\mathrm{FFN}+\mathrm{Attn}}||_2 = \Biggl| \Biggl| \frac{\partial (x + \mathrm{FFN}(\mathrm{LN}(x)) + \mathrm{Attn}(\mathrm{LN}(x)))}{\partial x} \Biggr| \Biggr|_2 = \Biggl| \Biggl| I + \frac{\partial \mathrm{FFN}(\mathrm{LN}(x))}{\partial x} + \frac{\partial\mathrm{Attn}(\mathrm{LN}(x))}{\partial x} \Biggr| \Biggr|_2 \leq 1 + \Biggl| \Biggl| \frac{\partial \mathrm{FFN}(\mathrm{LN}(x))}{\partial x} \Biggr| \Biggr|_2+ \Biggl| \Biggl| \frac{\partial\mathrm{Attn}(\mathrm{LN}(x))}{\partial x} \Biggr| \Biggr|_2

$

using the same procedure for Equations (8) and (9). For the upper bounds for $\biggl| \biggl| \frac{\partial \mathrm{FFN}(\mathrm{LN}(x))}{\partial x} \biggr| \biggr|_2$ and $\biggl| \biggl| \frac{\partial\mathrm{Attn}(\mathrm{LN}(x))}{\partial x} \biggr| \biggr|_2$, we had estimated them with Equations (15) and (20), respectively.

clear writing in section 2.1, 2.2 L130: What is the difference between W_* and W?

We interpret this comment as asking whether the $W$ in line 131 is a typo of $W_{*}$. Yes, we thank the reviewer for identifying this typo.

We thank the reviewer for reading our paper and feedback. Before addressing each question, we would like to explain the motivation behind our study and address main concerns as follows to clarify our contributions:

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.

Based on the weaknesses and questions raised in the review, we believe the reviewer may not have fully recognized the computational cost of LLM pre-training. Therefore, we would first like to explain the associated costs. As an example, we focus on the 1.7B parameter model used in our experiments. Training each model in our configuration required approximately 1,000 A100 80GB GPU hours. Given the AWS EC2 pricing of $40.96 per hour with 8 A100 80GB GPUs (as described in: https://aws.amazon.com/ec2/instance-types/p4/), the cost of training each model amounts to approximately$5,120.

Considering this substantial expense, it is crucial to construct a model that performs as well as possible while avoiding the failure of pre-training. This motivation underpins our approach. In fact, the worst-case scenario is a failure of pre-training due to the loss spike, as shown in Vanilla in Figures 1, 5 (d), and 9 (a). Such failures result in wasted costs.

In contrast, models that satisfy our provided conditions successfully completed pre-training using the same hyper-parameters.

We thank the reviewer for reading our paper and feedback. We would like to address the weaknesses and questions as follows:

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

We theoretically demonstrated that larger models are more unstable, as the upper bounds of the gradient norms for each sub-layer, shown in Equations (15) and (20), depend on the dimension size.

On the other hand, we would like to emphasize that it is impractical to conduct experiments with extremely large models, such as those exceeding 100B parameters, as suggested by the reviewer, due to the tremendous budget requirements for such experiments. For example, at least $25 million USD would be required to train two 100B parameter models (a baseline and a stable method) with 2T tokens, that is the compute-optimal training data size suggested by Hoffmann et al., 22. We estimated this cost as follows:

• Training a single 100B parameter model requires approximately 2,457,600 GPU hours on A100 80GB GPU, based on the Llama2 paper (https://arxiv.org/abs/2307.09288).
• The AWS EC2 cost, for instance, with 8 A100 80GB GPUs, is $40.96 per hour (source: https://aws.amazon.com/ec2/instance-types/p4/).

Therefore, the estimated cost to train two 100-billion-parameter models (baseline and stable method) is: $2 * 2457600 * 40.96 / 8 \simeq$ $25.2 million USD. This detailed estimation underscores the substantial cost involved in training such large models.

In addition, as described in Appendix C, we conducted experiments with 13B parameter models. This configuration represents the large model in studies focused on the stability of LLM pre-training. For example, Wortsman et al., 24 conducted their experiments with models under 10B parameters.

Moreover, we would like to emphasize that our study provides more than just empirical findings; the essential contribution is the theoretical approach to prevent the loss spike problem. Therefore, we believe that the experiments with models exceeding 100B parameters are not indispensable. We would appreciate it if the reviewer takes into account that most researchers lacks access to large-scale computing resources, as noted in the review guidelines of a recent conference: https://colmweb.org/ReviewGuide.html.

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.

We have read this comment multiple times, but we are still unable to understand its main point. Therefore, we will respond based on the assumption that the reviewer is raising a question about the learning rate scheduler.

We used the cosine learning rate scheduler, which is a standard choice in the pre-training of LLMs, as demonstrated in the Llama series (https://arxiv.org/abs/2307.09288) and Le Scao et al., 2022. In addition, we tested several schedulers in our preliminary experiments, and the Vanilla achieved the best performance with the cosine scheduler. If the reviewer has a suggestion for a more sophisticated scheduler, we would greatly appreciate it if the reviewer could share the idea with us.

Where approaches evaluated on benchmarks post-training?

This paper focuses on stabilizing the pre-training phase, and as such, post-training is outside its scope. However, as shown in Appendix E of Fedus et al., 2021, negative log-likelihood correlates with the performance of downstream tasks. Therefore, we believe that a model achieving better negative log-likelihood during pre-training, such as Scaled Embed in our experiments, is also likely to achieve better performance after post-training.

We thank the reviewer for reading our paper and feedback. We would like to address the weaknesses as follows:

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.

We appreciate the reviewer's recognition of the value of our theoretical analysis of the loss spike problem. We would like to emphasize that this paper does not aim to propose a method as stated on line 543 of our paper; rather, it aims to understand why the loss spike problem occurs during the pre-training of LLMs and to provide conditions to prevent it. As described in "Novelty is in the Eye of the Beholder" (https://drive.google.com/file/d/1ydN247sEXjnP0P_JByf287ifXNcdoIBM/view), introduced in the ICLR 2022 Reviewer Tutorial (https://icml.cc/Conferences/2022/ReviewerTutorial), we believe that the novelty of a research paper is not limited to proposing “a novel method” but also includes offering insights through analysis and understanding, as our paper does.

In addition, we would like to emphasize that using Scaled Embed or Embed LN is not currently standard practice in the pre-training of LLMs. This can be verified by checking widely used open-source implementations for LLM pre-training, such as Megatron-LM and LLMFoundry; neither of which includes Scaled Embed or Embed LN as part of their functionalities. This fact indicates that current LLM pre-training processes are generally conducted under conditions that conflict with the requirements we introduced to prevent the loss spike problem: (1) initializing the parameters of sub-layers with small values, and (2) ensuring the standard deviation of each embedding is close to 1. As a result, the publicly available pre-trained models may potentially suffer from reduced performance and pre-training efficiency due to, for example, the use of smaller-than-ideal learning rates aimed at avoiding loss spikes. Moreover, to show that the conditions we introduced for preventing loss spikes can be met easily without requiring major changes to the standard Transformer architecture, we intentionally used the simplest and most conventional methods available. The fact that loss spikes can be prevented with such simple modifications is expected to encourage many researchers/developers conducting LLM pre-training to implement our conditions, potentially resulting in significant practical and societal impact. This is why we chose straightforward methods. Once again, our contribution lies not in proposing novel methods but in providing a theoretical justification for previously unresolved observations regarding why loss spikes occur, along with simple solutions to prevent them.

Thus, we believe our theoretical analysis of the loss spike problem, along with the empirical validation, offers valuable insights for the research 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?

We are planning to work on the proofs for other activation functions as well. While setting aside the detailed proofs for the upper bounds, we believe the overall theory will likely still hold.

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?

As shown in Appendix E of Fedus et al., 2021, perplexity (i.e., negative log-likelihood) strongly correlates with the performance of downstream tasks. Consequently, when a loss spike degrades perplexity, it also reduces the accuracy of downstream tasks. Additionally, as demonstrated in our paper, loss spikes can sometimes ruin the pre-training process. Therefore, addressing the loss spike problem is crucial.

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

We have already demonstrated, from a theoretical perspective, that larger models are more unstable because the upper bounds of the gradient norms for each sub-layer, as shown in Equations (15) and (20), depend on the dimensionality. Due to this property, we consider it necessary to set a smaller learning rate to prevent the growth of each parameter in larger models. However, we would like to emphasize that we validated our theoretical findings for stabilizing the pre-training of LLMs through various experiments.

Thank you for the time and effort the authors have invested in this rebuttal. I appreciate the additional explanation that 'using Scaled Embed or Embed LN is not currently standard practice in the pre-training of LLMs,' with detailed examples such as Megatron-LM and LLMFoundry. I also agree with the point that 'the novelty of a research paper is not limited to proposing a novel method, but also includes offering insights through analysis and understanding.' Therefore, I am increasing my score from 3 to 5, as I believe this paper delivers a meaningful message for achieving effective LLM training.

However, I am still concerned about the lack of investigation into other activation functions and the somewhat limited scope of the experiments. It is now common practice to present results not only for PPL but also for downstream tasks. Although the authors provide a supporting reference, I am unsure if the differences in PPL between the models of this work are significant enough to impact downstream task performance. Considering that this paper focuses on small-scale models, conducting additional experiments on other benchmarks with further fine-tuning (using cost-effective techniques like LoRA) could enhance the experimental validity of the approach. Alternatively, the authors could present downstream task scores of the pretrained models without further training.

Additionally, concerning the remark, 'we consider it necessary to set a smaller learning rate to prevent the growth of each parameter in larger models,' it would be helpful to specify what range is considered "small" for learning rates. Detailed analysis or guidance on the relationship between learning rate and loss spikes would also be beneficial.

I feel that many recent papers with theoretical contributions are also supported by sufficient experimental evidence. However, it remains unclear to me if this paper meets such criteria.

We greatly appreciate your positive feedback and are delighted to hear that our findings are useful for a text-to-speech model!