How Do Vision Transformers Work?

In plain words

Vision Transformers cut an image into patches and let the attention step relate every patch to every other patch. This paper asks what that step actually does. Its answer is that attention works like a smoothing operation: each spot in the internal picture is replaced by a weighted average of other spots, with the weights learned from the data.

The authors report that this smoothing helps training. It flattens the landscape the optimiser has to descend. Standard convolutions do the opposite and emphasise fine detail. The two kinds of layer therefore seem to complement each other. The weak results of plain Vision Transformers on small datasets are traced to a bumpy training landscape rather than to the model simply memorising its training images.

Based on this view, the paper builds a hybrid. Most of a convolutional network stays as it is, but some of its blocks near the end of each stage are replaced with attention blocks that look only at small local windows. The authors report that this hybrid outperforms plain convolutional networks on CIFAR-100.

TL;DR

Multi-head self-attention (MSA) in vision works as a learned spatial smoothing: each output point is a weighted average of feature-map points with data-dependent importance weights §1.1 app. B. The authors report that MSAs flatten the loss landscape and act as low-pass filters that reduce feature-map variance, while convolutions do the opposite §2 §3 fig. 1 fig. 2 fig. 9. Plain ViTs fit small datasets poorly because their loss has many negative Hessian eigenvalues, not because they overfit §2 fig. 4 fig. 5. AlterNet, a pre-activation ResNet with some stage-end Conv blocks replaced by local MSA blocks, is reported to beat CNNs on CIFAR-100 and to be on par with Swin in ImageNet-1K accuracy §4.2 app. E fig. 12. All accuracy numbers are author-reported and appear only in figures.

Problem

The usual explanation for why MSAs work is their weak inductive bias and their ability to model long-range dependencies §1. The paper argues that this explanation is poorly supported §1. Vision Transformers are known to do poorly on small datasets such as CIFAR, which the authors attribute to optimisation rather than overfitting §2. The open questions are: which properties of MSAs help optimisation, whether MSAs behave like convolutions, and how to combine the two §1.

Method

The paper is mostly an analysis, and its model (AlterNet) follows from the analysis. Each strand below is one of the paper's claims, followed by the evidence the paper gives for it.

Interpretation: MSA as spatial smoothing. Self-attention computes, for each output point, a softmax-normalised weighted sum of value vectors over all points eq. 1 §1.1. The authors read this as an ensemble average over proximate points with positive importance weights, and show it has the same form as a Bayesian ensemble of predictions for neighbouring points app. B eq. 2. Under this reading a fixed 2×2 box blur is a special case with equal weights app. B. The paper states that spatial smoothing has the same three properties it attributes to MSAs: it helps optimisation, it acts as a low-pass filter, and it helps most when placed at the end of a stage app. B.

Claim 1: MSAs flatten the loss landscape, and data specificity matters more than range. Hessian max-eigenvalue spectra show that ViT has smaller positive eigenvalues than ResNet, which the authors read as a flatter loss §2 fig. 1 fig. 4. ViT has many negative eigenvalues in the early phase of training, which they read as non-convexity, and large datasets suppress them §2 fig. 4. Global average pooling, used instead of the CLS token, also suppresses the negative eigenvalues and improved CIFAR-100 accuracy by 2.7 percentage points §2 fig. 6. Local MSAs with a 5×5 window outperform global MSA on CIFAR-100 and have fewer negative eigenvalues than both global MSA and a 3×3 window §2 fig. 7. The authors read this as the restriction of unneeded degrees of freedom, and they conclude that data specificity, not long-range dependency, is the key feature §2. Convolutional ViT with 3×3, 5×5 and global windows is the test bed for this claim §2 fig. 7.

