[Paper Review] How Do Vision Transformers Work?

Namuk Park, Songkuk Kim|arXiv (Cornell University)|Feb 14, 2022

Visual Attention and Saliency Detection202 citations

TL;DR

The paper analyzes how multi-head self-attention in Vision Transformers (ViTs) flattens loss landscapes and acts as a data-specific spatial smoothing, proposing AlterNet which blends CNN blocks and MSAs to improve performance in both large and small data regimes.

ABSTRACT

The success of multi-head self-attentions (MSAs) for computer vision is now indisputable. However, little is known about how MSAs work. We present fundamental explanations to help better understand the nature of MSAs. In particular, we demonstrate the following properties of MSAs and Vision Transformers (ViTs): (1) MSAs improve not only accuracy but also generalization by flattening the loss landscapes. Such improvement is primarily attributable to their data specificity, not long-range dependency. On the other hand, ViTs suffer from non-convex losses. Large datasets and loss landscape smoothing methods alleviate this problem; (2) MSAs and Convs exhibit opposite behaviors. For example, MSAs are low-pass filters, but Convs are high-pass filters. Therefore, MSAs and Convs are complementary; (3) Multi-stage neural networks behave like a series connection of small individual models. In addition, MSAs at the end of a stage play a key role in prediction. Based on these insights, we propose AlterNet, a model in which Conv blocks at the end of a stage are replaced with MSA blocks. AlterNet outperforms CNNs not only in large data regimes but also in small data regimes. The code is available at https://github.com/xxxnell/how-do-vits-work.

Motivation & Objective

• Investigate why MSAs improve generalization and how they influence loss landscapes.
• Compare MSAs with Convs in terms of frequency behavior and data specificity.
• Explain how MSAs and Convs can be harmonized in hybrid architectures.
• Propose design guidelines for architectures that blend CNN blocks and MSA blocks.
• Demonstrate that a staged alternating CNN/MSA pattern can outperform CNNs across data regimes.

Proposed method

• Analyze MSAs as trainable spatial smoothing of feature maps using self-attention equations.
• Use loss landscape visualizations and Hessian spectra to assess optimization properties.
• Perform Fourier analysis on feature maps to compare high-frequency content between MSAs and Convs.
• Experiment with locality-constrained MSAs and multi-stage architectures (PiT, Swin) versus CNNs.
• Propose AlterNet by replacing end-of-stage Convs with MSAs following build-up rules.
• Evaluate AlterNet on CIFAR-100 and ImageNet to compare with CNNs and ViTs.

Experimental results

Research questions

• RQ1 Do MSAs act as long-range dependencies or as data-specific spatial smoothers?
• RQ2 How do MSAs and Convs differ in their effect on loss landscapes and optimization dynamics?
• RQ3 Can a hybrid architecture that alternates CNN blocks and MSA blocks improve performance across data regimes?
• RQ4 What architectural rules maximize the benefits of MSAs in multi-stage networks?

Key findings

• MSAs flatten loss landscapes and improve generalization, mainly due to data specificity rather than long-range dependency.
• MSAs act as low-pass filters, while Convs act as high-pass filters; they are complementary.
• Multi-stage networks behave like a series of small models, with MSAs at stage ends playing a key predictive role.
• Local MSAs with appropriate receptive fields improve optimization by reducing negative Hessian eigenvalues.
• Convolving MSAs with end-stage placement and more heads in late stages yields better performance in small and large data regimes (AlterNet).
• AlterNet outperforms CNNs and ViTs in CIFAR-100 and scales to ImageNet, even in small data regimes.