[Paper Review] Exploring Self-attention for Image Recognition

Hengshuang Zhao, Jiaya Jia+1 more|arXiv (Cornell University)|Apr 28, 2020

Domain Adaptation and Few-Shot Learning47 references33 citations

TL;DR

This paper compares pairwise and patchwise self-attention variants as alternatives to convolution for ImageNet-scale image recognition and shows vector and patchwise self-attention can match or exceed convolutional baselines with lower or comparable compute.

ABSTRACT

Recent work has shown that self-attention can serve as a basic building block for image recognition models. We explore variations of self-attention and assess their effectiveness for image recognition. We consider two forms of self-attention. One is pairwise self-attention, which generalizes standard dot-product attention and is fundamentally a set operator. The other is patchwise self-attention, which is strictly more powerful than convolution. Our pairwise self-attention networks match or outperform their convolutional counterparts, and the patchwise models substantially outperform the convolutional baselines. We also conduct experiments that probe the robustness of learned representations and conclude that self-attention networks may have significant benefits in terms of robustness and generalization.

Motivation & Objective

• Investigate two forms of self-attention (pairwise and patchwise) as building blocks for image recognition.
• Assess whether self-attention can match or surpass convolutional networks in accuracy, parameters, and FLOPs.
• Examine robustness and generalization properties of learned representations under self-attention.

Proposed method

• Formalize pairwise self-attention as a set operator with channel-adaptive weights and position encoding.
• Introduce patchwise self-attention that is strictly more powerful than convolution and computes location- and channel-aware weights.
• Construct self-attention blocks that interleave attention with a linear feature transformation via bottleneck layers.
• Build SAN architectures (SAN10, SAN15, SAN19) with residual-like blocks and multi-stage backbones.
• Compare against ResNet baselines on ImageNet, evaluating top-1/top-5 accuracy, parameters, and FLOPs.
• Perform ablations on relation functions, mapping functions, footprint size, and position encoding to understand design choices.

Experimental results

Research questions

• RQ1 Can pairwise self-attention achieve ImageNet-scale accuracy comparable to or better than convolutional baselines with similar resources?
• RQ2 Do patchwise self-attention variants substantially outperform convolutional networks and their parameter/FLOP budgets?
• RQ3 How do different relation functions, mapping architectures, footprint sizes, and position encodings impact performance and efficiency?
• RQ4 Do self-attention networks exhibit improved robustness and generalization relative to purely convolutional models?

Key findings

• Pairwise self-attention networks match or surpass convolutional counterparts with similar or lower parameter and FLOP budgets.
• Patchwise self-attention models substantially outperform convolutional baselines, with SAN15 achieving 78% top-1 accuracy versus 76% for ResNet50 while using fewer parameters and FLOPs.
• Vector (multi-channel) attention outperforms scalar attention in controlled experiments.
• Increasing the footprint size generally improves accuracy for self-attention until it saturates; patchwise attention shows strong gains independent of fixed kernel size.
• Relative position encoding significantly improves performance over no encoding; absolute encoding is better than none.
• SAN models with patchwise attention can outperform larger ResNets (e.g., SAN15 vs ResNet50) while being more parameter- and compute-efficient.
• Different transformation functions (phi, psi, beta) and multi-layer attention mappings improve performance, with three distinct transformations often yielding the best results.