Does the Gated Attention Unit (GAU) still need Warmup?

门控注意力单元（GAU）还需要Warmup吗？ 苏剑林 · 2022-03-11 · Machine Learning model optimization attention

Right after publishing the article "What's So Hard About Training a 1000-Layer Transformer?", several readers quickly asked what would happen if we applied its ideas to the "Gated Attention Unit (GAU)" from "FLASH: Perhaps the Most Interesting Efficient Transformer Design in Recent Times". How would the result compare to that of a standard Transformer? This post addresses exactly that question.

The conclusion first

As it turns out, GAU is an extremely easy model to train — even if we use "Post Norm + Xavier initialization" without any adjustments, we can effortlessly train a GAU with dozens of layers, and without Warmup at that. So many of the training tricks devised for standard Transformers may simply be unnecessary here...

Why can GAU pull this off? Simply put, under the default settings, $\text{GAU}(\boldsymbol{x}_l)$ is theoretically almost two orders of magnitude smaller than $\boldsymbol{x}_l$, so

\begin{equation}\boldsymbol{x}_{l+1} = \text{LN}(\boldsymbol{x}_l + \text{GAU}(\boldsymbol{x}_l))\approx \boldsymbol{x}_l\end{equation}more

As a result, GAU together with the residual connection is already very close to an identity function under standard initialization, and models with this property are generally very easy to train — they typically don't need Warmup at all. If we relate this back to the conclusion of "What's So Hard About Training a 1000-Layer Transformer?", these two orders of magnitude correspond to $\lambda=1,\alpha=100$, which means it automatically incorporates the equivalent of DeepNorm operations for a model with over a hundred layers. So in theory, we should be able to directly train a GAU model with hundreds of layers without any special adjustment tricks.

Model assumptions

All we really need to do is a magnitude analysis of the input and output of GAU. The standard GAU computation is as follows:

\begin{equation}\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}\end{equation}

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, Swish by default. If anything here is unclear, you can refer to "FLASH: Perhaps the Most Interesting Efficient Transformer Design in Recent Times".

We assume that each component of $\boldsymbol{X}$ independently follows a standard normal distribution $\mathcal{N}(0,1)$, and that the initialization distribution of $\boldsymbol{W}_u,\boldsymbol{W}_v,\boldsymbol{W}_z$ is $\mathcal{N}(0,1/d)$, while that of $\boldsymbol{W}_o$ is obtained by independently sampling repeatedly from $\mathcal{N}(0,1/e)$. This kind of initialization distribution is called LeCun initialization, and its characteristic feature is that it keeps the output mean at 0 while preserving the second moment between input and output. For related background, see the author's earlier post "A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers".

Basic integrals

Under these assumptions, let's estimate the distribution after each step of the computation one by one. Given our assumptions, since LeCun initialization preserves the second moment, $\boldsymbol{X}\boldsymbol{W}$ can also be approximately regarded as standard normal. So we can use the following expression to estimate the mean and second moment after applying the activation function $\phi$:

\begin{equation}\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}\end{equation}

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$ will be used later; for a rough estimate, we can just take $\nu=0.6$.

Self-attention

At the initial stage, we have $\mathcal{Q}(\boldsymbol{Z})=\mathcal{K}(\boldsymbol{Z})=\boldsymbol{Z}$, so at this stage $\boldsymbol{A}=\frac{1}{ns}\text{relu}^2\left(\boldsymbol{Z}\boldsymbol{Z}^{\top}\right)$ holds, that is (below, $i\neq j$)

\begin{equation}\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{sv^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}\end{equation}

Note that $\boldsymbol{A}_{i,i} / \boldsymbol{A}_{i,j} \approx \nu^4 / \mu^4 \approx 69 \gg 1$, i.e., the diagonal elements are far larger than the off-diagonal ones. So at the initial stage, $\boldsymbol{A}$ is in fact very close to $\frac{sv^4}{n}$ times the identity matrix, i.e., $\boldsymbol{A}\approx \frac{sv^4}{n}\boldsymbol{I}$, giving us

\begin{equation}\boldsymbol{O}=(\boldsymbol{U}\odot\boldsymbol{A}\boldsymbol{V})\boldsymbol{W}_o\approx \frac{sv^4}{n}(\boldsymbol{U}\odot\boldsymbol{V})\boldsymbol{W}_o\end{equation}

The remaining part

For $\boldsymbol{U}\odot\boldsymbol{V}$, it is approximately $\phi(\varepsilon_i)\phi(\varepsilon_j)$ computed from two i.i.d. variables $\varepsilon_i,\varepsilon_j$, so

\begin{equation}\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\end{equation}

which gives us ($\boldsymbol{W}_o$ doesn't change the second moment)

\begin{equation}\mathbb{E}[\boldsymbol{O}^2] \approx \mathbb{E}\left[\left(\frac{sv^4}{n}\boldsymbol{U}\odot\boldsymbol{V}\right)^2\right] = \mathbb{E}[\phi(\varepsilon_i)^2\phi(\varepsilon_j)^2] = \frac{s^2\nu^{12}}{n^2}\end{equation}

Hence the magnitude of $\boldsymbol{O}$ is

\begin{equation}\boldsymbol{O} = \mathcal{O}\left(\sqrt{\frac{s^2\nu^{12}}{n^2}}\right) = \mathcal{O}\left(\frac{s\nu^{6}}{n}\right) \end{equation}

Taking the typical pretraining setting $s=128,n=512$ as an example, $s\nu^6/n\approx 0.01$, so at the initial stage, after passing through $\text{GAU}(\boldsymbol{x}_l)$ the result comes out to roughly the level of $0.01\boldsymbol{x}_l$ — two orders of magnitude smaller. Of course, this is a theoretical result; in practice, due to random fluctuations, it could turn out larger or smaller. But even if it comes out larger, there's no need to worry, because GAU has yet another "crazy scale" property, described below.

Crazy scale

In the reference code in the appendix of the GAU paper, the initialization method the authors actually used isn't LeCun initialization but rather a normal distribution with standard deviation 0.02. For BERT base, $d=786$, LeCun initialization gives a standard deviation of $1/\sqrt{d}\approx 0.036$ — meaning that the initialization standard deviation used in the appendix is only about half that of LeCun initialization.

When we replace all instances of $\boldsymbol{W}$ in GAU with $\lambda \boldsymbol{W}$, we get

\begin{equation}\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}\end{equation}

In other words, if all initializations are scaled down to $\lambda$ of their original value, the output of GAU shrinks to $\lambda^7$ of its original value! This is a rather wild scaling behavior for GAU: computing it out with $\lambda=1/2$, $\lambda^7$ again comes out at the 0.01 level — another two orders of magnitude smaller! So, following the initialization choice from the original paper, we should in theory be able to directly train GAU models with tens of thousands of layers!

Summary

This post has given a simple analysis of the magnitude of GAU at the initial stage, showing that GAU under standard initialization is already close to an identity function, and consequently has the property of being remarkably easy to train — training a GAU model with hundreds of layers essentially requires no extra adjustments at all.