门控注意力单元（GAU）还需要Warmup吗？

在文章《训练1000层的Transformer究竟有什么困难？》发布之后，很快就有读者问到如果将其用到《FLASH：可能是近来最有意思的高效Transformer设计》中的“门控注意力单元（GAU）”，那结果是怎样的？跟标准Transformer的结果有何不同？本文就来讨论这个问题。

先说结论 #

事实上，GAU是非常容易训练的模型，哪怕我们不加调整地直接使用“Post Norm + Xavier初始化”，也能轻松训练个几十层的GAU，并且还不用Warmup。所以关于标准Transformer的很多训练技巧，到了GAU这里可能就无用武之地了...

为什么GAU能做到这些？很简单，因为在默认设置之下，理论上$\text{GAU}(\boldsymbol{x}_l)$相比$\boldsymbol{x}_l$几乎小了两个数量级，所以
\begin{equation}\boldsymbol{x}_{l+1} = \text{LN}(\boldsymbol{x}_l + \text{GAU}(\boldsymbol{x}_l))\approx \boldsymbol{x}_l\end{equation}
因此，GAU配合残差，在标准的初始化之下就已经很接近一个恒等函数，有这种性质的模型是非常容易训练的，通常都不需要Warmup。如果要对应上《训练1000层的Transformer究竟有什么困难？》的结论，这两个数量级相当于$\lambda=1,\alpha=100$，意味着它自动地包含了上百层的模型DeepNorm操作，因此理论上我们可以直接训练上百层的GAU模型而不需要特别的调整技巧。

模型假设 #

其实我们只需要对GAU的输入和输出做一个量级分析就行了。标准的GAU运算如下：
\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}
其中$\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}$是简单的仿射变换，$\phi$是激活函数，默认是Swish。如果还有不清楚的地方，可以参考《FLASH：可能是近来最有意思的高效Transformer设计》。

我们假设$\boldsymbol{X}$的各个分量独立地服从标准正态分布$\mathcal{N}(0,1)$，然后$\boldsymbol{W}_u,\boldsymbol{W}_v,\boldsymbol{W}_z$的初始化分布是$\mathcal{N}(0,1/d)$而$\boldsymbol{W}_o$的初始化分布则是$\mathcal{N}(0,1/e)$独立重复采样出来的，这种初始化分布被称为LeCun初始化，它的特点是能让输出的均值为0，并且保持输入输出的二阶矩一致，相关内容可以参考笔者之前的文章《浅谈Transformer的初始化、参数化与标准化》。

基本积分 #

在这些假设之下，我们来逐一估计每步运算之后的分布。结合假设，由于LeCun初始化能保持二阶矩不变，所以$\boldsymbol{X}\boldsymbol{W}$也可以近似认为是标准正态分布的，于是我们可以用下面的式子估计加了激活函数$\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}
换言之，$\boldsymbol{U},\boldsymbol{V},\boldsymbol{Z}$的分量均值和二阶矩分别是$\mu$和$\nu^2$，事实上后面只用到了二阶矩$\nu^2$，简单估计时，取$\nu=0.6$就行了。

自注意力 #

在初始阶段，我们有$\mathcal{Q}(\boldsymbol{Z})=\mathcal{K}(\boldsymbol{Z})=\boldsymbol{Z}$，所以初始阶段有$\boldsymbol{A}=\frac{1}{ns}\text{relu}^2\left(\boldsymbol{Z}\boldsymbol{Z}^{\top}\right)$，即（下面$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}
注意到$\boldsymbol{A}_{i,i} / \boldsymbol{A}_{i,j} \approx \nu^4 / \mu^4 \approx 69 \gg 1$，也就是对角线元素远远大于非对角线元素，因此初始阶段的$\boldsymbol{A}$其实很接近单位阵的$\frac{sv^4}{n}$倍，即$\boldsymbol{A}\approx \frac{sv^4}{n}\boldsymbol{I}$，于是
\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}

剩余部分 #

对于$\boldsymbol{U}\odot\boldsymbol{V}$，它近似于两个独立同分布的变量$\varepsilon_i,\varepsilon_j$算出来的$\phi(\varepsilon_i)\phi(\varepsilon_j)$，所以
\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}
于是有（$\boldsymbol{W}_o$不改变二阶矩）
\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}
因此$\boldsymbol{O}$的量级是
\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}
以常规的预训练设置$s=128,n=512$为例，$s\nu^6/n\approx 0.01$，因此在初始阶段经过$\text{GAU}(\boldsymbol{x}_l)$后出来的结果大致是$0.01\boldsymbol{x}_l$这个级别的，小两个数量级。当然，这是理论结果，实际上由于随机误差原因可能会更大或更小，不过就算更大了也不用担心，因为GAU还有下面的“疯狂尺度”性质。

疯狂尺度 #

在GAU论文的附录参考代码中，作者所用的初始化方法还不是LeCun初始化，而是0.02标准差的正态分布。对于BERT base来说$d=786$，LeCun初始化给出的标准差是$1/\sqrt{d}\approx 0.036$，也就是说附录所用的初始化标准差大约只有LeCun初始化的一半。

当我们将GAU中所有的$\boldsymbol{W}$都换成$\lambda \boldsymbol{W}$时，我们将有
\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}
也就是说，如果所有初始化都缩小到原来的$\lambda$倍，那么GAU的输出将会缩小到原来的$\lambda^7$倍！这是关于GAU的一个相当疯狂的Scale，按照$\lambda=1/2$算，$\lambda^7$同样是0.01级别，再次缩小了两个数量级！所以，如果按照原论文的初始化选择，我们理论上可以直接训练上万层的GAU模型！

本文小结 #

本文主要简单分析了一下GAU在初始阶段的数量级，得出标准初始化下的GAU其实已经接近恒等函数，因此具有相当容易训练的特点，基本上训练上百层的GAU模型也用不着额外的调整。

转载到请包括本文地址：https://kexue.fm/archives/8990

更详细的转载事宜请参考：《科学空间FAQ》