论文翻译-Beyond Self-attention_External Attention using Two Linear Layers for Visual Tasks

摘要

Attention机制，尤其是self-attention，在视觉任务的深度特征表征中起着越来越重要的作用。

self-attention通过使用所有位置的成对相似性计算特征的加权和来更新每个位置的特征，以捕捉单个样本内的长期相关性。然而，self-attention具有$n^2$复杂度，忽略了不同样本之间的潜在相关性。

本文提出了一种新的Attention机制，称为外部注意力机制，它基于两个外部的、小的、可学习的和共享的记忆单元，只需使用两个级联线性层和两个归一化层就可以很容易地实现，它可以方便地取代现有流行体系结构中的self-attention。

外部注意力机制具有线性复杂度，隐含地考虑所有样本之间的相关性。关于图像分类，语义分割，图像生成，点云分类和点云分割任务的广泛实验表明，我们的方法为self-attention和其一些变体提供了可比或卓越的性能，具有较低的计算和存储成本。

原文地址：arxiv cs.CL/2105.03432

1. 介绍

鉴于其捕获远程依赖性的能力，自注意力机制有助于提高各种自然语言处理[1,2]和计算机视觉[3,4]任务的性能。

自注意力侧重于通过聚集单个样本中所有其他位置的特征来细化每个位置的表示，这导致样本中关于位置数量的$n^2$复杂度。因此，一些变体试图近似自注意力机制降低计算成本[5,6,7,8]。

此外，自注意集中在单个样本中不同位置之间的自我相似性，而忽略了与其他样本的潜在相关性。很容易看出，不同样本之间的潜在相关性有助于更好地表示特征。

例如，属于同一类别但分布在不同样本中的特征应该在语义分割任务中得到一致的处理，类似的观察适用于图像分类和各种其他视觉任务。

本文提出了一种新的轻量级注意力机制，我们称之为外部注意力(见图1c)。

如图1A所示，为了计算自注意力，我们首先通过计算self query向量和self key向量之间的相似度来计算Attention 映射，然后通过将self value向量与Attention 映射进行加权来生成新的特征映射。

外部注意力不同。我们首先计算self-query向量与外部可学习记忆单元key之间的相似性来生成Attention映射，然后将此Attention映射乘以另一个外部可学习记忆单元value来获得精细特征映射。

实际上，两个记忆单元用线性层实现，因此可以以端到端的方式通过反向传播进行优化。它们独立于各个样本并在整个数据集中共享，该数据集发挥着强大的正则化作用，并提高了注意力机制的泛化能力。

外部注意力机制的轻量级特性的关键在于，记忆单元中的元素数量比输入特征中的数量小得多，从而产生线性计算复杂度，相对于输入中的元素数量。

外部记忆单元的设计是为了学习整个数据集中最具区分性的特征，捕捉最具信息量的部分，以及排除来自其他样本的干扰信息。

类似的想法可以在稀疏编码[9]或字典学习[10]中找到。然而，与这些方法不同的是，我们既不尝试重构输入特征，也不对Attention映射应用任何显式稀疏正则化。

所提出的外部注意力机制虽然简单，但对不同的视觉任务是有效的。由于其简单性，它可以很容易地集成到现有流行的基于自注意力的体系结构中，如DANET[4]、Sagan[11]和T2T-Transformer[12]。

图2展示了一个典型的编解码器式的架构，在图像语义分割任务中用我们的外部注意力代替self-attention。

我们在不同的输入方式(图像和点云)下，对分类、语义分割和生成等基本视觉任务进行了广泛的实验。实验结果表明，在较低的计算代价下，我们的方法获得了与原始的自我注意机制及其一些变体相当或更好的结果。

2. 相关工作

由于对注意力机制的全面回顾超出了本文的范围，因此我们只讨论视觉领域中最密切相关的文献。

2.1 视觉任务中的注意力机制

注意力机制可以被视为根据激活的重要性重新分配资源的机制。它在人类视觉系统中扮演着重要的角色。在过去的十年里，这一领域得到了蓬勃的发展[13，14，15，16，3，17，18]。Hu等人提出了SENet[15]，表明注意力机制可以降低噪声，提高分类性能。此后，许多其他论文将其应用于视觉任务。Wang等人介绍了用于视频理解的no-local网络[3]。Hu等人[19]在物体检测中使用注意力。Fu等人提出了DANet[4]的语义切分方法。张等人[11]证明了注意力机制在图像生成中的有效性。谢等人提出了用于点云处理的A-SCN[20]。

2.2 视觉任务中的自注意力机制

自注意是注意力的一种特殊情况，许多文献[3，17，4，11]都研究了视觉的自注意机制。自注意的核心思想是计算特征之间的相似度，以获取长期依赖关系。然而，随着特征图大小的增加，计算和内存开销呈二次曲线增加。

为了减少计算和存储成本，Huang等人[5]提出了交叉注意力，首先考虑行注意力，然后考虑列注意力，以捕捉全局上下文。Li等人[6]采用期望最大化(EM)聚类对自注意力进行优化。Yuan等人[7]提出使用对象-上下文向量来处理注意力。然而，这取决于语义标签。耿等人[8]表明在语义分割和图像生成中，矩阵分解是一种更好的全局上下文建模方法。

