Model Optimization Chat: Why Is BERT's Initialization Standard Deviation 0.02?
A few days ago there was a discussion in a group chat about the question "How does the Transformer solve the vanishing gradient problem?" Some answers mentioned residual connections, others mentioned LN (Layer Norm). Are these all correct answers? In fact, this is a rather interesting and multi-faceted question, and it turns out to be connected to quite a few model details, such as: "Why does BERT need warmup?", "Why is BERT's initialization standard deviation 0.02?", "Why does BERT add an extra Dense layer before making MLM predictions?", and so on. This post is a concentrated discussion of these questions.
What Does "Vanishing Gradient" Actually Mean?
In the post Also Discussing the Vanishing/Exploding Gradient Problem in RNNs], we discussed the vanishing gradient problem in RNNs. In fact, the vanishing gradient phenomenon in general models is similar: it refers to the fact that (mainly in the early stages of training) the layers closer to the input have smaller gradients, tending toward zero or even equal to zero. Since we mainly use gradient-based optimizers, vanishing gradients mean we lack a good signal for updating and optimizing the earlier layers.
In other words, the earlier layers might barely get updated at all, remaining essentially in their randomly initialized state; only the layers closer to the output get updated reasonably well. But the inputs to these later layers come from the outputs of the poorly-updated earlier layers, so the input quality may be quite bad (having passed through a nearly random transformation). As a result, even if the later layers are well-updated, the overall performance is still poor. In the end, we observe a very counter-intuitive phenomenon: the deeper the model, the worse the performance — even on the training set.
A standard method for addressing vanishing gradients is the residual connection, formally proposed in ResNet]. The idea behind residuals is very simple and direct: worried that the gradient of the input will vanish? Then just directly add a term with constant gradient! In the simplest case, the model becomes
\begin{equation}y = x + F(x)\end{equation}
This way, thanks to the extra "direct" path $x$, even if the gradient of $x$ inside $F(x)$ vanishes, the gradient of $x$ can still be largely preserved, allowing deep models to be trained effectively.
Does LN Really Alleviate Vanishing Gradients?
However, in BERT and the original Transformer, the design used is Post-Norm, which places the Norm operation after the residual addition:
\begin{equation}x_{t+1} = \text{Norm}(x_t + F_t(x_t))\end{equation}
The specific Norm method doesn't matter too much here — whether it's Batch Norm or Layer Norm, the conclusion is similar. In the post A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers], we already analyzed this Norm structure; let's go over it again here.
At initialization, since all parameters are randomly initialized, we can treat $x$ and $F(x)$ as two mutually independent random vectors. If we assume each has variance 1, then the variance of $x+F(x)$ is 2, and the $\text{Norm}$ operation is responsible for rescaling the variance back to 1. So during initialization, the $\text{Norm}$ operation is effectively equivalent to "dividing by $\sqrt{2}$":
\begin{equation}x_{t+1} = \frac{x_t + F_t(x_t)}{\sqrt{2}}\end{equation}
Recursing this gives
\begin{equation}\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}\end{equation}
We know that residual connections help with vanishing gradients, but in Post Norm, this residual channel is severely weakened — and the weakening gets worse the closer we get to the input, so the residual connection is "nominal but not real." So in the Post-Norm BERT model, LN not only fails to alleviate vanishing gradients, it is actually one of the "culprits" behind them.
So Why Do We Still Add LN?
Naturally, the next question arises: since LN makes vanishing gradients worse, why not just remove it?
We could remove it, but as mentioned above, the variance of $x+F(x)$ becomes 2, and the more residual connections there are, the larger the variance grows. So we still need some kind of Norm operation. We could instead place it at the input of each module, i.e., change it to $x+F(\text{Norm}(x))$, and just add one final $\text{Norm}$ at the very end of the whole output — this is the Pre-Norm structure. In this case each residual branch has equal weight, rather than the exponentially decaying trend seen in Post Norm. There are also designs that skip Norm entirely, but these require special initialization of $F(x)$ so that its initial output is closer to 0, such as ReZero, Skip Init, Fixup, etc. These have all been introduced in A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers].
But setting aside these improved variants, does Post Norm really have nothing going for it? Was the original Transformer/BERT design simply a complete failure from the start?
That seems unlikely. Although Post Norm does introduce some vanishing gradient issues, it also has benefits in other respects. The most obvious one is that it stabilizes the numerical scale of the forward pass and maintains consistency across modules. For example, with BERT base, we can attach a Dense layer on top of the last layer for classification, or equally attach it on top of layer 6 for classification. But with Pre Norm, if you take an intermediate layer's output, you need to add your own LN before the Dense layer; otherwise the variance grows larger for later layers, which is not favorable for optimization.
Second, vanishing gradients aren't purely "bad" — in fact, during fine-tuning, they can actually be beneficial. During fine-tuning, we usually want to preferentially adjust the parameters near the output layer, and avoid overly disturbing the parameters near the input layer, so as not to severely damage what was learned during pretraining. Vanishing gradients mean that the closer a layer is to the input, the weaker its influence on the final output — which is exactly what we want during fine-tuning. So a pretrained Post-Norm model often achieves better fine-tuning results than a Pre-Norm model, as we also mentioned in RealFormer: Moving the Residual onto the Attention Matrix].
Are We Really Worried About Vanishing Gradients?
The key reason, in fact, is that under the various adaptive optimization techniques used today, we no longer worry much about vanishing gradients.
This is because the mainstream optimizer in NLP today is Adam and its variants. For Adam, since it incorporates momentum and second-moment correction, its update, approximately speaking, is roughly
\begin{equation}\Delta \theta = -\eta\frac{\mathbb{E}_t[g_t]}{\sqrt{\mathbb{E}_t[g_t^2]}}\end{equation}
We can see that the numerator and denominator are on the same scale, so the ratio is essentially of the order of $\mathcal{O}(1)$, and the update size is on the order of $\mathcal{O}(\eta)$. In other words, in theory, as long as the gradient's absolute value is larger than random noise, the corresponding parameter will still get an update of constant order. This is different from SGD, where the update size is proportional to the gradient — if the gradient is small, the update will also be small, and if the gradient is too small, the parameter will barely be updated at all.
So, although the residual channel in Post Norm is severely weakened, at the base and large model scales it isn't weakened to the point of falling below the level of random noise. So, combined with optimizers like Adam, it can still receive effective updates, and thus training can potentially succeed. Of course, only "potentially" — in fact, the deeper the Post-Norm model, the harder it genuinely is to train, requiring careful tuning of the learning rate, warmup, and so on.
How Does Warmup Work?
You may have heard that warmup is a critical step in training Transformers — without it, training may fail to converge, or converge to a fairly poor point. Why is that? Didn't we just say that with Adam we no longer need to fear vanishing gradients?
Note that Adam addresses the problem of update magnitudes being too small due to vanishing gradients — that is, regardless of whether gradients vanish, the update size won't be too small. But for a Post-Norm model, vanishing gradients still exist; it's just that their meaning has changed. According to the Taylor expansion:
\begin{equation}f(x+\Delta x) \approx f(x) + \langle\nabla_x f(x), \Delta x\rangle\end{equation}
That is, the increment $f(x+\Delta x) - f(x)$ is proportional to the gradient — in other words, the gradient measures how much the output depends on the input. If the gradient vanishes, it means the model's output has become less dependent on the input.
Warmup gradually increases the learning rate from 0 up to the specified value at the start of training, rather than training with the specified learning rate right from the start. If we skip warmup, the model starts learning quickly right away. Because of vanishing gradients, the model is more sensitive to the later layers — meaning the later layers learn faster. But since these later layers take the earlier layers' outputs as their input, and the earlier layers haven't been trained well yet, the later layers — despite learning fast — are building on a poor foundation.
Very quickly, the later layers reach a poor local optimum based on this poor input, and at that point their learning starts to slow down (since they've arrived near what they consider optimal). At the same time, the gradient signal propagated back to the earlier layers becomes even weaker, causing the gradients of the earlier layers to become inaccurate. But as we said, Adam's update size is of constant order — so even with inaccurate gradients, the update magnitude remains constant order, meaning it could essentially become constant-scale random noise. The learning direction then becomes unreasonable, the earlier layers' outputs start to collapse, and this drags the later layers down with them.
So, if a Post-Norm model is trained without warmup, what we typically observe is: the loss quickly converges near some constant, and then after training for a while longer, the loss starts to diverge, eventually going to NaN. If warmup is used, the model is given enough time to "warm up." During this process, the learning speed of the later layers is mainly suppressed, while the earlier layers are given more time to optimize, promoting synchronized optimization across all layers.
The discussion here presupposes the existence of vanishing gradients. For structures like Pre Norm, where there's no obvious vanishing gradient phenomenon, training can often succeed without warmup as well.
Why Is the Initial Standard Deviation 0.02?
Readers who like to dig into details may have noticed that BERT's default initialization method is a truncated normal distribution with standard deviation 0.02. As we also mentioned in A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers], since it's a truncated normal distribution, the actual standard deviation is a bit smaller, roughly $0.02/1.1368472\approx 0.0176$. Is this standard deviation large or small? For Xavier initialization, a matrix of shape $n\times n$ should be initialized with variance $1/n$. With BERT base's $n$ equal to 768, the computed standard deviation would be $1/\sqrt{768}\approx 0.0361$. This means that BERT's initialization standard deviation is noticeably smaller — only about half of the common initialization standard deviation.
Why does BERT use such a smaller-than-usual standard deviation for initialization? In fact, this is again related to the Post-Norm design. A smaller standard deviation causes the function's output to be smaller overall, which makes the Post-Norm design closer to an identity function at initialization, and thus more favorable for optimization. Specifically, following our earlier assumption, if $x$ has variance 1 and $F(x)$ has variance $\sigma^2$, then at initialization the $\text{Norm}$ operation is effectively equivalent to dividing by $\sqrt{1+\sigma^2}$. If $\sigma$ is relatively small, the weight of the "direct" path in the residual gets closer to 1, meaning the model at initialization is closer to an identity function, and is thus less prone to vanishing gradients.
As the saying goes, "we're not afraid of vanishing gradients, but we still don't want them" — simply setting the initialization standard deviation a bit smaller makes $\sigma$ smaller too, alleviating the vanishing gradient problem somewhat while keeping the Post-Norm structure. Why not do it? So could we make it even smaller, or even all-zero? Generally speaking, initializing too small a scale sacrifices diversity, shrinking the model's space for trial and error, which also brings negative effects. All things considered, shrinking it to half the standard value is a fairly sensible choice.
Of course, there really are people who like to push things to the extreme. Recently the author came across a paper that attempted to initialize an entire model with almost all zeros, and still achieved decent results. Readers interested can check it out: ZerO Initialization: Initializing Residual Networks with only Zeros and Ones].
Why Does MLM Add an Extra Dense Layer?
Finally, there's a detail about BERT's MLM model: before predicting probabilities for MLM, BERT adds an extra Dense layer and LN layer. Why is that? Would it not work without them?
The explanation I've seen before is roughly this: the closer a layer is to the output, the more task-specified it is. By adding an extra Dense layer, we hope that this Dense layer becomes MLM-specified, so that during downstream fine-tuning, since it's no longer MLM-specified, it can simply be discarded. This explanation seems somewhat reasonable, but it feels a bit hand-wavy, since "task-specified" isn't something that's easy to analyze quantitatively.
Here I'll offer another, more concrete explanation, which is in fact still directly related to BERT's use of a 0.02 standard deviation for initialization. As mentioned earlier, this initialization is smaller than usual. If we didn't add an extra Dense layer and simply multiplied by the Embedding matrix to predict the probability distribution, the resulting distribution would be too uniform (before Softmax, every logit would be close to 0). So the model would want to scale up the numbers. Now the model has two options: first, scale up the values in the Embedding layer — but Embedding updates are sparse, so scaling them up one by one would be too cumbersome; second, scale up the input instead. We know that the final layer of the BERT encoder is an LN, and LN has a gamma parameter initialized to 1 — directly scaling up that parameter would do the trick.
Model optimization uses gradient descent, which we know will pick the fastest path — and clearly the second option is faster, so the model will preferentially go this route. This leads to a phenomenon: the gamma value of the final LN layer tends to become abnormally large. If we don't add a Dense+LN layer before predicting the MLM probability distribution, then the gamma of the last LN layer in the BERT encoder becomes disproportionately large, causing the variance of the last layer to be noticeably larger than that of other layers — which clearly isn't elegant. But by adding an extra Dense+LN, the inflated gamma gets shifted onto this new LN layer instead, and consistency is maintained across all the encoder's layers.
In fact, readers can go and observe the gamma values of each LN layer in BERT themselves, and they'll find that indeed the gamma of the last LN layer is noticeably larger — confirming our conjecture~
Please Be Generous with Criticism and Corrections
This post has attempted to answer several questions related to model optimization in the Transformer and BERT. Some of these are results I discovered in my own pretraining work; others are intuitive conjectures drawn from my own experience. In any case, consider this a reference answer of sorts. If there are any mistakes, please be understanding, and feel free to offer criticism and corrections~
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.