PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention Augmented Convolutional Networks

Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, Quoc V. Le

arXiv:1904.09925v5cs.CV

TL;DR

Convolutions miss global image context, while self-attention offers long-range interactions but has been used mainly for sequence and generative tasks. The paper introduces translation-equivariant two-dimensional relative self-attention and augments convolutions with it, yielding systematic gains across image classification and object detection. Attention Augmentation improves ImageNet and COCO performance across architectures and computational settings, including a 1.3% top-1 gain over ResNet50 and a 1.4 mAP gain over RetinaNet.

  • Problem

    Convolutional locality limits access to global image context, while self-attention had mainly been applied to sequence and generative modeling rather than discriminative vision.

  • Method

    The paper augments convolutions by concatenating convolutional feature maps with maps from translation-equivariant two-dimensional relative self-attention.

  • Results

    Systematic improvements occur across image classification and object detection; ImageNet gains 1.3% top-1 accuracy over ResNet50 and COCO gains 1.4 mAP over RetinaNet.

  • Takeaways & Limitations

    Attention Augmentation is a competitive computational primitive for vision models across architectures and computational settings.

  • Takeaways & Limitations

    Open questions remain about fully attentional regimes, computational efficiency versus representational power, and whether convolution-oriented architectural choices remain optimal.

Abstract

from arXiv · show

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.

1. Introduction

Convolutions provide useful locality but miss global image context, motivating a translation-equivariant two-dimensional self-attention mechanism that augments convolutional operators. Across classification and detection experiments, this combination improves performance with limited additional computational burden.

  • Convolutional layers impose locality through limited receptive fields, which can prevent models from capturing global image context needed for recognition.
  • Attention-augmented convolution computes multi-head attention over image locations, reshapes the resulting weighted value averages, and concatenates them with standard convolution outputs.
  • The proposed two-dimensional relative self-attention preserves translation equivariance while modeling longer-range dependencies, and its feature maps are concatenated with convolutional maps.
  • 1.3% top-1 accuracy improvement on ImageNet is achieved over a ResNet50 baseline, while COCO object detection gains 1.4 mAP over a RetinaNet baseline.
  • Experiments across CIFAR-100, ImageNet, and COCO show systematic improvements across architectures and computational budgets, with minimal additional computational burden.

2. Related Work

Prior work established attention for long-distance sequence interactions and introduced image mechanisms that recalibrate convolutional features. Attention Augmentation instead jointly models spatial and feature subspaces using two-dimensional relative self-attention without relying on pretrained fully convolutional counterparts.

  • Attention became widely used in sequence modeling because it captures long-distance interactions, including in machine translation architectures.
  • Visual attention methods such as Squeeze-and-Excitation and Gather-Excite reweight channels, while BAM and CBAM separately refine channel and spatial features.
  • Attention Augmented networks use self-attention throughout the architecture rather than relying on pretraining fully convolutional counterparts.
  • The method jointly attends to spatial and feature subspaces and extends relative self-attention to two-dimensional inputs to model translation equivariance.

3. Methods

The method extends self-attention to two-dimensional images with relative positional embeddings, then concatenates attentional and convolutional feature maps. It addresses global context while preserving translation equivariance, with memory-aware implementation choices and empirical gains over convolutional baselines.

  • Self-attention over images: Multi-head attention flattens an input tensor of shape (H, W, Fin), computes attention, concatenates head outputs, projects them, and reshapes the result to (H, W, dv).The formulation has complexity O((HW)^2dk) and memory cost O((HW)^2Nh).
  • Two-dimensional positional embeddings: Two-dimensional relative self-attention adds learned relative height and width embeddings to attention logits, preserving translation equivariance while preventing permutation equivariance.The embeddings are shared across heads but not layers, adding (2(H + W) − 2)dh_k parameters per layer.
  • Two-dimensional positional embeddings: The implementation replaces explicit relative-embedding storage with a memory-efficient algorithm whose memory cost is O(HWdh_k).This addresses the additional memory burden of storing relative embeddings for every pixel pair.
  • Attention augmented convolution: The augmented design adds attention-derived feature maps rather than merely recalibrating convolutional features, unlike channelwise and channel-spatial attention mechanisms.The method is contrasted with Squeeze-and-Excitation, Gather-Excite, BAM, and CBAM.
  • Attention augmented convolution: Attention Augmentation concatenates convolutional and attentional feature maps, retaining translation equivariance and supporting inputs with different spatial dimensions.The attentional-channel ratio is υ = dv/Fout and the key-depth ratio is κ = dk/Fout.
  • Parameter effects: Attention-augmented networks significantly outperform fully convolutional counterparts while using fewer parameters in experiments.Replacing 3x3 convolutions slightly decreases parameters, whereas replacing 1x1 convolutions slightly increases them.
  • Memory-aware architecture: Because attention memory grows with spatial dimensions, augmentation begins at the smallest spatial layers and may use smaller batches or input downsampling.Downsampling uses 3x3 average pooling with stride 2, followed by bilinear interpolation for concatenation.

4. Experiments

