Vision transformers

The architecture#

Dosovitskiy et al. (2021) applied the standard Transformer encoder to images with essentially no vision-specific modification.

Patch embedding. Reshape x in R^(H x W x C) into N flattened patches and project:

N = H * W / P^2
z_0 = [ x_class ; x_p^1 E ; x_p^2 E ; ... ; x_p^N E ] + E_pos

• P — patch side length, 16 in ViT-B/16.
• x_p^i in R^(P^2 * C) — patch i, flattened.
• E in R^(P^2 C x D) — the linear projection; D is the model width, 768 for ViT-B.
• x_class in R^D — a learned token prepended to the sequence.
• E_pos in R^((N+1) x D) — learned position embeddings.

For 224 x 224 input with P = 16, N = 196 and the sequence length is 197. The projection E is implemented as a Conv2d with kernel_size = stride = P, which is numerically identical and far faster.

Encoder block, pre-norm ordering:

z'_l = MSA( LN(z_{l-1}) ) + z_{l-1}
z_l  = MLP( LN(z'_l) )    + z'_l

• LN — layer normalisation.
• MSA — multi-head self-attention.
• The MLP is two linear layers with GELU, expansion ratio 4.

Attention:

Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V

• Q = z W_Q, K = z W_K, V = z W_V, each in R^(N x d_k).
• d_k = D / h — per-head dimension; h is the number of heads, 12 in ViT-B.
• The sqrt(d_k) divisor keeps the dot products from saturating the softmax as d_k grows.

Configurations:

Cost#

Per layer, with sequence length N and width D:

attention:  O(N^2 * D)     the  N x N  score matrix, per head
MLP:        O(N * D^2)     the two projections

The crossover is at N ≈ D. For ViT-B/16 at 224 pixels, N = 197 and D = 768, so the MLP dominates and attention is not the bottleneck. Double the resolution to 448 and N = 784, now exceeding D, and attention takes over — cost in the attention term rises sixteen-fold for a four-fold increase in pixels.

This is why plain ViT does not scale to dense prediction at high resolution, and why every hierarchical variant restricts attention somehow.

Memory for the score matrix alone is O(h * N^2) per layer in the naive implementation. FlashAttention (Dao et al., 2022) computes exact attention without materialising it, tiling the computation to stay in on-chip SRAM. It is an IO-complexity result, not an approximation, and it is the reason long-sequence attention became affordable.

Position embeddings#

The original paper compared learned 1-D, learned 2-D and relative embeddings, and found little difference — 1-D learned is the default. The result is less surprising than it first appears. A learned embedding can encode 2-D structure without being told the layout, and the paper's own visualisations show exactly that recovery.

Practical consequences:

• Resolution change requires interpolation. E_pos is learned for a specific N. Changing the input size means resampling the position grid, usually bicubically, with a measurable accuracy cost. Fine-tuning after interpolation recovers most of it.
• Relative position bias (used in Swin) generalises across resolutions more gracefully.
• Rotary embeddings, standard in language models, are increasingly used in vision and remove the fixed-length constraint.

Data scale is the whole story#

The paper's central experiment: ViT trained on ImageNet-1k underperforms a comparable BiT ResNet; on ImageNet-21k they are close; on JFT-300M the ViT wins outright and scales better.

The interpretation is about inductive bias. Convolution supplies locality and translation equivariance as hard architectural constraints. A transformer must learn both from data. When data is plentiful, learned structure beats imposed structure. When it is scarce, the imposed structure is worth roughly an order of magnitude of data.

Three lines of work reduced the data requirement substantially:

DeiT (Touvron et al., 2021) reached competitive ImageNet-1k accuracy with no external data, using heavy augmentation, stochastic depth, repeated augmentation, and a distillation token supervised by a convolutional teacher. The distillation token is the interesting part — the ViT learns from a CNN's inductive bias without inheriting its architecture.

AugReg (Steiner et al., 2021) systematically measured augmentation and regularisation against dataset size, and showed that a well-regularised ViT on ImageNet-21k matches a ViT trained on much more data. Their headline conclusion is uncomfortable and useful: many published architecture comparisons are confounded by training recipe.

Self-supervised pre-training. MAE (He et al., 2022) masks 75 percent of patches and reconstructs them; the asymmetric encoder sees only visible patches, making pre-training cheap. DINO and DINOv2 (Caron et al., 2021; Oquab et al., 2023) use self-distillation and produce features whose attention maps segment objects without any segmentation labels. DINOv2 features support linear probes that approach full fine-tuning on many dense tasks.

Hierarchical variants#

Plain ViT has one resolution throughout, which suits classification and not dense prediction.

Swin Transformer (Liu et al., 2021) computes attention within local windows and shifts the window partition between consecutive blocks so information crosses boundaries. Cost becomes linear in image size rather than quadratic, and the multi-scale feature maps drop into existing detection and segmentation necks.

PVT, MViT and others reduce sequence length progressively by pooling keys and values.

ViTDet (Li et al., 2022) is the counterargument: plain, non-hierarchical ViT backbones with a simple feature pyramid built afterwards are competitive for detection, given strong pre-training such as MAE. Hierarchy may be a convenience rather than a necessity.

The ConvNeXt correction#

Liu et al. (2022) took a ResNet-50 and applied, one at a time, the design and training choices that came with transformers: the training recipe, depthwise convolutions with large kernels, an inverted bottleneck, fewer normalisation layers, GELU, LayerNorm. The resulting pure ConvNet matched Swin at equal FLOPs across classification, detection and segmentation.

The honest reading is not "convolutions won". It is that the 2020-to-2021 transformer results conflated architecture with training methodology, and once the methodology was transferred, the architectural gap largely closed. Any claim that attention is intrinsically better for vision has to engage with this paper.

A known artefact#

Darcet et al. (2024), Vision Transformers Need Registers, identified high-norm tokens appearing in low-information background patches of trained ViTs, including DINOv2. The model appears to repurpose uninformative patches as scratch storage for global information, which corrupts the feature map for dense tasks and makes attention visualisations misleading.

The fix is small and effective: append a few extra learnable "register" tokens that are discarded at output. Attention maps become clean and dense-task performance improves. It is a useful reminder that trained networks develop internal conventions nobody designed, and that reading attention maps as explanation is hazardous.

Papers#

• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762
• Dosovitskiy, A. et al. (2021). An Image is Worth 16x16 Words. arXiv:2010.11929
• Touvron, H. et al. (2021). Training data-efficient image transformers and distillation through attention (DeiT). arXiv:2012.12877
• Caron, M. et al. (2021). Emerging Properties in Self-Supervised Vision Transformers (DINO). arXiv:2104.14294
• Liu, Z. et al. (2021). Swin Transformer. arXiv:2103.14030
• Steiner, A. et al. (2021). How to train your ViT? arXiv:2106.10270
• Liu, Z. et al. (2022). A ConvNet for the 2020s. arXiv:2201.03545
• He, K. et al. (2022). Masked Autoencoders Are Scalable Vision Learners. arXiv:2111.06377
• Dao, T. et al. (2022). FlashAttention. arXiv:2205.14135
• Li, Y. et al. (2022). Exploring Plain Vision Transformer Backbones for Object Detection (ViTDet). arXiv:2203.16527
• Oquab, M. et al. (2023). DINOv2. arXiv:2304.07193
• Darcet, T. et al. (2024). Vision Transformers Need Registers. arXiv:2309.16588

Where this stands#

The plain ViT is now infrastructure rather than a result. It is the default backbone for multimodal models, for open-vocabulary detection and segmentation, and for self-supervised feature learning, mainly because it shares an architecture and a scaling story with language models. Whether it is the best classifier at a given FLOP budget is no longer the interesting question.

Three problems remain genuinely open.

High-resolution dense prediction. Quadratic attention still constrains what is affordable. Window attention, token merging and linear-attention approximations all trade something away, and none is a clean win.

Small-data regimes. Most real projects have thousands of images, not millions. Self-supervised pre-training plus a linear probe is currently the best answer. It also presumes a pre-trained backbone matching your domain, which for medical, satellite and industrial imagery frequently does not exist.

Interpretability. Attention maps are not explanations, the register-token result shows that trained ViTs hide state in unexpected places, and faithful attribution for transformers remains unsolved. Anyone deploying one where a wrong answer has a cost should assume they cannot explain its decisions, and design the surrounding system accordingly.

What to learn next#

• Transformers — the architecture this borrows wholesale, explained from the language side.
• Attention — the single operation everything above is built on.
• Image segmentation — where transformer backbones now dominate.