GPUs. The learning rate is linearly scaled from 0 to
0:2B=256, whereBis the total batch size, for the ﬁrst 5%
training epochs and then annealed with cosine decay [30].
We use standard CIFAR preprocessing: mean normalizing,
random ﬂipping and cropping [55, 10, 48]. Non-augmented
architectures are trained with a batch size of 1024 and a
weight decay of 2e-4. Augmented architectures are trained
with batch size of 256 and a weight decay of 5e-4.
ImageNet classiﬁcation with ResNet We train all
ResNet architectures for 100 epochs using synchronous
SGD with momentum 0.9 across 8 TESLA V100 GPUs and
weight decay of 1e-4. We use the largest batch size per
workerB2f32;64;128;256gthat ﬁts in a minibatch. The
initial learning rate is scaled linearly according to the total
batch size using a base learning rate of 0.128 for total batch
size of 256. During training, we linearly scale the learning
rate from 0 to this value for the ﬁrst 5% of training epochs
and divide it by 10 at epochs 30, 60, 80 and 90. We use
Convolutional networks have been the paradigm of choice in many computer
vision applications. The convolution operation however has a significant
weakness in that it only operates on a local neighborhood, thus missing global
information. Self-attention, on the other hand, has emerged as a recent advance
to capture long range interactions, but has mostly been applied to sequence
modeling and generative modeling tasks. In this paper, we consider the use of
self-attention for discriminative visual tasks as an alternative to
convolutions. We introduce a novel two-dimensional relative self-attention
mechanism that proves competitive in replacing convolutions as a stand-alone
computational primitive for image classification. We find in control
experiments that the best results are obtained when combining both convolutions
and self-attention. We therefore propose to augment convolutional operators
with this self-attention mechanism by concatenating convolutional feature maps
with a set of feature maps produced via self-attention. Extensive experiments
show that Attention Augmentation leads to consistent improvements in image
classification on ImageNet and object detection on COCO across many different
models and scales, including ResNets and a state-of-the art mobile constrained
network, while keeping the number of parameters similar. In particular, our
method achieves a $1.3\%$ top-1 accuracy improvement on ImageNet classification
over a ResNet50 baseline and outperforms other attention mechanisms for images
such as Squeeze-and-Excitation. It also achieves an improvement of 1.4 mAP in
COCO Object Detection on top of a RetinaNet baseline.
Irwan Bello Barret Zoph Ashish Vaswani Jonathon Shlens Quoc V. Le