论文链接：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

上图中(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。

所以只需要用到需要的value即可，大可不必计算整张N-to-N的attention map。图（c）表示作者提出的external-attention module，他负责计算输入feature和external memory unit之间的attention：

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的计算式如下：

所以external attention的计算复杂度降为 O ( d S N ) O(dSN) O(dSN)，由于S是超参数，所以可以设置的远小于pixel数目N。经实验发现，S为64的时候效果就非常好了。所以此时external attention的计算复杂度与N线性相关，相比于经典的self-attention来说更加efficient。

external attention的伪代码如下：

关于上面的Norm，attention map对于输入特征的scale比较敏感，因此作者没有采用简单的softmax，而是利用了一种double-normalization，公式如下：

Experiment

作者在图像分类、语义分割、图像生成、点云分类、点云分割等任务上做了一系列实验。

Conclusion

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

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

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