论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》

论文链接:https://arxiv.org/pdf/2105.02358.pdf

Motivation

self-attention在各种视觉任务中大放异彩,相比于卷积这类局部感知的操作,self-attention可以获取更多的long-range dependency,从而学习到融合了全局特征的feature。但是self-attention自身存在两个缺点:(1)计算量太大,计算复杂度与pixel的平方相关;(2)没有考虑不同样本之间的潜在关联,只是单独处理每一个样本,在单个样本内去捕获这类long-range dependency。

针对这两个问题,作者提出了一个external attention模块,仅仅通过两个可学习的external unit,就可以简化self-attention的时间复杂度,简化到与pixel数量线性相关;同时由于两个unit是external的,对于整个数据集来说都是shared,所以还可以隐式地考虑到不同样本之间的关联。两个unit在实现的时候是两个linear layer,因此可以直接end2end优化。

Related Work

这部分写了两三年来最经典的一些视觉attention,self-attention,transformer相关的研究进展,墙裂推荐看一下。

Method

论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》
上图中(a)和(b)分别表示经典的self-attention和简化版self-attention,计算复杂度均为 O ( d N 2 ) O(dN^{2}) O(dN2),N表示特征图pixel数,d表示feature维度。经作者观察发现,这两种方法都是建立一个N-to-N的attention关联,即attention map的生成。而attention map自身是稀疏的,也就是说只有很少量的pixel之间有关联,所以N-to-N的attention map可能会显得redundant。

论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》

所以只需要用到需要的value即可,大可不必计算整张N-to-N的attention map。图(c)表示作者提出的external-attention module,他负责计算输入feature和external memory unit之间的attention:
论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》
M作为external unit,其维度是 S ∗ d S*d S∗d,S是超参数。M这里作为一个独立于input的可学习的参数,算是对整个数据集来说的一个memory。A是根据先验推断出的attention map,在feature map和M线性变换后要经过一个不太一样的Norm(下文会讲)。

在实际中,作者利用了两个这样的memory unit M k M_{k} Mk​和 M v M_{v} Mv​ ,作为key和value,来增加网络的capacity。因此整个external attention的计算式如下:

论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》
所以external attention的计算复杂度降为 O ( d S N ) O(dSN) O(dSN),由于S是超参数,所以可以设置的远小于pixel数目N。经实验发现,S为64的时候效果就非常好了。所以此时external attention的计算复杂度与N线性相关,相比于经典的self-attention来说更加efficient。

external attention的伪代码如下:
论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》
关于上面的Norm,attention map对于输入特征的scale比较敏感,因此作者没有采用简单的softmax,而是利用了一种double-normalization,公式如下:
论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》

Experiment

作者在图像分类、语义分割、图像生成、点云分类、点云分割等任务上做了一系列实验。
论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》
论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》
论文阅读《Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks》

Conclusion

External memory的设想感觉很棒,既摒弃了传统self-attention那种庞大冗余的计算,使得维度可控;又摆脱了单张image的束缚,作为一个游离的参数模块,有利于更好的学习到整个数据集上一些informative的特征,更好的把握不同sample之间的差别。

最近关于self-attention的大讨论太热闹了,再附一篇文章:

《MLP-Mixer: An all-MLP Architecture for Vision》

上一篇:【图形学】变换 (Transformations)


下一篇:逻辑回归(pytorch)