Table of Contents
Fetching ...

You Only Need Less Attention at Each Stage in Vision Transformers

Shuoxi Zhang, Hanpeng Liu, Stephen Lin, Kun He

TL;DR

Vision Transformers incur a quadratic self-attention cost in token count, posing efficiency and scalability challenges. The paper introduces Less-Attention Vision Transformer (LaViT), which splits each stage into Vanilla-Attention (VA) layers and Less-Attention (LA) layers that re-parameterize and reuse previously computed attention via linear transforms, complemented by a Residual-based Attention Downsampling and a Diagonality Preserving loss to maintain token relationships. Key contributions include (i) a two-phase per-stage encoder, (ii) an attention residual downsampling mechanism for cross-stage context, (iii) a diagonality-preserving objective to stabilize transformed attention, and (iv) empirical validation across ImageNet-1K, COCO, and ADE20K showing favorable efficiency and strong performance. The approach offers a practical path to scalable ViTs, delivering competitive or superior results at reduced FLOPs and with simple, optimized matrix-multiplication operations suitable for real-world deployment.

Abstract

The advent of Vision Transformers (ViTs) marks a substantial paradigm shift in the realm of computer vision. ViTs capture the global information of images through self-attention modules, which perform dot product computations among patchified image tokens. While self-attention modules empower ViTs to capture long-range dependencies, the computational complexity grows quadratically with the number of tokens, which is a major hindrance to the practical application of ViTs. Moreover, the self-attention mechanism in deep ViTs is also susceptible to the attention saturation issue. Accordingly, we argue against the necessity of computing the attention scores in every layer, and we propose the Less-Attention Vision Transformer (LaViT), which computes only a few attention operations at each stage and calculates the subsequent feature alignments in other layers via attention transformations that leverage the previously calculated attention scores. This novel approach can mitigate two primary issues plaguing traditional self-attention modules: the heavy computational burden and attention saturation. Our proposed architecture offers superior efficiency and ease of implementation, merely requiring matrix multiplications that are highly optimized in contemporary deep learning frameworks. Moreover, our architecture demonstrates exceptional performance across various vision tasks including classification, detection and segmentation.

You Only Need Less Attention at Each Stage in Vision Transformers

TL;DR

Vision Transformers incur a quadratic self-attention cost in token count, posing efficiency and scalability challenges. The paper introduces Less-Attention Vision Transformer (LaViT), which splits each stage into Vanilla-Attention (VA) layers and Less-Attention (LA) layers that re-parameterize and reuse previously computed attention via linear transforms, complemented by a Residual-based Attention Downsampling and a Diagonality Preserving loss to maintain token relationships. Key contributions include (i) a two-phase per-stage encoder, (ii) an attention residual downsampling mechanism for cross-stage context, (iii) a diagonality-preserving objective to stabilize transformed attention, and (iv) empirical validation across ImageNet-1K, COCO, and ADE20K showing favorable efficiency and strong performance. The approach offers a practical path to scalable ViTs, delivering competitive or superior results at reduced FLOPs and with simple, optimized matrix-multiplication operations suitable for real-world deployment.

Abstract

The advent of Vision Transformers (ViTs) marks a substantial paradigm shift in the realm of computer vision. ViTs capture the global information of images through self-attention modules, which perform dot product computations among patchified image tokens. While self-attention modules empower ViTs to capture long-range dependencies, the computational complexity grows quadratically with the number of tokens, which is a major hindrance to the practical application of ViTs. Moreover, the self-attention mechanism in deep ViTs is also susceptible to the attention saturation issue. Accordingly, we argue against the necessity of computing the attention scores in every layer, and we propose the Less-Attention Vision Transformer (LaViT), which computes only a few attention operations at each stage and calculates the subsequent feature alignments in other layers via attention transformations that leverage the previously calculated attention scores. This novel approach can mitigate two primary issues plaguing traditional self-attention modules: the heavy computational burden and attention saturation. Our proposed architecture offers superior efficiency and ease of implementation, merely requiring matrix multiplications that are highly optimized in contemporary deep learning frameworks. Moreover, our architecture demonstrates exceptional performance across various vision tasks including classification, detection and segmentation.
Paper Structure (19 sections, 9 equations, 3 figures, 7 tables)

This paper contains 19 sections, 9 equations, 3 figures, 7 tables.

Figures (3)

  • Figure 1: The architecture of our Less-Attention Vision Transformer (LaViT). The bottom part: the proposed Less-Attention layer, which together with conventional Transformer blocks in the preceding layers constitutes the feature extraction module of this stage.
  • Figure 2: The downsampling on attention across stages.
  • Figure 3: The similarity ratio of the generated self-attention maps of the current layer with its previous layer.