Vision transformers

The big idea

A transformer reads a sentence as a sequence of tokens and lets every token attend to every other one ( Self-attention, step by step ). Could the same machine read an image? Attending from every pixel to every pixel is far too expensive: a 224 by 224 image has 50,176 pixels. The vision transformerA transformer that classifies images by splitting them into fixed-size patches, embedding each patch as a token, and processing the sequence with a standard transformer encoder. (ViT) makes it practical with one move: cut the image into 16 by 16 patches and treat each patch as a word. The 2020 paper that introduced it is titled "An Image is Worth 16x16 Words".

Think of a jigsaw puzzle you are describing over the phone. You do not read out every pixel; you describe piece by piece, and you say where each piece sits ("top row, second from the left"). The listener can relate any piece to any other. ViT does the same: each patch becomes a token, a position embedding says where it was, and attention relates every patch to every other one from the very first layer.

How it works

Step 1: Patches as tokens

An image of height H, width W and C channels is reshaped into N = HW / P squared flattened patches of P by P by C numbers each. With P = 16, a 224 by 224 RGB image becomes 196 patches of 16 x 16 x 3 = 768 numbers. One trainable linear layer, shared by all patches, maps each patch to a vector of size D: the patch embeddingThe vector a vision transformer gets by multiplying a flattened image patch by a learned matrix; it plays the role of a word embedding.. The patch count is the sequence length, so smaller patches mean more tokens: the paper notes that sequence length is inversely proportional to the square of the patch size.

That linear layer is a convolution in disguise: a convolution with a 16 by 16 kernel and stride 16 ( Convolutions ) visits exactly the non-overlapping patches and applies the same weights to each. Libraries often implement the patch embedding that way, and the code below checks the equivalence.

Step 2: The class token and position embeddings

Like BERT's [class] token, ViT puts one extra learnable embedding in front of the patch tokens. After the last layer, its output serves as the representation of the whole image, and a small MLP head maps it to class scores. Pooling the patch outputs instead is a common alternative.

Attention by itself is order-blind: shuffle the patches and the outputs shuffle with them. So a position embeddingA vector added to each token that encodes its position in the sequence; in ViT it tells the model where in the image a patch came from. is added to every token. The paper uses standard learnable 1D position embeddings, one per sequence position, because it saw no significant gains from more advanced 2D-aware versions. The model learns the 2D layout on its own: the paper visualizes how embeddings of nearby patches end up similar.

Step 3: The transformer encoder, unchanged

The token sequence goes through a standard transformer encoder: L layers, each with multi-head self-attention and an MLP block, layer norm before each and a residual connection around each ( Vanishing and exploding gradients, residual connections ). Nothing is image-specific. The paper's variants reuse BERT sizes:

• ViT-Base: 12 layers, D = 768, 12 heads, 86 million parameters.
• ViT-Large: 24 layers, D = 1024, 16 heads, 307 million parameters.
• ViT-Huge: 32 layers, D = 1280, 16 heads, 632 million parameters.

Names add the patch size: ViT-L/16 is the Large model with 16 by 16 patches.

Step 4: ViT versus CNN: inductive bias and data

A CNN has two assumptions built in ( Pooling and CNN architecture ): inductive biasesAssumptions built into a model's architecture that steer what it learns, such as a CNN's locality and weight sharing.. Locality (a kernel looks at a small neighbourhood) and translation equivariance (the same kernel everywhere, so a shifted object gives a shifted response). ViT has almost none: apart from cutting patches, it must learn from data that nearby pixels matter together.

That cuts both ways. The ViT paper reports that trained on a mid-sized dataset such as ImageNet without strong regularization, ViTs land a few percentage points below ResNets of comparable size, and attributes this to the missing inductive biases. Pre-trained on 14 million to 300 million images (ImageNet-21k or the in-house JFT-300M), the picture flips: "large scale training trumps inductive bias", and the best model reached 88.55 percent on ImageNet. Later work narrowed the data gap: DeiT trained a competitive 86-million-parameter ViT on ImageNet alone, reporting 83.1 percent top-1 with no external data, using strong augmentation ( Data augmentation ) and distillation from a teacher.

The freedom also shows in what ViT learns. Self-attention lets it integrate information across the whole image even in the lowest layers, and the paper finds that some heads attend to most of the image already in layer 1, while others stay local, as CNN kernels do.

Going deeper

Code: a ViT front end in NumPy

import numpy as np
rng = np.random.default_rng(0)

H = W = 224; C = 3; P = 16; D = 768            # ViT-Base/16 sizes
img = rng.random((H, W, C))

