作者丨hdc

编辑丨极市平台

论文链接：https://arxiv.org/abs/2308.00442

代码链接：https://github.com/LeapLabTHU/FLatten-Transformer

在Transformer模型应用于视觉领域的过程中，降低自注意力的计算复杂度是一个重要的研究方向。线性注意力通过两个独立的映射函数来近似Softmax操作，具有线性复杂度，能够很好地解决视觉Transformer计算量过大的问题。然而，目前的线性注意力方法整体性能不佳，难以实际应用。本文深入分析了现有线性注意力方法的缺陷，并提出了一个全新的聚焦的线性注意力模块（Focused Linear Attention），同时具有高效性和很强的模型表达能力。该方法在ImageNet上使DeiT、PVT、PVT-v2、Swin Transformer、CSwin Transformer等模型架构取得了显著的性能提升，能够将模型在CPU端加速约2.0倍，在GPU端加速约1.5倍。

一、简介与研究动机

近年来，视觉Transformer模型得到了极大的发展，相关工作在分类、分割、检测等视觉任务上都取得了很好的效果。然而，将Transformer模型应用于视觉领域并不是一件简单的事情。与自然语言不同，视觉图片中的特征数量更多，由于自注意力是平方复杂度，直接进行全局自注意力的计算往往会带来过高的计算量。针对这一问题，先前的工作通常通过减少参与自注意力计算的特征数量的方法来降低计算量。例如，设计稀疏注意力机制（如PVT）或将注意力的计算限制在局部窗口中（如Swin Transformer）。尽管有效，这样的自注意力方法很容易受到计算模式的影响，同时也不可避免地牺牲了自注意力的全局建模能力。

与这些方法不同，线性注意力将Softmax解耦为两个独立的函数，从而能够将注意力的计算顺序从(query·key)·value调整为query·(key·value)，使得总体的计算复杂度降低为线性。然而，目前的线性注意力方法要么性能明显不如Softmax注意力，要么引入了过多的额外计算量导致模型推理速度很慢，难以实际应用。

（图注：Softmax注意力与线性注意力机制对比）

针对目前线性注意力机制的缺陷，论文创新地提出了聚焦的线性注意力模块，具有很强的模型表达能力并且推理速度很快。具体而言，文章从聚焦能力和特征多样性两个方面分析了现有线性注意力机制的缺陷，并分别设计了聚焦函数和矩阵秩恢复模块作为补偿，最终得到了比Softmax注意力性能更好的线性注意力模块。

二、方法

1. 聚焦能力

先前的一些工作中指出，在自注意力计算中，Softmax提供了一种非线性的权重生成机制，使得模型能够更好地聚焦于一些重要的特征。如下图所示，本文基于DeiT-tiny模型给出了注意力权重分布的可视化结果。可以看到，Softmax注意力能够产生较为集中、尖锐的注意力权重分布，能够更好地聚焦于前景物体；而线性注意力的分布则十分平均，这使得输出的特征接近所有特征的平均值，无法聚焦于更有信息量的特征。

针对这一问题，文章提出了一个简单的映射函数——聚焦函数（Focused Function）。公式如下：

其中，x**p表示对x逐位取p次方。可以证明，在很容易满足的假设下，该聚焦函数可以使得相似的向量内积更大，不相似的向量内积更小，从而实现聚焦。

（图注：聚焦函数效果的数学命题）

为了便于理解聚焦函数的作用，文中给出了一个示例。如下图所示，聚焦函数实际上会把每个特征向量拉向它最靠近的坐标轴。这样一来，模型就自然地把所有特征根据最靠近的坐标轴划分为若干组，不同组特征之间的相似度会减小，同一组特征之间的相似度会得到提升，从而实现模型的聚焦。可以看到，在该示例中，经过聚焦函数处理，注意力权重能够更好地集中于最相似的特征。

（图注：聚焦函数效果示例）

2. 特征多样性

除聚焦能力外，特征多样性也是限制线性注意力性能的一个因素。本文基于DeiT-tiny可视化了完整的注意力矩阵，并计算了矩阵的秩，将Softmax注意力与线性注意力进行对比。

（图注：注意力矩阵的秩）

从上图中可以看到，Softmax注意力可以产生满秩的注意力矩阵，这反映出模型提取到的特征具有多样性。然而，线性注意力无法得到满秩的注意力矩阵。事实上，线性注意力矩阵的秩会被每个head的维度d和特征数量N中的较小者所限制：

