Transformer新内核Synthesizer：低复杂度的attention代替点乘式的注意力机制

论文链接：https://arxiv.org/pdf/2005.00743.pdf
参考文档：https://mp.weixin.qq.com/s/e2Cor8amz7GiFfBGdLXULg

导读

今天阅读的是一篇来自Google的最新论文《SYNTHESIZER: Rethinking Self-Attention in Transformer Models》，该论文重新探索了Transformer中注意力机制的必要性，并引入了新的attention计算方法Synthesizer。实验显示，即使不进行token之间的attention交互计算，synthesizer在翻译、语言模型、GLUE等任务上也可以达到很好的效果。

1、引言

1.1、什么是自注意力机制

自注意力机制算是解释性比较强的模型之一，它通过直接把序列两两比较（代价是计算量变为 $O(n^2)$O(n2)，当然由于是纯矩阵运算，这个计算量相当也不是很严重），能够一步到位捕捉到全局的联系。相比之下，RNN 需要一步步递推才能捕捉到，而 CNN 则需要通过层叠来扩大感受野，这是 Self Attention 层的明显优势。

自注意力机制到底是如何生效的？这种“token”对“token”的注意力是必须的吗？$Attention(Q,K,V)=softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)V$Attention(Q,K,V)=softmax(dk​​QKT​)V$Self-Attention(X)=Attention(XW_Q,XW_K,XW_V)=softmax\left(\frac{XW_QW_K^TX^T}{\sqrt{d_k}}\right)XW_V$Self−Attention(X)=Attention(XWQ​,XWK​,XWV​)=softmax(dk​​XWQ​WKT​XT​)XWV​本质上来看，自注意力就是通过一个$n×n$n×n的矩阵$A$A和$d×d'$d×d′的矩阵$W_v$Wv​，将原本是$n×d$n×d的矩阵$X$X，变成了$n×d'$n×d′的矩阵$AXW_v$AXWv​。其中矩阵A是动态生成的，即$A=softmax(B),B=\frac{XW_QW_K^TX^T}{\sqrt{d_k}}$A=softmax(B),B=dk​​XWQ​WKT​XT​对于矩阵$B$B，本质上来说它就是 $X$X 里边两两向量的内积组合，所以称为“token对token”的Attention。

至于Multi-Head Attention，则不过是Attention运算在不同的参数下重复多次然后将多个输出拼接起来，属于比较朴素的增强。

那么，就到了前面提出的问题：“token对token”是必须的吗？能不能通过其他方式来生成这个矩阵$B$B？Google的这篇论文正是“天马行空”了几种新的形式并做了实验，这些形式统称为Synthesizer。

1.2、Transformer存在的问题

点积自注意力提供了强大的建模能力，但同时作者对点积自注意力提出了质疑，它的计算可能是不必要的。

点积的基本作用是确定单个 token 相对于序列中所有其他 token 的相对重要性。key、query、value 暗示自注意力模拟一个基于内容的检索过程，过程的核心是 pairwise 的交互。该文章对整个过程进行了反思。

2、模型简介

2.1、Synthesizer的关键思想

Synthesizer 的核心思想是用低复杂度的attention计算代替dot product式的注意力机制。传统 Transformer 的注意力机制需要进行 token 之间的两两交互，虽然可以获得更多的交互信息，但attention score会很依赖实例，难以保证模型学到更多的泛化特征。

因此，synthesizer提出了新的attention matrix学习方法，只通过简单的前馈神经网络即可得到注意力分数，完全省去了token之间的点积交互。

2.2、实现细节

Synthesizer 大体上可以理解为一个 Transformer，其中自注意力模块被 Synthetic Attention 模块替换。

上图表示了 Transformer、Dense Synthesizer 和 Random Synthesizer 的关键思想。

