PaperlayerView exact revision on arXiv

Source-linked AI summary

How Do Vision Transformers Work?

Namuk Park, Songkuk Kim

arXiv:2202.06709v4cs.CVcs.LG

TL;DR

MSAs are widely used in vision, but their behavior and relationship to convolutions remain poorly understood. This paper analyzes their optimization and filtering properties and introduces AlterNet, which outperforms CNNs in both large- and small-data regimes.

  • Problem

    The paper addresses limited understanding of how MSAs work, including whether long-range dependencies matter, how MSAs differ from Convs, and how to combine them.

  • Method

    It analyzes MSA optimization and filtering behavior, examines multi-stage architectures, and proposes AlterNet, alternating CNN blocks with MSA blocks at stage ends.

  • Results

    AlterNet outperforms CNNs in both large- and small-data regimes, while analyses characterize MSAs as low-pass filters complementary to Conv high-pass filtering.

  • Takeaways & Limitations

    MSAs are generalized spatial smoothings whose data specificity, rather than long-range dependency, is central to their benefits and complementarity with Convs.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper explains MSAs as data-specific spatial smoothers that flatten loss landscapes, differ fundamentally from convolutional filters, and complement Convs in alternating architectures. These insights motivate placing MSA blocks at stage ends, yielding models that outperform CNNs in both large- and small-data regimes.

  • The paper addresses how MSA properties affect optimization, whether MSAs act like Convs, and how to harmonize MSAs with Convs.
  • MSAs improve accuracy and generalization by flattening loss landscapes, primarily through data specificity rather than long-range dependency.Their weak inductive bias can disrupt training, whereas data-specific weighting provides an appropriate inductive bias.
  • In small-data regimes, MSA losses can be non-convex because they allow negative Hessian eigenvalues, while large training datasets suppress these eigenvalues.
  • MSAs act as low-pass filters by reducing high-frequency signals, whereas Convs act as high-pass filters by amplifying them, making the two complementary.Convs are vulnerable to high-frequency noise, while MSAs are not.
  • Alternating CNN blocks with one or a few MSA blocks at each stage’s end improves accuracy and outperforms CNNs on both large and small datasets.The design uses stage-end spatial smoothing to ensemble transformed feature-map outputs, unlike canonical ViTs that perform poorly with little data.

2 WHAT PROPERTIES OF MSAS DO WE NEED TO IMPROVE OPTIMIZATION?

MSAs and ViTs benefit from data-specific and locality inductive biases that strengthen representations without causing overfitting, while also flattening or convexifying difficult loss landscapes. ViT optimization improves with large datasets and loss-smoothing methods, and data specificity matters more than long-range dependency.

  • Inductive biases: Stronger inductive biases lower both test error and training NLL, showing that ViT does not overfit training datasets.Appropriate MSA locality constraints help neural networks learn strong representations.
  • Loss geometry: ViT’s non-convex loss disrupts optimization, whereas ResNet’s loss is strongly (near-)convex.ViT has many negative Hessian eigenvalues, while ResNet has few.
  • Loss geometry: Large datasets suppress negative Hessian eigenvalues early in training, helping ViT convexify its loss and learn strong representations.ResNet gains little from large datasets because its loss is already convex on small datasets.
  • Loss smoothing: GAP suppresses negative Hessian max eigenvalues and improves ViT accuracy by +2.7 percent point.The smoothing effect of GAP convexifies the loss by ensembling feature-map points.
  • Data specificity: Data specificity improves neural networks, whereas long-range dependency hinders optimization; input-independent MLP-Mixer underperforms ViTs.Data specificity without self-attention also improves performance.
  • Locality: Local MSAs outperform global MSA because locality improves optimization by reducing unnecessary degrees of freedom and convexifying the loss landscape.On CIFAR-100, the 5 × 5 kernel outperforms the 8 × 8 kernel on both training and test datasets.

3 DO MSAS ACT LIKE CONVS?