与通过计算self-query和self-key之间的相似度来获得Attention映射的self-attention不同，我们的外部注意力机制计算self-query和小得多的可学习记忆单元key之间的关系，以捕获数据集的全局上下文。

外部注意力不依赖于语义信息，可以通过端到端的反向传播算法来优化，而不是使用迭代的聚类算法。

2.3 视觉任务中你的transformer

基于变压器的模型在自然语言处理方面取得了巨大的成功[1，21，16，2，22，23，24]。最近，它们在视觉任务方面也显示出巨大的潜力。Carion等人[25]提出了一种端到端检测变压器，它以CNN特征为输入，通过变压器生成检测框。Chen等人[26]提出了一种基于变压器的IGPT图像生成方法。Dosovitski[18]提出了基于patch编码和transformer的VIT，表明在有足够的训练数据的情况下，transformer比传统的CNN提供更好的性能。

随后，transformer方法已被成功地应用于许多视觉任务，包括图像分类[27，12]、对象检测[28]、低级视觉[29]、语义分割[30]、跟踪[31]、视频实例分割[32]、图像生成[33]、多模态学习[34]、对象重新识别[35]、图像字幕[36]和点云学习[37]。

3. 方法

在这一部分中，我们首先分析了原有的自注意机制。然后我们将详细介绍我们的外部注意力，这是一种定义注意力的新方式。只需使用两个线性层和两个归一化层即可轻松实现，如清单1所示。

3.1 外部注意力机制

我们首先回顾一下自注意机制(见图1a)。给定输入特征图$F \in \mathbb { R } ^ { N \times d }$，其中$N$是像素的数量，$d$是特征维度，自注意力机制将输入线性投影到query矩阵$Q \in \mathbb { R } ^ { N \times d ^ { \prime } }$，key矩阵$K \in \mathbb { R } ^ { N \times d ^ { \prime } }$，value矩阵$V \in \mathbb { R } ^ { N \times d }$[16]。那么自注意力可以表述为：
$$
\begin{aligned}
A & = ( \alpha ) _ { i , j } = \operatorname { softmax } \left( Q K ^ { T } \right) &(1)\
F _ { \text {out } } & = A V &(2)
\end{aligned}
$$

其中，$A \in \mathbb { R } ^ { N \times N }$是注意力矩阵，$\alpha_{i,j}$是第$i,j$个像素的成对亲和力(相似性)。

自注意力的一种流行的简化变体(图1b)简化了QKV矩阵，并通过使用输入特征$F$直接计算Attention映射，如下所示：
$$
\begin{aligned}
A & = \operatorname { softmax } \left( F F ^ { T } \right) &(3) \
F _ { \text {out } } & = A F &(4)
\end{aligned}
$$
这里，Attention映射是通过在特征空间中计算像素级相似度来获得的，输出是输入的精化特征表示。

然而，即使简化了，计算复杂度$\mathcal { O } \left( d N ^ { 2 } \right)$也很高,这是使用自注意的一个重大缺陷。输入像素数的二次方复杂性使得对图像直接应用自注意是不可行的。

因此，以前的工作[18]利用patch上的自注意而不是像素上的自我注意来降低计算复杂度。在可视化Attention映射时，我们发现大多数像素与其他几个像素密切相关，N到N的注意力矩阵可能是多余的。

因此，可以通过使用所需的值来实现精细化特征。因此，我们提出了一种替代方案，即外部注意力模块，它计算输入像素和外部存储单元$M \in \mathbb { R } ^ { S \times d }$之间的关注度：
$$
\begin{aligned} A & = ( \alpha ) _ { i , j } = \operatorname { Norm } \left( F M ^ { T } \right) &(5)
\ F _ { \text {out } } & = A M &(6)
\end{aligned}
$$
与自注意力机制不同的是，在公式5中的$\alpha_{i,j}$是$M$的第$i$个像素和第$j$行之间的相似性。这里，$M$是一个独立于输入的可学习参数，作为整个训练数据集的记忆。

$A$是由先验知识推断出的注意力映射;它以类似于自注意力的方式被规范化(见第3.2节)。最后，我们通过$A$的相似性来更新$M$的输入特征。

在实际应用中，我们使用两个不同的存储单元$M_k,M_v$作为key和value，以提高网络的能力。计算方式现在变为：
$$
\begin{aligned}
A & = \operatorname { Norm } \left( F M _ { k } ^ { T } \right) &(7)\
F _ { \text {out } } & = A M _ { v } &(8)
\end{aligned}
$$
外部注意力机制的计算复杂度是$\mathcal { O } ( d S N )$。其中$d,S$是超参数，该算法在像素数上是线性的。

事实上，我们发现小S，例如64，在实验中效果很好。因此，外部注意力机制比自注意力机制更有效，可以直接应用于大规模的输入。

3.2 正则化

