大家知道Layer Normalization是Transformer模型的重要组成之一,它的用法有PostLN和PreLN两种,论文《On Layer Normalization in the Transformer Architecture》中有对两者比较详细的分析。简单来说,就是PreLN对梯度下降更加友好,收敛更快,对训练时的超参数如学习率等更加鲁棒等,反正一切都好但就有一点硬伤:PreLN的性能似乎总略差于PostLN。实际上大家都用相同的参数,学习相同的步数,最终PreLN的效果是要优于PostLN的,但是将两者各自调试到最优的情况下,PostLN的效果通常好些
最近Google的一篇论文《RealFormer: Transformer Likes Residual Attention》提出了RealFormer设计,成功地弥补了这个Gap,使得模型拥有与PreLN一样的优化友好性,并且效果比PostLN还好,可谓"鱼与熊掌兼得"了
形式
RealFormer全称为"Residual Attention Layer Transformer",即"残差式Attention层的Transformer模型",顾名思义就是把残差放到Attention里了(原有的残差结构仍然保留)
说回模型,如上图所示,RealFormer主要是把残差放到了Attention矩阵上面了,而整体依然保持了PostLN的结构,因此既保持了PostLN的性能,又融合了残差的友好。具体来说,就是原来第$n$层的Attention为
$$ \text{Attention}(\boldsymbol{Q}_n, \boldsymbol{K}_n, \boldsymbol{V}_n) = \text{Softmax}(\boldsymbol{A}_n)\boldsymbol{V}_n,\quad \boldsymbol{A}_n = \frac{\boldsymbol{Q}_n\boldsymbol{K}_n^T}{\sqrt{d_k}}\tag{1} $$
现在改为了
$$ \text{Attention}(\boldsymbol{Q}_n, \boldsymbol{K}_n, \boldsymbol{V}_n) = \text{Softmax}(\boldsymbol{A}_n)\boldsymbol{V}_n,\quad \boldsymbol{A}_n = \frac{\boldsymbol{Q}_n\boldsymbol{K}_n^T}{\sqrt{d_k}}+\boldsymbol{A}_{n-1}\tag{2} $$
分析
这一节我们对RealFormer做一个简单的思考分析。RealFormer对梯度下降更加友好,这不难理解,因为$\boldsymbol{A}_n = \frac{\boldsymbol{Q}_n\boldsymbol{K}_n^T}{\sqrt{d_k}}+\boldsymbol{A}_{n-1}$的设计确实提供了一条直通路,使得第一层的Attention能够直通最后一层,自然就没有什么梯度消失的风险了。相比之下,PostLN是$\text{LayerNorm}(x+f(x))$的结构,看上去$x+f(x)$防止了梯度消失,但是LayerNorm这一步会重新增加了梯度消失风险,造成的后果是初始阶段前面的层梯度很小,后面的层梯度很大,如果用大学习率,后面的层容易崩,如果用小学习率,前面的层学不好,因此PostLN更难训练,需要用小的学习率加warmup慢慢训
那么PreLN改善了梯度状况,为什么又比不上PostLN呢?按照笔者的猜测,PreLN每一步都是$x+f(x)$的形式,到了最后一层就变成了$x+f_1(x)+f_2(x)+⋯+f_n(x)$的形式,一层层累加,可能导致数值和方差都很大,所以最后"迫不得已"会强制加一层Layer Norm让输出稳定下来。这样,尽管PreLN改善了梯度状况,但它本身设计上就存在一些不稳定因素,也许这就是它效果略差的原因
读者可能会反驳,既然PreLN存在问题,那RealFormer的$\boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^T}{\sqrt{d_k}}+\boldsymbol{A}_{n-1}$不也是存在同样的叠加问题吗?如果只看$\boldsymbol{A}$,确实会有这样的问题,但别忘了$\boldsymbol{A}$后面还要做个softmax归一化后才参与运算,也就是说,模型对矩阵$\boldsymbol{A}$是自带归一化功能的,所以它不会有数值发散的风险。而且刚刚相反,随着层数的增加,$\boldsymbol{A}$的叠加会使得$\boldsymbol{A}$的元素绝对值可能越来越大,Attention逐渐趋于one hot形式,造成后面层梯度消失,但是别忘了,我们刚才说PostLN前面层梯度小后面层梯度大,而现在也进一步缩小了后面层梯度,反而使得两者更同步从而更好优化了;另一方面,Attention的概率值可能会有趋同的趋势,也就是说Attention的模式可能越来越稳定,带来类似ALBERT参数共享的正则化效应,这对模型效果来说可能是有利的。同时,直觉上来想,用RealFormer结构去做FastBERT之类的自适应层数的改进,效果会更好,因为RealFormer的Attention本身会有趋同趋势,更加符合FastBERT设计的出发点
此外,我们也可以将RealFormer理解为还是使用了常规的残差结构,但是残差结构只用在$\boldsymbol{Q},\boldsymbol{K}$而没有用在$\boldsymbol{V}$上:
$$ \begin{aligned} &\text{Attention}(\boldsymbol{Q}_n,\boldsymbol{K}_n,\boldsymbol{V}_n)=\text{Softmax}(\boldsymbol{A}_n)\boldsymbol{V}_n\\ &\boldsymbol{A}_n=\frac{\tilde{\boldsymbol{Q}}_n\tilde{\boldsymbol{K}}_n^T}{\sqrt{d_k}},\quad \tilde{\boldsymbol{Q}}_n=\boldsymbol{Q}_n+\tilde{\boldsymbol{Q}}_{n-1},\quad \tilde{\boldsymbol{K}}_n = \tilde{\boldsymbol{K}}_n+\tilde{\boldsymbol{K}}_{n-1} \end{aligned}\tag{3} $$
这在一定程度上与$\boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^T}{\sqrt{d_k}}+\boldsymbol{A}_{n-1}$是等价的,而PreLN相当于$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$都加了残差。为啥$\boldsymbol{V}$"不值得"一个残差呢?从近来的一些相对位置编码的改进中,似乎有一个共同的趋势,那就是去掉了$\boldsymbol{V}$的偏置,比如像NEZHA的相对位置编码,是同时在Attention矩阵(即$\boldsymbol{Q},\boldsymbol{K}$)和$\boldsymbol{V}$上施加的,而较新的XLNet和T5的相对位置编码则只施加在Attention矩阵上,所以,似乎去掉$\boldsymbol{V}$的不必要的偏置是一个比较好的选择,而RealFormer再次体现了这一点