A Brief Discussion on Transformer Initialization, Parameterization and Normalization
A few days ago while training a new Transformer model, I found that no matter how I trained it, it just wouldn't converge. After some debugging, it turned out that when doing self-attention, I had forgotten to divide by $\sqrt{d}$ after $\boldsymbol{Q}\boldsymbol{K}^{\top}$, so I went back and refreshed my understanding of why dividing by $\sqrt{d}$ is so important. Of course, Google's T5 genuinely doesn't divide by $\sqrt{d}$, yet it still converges fine — that's because it makes some adjustments to its initialization strategy, so this matter is also connected to initialization.
Taking this opportunity, this post walks through initialization, parameterization, normalization and related topics, with the discussion mostly centered on Transformers.
Sampling Distributions
Initialization is naturally done via random sampling, so let's first go over the commonly used sampling distributions. Generally speaking, we sample from a random distribution with a specified mean and variance to perform initialization. Three commonly used random distributions are: the normal distribution, the uniform distribution, and the truncated normal distribution. more
Clearly, both the normal and uniform distributions are very familiar. The normal distribution is usually denoted $\mathcal{N}(\mu,\sigma^2)$, where $\mu$ is the mean and $\sigma^2$ is the variance; the uniform distribution over the interval $[a,b]$ is generally denoted $U[a,b]$, with mean $\frac{a+b}{2}$ and variance $\frac{(b-a)^2}{12}$, so for a specified mean $\mu$ and variance $\sigma^2$, the corresponding uniform distribution is $U[\mu-\sqrt{3}\sigma,\mu + \sqrt{3}\sigma]$.
Generally speaking, samples from the normal distribution tend to be more diverse, but it is theoretically unbounded, so sampling values with excessively large absolute values may hurt optimization; conversely, the uniform distribution is bounded, but its samples tend to be more homogeneous. This gave rise to the "truncated normal distribution", which combines the advantages of both. The truncated normal distribution specifies a mean $\mu$ and variance $\sigma^2$, as well as an interval $[a,b]$; it samples from $\mathcal{N}(\mu,\sigma^2)$, and if the sampled result falls within $[a,b]$, it is kept, otherwise sampling is repeated until the result falls within $[a,b]$.
In TensorFlow's built-in tf.random.truncated_normal, $a=\mu-2\sigma,b=\mu+2\sigma$ is hard-coded. So, from the formula, we can work out that the actual mean of this function's sampling results is still $\mu$, but the actual variance is $\gamma\sigma^2$, where:
\begin{equation}\gamma=\frac{\int_{-2}^2 e^{-x^2/2}x^2 dx}{\int_{-2}^2 e^{-x^2/2} dx}=0.7737413\dots\end{equation}
If we want the sampled result to have variance $\sigma^2$, then the standard deviation passed into the function needs to be $\frac{\sigma}{\sqrt{\gamma}}=1.1368472\dots\sigma$.
Stabilizing the Second Moment
In an earlier post, Understanding Model Parameter Initialization Strategies from a Geometric Perspective, I analyzed existing initialization methods from a geometric point of view; the basic idea is that a particular random matrix approximately behaves like an orthogonal matrix, which guarantees the stability of the model at the initial stage. However, while the geometric perspective has the advantage of being intuitive, it is usually hard to generalize, so here let's instead try to understand initialization methods from an algebraic point of view.
In typical tutorials, the idea behind deriving initialization methods is to keep the input and output having the same mean and variance — usually one assumes the input is a random vector with mean 0 and variance 1, and then tries to make the output have mean 0 and variance 1. However, I think this is actually unnecessary, and moreover, for certain non-negative activation functions, a mean of 0 simply cannot be achieved. In fact, all we need is some metric that measures whether a given quantity "vanishes" or "explodes" — mean 0, variance 1 is not a necessary condition. Here we'll use the second (raw) moment instead, which can be viewed as a variant of the L2 norm; it plays a role similar to variance in that both can be used to measure whether a quantity vanishes or explodes, but it is comparatively more general and simpler.
Now, let's examine a fully-connected layer with no activation function (let the number of input nodes be $m$ and the number of output nodes be $n$):
\begin{equation} y_j = b_j + \sum_i x_i w_{i,j}\end{equation}
For simplicity, we usually initialize the bias term $b_j$ with all zeros, and also set the mean $\mathbb{E}[w_{i,j}]$ of $w_{i,j}$ to 0, which helps simplify the results below — though this isn't strictly required, it is indeed a fairly clean choice. Let's compute the second moment:
\begin{equation}\begin{aligned} \mathbb{E}[y_j^2] =&\, \mathbb{E}\left[\left(\sum_i x_i w_{i,j}\right)^2\right]=\mathbb{E}\left[\left(\sum_{i_1} x_{i_1} w_{i_1,j}\right)\left(\sum_{i_2} x_{i_2} w_{i_2,j}\right)\right]\\ =&\,\mathbb{E}\left[\sum_{i_1, i_2} (x_{i_1}x_{i_2}) (w_{i_1,j} w_{i_2,j})\right] = \sum_{i_1, i_2} \mathbb{E}[x_{i_1}x_{i_2}] \mathbb{E}[w_{i_1,j} w_{i_2,j}] \end{aligned}\end{equation}
Note that $w_{i_1,j},w_{i_2,j}$ are i.i.d., so when $i_1\neq i_2$, $\mathbb{E}[w_{i_1,j}w_{i_2,j}]=\mathbb{E}[w_{i_1,j}]\mathbb{E}[w_{i_2,j}]=0$, meaning we only need to consider the case $i_1=i_2=i$. Assuming the input has second moment 1, then
\begin{equation} \mathbb{E}[y_j^2] = \sum_{i} \mathbb{E}[x_i^2] \mathbb{E}[w_{i,j}^2]= m\mathbb{E}[w_{i,j}^2]\label{eq:m}\end{equation}
So for $\mathbb{E}[y_j^2]$ to equal 1, we need $\mathbb{E}[w_{i,j}^2]=1/m$; combined with the assumption of zero mean, we obtain the $w_{i,j}$ initialization strategy: "sample independently and repeatedly from a random distribution with mean 0 and variance $1/m$" — this is Lecun initialization. Note that in this derivation we made no assumption about the mean of the input, so it works fine even if the input is entirely non-negative.
Activation Functions
Of course, this is only the case with no activation function; once an activation function is added, the analysis must be done case by case. For example, if the activation function is $\text{relu}$, we can roughly assume that about half of the $y_j$ values get zeroed out, so the estimated second moment is half of the value from equation $\eqref{eq:m}$:
\begin{equation} \mathbb{E}[y_j^2] = \frac{m}{2}\mathbb{E}[w_{i,j}^2]\end{equation}
This gives an initialization variance of $2/m$ that keeps the second moment unchanged — this is the He initialization, specifically designed for $\text{relu}$ networks.
However, if the activation function is something like $\text{elu},\text{gelu}$, the analysis is not so simple; and if the activation function is $\tanh,\text{sigmoid}$, there is no initialization at all that can make the second moment equal to 1. In such cases, if we still want to keep the second moment unchanged, one option is to "slightly tweak the definition of the activation function."
Take $\text{sigmoid}$ as an example: assume the input has mean 0 and variance 1, and we still use "mean 0, variance $1/m$" initialization, so the pre-activation output also has mean 0 and variance 1. We can then estimate the second moment after applying $\text{sigmoid}$ using the standard normal distribution:
\begin{equation}\int_{-\infty}^{\infty} \frac{e^{-x^2/2}}{\sqrt{2\pi}}\text{sigmoid}(x)^2 dx = 0.2933790\dots\end{equation}
That is to say, under this assumption, the second moment of the model's output after activation is roughly $0.293379$. So, if we want to keep the second moment of the output roughly unchanged, we can simply divide the output by $\sqrt{0.293379}$ — in other words, change the activation function from $\text{sigmoid}(x)$ to $\frac{\text{sigmoid}(x)}{\sqrt{0.293379}}$; this is the "tweaked" activation function. If you think it necessary, you can also subtract a constant to bring the output mean back to 0.
I recall that back in 2017 there was a paper that caused quite a stir, Self-Normalizing Neural Networks, which proposed the activation function $\text{selu}$. It is actually also a "tweaked" version of the $\text{elu}$ function based on exactly the same idea, with the following form:
\begin{equation}\text{selu}(x)=\lambda\left\{\begin{aligned} &x,& (x > 0) \\ &\alpha e^{x}-\alpha, &(x\leq 0) \end{aligned}\right.\end{equation}
where $\lambda=1.0507\dots,\alpha=1.6732\dots$. It caused a stir at the time for two reasons: first, it claimed to achieve automatic self-normalization of the network without needing techniques like Batch Normalization; second, the accompanying dozens of pages of mathematical derivation seemed rather intimidating. But viewed through the lens above, all it does is introduce two parameters to tweak the $\text{elu}$ function so that, given a standard normal input, the output activation has mean 0 and variance 1 — so at best it's just a rather good initialization scheme, which is presumably why the stir was only momentary. We can likewise solve for its two parameters numerically using Mathematica:
f[x_] = Exp[-x^2/2]/Sqrt[2 Pi];
s[x_] = Piecewise[{{\[Lambda]*x,
x > 0}, {\[Lambda]*\[Alpha]*(Exp[x] - 1), x <= 0}}];
x1 = Integrate[f[x]*s[x], {x, -Infinity, Infinity}];
x2 = Integrate[f[x]*s[x]^2, {x, -Infinity, Infinity}];
N[Solve[{x1 == 0, x2 == 1}, {\[Lambda], \[Alpha]}], 20]
Direct Normalization
Of course, compared to this kind of simple "tweaking," a more direct approach is the various Normalization methods, such as Batch Normalization, Instance Normalization, Layer Normalization, and so on. These methods directly compute the mean and variance of the current data to normalize the output, without needing to estimate an integral in advance — sometimes we also call this "normalization" in a narrower sense. These three normalization methods are broadly similar, aside from Batch Normalization having an extra step for maintaining a running mean/variance used at prediction time; they otherwise just differ in which dimension is normalized. For example, the one most used in NLP, especially in Transformer models, is Layer Normalization:
\begin{equation}y_{i,j,k} = \frac{x_{i,j,k} - \mu_{i,j}}{\sqrt{\sigma_{i,j}^2 + \epsilon}}\times\gamma_k + \beta_k,\quad \mu_{i,j} = \frac{1}{d}\sum_{k=1}^d x_{i,j,k},\quad \sigma_{i,j}^2 = \frac{1}{d}\sum_{k=1}^d (x_{i,j,k}-\mu_{i,j})^2\end{equation}
I won't repeat the details for the others. For readers interested in the principles behind why these methods work, you can refer to my earlier post What Exactly Does BN Do? A Closed-Door Analysis.
Here's an interesting phenomenon I've noticed: Normalization generally consists of two parts — subtracting the mean ("center") and dividing by the standard deviation ("scale") — but some recent work has gradually started dropping the centering step, and in some cases the results even show a slight performance improvement after removing it.
For example, the 2019 paper Root Mean Square Layer Normalization compared Layer Normalization with the centering step removed, calling it RMS Norm, with the following form:
\begin{equation}y_{i,j,k} = \frac{x_{i,j,k}}{\sqrt{\sigma_{i,j}^2 + \epsilon}}\times\gamma_k,\quad \sigma_{i,j}^2 = \frac{1}{d}\sum_{k=1}^d x_{i,j,k}^2\end{equation}
As you can see, RMS Norm is really just a simple variant of L2 normalization, but this paper's overall results show that RMS Norm is faster than Layer Normalization while achieving essentially the same performance.
Besides this paper, RMS Norm was also used by Google in T5, and another paper, Do Transformer Modifications Transfer Across Implementations and Applications?, conducted fairly thorough comparative experiments demonstrating the advantages of RMS Norm. It looks like RMS Norm may well replace Layer Normalization as the standard choice for Transformers going forward.
Interestingly, another 2019 paper, Analyzing and Improving the Image Quality of StyleGAN, proposed StyleGAN2, an improved version of StyleGAN, and found that the Instance Normalization used in the original caused "water droplet" artifacts in some generated images. They ultimately removed Instance Normalization in favor of something called "weight demodulation," but they also found that keeping Instance Normalization while simply dropping the centering operation could also alleviate the artifacts. This provides further evidence that the centering operation in Normalization may have negative effects.
An intuitive guess is that the centering operation, much like the bias term in a fully-connected layer, ends up storing prior distributional information related to the pretraining task, and storing this kind of prior information directly in the model may actually hurt the model's transferability. This is presumably why T5 not only removes the centering operation from Layer Normalization, but also removes the bias terms from every layer.
NTK Parameterization
Back to Xavier initialization for a fully-connected layer: it tells us to initialize using "a random distribution with mean 0 and variance $1/m$." However, besides using this initialization directly, there's another way to parameterize things: initialize using "a random distribution with mean 0 and variance 1," but divide the output result by $\sqrt{m}$, i.e., the model becomes:
\begin{equation} y_j = b_j + \frac{1}{\sqrt{m}}\sum_i x_i w_{i,j}\end{equation}
In the context of Gaussian processes, this is called "NTK parameterization." Relevant references include Neural Tangent Kernel: Convergence and Generalization in Neural Networks and On the Infinite Width Limit of Neural Networks with a Standard Parameterization. For my own part, though, the first time I encountered this trick was in the PGGAN paper, Progressive Growing of GANs for Improved Quality, Stability, and Variation.
Clearly, with NTK parameterization we can initialize all parameters using a standard variance, while still keeping the second moment unchanged — and in fact, the "activation function tweaking" introduced earlier can also be seen as a form of NTK parameterization. A natural question is: what advantage does NTK parameterization have over directly using Xavier initialization?
In theory, there is a benefit. With NTK parameterization, all parameters can be initialized from a distribution with variance 1, meaning every parameter is roughly on the same order of magnitude, $\mathcal{O}(1)$. This means we can set a larger learning rate, say $10^{-2}$, and if we use an adaptive optimizer whose update magnitude is roughly $\frac{\text{gradient}}{\sqrt{\text{gradient}\otimes\text{gradient}}}\times\text{learning rate}$, then we know that with a learning rate of $10^{-2}$, the magnitude of each parameter update per step is roughly $1\%$. In short, NTK parameterization lets us treat every parameter on a more equal footing, and gives us a more concrete sense of the scale of training updates, helping us tune parameters more effectively.
Now we can come back to the question raised at the start of the post: why is dividing by $\sqrt{d}$ in attention so important? For two $d$-dimensional vectors $\boldsymbol{q},\boldsymbol{k}$, assuming both are sampled from a distribution with "mean 0, variance 1," the second moment of their inner product is:
\begin{equation}\begin{aligned} \mathbb{E}[(\boldsymbol{q}\cdot \boldsymbol{k})^2]=&\,\mathbb{E}\left[\left(\sum_{i=1}^d q_i k_i\right)^2\right] = \mathbb{E}\left[\left(\sum_i q_i k_i\right)\left(\sum_j q_j k_j\right)\right]\\ =&\,\mathbb{E}\left[\sum_{i,j} (q_i q_j) (k_i k_j)\right] = \sum_{i,j} \mathbb{E}[q_i q_j] \mathbb{E}[k_i k_j]\\ =&\,\sum_i \mathbb{E}[q_i^2] \mathbb{E}[k_i^2] = d \end{aligned}\end{equation}
That is, the second moment of the inner product is $d$; since the mean is also 0, this also means the variance is $d$. Attention applies softmax after the inner product — the key operation is $e^{\boldsymbol{q}\cdot \boldsymbol{k}}$ — and we can roughly say that the value after the inner product and before softmax ranges from $-3\sqrt{d}$ to $3\sqrt{d}$. Since $d$ is usually at least 64, $e^{3\sqrt{d}}$ is fairly large while $e^{-3\sqrt{d}}$ is quite small; hence, after softmax, the attention distribution ends up extremely close to a one-hot distribution, which causes a severe vanishing gradient problem and results in poor training.
Accordingly, there are two ways to fix this. One is, as in NTK parameterization, to divide by $\sqrt{d}$ after the inner product, so that the variance of $\boldsymbol{q}\cdot \boldsymbol{k}$ becomes 1, keeping $e^3,e^{-3}$ from being too large or too small, so that softmax doesn't collapse into a near one-hot distribution and cause gradient vanishing — this is the standard approach used in the self-attention of regular Transformers like BERT. The other option is not to divide by $\sqrt{d}$ at all, but instead, when initializing the fully-connected layers that produce $\boldsymbol{q},\boldsymbol{k}$, divide the initialization variance by an extra factor of $\sqrt{d}$, which likewise brings the initial variance of $\boldsymbol{q}\cdot \boldsymbol{k}$ to 1. T5 adopts this latter approach.
Residual Connections
Finally, we must discuss the design considerations around residual connections $x + F(x)$. It's easy to show that if the variance of $x$ (and likewise its second moment) is $\sigma_1^2$ and the variance of $F(x)$ is $\sigma_2^2$, and assuming the two are independent, then the variance of $x + F(x)$ is $\sigma_1^2 + \sigma_2^2$. In other words, residual connections further amplify the variance, so we also need a corresponding strategy to shrink it back down.
One rather naive approach is to simply add a Normalization operation right after the residual sum:
\begin{equation}x_{t+1} = \text{Norm}(x_t + F_t(x_t))\end{equation}
We can call this the Post-Norm structure — it's also the design used by the original Transformer and BERT. However, while this approach stabilizes the forward-pass variance, it actually severely weakens the identity branch of the residual connection, thereby losing the "easy to train" advantage that residual connections normally provide. Typically, one needs warmup and a sufficiently small learning rate to get it to converge.
How should we understand this? Suppose that at initialization, $x,F(x)$ both have variance 1; then the variance of $x+F(x)$ is 2, and the Normalization operation is responsible for bringing the variance back down to 1 — which means that, at the initial stage, Post-Norm is effectively equivalent to
\begin{equation}x_{t+1} = \frac{x_t + F_t(x_t)}{\sqrt{2}}\end{equation}
Recursing this down, we get
\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}
Do you see the problem? The whole point of a residual connection is to give the earlier layers a "fast lane" so gradients can flow back more directly, but in Post-Norm, this "fast lane" is severely weakened — the closer a layer is to the front, the smaller its weight ends up being, so the residual connection is a "name without substance," and training remains difficult. For more analysis along these lines, see the paper On Layer Normalization in the Transformer Architecture.
One targeted improvement is called Pre-Norm, whose idea is "only normalize right when it's needed," taking the form:
\begin{equation}x_{t+1} = x_t + F_t(\text{Norm}(x_t))\end{equation}
Similarly, if we iteratively expand this, we can see that, at the initial stage,
\begin{equation} x_l = x_0 + F_0(x_0) + F_1(x_1/\sqrt{2}) + F_2(x_2/\sqrt{3}) + \cdots + F_{l-1}(x_{l-1}/\sqrt{l})\end{equation}
This way, at least every residual channel is given equal weight, and the effect of the residual connection is much more pronounced than in Post-Norm, making it easier to optimize. Of course, this means the final variance of $x_l$ will end up quite large, so we still need to add a Normalization right before the prediction head, on $x_l$.
As I see it, neither Post-Norm nor Pre-Norm is entirely satisfactory, because neither can maintain an identity function at initialization. To my mind, the cleanest approach is to introduce a scalar parameter $\alpha_t $ initialized to 0, so that
\begin{equation}x_{t+1} = x_t + \alpha_t F_t(x_t)\end{equation}
and then gradually update $\alpha_t$. This way, at the initial stage, we can guarantee that the model is an identity function, and hence there's no variance problem to worry about. This trick later showed up in two papers: in Batch Normalization Biases Residual Blocks Towards the Identity Function in Deep Networks it's called SkipInit, and in ReZero is All You Need: Fast Convergence at Large Depth it's called ReZero. The two papers appeared less than a month apart, and both show that this trick can basically substitute directly for the Normalization operation in residual connections. In addition, Fixup Initialization: Residual Learning Without Normalization proposed a method called Fixup, which initializes the last layer of each residual branch with all zeros — this also has some things in common with SkipInit and ReZero.
As for updating $\alpha_t$, both SkipInit and ReZero treat it as a model parameter to be updated alongside all the others, and that's how I initially thought about it too. Later I realized that $\alpha_t$'s role is not on equal footing with the other parameters, and shouldn't be treated the same way. For instance, thanks to the NTK parameterization discussed earlier, we can use a large learning rate for other parameters, but clearly $\alpha_t$ should not use a large learning rate. Furthermore, we know that if training succeeds, both Post-Norm and Pre-Norm end up performing quite well (corresponding to $\alpha_t=1$), so the choice of residual pattern is purely an initialization issue rather than a matter of model capacity. Taking all this into account, I ended up simply letting $\alpha_t$ slowly increase with a fixed, small step size, until it reaches $\alpha_t=1$ and then holds fixed there. In my experiments, this update scheme achieved the best results.
A Long Road of Model Training Ahead
This post has discussed initialization, parameterization, and normalization in neural network models, and I hope it provides some useful reference for your own model tuning. The road of model training is a long one, and beyond what's covered here, there's plenty more that can be tuned — learning rate, optimizer choice, data augmentation, and so on. Wishing all readers smooth sailing on their own training journeys~
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.