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操作，提供了更高效的线性复杂度替代方案。然而，当前的线性注意力方法要么经历显著的性能下降，要么因映射函数引入额外的计算开销。本文提出一种新颖的聚焦线性注意力模块，以同时实现高效性和表现力。具体而言，我们首先从聚焦能力和特征多样性两个角度分析了导致线性注意力性能下降的因素。为克服这些局限性，我们引入一种简单高效的映射函数和一个有效的秩恢复模块，在保持低计算复杂度的同时增强自注意力的表现力。大量实验表明，我们的线性注意力模块适用于多种先进的视觉Transformer，并在多个基准测试上取得持续提升的性能。代码开源在https://github.com/LeapLabTHU/FLatten-Transformer。