[论文解读] FLatten Transformer: Vision Transformer using Focused Linear Attention

Dongchen Han, Xuran Pan+3 more|arXiv (Cornell University)|Aug 1, 2023

CCD and CMOS Imaging Sensors被引用 15

一句话总结

引入 Focused Linear Attention，以取代 vision transformers 中的 Softmax，实现线性复杂度，并通过聚焦映射和深度卷积来恢复特征多样性，从而提升表达能力。在分类、分割和检测基准上显示出一致的提升。

ABSTRACT

The quadratic computation complexity of self-attention has been a persistent challenge when applying Transformer models to vision tasks. Linear attention, on the other hand, offers a much more efficient alternative with its linear complexity by approximating the Softmax operation through carefully designed mapping functions. However, current linear attention approaches either suffer from significant performance degradation or introduce additional computation overhead from the mapping functions. In this paper, we propose a novel Focused Linear Attention module to achieve both high efficiency and expressiveness. Specifically, we first analyze the factors contributing to the performance degradation of linear attention from two perspectives: the focus ability and feature diversity. To overcome these limitations, we introduce a simple yet effective mapping function and an efficient rank restoration module to enhance the expressiveness of self-attention while maintaining low computation complexity. Extensive experiments show that our linear attention module is applicable to a variety of advanced vision Transformers, and achieves consistently improved performances on multiple benchmarks. Code is available at https://github.com/LeapLabTHU/FLatten-Transformer.

研究动机与目标

• 解决视觉 Transformer 中自注意力的高计算成本。
• 弥合线性注意力与 Softmax 注意力之间的性能差距。
• 通过提升聚焦和特征多样性的机制来增强线性注意力。
• 提供可应用于多种 Vision Transformer 架构的插件模块。

提出的方法

• 提出一个聚焦线性注意力模块，结合简单的聚焦映射和通过深度卷积（DWC）进行的秩恢复。
• 用一个映射函数 fp 来近似 Softmax，该函数调整查询/键的方向以锐化注意力分布。
• 对 V 应用额外的 DWC 以恢复秩并丰富特征。
• 将注意力公式表述为 O = Sim(Q,K)V = fp(Q) fp(K)^T V + DWC(V)。
• 通过将计算重新排序为 Q(K^T V) 而不是 (QK^T)V 来展示线性时间复杂度。
• 在 ImageNet、ADE20K 和 COCO 上，将该模块作为插件应用于 DeiT、PVT、PVT-v2、Swin 和 CSWin 骨干网络的评估。

实验结果

研究问题

• RQ1 聚焦线性注意力是否能够在视觉 Transformer 中以线性计算成本达到与 Softmax 注意力相当或更高的准确度？
• RQ2 简单的基于映射的聚焦调整和基于深度卷积的秩恢复是否能提高线性注意力的表达能力和特征多样性？
• RQ3 聚焦线性注意力模块是否可作为插件广泛兼容主流的视觉 Transformer 架构？
• RQ4 在将基线注意力替换为 FLatten 注意力时，在 ImageNet-1K、ADE20K 和 COCO 上的经验收益是多少？

主要发现

• 聚焦线性注意力在多种模型上优于原生线性注意力，并且可超过 Softmax 基线。
• 引入 fp 锐化与 DWC 恢复注意力秩和特征多样性，从而带来准确率提升（例如 DeiT-T 与 Swin-T 的比较）。
• 在 DeiT-Tiny、Swin-Tiny 及其他骨干网络上，FLatten 在 FLOPs 和参数量相近的条件下实现更高的 Top-1 准确率。
• 推理时延分析表明，在 CPU/GPU 硬件上，与基线相比，运行时可提速至最多 2.1 倍，且准确率具有竞争力。
• 在各基准测试（ImageNet-1K、ADE20K、COCO）上，FLatten 在相似计算预算下持续提升或匹配基线。
• 相比四种现有线性注意力设计，FLatten 取得更高的准确率（例如 DeiT-Tiny: 74.1 对 72.9–70.8；Swin-Tiny: 82.1 对 80.7–81.8）。