Several recent papers [1]}, [2]}, [3]}, [4]}, [5]} have attempted using self-attention primitives to improve image classification accuracy over the strong and commonly used ResNet backbones [6]}, [7]}.
Among them, the Stand-Alone Self-Attention (SASA) [2]} is a fully self-attentive model that replaces every spatial convolution with local self-attention, which improves the performance of ResNet backbones while having fewer parameters and floating point operations.
While conceptually promising, these models lag behind state-of-the-art convolutional models in image classification.
State-of-the-art convolutional models [9]}, [10]}, [11]} use a variety of scaling techniques to achieve strong performance across a range of computation and parameter regimes.
In this work, we aim to develop and understand techniques for scaling local self-attention models to outperform some of the best convolutional models.
Scaling self-attention models presents a unique set of challenges.
For example, convolutions have been very efficiently mapped to matrix accelerators such as TPUs and GPUs that drive most deep learning workloads, but fast implementations of local 2D self-attention do not currently exist.
To bridge this gap, we introduce a non-centered version of local attention that efficiently maps to existing hardware with haloing.
While our formulation breaks translational equivariance, it improves both throughput and accuracies over the centered local self-attention used in SASA.
We also introduce a strided self-attentive downsampling operation for multi-scale feature extraction.
