What are the Difficulties in Training a 1000-layer Transformer?

Update Explosion

The complete analysis in the original paper is quite long, and some assumptions or descriptions are not entirely reasonable upon closer inspection. Therefore, in this sharing, I will try to correct these issues and derive similar results in a more reasonable manner.

Suppose the loss function is \mathcal{L}(\boldsymbol{\theta}), where \boldsymbol{\theta} represents its parameters. Consider the increment of the loss function when the parameters change from \boldsymbol{\theta} to \boldsymbol{\theta}+\Delta\boldsymbol{\theta}: \Delta\mathcal{L} = \mathcal{L}(\boldsymbol{\theta}+\Delta\boldsymbol{\theta}) - \mathcal{L}(\boldsymbol{\theta}) \approx \langle\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}),\Delta\boldsymbol{\theta}\rangle For SGD, we have \Delta\boldsymbol{\theta}=-\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}), so \Delta\mathcal{L} \approx -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert^2. Suppose the model has N layers, and each layer has K parameter matrices (K is approximately constant). Combined with Xavier initialization and various Normalization techniques, we can make the gradient magnitude of each parameter matrix on the order of \mathcal{O}(1). Thus, we have \Delta\mathcal{L}=\mathcal{O}(\eta NK). Therefore, the update amount of the model at each step is proportional to the model depth N. If the model is deeper, the update amount is larger, which means that in the initial stage, the model is more likely to enter a poor local optimum, leading to training stagnation or even collapse. This is the "update explosion" problem.

There are two solutions at this point. One is to use a smaller learning rate during the initial stage (no more than the order of \eta/N) and then gradually increase it; this is the Warmup technique. The second is to adjust the initialization scheme so that the gradient of the parameters is on the order of \mathcal{O}(1/\sqrt{N}), which automatically offsets the influence of the model depth.

Magnitude Analysis

How can we achieve the second solution? We can attempt to analyze the gradients of the Transformer. However, calculating exact gradients is tedious, and in fact, we do not need exact gradients but rather a magnitude analysis. Thus, we can use the following "magnitude decomposition" trick to transform the problem into scalar derivatives.

For a matrix \boldsymbol{W}, we decompose it into the form \boldsymbol{W}=\lambda \boldsymbol{U}, where \lambda = \mathop{\text{argmin}}_{\kappa > 0} \Vert \boldsymbol{W}\boldsymbol{W}^{\top}/\kappa^2 - \boldsymbol{I}\Vert Simply put, we want to decompose a matrix into the product of a scalar \lambda and a matrix \boldsymbol{U} that is as orthogonal as possible. Since \boldsymbol{U} is close to an orthogonal matrix, it serves as a standard reference frame, while the corresponding \lambda represents the magnitude of the matrix \boldsymbol{W}. If \boldsymbol{W} uses Xavier initialization, then \lambda is equivalent to the gain parameter therein; that is, an additional \lambda is multiplied on top of the Xavier initialization. This is because the result of Xavier initialization is close to an orthogonal matrix (refer to "Understanding Model Parameter Initialization Strategies from a Geometric Perspective").

Under this decomposition, we have \frac{\partial \mathcal{L}(\lambda \boldsymbol{U})}{\partial \lambda} = \left\langle\frac{\partial \mathcal{L}(\lambda \boldsymbol{U})}{\partial (\lambda \boldsymbol{U})}, \boldsymbol{U}\right\rangle = \left\langle\frac{\partial \mathcal{L}(\boldsymbol{W})}{\partial \boldsymbol{W}}, \boldsymbol{U}\right\rangle This means that \frac{\partial \mathcal{L}}{\partial \lambda} is proportional in magnitude to \frac{\partial \mathcal{L}}{\partial \boldsymbol{W}}. Therefore, performing a magnitude analysis on \frac{\partial \mathcal{L}}{\partial \lambda} is equivalent to performing one on \frac{\partial \mathcal{L}}{\partial \boldsymbol{W}}. Thus, \frac{\partial \mathcal{L}}{\partial \lambda} acts as a simple "probe" for the magnitude of \frac{\partial \mathcal{L}}{\partial \boldsymbol{W}}, and the original matrix derivation can be converted into scalar derivation, reducing the difficulty of analysis.

