Flatten Transformer: 使用聚焦线性注意力的视觉Transformer

原文链接：

0.摘要

1.引言

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

2.相关工作

        近年来，Transformer和自注意力在计算机视觉领域得到了广泛的发展。随着Vision Transformer [11, 39]的出现，自注意力技术在包括图像分类[41,43, 30, 46]，语义分割[6, 49]，目标检测[4, 61, 22]和多模态任务[35, 31]在内的各种视觉任务中展示出了巨大的潜力。然而，将Transformer应用于视觉模型是一项不容易的任务。与轻量级卷积神经网络[37, 16, 44, 33]不同，自注意力在序列长度n方面的二次计算复杂度O(n2)导致了在使用具有全局感受野的自注意力时高计算成本。以前的研究通过将全局感受野限制在较小的区域内来缓解这一挑战，例如设计稀疏的全局注意力模式[41, 46]或应用较小的注意力窗口[24, 17]。虽然这些方法是有效的，但是它们要么容易忽略其他区域的有用特征，要么不可避免地牺牲了对长程依赖性的建模能力。

另一方面，线性注意力被认为是一种简单而有效的替代方案，可通过降低一般复杂度来解决计算困境。早期的研究利用一种局部敏感哈希方案[21]，将计算复杂度从O(n2)压缩到O(nlog(n))。然而，它引入了大的常数项，使得在常见情况下仍然无法承受。最近的研究已经注意到，自注意力操作中的Softmax函数的利用实际上强制要求所有查询和键之间进行一对一的计算，导致主要的O(n2)复杂度。为解决这个问题，一些方法采用简单的激活函数[19, 38]或定制的映射函数[7, 26]来近似原始的Softmax函数。如图1所示，通过将计算顺序从(query·key)·value更改为query·(key·value)，可以将整体计算复杂度降低到O(n)。然而，与Softmax注意力相比，当前的线性注意力方法仍然面临严重的性能下降，并可能涉及来自映射函数的额外计算开销，从而限制了它们的实际应用。

本文针对当前线性注意力方法的局限性，提出了一种新颖的Focused Linear Attention模块，既实现了高效性又保持了表达能力。具体而言，我们对导致线性注意力性能下降的因素进行了双重分析，并随后提出了相应的解决方案。首先，以前的线性注意力模块中注意力权重的分布相对平滑，缺乏处理最有信息的特征的聚焦能力。因此，我们提出了一个简单的映射函数来调整查询和键的特征方向，使注意力权重更具可区分性。其次，我们注意到注意力矩阵的降维会减少线性注意力中特征的多样性。为了解决这个问题，我们提出了一个降维恢复模块，通过对原始注意力矩阵应用额外的深度卷积（DWC），帮助恢复矩阵的秩并使不同位置的输出特征多样化。利用这些改进的技术，我们的模块展示出与其Softmax注意力对应物相当或更优的性能，同时享受低计算复杂度的好处。

我们在图像分类，语义分割和目标检测任务中使用了五个先进的视觉Transformer模型，经验验证了我们模块的有效性。结果表明，我们的模块在所有基线和其他线性注意力方法上都得到了一致的改进。

图1：Softmax注意力和线性注意力之间的差异。其中，Q、K、V ∈ R N×d 分别表示查询、键和值矩阵。Softmax注意力强制计算查询和键之间的成对相似度，导致复杂度为O(N 2 d)。线性注意力通过适当的近似和计算顺序变化将Softmax操作分离，先计算KT V，从而导致复杂度为O(N d2)。考虑到现代视觉Transformer设计中，通道维度d通常比令牌数量N小，例如DeiT[39]中d=64，N=196，Swin Transformer[24]中d=32，N=49，线性注意力模块实际上可以在节省总体计算成本的同时，享受更大的感受野和更高的吞吐量的好处。