利用Softmax对注意映射进行归一化，$\sum _ { j } \alpha _ { i , j } = 1$。然而，注意力映射是通过矩阵乘法计算出来的。与余弦相似度不同，注意力映射对输入特征的比例非常敏感。为了避免这个问题，我们选择[37]中提出的双重规范化，它分别规范化列和行。双规格化表示为:
$$
\begin{aligned} ( \tilde { \alpha } ) _ { i , j } & = F M _ { k } ^ { T } &(9) \
\alpha _ { i , j } & = \frac { \exp \left( \tilde { \alpha } _ { i , j } \right) } { \sum _ { k } \exp \left( \tilde { \alpha } _ { k , j } \right) } &(10)\
\alpha _ { i , j } & = \frac { \alpha _ { i , j } } { \sum _ { k } \alpha _ { i , k } } &(11)
\end{aligned}
$$

4. 实验

我们对图像分类、语义分割、图像生成、点云分类和点云分割任务进行了实验，以评估我们提出的外部注意方法的有效性。所有实验使用Jittor[40]和Pytorch[41]深度学习框架实现。

4.1 图像分类

ImageNet-1K[42]是一个广泛用于图像分类的数据集。我们将T2T-ViT[12]中的T2T-Transformer块替换为外部注意力。为了进行公平比较，其他超参数设置与T2T-ViT相同。表1的实验结果表明，外部注意与多头注意和performer[43]的结果相当。

4.2 语义分割

在这个实验中，我们采用了图2中的语义分割架构，使用了Pascal VOC数据集[56]、ADE20K数据集[57]和城市场景数据集[58]。

Pascal VOC包含10,582张用于训练的图像，1,449张用于验证的图像，1,456张用于测试的图像。它有20个前景对象类和一个用于分割的背景类。我们使用扩张的ResNet-101，输出步幅为8;后者在ImageNet-1K上进行了预先训练。在训练阶段采用多元学习速率策略。

初始学习率、批数和输入数分别设置为0.009、16和513 × 513。我们首先在训练集上训练了45k次迭代，然后在trainval集上微调为15k次迭代。最后对测试集进行了多尺度测试和翻转测试;结果如表2所示。

ADE20K是一个更具挑战性的数据集，有150个类，有20K, 2K和3K的图像用于训练，验证和测试。我们采用扩张的ResNet-101作为骨干，输出步幅为8。实验配置与mmsegmentation相同，训练ADE20K进行160k次迭代。表3中的结果表明，我们的方法在ADE20K val集上提供了更好的结果。

城市景观包含了5000个高质量像素级精细注释的标签，在19个语义类用于城市场景理解。每张图像都是1024 × 2048像素。分为2975张、500张和1525张，用于训练、验证和测试。(它还包含2万张粗略标注的图像，我们在实验中没有使用它们)。我们采用扩张的ResNet-101作为骨干，输出步幅为8。实验配置与mmsegmentation相同，用80k次迭代训练ADE20K。表4的结果显示，我们的方法在城市景观val集上取得了可比性的结果。

4.3 图像生成

自注意力是图像生成中常用的方法，典型的方法是SAGAN[11]。我们将SAGAN中的自注意力机制替换为我们在生成器和鉴别器中的外部注意力方法，从而得到我们的EAGAN模型。所有实验都基于流行的PyTorch-StudioGAN repo[78]。超参数使用缺省配置。我们选择Frechet Inception Distance (FID)[79]和Inception Score(IS)[80]作为我们的评估指标。生成的部分图像如图4所示，表5的定量结果显示，外部注意力机制比SAGAN和其他一些gan的性能更好。

4.4 点云分类

ModelNet40[81]是一种流行的3D形状分类基准，包含40个类别的12,311个CAD模型。有9843个训练样本和2468个测试样本。EA T取代PCT[37]中的所有自注意模块。我们在每个形状上取样1024个点，并使用随机平移、各向异性缩放和dropout (PCT[37])增加输入。表6表明，我们的方法优于所有其他方法，包括其他一些基于注意力的方法，如PCT。结果表明，所提出的方法为2D和3D视觉提供了一个突出的主干。

4.5 点云分割

我们在ShapeNet部分数据集上进行了点云分割实验[82]。训练集中有14006个3D模型，评估集中有2874个。每个形状被分割成零件，共16个对象类别和50个零件标签。我们遵循PCT[37]的实验设置。EA T在这个数据集上取得了最好的结果，如表7所示。

4.6 计算需求

与投入规模相关的线性复杂性带来了效率上的显著优势。我们比较了在1×512×128×128输入大小下，外部注意与标准自注意(SA)[16]及其几个变体在参数数量和推理操作方面的差异，并在表8中报告了结果。

外部注意力只需要自注意力所需参数的三分之一，其速度是自我关注的50倍。与最好的变体相比，外部注意力仍然是3倍快，而且更轻更高效。

5. 总结

本文提出了一种新型的、轻量级的、有效的、适用于各种视觉任务的注意机制——外部注意。外部注意所采用的两个外部存储单元可以看作是整个数据集的字典，能够在权衡计算代价的同时学习输入的更有代表性的特征。我们希望外界的关注能够激发它在其他领域的实际应用和研究，比如自然语言处理