Synthesizer 移除了 query-key-value 的概念，而直接合成一个对齐矩阵。具体而言，即去除了 $K$K、$Q$Q、$V$V，而使用一个行和列大小等同于序列长度$l$l 的矩阵 $B$B来表示任意 token 之间的关系。作者提出了两类 synthesizer，分别是 Dense Synthesizer 和 Random Synthesizer。

2.2.1、Dense Synthesizer

第一种形式在原论文中称为Dense：$B$B需要是$n×n$n×n大小的，而$X$X是$n×d$n×d的，所以只需要一个$d×n$d×n的变换矩阵$W_a$Wa​就可以将它变成$n×n$n×n了，即$B=XW_a$B=XWa​这其实就相当于把$K$K固定为常数矩阵$W_a^T$WaT​了。当然，原论文还做得更复杂一些，用到了两层Dense层:$B=relu(XW_1+b_1)W_2+b_2$B=relu(XW1​+b1​)W2​+b2​这可以理解为两个 Dense 层，$W$W和$b$b 用于 Dense 层的计算。而最后模型的输出$Y$Y 由表示 token 间关系的矩阵 $B$B得到。$Y=Softmax(B)G(X)$Y=Softmax(B)G(X)
其中，$G(X)$G(X)可类比为标准 Transformer 的 $V$V.

但思想上并没有什么变化。

2.2.2、Random Synthesizer

刚才说Dense形式相当于把$K$K固定为常数矩阵，我们还能不能更“异想天开”一些：把$Q$Q固定为常数矩阵？这时候整个$B$B相当于是一个常数矩阵，即$B=R$B=R原论文中还真是实验了这种形式，称之为Random，顾名思义，就是B是随机初始化的，然后可以选择随训练更新或不更新。据原论文描述，固定形式的Attention首次出现在论文《Fixed Encoder Self-Attention Patterns in Transformer-Based Machine Translation》，不同点是那里的Attention矩阵是由一个函数算出来的，而Google这篇论文则是完全随机初始化的。从形式上看，Random实际上就相当于可分离卷积（Depthwise Separable Convolution）运算。

Random Synthesizer 的方法中，注意力权重的初始化是不受任何输入 token 的影响，而是完全随机初始化。这些随机初始化的值可以被训练，或者保持固定。

以$R$R 表示一个随机初始化矩阵，则 Random Synthesizer 被定义为：$Y=Softmax(R)G(X)$Y=Softmax(R)G(X)即 $B$B初始化的值是 $R$R. 该方法不依赖 token 对之间的交互或者任何单个 token 的信息，而是学习一个能跨实例有效的任务特定的对齐。

换句话说，作者认为，学习一个跨实例有效的模型意味着在初始化时不直接依靠任何 token 信息。

2.3、低秩分解

上面两种新形式，往往会面对着参数过多的问题，所以很自然地就想到通过低秩分解来降低参数量。对于Dense和Random，原论文也提出并验证了对应的低秩分解形式，分别称为Factorized Dense和Factorized Random。

Factorized Dense通过Dense的方式，生成两个$n×a,n×b$n×a,n×b的矩阵$B_1,B_2$B1​,B2​，其中$ab=n$ab=n；然后将$B_1$B1​重复$b$b次、然后将$B_2$B2​重复$a$a次，得到对应的$n×n$n×n矩阵$\tilde{B}_1,\tilde{B}_2$B~1​,B~2​，最后将它们逐位相乘（个人感觉相乘之前$\tilde{B}_2$B~2​应该要转置一下比较合理，但原论文并没有提及），合成一个$n×n$n×n的矩阵：$B=\tilde{B}_1\otimes\tilde{B}_2$B=B~1​⊗B~2​至于Factorized Random就很好理解了，本来是一整个$n×n$n×n的矩阵$R$R，现在变成两个$n×k$n×k的矩阵$R_1,R_2$R1​,R2​，然后$B=R_1R_2^T$B=R1​R2T​

2.4、混合模式

