External Attention:计图团队提出两层线性层结构超越Self Attention
5月4日,谷歌团队在arXiv上提交了一篇论文“MLP-Mixer: An all-MLP Architecture for Vision”[1],引起了广大计算机视觉的研究人员的热烈讨论,MLP究竟有多大的潜力?
5月5日,清华大学图形学实验室Jittor团队在arXiv上提交论文“Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks”[2], 提出了一种新的注意机制,称之为“External Attention”,基于两个外部的、小的、可学习的和共享的存储器,只用两个级联的线性层和归一化层就可以取代了现有流行的学习架构中的“Self-attention”,揭示了线性层和注意力机制之间的关系。
同日,清华大学软件学院丁贵广团队在arXiv上提交了论文“RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition”[3],展示了结合重参数化技术的MLP也取得了非常不错的效果。
5月6日牛津大学的学者提交了一篇名为"Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet"的论文[4],也提出了Transformer中的attention是不必要的,仅仅使用Feed forward就可以在ImageNet上实现非常高的结果。
Part1
由于QKV是F的线性变换,简单起见,我们可以将自注意力计算公式简记如下:
这是 F 对 F 的注意力,也就是所谓的 Self-attention。如果希望注意力机制可以考虑到来自其他样本的影响,那么就需要一个所有样本共享的特征。为此,我们引入一个外部的Sxd维空间的记忆单元M,来刻画所有样本最本质的特征,并用M来表示输入特征。
我们称这种新的注意力机制为External-attention。我们可以发现,公式(5)(6)中的计算主要是矩阵乘法,就是常见的线性变换,一个自注意力机制就这样被两层线性层和归一化层代替了。我们还使用了之前工作[5]中提出的Norm方式来避免某一个特征向量的过大而引起的注意力失效问题。
下图形象地展示了External-attention与Self-attention的区别。
图1 Self Attention和External Attention的区别
Part2
其次,自注意力机制仅利用了自身样本内的信息,忽略了不同样本之间的潜在的联系,而这种联系在计算机视觉中是有意义的。打个比方,对于语义分割任务,不同样本中的相同类别的物体应该具有相似的特征。
External-attention通过引入两个外部记忆单元,隐式地学习了整个数据集的特征。这种思想同样在稀疏编码和字典学习中得到了应用。
计图团队在Pascal VOC 的Test set上,可视化了注意力图以及分割的结果,如图2所示,可以发现,使用两层线性层的External attention 的注意力图是合理的。
Part3
为了证明方法的通用性,我们在图像分类、分割、生成以及点云的分类和分割上均做了实验,证明了方法的有效性,External-attention在大大减少计算量的同时,可以取得与目前最先进方法相当,甚至更好的结果。
Part4
清华的External Attention则揭示了线性层和注意力机制之间的内在关联,证明了线性变换其实是一种特殊形式的注意力实现,如下公式所示:
Attention(x)=Linear(Norm(Linear(x)))
计图团队的工作和谷歌团队的工作都证明了线性层的有效性。值得注意的是,如果将External-attention不断级联堆叠起来,也是MLP的形式,就可以实现一个纯MLP的网络结构,但External-attention使用不同的归一化层,其更符合注意力机制。这与谷歌团队的工作有异曲同工之妙。
清华的External Attention的部分计图代码已经在Github开源。
https://github.com/MenghaoGuo/-EANet
后续将尽快开源全部计图代码。
参考文献
Tolstikhin et al., MLP-Mixer:An all-MLP Architecture forVision, https://arxiv.org/abs/2105.01601 Guo et al. Beyond Self-attention: External Attentionusing Two Linear Layers for Visual Tasks, https://arxiv.org/abs/2105.02358 Ding et al., RepMLP: Re-parameterizing Convolutions intoFully-connected Layers for Image Recognition, https://arxiv.org/abs/2105.01883 Luke Melas-Kyriazi, Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet, https://arxiv.org/abs/2105.02723 Guo et al. PCT: Point cloud transformer. Computational Visual Media, 7(2), 187–199 (2021).
GGC往期回顾