摘要

self-attention 的二次计算复杂度一直是将Transform模型应用于视觉任务时面临的一个持续挑战。另一方面，线性注意力通过精心设计的映射函数逼近Softmax操作，以其线性复杂性提供了一种更有效的替代方案。然而，当前的线性注意力方法要么遭受显着的性能下降，要么引入了映射函数的额外计算开销。在本文中，我们提出了一种新颖的聚焦线性注意力模块，以实现高效率和表现力。具体而言，我们首先从两个视角分析了导致线性注意力性能下降的因素：聚焦能力和特征多样性。为了克服这些限制，我们引入了一个简单而有效的映射函数和一个高效的秩恢复模型来增强自我注意力的表现力，同时保持较低的计算复杂度。广泛的实验表明，我们的线性注意力模块适用于各种高级视觉变压器，并在多个基准上实现了持续改进的性能。
代码：https://github.com/LeapLabTHU/FLatten-Transformer

引言

Preliminary

Vision Transformer and Self-Attention

Linear Attention

然而，目前的线性注意力方法也面临着模型复杂性和表达性之间的困境。一方面，简单的近似，例如使用ReLU激活[2]，过于宽松，导致显着的性能下降。另一方面，精心设计的核函数[7]或矩阵分解方法[26,50]可能会产生额外的计算开销。总的来说，线性注意力的实际性能与Softmax注意力之间仍然存在差距。

Focused Linear Attention

虽然享受线性计算复杂性，但前人的各种工作也证明，简单地用线性注意力替换Softmax注意力通常会导致严重的性能下降[34,2,7,27]。在本节中，我们首先从聚焦能力和特征多样性两个角度详细分析了线性注意力的劣势性能。然后，我们介绍了我们的聚焦线性注意力，它充分解决了这些问题，并实现了高效率和表达能力。

1.Focus ability

Softmax注意力实际上提供了一种非线性重新加权机制，这使得人们很容易专注于重要特征[34,2,58]。如图3所示，Softmax注意力的注意力图分布在某些区域特别清晰，例如前景物体。相比之下，线性注意力的分布相对较快平滑，使其输出更接近所有特征的平均值，并且未能关注更多信息区域。作为补救措施，我们提出了一个简单而有效的解决方案，通过调整每个查询和关键特征的方向，推动相似的查询键对更接近，同时推开不相似的查询键对。具体来说，我们提出了一个简单的映射函数fp，称为聚焦函数：

我们遵循前面的线性注意力模块，首先使用ReLU函数来保证输入的非负性和分母在方程（4）中的有效性。一个直接的观察是，映射后特征的范数保持不变，即 $||x||=||f_p(x)||$，表明只调整了特征方向。
在此基础上，我们证明了在温和假设下，所提出的映射函数$f_p$实际上反映了注意力的分布。

s.t. 受限于

本文的注意力函数 $f_p(\cdot )$，公式7代表了更相似的 query-key pairs； 公式8代表了不相似的 query-key pairs；

如图4所示，可以看出fp实际上将每个向量“拉”到它最近的轴上，p决定了这种“拉”的程度。通过这样做，fp有助于根据特征最近的轴将特征分成几组，提高每组内的相似度，同时减少组之间的相似度。可视化符合我们上面的分析。

2.Feature diversity

线性注意力受限于 注意力的能力 和 特征的差异性；

注意力矩阵的秩可能是导致特征多样性限制的原因之一。举例来说，文章提到了来自DeiT-Tiny模型的一个Transformer层，该层的尺寸为N=14×14。在图5（a）中可以看到，这个注意力矩阵具有完整的秩（196/196），显示了从值（values）中聚合特征时的多样性。

然而，在线性注意力的情况下，这很难实现。事实上，线性注意力中注意力矩阵的秩受标记数量N和每个头的通道维度d的限制：