# 1. Patchify: (224, 224, 3) -> (14 * 14 patches, 16 * 16 * 3 values)
patches = img.reshape(H // P, P, W // P, P, C).transpose(0, 2, 1, 3, 4).reshape(-1, P * P * C)
print("patches:", patches.shape)

# 2. Linear patch embedding (shared by all patches), 3. prepend [class], 4. add positions
E = rng.normal(0, 0.02, (P * P * C, D)); b = np.zeros(D)
tokens = patches @ E + b
cls = rng.normal(0, 0.02, (1, D))
pos = rng.normal(0, 0.02, (len(tokens) + 1, D))   # learned in a real model
x = np.vstack([cls, tokens]) + pos
print("sequence into the transformer:", x.shape)
print(f"patch embedding parameters: {E.size + b.size:,}")

# The patch embedding is a convolution with kernel size = stride = 16
r, c = 3, 5                                         # patch at row 3, column 5
window = img[r*P:(r+1)*P, c*P:(c+1)*P, :].reshape(-1)
print("conv at that window equals token", r * 14 + c, ":", np.allclose(window @ E + b, tokens[r * 14 + c]))

# 5. One self-attention head over all tokens: every patch can look at every other patch
Wq, Wk = rng.normal(0, 0.02, (2, D, 64))
scores = (x @ Wq) @ (x @ Wk).T / np.sqrt(64)
A = np.exp(scores - scores.max(1, keepdims=True)); A /= A.sum(1, keepdims=True)
print("attention matrix:", A.shape, " row sums:", np.round(A.sum(1)[:3], 6))

# Sequence length and attention cost grow fast as patches shrink
for p in (32, 16, 8):
    n = (224 // p) ** 2 + 1
    print(f"patch {p:2}: {n:4} tokens, {n * n:9,} attention scores per head per layer")

patches: (196, 768)
sequence into the transformer: (197, 768)
patch embedding parameters: 590,592
conv at that window equals token 47 : True
attention matrix: (197, 197)  row sums: [1. 1. 1.]
patch 32:   50 tokens,     2,500 attention scores per head per layer
patch 16:  197 tokens,    38,809 attention scores per head per layer
patch  8:  785 tokens,   616,225 attention scores per head per layer

The code patchifies a 224 by 224 image into 196 vectors of 768 numbers, embeds them, prepends the class token and adds positions, giving the 197 by 768 sequence a ViT-Base/16 sees. The patch embedding has 590,592 parameters (768 x 768 weights plus 768 biases). Applying the same weights to one 16 by 16 window, as a stride-16 convolution would, reproduces that patch's token exactly. The last lines show the cost of smaller patches: halving the patch size from 16 to 8 quadruples the tokens to 785 and multiplies the attention scores per head per layer by about 16, from 38,809 to 616,225.

Resolution, cost and variants

Self-attention costs grow with the square of the number of tokens, and tokens grow with the square of the resolution. So high-resolution inputs and small patches get expensive fast. Fine-tuning at a higher resolution than pre-training is common and helpful; ViT keeps the patch size, which lengthens the sequence, and 2D-interpolates the pre-trained position embeddings to the new grid. The paper also tests a hybrid: patches taken from a CNN feature map instead of raw pixels. Later designs bring back some CNN ideas, such as attention within local windows and stages that shrink the token grid, to handle dense tasks like detection and segmentation ( Image segmentation ) at high resolution.

Why it matters beyond classification

Turning an image into a sequence of tokens gives vision the same interface as language. Multimodal models encode an image into patch tokens and feed them to a language model next to text tokens; how many tokens an image costs is covered in Multimodal models . The same transformer stack, training tricks and hardware kernels serve both.

Trade-offs

Common mistakes

• Training a ViT from scratch on a few thousand images and expecting it to beat a pretrained CNN. Start from pretrained weights ( Transfer learning and fine-tuning ) or use a CNN.
• Changing the input resolution at fine-tuning without interpolating the position embeddings, so the sequence and the embeddings no longer match.
• Forgetting that halving the patch size quadruples the sequence length and multiplies attention cost by about 16.
• Thinking ViT ignores where patches are. Position embeddings carry location; without them it would see a bag of patches.
• Treating the [class] token as special magic. It is one more learned token that attention lets gather information from all patches.

In the wild

Vision transformers are now a standard backbone choice. They are the image encoders inside many multimodal models that read images and text together, they are pretrained on very large image collections and fine-tuned for classification, retrieval and medical imaging, and windowed or hierarchical variants serve as backbones for detection and segmentation. CNNs remain common on phones and embedded devices, where their locality and efficiency at high resolution still pay off.

Check yourself

Recap