Feed-forward Gradient

Many experimental results show that although Pre-Norm is easier to train than Post-Norm, the final performance of Post-Norm is often better. Therefore, the original paper retains the Post-Norm structure and considers a more general form (DeepNorm): \boldsymbol{x}_{l+1} = \text{LN}(\alpha\boldsymbol{x}_l + F(\boldsymbol{x}_l)) = \text{LN}(\boldsymbol{x}_l + F(\boldsymbol{x}_l)/\alpha) where \alpha > 0 is a constant. For simplicity, let’s first consider the FFN layer: \boldsymbol{x}_{l+1} = \text{LN}(\boldsymbol{x}_l + \phi(\boldsymbol{x}_l \boldsymbol{W}_1)\boldsymbol{W}_2/\alpha) Here \phi is the activation function, typically ReLU or its variants (Swish, GeLU, etc.), which (approximately) satisfy \phi(\lambda x) = \lambda \phi(x), \forall \lambda > 0. Using the magnitude decomposition probe from the previous section, we get: \boldsymbol{x}_{l+1} = \text{LN}(\underbrace{\boldsymbol{x}_l + \lambda_1 \lambda_2 \phi(\boldsymbol{x}_l \boldsymbol{U}_1)\boldsymbol{U}_2/\alpha}_{\text{denoted as } \boldsymbol{z}_{l+1}}) \label{eq:ffn} Calculating the gradient of \lambda: \begin{aligned} \frac{\partial \mathcal{L}}{\partial \lambda_1} &= \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\partial \boldsymbol{z}_{l+1}}{\partial \lambda_1} = \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\lambda_2 \phi(\boldsymbol{x}_l \boldsymbol{U}_1)\boldsymbol{U}_2}{\alpha} \\ \frac{\partial \mathcal{L}}{\partial \lambda_2} &= \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\partial \boldsymbol{z}_{l+1}}{\partial \lambda_2} = \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\lambda_1 \phi(\boldsymbol{x}_l \boldsymbol{U}_1)\boldsymbol{U}_2}{\alpha} \end{aligned} We assert that \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}} and \frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}} are both \mathcal{O}(1), and since \boldsymbol{U}_1 and \boldsymbol{U}_2 are both close to orthogonal matrices, \phi(\boldsymbol{x}_l \boldsymbol{U}_1)\boldsymbol{U}_2 is also \mathcal{O}(1). Therefore, we finally have: \frac{\partial \mathcal{L}}{\partial \lambda_1} = \mathcal{O}\left(\frac{\lambda_2}{\alpha}\right),\quad \frac{\partial \mathcal{L}}{\partial \lambda_2} = \mathcal{O}\left(\frac{\lambda_1}{\alpha}\right)

Self-Attention

