在将Transformer模型应用于视觉任务时，自注意力的二次计算复杂度一直是一个持续的挑战。而线性注意力则通过精心设计的映射函数来近似 操作，从而提供了更高效的替代方法，其计算复杂度是线性的。然而，当前的线性注意力方法要么遭受显著的性能下降，要么引入额外的计算开销来处理映射函数。在本文中，作者提出了一种新颖的“ Focused Linear Attention ”模块，以实现高效性和表达能力的平衡。
具体而言，作者首先从2个方面分析了导致线性注意力性能下降的因素：焦点能力和特征多样性。为了克服这些局限性，作者引入了一个简单而有效的映射函数和一个高效的排名恢复模块，以提升自注意力的表达能力同时保持低计算复杂度。
广泛的实验表明，作者的线性注意力模块适用于各种先进的视觉Transformer模型，并在多个基准测试中始终取得了改进的性能。
1、简介
近年来，Transformer和自注意力技术在计算机视觉领域取得了巨大发展。随着Vision Transformer的出现，自注意力技术在图像分类、语义分割和目标检测等各种视觉任务中展现出了巨大的潜力。
然而，将Transformer应用于视觉模型是一项不容易的任务。由于自注意力与序列长度相关，其二次计算复杂度导致在使用具有全局感受野的自注意力时计算成本很高。以前的研究试图通过将全局感受野限制在较小的区域来解决这个问题，例如设计稀疏的全局注意力模式或应用较小的注意力窗口。尽管这些方法有效，但它们要么对可能丢弃其他区域中有信息的特征的注意力模式敏感，要么不可避免地牺牲了对长距离依赖性建模的能力。
另一方面，线性注意力被认为是解决计算困境的一种简单而有效的替代方法，通过降低一般复杂性来实现。早期的研究利用局部敏感哈希方案将计算复杂度从降低到。然而，它引入了一个大的常数项，在常见情况下仍然是不可承受的。
更近期的研究注意到，自注意力操作中的 函数实际上强制要求在所有 Query 和键之间进行成对计算，导致了的复杂性，并成为主要的开销。
为了解决这个问题，一些方法采用简单的激活函数或特定的映射函数来近似原始的 函数。如图1所示，通过将计算顺序从改变为，整体计算复杂性可以降低到。然而，与 注意力相比，当前的线性注意力方法仍然存在严重的性能下降，并可能涉及来自映射函数的额外计算开销，限制了它们的实际应用。
在本文中，作者针对当前线性注意力方法的局限性，提出了一种新颖的“ Focused Linear Attention ”模块，旨在实现高效性和表达能力的平衡。
具体而言，作者从两个方面分析了线性注意力性能下降的原因，并提出了相应的解决方案:
- 首先，以前的线性注意力模块中注意力权重的分布相对平滑，缺乏处理最具信息的特征的能力。因此，作者提出了一个简单的映射函数来调整 Query 和键的特征方向，使得注意力权重更加可区分。
- 其次，作者注意到注意力矩阵的降低秩限制了线性注意力的特征多样性。为了解决这个问题，作者提出了一个排名恢复模块，通过对原始注意力矩阵应用额外的深度卷积（DWC），帮助恢复矩阵的秩并保持不同位置的输出特征的多样性。
通过这些改进技术，作者的模块能够在性能上与其 对应物相媲美，同时享受低计算复杂性的好处。
作者通过使用5个先进的视觉Transformer模型在图像分类、语义分割和目标检测任务上对作者的模块的有效性进行了实证验证。实验结果表明，在所有基线和其他线性注意力方法上都取得了一致的改进。
2、相关工作
2.1 Vision Transformer
Transformer和自注意力机制最初是在自然语言处理领域引入的，并在计算机视觉领域引起了广泛的研究兴趣。然而，自注意力的高计算复杂性限制了其直接应用于视觉任务。以前的研究试图从几个方面解决这个问题。先驱性的Vision Transformer 考虑通过将相邻像素合并成一个单独的Token来减少输入分辨率。类似的见解在后续的研究中得到了采用，并扩展到下游任务。
另一条研究线是受卷积神经网络中的金字塔架构启发，考虑逐渐降低特征分辨率，同时采用精心设计的注意力模式来限制关注Token的数量。例如，PVT使用稀疏注意力模式，并从全局角度选择关注Token。DAT遵循这一路径，并设计了可变形注意力模块以实现数据相关的注意力模式。Swin Transformer通过将输入划分为独立的窗口，从局部选择关注Token。NAT遵循卷积的 Query 中心模式，并为所有 Query 设计独立的关注Token。一些研究还注意到卷积操作对于Transformer模型是有价值的，并且可能有助于提高整体效率。CMT将Transformer块与高效卷积算子（如深度卷积）相结合，实现更好的效率-性能权衡。ACmix共享了卷积和自注意力的计算开销，并以较低的代价将两个模块整合在一起。
对于需要高效率的应用场景，MobileFormer维护了两条路径，分别用于卷积和Transformer，并从两个模块中受益。MobileViT利用MobileNets的成功经验，使用mobilenet块和Transformer块的组合来实现轻量级和低延迟。
然而，这些方法仍然依赖于 操作，其固有的高计算复杂性不可避免地导致在模型架构设计和实际应用中的不便。
2.2 线性注意力
除了以上方法，另一条研究线使用线性注意力来解决高计算复杂性的问题。具体而言，线性注意力使用独立的核函数替换自注意力中的 函数。在这种情况下，线性注意力不必先计算的成对相似性。
如图1所示，基于矩阵乘法的关联性质，线性注意力可以通过先计算来改变计算顺序，从而将计算复杂度从降低到。尽管高效，如何设计与 注意力同样有效的线性注意力模块是一个非常复杂的问题。
Performer用正交随机特征来近似 操作。Efficient attention将 函数分别应用于Q和K，这自然地保证了的每一行和为1。Nystromformer和SOFT通过矩阵分解来近似完整的自注意力矩阵。Hydra attention用余弦相似性替换了 ，并提出了降低计算复杂性到的hydra技巧。EfficientVit使用深度卷积来提高线性
注意力的局部特征提取能力。Castling-ViT提出线性角核来衡量每个和之间的频谱相似性。
然而，当前的线性注意力设计要么没有足够的表达能力与 注意力相媲美，要么涉及复杂的核函数导致额外的计算开销。
在本文中，作者从关注能力和特征多样性的角度分析了线性注意力性能下降的原因。基于这些分析，作者提出了一种新颖的线性注意力模块，称为 Focused Linear Attention ，它在低计算复杂性的同时实现了比 注意力更好的性能（图2）。
3、准备工作
3.1、Vision Transformer and Self-Attention
首先，作者重新审视Vision Transformers中自注意力的一般形式。给定输入的N个Token ，在每个 Head 内，自注意力可以表示为：
其中， 是投影矩阵，表示相似性函数。现代视觉Transformer主要采用 attention，其中相似性度量为。在这种情况下，注意力图通过计算所有 Query -键对之间的相似性得到，从而导致计算复杂性为。
由于二次计算复杂性，简单地使用具有全局感受野的自注意力变得难以处理，通常会导致过高的计算成本。先前的研究通过设计稀疏的全局注意力模式或应用较小的注意力窗口来解决这个问题。虽然有效，但这些方法变得对精心设计的注意力模式敏感，或者不可避免地牺牲了对长程依赖性的建模能力。
3.2、Linear Attention
相比之下，线性注意力被认为是一种有效的替代方法，将计算复杂性从限制到。具体来说，引入了精心设计的核函数作为原始相似性函数的近似，即
其中自注意力模块可以重写为：
这样，作者可以基于矩阵乘法的结合性质（如图1所示）将计算顺序从改变为，其中相对于Token数的计算复杂性被降低为。
然而，当前的线性注意力方法在模型复杂性和表达能力之间面临着困境。一方面，简单的近似方法，例如使用ReLU激活，过于宽松，导致性能显著下降。另一方面，精心设计的核函数或矩阵分解方法可能导致额外的计算开销。
总体上，线性注意力和 注意力之间仍然存在实际性能的差距。
4、Focused Linear Attention
尽管线性注意力具有线性的计算复杂性，但以往的各种研究也证明，简单地用线性注意力替换 注意力通常会导致性能显著下降。
在本节中，作者首先从 Focus 能力和特征多样性两个角度对线性注意力的较差性能进行了详细分析。然后，作者介绍了作者的 Focused Linear Attention 方法，它充分解决了这些问题，并实现了高效率和表达能力。
4.1、Focus ability
实际上提供了一种非线性重新权重机制，能够轻松地集中于重要的特征。如图3所示，从 得到的注意力图分布在某些区域上尤其突出，例如前景对象。相比之下，线性注意力的分布相对平滑，使其输出更接近所有特征的平均值，并未能集中于更具信息量的区域。
为了解决这个问题，作者提出了一个简单而有效的解决方案，通过调整每个 Query 和关键特征的方向，使相似的 query-key 对靠近，同时将不相似的 query-key 对远离。具体来说，作者提出了一个称为"Focused Function"的简单映射函数：
其中表示的逐元素次方。作者首先使用ReLU函数来确保输入的非负性以及公式(4)中分母的有效性，这与以前的线性注意力模块相同。直观观察发现，在映射之后，特征的范数保持不变，即，这表明只有特征方向被调整。
在此基础上，作者证明了在温和的假设下，提出的映射函数实际上影响了注意力的分布。
「命题1（使用进行特征方向调整）」设，。假设和分别具有最大值和。对于一对特征，，其中：
因此，通过适当选择，作者的Focused Function 实际上在相似的 query-key 对（公式(15)）和不相似的 query-key 对（公式(16)）之间实现了更显著的差异，从而恢复了原始 函数的注意力分布。
为了更好地理解，作者在图4中给出了一个例子展示了的效果。可以看到， 实际上将每个向量“拉”到最近的轴线上，而p决定了这种“拉”的程度。通过这样做， 有助于根据最近的轴线将特征分成几组，提高每个组内的相似性，同时降低组之间的相似性。这些可视化结果与作者上面的分析是一致的。
4.2、Feature diversity
除了注意力集中能力之外，特征多样性也是限制线性注意力表达能力的因素之一。其中可能的原因之一是要归功于注意力矩阵的秩，在这里可以看到明显的差异。以DeiT-Tiny中的一个Transformer层为例，假设，从图5 (a)可以看出，注意力矩阵具有完全秩（196中的196），显示出从values聚合特征时的多样性。
然而，在线性注意力的情况下，很难达到这一点。实际上，线性注意力的注意力矩阵秩受到 Token 数量N和每个 Head 的通道维度的限制：
其中通常比常见的视觉Transformer设计中的N小，例如DeiT中，，Swin Transformer中，。
在这种情况下，注意力矩阵秩的上限受到较低比例的限制，这意味着许多注意力图的行严重同质化。由于自注意力的输出是相同的V的加权和，注意力权重的同质化不可避免地导致聚合特征之间的相似性。
为了更好地说明这一点，作者将DeiT-Tiny中的原始 替换为线性注意力，并在图5 (b)中显示注意力图的秩。可以观察到秩大大降低（从196降为54），而且注意力矩阵的许多行是相似的。
为了解决线性注意力的这一局限性，作者提出了一个简单而有效的解决方案。具体而言，作者添加了一个深度卷积（DWC）模块到注意力矩阵，并且输出可以表示为：
为了更好地理解这个DWC模块的效果，作者可以将它视为一种注意力，其中每个 Query 只关注空间中的几个相邻特征，而不是所有特征V。
这种局部性确保即使与两个 Query 对应的线性注意力值相同，作者仍然可以从不同的局部特征中得到不同的输出，从而保持特征的多样性。DWC的效果也可以从矩阵秩的角度解释。根据方程（10），作者有：
其中作者用表示对应于深度卷积函数的稀疏矩阵，用表示等效的完全注意力图。由于有可能是一个完全秩矩阵，作者实际上增加了等效注意力矩阵秩的上限，这在计算开销很小的情况下极大地提高了线性注意力的性能。
为了更好地说明，作者对DeiT-Tiny进行类似的修改。通过添加额外的DWC模块，线性注意力中注意力图的秩可以恢复到完全秩（如图5 (c)中所示，从196恢复为196），保持了原始 的特征多样性。
4.3、Focused linear attention module
基于上述分析，作者提出了一种新颖的线性注意力模块，称为“Focused Linear Attention”，它在保持表达能力的同时降低了计算复杂性。具体而言，作者首先设计了一种新颖的映射函数，以模拟原始 的尖锐分布。
在此基础上，作者关注了先前线性注意力模块中的低秩困境，并采用了简单的深度卷积来恢复特征多样性。通过这种方式，作者的新模块可以兼具线性复杂性和高表达能力的优势。
具体而言，作者的模块可以表示为：
一般来说，作者的模块具有以下优点：
1、计算复杂性低，类似线性注意力
通过改变自注意力的计算顺序，复杂性从转换为，其中和分别表示 Token 数量和每个 Head 的通道维度。在常见的视觉Transformer设计中，通常比小，例如DeiT中，，Swin Transformer中，，整体计算复杂性实际上得到了降低。
此外，与先前设计复杂的kernel函数的线性注意力模块相比，作者提出的Focus函数 只采用简单的运算符，以最小的计算开销实现了近似。
2、表达能力强，类似
正如作者上面分析的，先前基于kernel的线性注意力设计在注意力集中能力和特征多样性方面通常不如 。借助于作者提出的 Focus 函数 和深度卷积，作者的 Focused Linear Attention 可以实现比 注意力更好的性能。
此外，作者的模块还具有适应更大感受野和不同模型架构的潜力。基于 的现代Transformer模型主要使用有限数量的键值对，因为 Token 数量带来的二次复杂性限制了其数量。然而，作者模块的线性复杂性使得作者可以将感受野扩展到更大的区域，同时保持相同数量的计算，并享受建模长距离依赖性的优势。
此外，作者的模块可以作为一个插件模块，轻松地应用于各种现代视觉Transformer架构。作者在5个先进的模型上进行了实验，包括DeiT，PVT，PVT-v2，Swin Transformer和CSwin Transformer。考虑到扩大感受野的优势，作者在视觉Transformer的早期阶段采用了 Focused Linear Attention 块，其余的块保持不变。
4.4、模型架构说明
1、FLatten-DeiT
2、FLatten-PVT
3、FLatten-PVTv2
4、FLatten-Swin
5、FLatten-CSwin
5、实验
5.1、ImageNet-1K分类
分类结果如图6所示。可以看到，在相当的FLOPs或参数下，作者的方法相对于基准模型实现了一致的改进。例如，作者的FLatten-PVT-T/S相对于PVTT/S分别提高了2.7%和1.9%，而FLOPs相似。
基于Swin，作者的模型在60%的FLOPs下达到了可比较的性能。基于PVT-v2和CSwin的模型在计算成本和模型性能之间实现了更好的平衡。这些结果表明，作者的模块具有很高的表达能力，并适用于各种模型结构。
5.2、语义分割
ADE20K是一种广泛采用的20K/2K训练/验证图像语义分割基准。作者在Semantic FPN和UperNet这两个具有代表性的分割模型上使用了作者的模型。
如表1所示，作者的模型在所有设置下都能持续获得更好的结果。具体来说，作者可以看到0.5～1%mIoU的改进，计算成本和参数相当。mAcc的改进更为显著。
5.3、目标检测
COCO目标检测和实例分割数据集具有118K个训练图像和5K个验证图像。作者在Mask R-CNN和Cascade Mask R-CN框架中使用ImageNet预训练模型作为Backbone来评估有效性。
作者用不同的检测Head对1x和3x时间表进行了实验，结果如表2所示。利用更大的感受野，作者的模型在所有设置下都显示出更好的结果。
5.4、与其他线性注意点的比较
为了与其他线性注意力模块进行公平的比较，作者分别基于两种具有代表性的视觉Transformer结构DeiT和Swin-Transformer进行了实验。
基于这两个模型，作者将作者的 Focused Linear Attention 模块与之前的4种线性注意力设计进行了比较，包括Hydra Attention, Efficient Attention, Linear Angular Attention和Enhanced Linear Attention 。
如表3所示，以前的线性注意力模块通常不如 对应模块，而作者的模型显著优于所有其他设计和 基线。
这表明作者的模块具有很高的表达能力，可以在较低的计算复杂度下获得比 更好的性能。
5.5、Inference Time
作者进一步评估了作者模型的实际效率，并将其与两个竞争基线进行了比较。结果如图7所示。作者在多个硬件平台上测试了推理延迟，包括一个台式CPU（Intel i5-8265U）和2个服务器GPU（RTX2080Ti和RTX3090）。
可以观察到，作者的模型在CPU和GPU上的运行时间和准确性之间实现了更好的权衡，推理速度快2.1倍，性能相当甚至更好。
5.6、消融实验
1、Focused function and DWC
作者首先评估了作者提出的 Focus 函数 和深度卷积的有效性。作者从基本的线性注意力开始，依次引入 和DWC。
如表4所示，采用作者提出的 Focus 函数 提供了+1.3的改进。使用DWC来保持特征多样性进一步带来了+2.3的准确度增益，达到了总体准确率74.1。这些结果证明作者提出的 和DWC可以极大地提高线性注意力的表达能力，从而帮助作者的 Focused Linear Attention 模块实现比 注意力更好的性能。
2、Ablation on different
在表5中，作者研究了 Focus 因子对模型性能的影响。作者发现当p在2和32之间变化时，Top-1分类准确率并没有发生很大变化，这表明作者的模块对这个超参数具有鲁棒性。实际上，在本文中作者选择作为所有模型的 Focus 因子，无需额外调整。
3、Receptive field
作者还研究了基于FLatten-Swin-tiny模型的感受野对性能的影响。如表6所示，随着窗口大小的增加，作者的模型性能逐渐提高。这进一步证明了虽然窗口注意力是有效的，但它从全局的角度来看，不可避免地牺牲了自注意力的长距离依赖性，仍然不如全局注意力。
借助线性复杂性，作者的模块有可能在保持相同计算量的情况下实现更大甚至全局的感受野。
4、Focused linear attention at different stages
作者在不同阶段用作者的模块替换了Swin-T的移位窗口注意力。如表7所示，作者可以看到替换前两个阶段会带来0.8的性能提升，而替换最后两个阶段会略微降低整体准确率。作者将这个结果归因于Swin的前两个阶段具有较大的分辨率，更适合作者具有大感受野的模块。
5、可视化对比
在图8中，作者可视化了更多的注意力权重示例。为了更好地展示作者的 Focus 函数和DWC的贡献，作者从基本的线性注意力开始，分别引入 和DWC。如最后三行所示，DWC提高了局部注意力能力，但无法集中在任何位置，而 实际上增强了模型的注意力能力，帮助模型集中在更具信息量的区域。结合 和DWC，作者的 Focused Linear Attention 模块恢复了与原始 相同的尖锐分布。
6、参考
[1].FLatten Transformer: Vision Transformer using Focused Linear Attention.
[2].https://github.com/LeapLabTHU/FLatten-Transformer.
