Efficient Attention: Attention with Linear Complexities

paper: https://arxiv.org/pdf/1812.01243.pdf

一、引言

注意机制在计算机视觉和自然语言处理中有着广泛的应用。最近的工作开发了点积注意力机制，并将其应用于各种视觉和语言任务。然而，点积注意力的内存和计算成本随着输入的时空复杂度大小呈二次增长。这样的增长阻碍了该机制在大输入上的应用，例如长序列、高分辨率图像或大视频。为了弥补这一缺陷，本文提出了一种新的有效注意力机制，它相当于点积注意力，但大大减少了内存和计算成本。资源效率允许更广泛和灵活地将有效的注意力模块合并到神经网络中，从而提高准确性。由于高昂的成本，这些模型无法包含原始的点积注意力。
对大输入(如长序列、高清图像和大视频)的全局依赖建模仍然是一个未解决的问题。现有点积注意力模块的二次内存和输入大小的计算复杂性限制了它们在如此大输入上的应用。例如，一个非本地模块对于一个64通道128 × 128特征映射使用超过1GB的GPU内存和超过50 GMACC（MACC代表乘法积累，MACC表示1个乘法或加法运算。）的计算。
对大输入的全局依赖关系建模的需求极大地激发了对资源高效注意力算法的探索。对非局部模块的研究揭示了一个有趣的现象。每个位置的注意图尽管是独立生成的，但却是相关的。位置的注意力图主要集中在语义相关的区域。下图显示了在非局部模块中学习到的注意映射。当生成灌木丛前的鸟的图像时，腿上的像素倾向于关注其他腿上的像素以获得结构一致性。同样，主体像素主要关注主体，背景像素关注灌木。

这一观察启发了本文提出的有效注意机制的设计。类似于dotproduct注意，该机制首先从输入生成一个关键特征映射、一个查询特征映射和一个值特征映射。它将键值特征图的每个通道解释为全局注意力图。有效注意力将每个全局注意图作为权重，聚合键值特征图，生成一个全局上下文向量，该向量概括了全局特征的一个方面。然后，在每个位置，模块将查询特征视为全局上下文向量上的一组系数。最后，该模块以查询特征为权重，计算全局上下文向量的和，生成该位置的输出特征。
该算法避免了成对注意矩阵的生成，其大小在输入的时空大小中是二次的。因此，它实现了与输入大小相关的线性复杂性，并获得了显著的效率改进。
本文的主要贡献是有效的注意机制，即:
1。相对于输入的时空大小，具有线性记忆和计算复杂性;
2. 具有与广泛采用的点积注意机制相同的表示力;
3. 允许将更多的注意力模块合并到神经网络中。
除了点积注意力，文献中还有一组单独的技巧被称为注意力机制。这部分将它们称为缩放注意力。点积注意力对于全局依赖关系建模是有效的，而缩放注意力则侧重于强调重要的特征并抑制无信息的特征。例如，挤压-激励(SE)模块使用全局平均池化和线性层来计算每个通道的缩放因子，然后相应地缩放通道。增强后的模型在图像分类方面取得了最先进的性能，在场景分割和目标检测方面取得了实质性的改进。在SE的基础上，CBAM在全局平均池的基础上增加了全局最大池和一个额外的空间注意力子模块。这进一步提高了SE的性能。
尽管两个名字都包含注意力，点积注意力和缩放注意力是两套完全不同的技术，它们的目标非常不同。在适当的时候，可以同时使用这两种技术并让它们一起工作。因此，没有必要将有效注意力与缩放注意技术进行比较。

二、方法实现

将所有n个位置的查询、键、值特征分别表示为$Q\in R^{n\times d_k}$, $K\in R^{n\times d_k}$,$V\in R^{n\times d_v}$，则点积注意力的输出为:

归一化函数有两种常见的选择:

其中${\sigma }_{row}$表示沿矩阵Y的每一行应用softmax函数。下图(左)展示了点积注意力模块。

