model parameters (49MB).
Layer Memory Params FLOPS
fStage 2 - H=W=14g* 4 600KB 43k 22M
fStage 3 - H=W=14g* 6 600KB 90k 40M
fStage 4 - H=W=7g* 3 37.5KB 190k 19M
Training 6MB (total) 1.3M 390M
Inference 600KB (max) 1.3M 390M
Table 9. Computational costs associated with self-attention in the
forward pass of the ResNet50. During inference, we only consider
the largest memory cost since activations are not stored.
Figures 5 and 6 show the accuracies of our attention
augmented networks across FLOPS counts, which correlate
with running times across hardware platforms.
5 10 15 20 25
computational demand (GFlops).72737475767778798081top-1 accuracy.
73.6%76.4%77.9%78.4%
74.3%77.5%78.4%78.9%
74.7%77.7%78.7%79.1%
ResNet
SE-ResNet
AA-ResNet (ours)
Figure 5. ImageNet top-1 accuracy as a function of computational
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