[Paper Review] FLatten Transformer: Vision Transformer using Focused Linear Attention

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

CCD and CMOS Imaging Sensors15 citations

TL;DR

Introduces Focused Linear Attention to replace Softmax in vision transformers, achieving linear complexity with improved expressiveness via a focused mapping and a depthwise convolution to restore feature diversity. Demonstrates consistent gains across classification, segmentation, and detection benchmarks.

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.

Motivation & Objective

• Address the high computational cost of self-attention in vision Transformers.
• Bridge the performance gap between linear attention and Softmax attention.
• Enhance linear attention with mechanisms to improve focus and feature diversity.
• Provide a plug-in module applicable to multiple Vision Transformer architectures.

Proposed method

• Propose a Focused Linear Attention module combining a simple focus mapping and a rank restoration via depthwise convolution (DWC).
• Approximate Softmax with a mapping function fp that adjusts query/key directions to sharpen attention distribution.
• Apply an additional DWC to V to restore the rank and diversify features.
• Formulate attention as O = Sim(Q,K)V = fp(Q) fp(K)^T V + DWC(V).
• Demonstrate linear-time complexity by reordering computations to Q(K^T V) instead of (QK^T)V.
• Evaluate the module as a plug-in across DeiT, PVT, PVT-v2, Swin, and CSWin backbones on ImageNet, ADE20K, and COCO.

Experimental results

Research questions

• RQ1 Can focused linear attention achieve comparable or superior accuracy to Softmax attention with linear computational cost in vision transformers?
• RQ2 Do simple mapping-based focus adjustments and depthwise convolution-based rank restoration improve expressiveness and feature diversity of linear attention?
• RQ3 Is the Focused Linear Attention module broadly compatible as a plug-in across major vision Transformer architectures?
• RQ4 What are the empirical gains on ImageNet-1K, ADE20K, and COCO when replacing baseline attention with FLatten attention?

Key findings

• Focused Linear Attention improves over vanilla linear attention and can surpass Softmax baselines on several models.
• Introducing fp sharpening plus DWC restores attention rank and feature diversity, yielding accuracy gains (e.g., DeiT-T and Swin-T comparisons).
• On DeiT-Tiny, Swin-Tiny, and other backbones, FLatten achieves higher Top-1 accuracy with comparable FLOPs and parameters.
• Inference latency analyses show up to 2.1x faster runtime with competitive accuracy on CPU/GPU hardware compared to baselines.
• Across benchmarks (ImageNet-1K, ADE20K, COCO), FLatten consistently improves or matches baselines under similar computation budgets.
• Compared to four existing linear attention designs, FLatten achieves higher accuracy (e.g., DeiT-Tiny: 74.1 vs 72.9–70.8; Swin-Tiny: 82.1 vs 80.7–81.8).