Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks
标题 | Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks |
---|---|
年份: | 2021 年 5 月 |
GB/T 7714: | Guo M H, Liu Z N, Mu T J, et al. Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks[J]. arXiv preprint arXiv:2105.02358, 2021. |
原作者:国孟昊 清华大学 工学博士在读
"本文是对最近我们更新到 arxiv 的 paper :Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks (External attention) 的解读论文,也分享一下在做这篇论文时候的实验过程中一些问题和想法。"
本次更新主要包含了三个方面:
- 加入了
multi-head external attention
机制,multi-head external attention 也可以使用两个线性层实现,由于有了 multi-head external attention 结构,我们实现了一个 MLP 结构,我们把它叫做 EAMLP。 - 补充了一个 ablation study 的实验以及一些分析,可以更清楚的理解这种 external attention 机制。
- 补充了 COCO 上的 object detection 和 instance segmentation 的实验和 Tiny ImageNet 上的生成实验。现在已经在图像分类、检测、分割、实例分割、图像生成、点云的分割和分割上证明了 external attention 的有效性。
1 Self-Attention
自注意力机制,通过计算两两之间的相似度
,然后更新根据相似度更新特征,达到特征增强的作用。具体的,对于输入特征 $F \in N \times d$,QKV 是 F 的一种线性变换,自注意力可以写成如下的形式。
简要记为:
2 External attention
对于简记的注意力机制,我们可以认为这是 F 对 F 的注意力,也就是 self-attention。这种注意力是非常有效的,但是也有它的不足。首先,它使用的是一个 F 对 F 的注意力形式,这种注意力只会考虑单个样本内部的关联,而会忽略样本之间的潜在联系,这种联系对于视觉任务来说是有用的,比如对于语义分割这个任务来说,我们希望分布在不同样本中的同类物体能有着相似的表征
。其次,由于是算 F 对 F 的注意力,它不够灵活,这种不灵活带来的是$N^2$级别的复杂度,难以用到高分辨率图像场景。
下面考虑这样一个问题,如果有一个共享的矩阵$M \in S \times d$ ,M是随机初始化的,我们可以把上面那个形式,记成一种 M 对 F的注意力,即如下公式所示。
对于上述形式,首先,M 是共享的,所以可以隐式的考虑不同样本之间的关联,其次,由于 S 是非常灵活的,我们可以通过控制 S 的大小,来使得整个 external attention 变得灵活,使得 attention 变成 N 复杂度,以便用于高分辨率的情况。
对于 self-attention, 简化的 self-attention 以及这个 M 对 F的 attention (external attention),可以用下图来表示。
3 External attention 和 线性层
进一步考虑公式,可以发现,公式中的$FM^T$是什么呢 ?
是矩阵乘法,也就是我们常用的线性层 (Linear Layer)。这就是解释了为什么说线性可以写成是一种注意力机制。写成代码就是下面这几行, 就是线性层。
4 Multi-head external attention
到这里,我们已经把 attention 写成了 Attention = Linear(Norm(Linear(F))) 的形式,不妨把该式子记为 Attention=G(F) 。考虑多层 attention 的堆叠,也就是现在常用的 Transformer。那么就是 Transformer = G(G’…..G(x)) 。 没错,是可以这么做,我们也做了实验,在 ImageNet 仅获得了 63.2 的准确率(对标同级别的 transformer 71.7),我们反思了一下,到底是哪出了问题,发现 multi-head 机制在这里面起到了关键的作用。我们吧 multi-head self-attention 改成 single head 之后发现了其性能也只能达到 67.4,直接降低了 4.3 个点。有了这个 insight 之后,我们设计了我们的 multi-head external attention。如下所示:
和 self-attention 的 multi-head 不同,我们的 multi-head 的每个 head 是共享 $M_k$和$M_v$的,multi-head self-attention 每个 head 是独立的。这个multi-head external attention 也可以写成线性层,具体写法如下:
5 Normalization
在通常的 attention 中,我们常常只使用一个 softmax 作为中间的归一化层,这种归一化的目的是使得 attention map 中的某一行或者某一列和为1。这种只用 softmax 有什么问题呢 ?问题在于,当某一个特征值特别大(特别小)的时候,他对其他特征的点乘也会变得特别大(特别小),这种情况下,只使用 softmax 会破坏 attention 原始的含义。我们使用了 softmax + L1 norm 的这种 double normalization 的方式,去避免这个问题,公式如下:
值得注意的是,external attention 非常依赖这种 double norm 的方式,这种 norm 也对self-attention 有提升作用,实验结果如下:
这种 double norm 的方式在 ImageNet 上面也会让 T2T-VIT 有 0.1 - 0.3 个百分点的提升。
6 实验中的一些感想
我们在 ImageNet 上做了一个纯 MLP 的结构,我们叫它 EAMLP,实验结果如下:
同时我们可视化了最后一层 external attention 的 attention map,如下所示:
关于这个实验,说几点感想
- 我们采用了 MoCo V3 中的将 LayerNorm 换成 BatchNorm 的尝试,也在 EAMLP-7 上取得了 1.1 个点的提升,但是在大模型上梯度会崩掉。这个实验让我感觉,在视觉里面,Transformer 中的 LayerNorm 在视觉中可能不如 BatchNorm 好,这可能是 Norm 本身造成的,也有可能是优化器不适合造成。
- multi-head 这种结构是重要的,之前也看过一些轻量级结构,比如 EMANet, HamNet、OCRNet、CCNet、ANN…这些结构如果设计一个合理的multi-head 结构应该也可能用到 transformer 里面代替 self-attention。
- 另外,值得注意的是,softmax 不一定就是最合适的归一化方式。当然 double norm 也极大可能不是最合适的归一化方式。external attention 非常依赖这种 double norm 的方式,这种 norm 也对self-attention 有提升作用。attention 里面的归一化方式应该是一个非常值得探索的 topic。
另外,简单写了一些我们对最近的 MLP 的一些观点和一些未来可能的方向,一个 4 页纸的short paper : Can Attention Enable MLPs To Catch Up With CNNs?
- 文章作者:Ogo
- 文章链接:https://wyxogo.top/eamlp/
- 版权声明:本站所有文章除特别声明外,均采用 转载请注明来自