到目前为止，连同标准的自注意力，我们有5种不同的生成矩阵$B$B的方案，它们也可以混合起来，即$B=\sum_{i=1}^N\alpha_iB_i$B=i=1∑N​αi​Bi​其中$B_i$Bi​是不同形式的自注意力矩阵，而$\sum_{i=1}^N\alpha_i=1$∑i=1N​αi​=1是可学习参数。

3、实验

前面介绍了统称为Synthesizer的几种新型自注意力形式，它们的共同特点是没有保持“token对token”形式，尤其是Random，则完全抛弃了原有注意力的动态特点，变成了静态的矩阵。

那么，这些新型自注意力的效果如何呢？它们又怎样冲击我们对自注意力机制的认识呢？

3.1、机器翻译与语言建模

第一个评测任务是机器翻译与语言建模，详细地比较了各种自注意力形式的效果：

关于机器翻译任务，可以看到相同参数量的 Synthesizer(Random + Vanilla) 与其他模型拉开了一定差距，也比相同参数量的 Transformer (Control) 表现更好。值得注意的是，两个分解方法取得的提升并不如混合模型取得的提升更多，但在一定程度上减少了参数量。

关于语言建模任务，使用的数据集是 LM1B，取得最好效果的是 Synthesizer (Dense + Vanilla)，它仍然是一个混合模型，同样是 Synthesizer 的各种设置中唯一超过 Transformer 的模型。

不知道读者怎么想，反正Synthesizer的这些结果是冲击了笔者对自注意力的认知的。表格显示，除了固定的Random外，所有的自注意力形式表现基本上都差不多，而且就算是固定的Random也有看得过去的效果，这表明我们以往对自注意力的认知和解释都太过片面了，并没有揭示自注意力生效的真正原因。

3.2、摘要生成

接下来在摘要和对话生成任务上的结果：

在自动摘要这个任务上，标准注意力效果比较好，但是对话生成这个任务上，结果则反过来：标准的自注意力是最差的，Dense（D）和Random（R）是最好的，而当Dense和Random混合了标准的自注意力后（即 D+V 和 R+V），效果也变差了。这说明标准注意力并没有什么“独占鳌头”的优势，而几个Synthesizer看起来是标准注意力的“退化”，但事实上它们互不从属，各有优势。

3.3、预训练+微调

最后，对于我们这些普通读者来说，可能比较关心是“预训练+微调”的效果怎样，也就是说，将BERT之类的模型的自注意力替换之后表现如何？原论文确实也做了这个实验，不过Baseline不是BERT而是T5(https://github.com/google-research/text-to-text-transfer-transformer)，结果如下：

在这个结果中，相比标准自注意力，Dense和Random就显得逊色了，这表明Dense和Random也许会在单一任务上表现得比较好，而迁移能力则比较弱。但是不能否定的是，像Random这样的自注意力，由于直接省去了$QK^T$QKT这个矩阵运算，因此计算效率会有明显提升，因此如果能想法子解决这个迁移性问题，说不准Transformer模型家族将会迎来大换血。

4、结论

该文提出了 Synthesizer，一个新的 Transformer 模型，它采用了合成注意力（Synthetic Attention）。作者试图更好地理解和评估全局对齐和局部、实例对齐（独立 token 和 token 到 token 的对齐）在自注意力中的效用。

在机器翻译、语言建模和对话生成等多个任务上，合成注意力与原有自注意力相比，表现出有竞争力的性能，点积自注意力的作用与效率值得怀疑与进一步研究。此外，在对话生成任务上，token 之间的交互信息实际会损害性能。

实际上，Synthesizer 的不同设置没有绝对的优劣，也和具体的任务相关。个人认为为何在各种任务中 Synthesizer 的表现存在明显差异、它与点积注意力分别适合于哪些任务，以及背后的成因是值得深究的。

论文中提出了几种新型的自注意力机制，并做了相当充分的实验，而实验结果很可能会冲击我们对自注意力机制的已有认知，值得大家读读～