一般来说，d比N小，例如在DeiT中d=64，N=196；在Swin Transformer中d=32，N=49。因此，通常线性注意力矩阵无法达到满秩N，这意味着不同行的权重之间存在冗余性。由于自注意力的输出是这些权重对同一组value加权组合得到的，权重的同质化就必然会导致模型输出的多样性下降，进而影响模型性能。

针对这一问题，文中使用轻量化的depthwise convolution来恢复注意力矩阵的秩：

为了便于理解DWC的作用，可以将DWC视为一种特殊的注意力，其中每个query都只会关注到空间上相邻的几个特征，而不是所有的value。此时，即便线性注意力中两个queries对应的输出相似，模型仍然可以从DWC提出的局部特征中得到不同的输出，从而保证特征多样性。DWC的作用也可以从矩阵秩的角度进行解释：

如该公式所示，此时等效的注意力矩阵是原始线性注意力矩阵和DWC对应的稀疏矩阵的和，其中DWC对应的稀疏矩阵有能力取得满秩，因而整个等效矩阵可以取得满秩。如上图(c)中所示，引入DWC之后，等效的注意力矩阵达到了满秩，能够保持特征多样性。

3. 聚焦的线性注意力模块

基于上述分析，论文提出了一种新型线性注意力模块，即聚焦的线性注意力。该模块借助聚焦函数获得更加聚焦的注意力分布，借助DWC模块保持特征多样性，既高效又具有很强的模型表达能力。总体来说，该模块具有以下几个优势：

(1) 计算复杂度低。通过改变自注意力机制的矩阵乘法顺序，本文提出的模块能够将计算复杂度降低为线性。此外，不同于以前一些线性注意力模块设计的复杂核函数，本文使用的聚焦函数和DWC的计算开销很小。

(2) 模型表达能力强。以前的线性注意力模块的性能通常不如Softmax注意力机制。但是，在使用聚焦函数和DWC解决两个性能瓶颈后，本文提出的聚焦的线性注意力可以获得比Softmax注意力机制更好的性能。

(3) 能够采用更大的感受野。得益于线性计算复杂度，本文的模块可以自然地采用更大的感受野，而不会增加模型计算量。例如，可以将Swin Transformer的window size由7扩大为56，即直接采用全局自注意力，而完全不引入额外计算量。

(4) 应用的灵活性强。本文提出的模块是常用的Softmax注意力的一个更优的替代， 可以作为一个插件模块应用于各种各样的ViT模型。

三、实验结果

聚焦的线性注意力是一个通用的注意力模块，文章基于DeiT、PVT、PVT-v2、Swin Transformer、CSwin Transformer等模型架构进行了实验。如下图所示，在ImageNet分类任务中，基于聚焦的线性注意力构建的模型能够取得显著的性能提升。例如，FLatten-Swin-S可以取得与Swin-B相同的性能，而其参数量和计算量不到后者的60%。

（图注：ImageNet图片分类结果）

在检测和分割任务中，相较于基础模型，FLatten Transformer也能够取得一致的性能提升。

（图注：COCO物体检测与分割结果）

（图注：ADE20K语义分割结果）

此外，论文将聚焦的线性注意力模块与先前的一些线性注意力方法进行了对比。如下图所示，先前的线性注意力方法的性能往往不如Softmax注意力，而聚焦的线性注意力模块能够取得超越其他线性注意力和Softmax注意力的性能。

（图注：不同线性注意力方法效果对比）

论文还将模型部署于不同设备上进行了实际测速。如下图所示，在CPU端，FLatten模型能够取得2.0倍左右的加速；在GPU端，能够取得1.4-1.7倍的加速。

（图注：实际测速结果）

四、总结

这篇工作的价值至少体现在三个方面：

(1) 从聚焦能力和特征多样性两个角度分析了线性注意力性能不佳的原因。

(2) 针对线性注意力两方面的缺陷，分别提出了新颖、自然、有效且高效的解决方案，并且给出了充分的理论证明和实验验证。

(3) 提出了聚焦的线性注意力模块，取得了显著的实际效果。该模块具有高效、模型表达能力强等优势，并且能够应用于大感受野和诸多模型架构，这或为开发大尺度、细粒度的视觉、语言大模型提供了新的思路。