PaperlayerView exact revision on arXiv

Source-linked AI summary

FLatten Transformer: Vision Transformer using Focused Linear Attention

Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, Gao Huang

arXiv:2308.00442v2cs.CV

TL;DR

Vision self-attention is costly because its computation scales quadratically with sequence length, while current linear attention methods either lose performance or add mapping overhead. Focused Linear Attention addresses focus ability and feature diversity with a mapping function and rank restoration module. Experiments show consistent improvements across image classification, semantic segmentation, and object detection.

  • Problem

    Vision self-attention has quadratic computation complexity, while existing linear attention methods either suffer performance degradation or introduce mapping-function overhead.

  • Method

    Focused Linear Attention improves attention focus with a mapping function and restores feature diversity through an efficient rank restoration module.

  • Results

    The module consistently improves performance across image classification, semantic segmentation, and object detection using five advanced vision Transformer models.

  • Takeaways & Limitations

    Focused linear attention offers a better trade-off between computation efficiency and model performance across multiple vision Transformer architectures.

  • Takeaways & Limitations

    Current linear attention still faces a trade-off between model complexity and expressiveness, with simple approximations losing performance and complex functions adding overhead.

Abstract

from arXiv · show

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.

1. Introduction

Vision Transformers offer broad applicability, but global self-attention is expensive because its computation grows quadratically with sequence length. Linear attention reduces this cost through reordered computation, yet existing designs can lose performance or add mapping overhead; Focused Linear Attention addresses these limitations and improves results across vision tasks.

  • Vision Transformers support image classification, semantic segmentation, object detection, and multi-modal tasks.
  • Global self-attention incurs quadratic computation O(n^2) with sequence length, creating high costs for vision models.
  • Linear attention changes computation from (query·key)·value to query·(key·value), reducing complexity to O(n) while preserving a global receptive field.
  • Existing linear attention methods can suffer severe performance drops or incur additional computation overhead from mapping functions.
  • Focused Linear Attention improves focus ability through a mapping function and restores feature diversity with a rank restoration module.The module adjusts query and key directions and applies depthwise convolution to the attention matrix.
  • Experiments across image classification, semantic segmentation, and object detection show consistent improvements over baselines and other linear attention approaches.

2. Related Works

Vision Transformer research has reduced attention costs through token reduction and structured attention patterns, while linear attention changes the computation order to avoid pairwise similarity calculations. The paper targets the remaining expressiveness–complexity gap with focused linear attention.

  • Token merging, sparse or deformable patterns, and local windows reduce attention costs but can constrain resolution, regional coverage, or long-range dependencies.
  • Softmax-based attention remains computationally expensive, limiting architectural design and practical application.
  • Linear attention replaces Softmax with kernel functions and computes K^T V before multiplying by Q, reducing complexity from O(N^2d) to O(Nd^2).
  • Existing linear attention designs use approximations ranging from random features and separate normalization to matrix decomposition and cosine similarity.
  • Figure 2 compares different linear attention designs on DeiT-Tiny and Swin-Tiny structures.
  • Focused linear attention addresses limited expressiveness and kernel overhead, aiming for better performance than Softmax attention with lower computation complexity.

3. Preliminaries

Vision Transformer self-attention forms attention maps from query–key similarities, but global pairwise computation is quadratic in token count. Linear attention approximates the similarity function and reorders multiplication to obtain linear token complexity, while facing an expressiveness–complexity trade-off.

  • Self-attention projects input tokens into queries, keys, and values, then computes attention from a similarity function.
  • Softmax attention computes similarities between all query–key pairs, leading to computation complexity O(N^2).
  • Global self-attention can become intractable because quadratic complexity produces excessive computation costs.
  • Linear attention introduces kernels as approximations of the original similarity function.
  • Associative matrix multiplication changes computation from (QK^T)V to Q(K^T V), avoiding explicit pairwise attention computation.
  • Reordered linear attention reduces complexity with respect to token number to O(N).
  • Simple approximations may cause performance drops, whereas complex kernels or matrix decompositions can add computation overhead.

4. Focused Linear Attention

Focused Linear Attention addresses linear attention’s performance degradation through two complementary mechanisms: a focused mapping function that sharpens attention distributions and depthwise convolution that restores feature diversity. Together, these designs retain linear computational complexity while improving expressive capability.

  • Motivation: Previous linear attention approaches suffer performance degradation because their attention distributions are relatively smooth and their attention maps have limited rank.These limitations reduce focus on informative features and can homogenize output features.
  • 4.1. Focus ability: The Focused Function fp adjusts query and key directions to bring similar pairs closer and dissimilar pairs farther apart.It preserves feature norms while changing feature directions, and the exponent p controls the degree of directional adjustment.
  • 4.1. Focus ability: With a suitable p, fp increases the distinction between similar and dissimilar query-key pairs, restoring a sharp attention distribution like Softmax attention.In the illustrated example, scores change from [0.37, 0.19, 0.26, 0.18] to [0.75, 0.11, 0.09, 0.05] after f3.
  • 4.2. Feature diversity: Linear attention’s rank is bounded by the token number N and head channel dimension d, with d typically smaller than N in vision Transformers.For DeiT, d = 64 and N = 196, contributing to homogenized attention rows and resembling aggregated features.
  • 4.3. Focused linear attention module: The resulting Focused Linear Attention combines linear complexity with Softmax-like expressive capability by pairing fp with depthwise convolution.Its computation changes from O(N^2d) to O(Nd^2), where d is typically smaller than N in common vision Transformer designs.