Claim 2: MSAs are low-pass filters and convolutions are high-pass filters. The authors compute the Fourier amplitude of feature maps. At the highest normalised frequency, the log-amplitude change is negative for MSA blocks and positive for MLP blocks in most stages; MSAs in the early stages are the exception §3 fig. 2 fig. 8. ViT is vulnerable to low-frequency random noise and ResNet to high-frequency noise, which the authors relate to shape-biased versus texture-biased behaviour §3 fig. 2.

Claim 3: MSAs reduce feature-map variance, convolutions increase it. Measured feature-map variance falls across MSA blocks in ViT and rises across convolutions in ResNet and MLP blocks in ViT §3 fig. 9. In ResNet the variance peaks at the end of each stage §3 fig. 9. The authors therefore predict that MSAs placed at stage ends should help ResNet §3.

Claim 4: multi-stage networks behave like series of small models. Representational similarity (mini-batch CKA) shows block structure for multi-stage ViTs such as PiT and Swin, but not for vanilla ViT §4.1 fig. 10. A lesion study (removing one block from a trained network) shows that removing a block at the start of a stage hurts ResNet more than one at the end; in Swin, removing an MSA at the end of a stage hurts most §4.1 fig. 10. The authors conclude that MSAs near the end of a stage matter, which contrasts with the popular view that MSAs near the end of the whole model matter §4.1.

AlterNet build-up rules. The authors state three rules §4.1:

• Alternately replace Conv blocks with MSA blocks, starting from the end of a baseline CNN.
• If an added MSA does not improve accuracy, replace a Conv block at the end of an earlier stage instead.
• Use more heads and larger hidden dimensions for MSA blocks in later stages.

Applied to ResNet-50 on CIFAR-100, MSAs in the last stage improve accuracy, one MSA in the second stage also improves it, and two or more MSAs in the third stage reduce it §4.1 fig. 12.

AlterNet architecture. The backbone is pre-activation ResNet-50 with one extra CNN block added to the last stage app. E. The stage layout in the released code is: stage 1 all Conv; stage 2 one MSA at the end; stage 3 six blocks alternating Conv and MSA, three of them MSA; stage 4 four blocks alternating Conv and MSA, two of them MSA alternet.py alternet.py. Each stage's last block is an MSA except stage 1 app. E fig. 11. MSAs use 3, 6, 12 and 24 heads across the four stages app. E fig. 11. Attention is local, with a relative positional encoding taken from Swin; the paper says it omits Swin's cyclic shift and layer-specific initialisation for simplicity app. E. On CIFAR the MSA patch is 1×1 and the window 4×4 app. E. For ImageNet the model uses 6 MSAs app. E.

Local MSA mechanics as implemented. The feature map is cut into non-overlapping windows; attention runs inside each window; a learned bias indexed by relative offset is added to the logits alternet.py alternet.py. The attention scale is the inverse square root of the per-head dimension attentions.py.

# Local MSA with a learned relative position bias  [app. E] [§4.1]
# Input:  x (B, C, H, W); window p; bias_table (2p-1, 2p-1)
# Output: (B, C, H, W)

def local_msa(x, bias_table, p):
    # Cut the feature map into non-overlapping p x p windows
    # [code: models/alternet.py#L41]
    windows = partition_windows(x, p)           # (B*H/p*W/p, C, p, p)
    # Project each window to queries, keys and values
    # [code: models/attentions.py#L39-L43]
    queries, keys, values = project(windows)
    # Scaled dot-product logits  [eq. 1]
    # [code: models/attentions.py#L92]
    logits = dot(queries, keys) * head_dim ** -0.5
    # Add one learned bias per relative offset, shared by all windows
    # [code: models/alternet.py#L38-L39]
    logits = logits + bias_table[relative_index]
    # Softmax over the keys of each query point  [eq. 1]
    # [code: models/attentions.py#L94]
    weights = softmax(logits, axis=-1)
    # Weighted sum of values, then restore the feature-map layout
    # [code: models/alternet.py#L42-L43]
    return unpartition(weights @ values, x)