MSAs and Convs exhibit complementary behaviors: MSAs are data-specific, low-pass, shape-biased, and aggregating, whereas Convs are data-agnostic, high-pass, texture-biased, and non-aggregating. These differences motivate placing MSAs at the ends of stages, especially later stages, to improve prediction.

  • Complementary behaviors: MSAs are data-specific and channel-agnostic, whereas Convs are data-agnostic and channel-specific, making them complementary.Convs mix channel information without exploiting data information; MSAs do the opposite.
  • Frequency behavior: MSAs generally reduce high-frequency feature-map amplitude, whereas Convs or MLPs increase it except in early stages.At high frequency 1.0π on ImageNet, MSAs decrease amplitude almost always, while MLPs increase it; early-stage MSAs are the exception.
  • Frequency behavior: ViT is vulnerable to low-frequency noise, while ResNet is vulnerable to high-frequency noise, indicating that MSAs are shape-biased and Convs are texture-biased.Low- and high-frequency signals correspond to image shape and texture, respectively.
  • Feature-map variance: MSAs reduce feature-map variance and ensemble predictions, whereas Convs and MLPs increase variance and do not ensemble.Reducing feature-map uncertainty can stabilize transformed features and aid optimization.
  • Stage-wise design: MSAs inserted at the ends of ResNet stages can improve predictive performance, with larger head counts especially beneficial in late stages.ResNet feature-map variance peaks at stage ends and tends to accumulate with depth.

4 HOW CAN WE HARMONIZE MSAS WITH CONVS?

MSAs and Convs are complementary, motivating AlterNet’s stage-aware alternating design that replaces selected Conv blocks with MSAs. The resulting models outperform CNNs in both large- and small-data regimes, while MSAs near stage ends are especially important for prediction.

  • Multi-stage architecture: Multi-stage CNNs and ViTs behave like series-connected small individual models, exhibiting repeating feature-map variance and block-structured representational similarities across stages.Feature-map similarities show block structures in both CNNs and multi-stage ViTs, unlike vanilla ViT.
  • Multi-stage architecture: Lesion studies indicate that removing layers at stage beginnings harms accuracy more than removing layers later, motivating MSAs closer to stage ends.The paper contrasts this stage-level effect with the belief that MSAs near the model’s end are most beneficial.
  • AlterNet design: AlterNet alternately replaces Conv blocks from a baseline CNN’s end, shifts replacements to earlier stage ends when needed, and uses larger late-stage MSA capacity.The model adjusts the ratio of MSAs and Convs according to these build-up rules.
  • AlterNet design: MSAs at stage ends improve prediction: MSAs in the last stage significantly improve accuracy, an MSA in the second stage helps, while multiple third-stage MSAs reduce accuracy.These findings come from Alter-ResNet-50 evaluated on CIFAR-100.
  • Performance: MSAs in AlterNet suppress large eigenvalues while allowing only a few negative eigenvalues, and larger datasets permit models to use more MSAs.The paper links this spectral behavior to the ability of large datasets to compensate for MSA shortcomings.
  • Performance: Alter-ResNet outperforms CNNs in the small-data regime and remains competitive on larger datasets, demonstrating that appropriately placed MSAs complement Convs.The comparison includes accuracy and corruption robustness on CIFAR-100 and CIFAR-100-C, with additional larger-dataset evidence from ImageNet.

5 DISCUSSION

MSAs act as generalized spatial smoothings rather than generalized Convs, complementing Convs by ensembling feature-map points and flattening the loss landscape. AlterNet preserves the Conv and MSA block architectures, leaving strong potential for future improvements.

  • 5 DISCUSSION: MSAs are generalized spatial smoothings, not merely generalized Convs, and complement Convs.They help neural networks learn strong representations by ensembling feature-map points and flattening the loss landscape.
  • 5 DISCUSSION: AlterNet preserves the architectures of Conv and MSA blocks, giving it strong potential for future improvements.

A EXPERIMENTAL DETAILS · A.1 SETUPS · A.2 BACKGROUND INFORMATION

