众所周知,现在的Transformer越做越大,但这个“大”通常是“宽”而不是“深”,像GPT-3虽然参数有上千亿,但也只是一个96层的Transformer模型,与我们能想象的深度相差甚远。是什么限制了Transformer往“深”发展呢?可能有的读者认为是算力,但“宽而浅”的模型所需的算力不会比“窄而深”的模型少多少,所以算力并非主要限制,归根结底还是Transformer固有的训练困难。一般的观点是,深模型的训练困难源于梯度消失或者梯度爆炸,然而实践显示,哪怕通过各种手段改良了梯度,深模型依然不容易训练。

近来的一些工作(如Admin)指出,深模型训练的根本困难在于“增量爆炸”,即模型越深对输出的扰动就越大。上周的论文《DeepNet: Scaling Transformers to 1,000 Layers》则沿着这个思路进行尺度分析,根据分析结果调整了模型的归一化和初始化方案,最终成功训练出了1000层的Transformer模型。整个分析过程颇有参考价值,我们不妨来学习一下。

增量爆炸 #

原论文的完整分析比较长,而且有些假设或者描述细酌之下是不够合理的。所以在本文的分享中,笔者会尽量修正这些问题,试图以一个更合理的方式来得到类似结果。

假设损失函数为$\mathcal{L}(\boldsymbol{\theta})$,$\boldsymbol{\theta}$是它的参数,考虑参数由$\boldsymbol{\theta}$变为$\boldsymbol{\theta}+\Delta\boldsymbol{\theta}$时损失函数的增量:
\begin{equation}\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\end{equation}
对于SGD有$\Delta\boldsymbol{\theta}=-\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$,那么$\Delta\mathcal{L} \approx -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert^2$。设模型有$N$层,每层的平均参数量为$K$,配合Xavier初始化以及各种Normalization手段,我们可以使得多数参数的梯度是$\mathscr{O}(1)$量级,所以有$\Delta\mathcal{L}=\mathscr{O}(\eta NK)$。因此,模型每一步的更新量是正比于模型深度$N$的(宽度不在本文讨论范围),如果模型越深,那么更新量就越大,这意味着初始阶段模型越容易进入不大好的局部最优点,然后训练停滞甚至崩溃,这就是“增量爆炸”问题。

这时候解决方法有两个,一是初始阶段用更小的学习率进行训练(不超过$\eta/N$量级),然后慢慢增大学习率,这就是Warmup技巧;二就是调整初始化方案,使得参数的梯度是$\mathscr{O}(1/\sqrt{N})$量级,这样就自动抵消掉模型深度的影响。

量级分析 #

怎么做到第二种方案呢?我们可以尝试分析Transformer的梯度。然而,精确的梯度求起来比较繁琐,并且事实上我们也不需要精确的梯度,而只是要对梯度做一个量级分析,所以我们可以用如下的“量级分解”技巧转化为标量的导数问题。

对于一个矩阵$\boldsymbol{W}$,我们将其分解为$\boldsymbol{W}=\lambda \boldsymbol{U}$的形式,其中
\begin{equation}\lambda = \mathop{\arg\min}_{\kappa > 0} \Vert \boldsymbol{W}\boldsymbol{W}^{\top}/\kappa^2 - \boldsymbol{I}\Vert,\quad \end{equation}
说白了,我们就是要将一个矩阵分解为一个标量$\lambda$与一个尽可能正交的矩阵$\boldsymbol{U}$之积。由于$\boldsymbol{U}$接近正交矩阵,它起到了一个标准参考系的作用,而对应的$\lambda$则代表了矩阵$\boldsymbol{W}$的量级。如果$\boldsymbol{W}$使用Xavier初始化,那么$\lambda$相当于其中的gain参数,即在Xavier初始化的基础上还要再乘一个$\lambda$。这是因为Xavier初始化的结果就接近一个正交矩阵,这一点可以参考《从几何视角来理解模型参数的初始化策略》

