Source-linked AI summary
Stand-Alone Self-Attention in Vision Models
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, Jonathon Shlens
TL;DR
The paper asks whether attention can replace convolution as a primary vision primitive for modeling content-based interactions. It develops a stand-alone local self-attention layer, uses it to build fully attentional models, and evaluates them on classification and detection. These models achieve competitive or improved performance with fewer parameters and floating-point operations, while attention is especially effective in later network layers.
Problem
Convolutions struggle with long-range interactions, while prior vision uses of attention mainly augment convolutional models rather than replacing them.
Method
The paper replaces ResNet spatial convolutions with a local self-attention layer while preserving the remaining architecture and downsampling schedule.
Results
Fully attentional models outperform convolutional baselines on ImageNet classification and match RetinaNet mAP on COCO detection with fewer parameters and FLOPS.
Takeaways & Limitations
Stand-alone self-attention is a viable vision primitive, with ablations showing particular effectiveness in later network parts.
Takeaways & Limitations
Although computational demand is favorable to convolution, the attention-based network is slower in wall-clock time because optimized accelerator kernels are lacking.
Abstract
from arXiv · showhide
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.
1 Introduction
Convolutional neural networks became central to vision because they offer parameter-efficient image representations and translation equivariance, but large receptive fields make long-range interactions difficult. Attention has improved convolutional models, motivating the question of whether content-based interactions can instead serve as the primary vision primitive.
- CNNs became a backbone for vision applications as datasets and compute resources expanded, spanning recognition, detection, and segmentation.
- Convolutions provide translation equivariance, but capturing long-range interactions is challenging because large receptive fields scale poorly.
- Attention has been used to augment CNNs through channel-based and spatially aware mechanisms for image classification and object detection.
- Global attention over all spatial locations is limited to small inputs that typically require substantial downsampling.
- The paper asks whether content-based interactions can serve as the primary primitive of vision models rather than merely augmenting convolution.
2 Background
Convolutions aggregate local image neighborhoods with shared weights, while self-attention enables content-based interactions and direct long-distance modeling. The paper develops local self-attention with relative position information as a convolution-replacing layer for fully attentional vision models.
- 2.1 Convolutions: CNNs extract k × k neighborhoods around each pixel and apply learned weights shared across spatial positions.Weight sharing supports translation equivariance and keeps convolution parameters independent of input size.
- 2.2 Self-Attention: Self-attention models content-based interactions and can directly represent long-distance dependencies while remaining parallelizable.Unlike convolutional aggregation, self-attention uses query, key, and value representations from a shared context.
- 2.2 Self-Attention: Local self-attention replaces global all-to-all interaction with a neighborhood memory block, making attention usable across both small and large inputs.Global attention is computationally expensive before substantial spatial downsampling, whereas the proposed layer operates locally around each pixel.
- 2.2 Self-Attention: Local attention aggregates neighborhood values using content-dependent softmax weights derived from queries, keys, and values.Multiple heads partition pixel features into groups and concatenate their separately computed representations.
- 2.2 Self-Attention: Relative position embeddings modulate attention logits by spatial distance and preserve translation equivariance.The relative distance is factorized into row and column offsets whose embeddings are concatenated.
- 2.2 Self-Attention: Attention parameters are independent of spatial extent, and with d_in = d_out = 128, k = 3 convolution costs as much as k = 19 attention.Convolution parameter count grows quadratically with spatial extent, while attention computational cost grows more slowly for typical channel dimensions.
3 Fully Attentional Vision Models
The fully attentional architecture replaces spatial convolutions in a ResNet while preserving the surrounding bottleneck structure and downsampling schedule. The stem remains convolutional because early RGB inputs are poorly suited to self-attention, with spatially aware value features used to connect convolutional locality and attention.
- Replacing Spatial Convolutions: A spatial convolution is defined as a convolution with spatial extent k > 1, excluding 1 × 1 convolutions.The excluded 1 × 1 convolutions can be viewed as independently applied fully connected layers.
- Replacing Spatial Convolutions: The construction replaces each 3 × 3 spatial convolution in ResNet bottleneck blocks with self-attention while preserving other layers and downsampling locations.The transformation keeps the 1 × 1 down-projection, 1 × 1 up-projection, residual connections, layer count, and downsampling schedule unchanged.
- Replacing the Convolutional Stem: The stem uses convolutional processing because self-attention underperforms there on spatially correlated, individually uninformative RGB pixels.The stem is the initial CNN stage responsible for learning local features such as edges before later layers identify global objects.
- Replacing the Convolutional Stem: The stem combines an attention layer with spatially aware value features and max pooling, aligning the attention receptive field with the pooling window.Position-dependent factors combine multiple value matrices through spatially varying linear transformations, injecting distance-based information into the value transformation.
4 Experiments
Experiments show that fully attentional models match or exceed convolutional baselines on ImageNet classification and COCO detection with lower computational and parameter costs. Ablations indicate that attention is most effective in later layers, while positional encoding, spatial extent, and a spatially-aware stem materially affect performance.
- ImageNet Classification: Attention models outperform the ImageNet baseline across network depths while retaining 12% fewer FLOPS and 29% fewer parameters.Full Attention applies attention in the stem and throughout the network; Conv-stem + Attention retains spatial convolution in the stem.
- ImageNet Classification: 0.5% higher classification accuracy, 12% fewer FLOPS, and 29% fewer parameters are achieved by full attention versus the ResNet-50 baseline.The performance gain is consistent across most depth- and width-scaled model variations.
- COCO Object Detection: 34% fewer parameters and 39% fewer FLOPS are obtained while matching the baseline RetinaNet mAP when attention replaces the backbone, FPN, and detection heads.An attention-based backbone alone matches mAP with 22% fewer parameters.
- Where is stand-alone attention most useful?: The best-performing layer arrangements use convolutions in early groups and stand-alone attention in later groups, with costs similar to the fully attentional model.Reversing this arrangement degrades performance despite a large increase in parameter count, suggesting complementary low-level and global-information roles.
- Ablations: Performance improves with larger attention spatial extent until plateauing around k = 11, whereas small settings such as k = 3 substantially hurt performance.The plateau depends on settings such as feature size and number of attention heads; the parameter count remains constant across the spatial-extent ablation.
- Ablations: Relative positional encodings perform 2% better than absolute encodings, while removing content-content interactions reduces accuracy by only 0.5%.Any positional encoding improves over using none, and content-relative interactions remain especially important in the ablation.
5 Discussion
The paper finds that fully attentional networks achieve competitive performance on ImageNet and COCO while using fewer parameters and floating-point operations than convolutional baselines. Attention is especially effective in later network layers, while wall-clock speed and architecture design remain areas for improvement.
- Fully attentional networks achieve competitive predictive performance on ImageNet classification and COCO object detection with fewer parameters and floating-point operations than convolutional baselines.
- Attention is especially effective in the later parts of the network.
- Future work includes improving geometry capture, searching architectures designed for attention, and developing attention forms that capture low-level features earlier in networks.
- The resulting attention-based network is slower in wall-clock time despite favorable training efficiency and computational demand.The paper attributes this discrepancy to a lack of optimized kernels on various hardware accelerators.
- The authors also propose exploring unified convolution and self-attention models and applying attention to additional vision tasks.
A.1 Attention Stem
The attention stem defines a local self-attention layer that produces per-pixel queries, keys, values, and outputs, then replaces pointwise values with spatially aware transformations. Its mixture weights depend on absolute positions within a 4 × 4 pooling window and are shared across four heads.
- The standard single-headed self-attention layer maps each input xij to per-pixel queries, keys, values, and a final output.The projections use WQ, WK, and WV, while attention operates over the neighborhood Nk(i, j).
- The attention stem replaces pointwise values with spatially aware linear transformations.
- A 4 × 4 pooling-window receptive field aligns the query, key, and value regions to inject distance-aware value features.
- The value transformation is a convex combination of multiple value matrices whose weights depend on absolute positions within the pooling window.Equation 9 computes logits between absolute position embeddings and mixture embeddings.
- The position-dependent mixture weights are shared across the four attention heads in the mixture stem layer.
A.2 ImageNet Training Details
ImageNet training uses a 4% validation subset and 130 epochs on 128 Cloud TPUv3 cores, with warmup followed by cosine learning-rate decay.
- A validation set containing a 4% random subset of the training set is used for tuning.
- Training runs for 130 epochs with Nesterov’s Accelerated Gradient, a learning rate of 1.6, 10-epoch linear warmup, and cosine decay.
- A total batch size of 4096 is distributed across 128 Cloud TPUv3 cores.
- The setup uses batch normalization with decay 0.9999 and exponential moving average with weight 0.9999 over trainable parameters.
A.3 Object Detection Training Details
The object-detection models use fully attentional classification backbones and replace RetinaNet’s spatial convolutions with self-attention, while training uses 150 epochs, batch size 64, cosine decay, and image augmentation.
- The fully attentional object-detection architecture uses the fully attentional classification models as its backbone.
- RetinaNet’s 3 × 3 convolutions are replaced with same-width self-attention layers using dout = 256.
- Strided convolutions are replaced with 2 × 2 average pooling using stride 2.
- Training applies multiscale jitter, crops images to a maximum dimension of 640, and randomly flips images horizontally with 50% probability.