Transformer升级之路:从Performer到线性Attention

Transformer升级之路:从Performer到线性Attention_第1张图片

©PaperWeekly 原创 · 作者|苏剑林

单位|追一科技

研究方向|NLP、神经网络

看过笔者之前的文章线性Attention的探索:Attention 必须有个 Softmax 吗?和 Performer:用随机投影将 Attention 的复杂度线性化的读者,可能会觉得本文的标题有点不自然,因为是先有线性 Attention 然后才有 Performer 的,它们的关系为“Performer 是线性 Attention 的一种实现,在保证线性复杂度的同时保持了对标准 Attention 的近似”,所以正常来说是“从线性 Attention 到 Performer”才对。

然而,本文并不是打算梳理线性 Attention 的发展史,而是打算反过来思考 Performer 给线性 Attention 所带来的启示,所以是“从 Performer 到线性 Attention”。

激活函数

线性 Attention 的常见形式是:

Transformer升级之路:从Performer到线性Attention_第2张图片

其中 、 是值域非负的激活函数。那么如何选取这个激活函数呢?Performer 告诉我们,应该选择指数函数:

首先,我们来看它跟已有的结果有什么不一样。在《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》[1] 给出的选择是:

Transformer升级之路:从Performer到线性Attention_第3张图片

我们知道 正是 在 x=0 处的一阶泰勒展开,因此 这个选择其实已经相当接近 了。

此外, 这个方案还跟《Efficient Attention: Attention with Linear Complexities》[2] 一文中引入的双重 softmax 来构建线性 Attention 的设计很相似,在那种设计中有 ,相比直接 只不过归一化的位置有所不同。

简单推导

为什么说 Performer 告诉我们激活函数的最佳选择是 呢?我们来看 Performer 找到的将标准 Attention 线性化的映射:

Transformer升级之路:从Performer到线性Attention_第4张图片

简单来说,Performer 找到了一个映射,使得 d 维向量 被映射为了 m 维向量 ,并且满足近似关系 ,此时:

Transformer升级之路:从Performer到线性Attention_第5张图片

最后一个等式表明,往 里边乘以一个常数(哪怕这个常数跟 有关),Performer 的结果完全不改变,这意味着将映射改为:

Transformer升级之路:从Performer到线性Attention_第6张图片

Performer 的结果不会有任何变化。当然,这里 这一项还不能去掉,但是如果我们假设 不会波动太大,它并不是 Attention 的主要因素,那么这一项也相当于一个常数,于是最终的映射(近似地)等价为:

Transformer升级之路:从Performer到线性Attention_第7张图片

这个看上去已经简化很多的映射该怎么理解呢?其实 m 个随机向量 拼成了一个 的矩阵,它将 d 维的 映射为了 m 维的向量,然后加上激活函数 得到了 。我们知道 Attention 的 都有一个全连接层变换,如果我们将这个 的映射矩阵整合到全连接层中,那么剩下的就是一个激活函数 了!

所以这就是最优激活函数 的来源了,只要我们将 的输出维度从 d 维改为 m 维,然后配合 的激活函数,那么理论上它就有 Performer 的拟合能力,甚至更强,因为 Performer 的 矩阵是一个固定的随机矩阵,而这里我们相当于把该矩阵也设为可训练了,还去掉了低秩约束,空间是比 Performer 更大的。

低秩问题

不管是本文的主角 Performer,还是之前在 Nyströmformer:基于矩阵分解的线性化 Attention 方案介绍的 Nyströmformer,它们的思路都是“寻找一个能逼近标准 Attention 的线性 Attention”。那么一个很自然的问题就是:标准 Attention 有什么好的?哪里值得大家向它对齐?

从信息损失的角度来看,标准 Attention 矩阵的“秩”可能更大,即更接近可逆矩阵,这意味着它能保留更多有效信息。具体来说,Attention 矩阵是一个 的矩阵,它由 通过 而来,要注意的是,这里的 d 是 Attention 的key_size,比如对于 BERT base 来说它只是 64,而 n 往往比较大,这说明 的秩不超过d,而且 ,即离满秩还远得很。

不过,softmax 的关键运算是 ,一个矩阵如果每个元素取指数的话,那么新矩阵的秩是可能增加的!所以标准 Attention 矩阵有升秩的可能性,意味着它蕴含了更有效处理信息的能力。

相比之下,线性 Attention 矩阵是 的形式,所以线性 Attention 矩阵的秩一定不超过 m,而为了弥补秩的损失,所以一般要设置 m > d,在 Performer 的实验中选择的是 m = 4d,也就是 key_size 扩大为 4 倍,秩的重要性可见一斑。当然,扩大了 key_size,一个直接的后果是处理短序列的时候,线性 Attention 还比标准 Attention 要慢,这是线性 Attention 的固有瓶颈。

关于 Attention 矩阵的秩的理论分析,也有一些论文可以参考,比如《Low-Rank Bottleneck in Multi-head Attention Models》[3] 就指出哪怕在标准 Attention 中,低秩性也是一个严重的瓶颈,增大 key_size 可以提升性能;