在此分解之下,我们有
\begin{equation}\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\end{equation}
这意味着$\frac{\partial \mathcal{L}}{\partial \lambda}$跟$\frac{\partial \mathcal{L}}{\partial \boldsymbol{W}}$在量级上是成正比的,所以对$\frac{\partial \mathcal{L}}{\partial \lambda}$做量级分析就相当于对$\frac{\partial \mathcal{L}}{\partial \boldsymbol{W}}$做量级分析。这样$\frac{\partial \mathcal{L}}{\partial \lambda}$就相当于$\frac{\partial \mathcal{L}}{\partial \boldsymbol{W}}$量级的一个简单的“探针”,原来的矩阵求导就可以转化为标量求导,降低了分析难度。

前馈梯度 #

很多实验结果都显示虽然Pre Norm比Post Norm更容易训练,但Post Norm的最终效果往往更好些,所以原论文保留了Post Norm结构,并考虑了更一般的形式(DeepNorm):
\begin{equation}\boldsymbol{x}_{l+1} = \text{LN}(\alpha\boldsymbol{x}_l + F(\boldsymbol{x}_l)) = \text{LN}(\boldsymbol{x}_l + F(\boldsymbol{x}_l)/\alpha)\end{equation}
其中$\alpha > 0$是一个常数。简单起见,我们先考虑FFN层,此时
\begin{equation}\boldsymbol{x}_{l+1} = \text{LN}(\boldsymbol{x}_l + \phi(\boldsymbol{x}_l \boldsymbol{W}_1)\boldsymbol{W}_2/\alpha)\end{equation}
这里的$\phi$是激活函数,一般为ReLU或其变体(Swish、GeLU等),它们(近似)满足$\phi(\lambda x) = \lambda \phi(x),\forall \lambda > 0$。使用前一节的量级分解探针,我们得到
\begin{equation}\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{记为}\boldsymbol{z}_{l+1}})\label{eq:ffn}\end{equation}
求$\lambda$的梯度:
\begin{equation}\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}\end{equation}
我们断言$\frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}$、$\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}$都是$\mathscr{O}(1)$的,并且由于$\boldsymbol{U}_1$、$\boldsymbol{U}_2$都接近正交矩阵,所以$\phi(\boldsymbol{x}_l \boldsymbol{U}_1)\boldsymbol{U}_2$也是$\mathscr{O}(1)$的,因此最终有
\begin{equation}\frac{\partial \mathcal{L}}{\partial \lambda_1} = \mathscr{O}\left(\frac{\lambda_2}{\alpha}\right),\quad \frac{\partial \mathcal{L}}{\partial \lambda_2} = \mathscr{O}\left(\frac{\lambda_1}{\alpha}\right)\end{equation}

自注意力 #

现在考虑自Self Attention,作为量级分析,我们考虑单头注意力即可,其形式为
\begin{equation}\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)\end{equation}
其中$\sigma(\cdot)$是softmax操作的简写,这里省略了Attention的scale操作。对上式进行量级分解后的形式为
\begin{equation}\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{记为}\boldsymbol{z}_{l+1}})\label{eq:sa}\end{equation}
现在我们可以对各个$\lambda$分别求梯度,而由于softmax的存在,事实上$\lambda_q,\lambda_k$的梯度本身会很小,不会明显影响最终的更新量,所以其实我们考虑$\lambda_v,\lambda_o$的更新量足矣:
\begin{equation}\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}\end{equation}
同样断言$\frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}$、$\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}$都是$\mathscr{O}(1)$的,并且注意softmax出来是一个概率分布,然后对$\boldsymbol{x}_l$的各个token做加权平均,通常而言,平均前后的向量会在同一数量级,所以我们认为$\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$也是$\mathscr{O}(1)$的,因此结果跟FFN层的类似:
\begin{equation}\frac{\partial \mathcal{L}}{\partial \lambda_v} = \mathscr{O}\left(\frac{\lambda_o}{\alpha}\right),\quad \frac{\partial \mathcal{L}}{\partial \lambda_o} = \mathscr{O}\left(\frac{\lambda_v}{\alpha}\right)\end{equation}

