How Do Vision Transformers Work?

TL;DR

The paper challenges the view that Vision Transformers rely primarily on long-range dependencies by showing that multi-head self-attention acts as data-specific spatial smoothing that flattens loss landscapes. It demonstrates that MSAs function as low-pass filters and tend to aggregate feature maps, making them complementary to Convs (which are high-pass). Through extensive analysis of loss landscapes, Hessian spectra, and Fourier properties, the authors argue that stronger inductive biases and locality improve representations, particularly in small-data regimes, while large datasets mitigate non-convexities. They introduce AlterNet, an alternating CNN–MSA architecture with MSAs placed at the end of stages, showing improved performance on CIFAR and ImageNet and offering design rules for harmonizing MSAs with Convs. The work thus provides a unified perspective on how to combine CNNs and ViTs to exploit the strengths of both, along with practical guidance for architecture and training to achieve robust performance.

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.

Figures (26)

• Figure 1: Global and local aspects consistently show that MSAs flatten loss landscapes. Left: Loss landscape visualizations show that ViT has a flatter loss ($\text{NLL} + \ell_{2} \ \text{regularization}$) than ResNet. Right: The magnitude of the Hessian eigenvalues of ViT is smaller than that of ResNet during training phases. Since the Hessian represents local curvature, this also suggests that the loss landscapes of ViT is flatter than that of ResNet. To demonstrate this point, we present the Hessian max eigenvalue spectra at the end of the warmup phases and of 100$^\text{th}$, 200$^\text{th}$, and 300$^\text{th}$ epochs. See \ref{['fig:negative-hessian']} for a more detailed analysis.
• Figure 2: The Fourier analysis shows that MSAs do not act like Convs. Left: Relative log amplitudes of Fourier transformed feature map show that ViT tends to reduce high-frequency signals, while ResNet amplifies them. $\Delta$ Log amplitude of high-frequency signals is the difference between the log amplitude at normalized frequency $0.0 \pi$ (center) and at $1.0 \pi$ (boundary). For better visualization, we only provide the half-diagonal components of two-dimensional Fourier transformed feature maps. See \ref{['fig:log-amplitude']} for a more detailed analysis. Right: We measure the decrease in accuracy against frequency-based random noise. ResNet is vulnerable to high-frequency noise, while ViT is robust against them. We use a frequency window size of $0.1 \pi$ for frequency-based noise.
• Figure 3: Comparison of three different repeating patterns. Left: Spatial smoothings are located at the end of CNN stages. Middle: The stages of ViTs consist of repetitions of canonical Transformers. "D" is the hidden dimension and "H" is the number of heads. Right: The stages using alternating pattern consists of a number of CNN blocks and an MSA block. For more details, see \ref{['fig:alter-resnet']}.
• Figure 4: Hessian max eigenvalue spectra show that MSAs have their advantages and disadvantages. The dotted line is the spectrum of ViT using 6% dataset for training. Left: ViT has a number of negative Hessian eigenvalues, while ResNet only has a few. Right: The magnitude of ViT's positive Hessian eigenvalues is small. See also \ref{['fig:loss-landscape:hmes']} for more results.
• Figure 5: ViT does not overfit training datasets. "R" is ResNet and "RX" is ResNeXt. Left: Weak inductive bias disturbs NN optimization. The lower the $\text{NLL}_\text{train}$, the lower the error. Right: The lack of dataset also disturbs NN optimization.