Vision Transformers, from pixels to prediction

See it in action: one face through a tiny ViT

This is a complete Vision Transformer, shrunk so every number fits on screen: a 224×224 image, 32×32 patches (so 7×7 = 49 patches), token width D = 8, two attention heads, two encoder layers. The structure and operations are the ones ViT-Base uses, which has D = 768, 12 heads, 12 layers and 16×16 patches. Step through it, and hover a patch on the face (or a row of the token matrix) to follow that token.

Why use a transformer for images at all?

Transformers took over language modelling because self-attention lets every word look at every other word in one step, and the architecture scales well with data and compute. The ViT paper asked a blunt question: what happens if you apply a standard Transformer to images with as few image-specific changes as possible? The answer: split the image into fixed-size patches, treat each patch like a word, and feed the sequence to an unmodified Transformer encoder.

Inductive bias: what the architecture assumes before seeing data

A CNN hard-codes two strong assumptions into every layer:

• Locality. A 3×3 filter only looks at neighbouring pixels, so the network assumes nearby pixels belong together.
• Translation equivariance. The same filter slides everywhere, so a pattern means the same thing wherever it appears.

ViT keeps almost none of this. Only the patch cutting is local, and only the patch embedding weights are shared across positions. Self-attention is global from the first layer, and the model is not told that patch 15 sits directly below patch 1; it has to learn spatial relations from the position embeddings.

What the paper found

• Trained only on ImageNet-1k (about 1.3M images) with standard regularization, ViT models land a few points below comparable ResNets, and larger ViTs do no better than smaller ones.
• Pre-trained on ImageNet-21k (about 14M images) or the private JFT-300M (about 300M images) and then fine-tuned, ViT matches or beats strong ResNet baselines while needing considerably less pre-training compute.
• The paper's own summary: large-scale training trumps inductive bias.

From image to sequence: patches

A Transformer consumes a sequence of vectors. An image x ∈ ℝ^H×W×C is reshaped into a sequence of N flattened patches x_p ∈ ℝ^N×(P²·C), where P is the patch size:

Each patch has P·P pixels with C = 3 colours, so its flattened vector has P²·C numbers: 16·16·3 = 768 for ViT-B/16. The "/16" in model names is the patch size. Patches are ordered row by row (raster order); that order means nothing to the Transformer until position embeddings (Chapter 4) add it back.

Choosing P: the central cost trade-off

Smaller patches keep finer detail but make the sequence longer, and self-attention compares every token with every other, so its cost grows with N². Halving P from 16 to 8 gives 4× the tokens (784) and 16× the attention entries per head (785² vs 197²). This is why ViT-H uses /14 rather than /8, and why CvT and MobileViT shrink the token count early.

Patch embedding: a learned linear map

Raw 768-number patches are mapped to the model width D by one trainable matrix E ∈ ℝ^(P²·C)×D (plus a bias), applied to every patch:

The same E is shared by all patches. So this is exactly a convolution with kernel size P, stride P and D output channels, which is how every real implementation computes it: nn.Conv2d(3, D, kernel_size=P, stride=P). Parameter count: P²·C·D + D = 768·768 + 768 = 590,592 for ViT-B/16.

What does E learn? Each of its D columns is a 16×16×3 pattern; the embedding value is the dot product of the patch with that pattern. The paper shows that the top principal components of the learned filters look like plausible basis functions (smooth gradients, edges and stripes at different frequencies) for describing the fine structure inside a patch. The widget computes that kind of basis directly from your image with PCA.

The CLS token and position embeddings

Two learnable additions turn the patch embeddings into the encoder's input z₀:

The [class] token

Borrowed from BERT, x_class is one extra learnable vector placed at position 0. It has no image content of its own. Through attention it gathers information from all patches, layer by layer, and its final state z_L⁰ is the only token sent to the classifier: y = LN(z_L⁰). An alternative is global average pooling over the patch tokens; the paper reports both work when the learning rate is tuned for each.

Why position embeddings are necessary

Self-attention is permutation-equivariant: shuffle the input tokens and the outputs are shuffled the same way, with no other change. With a CLS token and no position information, the CLS output is identical for every shuffling of the patches. A face and the same face cut up and rearranged would get exactly the same prediction. Adding a distinct vector to each position breaks that symmetry.

ViT uses learned 1D position embeddings, one vector per index 0…N, added (not concatenated) to the tokens. The paper also tried 2D-aware and relative embeddings and found no significant difference. Even though the 1D embeddings only know the index, the learned ones end up encoding 2D layout: patches in the same row or column get similar embeddings. The widget uses fixed sinusoidal embeddings to show both kinds of structure.

