A Brief Discussion on Initialization, Parameterization, and Normalization of Transformers

Sampling Distributions

Initialization naturally involves random sampling, so let’s first introduce commonly used sampling distributions. Generally, we sample from a random distribution with a specified mean and variance. Three distributions are commonly used: Normal, Uniform, and Truncated Normal.

The Normal and Uniform distributions are very common. The Normal distribution is usually denoted as \mathcal{N}(\mu, \sigma^2), where \mu is the mean and \sigma^2 is the variance. A Uniform distribution on the interval [a, b] is denoted as U[a, b], with a mean of \frac{a+b}{2} and a variance of \frac{(b-a)^2}{12}. Therefore, if we specify a mean \mu and variance \sigma^2, the corresponding Uniform distribution is U[\mu - \sqrt{3}\sigma, \mu + \sqrt{3}\sigma].

Generally, sampling from a Normal distribution is more diverse, but it is theoretically unbounded. Sampling values with excessively large absolute values might be detrimental to optimization. Conversely, the Uniform distribution is bounded but often results in more "singular" samples. Thus, the "Truncated Normal distribution" emerged, combining the advantages of both. It requires a specified mean \mu, variance \sigma^2, and an interval [a, b]. It samples from \mathcal{N}(\mu, \sigma^2); if the result is within [a, b], it is kept; otherwise, it resamples until the result falls within the interval.

In TensorFlow’s tf.random.truncated_normal, the interval is hardcoded as a = \mu - 2\sigma, b = \mu + 2\sigma. Based on the formula, the actual mean of the sampled results remains \mu, but the actual variance is \gamma\sigma^2, where: \gamma = \frac{\int_{-2}^2 e^{-x^2/2}x^2 dx}{\int_{-2}^2 e^{-x^2/2} dx} = 0.7737413\dots To obtain a sampling result with variance \sigma^2, the standard deviation passed to the function should be \frac{\sigma}{\sqrt{\gamma}} = 1.1368472\dots\sigma.

Stable Second Moments

In a previous article, "Understanding Model Parameter Initialization Strategies from a Geometric Perspective", I analyzed existing initialization methods from a geometric viewpoint. The general idea is that specific random matrices approximate orthogonal matrices, ensuring model stability during the initial phase. While the geometric perspective is intuitive, it is often difficult to generalize. Therefore, we will now understand initialization from an algebraic perspective.

In standard tutorials, the idea behind deriving initialization methods is to keep the input and output means and variances as similar as possible. Usually, it is assumed that the input is a random vector with mean 0 and variance 1, and the goal is to make the output mean 0 and variance 1. However, I believe this is not strictly necessary. For certain non-negative activation functions, a mean of 0 is impossible to achieve. In fact, we only need a metric to measure whether a value "vanishes" or "explodes." Zero mean and unit variance are not essential; here, we use the second (raw) moment instead. It can be seen as a variant of the L2 norm and serves a similar purpose to variance in measuring whether values "vanish" or "explode," but it is more universal and simpler.

Now, let’s examine a fully connected layer without an activation function (let m be the number of input nodes and n be the number of output nodes): y_j = b_j + \sum_i x_i w_{i,j} For simplicity, we usually initialize the bias b_j to zero and set the mean of w_{i,j}, \mathbb{E}[w_{i,j}], to 0. This simplifies the results, though it is not strictly required—it is just a clear choice. We calculate the second moment: \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} Note that w_{i_1,j} and w_{i_2,j} are independent and identically distributed (i.i.d.). 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. Therefore, we only need to consider the case where i_1 = i_2 = i. Assuming the second moment of the input is 1, then: \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} To make \mathbb{E}[y_j^2] equal to 1, we need \mathbb{E}[w_{i,j}^2] = 1/m. Combined with the assumption of zero mean, we get the initialization strategy for w_{i,j}: "sample independently from a random distribution with mean 0 and variance 1/m." This is LeCun initialization. Note that we made no assumptions about the mean of the input, so it works even if the inputs are all non-negative.

Activation Functions

Of course, this only applies to scenarios without activation functions. If an activation function is included, we need to analyze it case-by-case. For example, if the activation function is \text{ReLU}, we can assume that roughly half of the y_j values are set to zero. Thus, the estimated second moment is half of that in Equation [eq:m]: \mathbb{E}[y_j^2] = \frac{m}{2}\mathbb{E}[w_{i,j}^2] To keep the second moment unchanged, the initialization variance should be 2/m. This is He initialization, specifically designed for \text{ReLU} networks.

However, if the activation function is \text{ELU} or \text{GELU}, the analysis is not as simple. If the activation function is \tanh or \text{sigmoid}, no initialization can keep the second moment at 1. In such cases, if we still want to maintain the second moment, one option is to "fine-tune" the definition of the activation function.