The experiments use standardized hardware, PyTorch implementations, and controlled optimization settings across CIFAR and ImageNet, with ResNet-50, ViT-Ti, PiT-Ti, and Swin-Ti as primary models. The paper evaluates predictive performance, optimization landscapes, and feature-map frequencies using metrics and analyses defined for augmented training data or test data as appropriate.

  • A.1 SETUPS: The main experiments use separate hardware configurations for CIFAR and ImageNet, with models implemented in PyTorch.CIFAR uses either one GeForce RTX 2080 Ti or one NVIDIA T4; ImageNet uses four GeForce RTX 2080 Ti GPUs.
  • A.1 SETUPS: Networks are trained with categorical cross-entropy (NLL), AdamW, cosine annealing, 300 epochs, and batch sizes of 96 for CIFAR and 128 for ImageNet.The initial learning rate is 1.25 × 10−4, weight decay is 5 × 10−2, and learning-rate warm-up lasts 5 epochs.
  • A.1 SETUPS: The primary comparison includes ResNet-50, ViT-Ti, PiT-Ti, and Swin-Ti, whose CIFAR-100 throughputs are 320, 434, 364, and 469 image/sec, respectively.Additional results cover ResNeXt-50, Twins-S, MLP-Mixer-Ti, ConViT-Ti, and LeViT-128S.
  • A.1 SETUPS: Hessian analysis uses 10% of the training data and power iteration with batch size 16 to estimate the top-5 largest eigenvalues for ℓ2-regularized NLL on augmented data.The implementation is based on Yao et al. (2020), with the distinction between weight decay and ℓ2 regularization neglected.
  • A.2 BACKGROUND INFORMATION: Test errors are measured on clean test datasets, while training NLLs are measured on augmented training datasets because NLL is optimized during training and evaluates accuracy and uncertainty.The paper uses error as the predictive-performance metric on test data.
  • A.2 BACKGROUND INFORMATION: Lower training NLL indicates stronger representations, whereas compromised training NLL indicates regularization; overfitting occurs when test error worsens as training NLL improves.These definitions distinguish representation learning from regularization effects.
  • A.2 BACKGROUND INFORMATION: NEP measures loss non-convexity and APE measures sharpness, with lower values preferred for both metrics.Negative Hessian eigenvalues indicate non-convexity, while eigenvalue magnitude indicates sharpness.
  • A.2 BACKGROUND INFORMATION: Fourier analysis transforms feature maps into a normalized two-dimensional frequency domain to demonstrate that MSA acts as a low-pass filter.The lowest frequency is f = 0 and the highest-frequency components are at f = {−π, +π}; feature-map amplitudes and variances are averaged over channels.

B MSAS BEHAVE LIKE SPATIAL SMOOTHINGS

MSAs can be understood as spatial smoothing or ensemble averaging over proximate feature-map points. This perspective explains their low-pass behavior, optimization benefits, robustness, and stronger effect when placed near the output.

  • MSAs as spatial smoothing: Self-attention and ensemble averaging for proximate data points have identical formulations, with Softmax and V corresponding to importance weights and predictions.The weight samples in ensemble averaging correspond to MSA multi-heads.
  • Robustness and uncertainty: Spatial smoothing improves accuracy, predicts calibrated uncertainty, and is robust to MC dropout, data corruption, adversarial attacks, and especially high-frequency noise.MC dropout is described as equivalent to image occlusion.
  • Layer placement: Spatial-smoothing layers closer to the output layer significantly improve predictive performance.This observation motivates examining the placement of MSA blocks within stages.
  • MSAs as spatial smoothing: MSAs implement spatial smoothing by complementing each prediction with predictions from similar, proximate feature-map points.A 2 × 2 box blur spatially ensembles four neighboring points with weight 1⁄4.
  • Optimization and filtering: Spatial smoothing improves neural-network optimization by flattening loss landscapes, and even a small 2×2 box-blur filter significantly improves performance.The passage also characterizes spatial smoothing as a low-pass filter.

C VITS FROM A LOSS LANDSCAPE PERSPECTIVE