# Stage layout of AlterNet  [app. E]
# [code: models/alternet.py#L228]
# Input:  n_blocks (blocks in a stage), n_msa (MSA blocks in that stage)
# Output: list with one block type per block

def stage_layout(n_blocks, n_msa):
    # Leading blocks stay convolutional
    conv_prefix = ["conv"] * (n_blocks - 2 * n_msa)
    # Each remaining (conv, msa) pair ends the stage with an MSA
    alternating = ["conv", "msa"] * n_msa
    return conv_prefix + alternating

Training recipe (used for the ViT and AlterNet experiments on CIFAR-100). The paper states: AdamW with initial learning rate 1.25×10⁻⁴ and weight decay 5×10⁻²; cosine annealing; 300 epochs with batch size 96 on CIFAR and 128 on ImageNet; a 5-epoch gradual warmup; and DeiT-style strong augmentation (RandAugment, Random Erasing, label smoothing, mixup, CutMix) plus stochastic depth app. A. Learning-rate schedule as implemented:

# Learning rate  [app. A]
# Input:  base_lr, warmup_iters (5 epochs in iterations), epoch after warmup
# Output: learning rate

def warmup_lr(iteration, warmup_iters, base_lr):
    # Linear ramp from 0 to base_lr, advanced every iteration
    # [code: ops/schedulers.py#L12]
    return base_lr * iteration / warmup_iters

def cosine_lr(epoch, base_lr, t_max):
    # Cosine decay to zero over t_max epochs (t_max = 300 in the config)
    # [code: configs/cifar100_vit.yaml#L44-L47]
    return base_lr * (1 + cos(pi * epoch / t_max)) / 2

Hyperparameters

Values in the table come from the paper unless the notes say "code only". Where the paper gives no value, the code value is listed and marked.

Hardware used (for timing comparisons): two sets of machines for CIFAR (one with a GeForce RTX 2080 Ti, one with an NVIDIA T4), and four GeForce RTX 2080 Ti for ImageNet app. A.

Results

Author-reported. Setup: CIFAR-100 and CIFAR-10 (small data) and ImageNet-1K (large data), with corruption benchmarks CIFAR-100-C and ImageNet-C §2 §4.2 app. A. The text gives few numbers for AlterNet; most accuracy results are plotted only.

Notes:

• Small-data result: the paper reports that ViT does not overfit even at 2% of CIFAR-100 training data §2 fig. 5.
• Lower training NLL goes with lower test error across architectures; ResNet's stronger inductive bias gives better test error at the same training NLL app. C.
• Large-data result: ViTs with local MSAs beat CNNs on ImageNet in both training NLL and error app. C.

Paper vs code

Code: official repository github.com/xxxnell/how-do-vits-work (linked from the abstract), Apache-2.0, commit 9571ff1370f31af2f433558fbfa7bc87a97b0bab.

• Local window size. Paper: window 4×4 for AlterNet on CIFAR app. E. Code: LocalAttention and AttentionBlockB default to window 7 alternet.py alternet.py; dnn_50 and the CIFAR configs do not override it alternet.py cifar100_vit.yaml. Unless the caller passes window_size=4, the released model does not use the paper's CIFAR window.
• MSA placement in stage 3. Paper text: two or more MSAs in stage 3 reduce CIFAR-100 accuracy §4.1 fig. 12. Code default dnn_50 puts three MSAs in stage 3 and two in stage 4 alternet.py alternet.py. The paper does not say which placement produced the reported numbers.
• Random Erasing. Paper lists Random Erasing as used app. A. CIFAR config sets it to 0.0 cifar100_vit.yaml; ImageNet config uses 0.25 imagenet_vit.yaml.
• Total epochs. Paper: 300 epochs including a 5-epoch warmup app. A. Code: 5 warmup epochs followed by 300 scheduled epochs, 305 in total trains.py trains.py.
• Gradient clipping. Not in the paper. CIFAR config clips the gradient norm to 5 cifar100_vit.yaml trains.py; the ImageNet config leaves it commented out.
• Warmup. Paper: gradual warmup app. A. Code: linear per iteration over all warmup iterations schedulers.py.
• Attention formula. Paper's Eq. 1 writes the query-key product as QK eq. 1. The code computes Q times K-transpose, scaled by the inverse square root of the per-head dimension attentions.py eq. 1.
• Loss for landscapes. Paper: landscapes use ℓ2-regularised NLL on augmented training data app. A. Notebook: loss is NLL plus weight decay times an ℓ2 term losslandscape.ipynb (cell line 286 of the notebook JSON).
• Head dimension. Paper does not state it. Code default is 64 per head alternet.py, so the inner dimension grows with the head count.
• Stochastic depth, label smoothing, mixup, CutMix. Values are code only (see the Hyperparameters table) cifar100_vit.yaml alternet.py.