Taking \text{sigmoid} as an example, assume the input has mean 0 and variance 1, and we use the "mean 0, variance 1/m" initialization. The output before activation also has mean 0 and variance 1. We can estimate the second moment after \text{sigmoid} using a standard normal distribution: \int_{-\infty}^{\infty} \frac{e^{-x^2/2}}{\sqrt{2\pi}}\text{sigmoid}(x)^2 dx = 0.2933790\dots In other words, under this assumption, the second moment after activation is approximately 0.293379. To keep the output second moment roughly constant, we can divide the output by \sqrt{0.293379}. That is, change the activation function from \text{sigmoid}(x) to \frac{\text{sigmoid}(x)}{\sqrt{0.293379}}. This is the "fine-tuned" activation function. If necessary, you can also subtract a constant to make the output mean 0.

In 2017, a "sensational" paper "Self-Normalizing Neural Networks" proposed the \text{SELU} activation function. It is essentially an \text{ELU} function "fine-tuned" using the same logic: \text{SELU}(x) = \lambda \begin{cases} x, & (x > 0) \\ \alpha e^x - \alpha, & (x \leq 0) \end{cases} where \lambda = 1.0507\dots and \alpha = 1.6732\dots. It was sensational because it claimed to achieve automatic normalization without Batch Normalization and because its dozens of pages of mathematical derivation were quite "intimidating." From the perspective above, it simply introduces two parameters to fine-tune the \text{ELU} function so that when the input is a standard normal distribution, the output activation values have mean 0 and variance 1. Thus, it is at most a good initialization method, which is why its fame was short-lived. We can solve for these 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 simple "fine-tuning," a more direct approach is various Normalization methods, such as Batch Normalization, Instance Normalization, and Layer Normalization. These methods directly calculate the mean and variance of the current data to standardize the output, without needing to estimate integrals beforehand. These three methods are generally similar; apart from Batch Normalization using a moving average for prediction, they differ only in the dimensions they normalize. For example, Layer Normalization, commonly used in NLP and Transformers, is: 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 I won’t repeat the others. For the principles behind these methods, readers can refer to my previous post "What Does BN Actually Do? A ’Closed-Door’ Analysis".

I have noticed an interesting phenomenon: Normalization generally includes two parts: subtracting the mean (centering) and dividing by the standard deviation (scaling). However, some recent works have attempted to remove the centering step, and some results even show a slight performance improvement without it.

For instance, the 2019 paper "Root Mean Square Layer Normalization" compared Layer Normalization without centering, calling it RMS Norm: 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 RMS Norm is essentially a simple variant of L2 Normalization. The paper’s results show that RMS Norm is faster than Layer Normalization and yields basically the same effect.

In addition to that paper, RMS Norm was used by Google in T5 and was thoroughly compared in another paper "Do Transformer Modifications Transfer Across Implementations and Applications?", which demonstrated its superiority. It seems likely that RMS Norm will replace Layer Normalization as the standard for Transformers in the future.

Coincidentally, the 2019 paper "Analyzing and Improving the Image Quality of StyleGAN" proposed StyleGAN2, where they found that Instance Normalization caused "water droplets" in some generated images. They eventually replaced Instance Normalization with "Weight demodulation," but they also found that keeping Instance Normalization while removing the centering operation could also mitigate the phenomenon. This provides further evidence that the centering operation in Normalization might have negative effects.

An intuitive guess is that the centering operation, similar to the bias term in fully connected layers, stores prior distribution information about the pre-training task. Storing this prior information directly in the model might reduce its transferability. Therefore, T5 not only removed the centering operation in Layer Normalization but also removed the bias terms in every layer.

NTK Parameterization

Returning to Xavier initialization for fully connected layers, it suggests using a "random distribution with mean 0 and variance 1/m." However, instead of this initialization, we can use another parameterization method: initialize with a "random distribution with mean 0 and variance 1," but divide the output by \sqrt{m}. The model becomes: y_j = b_j + \frac{1}{\sqrt{m}}\sum_i x_i w_{i,j} In Gaussian processes, this is known as "NTK parameterization". References include "Neural Tangent Kernel: Convergence and Generalization in Neural Networks" and "On the infinite width limit of neural networks with a standard parameterization". Personally, I first saw this in the PGGAN paper "Progressive Growing of GANs for Improved Quality, Stability, and Variation".

Clearly, with NTK parameterization, we can initialize all parameters with unit variance while maintaining the second moment. Even the "fine-tuned activation functions" mentioned earlier can be seen as a form of NTK parameterization. A natural question is: What are the benefits of NTK parameterization compared to direct Xavier initialization?

Theoretically, there is a slight advantage. With NTK parameterization, all parameters are initialized with variance 1, meaning every parameter is roughly of the same magnitude \mathcal{O}(1). This allows us to set a larger learning rate, such as 10^{-2}. If using an adaptive optimizer, the update is roughly \frac{\text{gradient}}{\sqrt{\text{gradient} \otimes \text{gradient}}} \times \text{learning rate}. We then know that a 10^{-2} learning rate adjusts parameters by about 1\% per step. Overall, NTK parameterization allows us to treat every parameter more equally and provides a more intuitive understanding of the update magnitude, helping us tune parameters better.