这种机制的主要缺点是其资源需求。由于它计算每对位置之间的相似性，因此有$n^2$个这样的相似性，这导致$O(n^2)$内存复杂度和$O(d_k{n}^2)$计算复杂度。因此，点积注意力的资源需求在大输入时变得高得令人望而却步。在实践中，该机制的应用只可能在低分辨率特征上。
针对点积注意的缺陷，提出了一种数学上等同于点积注意的高效注意机制。这速度更快，内存效率更高。在高效关注中，单个特征向量$X\in R^{n\times d}$仍然经过三个线性层，形成查询特征$Q\in R^{n\times d_k}$，键特征$K\in R^{n\times d_k}$，值特征$V\in R^{n\times d_v}$。但是，该模块没有将键值特征解释为$R^{d_k}$中的n个特征向量，而是将它们视为$d_k$个单通道特征映射。高效注意力使用这些特征映射中的每一个作为对所有位置的加权，并通过加权求和将所有位置的值特征聚合起来，形成全局上下文向量。这个名字反映了这样一个事实，即向量不对应于特定的位置，而是对输入特征的全局描述。
高效注意机制的特征如下:

其中，${\rho }_q$和${\rho }_k$分别为查询和键特征的归一化函数。实现了与点积注意相同的两种归一化方法。

其中${\sigma }_{row}$， ${\sigma }_{col}$表示分别沿矩阵Y的每一行或每一列应用softmax函数。
高效注意模块是计算机视觉数据处理机制的具体实现。对于输入特征映射$x\in R^{h\times w\times d}$，模块将其扁平化为矩阵$x\in R^{hw\times d}$，对其应用高效关注机制，并将结果重塑为$h\times w\times d_v$。如果$d_v!=d$，则进一步应用1x1卷积将维数恢复为d。最后，将得到的特征与输入特征相加，形成残差结构。

下面是在使用缩放归一化时，点积和高效注意之间等价的正式证明。将式(2)中的比例归一化公式代入式(1)得到：

同理，将式(4)中的比例归一化公式代入式(3)得到:

因为标量乘法和矩阵乘法是符合乘法结合律，则有：

比较式(5)和式(7)，我们得到:

上面的证明适用于softmax规范化变量。QK上的两个softmax操作并不完全等同于QKT上的单个softmax操作。然而，它们非常接近原始softmax函数的效果。${\sigma }_{row}(Q{K}^T)$的临界性质它的每一行之和为1，代表了所有位置的标准化注意力分布。

矩阵${\sigma }_{row}(Q){\sigma }_{col}(K{)}^T$具有这个性质。因此，高效注意的softmax变体是点积注意变体的近似。

高效注意力的解释

高效注意为注意机制提供了新的解释。在点积注意力中，选择位置i作为参考位置，可以收集所有位置与位置i的相似度，形成该位置的注意图$s_i$。注意图$s_i$表示位置i对输入中每个位置j的关注程度。位置j在$s_i$上的值越高，意味着位置i更关注位置j。在点积注意中，每个位置i都有这样一个注意图$s_i$，该机制利用该注意图来聚合值特征V来产生位置i的输出。
相反，高效的注意力不会为每个位置生成一个注意力图。相反，它将键特征$K\in R^{n\times d_k}$解释为$d_k$注意力图$k_j^T$。每个$k_j^T$都是一个全局注意力图，不对应于任何特定的位置。相反，它们中的每一个都对应于整个输入的一个语义方面。例如，一个这样的注意图可能会覆盖输入中的人。另一个可能对应于背景。高效注意力使用每个$k_j^T$来聚合值特征V，并产生一个全局上下文向量$g_j$。由于$k_j^T$描述了输入的全局语义方面，$g_j$也总结了输入的全局语义方面。然后，位置i使用$q_i$作为$g_0,g_1，\dots ,g_{dk-1}$的一组系数。使用前面的示例，人物像素可能会对人物的全局上下文向量施加较大权重，以改进其表示。物体边界的像素可能对物体和背景的全局上下文向量具有较大的权重，以增强对比度。