上个月的《Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth》[4] 则指出,如果没有残差和 FFN,那么标准 Attention 有极大的风险退化为秩等于 1 的简单变换。连标准 Attention 这个有“升秩潜力”的模型都有低秩问题,更不用说线性 Attention 这种本身秩就有上限的模型了。

所以,一句话就是:用线性 Attention 需要用更大的 key_size 来维持矩阵的秩。

集中注意

我们还可以从稀疏性角度来理解标准 Attention 的好处。直观来想,既然是“注意力机制”,那么肯定需要“集中注意力”,如果太分散,那么可能就相当于平均池化了,而“集中注意力”,意味着每个 token 应该只能显著地关联到若干个 token,用数学的话说,那就是意味着 Attention 矩阵是稀疏的,或者说至少要具备变得稀疏的可能性。

对于标准 Attention 来说,它通过 softmax 来归一化:

Transformer升级之路:从Performer到线性Attention_第8张图片

其中指数函数 起到了一个放大的作用,只要各个 本身能拉开一定差距,那么 会进一步放大这种差距,结果就是归一化之后除了最大值的那几个位置之外,剩下的概率都很接近于 0 了,这说明标准 Attention 是有潜力“集中注意力”的。

而对于线性 Attention 来说,它是直接内积的结果,没有得到 的进一步放大,所以它的注意力是比较稠密的,在序列长度较大的时候,它往往就很接近平均池化了。要缓解这一点,还是需要增大 key_size,来放大差距,直观来说,就是 n 向量放到一个低维空间太“挤”了,换到更高维的空间就“松”一些了。

怎么样验证稀疏的重要性呢?笔者曾经尝试过,将线性 Attention 的 Attention 矩阵先算出来,然后强行截断 Attention 矩阵(也就是每个token只跟前后几个 token 做 attention,变成局部形式的 Attention)让它变得稀疏,结果发现这种截断后的线性 Attention 效果明显好于全矩阵的线性 Attention。

这就肯定了稀疏的重要性了,当然,这样把 Attention 矩阵先算出来然后前行截断的方式,使得线性 Attention 的复杂度不再是线性的了,因此不具备实用价值,仅用于理论验证。

还有一个实验现象可以辅助证明稀疏的重要性,那就是线性 Attention 做语言模型或者解码器的时候,效果是跟标准 Attention 差不了多少的,这时候线性 Attention 变成了单向的 RNN(参考《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》[1] ),等价于 Attention 矩阵变成了下三角阵,也是更稀疏了。相比之下,如果用不稀疏的双向的线性 Attention 直接做 MLM 模型,则掉点会相当明显。

更重要的是,稀疏性和前一节提到的秩是有密切关联的,甚至可以说它们是“一体两面”:适当的稀疏化方法能提高矩阵的秩!比如做语言模型的下三角 Attention 矩阵,只要对角线元素非零(往往都能达到),那么这时候的矩阵直接就是满秩可逆阵了!还有笔者实验的局部 Attention 截断,也能增加矩阵的秩,比如极端情况下,每个 token 只跟自身做 attention,那么 Attention 矩阵就是满秩的单位阵了!

文章小结

本文从 Performer 出发思考了线性 Attention 的一些问题,包括关于线性 Attention 的激活函数选择,以及线性 Attention 的瓶颈所在(低秩性、稀疏性),总的结论是,线性 Attention 的最佳激活函数应当是指数函数,而有效的 Attention 机制应当具备更高的秩和更大的稀疏性。

参考文献

[1] https://arxiv.org/abs/2006.16236

[2] https://arxiv.org/abs/1812.01243

[3] https://arxiv.org/abs/2002.07028

[4] https://arxiv.org/abs/2103.03404

更多阅读

Transformer升级之路:从Performer到线性Attention_第9张图片

Transformer升级之路:从Performer到线性Attention_第10张图片

Transformer升级之路:从Performer到线性Attention_第11张图片

#投 稿 通 道#

 让你的论文被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。

总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学习心得技术干货。我们的目的只有一个,让知识真正流动起来。

???? 来稿标准:

• 稿件确系个人原创作品,来稿需注明作者个人信息(姓名+学校/工作单位+学历/职位+研究方向) 

• 如果文章并非首发,请在投稿时提醒并附上所有已发布链接 

• PaperWeekly 默认每篇文章都是首发,均会添加“原创”标志

???? 投稿邮箱:

• 投稿邮箱:[email protected] 

• 所有文章配图,请单独在附件中发送 

• 请留下即时联系方式(微信或手机),以便我们在编辑发布时和作者沟通

????

现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

关于PaperWeekly

PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域,欢迎在公众号后台点击「交流群」,小助手将把你带入 PaperWeekly 的交流群里。

你可能感兴趣的:(人工智能,深度学习,bmp,机器学习,办公软件)