5. Experiments

Experiments evaluate FLatten across classification, semantic segmentation, object detection, comparisons with linear-attention designs, efficiency, and ablations. Results consistently show improved accuracy or trade-offs under comparable computation, with up to 2.1x faster inference.

  • Experimental setup: Experiments cover ImageNet-1K classification, ADE20K semantic segmentation, COCO detection, comparisons with four linear-attention designs, and component ablations.The evaluations use five advanced Vision Transformer models and representative segmentation and detection frameworks.
  • ImageNet-1K classification: 2.7% and 1.9% Top-1 accuracy gains are reported for FLatten-PVT-T/S over PVT-T/S with similar FLOPs.The Swin-based model achieves comparable performance with 60% FLOPs, while PVT-v2- and CSwin-based models improve the computation-performance trade-off.
  • Semantic segmentation: 0.5 ∼1% mIoU improvements are achieved on ADE20K with comparable computation cost and parameters.The passage also reports more significant improvements in mAcc across SemanticFPN and UperNet settings.
  • Object detection: Better results are obtained under all reported COCO detection and instance-segmentation settings.Experiments use Mask R-CNN and Cascade Mask R-CNN with 1x and 3x schedules and different detection heads.
  • Linear-attention comparison: FLatten significantly outperforms other linear-attention designs and the Softmax baseline while using lower computation complexity.The comparison is conducted on DeiT-Tiny and Swin-Tiny structures.
  • Practical efficiency: Up to 2.1x faster inference is achieved with on-par or better performance across tested CPU and GPU platforms.The runtime-accuracy trade-off is evaluated on an Intel i5-8265U CPU, RTX2080Ti, and RTX3090.
  • Ablation study: Replacing Swin-T attention in the first two stages yields a 0.8 performance gain, whereas replacing the last two stages slightly decreases accuracy.The authors associate the first-stage benefit with larger resolutions and suitability for a large receptive field.

6. Conclusion

The paper concludes that focused linear attention addresses focus ability and feature diversity to combine efficiency with expressive capability. Extensive experiments support broad applicability across vision Transformers and improved computation-performance trade-offs.

  • 6. Conclusion: The focused linear attention module addresses limitations of previous linear attention from focus ability and feature diversity perspectives.The proposed module is characterized as combining high efficiency with expressive capability.
  • 6. Conclusion: Experiments across image classification, object detection, and semantic segmentation demonstrate broad applicability and better computation-performance trade-offs.The conclusion covers a variety of vision Transformer architectures.

A. Proof of Proposition 1

The proof explains how the focused function f_p adjusts feature directions and separates similar from dissimilar query-key pairs. This restores a sharper attention distribution by grouping features according to their nearest axes.

  • Focused function: The focused function is defined after ReLU-based non-negativity preprocessing and uses a power p applied bit by bit to x.The function is introduced to restore the sharp distribution in linear attention.
  • Proposition assumptions: The proposition analyzes nonnegative feature pairs with single largest values and distinguishes cases where their largest-value coordinates coincide or differ.The assumptions include 0 < ⟨x, y⟩ < ∥x∥∥y∥ and coordinatewise nonnegative features.
  • Proposition result: With a proper p, f_p creates a larger difference between similar and dissimilar query-key pairs.This is stated through the proposition’s two cases and corresponding equations.
  • Interpretation: f_p groups features by nearest axes, increasing within-group similarity and reducing between-group similarity to restore a Softmax-like sharp attention distribution.Figure 8 is used to visualize attention-weight distributions for this behavior.

B. More Visualizations

Additional attention-weight visualizations isolate the effects of the focused function and depthwise convolution. Their combination restores a sharp distribution resembling Softmax attention.

  • B. More Visualizations: Depthwise convolution improves local focus but does not focus on any particular position, whereas f_p enhances focus on informative regions.The comparison starts from vanilla linear attention and introduces f_p and DWC separately.
  • B. More Visualizations: Combining f_p and DWC restores the sharp attention-weight distribution associated with Softmax attention.The visualizations are presented in Figure 8.

C. Full Classification Results

The supplementary material provides complete ImageNet-1K classification results for the focused linear attention module across multiple sizes of five baseline models, complementing representative main-paper results.

  • C. Full Classification Results: Complete ImageNet-1K classification results are reported for the focused linear attention module across multiple sizes of five baseline models.The main paper presents only representative results in Figure 6, while the supplementary material gives the full results in Table 8.

D. Model Architectures

The supplementary material describes FLatten variants built from five Transformer backbones and details their architectures. The focused linear attention block replaces self-attention throughout DeiT, PVT, and PVTv2, but is used only in early Swin and CSwin stages, with specified depth changes for CSwin-T and CSwin-B.

  • D. Model Architectures: The five adopted Transformer backbones are DeiT, PVT, PVTv2, Swin Transformer, and CSwin Transformer.Their architectures are summarized in the supplementary tables.
  • D. Model Architectures: Focused linear attention replaces the original self-attention blocks at all stages of DeiT, PVT, and PVTv2.The model structure remains unchanged for these variants.
  • D. Model Architectures: For Swin and CSwin, the module is adopted only at early stages; CSwin-T and CSwin-B also shift depth toward the first two stages.Their third-stage depth is correspondingly reduced to reflect the module’s enlarged receptive-field advantage.
Loading 2308.00442v2…