Now consider Self-Attention. For magnitude analysis, we can consider single-head attention: \boldsymbol{x}_{l+1} = \text{LN}(\boldsymbol{x}_l + \sigma(\boldsymbol{x}_l \boldsymbol{W}_q\boldsymbol{W}_k^{\top}\boldsymbol{x}_l^{\top})\boldsymbol{x}_l\boldsymbol{W}_v\boldsymbol{W}_o/\alpha) where \sigma(\cdot) is shorthand for the softmax operation, and the scale operation of Attention is omitted here. The form after magnitude decomposition is: \boldsymbol{x}_{l+1} = \text{LN}(\underbrace{\boldsymbol{x}_l + \lambda_v\lambda_o \sigma (\lambda_q\lambda_k\boldsymbol{x}_l \boldsymbol{U}_q\boldsymbol{U}_k^{\top}\boldsymbol{x}_l^{\top})\boldsymbol{x}_l\boldsymbol{U}_v\boldsymbol{U}_o/\alpha}_{\text{denoted as } \boldsymbol{z}_{l+1}}) \label{eq:sa} Now we can calculate the gradients for each \lambda. Due to the presence of softmax, the gradients of \lambda_q, \lambda_k themselves will be very small and will not significantly affect the final update amount. Thus, considering the updates of \lambda_v, \lambda_o is sufficient: \begin{aligned} \frac{\partial \mathcal{L}}{\partial \lambda_v} &= \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\partial \boldsymbol{z}_{l+1}}{\partial \lambda_v} = \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\lambda_o \sigma (\lambda_q\lambda_k\boldsymbol{x}_l \boldsymbol{U}_q\boldsymbol{U}_k^{\top}\boldsymbol{x}_l^{\top})\boldsymbol{x}_l\boldsymbol{U}_v\boldsymbol{U}_o}{\alpha} \\ \frac{\partial \mathcal{L}}{\partial \lambda_o} &= \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\partial \boldsymbol{z}_{l+1}}{\partial \lambda_o} = \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}\frac{\lambda_v \sigma (\lambda_q\lambda_k\boldsymbol{x}_l \boldsymbol{U}_q\boldsymbol{U}_k^{\top}\boldsymbol{x}_l^{\top})\boldsymbol{x}_l\boldsymbol{U}_v\boldsymbol{U}_o}{\alpha} \end{aligned} Similarly, we assert that \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}} and \frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}} are \mathcal{O}(1). Note that the softmax output is a probability distribution, which performs a weighted average of the tokens in \boldsymbol{x}_l. Generally, the vectors before and after the average are of the same order of magnitude, so we consider \sigma (\lambda_q\lambda_k\boldsymbol{x}_l \boldsymbol{U}_q\boldsymbol{U}_k^{\top}\boldsymbol{x}_l^{\top})\boldsymbol{x}_l\boldsymbol{U}_v\boldsymbol{U}_o to be \mathcal{O}(1). Thus, the result is similar to the FFN layer: \frac{\partial \mathcal{L}}{\partial \lambda_v} = \mathcal{O}\left(\frac{\lambda_o}{\alpha}\right),\quad \frac{\partial \mathcal{L}}{\partial \lambda_o} = \mathcal{O}\left(\frac{\lambda_v}{\alpha}\right)

Preliminary Conclusion

Whether for FFN or Self-Attention, we have reached similar conclusions. For simplicity, assume the magnitude of each parameter (at least during initialization) is consistent, i.e., all \lambda take the same value. Then the overall conclusion is: \frac{\partial \mathcal{L}}{\partial \lambda} = \mathcal{O}\left(\frac{\lambda}{\alpha}\right) That is, the magnitude of the gradient is \mathcal{O}(\lambda/\alpha). On the other hand, an N-layer Transformer model generally consists of N layers of Self-Attention plus N layers of FFN, so strictly speaking, the number of layers is 2N. Therefore, according to the analysis in the "Update Explosion" section, we need to adjust the gradient to \mathcal{O}(1/\sqrt{2N}). The above equation tells us we can achieve this by letting \lambda/\alpha=1/\sqrt{2N}. The original paper’s scaling is slightly more relaxed, resulting in \lambda/\alpha = 1/\sqrt{4N}, which is equivalent in magnitude.

Now we have a proportional relationship between \lambda and \alpha, but we cannot directly obtain specific values for \lambda and \alpha. According to the paper, from a symmetry perspective, letting \lambda=1/\alpha allows us to solve: \alpha = (2N)^{1/4},\quad \lambda = (2N)^{-1/4} \label{eq:result} However, a purely symmetrical explanation is clearly not convincing enough. We need to understand what different choices actually lead to. To this end, we can compare two other sets of solutions:

Alternative 1: \alpha=1, \lambda=(2N)^{-1/2}. In this case, the parameter initialization is scaled down to (2N)^{-1/2} times the original, and the gradient is also scaled down to (2N)^{-1/2} times. According to SGD’s \Delta\boldsymbol{\theta}=-\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}), the update amount per step is also (2N)^{-1/2} times the original. This means the relative learning magnitude before and after adjustment remains unchanged. Therefore, it is possible that \lambda starts at the \mathcal{O}((2N)^{-1/2}) level but deviates from this magnitude after a few steps on the training set.

