Global Self-Attention Networks for Image Recognition

TL;DR

This work tackles the prohibitive cost of global self-attention in vision by introducing the Global Self-Attention (GSA) module, which pairs a global content-focused branch with an axial positional branch whose outputs are summed. The GSA module can replace standard convolutions to form standalone GSA networks that model long-range pixel interactions throughout a network, achieving superior accuracy with fewer parameters and FLOPs on ImageNet and CIFAR-100 compared to CNNs and several attention-based baselines. Key contributions include the non-axial global content attention, the axial positional attention with relative position embeddings, and extensive ablations that justify design choices and demonstrate where global attention provides the most benefit. Overall, GSA networks demonstrate the viability of global attention as a backbone operation for efficient, high-performance image recognition, with potential applicability to other CV tasks.

Abstract

Recently, a series of works in computer vision have shown promising results on various image and video understanding tasks using self-attention. However, due to the quadratic computational and memory complexities of self-attention, these works either apply attention only to low-resolution feature maps in later stages of a deep network or restrict the receptive field of attention in each layer to a small local region. To overcome these limitations, this work introduces a new global self-attention module, referred to as the GSA module, which is efficient enough to serve as the backbone component of a deep network. This module consists of two parallel layers: a content attention layer that attends to pixels based only on their content and a positional attention layer that attends to pixels based on their spatial locations. The output of this module is the sum of the outputs of the two layers. Based on the proposed GSA module, we introduce new standalone global attention-based deep networks that use GSA modules instead of convolutions to model pixel interactions. Due to the global extent of the proposed GSA module, a GSA network has the ability to model long-range pixel interactions throughout the network. Our experimental results show that GSA networks outperform the corresponding convolution-based networks significantly on the CIFAR-100 and ImageNet datasets while using less parameters and computations. The proposed GSA networks also outperform various existing attention-based networks on the ImageNet dataset.

Figures (3)

• Figure 1: Proposed GSA module: The keys, queries, and values (generated using $1\times 1$ convolutions) are processed by the content attention and positional attention layers in parallel. The positional attention layer is split into column-only and row-only positional attention layers, which use learned relative position embeddings ${\bm{R}}^c$ and ${\bm{R}}^r$ as keys. Finally, the outputs of the content and positional attention layers are summed to generate the output of the GSA module. Here, BN denotes batch normalization batchnorm, and PA stands for positional attention.
• Figure 2: Comparison between ResNet-{38, 50, 101} structure-based CNNs and GSA networks. GSA networks clearly outperform CNNs while using less parameters, computations, and runtime.
• Figure 3: Comparison between ResNet-{38, 50, 101} structure-based CNNs and GSA networks. GSA networks clearly outperform CNNs while using less parameters, computations, and runtime.