Pitfalls

• Eq. 1 notation. Eq. 1 as printed multiplies Q by K without a transpose eq. 1. Implement the standard form, Q times K-transpose over the square root of d, with softmax over keys attentions.py.
• Bias placement. The relative position bias is added to the attention logits before softmax, not to the values alternet.py.
• Window size on CIFAR. See Paper vs code. The paper's 4×4 and the code's default 7 disagree; set the window explicitly and state which one you used.
• Throughput numbers conflict. Appendix A lists ResNet-50 at 320 images/s and Alter-ResNet-50 at 473 images/s on CIFAR-100 app. A app. E. The ratio is about 1.48, yet appendix E says Alter-ResNet-50 is 20% faster than pre-activation ResNet-50. The text does not reconcile these, so do not use either as a benchmark.
• ImageNet baseline is not matched. PiT and Swin on ImageNet are pretrained timm checkpoints, not trained with the paper's recipe app. A. The "on par with Swin" claim compares different training pipelines app. E.
• Strong augmentation changes calibration. In the reliability diagram on CIFAR-100, ResNet and ViT trained without strong augmentation are over-confident, and with strong augmentation they become under-confident app. F. The paper leaves this for future work app. F.
• Strong augmentation and the Hessian. The paper reports that strong augmentation reduces the magnitude of Hessian eigenvalues but also produces many negative ones app. F. It contrasts this with large datasets, which convexify the loss app. F §2, so the two effects are not the same mechanism.
• Hessian measurement. Landscapes and eigenvalues must use the regularised objective on augmented data; measuring NLL alone on clean data gives different and, per the paper, incorrect results app. A.
• Weight decay is not ℓ2 in the strict sense. The paper says it neglects the difference app. A. The optimiser is AdamW, which decouples weight decay trains.py.
• Normalisation layers. The default normalisation in AlterNet blocks is BatchNorm alternet.py. The paper does not mention normalisation.
• Simplifications in AlterNet. The paper omits Swin's cyclic shift and layer-specific initialisation app. E. Replicating the reported ImageNet numbers needs these to be considered, since the paper says the omission is for simplicity app. E.
• Reduced-depth ablations. ViT depth is halved for some multi-head and embedding-dimension experiments due to memory app. A. Those numbers do not describe the full-depth model.
• Figures only. Most accuracy results are in figures without values in the text fig. 12 fig. 6 app. E. Recover numbers from the authors' repository or re-run; do not infer them from plots.
• Claims are often qualitative. The text repeatedly uses "MSAs flatten the loss landscape" and "low-pass filter" as conclusions drawn from spectra and Fourier plots §2 §3. Treat them as the authors' interpretation, supported by the cited figures.

Lineage

• Builds on: the Vision Transformer §1 and multi-stage and local-window variants PiT and Swin §2; the view of blur and spatial smoothing as ensembling §1.1 app. B. These works are named in the paper and have no card here.
• Compared against: ResNet, ResNeXt, PiT, Swin, ViT, Twins, MLP-Mixer, ConViT and LeViT app. A app. E.
• Superseded by: not stated in this wiki.