After the publication of the article "What are the difficulties in training a 1000-layer Transformer?", readers quickly asked: what happens if we apply those findings to the "Gated Attention Unit (GAU)" from "FLASH: Probably the Most Interesting Efficient Transformer Design Recently"? How do the results differ from a standard Transformer? This article discusses this issue.
Conclusion First
In fact, GAU is a very easy model to train. Even if we directly use "Post-Norm + Xavier initialization" without any adjustments, we can easily train a GAU with dozens of layers without needing Warmup. Therefore, many training techniques for standard Transformers might find no use when it comes to GAU.
Why can GAU achieve this? It is simple: under default settings, theoretically, \text{GAU}(\boldsymbol{x}_l) is nearly two orders of magnitude smaller than \boldsymbol{x}_l. Thus: \boldsymbol{x}_{l+1} = \text{LN}(\boldsymbol{x}_l + \text{GAU}(\boldsymbol{x}_l)) \approx \boldsymbol{x}_l Therefore, GAU combined with residual connections is already very close to an identity function under standard initialization. Models with this property are very easy to train and typically do not require Warmup. To relate this to the conclusions of "What are the difficulties in training a 1000-layer Transformer?", these two orders of magnitude are equivalent to \lambda=1, \alpha=100, meaning it automatically incorporates the effect of DeepNorm for a hundred-layer model. Thus, theoretically, we can directly train a hundred-layer GAU model without special tuning techniques.
Model Assumptions
We only need to perform a magnitude analysis of the input and output of the GAU. The standard GAU operations are as follows: \begin{aligned} &\boldsymbol{O}=(\boldsymbol{U}\odot\boldsymbol{A}\boldsymbol{V})\boldsymbol{W}_o,\quad \boldsymbol{A}=\frac{1}{ns}\text{relu}^2\left(\mathcal{Q}(\boldsymbol{Z})\mathcal{K}(\boldsymbol{Z})^{\top}\right)\\ &\boldsymbol{U}=\phi(\boldsymbol{X}\boldsymbol{W}_u),\quad\boldsymbol{V}=\phi(\boldsymbol{X}\boldsymbol{W}_v),\quad\boldsymbol{Z}=\phi(\boldsymbol{X}\boldsymbol{W}_z) \end{aligned} where \boldsymbol{X}\in\mathbb{R}^{n\times d}, \boldsymbol{W}_u,\boldsymbol{W}_v\in\mathbb{R}^{d\times e}, \boldsymbol{W}_z\in\mathbb{R}^{d\times s}, \boldsymbol{W}_o\in\mathbb{R}^{e\times d}. \mathcal{Q}, \mathcal{K} are simple affine transformations, and \phi is the activation function, which defaults to Swish. For further details, please refer to "FLASH: Probably the Most Interesting Efficient Transformer Design Recently".
We assume that the components of \boldsymbol{X} are independently distributed according to the standard normal distribution \mathcal{N}(0,1). The initialization distributions for \boldsymbol{W}_u, \boldsymbol{W}_v, \boldsymbol{W}_z are \mathcal{N}(0,1/d), and for \boldsymbol{W}_o it is \mathcal{N}(0,1/e), sampled independently. This is known as LeCun initialization. Its characteristic is that it keeps the output mean at 0 and maintains consistency between the second moments of the input and output. Related content can be found in the author’s previous article "A Brief Discussion on Transformer Initialization, Parameterization, and Normalization".
Basic Integrals
Under these assumptions, we estimate the distribution after each operation step. Given that LeCun initialization maintains the second moment, \boldsymbol{X}\boldsymbol{W} can be considered approximately standard normal. We can then estimate the mean and second moment after applying the activation function \phi: \begin{aligned} \mu\triangleq\mathbb{E}[\phi(\varepsilon)] =&\, \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi}}\exp\left(-\frac{1}{2}\varepsilon^2\right)\phi(\varepsilon)d\varepsilon = 0.2066\cdots \\ \nu^2\triangleq\mathbb{E}[\phi(\varepsilon)^2] =&\, \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi}}\exp\left(-\frac{1}{2}\varepsilon^2\right)\phi(\varepsilon)^2d\varepsilon = 0.3557\cdots \end{aligned} In other words, the component-wise mean and second moment of \boldsymbol{U}, \boldsymbol{V}, \boldsymbol{Z} are \mu and \nu^2, respectively. In fact, only the second moment \nu^2 is used later. For a simple estimation, one can take \nu \approx 0.6.
Self-Attention
In the initial stage, we have \mathcal{Q}(\boldsymbol{Z})=\mathcal{K}(\boldsymbol{Z})=\boldsymbol{Z}, so \boldsymbol{A}=\frac{1}{ns}\text{relu}^2\left(\boldsymbol{Z}\boldsymbol{Z}^{\top}\right). That is (for i \neq j): \begin{aligned} &\boldsymbol{A}_{i,i} = \frac{1}{ns}\text{relu}^2\big(\left\langle\boldsymbol{Z}_i, \boldsymbol{Z}_i\right\rangle\big) \approx \frac{1}{ns}\text{relu}^2\big(s\mathbb{E}[\phi(\varepsilon)^2]\big) = \frac{s\nu^4}{n} \\ &\boldsymbol{A}_{i,j} = \frac{1}{ns}\text{relu}^2\big(\left\langle\boldsymbol{Z}_i, \boldsymbol{Z}_j\right\rangle\big) \approx \frac{1}{ns}\text{relu}^2\big(s\mathbb{E}[\phi(\varepsilon)]^2\big) = \frac{s\mu^4}{n} \end{aligned} Note that \boldsymbol{A}_{i,i} / \boldsymbol{A}_{i,j} \approx \nu^4 / \mu^4 \approx 69 \gg 1. This means the diagonal elements are much larger than the non-diagonal elements. Therefore, in the initial stage, \boldsymbol{A} is actually very close to \frac{s\nu^4}{n} times the identity matrix \boldsymbol{I}: \boldsymbol{O}=(\boldsymbol{U}\odot\boldsymbol{A}\boldsymbol{V})\boldsymbol{W}_o \approx \frac{s\nu^4}{n}(\boldsymbol{U}\odot\boldsymbol{V})\boldsymbol{W}_o
Remaining Part
For \boldsymbol{U}\odot\boldsymbol{V}, it is approximately \phi(\varepsilon_i)\phi(\varepsilon_j) calculated from two independent and identically distributed variables \varepsilon_i, \varepsilon_j. Thus: \mathbb{E}[(\boldsymbol{U}\odot\boldsymbol{V})^2] \approx \mathbb{E}[\phi(\varepsilon_i)^2\phi(\varepsilon_j)^2] = \mathbb{E}[\phi(\varepsilon_i)^2]\mathbb{E}[\phi(\varepsilon_j)^2] = \nu^4 Then we have (since \boldsymbol{W}_o does not change the second moment): \mathbb{E}[\boldsymbol{O}^2] \approx \mathbb{E}\left[\left(\frac{s\nu^4}{n}\boldsymbol{U}\odot\boldsymbol{V}\right)^2\right] = \frac{s^2\nu^8}{n^2} \mathbb{E}[(\boldsymbol{U}\odot\boldsymbol{V})^2] = \frac{s^2\nu^{12}}{n^2} Therefore, the magnitude of \boldsymbol{O} is: \boldsymbol{O} = \mathcal{O}\left(\sqrt{\frac{s^2\nu^{12}}{n^2}}\right) = \mathcal{O}\left(\frac{s\nu^{6}}{n}\right) Taking a conventional pre-training setting of s=128, n=512 as an example, s\nu^6/n \approx 0.01. Thus, in the initial stage, the result coming out of \text{GAU}(\boldsymbol{x}_l) is roughly at the level of 0.01\boldsymbol{x}_l, which is two orders of magnitude smaller. Of course, this is a theoretical result; in practice, it might be larger or smaller due to random errors. However, even if it were larger, there is no need to worry because GAU also possesses the following "Crazy Scale" property.
Crazy Scale
In the reference code provided in the GAU paper’s appendix, the initialization method used by the authors is not LeCun initialization, but a normal distribution with a standard deviation of 0.02. For BERT-base, d=768, and LeCun initialization gives a standard deviation of 1/\sqrt{d} \approx 0.036. This means the standard deviation used in the appendix is only about half of that of LeCun initialization.
When we replace all \boldsymbol{W} in the GAU with \lambda \boldsymbol{W}, we have: \begin{aligned} &\tilde{\boldsymbol{U}}=\phi(\boldsymbol{X}\lambda\boldsymbol{W}_u) \approx \lambda\phi(\boldsymbol{X}\boldsymbol{W}_u)=\lambda \boldsymbol{U}\\ &\tilde{\boldsymbol{V}}=\phi(\boldsymbol{X}\lambda\boldsymbol{W}_v) \approx \lambda\phi(\boldsymbol{X}\boldsymbol{W}_v)=\lambda \boldsymbol{V}\\ &\tilde{\boldsymbol{Z}}=\phi(\boldsymbol{X}\lambda\boldsymbol{W}_z) \approx \lambda\phi(\boldsymbol{X}\boldsymbol{W}_z)=\lambda \boldsymbol{Z}\\ &\tilde{\boldsymbol{A}}=\frac{1}{ns}\text{relu}^2\left(\lambda^2\mathcal{Q}(\boldsymbol{Z})\mathcal{K}(\boldsymbol{Z})^{\top}\right) = \lambda^4 \boldsymbol{A}\\ &\tilde{\boldsymbol{O}}=(\tilde{\boldsymbol{U}}\odot\tilde{\boldsymbol{A}}\tilde{\boldsymbol{V}})\lambda\boldsymbol{W}_o \approx \lambda^7 \boldsymbol{O} \end{aligned} In other words, if all initializations are scaled down to \lambda times the original, the output of the GAU will be scaled down to \lambda^7 times the original! This is a quite "crazy" scale for GAU. Calculating with \lambda=1/2, \lambda^7 is also at the 0.01 level, shrinking the output by another two orders of magnitude! Therefore, if we follow the initialization choices of the original paper, we can theoretically train GAU models with tens of thousands of layers directly!
Summary
This article briefly analyzed the magnitude of GAU in the initial stage and concluded that GAU under standard initialization is actually close to an identity function. Therefore, it is characterized by being quite easy to train, and basically, no additional adjustment techniques are needed to train a GAU model with hundreds of layers.
Reprinting: Please include the original link: https://kexue.fm/archives/8990