Flatten Attention

本文作者提出一种聚焦线性注意力机制 (Focused Linear Attention)，力求实现高效率和高表达力。

在将 Transformer 模型应用于视觉任务时，自注意力机制 (Self-Attention) 的计算复杂度随序列长度的大小呈二次方关系，给视觉任务的应用带来了挑战。各种各样的线性注意力机制 (Linear Attention) 的计算复杂度随序列长度的大小呈线性关系，可以提供一种更有效的替代方案。线性注意力机制通过精心设计的映射函数来替代 Self-Attention 中的 Softmax 操作，但是这种技术路线要么会面临比较严重的性能下降，要么从映射函数中引入额外的计算开销。

本文作者提出一种聚焦线性注意力机制 (Focused Linear Attention)，力求实现高效率和高表达力。作者首先分析了是什么导致了线性注意力机制性能的下降？然后归结为了两个方面：聚焦能力 (Focus Ability) 和特征丰富度 (Feature Diversity)，然后提出一个简单而有效的映射函数和一个高效的秩恢复模块来增强自我注意的表达能力，同时保持较低的计算复杂度。

Flatten Attention：聚焦的线性注意力机制构建视觉 Transformer

论文名称：FLatten Transformer: Vision Transformer using Focused Linear Attention (ICCV 2023)

论文地址：

http://arxiv.org/pdf/2308.00442.pdf

现有线性注意力机制的不足之处

 图1：标准的带 Softmax 的 Self-Attention 和 Linear Attention 的区别

但是，与 Softmax 注意力相比，当前的 Linear Attention 方法仍然存在严重的性能下降，并且可能涉及映射函数的额外计算开销，限制了它们的实际应用。

作者从两个角度分析了线性注意力的性能下降的原因，并提出了相应的解决方案。

1. Linear Attention 的注意力分布相对平滑，缺乏解决信息量最大的特征的聚焦能力。
2. Linear Attention 矩阵的秩相对较低，限制了特征多样性。

线性注意力机制的聚焦能力不够及其解决方案

Softmax Attention 实际上提供了一种非线性重新加权机制，这使得很容易聚焦于重要的特征。如下图1所示，可以看到 Softmax 注意力图的分布在某些区域 (如前景对象) 上比较集中，但是线性注意力图的分布却比较平滑，未能对包含更多信息的区域提供更有效的关注。

图1：线性注意力机制的聚焦能力不够如何为线性注意力机制带来更多的聚焦能力？作者提出了一种聚焦线性注意力机制 (Focused Linear Attention)，保留3式的计算方法，同时映射函数写成：

 

图2：黑色的向量代表 Query，其他几个彩色的向量代表不同的 Key，左图是映射函数作用前的结果，右图是映射函数作用之后的结果

命题结论和可视化结果基本一致。

线性注意力机制的特征丰富度不够及其解决方案

除了焦点能力之外，特征多样性也是对线性注意力表达能力造成限制的因素之一。如下图3所示，以 DeiT-Tiny 中的一个 Transformer 层为例，可以看到 Self-Attention 是满秩矩阵，特征的多样性比较丰富。

图3：不同 Attention 注意力图可视化和对应的秩，左：Softmax Attention 矩阵，中：Linear Attention 矩阵，带有 DWC 模块的 Linear Attention 矩阵但是，在线性注意的情况下，满秩矩阵很难实现。因为 Linear Attention 的秩受下式的制约：

如上图3所示，使用额外的 DWC 模块，线性注意力中注意力图的排名可以恢复到满秩 196，提升了特征的多样性。

聚焦线性注意力机制

基于上面两节的分析，本文提出一种聚焦线性注意力机制，在保持表达能力的同时降低了计算复杂度，可以表述为：

它有2个主要的优点：

• 计算复杂度很低，和线性注意力机制相当。

与前任设计的复杂核函数的线性注意模块相比，本文提出的聚焦函数只采用了简单的算子，以最小的计算开销实现了近似。

• 较高的表达能力，和 Softmax 注意力一致。

前人的基于核函数的线性注意力设计通常不如 Softmax Attention。通过本文所提出的聚焦函数和深度卷积，聚焦线性注意力机制可以实现比 Softmax Attention 更好的性能。

实验结果

ImageNet-1K 图像分类

实验结果如下图4所示，本文方法在差不多的 FLOP 或 Params 下与基线模型实现了一致的改进。比如，FLatten-PVT-T/S 在相似的 FLOP 下分别比 PVT-T/S 高出 2.7% 和 1.9%。基于 Swin，本文模型实现了与 60% FLOPs 相当的性能。这些结果说明 Flatten 方法对不同模型具有泛化能力。

图4：ImageNet 图像分类实验结果

ADE20K 语义分割

实验结果如下图5所示，本文在两个具有代表性的分割模型 SemanticFPN 和 UperNet 上使用 Flatten 方法。如表中所示。如图 1 所示，我们的模型在所有设置下都取得了始终更好的结果。具体来说，我们可以看到 0.5 ∼ 1% mIoU 的改进，计算成本和参数相当。mAcc 的改进更加显著。

图5：ADE20K 语义分割实验结果

COCO 目标检测

实验结果如下图6所示，作者使用 ImageNet 预训练模型和 Mask R-CNN 和 Cascade Mask R-CNN 检测头框来评估有效性，在不同检测头的 1x 和 3x schedule 设置下进行了实验。结果显示，利用更大的感受野，带有 Flattened Attention 的模型在所有设置下都显示出更好的结果。

图6：COCO 目标检测实验结果

推理时间

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

图7：不同硬件设备的推理时间对比