Source-linked AI summary
Scaling Local Self-Attention for Parameter Efficient Visual Backbones
Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake Hechtman, Jonathon Shlens
TL;DR
The paper addresses whether self-attention can scale efficiently enough in vision to compete with high-performing convolutional models. It introduces efficient local attention and downsampling techniques, yielding HaloNets with strong ImageNet, transfer-learning, detection, and segmentation results, while retaining identified scope limitations.
Problem
The paper asks whether self-attention models can outperform not only canonical convolutional baselines such as ResNet but also high-performing convolutional models.
Method
The paper develops blocked or non-centered local attention with haloing and strided self-attentive downsampling, then builds multiscale HaloNet backbones.
Results
HaloNets achieve state-of-the-art performance across parameter regimes, including 84.9% top-1 accuracy on ImageNet, with promising transfer and harder-task results.
Takeaways & Limitations
Self-attention is competitive with leading convolutional models for ImageNet training from scratch, while hybrids match EfficientNets' speed-accuracy tradeoff and transfer models outperform BiT and ViT on accuracy and speed.
Takeaways & Limitations
Pure self-attention HaloNets are slower to train than corresponding EfficientNets and require further optimization for large-batch training; larger HaloNet transfer studies remain future work.
Abstract
from arXiv · showhide
Self-attention has the promise of improving computer vision systems due to parameter-independent scaling of receptive fields and content-dependent interactions, in contrast to parameter-dependent scaling and content-independent interactions of convolutions. Self-attention models have recently been shown to have encouraging improvements on accuracy-parameter trade-offs compared to baseline convolutional models such as ResNet-50. In this work, we aim to develop self-attention models that can outperform not just the canonical baseline models, but even the high-performing convolutional models. We propose two extensions to self-attention that, in conjunction with a more efficient implementation of self-attention, improve the speed, memory usage, and accuracy of these models. We leverage these improvements to develop a new self-attention model family, HaloNets, which reach state-of-the-art accuracies on the parameter-limited setting of the ImageNet classification benchmark. In preliminary transfer learning experiments, we find that HaloNet models outperform much larger models and have better inference performance. On harder tasks such as object detection and instance segmentation, our simple local self-attention and convolutional hybrids show improvements over very strong baselines. These results mark another step in demonstrating the efficacy of self-attention models on settings traditionally dominated by convolutional models.
1. Introduction
The paper develops efficient local self-attention models for vision to address the gap between promising self-attention results and stronger convolutional networks. It introduces hardware-efficient attention and downsampling techniques, then uses them to build HaloNets with strong ImageNet performance and transfer results.
- Self-attention offers content-based interactions and parameter-independent receptive-field scaling, motivating its use in vision alongside convolutional processing.It can also capture long-range dependencies and integrate multiple data types.
- Existing self-attention models improve ResNet backbones with fewer parameters and floating-point operations but lag behind state-of-the-art convolutional models.The paper therefore targets self-attention models that can outperform strong convolutional architectures across computation and parameter regimes.
- The paper introduces non-centered local attention with haloing to map efficiently to existing hardware, improving throughput and accuracy over centered local attention.The formulation relaxes translational equivariance.
- The paper also introduces strided self-attentive downsampling for multi-scale feature extraction.
- HaloNet achieves state-of-the-art performance across parameter regimes, including 84.9% top-1 accuracy on ImageNet for its largest model.The models also show promising transfer, object detection, and instance segmentation results.
2. Models and Methods
The paper develops local self-attention mechanisms that generate spatially varying filters while addressing the memory and computational costs of neighborhood extraction. HaloNet combines blocked local attention, strided attention downsampling, and implementation optimizations to improve efficiency and parameter-accuracy trade-offs.
- Self-attention can generate spatially varying convolutional filters: Self-attention produces spatially varying filters through shared linear projections whose weights are modulated by content-content and content-geometry interactions.Relative position embeddings capture geometry, while dot-product terms capture content relationships.
- Improving the speed-memory tradeoff by relaxing translational equivariance: Haloing computes queries from non-overlapping blocks and keys and values from expanded halo regions, preserving the query spatial dimensions after attention.For block size b = 2 and halo h = 1, a [4, 4, c] memory is extracted around each [2, 2, c] query block.
- Self-attention can generate spatially varying convolutional filters: Increasing the local attention window k adds relatively few parameters because projection parameters are independent of k, unlike convolutional parameters that scale quadratically with k.The relative-position parameters constitute only a small fraction of the layer’s parameters.
- Improving the speed-memory tradeoff by relaxing translational equivariance: Blocked local self-attention extracts one shared neighborhood for a block of pixels, trading off memory and computation between per-pixel local attention and global attention.Neighboring pixels share most of their neighborhoods, allowing the extraction to be reused across query blocks.
- Improving the speed-memory tradeoff by relaxing translational equivariance: Non-masked blocked attention increases receptive fields without additional computation but loses pixel-level translational equivariance because outputs depend on block position.The resulting operation is translationally equivariant only to shifts of size b.
- HaloNet: HaloNet uses the blocked local self-attention implementation with residual bottleneck structure and matches state-of-the-art convolutional models on the parameter-accuracy trade-off curve.The family incorporates minor modifications to the ResNet structure, including a final 1 × 1 convolution.
3. Related Work
Attention has expanded from sequence modeling and image-captioning applications toward vision backbones, while transformer-based vision models achieve strong accuracy under specific pretraining conditions. Research also continues to improve attention and convolution through algorithmic, implementation, and structural changes.
- Attention originated in sequence modeling and was later used to attend to image features in image-captioning systems.
- Non-local means introduced pairwise-weighted global pixel aggregation for image denoising and was later applied to downsampled convolutional features.
- HaloNets are compared with EfficientNets on accuracy versus parameter count, with EfficientNet B5 and B7 accuracies obtained using RandAugment.
- Transformers applied to non-overlapping image patches achieve accuracies comparable to state-of-the-art models when pretrained on JFT-300M or ImageNet-21k.These models do not use a multiscale architecture, whereas this work focuses on training from scratch on ImageNet.
- Computational primitives improve through algorithmic changes, better software implementations, and structural designs such as grouped and depthwise separable convolution.
4. Experiments
Experiments show that HaloNets match or exceed strong convolutional baselines across classification, transfer, and structured prediction, while architectural choices expose accuracy–efficiency trade-offs. Regularization, larger inputs and windows, and selective convolution–attention hybrids substantially affect results.
- 4.1. HaloNets are competitive with state-of-the-art convolutional models: 84.9% top-1 ImageNet validation accuracy is achieved by HaloNet H7, which performs at par or slightly better than EfficientNet models with the same parameters.HaloNets outperform other model families, with training conducted from scratch on ImageNet.
- 4.2.1. Transfer of convolutional components to self-attention: 1.3% versus 0.8%: label smoothing and RandAugment improve HaloNet accuracy more than ResNet accuracy despite roughly 30% fewer parameters.The comparison suggests self-attention models may require regularizations typical of larger convolutional models.
- 4.2.1. Transfer of convolutional components to self-attention: 0.4% versus 1.3%: Squeeze-and-Excitation and SiLU improve HaloNet less than ResNet, while HaloNet benefits from label smoothing, SiLU, RandAugment, and stochastic depth in larger models.The authors speculate that self-attention’s gating and multiplicative interactions reduce the need for explicit gating such as SE.
- 4.2. Model study 1: comparing self-attention and convolutions: Larger images consistently improve HaloNet accuracy, while the accuracy gap between HaloNets and ResNets is maintained.Convolutional models also improve with larger images, but the relative gap remains.
- 4.3. Model study 2: HaloNet architecture study: 1.3% accuracy gain results when the window size doubles from 6 × 6 to 12 × 12, without increasing the number of parameters.For a fixed window size, query block size does not impact results, enabling larger query blocks.
- 4.4. Convolution-Attention hybrids improve the speed-accuracy tradeoff: 1% top-1 accuracy is lost when convolutions replace attention in all but the final stage, whereas stages 1–2 convolutions with stages 3–4 attention minimally hurt accuracy and significantly improve step times.The fastest model uses convolutions in all stages except the last; detailed improved hybrid studies are left for future work.
5. Discussion
The work develops multiscale self-attention models using blocked local attention and attention downsampling, achieving competitiveness with strong convolutional models while retaining important speed limitations.
- Blocked local attention and attention downsampling produced multiscale self-attention models competitive with the best convolutional models.The paper also used ablations to study how self-attention scaling could be improved.
- Self-attention was competitive in accuracy when trained from scratch on ImageNet.
- Pure self-attention HaloNets were slower to train than corresponding EfficientNets and need further large-batch optimizations.
- Hybrid models matched EfficientNets in speed-accuracy tradeoffs, while transfer experiments showed higher accuracy and speed than BiT and ViT.
- Relative embeddings add few parameters because their size is small compared with the three linear projections.
- Increasing value-channel computation and layer count had the greatest accuracy impact in the reported scaling study.
B. Study of enlarging self-attention models
The study enlarges HaloNet-50 under a fixed 30-million-parameter budget to identify which scaling dimensions most improve accuracy.
- Increasing attention computation through rv and adding more layers were the most fruitful dimensions for improving accuracy.The comparison scaled different hyperparameters of HaloNet-50 to reach 30 million parameters.
C. Experimental details, hyperparameters
This section records experimental details and model configurations omitted from the main paper for space reasons.
- The appendix lists experimental details and model configurations omitted from the main text for space.
C.1. Experimental details for model studies
The model studies use controlled training settings and define the HaloNet family through image, attention, bottleneck, depth, and final-width hyperparameters.
- HaloNet-50 and ResNet-50 used the same layer allocations and channel widths, with 150-epoch training on 256 × 256 images.The learning rate was 0.1; weight decay differed according to RandAugment usage.
- HaloNet configurations H1–H7 vary image size, query block size, halo size, attention and bottleneck widths, depth, and final 1 × 1 convolution width.Each HaloNet used an image size comparable to its corresponding EfficientNet.
C.3. Classification hyperparameters
Classification training uses weight decay, cosine-annealed learning rates, and model-scaled RandAugment magnitudes. The largest models receive a nonzero final learning rate to address end-of-training overfitting.
- Classification uses weight decay of 2e−5 and cosine annealing with an initial learning rate of 0.1.
- The largest models consistently overfit at the end of training, attributed to the learning rate reaching 0.
- The cosine annealing schedule ends at 1.0 128 of the original learning rate instead of 0.
- RandAugment magnitudes increase across HaloNet sizes H0 through H7 from 6 to 31, without extensive tuning.
C.4. Detection and instance segmentation hyperparameters
Detection and instance segmentation use Mask-RCNN with ImageNet-pretrained backbones and a TPU detection framework. Pretraining and downstream training use specified image sizes, regularization, initialization, and optimization schedules.
- All detection and instance segmentation experiments use Mask-RCNN with backbones pretrained on ImageNet.
- Backbones are pretrained for 350 epochs at image size 512, with RandAugment magnitude 15 and stochastic depth probability 0.1.
- Detection training initializes the backbone from pretrained weights and other parameters from scratch.
- The detection code and hyperparameters directly use an open-source TPU detection and segmentation framework.
- The model trains for 67500 steps with SGD learning rate 0.1, momentum 0.9, and learning-rate decays at 60000 and 65000 steps.
D. Optimizations
The implementation avoids repeated data-format conversions by retaining blocked representations through the network and performing halo gathering with convolutions. These choices target speed improvements.
- The implementation avoids data-formatting operations whenever possible because they can slow down the model.
- Persistent blocking: Persistent blocking keeps flattened (b, b) blocks as sequences of length b^2 until the network end, processing images in 5D.
- Persistent blocking: Batch normalization and other operations are implemented directly on the blocked format instead of reshaping back to 4D.
- Gathers with convolutions: Haloing is performed in 5D with flattened neighborhoods, using 3D convolutions as gathering operations instead of slices and concatenations.
E. ImageNet-21k Models
The ImageNet-21k transfer setup enlarges HaloNet H4 and adapts block and halo sizes to each image resolution. Relative position encodings are important for accuracy, while absolute factorized encodings reduce it.
- The ImageNet-21k transfer setup modifies HaloNet H4 by increasing base width, increasing rb, removing the final 1×1 convolution, and adjusting label embeddings.
- Pretraining uses b = 8 and h = 2 at 256 × 256 images, while finetuning uses b = 12, h = 2 at 384 × 384 and b = 16, h = 1 at 512 × 512.
- Transfer initializes parameters from the final pretraining checkpoint, except for zero-initialized label embeddings and linearly interpolated relative embeddings.
- Absolute factorized position encodings reduce accuracy from 78.6% to 77.5%, supporting the importance of relative position encoding.