Initialization of Large Language Models via Reparameterization to Mitigate Loss Spikes (2024)

Abstract: Loss spikes, a phenomenon in which the loss value diverges suddenly, is a fundamental issue in the pre-training of large language models. This paper supposes that the non-uniformity of the norm of the parameters is one of the causes of loss spikes. Here, in training of neural networks, the scale of the gradients is required to be kept constant throughout the layers to avoid the vanishing and exploding gradients problem. However, to meet these requirements in the Transformer model, the norm of the model parameters must be non-uniform, and thus, parameters whose norm is smaller are more sensitive to the parameter update. To address this issue, we propose a novel technique, weight scaling as reparameterization (WeSaR). WeSaR introduces a gate parameter per parameter matrix and adjusts it to the value satisfying the requirements. Because of the gate parameter, WeSaR sets the norm of the original parameters uniformly, which results in stable training. Experimental results with the Transformer decoders consisting of 130 million, 1.3 billion, and 13 billion parameters showed that WeSaR stabilizes and accelerates training and that it outperformed compared methods including popular initialization methods.

## introduction transformer-based large language models (llms) have attracted remarkable attention @xcite . the discovery of a scaling-law @xcite has been driving the model and corpus sizes ever larger, causing huge computational costs for pre-training. during pre-training of llms, the loss value often diverges suddenly @xcite , as illustrated at the top of figure 1 . this phenomenon, known as loss spikes, is a fundamental issue in the llm pre-training because it not only increases the final loss value, but also causes the pre-training to fail if the loss diverges completely. here, let ∆w be the update of the parameter w at an optimization step. ∥∆w ∥/∥w ∥ represents the magnitude of the parameter update relative to the parameter itself, and we call it the update ratio. the bottom of figure 1 shows the update ratios. we consider that different scales of update ratios among parameter matrices can lead to unstable training. indeed, before the loss spike, the update ratio of w d is larger than that of w u . that is, w d undergoes a more pronounced change. after the spike, the difference between the update ratios decreases. this observation motivated us to regulate the update ratios in the model in a certain range. we consider that uneven and large update ratios are due to non-uniformity of the norm of the parameters. with the current initialization methods, w d is set smaller than other parameters, which is required to avoid the vanishing and exploding gra-dients problem. consequently, by definition, the update ratio of w d tends to be larger. to address this issue, we propose a novel technique, called weight scaling as reparameterization @xcite . wesar introduces a gate parameter α ∈ r for each parameter matrix w and uses αw instead of w inside the model. wesar relieves the parameter w of non-uniformity by adjusting α to the values required to avoid the vanishing and exploding gradients problem. moreover, wesar enables an arbitrary small common standard deviation to set be for all parameters, which results in not only stable, but also accelerated, training. we conducted pre-training of transformer decoders consisting of 130 million (13m), 1.3 billion (1.3b), and 13b parameters. our experimental results show that wesar stabilized and accelerated their training due to the stable and equal-scale update ratios, as shown in figure 1 . we also confirmed that wesar outperformed compared methods, including a initialization method widely used for pre-training llms @xcite and the existing reparameterization methods @xcite @xcite . our contributions can be summarized as follows: ## preliminaries we consider transformer models @xcite consisting of the following layers: an embedding layer with w e , self-attention layers with w q , w k , w v , and w o (query, key, value, and output projections), feed-forward layers with w u and w d (up and down projections) 1 , and a prediction 1 we did not use glu @xcite for simplicity. in this section, we first review the backpropagation algorithm @xcite . then, we describe the initialization strategies of the transformer models to avoid the vanishing and exploding gradients problem. ## existing methods and their problems here, we review two of the existing initialization methods and their problems. the methods are summarized in ## proposed method we propose wesar as a way to meet the two conflicting aforementioned requirements: (i) the criteria of any initialization method designed to avoid the vanishing and exploding gradients problem, as discussed in §2.2, and (ii) the common scales of all parameters to keep stable and uniform update ratios for mitigating loss spikes, as discussed in §3.3. in addition to stabilizing the training, wesar enables a hyperparameter setting that achieves a rapid decrease in loss. 2 we approximate the gain of the activation function used in the feed-forward layer to that of relu (i.e., √ 2). ## conclusion loss spikes are a fundamental issue in pre-training of llms because they increase the pre-training cost and degrade the performance of the model. to address this problem, we identified one of the causes as the non-uniformity of the norm of the model parameters. we proposed a novel reparameterization method, wesar, that addresses the non-uniformity problem by adjusting the gate parameter to the required scale and initializing the actual parameters with a common standard deviation. wesar not only stabilizes the pre-training, but also accelerates the pre-training by setting a standard deviation smaller than in the conventional setting. experimental results showed that wesar outperformed the compared methods, and the parameters and their update ratios were stable during pre-training. the use of llms has been spreading. we believe this study to be a significant contribution that increases both the efficiency of the llm's pretraining and the effectiveness of the pre-trained llms.