效率优势

本节分析了高效注意力相对于点积注意力在内存和计算方面的效率优势。高效的注意力不计算每对位置之间的相似性，这将占用$O(n^2)$内存并需要$O(d_k{n}^2)$计算来生成。相反，它只在$R^{d_v}$中生成$d_k$个全局上下文向量。这一改变消除了模块内存和计算复杂性中的$O(n^2)$项。
因此，在$d_v=d$的通用设置下，有效注意具有$O((d_k+d)n+d_{k}d)$内存复杂度和$O((d_{k}d+d^2)n)$计算复杂度。下表给出了高效注意模块和非局部模块(使用点积注意)的复杂度公式。

在计算机视觉中，这种复杂度差异是非常显著的。首先，n本身是图像边长的二次方，在实际应用中往往很大。其次，$d_k$是模块的一个参数，网络设计者可以对其进行调整，以满足不同的资源需求。
在合理的范围内，该参数对性能的影响最小。这个结果意味着一个有效的注意力模块通常可以有一个小的$d_k$，这进一步增加了它比点积注意力的效率优势。下表比较了高效注意力模块与ResBlock的复杂性。

该表显示，高效注意力模块的资源需求与ResBlock相当(在大多数情况下低于ResBlock)，这给了该模块效率水平的直观想法。

三、实验

图4说明了本文的体系结构。从ImageNet预训练中初始化主干。

上表报告了与非局部模块的比较。如res3和fpn3所示，在网络的相同位置插入一个有效的注意力模块或一个非局部模块对神经网络的影响几乎相同的性能。然而，在输入较大的层中，插入非局部模块会导致内存不足错误，而插入有效的注意力模块则会显著提高性能并消耗很少的额外资源。因此，使用每个模块的最佳插入方案，高效注意力模块比非局部模块高出0.9AP和0.8AP。

上表比较了MS-COCO 2017数据集上有效的注意力增强检测器与其他最先进的方法。

消融

插入层

下表显示了在不同层插入EA模块的模型对比实验。当添加单个模块时，读者可以推断出在相同类型的层中(例如在FPN层或ResNet层中)的一般趋势，一层的分辨率越高，插入高效注意力模块带来的性能收益越多。这一结果重申了有效关注的重要性。由于注意力对高分辨率的输入更有利，因此有效注意力的效率优势具有很强的影响力。res4层是唯一一个插入会导致性能下降的层。一个假设是，低分辨率和输入(2048)和键(64)之间的大差距可能是原因。由于实现上的缺陷导致res1和res2的实验失败，因此本文无法获得实验结果。

随后的实验探索了在网络中添加多个有效的注意力模块。在res3、res4和每个FPN级别中添加一个有效的注意力模块，导致1.8AP和1.6AP的改进。
尽管在res4处插入单个模块会降低性能，但从fpn1-5和res3,4模型中删除res4处的模块会导致性能显著降低，如fpn1-5和res3所示。这一结果进一步强化了多跳连接依赖建模的重要性，因为即使是不能独立提高性能的高效注意力模块，也可以在与其他高效注意力模块协作时做出贡献。高效注意力模块的低复杂性使多次插入成为可能，并实现了先前假设的注意力优势。

键的维度

这些实验测试了键的维度对有效注意效果的影响。如下表所示，将键的维数从128降至32只会导致最小的性能变化。这一结果强化了第1节中的假设，即大多数注意力图可表示为一组有限的基本注意力图的线性组合。因此，如果需要进一步节省资源，研究人员可以降低高效注意力模块中的关键字和查询的维数。

骨干架构

本节报告在不同骨干网上的实验。下表报告了使用三种选择的架构ResNet-50、ResNet-101和ResNeXt-101的结果