PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

Meng-Hao Guo, Zheng-Ning Liu, Tai-Jiang Mu, Shi-Min Hu

arXiv:2105.02358v2cs.CV

TL;DR

Self-attention captures long-range visual dependencies but incurs quadratic cost and overlooks correlations between samples. The paper introduces external attention with shared learnable memories and extends it to multi-head EAMLP. Across visual tasks, it reports comparable or better results than self-attention and variants with lower computational and memory costs.

  • Problem

    Self-attention has quadratic complexity and ignores potential correlations between different samples in visual feature representation.

  • Method

    External attention uses two small, learnable, shared memories to replace self-attention, with multi-head external attention forming the all-MLP EAMLP architecture.

  • Results

    Across image, point-cloud, and other visual tasks, external attention achieves comparable or better results than self-attention and variants with much lower computational and memory costs.

  • Takeaways & Limitations

    External attention provides a lightweight attention mechanism that can be incorporated into existing self-attention-based architectures and used for diverse visual tasks.

Abstract

from arXiv · show

Attention mechanisms, especially self-attention, have played an increasingly important role in deep feature representation for visual tasks. Self-attention updates the feature at each position by computing a weighted sum of features using pair-wise affinities across all positions to capture the long-range dependency within a single sample. However, self-attention has quadratic complexity and ignores potential correlation between different samples. This paper proposes a novel attention mechanism which we call external attention, based on two external, small, learnable, shared memories, which can be implemented easily by simply using two cascaded linear layers and two normalization layers; it conveniently replaces self-attention in existing popular architectures. External attention has linear complexity and implicitly considers the correlations between all data samples. We further incorporate the multi-head mechanism into external attention to provide an all-MLP architecture, external attention MLP (EAMLP), for image classification. Extensive experiments on image classification, object detection, semantic segmentation, instance segmentation, image generation, and point cloud analysis reveal that our method provides results comparable or superior to the self-attention mechanism and some of its variants, with much lower computational and memory costs.

1 INTRODUCTION

The paper proposes external attention as a lightweight alternative to self-attention, addressing quadratic cost and missing cross-sample correlations. It extends the mechanism with multiple heads to create the all-MLP EAMLP architecture and evaluates it across visual tasks.

  • Self-attention captures long-range dependencies but has quadratic complexity in the number of sample locations.
  • External attention uses learnable external key and value memories to replace self-attention while modeling relationships across the dataset.The memories are shared across samples and implemented with linear layers.
  • Across classification, detection, segmentation, generation, and point-cloud tasks, external attention achieves comparable or better results than self-attention and variants at lower computational effort.
  • O(n) complexity enables external attention to reduce computation while its shared memories provide a regularizing role across the dataset.
  • Multi-head external attention supports EAMLP, an all-MLP architecture for image classification.The paper reports a top1 accuracy of 79.4% on ImageNet-1K.

2 RELATED WORK

Related work frames self-attention as effective for capturing global context but costly at larger feature-map sizes. External attention instead uses a smaller learnable memory to model dataset-level context without semantic labels or iterative optimization.

  • Self-attention captures long-range dependencies by computing affinities between features, but its computation and memory overhead grow quadratically with feature-map size.
  • External attention computes relations between self queries and a much smaller learnable key memory that captures global dataset context.
  • Unlike approaches relying on semantic labels, external attention is optimized end-to-end by back-propagation rather than an iterative algorithm.
  • Vision research has developed attention mechanisms for classification, detection, segmentation, generation, and point-cloud learning.

3 METHODOLOGY

The methodology replaces quadratic self-attention with external attention, which computes affinities between input features and shared learnable memories. A multi-head extension uses shared memories across heads while retaining linear complexity in the number of input elements.

  • Self-Attention: Self-attention projects input features into query, key, and value matrices before forming pair-wise affinities across positions.Its attention matrix has N × N entries, where N is the number of input elements.
  • Self-Attention: O(dN^2) complexity makes direct pixel-level self-attention on images infeasible, motivating patch-based or more efficient alternatives.
  • External Attention: External attention computes affinities between input pixels and an external learnable memory M ∈ R^S×d, then updates features using similarities to that memory.The memory is independent of the input and acts as a learned prior over the training dataset.
  • External Attention: Two distinct memory units, M_k and M_v, serve as external key and value memories, increasing the capability of the external-attention module.
  • External Attention: O(dSN) complexity makes external attention linear in the number of input pixels, with small memory size S such as 64 working well experimentally.The computation load is described as roughly equivalent to a 1 × 1 convolution.
  • Multi-Head External Attention: Multi-head external attention applies separate heads to capture different relations while using shared memory units and an output projection for dimension consistency.The architecture can trade off the number of heads H against the number of memory elements S.

