WYXOGO
Posts Tags Categories Comment Friends
WYXOGO

Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

Yasin avatarYasin  included in  categories Reading Backbone
     510 words  3 minutes    views  
标题 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） 的解读论文，也分享一下在做这篇论文时候的实验过程中一些问题和想法。"

原文：External attention 和 EAMLP

论文：https://arxiv.org/abs/2105.02358

代码：https://github.com/MenghaoGuo/-EANet

本次更新主要包含了三个方面：

  1. 加入了 multi-head external attention 机制，multi-head external attention 也可以使用两个线性层实现，由于有了 multi-head external attention 结构，我们实现了一个 MLP 结构，我们把它叫做 EAMLP。
  2. 补充了一个 ablation study 的实验以及一些分析，可以更清楚的理解这种 external attention 机制。
  3. 补充了 COCO 上的 object detection 和 instance segmentation 的实验和 Tiny ImageNet 上的生成实验。现在已经在图像分类、检测、分割、实例分割、图像生成、点云的分割和分割上证明了 external attention 的有效性。

自注意力机制，通过计算两两之间的相似度，然后更新根据相似度更新特征，达到特征增强的作用。具体的，对于输入特征 $F \in N \times d$，QKV 是 F 的一种线性变换，自注意力可以写成如下的形式。

$$ A = (\alpha)_{i,j} = softmax(QK^T) \\ F_{out} = AV $$

简要记为：

$$ A = softmax(FF^T) \\ F_{out} = AF $$

对于简记的注意力机制，我们可以认为这是 F 对 F 的注意力，也就是 self-attention。这种注意力是非常有效的，但是也有它的不足。首先，它使用的是一个 F 对 F 的注意力形式，这种注意力只会考虑单个样本内部的关联，而会忽略样本之间的潜在联系，这种联系对于视觉任务来说是有用的，比如对于语义分割这个任务来说，我们希望分布在不同样本中的同类物体能有着相似的表征。其次，由于是算 F 对 F 的注意力，它不够灵活，这种不灵活带来的是$N^2$级别的复杂度，难以用到高分辨率图像场景。

下面考虑这样一个问题，如果有一个共享的矩阵$M \in S \times d$ ,M是随机初始化的，我们可以把上面那个形式，记成一种 M 对 F的注意力，即如下公式所示。

$$ A = (\alpha)_{i,j} = softmax(FM^T) \\ F_{out} = AM $$

对于上述形式，首先，M 是共享的，所以可以隐式的考虑不同样本之间的关联，其次，由于 S 是非常灵活的，我们可以通过控制 S 的大小，来使得整个 external attention 变得灵活，使得 attention 变成 N 复杂度，以便用于高分辨率的情况。

对于 self-attention, 简化的 self-attention 以及这个 M 对 F的 attention (external attention)，可以用下图来表示。

Fig. 1.
Fig. 1. Self-attention versus external-attention

A Tip

进一步考虑公式，可以发现，公式中的$FM^T$是什么呢 ？

是矩阵乘法，也就是我们常用的线性层 (Linear Layer)。这就是解释了为什么说线性可以写成是一种注意力机制。写成代码就是下面这几行, 就是线性层。

Algorithm 1
Algorithm 1 Pseudo-code for 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。如下所示：

$$ \begin{aligned} h_{i} &=\text { ExternalAttention }\left(F_{i}, M_{k}, M_{v}\right) \\ F_{\text {out }} &=\text { MultiHead }\left(F, M_{k}, M_{v}\right) \\ &=\operatorname{Concat}\left(h_{1}, \ldots, h_{H}\right) W_{o} \end{aligned} $$

和 self-attention 的 multi-head 不同，我们的 multi-head 的每个 head 是共享 $M_k$和$M_v$的，multi-head self-attention 每个 head 是独立的。这个multi-head external attention 也可以写成线性层，具体写法如下：

Algorithm 2
Algorithm 2 Pseudo-code for multi-head external attention.

在通常的 attention 中，我们常常只使用一个 softmax 作为中间的归一化层，这种归一化的目的是使得 attention map 中的某一行或者某一列和为1。这种只用 softmax 有什么问题呢 ？问题在于，当某一个特征值特别大（特别小）的时候，他对其他特征的点乘也会变得特别大（特别小），这种情况下，只使用 softmax 会破坏 attention 原始的含义。我们使用了 softmax + L1 norm 的这种 double normalization 的方式，去避免这个问题，公式如下：

$$ \begin{aligned} (\tilde{\alpha})_{i, j} &=F M_{k}^{T} \\ \hat{\alpha}_{i, j} &=\exp \left(\tilde{\alpha}_{i, j}\right) / \sum_{k} \exp \left(\tilde{\alpha}_{k, j}\right) \\ \alpha_{i, j} &=\hat{\alpha}_{i, j} / \sum_{k} \hat{\alpha}_{i, k} \end{aligned} $$

值得注意的是，external attention 非常依赖这种 double norm 的方式，这种 norm 也对self-attention 有提升作用，实验结果如下：

TABLE 1

这种 double norm 的方式在 ImageNet 上面也会让 T2T-VIT 有 0.1 - 0.3 个百分点的提升。

我们在 ImageNet 上做了一个纯 MLP 的结构，我们叫它 EAMLP，实验结果如下：

TABLE 2

同时我们可视化了最后一层 external attention 的 attention map，如下所示：

Fig. 5.
Fig. 5. Multi-head attention map in the last layer of EAMLP-14 on ImageNet val set. Left: Input image Others: 24 head attention map in the last layer of EAMLP-14 for the ImageNet val set. Last two rows: attention of two different rows of $M_k$ to the image patches.

关于这个实验，说几点感想

  1. 我们采用了 MoCo V3 中的将 LayerNorm 换成 BatchNorm 的尝试，也在 EAMLP-7 上取得了 1.1 个点的提升，但是在大模型上梯度会崩掉。这个实验让我感觉，在视觉里面，Transformer 中的 LayerNorm 在视觉中可能不如 BatchNorm 好，这可能是 Norm 本身造成的，也有可能是优化器不适合造成。
  2. multi-head 这种结构是重要的，之前也看过一些轻量级结构，比如 EMANet， HamNet、OCRNet、CCNet、ANN…这些结构如果设计一个合理的multi-head 结构应该也可能用到 transformer 里面代替 self-attention。
  3. 另外，值得注意的是，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?


  • Author: Yasin
  • Link: https://wyxogo.top/eamlp/
  • Copyright: This article is adopted , if reprint please indicate from Yasin

    • Related Content

    关于MLP的一些观点和一些未来可能的方向 【Can Attention Enable MLPs To Catch Up With CNNs】

    Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

    RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition

    【多层感知机混合器】MLP-Mixer: An all-MLP Architecture for Visio

    Updated on 2021-06-03
     EAMLP, MLP
    Back | Home