Discussion on Model Optimization: Why is BERT's Initial Standard Deviation 0.02?

What Does Vanishing Gradient Mean?

In the article "Also on the Vanishing/Exploding Gradient Problem of RNN", we discussed the vanishing gradient problem in RNNs. In fact, the phenomenon of vanishing gradients in general models is similar: it refers to the fact that (primarily in the initial stages of the model) the closer a layer is to the input, the smaller its gradient becomes, tending toward zero or even equaling zero. Since we primarily use gradient-based optimizers, vanishing gradients mean we lack good signals to adjust and optimize the earlier layers.

In other words, the earlier layers may receive almost no updates and remain in a state of random initialization; only the layers closer to the output are updated well. However, the inputs to these layers are the outputs of the unoptimized earlier layers, so the input quality may be terrible (having passed through a near-random transformation). Therefore, even if the later layers are updated well, the overall effect is poor. Ultimately, we observe a counter-intuitive phenomenon: the deeper the model, the worse the performance, even on the training set.

A standard method to solve vanishing gradients is residual connections, formally proposed in ResNet. The idea of residuals is very simple and direct: if you are worried that the input gradient will vanish, why not just add a term with a constant gradient? Simply, the model becomes: y = x + F(x) In this way, because there is a "shortcut" path x, even if the gradient of x within F(x) vanishes, the gradient of x itself can basically be preserved, allowing deep models to be trained effectively.

Does LN Really Alleviate Vanishing Gradients?

However, in BERT and the original Transformer, a Post-Norm design is used, where the Norm operation is added after the residual: x_{t+1} = \text{Norm}(x_t + F_t(x_t)) Actually, the specific Norm method is not very important; whether it is Batch Norm or Layer Norm, the conclusion is similar. In the article "A Brief Discussion on Initialization, Parameterization, and Normalization of Transformer", we analyzed this Norm structure; let’s repeat it here.

In the initialization phase, since all parameters are randomly initialized, we can assume that x and F(x) are two independent random vectors. If we assume their respective variances are 1, then the variance of x + F(x) is 2. The \text{Norm} operation is responsible for resetting the variance to 1. Thus, in the initialization phase, the \text{Norm} operation is equivalent to "dividing by \sqrt{2}": x_{t+1} = \frac{x_t + F_t(x_t)}{\sqrt{2}} Recursively, this becomes: \begin{aligned} x_l =&\, \frac{x_{l-1}}{\sqrt{2}} + \frac{F_{l-1}(x_{l-1})}{\sqrt{2}} \\ =&\, \frac{x_{l-2}}{2} + \frac{F_{l-2}(x_{l-2})}{2} + \frac{F_{l-1}(x_{l-1})}{\sqrt{2}} \\ =&\, \cdots \\ =&\, \frac{x_0}{2^{l/2}} + \frac{F_0(x_0)}{2^{l/2}} + \frac{F_1(x_1)}{2^{(l-1)/2}} + \frac{F_2(x_2)}{2^{(l-2)/2}} + \cdots + \frac{F_{l-1}(x_{l-1})}{2^{1/2}} \end{aligned} We know that residuals help solve vanishing gradients, but in Post-Norm, the residual path is severely weakened. The closer to the input, the more severe the weakening; the residual exists "in name only." Therefore, in a Post-Norm BERT model, LN not only fails to alleviate vanishing gradients, but it is also one of the "culprits" behind them.

Then Why Do We Still Add LN?

The question naturally arises: since LN exacerbates vanishing gradients, wouldn’t it be better to just remove it?

It can be removed, but as mentioned before, the variance of x + F(x) is 2, and the more residuals there are, the larger the variance becomes. So a Norm operation is still necessary. We can add it to the input of each module, i.e., x + F(\text{Norm}(x)), and then add a final \text{Norm} at the total output. This is the Pre-Norm structure, where each residual branch has equal weight, rather than the exponential decay trend seen in Post-Norm. Of course, there are also methods that do not use Norm at all, but they require special initialization for F(x) to make its initial output closer to 0, such as ReZero, Skip Init, Fixup, etc., which were also introduced in "A Brief Discussion on Initialization, Parameterization, and Normalization of Transformer".

However, putting these improvements aside, does Post-Norm have no merit? Did Transformer and BERT start with a completely failed design?