Alternative 2: \alpha=(2N)^{1/2}, \lambda=1. Here, the parameter initialization is not scaled down, but the gradient is scaled down to (2N)^{-1/2} times. According to SGD’s \Delta\boldsymbol{\theta}=-\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}), the update amount per step is also (2N)^{-1/2} times. The relative learning magnitude is significantly reduced, which might lead to very slow learning.

Both cases seem to have drawbacks, so Equation [eq:result], which lies between them, seems to make sense. it maintains the gradient scaling at (2N)^{-1/2} while making the initial learning step slightly slower, but not too slow, implicitly acting as a Warmup.

Various Optimizers

The above analyses are based on SGD, but in fact, we rarely use SGD directly to train NLP models. We mostly use adaptive learning rate optimizers, which fall into two main categories: one uses second moments to correct the learning rate, such as Adam and AdamW; the other further corrects the learning rate through parameter magnitudes, such as LAMB and AdaFactor. The original paper states, "We derive on SGD and then verify on Adam and find it works," but theoretically, they are not completely universal. In this section, we will perform a targeted analysis.

For Adam-like optimizers, the update amount per step is approximately \Delta\boldsymbol{\theta}=-\eta\,\text{sign}(\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})), so \Delta\mathcal{L} \approx -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert_1. This is proportional to the 1st power of the gradient rather than the 2nd power. Therefore, to make the update amount independent of the number of layers, the gradient should be scaled down to 1/(2N) times the original. That is, we should have \lambda/\alpha=1/(2N). If we also let \lambda=1/\alpha, then: \alpha = (2N)^{1/2},\quad \lambda = (2N)^{-1/2}

For LAMB-like optimizers, the update amount per step is approximately \Delta\boldsymbol{\theta}=-\eta\Vert\theta\Vert\,\text{sign}(\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})), so \Delta\mathcal{L} \approx -\eta\Vert\theta\Vert\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert_1. Noting that the parameter scaling ratio is \lambda and the gradient scaling ratio is \lambda/\alpha, we have \Delta\mathcal{L}=\mathcal{O}(2N\lambda^2/\alpha), thus \lambda^2/\alpha=1/(2N). Note that for this type of optimizer, the relative update amount per step is the same (equal to the learning rate \eta); no matter how \alpha, \lambda are adjusted, the relative update size will not change. Therefore, we can directly take \alpha=1, \lambda=(2N)^{-1/2}.

The comparison of results is summarized below:

Post-hoc Analysis

The derivation process in the previous two sections used the assertion that "\frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}} and \frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}} are both \mathcal{O}(1)." Is this true? Let’s perform a post-hoc analysis.

It is actually quite simple. After the aforementioned adjustments, whether for the FFN layer [eq:ffn] or the Self-Attention layer [eq:sa], the weight of each residual branch in the initial stage is scaled to \lambda^2/\alpha times the original. Regardless of the optimizer’s result, \lambda^2/\alpha is a relatively small number. This means that in the initial stage, the entire model is actually close to an identity function. Therefore, \frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}} and \frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}} are naturally \mathcal{O}(1), so the conclusion and the assertion are self-consistent.

Additionally, some readers might wonder if the same analysis can be applied to the Pre-Norm structure. The answer is yes, and the conclusion is basically consistent. However, because the Norm is placed before the residual branch, there is no need to set the \alpha parameter. Thus, the conclusion is that all \alpha in the Post-Norm results above are equal to 1, and the corresponding \lambda is recalculated.

Finally, readers might question whether model depth is really that important after spending so much effort discussing how to make the model deep. Yes, it is. The original paper provides a beautiful experimental result: a 200-layer "deep and narrow" model (3.2 billion parameters) outperformed the previous 48-layer "shallow and wide" SOTA model (12 billion parameters):

Summary

This article analyzed the bottlenecks in making Transformers "deep" and provided corresponding solutions. The main ideas of the article originate from Microsoft’s new DeepNet, and the analysis process of the original paper has been simplified and refined.

Original address: https://kexue.fm/archives/8978