在这种情况下，注意力矩阵秩的上界被限制在较低的比率，这表明注意力图的许多行被严重均匀化。
而自注意力权重是同一组V的加权和，注意力权重的均匀化必然导致聚合特征之间的相似性。为了更好地说明，我们用线性注意力替换了DeiT-Tiny中最初的Softmax注意力，并在图5（b）中显示了注意力图的秩。可以观察到秩大大降低（196个中的54个），注意力矩阵的许多行是相似的。作为补救措施，我们提出了一个简单而有效的解决方案来解决线性注意力的这种限制。具体来说，在注意力矩阵中添加了一个深度卷积（DWC）模块，输出可以表述为

为了更好地理解这个DWC模块的效果，我们可以将其视为一种注意，其中每个查询将只关注空间中的几个相邻特征，而不是所有特征V。这种局部性保证了即使两个查询对应的线性注意值相同，我们仍然可以从不同的局部羽毛中得到不同的输出，从而保持特征多样性。DWC的效果也可以从矩阵秩的角度来解释。基于等式（10），我们有了

其中我们将$M_{DWC}$表示为对应深度卷积函数的稀疏矩阵，并将$M_{eq}$表示为等价的全注意力图。由于$M_{DWC}$具有全秩矩阵的势，我们实际上增加了等价注意力矩阵秩的上界，这在极大地提高线性注意力性能的同时产生了极少的计算开销。
为了更好地说明，我们对DeiT-Tiny进行了类似的修改。通过额外的DWC模块，线性注意力中注意力图的等级可以恢复到满等级（196 out 196，如图5（c）所示），这保持了特征多样性作为原始Softmax注意力。

3. Focused linear attention module

基于上述分析，我们提出了一种新的线性注意力模块，称为聚焦线性注意力模块，它在保持表达能力的同时降低了计算复杂度。具体来说，我们首先设计了一个新颖的映射函数来模仿原始Softmax注意力的急剧分布。在此基础上，我们关注以前线性注意力模块中的低秩困境，并采用简单的深度卷积来重新存储特征多样性。这样，我们的新模块既可以享受线性复杂性又可以享受高表现力的好处。具体来说，我们的模块可以制定为：

总的来说，我们的模块具有以下优点：
1.作为线性注意力的计算复杂度低。通过改变自注意力的计算顺序，复杂度从$O(N^{2}d)$转换为$O(N{d}^2)$
2.如Softmax关注的高表达能力。正如我们上面所分析的，以前基于核的线性注意力设计在聚焦能力和特征多样性透视方面普遍不如Softmax同类。通过提出的聚焦函数fp和深度卷积，我们的聚焦线性注意力可以获得比Softmax注意力更好的性能。此外，我们的模块还具有适应更大的感受野和不同模型架构的潜力。
基于Softmax注意力的现代Transform模型主要使用有限数量的键/值对，因为对令牌数字的二次复杂度。尽管如此，我们模块的线性复杂度使我们能够将感受野扩展到更大的区域，同时保持相同的计算量，并享受建模远程依赖关系的先进。此外，我们的mod-ule可以作为插件模块，很容易被各种现代视觉变压器架构采用。我们在包括DeiT[39]、PVT[41]、PVT-v2[42]、Swin Trans-form[24]和CSwin Transform[10]在内的五个高级模型上经验性地实现了我们的模块。考虑到扩大感受野的优势，我们在视觉变压器的早期阶段采用聚焦线性注意力块，并保持其余块不变。详细的模型架构如附录所示

实验

推理时间

我们进一步评估了我们模型的实际效率，并将其与两个具有竞争力的基线进行了比较。结果如图7所示。我们在多个硬件平台上测试了推理延迟，包括桌面CPU（Intel i5-8265U）和两个服务器GPU（RTX2080Ti和RTX3090）。可以观察到，我们的模型在CPU和GPU上实现了运行时间和准确性之间的更好权衡，推理速度提高了2.1倍，性能相当甚至更好。

总结

在本文中，我们提出了一种新颖的聚焦线性注意力模块。通过从聚焦能力和特征多样性的角度解决以前线性注意力方法的局限性，我们的模块实现了高效和表达能力的令人印象深刻的结合。图像分类、目标检测和语义分割的广泛实验表明，我们的模块可以广泛应用于各种视觉转换器，并在计算效率和模型性能之间取得更好的权衡。