Clearly, that is unlikely. Although Post-Norm brings a certain degree of vanishing gradient issues, it also has benefits in other aspects. Most obviously, it stabilizes the numerical values of forward propagation and maintains consistency across modules. For example, with BERT-base, we can connect a Dense layer to the last layer for classification, or take the 6th layer and connect a Dense layer for classification. But if you use Pre-Norm, after taking an intermediate layer, you need to manually add an LN before the Dense layer; otherwise, the variance of later layers will be larger, which is unfavorable for optimization.

Secondly, vanishing gradients are not entirely a "bad thing"; in fact, for the Fine-tuning stage, they are an advantage. During Fine-tuning, we usually want to prioritize adjusting parameters closer to the output layer and avoid over-adjusting parameters closer to the input layer to prevent severely damaging the pre-training effects. Vanishing gradients mean that the closer a layer is to the input, the weaker its influence on the final output, which is exactly what is desired during Fine-tuning. Therefore, pre-trained Post-Norm models often have better Fine-tuning effects than Pre-Norm models, as we also mentioned in "RealFormer: Moving Residuals to the Attention Matrix".

Do We Really Worry About Vanishing Gradients?

In fact, the most critical reason is that with current adaptive optimization techniques, we are no longer very worried about the vanishing gradient problem.

This is because the mainstream optimizers in NLP today are Adam and its variants. For Adam, since it includes momentum and second-moment correction, its update amount is approximately: \Delta \theta = -\eta\frac{\mathbb{E}_t[g_t]}{\sqrt{\mathbb{E}_t[g_t^2]}} As we can see, the numerator and denominator have the same dimensions, so the result of the fraction is actually of the order \mathcal{O}(1), and the update amount is of the order \mathcal{O}(\eta). That is to say, theoretically, as long as the absolute value of the gradient is greater than the random error, the corresponding parameters will have a constant-order update amount. This is different from SGD; the update amount of SGD is proportional to the gradient. As long as the gradient is small, the update amount will also be small. If the gradient is too small, the parameters will hardly be updated.

Therefore, although the residuals in Post-Norm are severely weakened, in base and large-scale models, they are not weakened to the point of being smaller than random error. Thus, in conjunction with optimizers like Adam, they can still be updated effectively, making successful training possible. Of course, it is only a possibility; in fact, deeper Post-Norm models are indeed harder to train, requiring careful adjustment of learning rates and Warmup.

How Does Warmup Work?

You may have heard that Warmup is a key step in Transformer training; without it, the model may not converge or may converge to a poor position. Why is this? Didn’t we say that with Adam we don’t fear vanishing gradients?

It should be noted that Adam solves the problem of small parameter updates caused by vanishing gradients—meaning that regardless of whether gradients vanish, the update amount will not be too small. However, for models with a Post-Norm structure, vanishing gradients still exist, but their meaning has changed. According to the Taylor expansion: f(x+\Delta x) \approx f(x) + \langle\nabla_x f(x), \Delta x\rangle That is, the increment f(x+\Delta x) - f(x) is proportional to the gradient. In other words, the gradient measures the dependence of the output on the input. If the gradient vanishes, it means the model’s output’s dependence on the input has weakened.

Warmup involves gradually increasing the learning rate from 0 to a specified size at the beginning of training, rather than training from the specified size from the start. Without Warmup, the model starts learning rapidly from the beginning. Due to vanishing gradients, the model is more sensitive to later layers, meaning later layers learn faster. Since later layers take the outputs of earlier layers as input, and the earlier layers haven’t learned well at all, the later layers learn quickly but on a terrible input foundation.

Soon, the later layers reach a poor local optimum based on terrible inputs, at which point their learning slows down (because they believe they are near the optimum). Meanwhile, the gradient signal back-propagated to the earlier layers weakens further, leading to inaccurate gradients for the earlier layers. But as we said, Adam’s update amount is of a constant order; if the gradient is inaccurate but the update amount remains significant, it essentially becomes constant-order random noise. Consequently, the learning direction becomes unreasonable, the earlier outputs begin to collapse, and the later layers collapse along with them.

Therefore, if a Post-Norm model does not use Warmup, the phenomenon we often observe is: the loss quickly converges to near a constant, and after training for a while, the loss begins to diverge until it reaches NAN. If Warmup is used, the model is given enough time to "preheat." During this process, the learning speed of the later layers is primarily suppressed, giving the earlier layers more optimization time to promote synchronized optimization across all layers.

The premise of this discussion is vanishing gradients. If it were a Pre-Norm structure where there is no obvious vanishing gradient phenomenon, successful training is often possible without Warmup.