初步结论 #

现在不管是FFN还是Self Attention,我们都得到了相似的结论,现在简单起见,假设每个参数的量级(至少在初始化阶段)是一致的,即所有的$\lambda$取同一个值,那么总的结论是
\begin{equation}\frac{\partial \mathcal{L}}{\partial \lambda} = \mathscr{O}\left(\frac{\lambda}{\alpha}\right)\end{equation}
即梯度的量级是$\mathscr{O}(\lambda/\alpha)$。另一方面,我们说$N$层的Transformer模型,一般是$N$层的Self Attention加$N$层的FFN,所以严格来说层数是$2N$。因此,按照“增量爆炸”一节的分析,我们需要将梯度调整到$\mathscr{O}(1/\sqrt{2N})$,上式告诉我们可以通过让$\lambda/\alpha=1/\sqrt{2N}$来实现。原论文的放缩更为宽松一些,得到的结果是$\lambda/\alpha = 1/\sqrt{4N}$,量级上是等价的。

现在我们得到了$\lambda$与$\alpha$的一个比例关系,但无法直接得到$\lambda$和$\alpha$的具体值。按照论文的说法,是从对称角度出发,让$\lambda=1/\alpha$,从而可以解得
\begin{equation}\alpha = (2N)^{1/4},\quad \lambda = (2N)^{-1/4}\label{eq:result}\end{equation}
然而,单纯对称的解释显然是不够说服力的,我们需要搞清楚不同的选择究竟有什么不同的结果。为此,我们可以比较另外两组解:

另解一:$\alpha=1,\lambda=(2N)^{-1/2}$,此时参数的初始化缩小到原来的$(2N)^{-1/2}$倍,梯度也被缩小到原来的$(2N)^{-1/2}$倍,根据SGD的$\Delta\boldsymbol{\theta}=-\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$得出每步的更新量也是原来的$(2N)^{-1/2}$倍,也就是说,调整前后的相对学习幅度是没有变化的,因此有可能刚开始$\lambda=\mathscr{O}((2N)^{-1/2})$级别,但训练集几步后就脱离了这个量级了。

另解二:$\alpha=(2N)^{1/2},\lambda=1$,此时参数的初始化没有缩小,但梯度也被缩小到原来的$(2N)^{-1/2}$倍,根据SGD的$\Delta\boldsymbol{\theta}=-\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$得出每步的更新量也是原来的$(2N)^{-1/2}$倍,调整前后的相对学习幅度是明显缩小了,因此有可能出现学习得非常慢的情况。

这两种情况看上去都各有缺点,因此介乎两者之间的式$\eqref{eq:result}$似乎就能解释得通了,它就是保持梯度缩放到原来的$(2N)^{-1/2}$倍的同时,让初始学习步伐稍微慢一些,但又不至于太慢,隐式地起到了Warmup的作用。

多种优化 #

上面的分析都是基于SGD进行的,但事实上我们很少直接用SGD去训练NLP模型,我们更多是自适应学习率优化器,主要有两大类:一是用二阶矩来校正学习率,Adam、AdamW等都属此类;另一类是通过参数模长进一步校正学习率,比如LAMBAdaFactor。原论文的说法是“我们在SGD上进行推导,然后在Adam上验证发现也还可以”,但从理论上来讲,它们并不完全通用,这一节我们就来针对性地做一下分析。

对于Adam类优化器来说,每一步的更新量大约为$\Delta\boldsymbol{\theta}=-\eta\,\text{sign}(\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}))$,所以$\Delta\mathcal{L} \approx -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert_1$,它是正比于梯度的1次方而不是2次方,因此要过要让更新量跟层数无关,那么梯度应该缩小到原来的$1/(2N)$倍才对,即应该有$\lambda/\alpha=1/(2N)$,如果同样让$\lambda=1/\alpha$,那么有
\begin{equation}\alpha = (2N)^{1/2},\quad \lambda = (2N)^{-1/2}\end{equation}

