The Linear Attention Resurrection in Vision Transformer

TL;DR

This paper introduces L²ViT, a novel vision transformer architecture that combines linear attention with local window attention, enabling efficient high-resolution image processing while maintaining strong global and local feature representation.

Contribution

The paper proposes a new linear attention method with a local concentration module, creating L²ViT, which captures global and local features efficiently without sacrificing performance.

Findings

L²ViT achieves 84.4% Top-1 accuracy on ImageNet-1K.

Pre-training on ImageNet-22k boosts accuracy to 87.0%.

L²ViT performs well on object detection and semantic segmentation.

Abstract

Vision Transformers (ViTs) have recently taken computer vision by storm. However, the softmax attention underlying ViTs comes with a quadratic complexity in time and memory, hindering the application of ViTs to high-resolution images. We revisit the attention design and propose a linear attention method to address the limitation, which doesn't sacrifice ViT's core advantage of capturing global representation like existing methods (e.g. local window attention of Swin). We further investigate the key difference between linear attention and softmax attention. Our empirical results suggest that linear attention lacks a fundamental property of concentrating the distribution of the attention matrix. Inspired by this observation, we introduce a local concentration module to enhance linear attention. By incorporating enhanced linear global attention and local window attention, we propose a new ViT architecture, dubbed L$^2$ViT. Notably, L$^2$ViT can effectively capture both global interactions and local representations while enjoying linear computational complexity. Extensive experiments demonstrate the strong performance of L$^2$ViT. On image classification, L$^2$ViT achieves 84.4% Top-1 accuracy on ImageNet-1K without any extra training data or label. By further pre-training on ImageNet-22k, it attains 87.0% when fine-tuned with resolution 384$^2$. For downstream tasks, L$^2$ViT delivers favorable performance as a backbone on object detection as well as semantic segmentation.