# 2) repeat the above steps and
# 3) transpose to eventually put the logits
# in their right positions.
relembeddings h = tf.get variable(
’rheight’, shape=(2 H1, dkh),
initializer=tf.random normal initializer(dkh 0.5))
# [B, Nh, HW, HW]
rellogits h = relative logits 1d(
tf.transpose(q, [0, 1, 3, 2, 4]),
relembeddings h, W, H, Nh, [0, 1, 4, 2, 5, 3])
return rel logits h, rel logits w
def self attention 2d(inputs, dk, dv, Nh, relative=True):
"""2d relative self attention."""
, H, W, = shape list(inputs)
dkh = dk // Nh
dvh = dv // Nh
flatten hw = lambda x, d: tf.reshape(x, [ 1, Nh, H W, d])
# Compute q, k, v
kqv = tf.layers.conv2d(inputs, 2 dk + dv, 1)
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