对于LAMB类优化器来说,每一步更新量大约为$\Delta\boldsymbol{\theta}=-\eta\Vert\theta\Vert\,\text{sign}(\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}))$,所以$\Delta\mathcal{L} \approx -\eta\Vert\theta\Vert\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert_1$,注意到参数的缩放比例是$\lambda$、梯度的缩放比例是$\lambda/\alpha$,所以$\Delta\mathcal{L}=\mathscr{O}(2N\lambda^2/\alpha)$,从而是$\lambda^2/\alpha=1/(2N)$。注意这类优化器每步的相对更新量是一样的(等于学习率$\eta$),不管怎么调整$\alpha,\lambda$其相对更新大小都不会变化,所以我们可以直接取$\alpha=1,\lambda=(2N)^{-1/2}$。

结果汇总对比如下:
\begin{array}{c|cc|cc}
\hline
\text{优化器} & \Delta\boldsymbol{\theta} & \Delta\mathcal{L} & \alpha & \lambda \\
\hline
\text{SGD} & -\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}) & -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert^2 & (2N)^{1/4} & (2N)^{-1/4}\\
\text{Adam} & -\eta\,\text{sign}(\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})) & -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert_1 & (2N)^{1/2}& (2N)^{-1/2}\\
\text{LAMB} & -\eta\Vert\theta\Vert\,\text{sign}(\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})) & -\eta\Vert\theta\Vert\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert_1 & 1 & (2N)^{-1/2}\\
\hline
\end{array}

事后分析 #

前面的两节推导过程都用到了断言“$\frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}$、$\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}$都是$\mathscr{O}(1)$的”,那么它是否成立呢?这里我们事后分析一下。

其实也很简单,经过前述调整后,不管是FFN层$\eqref{eq:ffn}$还是Self Attention层$\eqref{eq:sa}$,初始阶段每个残差分支的权重被缩放到原来的$\lambda^2/\alpha$倍,不管是哪种优化器的结果,$\lambda^2/\alpha$都是一个比较小的数字,这意味着初始阶段整个模型其实接近一个恒等函数,因此$\frac{\partial \mathcal{L}}{\partial \boldsymbol{x}_{l+1}}$、$\frac{\partial \boldsymbol{x}_{l+1}}{\partial \boldsymbol{z}_{l+1}}$自然都是$\mathscr{O}(1)$的,所以结论和断言是自洽的。

另外,可能有读者想问同样的分析是否可以用到Pre Norm结构上呢?答案是可以的,并且结论是基本一致的,只是因为Norm放在了残差分支之前,所以就没必要设置$\alpha$参数了,所以结论就是上述关于Post Norm的结果中所有的$\alpha$都等于为1,然后重新计算相应的$\lambda$。

最后,读者可能有疑问的是花了那么多功夫讨论把模型做深,那么模型深度真有那么重要吗?有,原论文给出了一个漂亮的实验结果,用一个200层的“深而窄”的模型(32亿参数),战胜了之前48层“浅而宽”的SOTA模型(120亿参数):

“深而窄”的模型胜于“浅而宽”的模型

“深而窄”的模型胜于“浅而宽”的模型

文章小结 #

本文分析了将Transformer做“深”的瓶颈所在并给出了相应的解决方案,文章的主要思路源于微软新出的DeepNet,并对原论文的分析过程做了一定的简化和完善。

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

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

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Mar. 09, 2022). 《训练1000层的Transformer究竟有什么困难? 》[Blog post]. Retrieved from https://kexue.fm/archives/8978

@online{kexuefm-8978,
        title={训练1000层的Transformer究竟有什么困难?},
        author={苏剑林},
        year={2022},
        month={Mar},
        url={\url{https://kexue.fm/archives/8978}},
}