Experiments evaluate Attention Augmentation across CIFAR-100, ImageNet, MnasNet, and COCO, showing improvements across architectures and computational budgets. Ablations examine parameter efficiency, stand-alone attention, and relative position encodings.

  • Experimental setup: Attention Augmentation produces systematic improvements on image classification and object detection across architectures and computational demands.Experiments cover CIFAR-100, ImageNet, and COCO using architectures including Wide-ResNet, ResNet, MnasNet, and RetinaNet.
  • ImageNet with ResNet: Attention Augmentation offers a competitive accuracy/computational trade-off against channel and spatial attention mechanisms despite lacking specialized kernels.Table 2 benchmarks it against BAM, CBAM, and GALA on ResNet50.
  • ImageNet with ResNet: Attention Augmentation significantly improves ImageNet performance over non-augmented ResNet baselines and outperforms Squeeze-and-Excitation while using parameters efficiently.AA-ResNet-50 performs comparably to baseline ResNet-101, while AA-ResNet-101 outperforms baseline ResNet-152.
  • ImageNet with MnasNet: Attention Augmentation improves MnasNet accuracy across all tested width multipliers, and Figure 3 verifies that gains are not explained solely by increased parameters.MnasNet uses Squeeze-and-Excitation at multiple architecture-search-selected locations, further supporting the observed benefits.
  • Object detection with COCO: 1.4% mAP improvement over a strong RetinaNet baseline is obtained on COCO, while Squeeze-and-Excitation significantly hurts detection performance.Relative self-attention improves RetinaNet with both ResNet-50 and ResNet-101 backbones; the authors hypothesize that localization benefits from preserved spatial information.
  • Ablation study: Fully attentional models remain viable for image classification, with AA-ResNet-50 using exclusively attentional channels only 2.5% below its fully convolutional counterpart.The fully attentional architecture also outperforms ResNet-34 while using fewer parameters and FLOPs.
  • Ablation study: Relative position encodings become increasingly important as the fraction of attentional channels rises, giving fully self-attentional ResNet-50 a 2.8% top-1 ImageNet gain.Additional ablations find that sinusoidal encodings and coordinate convolution do not improve over position-unaware Attention Augmentation, whereas two-dimensional relative attention does.

5. Discussion and future work

The paper proposes augmenting convolutional operators with self-attention and reports systematic gains across image classification and object detection. It identifies fully attentional models and architecture design as directions for future work.

  • Attention Augmentation combines convolutional operators with the proposed self-attention mechanism.
  • Systematic improvements appear on image classification and object detection across architectures and computational settings.
  • Future work will examine fully attentional models and how attention mechanisms trade off computational efficiency against representational power.

A.1. Experimental details

The experiments use task-specific training setups for CIFAR-100, ImageNet classification, and COCO object detection. Training choices include fixed schedules, baseline code settings, and architecture-specific optimizers.

  • Hyperparameters generally follow default settings from baseline codebases without tuning, while κ, υ, and attention heads receive constrained searches.
  • CIFAR-100: CIFAR-100 models train for 500 epochs with synchronous momentum SGD across eight Tesla V100 GPUs and task-specific batch sizes and weight decay.
  • ImageNet classification with ResNet: ImageNet ResNets train for 100 epochs with synchronous momentum SGD, weight decay of 1e-4, and batch sizes selected to fit each worker.
  • ImageNet classification with MnasNet: ImageNet MnasNets train for 350 epochs with RMSProp and exponential learning-rate decay, with augmented models using a halved learning rate.
  • Object Detection with COCO dataset: COCO RetinaNet models train from scratch for 150 epochs without ImageNet pretraining for the ResNet backbone.

A.2. Computational & Memory costs

The computational analysis covers attention-map memory, inference storage, and accuracy across FLOPS budgets. The figures compare attention-augmented ResNets and MnasNets across architectures and width multipliers.

  • Memory costs: Self-attention stores maps with memory cost O((HW)^2Nh) per layer.
  • Memory costs: At inference, attention-map storage requires 49MB, equal to 1.2% of the memory required for model parameters.
  • Accuracy versus computational demand: Accuracy is plotted across FLOPS counts, which correlate with running times across hardware platforms.
  • Accuracy versus computational demand: Figure 5 compares ImageNet top-1 accuracy and computational demand across ResNet-34, ResNet-50, ResNet-101, and ResNet-152.
  • Accuracy versus computational demand: Figure 6 compares MnasNet and Attention-Augmented-MnasNet across width multipliers 0.75, 1.0, 1.25, and 1.4.

A.3. 2D Relative Self-Attention implementation

The implementation uses basic tensor operations for 2D relative self-attention but requires non-trivial data-layout transformations. The paper attributes current latency partly to missing hardware kernels.

  • The method uses matrix multiplication, addition, and softmax, while implementation requires tiling, transposing, and reshaping.
  • Current latency reflects the lack of dedicated low-level kernels for the implementation's non-trivial operations, according to the paper.
  • Figures 7 and 8 provide TensorFlow helper functions and code for 2D relative self-attention.

A.4. Attention visualizations.

Attention maps for an augmented convolution show that the eight heads specialize to different content and can delineate object boundaries across four queried image locations.

  • Eight attention heads specialize to different content in the input image.
  • The visualizations show that attention can delineate object boundaries.
  • Rows represent the eight heads, while columns represent four pixel locations marked in the input image.
Loading 1904.09925v5…