ViT optimization is shaped by non-convex, sometimes flat loss landscapes: small data and weak inductive biases hinder training, while large datasets, architectural smoothing, and more attention capacity improve optimization. Multi-stage designs and local MSAs flatten landscapes, but ViTs do not overfit small datasets merely by training longer.

  • Data regime: Lower training NLL consistently corresponds to lower test error on CIFAR-100, CIFAR-10, and ImageNet.The relationship is demonstrated directly on CIFAR-100 and reported similarly for CIFAR-10 and ImageNet.
  • Data regime: In small-data regimes, ViTs have higher training NLL and error than CNNs because non-convex losses disturb optimization.This occurs on CIFAR-10 and CIFAR-100, whereas large datasets convexify losses and let MSA benefits outweigh their drawbacks.
  • Architectural smoothing: NLLtrain 2.29 and error 24.1 % characterize PiT-Ti, whose accuracy is only 1.9 percent point lower than ResNet.Most improvement comes from PiT’s multi-stage architecture rather than convolutional inductive bias; small receptive fields also regularize ViT.
  • Training dynamics: Both NLLtrain and error decrease as training extends from 75 to 1200 epochs, yet ViT remains inferior to ResNet throughout and does not overfit CIFAR.The absence of overfitting persists even with large numbers of epochs.
  • Inductive bias: Weak inductive biases from small patches make loss landscapes flat but non-convex, while large patches suppress negative eigenvalues but sharpen landscapes and limit expression.Performance is best when these effects are balanced; small patches also reduce Hessian eigenvalue magnitude while producing negative eigenvalues.
  • Architectural smoothing: Multi-stage architectures and local MSAs flatten ViT loss landscapes, with PiT flatter than ViT near the optimum and all ViT surfaces smoother than ResNet.PiT’s multi-stage architecture suppresses negative Hessian eigenvalues, while Swin’s local MSA significantly reduces eigenvalue magnitudes.
  • Attention capacity: Increasing MSA heads convexifies and flattens loss landscapes, with APE decreasing according to an exponent of −0.562 for the number of heads.Higher embedding dimensions per head produce the same convexifying and flattening effect.
  • Model scale: Large models explore low NLLs while remaining globally sharp, because their Hessian eigenvalues are smaller than those of small models early in training.Thus, large models have locally flat loss functions during the early training phase.

D VITS FROM A FEATURE MAP PERSPECTIVE

Across multi-stage ViTs, MSAs generally ensemble feature maps and act as low-pass filters, while subsampling creates block-structured representations. Lesion studies further show that Convs at stage beginnings and MSAs at stage ends are especially important for accuracy.

  • Feature-map ensembling: MSAs in PiT and Swin reduce feature-map variance, indicating that they ensemble feature maps, except at Swin’s third-stage end.In Swin’s third stage, MSAs suppress variance growth at the beginning but not at the end.
  • Frequency filtering: MSAs in PiT and Swin are low-pass filters that generally decrease high-frequency amplitudes, whereas MLPs increase them.This matches the reported behavior of MSAs in ViTs and MLPs in ViT and Convs in ResNet.
  • Representation structure: Subsampling layers create a characteristic block structure in multi-stage ViT representations.This contrasts with the uniform representations reported across all layers of standard ViTs.
  • Stage-wise importance: Removing Convs at stage beginnings or MSAs at stage ends significantly harms accuracy across ResNets and ViTs.The lesion study removes one 3 × 3 Conv from ResNet bottleneck blocks and one MSA or MLP block from ViTs, producing periodic accuracy variation.

E EXTENDED INFORMATION OF ALTERNET · F DISTINCTIVE PROPERTIES OF DATA AUGMENTATION

