[Paper Review] Stand-Alone Self-Attention in Vision Models

Prajit Ramachandran, Niki Parmar+4 more|arXiv (Cornell University)|Jun 13, 2019

Visual Attention and Saliency Detection63 references221 citations

TL;DR

The paper shows that self-attention can serve as a stand-alone primitive for vision models, replacing convolutions to create fully attentional networks that outperform convolution baselines on ImageNet with fewer parameters and FLOPS, and match RetinaNet on COCO with fewer resources.

ABSTRACT

Convolutions are a fundamental building block of modern computer vision systems. Recent approaches have argued for going beyond convolutions in order to capture long-range dependencies. These efforts focus on augmenting convolutional models with content-based interactions, such as self-attention and non-local means, to achieve gains on a number of vision tasks. The natural question that arises is whether attention can be a stand-alone primitive for vision models instead of serving as just an augmentation on top of convolutions. In developing and testing a pure self-attention vision model, we verify that self-attention can indeed be an effective stand-alone layer. A simple procedure of replacing all instances of spatial convolutions with a form of self-attention applied to ResNet model produces a fully self-attentional model that outperforms the baseline on ImageNet classification with 12% fewer FLOPS and 29% fewer parameters. On COCO object detection, a pure self-attention model matches the mAP of a baseline RetinaNet while having 39% fewer FLOPS and 34% fewer parameters. Detailed ablation studies demonstrate that self-attention is especially impactful when used in later layers. These results establish that stand-alone self-attention is an important addition to the vision practitioner's toolbox.

Motivation & Objective

• Demonstrate that content-based interactions via stand-alone self-attention can replace spatial convolutions in vision models.
• Build fully attentional vision architectures by substituting convolutions with local self-attention.
• Show that stand-alone attention can outperform convolutional baselines on ImageNet with fewer parameters and FLOPS.
• Characterize where in the network stand-alone attention yields the most gains through ablations.
• Identify practical design considerations for stem layers and relative positional encoding in vision transformers.

Proposed method

• Develop a local self-attention layer that attends over a memory block of size k×k for each pixel.
• Use multi-head self-attention with queries, keys, and values linear projections (W_Q, W_K, W_V).
• Incorporate 2D relative position embeddings to enable translation equivariance and improve expressivity.
• Replace all spatial convolutions in ResNet bottleneck blocks with the proposed attention layer, preserving downsampling structure.
• Experiment with a convolutional stem versus an attention stem, including a stem with spatially varying value transformations to bridge stem performance.
• Evaluate on ImageNet classification (ResNet variants) and COCO object detection (RetinaNet), plus ablations on spatial extent, positional encodings, and stem design.

Experimental results

Research questions

• RQ1 Can stand-alone local self-attention replace spatial convolutions to form a fully attentional vision model?
• RQ2 How does a fully attentional ResNet perform on ImageNet and COCO compared to convolutional baselines in terms of accuracy, FLOPS, and parameter count?
• RQ3 Where in the network is stand-alone attention most beneficial (stem vs later layers) and how do architectural choices affect performance?
• RQ4 What is the impact of spatial extent (k) and positional encoding type on performance?
• RQ5 Do stem modifications (spatially aware values) improve attention-based networks compared to a plain attention stem?

Key findings

• A fully attentional model by replacing convolutions with local self-attention achieves higher ImageNet Top-1 accuracy with 12% fewer FLOPS and 29% fewer parameters than the convolution baseline.
• On COCO detection, a fully attentional backbone matches RetinaNet mAP while using 39% fewer FLOPS and 34% fewer parameters.
• Ablations show larger gains when attention layers are placed in later network stages rather than early ones, suggesting convolutions capture low-level features better while attention integrates global information.
• Relative 2D positional encodings significantly improve performance over absolute or no positional encodings (relative best among tested options).
• Spatially aware value transformations in the stem outperform both a plain stand-alone attention stem and a stem using spatial convolutions for values.