Now we can discuss the question from the beginning: Why is dividing by \sqrt{d} in Attention so important? For two d-dimensional vectors \boldsymbol{q} and \boldsymbol{k}, assuming they are sampled from a distribution with "mean 0 and variance 1," the second moment of their dot product is: \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} The second moment of the dot product is d. Since the mean is 0, the variance is also d. Attention involves a softmax after the dot product, primarily e^{\boldsymbol{q}\cdot \boldsymbol{k}}. We can roughly assume the values before softmax fall within the range [-3\sqrt{d}, 3\sqrt{d}]. Since d is usually at least 64, e^{3\sqrt{d}} is very large and e^{-3\sqrt{d}} is very small. Consequently, after softmax, the Attention distribution becomes very close to a one-hot distribution. This leads to severe gradient vanishing, resulting in poor training.

Accordingly, there are two solutions. One is like NTK parameterization: divide by \sqrt{d} after the dot product to make the variance of \boldsymbol{q}\cdot \boldsymbol{k} equal to 1. This ensures e^3 and e^{-3} are not too extreme, preventing the softmax from becoming one-hot and avoiding gradient vanishing. This is the standard approach in Transformers like BERT. The other is not to divide by \sqrt{d} but to divide the initialization variance of the fully connected layers for \boldsymbol{q} and \boldsymbol{k} by an additional \sqrt{d}. This also makes the initial variance of \boldsymbol{q}\cdot \boldsymbol{k} equal to 1. T5 adopts this latter approach.

Residual Connections

Finally, we must discuss the design of the residual connection x + F(x). It is easy to prove that if the variance (or second moment) of x is \sigma_1^2 and that of F(x) is \sigma_2^2, and assuming they are independent, the variance of x + F(x) is \sigma_1^2 + \sigma_2^2. In other words, residual connections further amplify variance, so we need strategies to reduce it.

A simple solution is to add a Normalization operation after the residual: x_{t+1} = \text{Norm}(x_t + F_t(x_t)) This is called the Post Norm structure, used in the original Transformer and BERT. However, while this stabilizes the variance of forward propagation, it severely weakens the identity branch of the residual, losing the "ease of training" advantage. Usually, warmup and a sufficiently small learning rate are required for convergence.

How to understand this? Assume that initially, the variances of x and F(x) are both 1. Then the variance of x + F(x) is 2, and the Normalization operation reduces it back to 1. This means that at the initial stage, Post Norm is equivalent to: x_{t+1} = \frac{x_t + F_t(x_t)}{\sqrt{2}} Recursively, we get: \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}} \\ &= \dots \\ &= \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}} + \dots + \frac{F_{l-1}(x_{l-1})}{2^{1/2}} \end{aligned} See the problem? The purpose of a residual connection is to provide a "green channel" for the gradient to flow back directly. In Post Norm, this channel is severely weakened; the earlier the channel, the smaller its weight. The residual exists "in name only," making it difficult to train. For more analysis, refer to the paper "On Layer Normalization in the Transformer Architecture".

An improvement is Pre Norm, which follows the idea of "normalize only when needed": x_{t+1} = x_t + F_t(\text{Norm}(x_t)) Similarly, after iteration, we can assume the initial stage looks like: x_l = x_0 + F_0(x_0) + F_1(x_1/\sqrt{2}) + F_2(x_2/\sqrt{3}) + \dots + F_{l-1}(x_{l-1}/\sqrt{l}) In this case, at least every residual channel is equal, making the residual effect more pronounced and optimization easier. Of course, the variance of the final x_l will be very large, so another Normalization is needed before the prediction layer.

In my view, neither Post Norm nor Pre Norm is perfect because they cannot maintain an identity function at the initial stage. The most elegant method would be to introduce a scalar parameter \alpha_t initialized to 0: x_{t+1} = x_t + \alpha_t F_t(x_t) and then gradually update \alpha_t. This ensures the model is an identity function at the start, eliminating variance issues. This trick appeared in two papers: "Batch Normalization Biases Residual Blocks Towards the Identity Function in Deep Networks" (called SkipInit) and "ReZero is All You Need: Fast Convergence at Large Depth" (called ReZero). Both papers, published less than a month apart, showed that this approach could basically replace Normalization in residuals. Additionally, "Fixup Initialization: Residual Learning Without Normalization" proposed Fixup, which initializes the last layer of each residual branch to zero, sharing similarities with SkipInit and ReZero.

Regarding the update of \alpha_t, both SkipInit and ReZero treat it as a model parameter updated alongside others. I initially thought the same. Later, I realized that \alpha_t is not on equal footing with other parameters. For example, through NTK parameterization, other parameters can use a large learning rate, but \alpha_t clearly should not. Furthermore, we know that if training is successful, both Post Norm and Pre Norm work well (corresponding to \alpha_t=1). Thus, the choice of residual mode is an initialization issue rather than a capacity issue. Considering these points, I eventually let \alpha_t increase slowly with a fixed, small step size until it reaches \alpha_t=1, where it stays fixed. In my experiments, this update mode achieved the best results.

The Long Road of Model Training

This article discussed issues related to model initialization, parameterization, and normalization, hoping to provide some reference for your model tuning. The road of "alchemy" (model training) is endless; besides these topics, there are many other things to tune, such as learning rates, optimizers, and data augmentation. I wish all readers smooth sailing on their journey of model training!

Original URL: https://kexue.fm/archives/8620