Global Self-Attention Networks for Image Recognition
https://arxiv.org/abs/2010.03019v2
6 Oct 2020
作者：Zhuoran Shen, Irwan Bello, Raviteja Vemulapalli, Xuhui Jia, Ching-Hui Chen
最近，计算机视觉领域的一系列工作在利用自注意的各种图像和视频理解任务中显示出了有希望的结果。然而，由于自注意的二次计算和内存复杂性，这些工作要么将注意力仅应用于深度网络后期的低分辨率特征映射，要么将每层的注意感受域限制在一个小的局部区域。为了克服这些限制，这项工作引入了一个新的全局自注意模块，称为GSA模块，它的效率足以作为深度网络的主干组件。该模块由两个平行层组成：一个内容注意层（仅基于内容注意像素）和一个位置注意层（基于像素的空间位置注意像素）。该模块的输出是两层输出的总和。基于提出的GSA模块，我们引入了新的独立的基于全局注意的深度网络，该网络使用GSA模块而不是卷积来模拟像素交互。由于提出的GSA模块的全局范围，GSA网络能够模拟整个网络中的远程像素交互。我们的实验结果表明，在CIFAR-100和ImageNet数据集上，GSA网络的性能显著优于相应的基于卷积的网络，同时使用较少的参数和计算。所提出的GSA网络也优于基于ImageNet数据集的各种现有的基于注意的网络。
1 导言
自注意是神经网络中的一种机制，它专注于建模长期依赖关系。与其他机制（如卷积和递归）相比，它在建立全局依赖性方面的优势使其在现代深度学习中流行。在计算机视觉领域，最近的几项工作已经用全局自注意模块增强了卷积神经网络（CNN），并在各种图像和视频理解任务中显示了有希望的结果（Bello等人，2019年；Chen等人，2018年；Huang等人，2019年；Shen等人，2018年；Wang等人，2018年；Yue等人，2018年）。为了简洁起见，在本文的其余部分中，我们将自注意简单地称为注意。
在计算机视觉任务中使用全局注意机制的主要挑战是输入的大空间维度。计算机视觉任务中的输入图像通常包含数万像素，而注意力机制的二次计算和内存复杂性使得全局注意力对于如此大的输入来说代价高昂。因此，早期作品如Bello等人（2019年）；Wang等人（2018年）将全局注意机制的使用限制在深度网络后期的低分辨率特征地图上。或者，其他近期作品，如胡等人（2019年）；Ramachandran等人（2019年）；赵等人（2020年）将注意力操作的感受野限制在小的局部区域。虽然这两种策略都能有效地限制注意力模块的资源消耗，但它们剥夺了网络在早期和中期模拟远距离像素交互的能力，从而阻碍了注意力机制发挥其全部潜力。
与上述工作不同，Chen等人（2018）；黄等（2019）；沈等人（2018）；Yue et al.（2018）通过移除查询和键乘积上的softmax归一化并改变注意计算中涉及的矩阵乘法顺序（Chen et al.，2018；Shen et al.，2018；Yue et al.，2018）或将一个全局注意层分解为一系列多个轴向注意层（Huang等人，2019年）。然而，所有这些作品都只注意内容，而没有考虑像素的空间排列。由于图像是空间结构的输入，忽略空间信息的注意机制本身并不适合图像理解任务。因此，这些工作将注意力模块作为辅助模块纳入标准CNN。
为了解决上述问题，我们引入了一个新的全局自注意模块，称为GSA模块，该模块在考虑像素的内容和空间位置的情况下执行注意。该模块由两个平行层组成：内容注意层和位置注意层，其输出在末尾汇总。内容注意层仅基于所有像素的内容一次注意它们。它使用了一种有效的全局注意机制，类似于Chen等人（2018年）；Shen等人（2018年），其计算和内存复杂性在像素数量上是线性的。位置注意层基于每个像素自身的内容及其相对于其他像素的相对空间位置来计算每个像素的注意图。按照轴向公式（Ho et al.，2019；Huang et al.，2019），位置注意层实现为仅列注意层，后跟仅行注意层。这个轴向位置注意层的计算和内存复杂性为O（N）√ N） 在像素的数量上。
所提出的GSA模块的效率足以作为深度网络的主干组件。在此模块的基础上，我们引入了新的独立的基于全局注意的深度网络，称为全局自注意网络。GSA网络使用GSA模块而不是卷积来模拟像素交互。借助于GSA模块的全局范围，GSA网络能够模拟整个网络中的远程像素交互。最近，Wang et al.（2020）还引入了独立的基于全局注意的深度网络，该网络使用轴向注意机制来处理内容和位置注意。与Wang et al.（2020）不同，提出的GSA模块使用了一种非轴向的全局内容注意机制，该机制一次注意整个图像，而不仅仅是一行或一列。我们的实验结果表明，GSA ResNet是一种采用ResNet（He et al.，2016）结构的GSA网络，在广泛使用的ImageNet数据集上，其性能优于原始的基于卷积的ResNet和各种最近基于全局或局部注意的ResNet。
主要贡献
• GSA模块：
我们引入了一个新的全局注意力模块，它的效率足以充当深度网络的主干组件。与Wang等人（2018）不同；Yue等人（2018）；陈等（2018）；沈等人（2018）；Huang等人（2019），提出的模块注意基于内容和空间位置的像素。与赵等人（2020）不同；胡等（2019）；Ramachandran等人（2019年），提议的模块注意整个输入，而不是一个小的局部邻居。与Wang et al.（2020）不同，提出的GSA模块使用了一种非轴向的全局内容注意机制，该机制一次注意整个图像，而不仅仅是一行或一列。
• GSA网络：
我们引入了新的基于全局注意力的独立网络，该网络使用GSA模块而不是空间卷积来模拟像素交互。这是探索独立的基于全局注意力的图像理解任务网络的首批工作之一（Wang等人（2020年）是唯一的其他工作）。现有的基于全局注意的作品将其注意模块插入CNN，作为网络后期的辅助块，而现有的基于独立注意的网络使用局部注意模块。
• 实验：
我们表明，在CIFAR-100和ImageNet数据集上，所提出的GSA网络的性能显著优于相应的CNN，同时使用较少的参数和计算。我们还表明，GSA网络在ImageNet数据集上的表现优于各种现有的基于注意的网络，包括Wang等人（2020年）的最新独立的基于全局注意的网络。
2 相关工作
2.1 辅助视觉注意
Wang et al.（2018）提出了非局部块，这是点状产品注意机制首次适应计算机视觉中的远程依赖建模。他们通过实验验证了它在视频分类和目标检测方面的有效性。后续工作将其扩展到不同的任务，如生成性对抗性图像建模（Zhang等人，2019；Brock等人，2019）、视频人物再识别（Lio等人，2018）、图像去训练（Li等人，2018）等。最近的几项工作专注于降低Wang等人（2018）的高计算成本。陈等（2018）；Shen et al.（2018）利用矩阵乘法的关联性将复杂性从二次降低到线性。Huang等人（2019）提出将全局注意分解为行注意和列注意，以节省资源。
最近，一系列作品（Sun等人，2019年；Carion等人，2020年）将Transformer（Vaswani等人，2017年）用于各种计算机视觉应用。这些工作首先使用深度CNN来提取语义特征，然后使用Transformer来建模高层语义特征之间的交互。例如，Carion等人（2020年）使用Transformer对对象检测的对象级交互进行建模，Sun等人（2019年）使用Transformer对视频表示学习的帧间相关性进行建模。
所有这些方法都使用注意模块作为辅助模块来增强CNN的长期依赖性建模，并将大部分特征提取工作简化为卷积运算。相比之下，GSA网络使用注意作为原始操作，而不是空间卷积。
2.2 视觉注意
Bello等人（2019年）首次将注意力作为计算机视觉任务的原始操作进行测试。然而，他们使用了昂贵的非局部块（Wang等人，2018），这使得他们无法完全替换卷积层。Ramachandran et al.（2019）、Hu et al.（2019）和Zhao et al.（2020）通过将注意力的接受域限制在一个当地社区来解决这个问题。与这些工作相比，提出的GSA网络在整个网络中使用全局注意力，并且仍然有效。最近，Wang等人（2020年）使用轴向分解来提高全局注意力的效率。与之不同的是，所提出的GSA网络使用了一种非轴向的全局内容注意机制，该机制优于轴向机制，如后面的实验所示。
3 全局自注意网络
3.1 全局自注意模块
让我∈ RWH×din和fo∈ RWH×dout分别表示所提出的GSA模块的（空间）平坦输入和输出特征映射。这里，W，H表示空间尺寸，din，dout表示通道尺寸。输出特征图中的每个像素都是通过基于其内容和空间位置聚合来自输入特征图中每个像素的信息来生成的。设K=[kij]∈ RWH×dk，Q=[qij]∈ RWH×dk和V=[vij]∈ RWH×dout分别表示使用输入特征映射fi上的三个1×1卷积生成的键、查询和值的矩阵。这里，dk表示用于键和查询的通道数。这些矩阵中的每一行对应一个输入像素。提出的GSA模块（见图1）由两个平行层组成：内容注意层和位置注意层。
图1
3.1.1 内容注意层
该层使用键、查询和值来生成新特性fc=[fcij]∈ RWH×dout使用以下基于内容的全局注意操作：其中K>表示K的矩阵转置，ρ表示分别对每行应用softmax归一化的操作。此注意操作可解释为首先使用ρK>中的权重将V中的像素特征聚合为dk全局上下文向量，然后使用Q中的权重将全局上下文向量重新分配回各个像素。此操作的计算和内存复杂性为像素数的O（N）。
该注意操作类似于Chen等人（2018）中使用的注意操作；Shen et al.（2018）除了不在查询上使用softmax归一化之外。归一化查询将输出特征约束为全局上下文向量的凸组合。由于这些约束可能会限制注意机制的表达能力，我们删除了查询的softmax归一化。这允许输出特征跨越dk全局上下文向量的整个子空间。当我们对查询进行softmax归一化实验时，ImageNet验证数据集的前1位准确率显著降低（1%）。
3.1.2 位置注意层
内容注意层不考虑像素的空间位置，因此与像素洗牌相同。因此，就其本身而言，它并不最适合处理空间结构数据（如图像）的任务。灵感来源于Bello等人（2019年）；Ramachandran等人（2019年）；Shaw等人（2018），我们通过使用位置注意层来解决这个问题，该位置注意层基于像素自身的内容及其相对于相邻像素的相对空间位置来计算像素的注意图。对于每个像素，我们的位置注意层注意其L×L空间邻居。受轴向公式（Ho et al.，2019；Huang et al.，2019）的启发，我们将此注意层实现为仅列注意层，然后是仅行注意层。在仅列注意层中，输出像素仅注意其列上的输入像素，而在仅行注意层中，输出像素仅注意其行上的输入像素。请注意，仅列注意层后跟仅行注意层有效地导致信息在整个L×L邻域上传播。
允许∆ = {− L−1 2 , .., 0, .., L−1 2}是一组L偏移量，Rc=[Rcδ]∈ rl×dk表示对应于L个空间偏移δ的L个可学习相对位置嵌入的矩阵∈ ∆ 沿着一根柱子。设V c ab=[va+δ，b]∈ rl×dout是由像素（a，b）的L列邻居处的值组成的矩阵。设fcab表示像素（a，b）处仅列位置注意层的输出。然后，我们的专栏式位置注意机制使用相对位置嵌入Rc作为键，可以使用
其中qab是像素（a，b）处的查询。由于每个像素只注意L列邻居，因此该仅列位置注意层的计算和内存复杂性为O（NL），其中N是像素数。类似地，具有O（NL）计算和内存复杂性的仅行位置注意层可以使用L可学习的相对位置嵌入Rr=[rδ]∈ R L×dk对应于L行邻居。在全局轴向注意的情况下，邻域跨越整个列或行，导致O（N√ N） 计算和内存复杂性
GSA模块的最终输出特征图是内容和位置注意层输出的总和。
3.2 GSA网络
GSA网络是一种深度网络，它使用GSA模块而不是空间卷积来模拟像素交互。表1显示了GSA网络与最近各种基于注意力的网络的区别。除Wang等人（2020年）外的所有现有作品要么在网络的后期将其注意力模块作为辅助块插入CNN（Bello等人，2019年；Chen等人，2018年；Huang等人，2019年；Shen等人，2018年；Wang等人，2018年；Yue等人，2018年；Carion等人，2020年；Sun等人，2019年），要么将其注意力机制限制在小范围内局部地区（胡等人，2019年；拉马钱德兰等人，2019年；赵等人，2020年）。相比之下，GSA网络将深度网络中的空间卷积层替换为全局注意模块，并且能够在整个网络中模拟远距离像素交互。Wang等人（2020年）也引入了一个全局注意模块作为空间卷积的替代方案，他们的模块对内容和位置注意都使用轴向注意机制。相比之下，提出的GSA模块使用非轴向全局内容注意机制，该机制一次注意整个图像，而不仅仅是一行或一列。
表1
3.3 理由
建议的GSA模块使用直接全局注意操作进行内容注意，使用轴向注意机制进行位置注意。
为什么不使用轴向注意力做内容交互呢？
轴向注意是一种近似于直接全局注意的机制，只注意列，然后只注意行。在建议的全局内容注意层中，两个像素（i，j）和（p，q）仅基于其内容直接交互。相反，在紧跟行轴内容注意层的列中，像素（i，j）和（p，q）将通过像素（p，j）相互作用，因此，它们的相互作用将不希望地由（p，j）处的内容控制。因此，对于内容注意，提出的直接全局注意机制优于轴向注意机制。表2中的实验结果也证实了这一点。实验结果表明，使用直接全局内容注意的提出GSA模块明显优于轴向注意。
为什么不直接用全局注意力做位置交互呢？
重要的是注意基于相对分辨率的像素位置（而不是绝对位置）以保持平移等变。在内容注意的情况下，每个像素都有一个唯一的关键点，因此，我们可以首先将关键点和值相乘，以使注意机制有效。这在位置注意的情况下是不可能的，因为像素处的关键点基于其相对于查询像素的相对位置而变化。因此，我们使用轴向机制来提高位置注意的效率。虽然轴向注意不利于内容注意（如上所述），但它适合于位置注意。像素（i，j）和（p，q）之间的相对位置与像素（i，j）和（p，j）之间以及像素（p，q）和（p，j）之间的相对位置强相关。因此，（i，j）和（p，q）到（p，j）之间基于路由位置的交互工作良好。
4 实验
模型
除非另有规定，否则我们使用GSA-ResNet-50，这是一种通过将ResNet-50中的所有3×3卷积层（He等人，2016年）替换为建议的GSA模块而获得的网络。我们使用224×224的输入大小，为了减少空间维度，我们在第二、第三和第四残差组的第一个GSA模块之后立即使用2×2平均池层（步幅2）。每个GSA模块中K、Q、V的通道数设置为与相应的输入特性相同。我们使用一种多头注意机制（Ramachandran等人，2019年；Vaswani等人，2017年），每个GSA模块中有8个头。相对位置嵌入在模块内的所有头之间共享，但不在模块之间共享。所有1×1卷积和GSA模块随后进行批量归一化（Ioffe&Szegedy，2015）。
训练和评价
所有模型分别在ImageNet数据集的训练集和验证集上进行训练和评估（Russakovsky等人，2015）。使用动量为0.9的随机梯度下降，基本学习率为0.1的余弦学习率计划，权重衰减为10，从零开始训练他们90个时代−4，最小批量为2048。我们使用标准的数据扩充，如随机裁剪和水平翻转。根据最近基于注意力的研究（Ramachandran等人，2019年；Zhao等人，2020年；Wang等人，2020年），我们还使用了系数为0.1的标签平滑正则化。对于评估，我们使用单个224×224中心裁剪。在计算FLOPs时，乘法和加法是分开计算的。对于报告运行时，我们在TPUv3加速器上测量单个图像的推断时间。
4.1 与卷积运算的比较
图2比较了基于ResNet-{38,50101}结构的CNN和GSA网络。GSA网络在使用较少的参数、计算和运行时间的情况下，性能显著优于CNN。这些结果清楚地显示了所提出的全局注意模块相对于广泛使用的卷积运算的优越性。随着基于注意力的模型的日益普及，我们希望硬件加速器将进一步优化以用于基于注意力的操作，并且GSA网络在不久的将来将比CNN快得多。
图2
4.2 与轴向注意的比较
GSA模块使用一种全局内容注意机制，一次注意整个图像。为了验证这种注意机制相对于轴向注意的优越性，在表2中，我们将提出的GSA模块与基于内容和位置的全局注意模块进行了比较，该模块与Ramachandran等人（2019）类似，但以轴向方式进行。GSA模块明显优于轴向替代方案。此外，我们单独的轴向位置注意的表现与同时使用内容和位置的轴向注意相当，这表明轴向机制无法利用仅内容的交互作用（参见第3.3节的理由）。
表2
4.3 与现有基于注意力的方法的比较
表3比较了GSA网络与最近基于注意力的网络。GSA网络比现有的基于全局和局部注意力的网络实现更好的性能，同时使用相似或更少数量的参数和FLOPs，但与Zhao等人（2020）相比除外；Wang等人（2020年），他们使用了更少的翻牌。与基于局部注意力的作品相比（Hu等人，2019年；Ramachandran等人，2019年；Zhao等人，2020年），提出的GSA网络利用了整个网络的全局注意力，并产生了更好的结果。与Shen等人（2018）相比；Yue等人（2018）；Bello等人（2019年）；Chen等人（2018年）将一些注意力模块作为辅助模块插入CNN，提议的GSA网络通过网络使用全局注意力。与Wang等人（2020年）相比，提出的GSA网络使用了一种非轴向的全局内容注意机制，该机制优于轴向机制。为了报告其他方法的运行时间，我们使用相应作者提供的代码测量TPUv3加速器上的单个图像推断时间。
表3
4.4 消融研究
4.4.1 单个组件的重要性
如第3节所述，GSA模块由三个组件组成：内容注意层、仅列位置注意层和仅行位置注意层。表4显示了通过移除一个或多个组件获得的提出GSA模块不同变体的结果。正如预期的那样，包含所有三个组件的模块性能最好，而仅注意内容的模块性能较差（最高精度下降7.7%），因为它将整个图像视为一袋像素。这清楚地表明了对位置注意的需要，这在许多现有的基于全局注意的作品中是缺失的（Chen等人，2018；Wang等人，2018；Yue等人，2018）。有趣的是，对于位置注意，仅列注意比仅行注意（第3行对第4行，第5行对第6行）表现更好，这表明对ImageNet数据集中类别沿垂直维度的像素交互建模比水平维度更重要。
左边是表4，右边是表5
4.4.2 全局注意在哪里最有帮助？
我们默认的GSA-ResNet-50在ResNet-50的所有残差组中用建议的全局注意模块替换空间卷积。表5显示了当全局注意力仅在某些残差组中取代空间卷积时，性能如何变化。从最后一个残差组开始，当我们进入网络的早期阶段时，用注意力代替卷积可以持续提高性能，直到第二个残差组。替换第一个残差组中的卷积会导致性能略有下降。这些结果表明，除了前几层之外，全局注意机制在整个网络中都是有用的。这是一种预期的行为，因为深度网络的前几层通常专注于学习底层特性。值得注意的是，通过在ResNet-50的第二、第三和第四个残差块中使用建议的GSA模块替换卷积，我们能够实现与基于卷积的ResNet-101相同的top1精度，同时速度显著加快。
4.5 关于CIFAR-100的结果（KRIZHEVSKY&HINTON，2009）
与ImageNet数据集类似，建议的GSA网络在CIFAR-100数据集上的性能显著优于相应的CNN，同时使用较少的参数、计算和运行时间。ResNet-{38,50,101}结构的top-1精度分别提高了2.5%、2.7%和1.6%。更多详情请参考附录中的图3和表6。
5结论
在这项工作中，我们引入了一个新的全局自注意模块，该模块同时考虑了像素的内容和空间位置。该模块由平行内容和位置注意分支组成，其输出在最后汇总。当内容分支使用有效的全局注意机制共同注意所有像素时，位置注意分支遵循轴向公式并执行仅列注意和仅行注意。总的来说，所提出的GSA模块的效率足以成为深度网络的主干组件。基于提出的GSA模块，我们引入了使用GSA模块而不是空间卷积的GSA网络。由于提出GSA模块的全局范围，这些网络能够模拟整个网络中的远程像素交互。我们在CIFAR-100和ImageNet数据集上进行了实验，结果表明，GSA网络在使用较少参数和计算量的情况下明显优于基于卷积的网络。我们还发现，GSA网络的表现优于最近各种基于局部和全局注意力的网络。在不久的将来，我们计划将这项工作扩展到其他计算机视觉任务。
致谢：我们衷心感谢王惠宇、朱玉坤和陆志超的讨论和支持。
CIFAR-100实验
所有模型分别在CIFAR-100的训练和测试部分进行训练和评估。从ImageNet预训练权重开始，使用动量为0.9、权重衰减为10的随机梯度下降，对他们进行10K步的训练−4，最小批量为128。我们使用5×10的初始学习率−3，每3K步减少10倍。对于训练和评估，我们使用224×224个输入图像。
图3比较了CIFAR-100数据集上基于ResNet-{38,50101}结构的CNN和GSA网络。与ImageNet结果类似，GSA网络在CIFAR100数据集上的性能显著优于CNN，同时使用的参数、计算和运行时间更少。表6报告了与图2和图3中的图对应的所有数字。
图3
表6
B.数学实施细节
本节介绍了全局自注意（GSA）模块的数学实现细节，以补充本文第3节中的高层描述。
为了简洁和更接近实际实现，本节使用爱因斯坦符号1。请注意，TensorFlow Abadi等人（2015年）和PyTorch Paszke等人（2019年）通过tf直接支持爱因斯坦符号。einsum（）和火炬。分别是einsum（）。因此，本节中所有方程都有直接的TensorFlow/PyTorch转录。
假设输入X是一个3阶张量，形状为h×w×d，h表示高度，w表示宽度，d表示通道数。
KQV层
第一步是使用3个单独的1×1（即逐点）卷积层从X计算密钥K、查询Q和值V。然后，模块沿通道维度将K、Q、V各分割为n个注意头的n个大小相等的切片。一个高效的实现将这两个步骤融合到一起
其中，W（K）、W（Q）、W（V）是相应的权重，x、y是空间维度，n是头部维度，d、K、V分别是输入、键和查询以及值的通道维度。
内容注意
正如本文第3节所述，在每个头部中，模块使用矩阵乘法来实现内容注意。实际的实现是通过计算在所有头上并行化过程
其中σ表示沿空间尺寸（x，y）的softmax。
位置注意
位置注意层由仅列注意子层、批量归一化层和仅行注意子层组成。由于仅列和仅行子层是对称的，因此本节仅介绍仅列子层的实现。
该层保持嵌入矩阵R的相对位置∈ R（2h）−1） ×k，h为图像高度，k为通道数。每两小时− 1行对应于可能的垂直相对位移，从−（h）− 1） 到h− 1.第一步是重新索引该矩阵，从使用相对移位到使用绝对移位。为了实现这个目标，模块创建了一个重新索引的张量I，其中
式中，L是需要注意的最大相对位移。GSA的默认版本将L=max{h，w}设置为全局位置注意。
然后，该模块计算位置嵌入张量，其指数为绝对位移，如下所示：
现在，仅列注意子层的输出为
在获得Y H后，模块对其应用批处理归一化，并将其用作仅行子层的输入，以生成Y W作为位置注意层的最终输出。
终熔合
在计算了内容和位置注意层的输出之后，最终的输出很简单
与竞争方法的比较
GSA模块的实现仅包括8个爱因斯坦符号方程和5个其他方程，每个方程对应于TensorFlow或PyTorch中的一行代码。与Ramachandran等人（2019）提出的竞争方法相比，该方法的实施更为简单；赵等人（2020年）使用需要定制内核的局部注意。
网友评论