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
standard Inception data augmentation as described in [41].
ImageNet classiﬁcation with MnasNet We follow the
training setup described in [42] and train all networks for
350 epochs with the RMSProp optimizer using exponential
learning rate decay. When training our augmented MnasNets, we divide the learning rate by 2 and adjusted the learning rate decay so that the ﬁnal learning rate stays the same.
Object Detection with COCO dataset We follow the
setup described in [26, 11] and train the RetinaNet from
scratch for 150 epochs without using ImageNet pretraining
for the ResNet backbone.
A.2. Computational & Memory costs
Table 9 provides the breakdown of self-attention related computational costs per image. Storing attention
maps in each layer induces a memory cost of Nh(HW)2bfloat16 . At inference, the memory cost for storing attention maps is only 1.2% of the memory required to store
model parameters (49MB).
Layer Memory Params FLOPS
fStage 2 - H=W=14g* 4 600KB 43k 22M
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