Self-attention, from the equations

Given the token matrix z ∈ ℝ^(N+1)×D, three learned matrices project every token into a query, a key and a value:

• Query: what this token is looking for. Key: what this token offers to be found by. Value: what this token passes on if attended to.
• Row i of q kᵀ scores token i's query against every key. Softmax turns each row into non-negative weights summing to 1.
• The output for token i is a weighted average of all value vectors. Every token can draw on every other token in one layer, so the receptive field is global from the start.

Why divide by √D_h?

If the entries of q and k are roughly independent with mean 0 and variance 1, their dot product is a sum of D_h such products, with variance D_h and standard deviation √D_h. Without scaling, larger heads produce larger scores, softmax pushes nearly all the weight onto one token, and its gradients vanish. Dividing by √D_h brings the variance back to 1 for any head size.

Multi-head attention

One attention map can only express one way of weighting the other tokens. Multi-head self-attention (MSA) runs k heads in parallel, each with its own U_qkv, then concatenates them and mixes them with an output projection:

To keep cost constant, the head size is D_h = D/k: ViT-B has 12 heads of 64. In code, one Linear(D, 3D) computes all queries, keys and values at once; the result is reshaped to (B, heads, N+1, D_h) so each head attends independently. Parameters: 3D² + 3D for qkv and D² + D for the output projection, so 4D² + 4D = 2,362,368 for ViT-B.

Heads specialise. The paper measured each head's mean attention distance (how far, in pixels, it looks on average): even in the first layer, some heads attend across most of the image while others stay local, behaving like early convolution filters. With depth, attention distance grows for all heads. Hover the face below to see four hand-set heads that mimic these roles.

The encoder block: LayerNorm, MLP, residuals

ViT stacks L identical blocks. Each has two sub-layers, each wrapped in a pre-norm residual:

• LayerNorm normalizes each token across its D features (mean 0, variance 1, then learned γ, β). It doesn't depend on the batch, unlike BatchNorm.
• Pre-norm (LN before each sub-layer, as in ViT) keeps an untouched identity path from input to output, which makes deep stacks train stably.
• Residuals mean each block only adds a refinement to the token. Gradients flow straight through the additions.
• MLP: two linear layers with GELU, applied to each token separately: D → 4D → D (768 → 3072 → 768 in ViT-B). Attention moves information between tokens; the MLP processes it within each token. The MLP holds about two-thirds of each block's parameters.

The whole model: shapes, parameters and compute

Put the pieces together and every size follows from a handful of hyperparameters. Pick a preset or change any value; the shape trace, the parameter breakdown and the compute are all calculated, not looked up.

Training, fine-tuning and higher resolution

Pre-training

The paper pre-trains on a large dataset with Adam (β₁ = 0.9, β₂ = 0.999), a large batch (4096), high weight decay (0.1), and linear warm-up followed by learning-rate decay. During pre-training the head is an MLP with one hidden layer; for fine-tuning it is replaced by a single zero-initialized linear layer with one output per class.

Fine-tuning, and why resolution needs special handling

Fine-tuning at a higher resolution than pre-training often helps. With the patch size fixed, more pixels mean more patches: at 384×384 with /16, the grid is 24×24 = 576 patches instead of 14×14 = 196. The Transformer itself handles any sequence length, but the learned position embeddings only exist for 196 positions. ViT reshapes them back into their 14×14 grid and interpolates them in 2D to 24×24. This is one of only two places where ViT is explicitly told the image is 2D; the other is patch extraction.

Applying this to a small dataset such as BioVid

• Start from pre-trained weights. ImageNet-21k pre-training supplies the visual knowledge that 87 participants can't.
• Replace the head with Linear(D, 2) for pain / no pain (or a single sigmoid unit).
• Split by participant. Frames from one person must never appear in both training and test data, or the model can recognise people instead of pain.
• Use a lower learning rate for the backbone than the head, plus augmentation and early stopping.

Looking inside a trained ViT, and what came next

Three things the paper inspected

• Embedding filters: the columns of E resemble basis functions for fine structure inside a patch (Chapter 3).
• Position embeddings: closer patches have more similar embeddings, with clear row and column structure, so 2D layout is learned (Chapter 4).
• Attention distance: some heads are global from the first layer, others are local, and all become more global with depth (Chapter 6).

Attention rollout: what does the CLS token look at?

One layer's attention map ignores the residual stream and everything before it. Attention rollout (Abnar & Zuidema, 2020; used in the ViT paper) combines all layers: average the heads, add the identity to account for the residual, re-normalize the rows, and multiply through the layers:

From ViT to the models in your study