Why is the Initial Standard Deviation 0.02?

Students who like to pay attention to details will notice that BERT’s default initialization method is a truncated normal distribution with a standard deviation of 0.02. As mentioned in "A Brief Discussion on Initialization, Parameterization, and Normalization of Transformer", because it is a truncated normal distribution, the actual standard deviation is smaller, approximately 0.02/1.1368472 \approx 0.0176. Is this standard deviation large or small? For Xavier initialization, an n \times n matrix should be initialized with a variance of 1/n. For BERT-base, where n=768, the calculated standard deviation is 1/\sqrt{768} \approx 0.0361. This means that this initial standard deviation is significantly smaller, only about half of the common initialization standard deviation.

Why does BERT use a smaller standard deviation for initialization? In fact, this is still related to the Post-Norm design. A smaller standard deviation leads to overall smaller function outputs, making the Post-Norm design closer to an identity function in the initialization phase, which is more conducive to optimization. Specifically, according to the previous assumption, if the variance of x is 1 and the variance of F(x) is \sigma^2, then in the initialization phase, the \text{Norm} operation is equivalent to dividing by \sqrt{1+\sigma^2}. If \sigma is small, the weight of the "shortcut" in the residual is closer to 1, so the model is closer to an identity function in the initial stage, making it less prone to vanishing gradients.

As the saying goes, "We are not afraid of vanishing gradients, but we don’t want them either." Simply setting the initialization standard deviation smaller can make \sigma smaller, thereby alleviating vanishing gradients while maintaining Post-Norm. Why not do it? Could it be set even smaller or even to zero? Generally, an initialization that is too small will result in a loss of diversity and reduce the model’s trial-and-error space, which also brings negative effects. Overall, reducing it to 1/2 of the standard is a relatively reliable choice.

Of course, some people do like to push the limits. Recently, I saw an article attempting to use almost all-zero initialization for the entire model and achieved good training results. Interested readers can check it out: "ZerO Initialization: Initializing Residual Networks with only Zeros and Ones".

Why Add an Extra Dense Layer for MLM?

Finally, regarding a detail of the BERT MLM model: why does BERT add an extra Dense layer and LN layer before the MLM probability prediction? Is it okay not to add them?

The answers I’ve seen before generally suggest that layers closer to the output are more Task-Specific. By adding an extra Dense layer, we hope this layer is MLM-Specific, and since it is removed during downstream task fine-tuning, it doesn’t affect other tasks. This explanation seems somewhat reasonable but feels a bit like "metaphysics," as Task-Specificity is not easy to analyze quantitatively.

Here, I provide another more specific explanation. In fact, it is directly related to BERT’s use of a 0.02 standard deviation for initialization. As we just said, this initialization is on the smaller side. If we don’t add an extra Dense layer and directly multiply by the Embedding to predict the probability distribution, the resulting distribution will be too uniform (before Softmax, every logit is close to 0), so the model wants to amplify the values. Now the model has two choices: first, amplify the values of the Embedding layer, but updates to the Embedding layer are sparse, and amplifying them one by one is too troublesome; second, amplify the input. We know the last layer of the BERT encoder is an LN, and LN has a gamma parameter initialized to 1; the model can just amplify that parameter.

Model optimization uses gradient descent, and we know it chooses the fastest path. Obviously, the second choice is faster, so the model will prioritize the second path. This leads to a phenomenon: the gamma value of the last LN layer will be relatively large. If a Dense+LN is not added before predicting the MLM probability distribution, the gamma value of the LN in the last layer of the BERT encoder will be large, causing the variance of the last layer to be significantly larger than that of other layers, which is clearly not elegant. By adding a Dense+LN, the large gamma is transferred to the newly added LN, while each layer of the encoder maintains consistency.

In fact, readers can observe the gamma values of each LN layer in BERT themselves, and they will find that the gamma value of the last LN layer is indeed significantly larger, which verifies our hypothesis!

Conclusion

This article has attempted to answer several questions related to the model optimization of Transformers and BERT. Some are results I discovered in my own pre-training work, and some are intuitive imaginations combined with my own experience. In any case, consider this a shared reference answer. If there are any improprieties, please be patient and feel free to offer criticisms and corrections.

Reprinting please include the original address: https://kexue.fm/archives/8747

For more details on reprinting, please refer to: "Scientific Space FAQ"