2.1.视觉Transformer

        Transformer和自注意机制首先在自然语言处理领域引入，后来在计算机视觉领域也受到广泛研究。然而，自注意机制的高计算复杂度制约了其直接应用于视觉任务的可行性。先前的研究从几个方面尝试解决这个问题。先锋性的Vision Transformer [11]考虑将输入分辨率降低，将相邻像素合并为一个令牌。这种洞察力也被应用于后续的研究[55，54]并扩展到下游任务[22]。另一条研究路线逐渐降低特征分辨率，并采用精心设计的注意力模式来限制关注令牌的数量。例如，PVT [41,42]使用稀疏注意力模式并从全局角度选择关注的令牌。DAT [46] 沿着这条路线设计了一种可变形的注意力模块，以实现数据依赖的注意力模式。Swin Transformer [24]通过将输入划分为隔离的窗口，局部地选择关注令牌。NAT [17]遵循卷积中的查询中心模式，并为所有查询设计独立的关注令牌。一些研究也注意到卷积操作对于Transformer模型是有价值的，并可能有助于提高整体效率[48]。CMT [12]将Transformer块与有效的卷积操作符（例如深度卷积[37]）相结合，并实现了更好的效率-性能权衡。ACmix [30]共享卷积和自注意力的计算开销，并以有限的成本集成两个模块。已经提出了用于Transformer高效训练的方法[45，29]。在需要高效率的应用场景中，MobileFormer [5]分别维护卷积和Transformer两个路径，并从两个模块中受益。Dyn-Perceiver [13]通过动态早期退出[15，14，51]实现了高效的视觉识别。MobileViT [28]利用MobileNets[37]的成功经验，使用mobilenet块和Transformer块的组合来实现轻量级和低延迟。然而，这些方法仍然依赖于Softmax运算符，其固有的高计算复杂度不可避免地导致模型架构设计和实际应用上的不便。

2.2.线性注意力

        除了上述方法之外，另一条研究路线通过线性注意力[19]来解决高计算复杂度的问题。具体而言，线性注意力用单独的内核函数替换自注意力中的Softmax函数。在这种情况下，线性注意力不必先计算两个矩阵之间的相似度QKT。正如图1所示，基于矩阵乘法的结合律，线性注意力可以通过先计算KT V来改变计算顺序，从而将计算复杂度从O(N2d)降低到O(N d2)。虽然高效，但如何设计与Softmax注意力一样有效的线性注意力模块是一个不容易的问题。Performer [7]用正交随机特征逼近Softmax操作。Efficient attention [38]分别将Softmax函数应用于Q和K，从而自然地确保QKT的每一行之和为1。Nystromformer [50]和SOFT [26]通过矩阵分解逼近全自注意力矩阵。Hydra attention [1]用余弦相似度替换Softmax，并提出hydra trick，将计算复杂度降低到O(N d)。EfficientVit [2]使用深度卷积来提高线性注意力的局部特征提取能力。Castling-ViT [52]提出线性角核来衡量每个Qi和Kj之间的谱相似度。然而，当前的线性注意力设计要么没有足够的表达能力来赶上Softmax注意力，要么涉及到复杂内核函数的额外计算开销。在本文中，我们从关注能力和特征多样性的角度分析了线性注意力性能下降的原因。基于这些分析，我们提出了一种新的线性注意力模块，称为聚焦线性注意力，它在比Softmax注意力更低的计算复杂度下实现了更好的性能（图2）。

图2. DeiT-Tiny和Swin-Tiny结构上不同线性注意力设计的比较。

3.准备工作

3.1.视觉Transformer和自注意力

        首先，我们重新审视Vision Transformers中自注意力的一般形式。在每个头部中，给定输入N个形如x ∈ R N×C的token，自注意力可以表示为：

其中，WQ，WK，WV ∈ R C×C是投影矩阵，Sim(·,·)表示相似性函数。现代Vision Transformers主要采用Softmax attention[40]，其中相似性函数的度量方式是Sim(Q, K) = exp(QKT/√d)。在这种情况下，通过计算所有查询-键值对之间的相似性得到了注意力图，这导致了O(N2)的计算复杂度。

由于二次计算复杂度，仅仅使用具有全局感受野的自注意力就变得难以处理，通常会导致过高的计算成本。以前的研究通过设计稀疏的全局注意力模式[41，46]或应用较小的注意力窗口[24，10]来解决这个问题。虽然有效，但这些方法对精心设计的注意力模式变得敏感，或者不可避免地牺牲了模拟长距