4 EXPERIMENTS

Experiments evaluate external attention across visual tasks and analyze its attention maps. The method provides competitive results across segmentation benchmarks and supports multi-head visual representations.

  • Experiments cover image classification, object detection, semantic segmentation, instance segmentation, image generation, and point cloud tasks.
  • External attention provides better accuracy than self-attention on the Pascal VOC segmentation ablation.
  • Attention maps focus on meaningful objects or background in segmentation, while different memory rows attend to different regions.
  • The learned attention representations support comparisons on PASCAL VOC, ADE20K, and Cityscapes segmentation benchmarks.

4.3 Image classification

Image classification experiments compare external attention, multi-head external attention, and EAMLP on ImageNet. Multi-head attention is necessary for the proposed architecture, while batch normalization has mixed effects.

  • ImageNet experiments replace Performer and multi-head self-attention blocks in T2T-ViT with external-attention variants.
  • External attention outperforms Performer but achieves about 2% point lower results than multi-head attention.
  • Multi-head mechanism is necessary for both self-attention and external attention.
  • Replacing LayerNorm with BatchNorm improves EAMLP-7 by 1% but causes failed cases in EAMLP-14 and EAMLP-19.

4.4 Object detection and instance segmentation

External attention is evaluated on COCO object detection and instance segmentation using a ResNet-50 backbone. Adding it at the end of stage 4 improves accuracy in both tasks.

  • The evaluation uses COCO, a benchmark for object detection and instance segmentation with 80 categories.
  • COCO experiments use MMDetection with a ResNet-50 backbone and add external attention at the end of ResNet stage 4.
  • About 1% improvement in accuracy is reported for both object detection and instance segmentation.

4.5 Semantic segmentation

EANet applies external attention to semantic segmentation on Pascal VOC, ADE20K, and Cityscapes. Results are comparable to state-of-the-art methods on Pascal VOC and Cityscapes and outperform others on ADE20K.

  • EANet is evaluated on Pascal VOC, ADE20K, and Cityscapes using the semantic segmentation architecture shown in Fig. 3.
  • Pascal VOC results are comparable to state-of-the-art methods on the test set.
  • ADE20K results outperform other methods on the validation set.
  • Cityscapes results are comparable to DANet on the validation set.

4.6 Image generation

The authors replace self-attention in both the generator and discriminator of SAGAN to create EAGAN. Experiments use default SAGAN hyperparameters and evaluate generated images with FID and IS.

  • EAGAN replaces SAGAN's self-attention mechanism in both the generator and discriminator.
  • The experiments use the default hyperparameter configuration from SAGAN.
  • Generated images are evaluated using Frechet Inception Distance and Inception Score.

4.7 Point cloud classification

On ModelNet40 point-cloud classification, the proposed EAT model replaces all self-attention modules in PCT. Under the reported experimental setup, it outperforms all compared methods and is presented as a backbone for 2D and 3D vision.

  • ModelNet40 contains 12,311 CAD models across 40 categories for 3D shape classification.
  • EAT replaces all self-attention modules in the PCT architecture.
  • EAT outperforms all compared methods, including the attention-based PCT model, on ModelNet40.
  • The authors present EAT as an outstanding backbone for both 2D and 3D vision.

4.8 Point cloud segmentation

On ShapeNet part segmentation, EAT achieves the best reported results. External attention also offers linear input-size complexity and substantially lower computational requirements than self-attention and its variants.

  • ShapeNet part contains 14,006 training models and 2,874 evaluation models across 16 object categories and 50 part labels.
  • EAT achieves the best reported results on the ShapeNet part dataset.
  • External attention has linear complexity with respect to input size.
  • Table 12 compares computational requirements using MACs, defined as multiply-accumulate operations.
  • External attention requires half as many parameters as self-attention and runs 32 times faster for the reported input size.Compared with the best variant, it remains about twice as fast.

5 CONCLUSIONS

The paper presents external attention as a lightweight mechanism for varied visual tasks. Its two external memory units function as dataset-wide dictionaries that learn representative features while reducing computational cost.

  • External attention is presented as a lightweight and effective mechanism for various visual tasks.
  • Its two external memory units can be viewed as dictionaries for the whole dataset.
  • The external memories learn more representative input features while reducing computational cost.
  • The authors hope external attention will inspire applications and research in domains such as NLP.
Loading 2105.02358v2…