Press "Enter" to skip to content

听说Attention与Softmax更配哦~

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

 

©PaperWeekly 原创 · 作者 | 苏剑林

 

单位 | 追一科技

 

研究方向 | NLP、神经网络

 

不知道大家留意到一个细节没有,就是当前 NLP 主流的预训练模式都是在一个固定长度(比如 512)上进行,然后直接将预训练好的模型用于不同长度的任务中。大家似乎也没有对这种模式有过怀疑,仿佛模型可以自动泛化到不同长度是一个“理所应当”的能力。

 

当然,笔者此前同样也没有过类似的质疑,直到前几天笔者做了 Base 版的 GAU 实验后才发现 GAU 的长度泛化能力并不如想象中好。经过进一步分析后,笔者才明白原来这种长度泛化的能力并不是“理所当然”的……

 

 

模型回顾

 

FLASH:可能是近来最有意思的高效Transformer设计 中,我们介绍了“门控注意力单元 GAU”,它是一种融合了 GLU 和 Attention 的新设计。

 

除了效果,GAU 在设计上给我们带来的冲击主要有两点:一是它显示了单头注意力未必就逊色于多头注意力,这奠定了它“快”、“省”的地位;二是它是显示了注意力未必需要 Softmax 归一化,可以换成简单的除以序列长度:

 

 

这个形式导致了一个有意思的问题:如果我们预训练的时候尽量将样本整理成同一长度(比如 512),那幺在预训练阶段 n 几乎一直就是 512,也就是说 n 相当于一个常数,如果我们将它用于其他长度(比如 64、128)微调,那幺这个 n 究竟要自动改为样本长度,还是保持为 512 呢?

 

直觉应该是等于样本长度更加自适应一些,但答案很反直觉:n 固定为 512 的微调效果比 n 取样本长度的效果要明显好!这就引人深思了……

 

 

问题定位

 

如果单看 GAU 的预训练效果,它是优于标准 Attention 的,所以 GAU 本身的拟合能力应该是没问题的,只是在样本长度方面的迁移能力不好。为了确认这一点,笔者也尝试了混合不同长度的样本来做 GAU 的预训练,发现结果会有明显的改善。

 

那幺,可能是 GAU 的什幺地方出了问题呢?其实这不难猜测,GAU 的整体运算可以简写成,其中都是 token-wise 的,也就是说它们根本不会受到长度变化的影响,所以问题只能是出现在中。

 

以前我们用标准的 Attention 时,并没有出现类似的问题,以至于我们以前都无意识地觉得这是一个“理所当然”的性质。所以,我们需要从 GAU 的 Attention 与标准 Attention 的差异中发现问题。前面说了,两者不同的地方有两点,其一是多头 Attention 变成单头 Attention,但是这顶多会让效果有一定波动,而我们测出来的结果是大幅下降,所以问题就只能出现在另一点,也就是归一化方式上,即 Attention 的 softmax 换成所带来的。

 

验证这个猜测很简单,笔者将 GAU 中 Attention 的归一化方式换回 Softmax 后重新训练一个 GAU 模型,然后微调测试不同长度的任务,发现其效果比时明显要好。所以,我们得出结论: Attention 还是与 Softmax 更配~

 

 

原因分析

 

为什幺更符合直觉的、自适应长度的 n 反而表现不如固定的 n 呢?既然我们已经以往用 Softmax 是没有这个问题的,所以我们不妨从 Softmax 出发找找灵感。Softmax 的操作是:

 

 

一个直接的问题就是: 跟 n 的关系是怎样的呢? 如果真有,那幺理论上将换成 n 应该能取得相近的效果,至少不会是特别差的那种。

 

然而,我们知道注意力的重点是“注意”,它应该有能力“聚焦”到它认为比较重要的几个 token 上。同时,以往关于高效 Transformer 的一些实验结果显示,把标准 Attention 换成 Local Attention 后结果并不会明显下降,所以我们可以预计位置为 i 的 Attention 基本上就聚焦在 i 附近的若干 token 上,超出一定距离后就基本为 0 了。事实上,也有很多事后的可视化结果显示训练好的 Attention 矩阵其实是很稀疏的。

 

综合这些结果,我们可以得出,存在某个常数 k,使得时都相当接近于 0,这样一来应该更接近而不是,这就意味着很可能跟 n 是无关的,或者说跟 n 的数量级关系至少是小于的!因此,我们如果要将替换成别的东西,那应该是一个比 n 的一次方更低阶的函数,甚至是一个常数。

 

现在回看 GAU,它的激活函数换成了时,其 Attention 情况是类似的,甚至会更稀疏。这是因为操作有直接置零的作用,不像总是正的,同时 GAU“标配”旋转位置编码 RoPE,在 Transformer升级之路:博采众长的旋转式位置编码 中我们就推导过,RoPE 本身自带一定的远程衰减的能力。综合这些条件,GAU 的归一化因子也应该是低于的阶甚至是常数级别的。

 

 

熵不变性

 

由此,我们可以总结出 GAU 的三个解决方案,一是预训练和微调都用同一个固定的 n;二是依然使用动态的样本长度 n,但是预训练时需要用不同长度的样本来混合训练,不能只使用单一长度的样本;三就是像 Softmax 那样补充上一个归一化因子,让模型自己去学:

 

 

既然存在这些解决方案,那为什幺我们还说“Attention 与 Softmax 更配”呢?GAU 的哪里不够配呢?首先,我们看 GAU 原论文的消融实验,显示出换成 Softmax,效果基本是一致的:

 

 

▲ GAU的squared_relu换成softmax效果是相近的

 

有了这个基本保证之后,我们就可以看 Softmax 比好在哪里了。我们看刚才提到的 GAU 三个解决方案,方案一总让人感觉不够自适应,方案二必须用多种长度训练显得不够优雅,至于方案三补充了归一化因子后形式上相比 Softmax 反而显得“臃肿”了。所以,总体来说还是用 Softmax 显得更为优雅有效。

 

此外,泛化能力可以简单分为“内插”和“外推”两种,在这里内插(外推)指的是测试长度小于(大于)训练长度。我们刚才说归一化因子是常数量级,更多是在内插范围内说的。对于外推来说,如果长度足够长,都“挤”在一起,所以很难保持距离超过某个范围就很接近于 0 的特性。而如果我们用 Softmax 的话,就是它可以推导出一个“熵不变性”的版本,来增强模型的外推能力:

 

 

从熵不变性看Attention的Scale操作 中我们做过简单的对比实验,显示该版本确实能提高模型在超出训练长度外的效果。

 

那幺,能否推一个“熵不变性”的版本呢?答案是不能,因为它相当于是通过温度参数来调节分布的熵,这要求激活函数不能是具备正齐次性,比如对于幂函数有,归一化后就抵消了,不起作用。激活函数最好比幂函数高一阶,才比较好实现这个调控,而比幂函数高阶的函数,最常见就是指数函数了,而指数归一化就是 Softmax 了。

 

 

实验结果

 

本文分析了 GAU 在微调效果不佳的原因,发现 Attention 的归一化因子应该是接近常数量级的,所以 GAU 用 n 或者做归一化因子会表现不佳。总的来说,笔者认为 Attention 还是跟 Softmax 更配,它是一个不错的基准,并且还可以通过“熵不变性”的拓展来进一步增强外推能力。

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。