[Paper Review] Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

Meng-Hao Guo, Zheng-Ning Liu+2 more|ArXiv.org|May 5, 2021

Advanced Neural Network Applications98 references79 citations

TL;DR

Proposes external attention with two external memories implemented as tiny learnable linear layers, achieving linear complexity and competitive results across vision tasks, including an all-MLP variant called EAMLP.

ABSTRACT

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.

Motivation & Objective

• Motivate and address the quadratic complexity and sample-invariant nature of self-attention in visual tasks.
• Introduce external attention with small shared memory units to capture dataset-level correlations.
• Show that external attention can replace self-attention in popular architectures with lower computation and memory costs.
• Demonstrate the versatility of external attention across image classification, detection, segmentation, generation, and 3D point cloud tasks.
• Propose multi-head external attention (EAMLP) to create an all-MLP architecture with competitive performance.

Proposed method

• Define external attention using two external memory units (M_k and M_v) as key and value memories.
• Compute attention as A = Norm(F M_k^T) and F_out = A M_v with linear layers implementing M_k and M_v.
• Use double normalization to stabilize attention scores across rows and columns.
• Extend to multi-head external attention for richer representations.
• Incorporate external attention into existing architectures and build an all-MLP model (EAMLP).

Experimental results

Research questions

• RQ1 Can external attention replace self-attention in vision architectures with linear computational cost?
• RQ2 Does incorporating dataset-level external memories improve generalization and performance across diverse vision tasks?
• RQ3 How does multi-head external attention (MEA) compare to self-attention and other attention variants in accuracy and efficiency?
• RQ4 Can external attention enable an all-MLP vision model that matches CNN/Transformer performance on ImageNet?
• RQ5 What is the impact of normalization strategies on external attention stability and performance?

Key findings

• External attention achieves comparable or superior results to self-attention across tasks with lower compute and memory usage.
• Using small shared memories (e.g., S ~ 64) yields linear complexity in input size (O(dSN)).
• Multi-head external attention enables an all-MLP architecture (EAMLP) with competitive ImageNet accuracy (up to 79.4% Top-1 in reported setup).
• Replacing self-attention with external attention improves segmentation and detection metrics in several benchmarks (e.g., VOC, COCO) when integrated into backbone networks.
• External attention provides interpretable attention maps showing focus on meaningful objects and regions, with heads attending to different regions.