The extended information specifies AlterNet’s architecture, stage-wise MSA allocation, computational efficiency, and ImageNet benefits. It also defines a comparison between strong and weak data augmentation to examine their distinct training effects.

  • E EXTENDED INFORMATION OF ALTERNET: AlterNet uses a pre-activation ResNet baseline, adds one CNN block to the last stage, and employs local MSAs with Swin-style relative positional encoding.Cyclic shifts and layer-specific initialization techniques are omitted for implementation simplicity.
  • E EXTENDED INFORMATION OF ALTERNET: AlterNet uses 3, 6, 12, and 24 MSA heads across its stages to strongly aggregate feature maps near the model’s end.The design follows the paper’s discussion that feature aggregation at the end of models improves performance.
  • E EXTENDED INFORMATION OF ALTERNET: MSAs in PiT and Swin reduce high-frequency signals, and multi-stage ResNet, PiT, and Swin models exhibit block structures in representational similarities.The figures identify Convs/MLPs, MSAs, and subsampling layers as distinct components in these patterns.
  • E EXTENDED INFORMATION OF ALTERNET: A lesion study finds that Convs at stage beginnings and MSAs at stage ends are important for prediction, with accuracy changes occurring periodically when units are removed.The experiment measures accuracy decrease after removing one unit from the trained model.
  • F DISTINCTIVE PROPERTIES OF DATA AUGMENTATION: Strong data augmentation makes predictions underconfident and significantly reduces Hessian maximum eigenvalues, indicating flatter loss landscapes but also producing a lot of additional effects.The supplied passage states that these observations occur on CIFAR-100 and that the underconfidence phenomenon also appears on ImageNet-1K.
  • E EXTENDED INFORMATION OF ALTERNET: The computational costs of Conv and MSA blocks are almost identical, while Alter-ResNet-50 reaches 473 image/sec on CIFAR-100, 20% faster than pre-activation ResNet-50.The optimal MSA count depends on the model and dataset; ImageNet uses 6 MSAs because large datasets alleviate MSA shortcomings.
  • E EXTENDED INFORMATION OF ALTERNET: On ImageNet-1K, appropriate MSA build-up rules significantly improve ResNet’s predictive performance, while robustness is measured as mean accuracy on ImageNet-C.The section attributes the improvement to MSAs complementing Convs and notes that some ViTs outperform CNNs on the large ImageNet dataset.
  • F DISTINCTIVE PROPERTIES OF DATA AUGMENTATION: The data-augmentation analysis compares DeiT-style strong augmentation with weak augmentation consisting of resize and crop.“Without data augmentation” denotes training with weak augmentation only.

F.1 DATA AUGMENTATION CAN HARM UNCERTAINTY CALIBRATION

On CIFAR-100, ResNet and ViT models with only weak augmentation are overconfident, whereas strong augmentation makes their predictions underconfident. This calibration shift contrasts with observations on large datasets such as ImageNet, where predictions without augmentation are not underconfident.

  • Calibration effects: On CIFAR-100, both ResNet and ViT models using only weak augmentation produce overconfident predictions.This finding is shown in a reliability diagram comparing models with and without strong augmentation.
  • Calibration effects: Strong data augmentation changes the predictive results of ResNet and ViT models from overconfident to underconfident on CIFAR-100.The passage characterizes this as an unexpected calibration effect.
  • Dataset contrast: Without data augmentation, predictions on large datasets such as ImageNet are not underconfident, unlike the strong-augmentation result on CIFAR-100.The passage notes that the detailed investigation of this difference remains incomplete.

F.2 DATA AUGMENTATION REDUCES THE MAGNITUDE OF HESSIAN EIGENVALUES

Strong data augmentation reduces the magnitude of Hessian eigenvalues, flattening the loss landscape early in training. These flatter losses improve generalization and help MSAs avoid overfitting while achieving better test accuracy.

  • Data augmentation and loss landscapes: Strong data augmentation reduces the magnitude of Hessian eigenvalues during the early phase of training.Figure F.1b compares Hessian max eigenvalue spectra with and without strong data augmentation.
  • Data augmentation and loss landscapes: By flattening the loss landscape, strong data augmentation produces flatter losses that lead to better generalization.
  • Data augmentation and generalization: Data augmentation helps an MSA avoid overfitting on the training dataset and achieve better accuracy on the test dataset.
Loading 2202.06709v4…