Vision Transformers (ViT) and When to Use Them

A shipping container with a dented side panel looks completely normal in close-up. The dent only reads as damage when you compare that panel against the straight edges of the container three metres away. A convolution kernel sees a 3×3 neighbourhood. That comparison is not available to it until many layers deep.

Prerequisites: CNN architectures and your first image classifier.

Why convolution struggles here

A container yard inspector flags damage. The visual cue for a dent is not local texture — it is that one panel is no longer coplanar with the panels far from it.

A 3×3 convolution’s receptive field grows slowly. It takes many stacked layers before any single unit can see both the dent and the reference edge, and by then the signal has passed through dozens of nonlinearities.

Convolution’s bias is usually right, which is why CNNs work so well on limited data. But when a task genuinely depends on long-range relationships, that bias is a constraint. ViT removes it and pays for the removal with data.

An image as a sequence

The whole idea is one preprocessing step.

Each patch flattens into a vector of length 16 and goes through one shared linear layer. Now you have four vectors, and everything after this point is the standard transformer used for text.

The real numbers for ViT-Base/16

Input 224×224×3, patch size 16.

$N = \left(\frac{224}{16}\right)^2 = 14^2 = 196 \text{ patches}$

Each patch flattened:

$16 \times 16 \times 3 = 768 \text{ values}$

That flattened vector goes through a linear projection to the model dimension $D = 768$:

The CLS token is an extra learned vector prepended to the 196 patches, giving 197 tokens. It belongs to no patch. Its job is to gather information from all of them through attention, and its final state is what the classifier reads.

This matters more than it sounds. Shuffle the patches and a transformer without positional embeddings produces exactly the same output. A CNN could never be confused this way, because position is implicit in its structure.

Attention, computed

The mechanism, in one formula:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$

Work it with $d_k = 4$ and three tokens. Take one query vector and three key vectors:

$q = [1, 0, 1, 0]$ $k_1 = [1, 0, 1, 0], \quad k_2 = [0, 1, 0, 1], \quad k_3 = [1, 1, 0, 0]$

Step 1 — dot products.

$q \cdot k_1 = 1 + 0 + 1 + 0 = 2$ $q \cdot k_2 = 0 + 0 + 0 + 0 = 0$ $q \cdot k_3 = 1 + 0 + 0 + 0 = 1$

Step 2 — scale by $\sqrt{d_k} = \sqrt{4} = 2$.

$[2, 0, 1] / 2 = [1.0,\; 0.0,\; 0.5]$

Step 3 — softmax.

$e^{1.0} = 2.7183, \quad e^{0.0} = 1.0000, \quad e^{0.5} = 1.6487$ $\text{sum} = 5.3670$ $\text{weights} = [0.5065,\; 0.1863,\; 0.3072]$

Step 4 — weighted sum of values. With $v_1 = [10, 0]$, $v_2 = [0, 10]$, $v_3 = [5, 5]$:

$0.5065[10,0] + 0.1863[0,10] + 0.3072[5,5]$ $= [5.065, 0] + [0, 1.863] + [1.536, 1.536] = [6.601,\; 3.399]$

The output is a blend, weighted by how similar this query was to each key. Token 1 matched best and contributed half.

Why divide by √d

Repeat step 3 without the scaling, using the raw scores 2, 0, 1:

$e^2 = 7.389, \quad e^0 = 1.000, \quad e^1 = 2.718, \quad \text{sum} = 11.107$ $\text{weights} = [0.6653,\; 0.0900,\; 0.2447]$

Dot products of $d$-dimensional random vectors have variance proportional to $d$. At $d_k = 64$ per head, raw scores routinely reach ±20, and $e^{20}$ against $e^{-20}$ is a ratio of $10^{17}$ — softmax becomes a hard argmax, and the gradient through it is essentially zero for every other token. Dividing by $\sqrt{d_k}$ normalises the variance back to roughly 1 and keeps the distribution soft enough to learn from.

Multi-head attention and the full block

One attention operation produces one weighting. ViT-Base uses 12 heads, each with $d_k = 768/12 = 64$, run in parallel and concatenated back to 768. Different heads reliably specialise: some track texture, some track position, some latch onto the CLS token.

graph TD
  A["Input tokens 197×768"] --> B["LayerNorm"]
  B --> C["Multi-head attention<br/>12 heads × 64 dims"]
  C --> D["+ residual"]
  A --> D
  D --> E["LayerNorm"]
  E --> F["MLP: 768 → 3072 → 768<br/>GELU"]
  F --> G["+ residual"]
  D --> G
  G --> H["Output 197×768"]

Note the layer norm comes before each sublayer, not after. The original text transformer put it after; ViT and everything since moved it in front because post-norm needs a careful learning-rate warmup to train at all, while pre-norm just works.

The MLP expands 768 to 3072 and back. Per block:

Most people assume attention dominates. It does not — it dominates the compute at long sequences, but the MLP holds most of the weights.

The quadratic cost

Every token attends to every token, so the attention matrix is $N \times N$.

This is the reason Swin Transformer exists. It computes attention only within local 7×7 windows and shifts the windows between layers, making cost linear in the number of tokens while still mixing information globally over several layers. For detection and segmentation at high resolution, Swin-style models are the practical choice.

The data requirement — the part that decides your project

The original ViT paper’s most useful result is not that ViT wins. It is when it wins.

The interpretation is clean. Convolution’s inductive bias is free knowledge — locality and translation invariance are true facts about images that the architecture gets without learning. ViT must learn them from examples. With enough examples it learns something better; without them it learns something worse.

Fine-tuning in practice

import timm, torch

model = timm.create_model('vit_base_patch16_224.augreg_in21k',
                          pretrained=True, num_classes=4)

# Small dataset: freeze the backbone, train the head only
for p in model.parameters():
    p.requires_grad = False
for p in model.head.parameters():
    p.requires_grad = True

opt = torch.optim.AdamW(model.head.parameters(), lr=1e-3, weight_decay=0.05)

If you unfreeze everything, drop the learning rate hard — ViT is far more sensitive than a CNN:

Practice task

Use the same four-class dataset from the classifier post.

1. Compute the token count for patch sizes 8, 16, and 32 at 224×224. Then compute the attention pairs for each.
2. Fine-tune vit_base_patch16_224 with only the head unfrozen. Record accuracy.
3. Fine-tune a resnet50 the same way. Compare on your dataset size.
4. Unfreeze the full ViT at learning rate 1e-3. Watch it diverge. Then retry at 1e-5.
5. Add 10% warmup and gradient clipping at 1.0. Compare final accuracy.
6. Extract the CLS token’s attention map from the last block and overlay it on a few images.
7. Shuffle the patch order at inference. See how badly it breaks, which shows what positional embeddings were doing.
8. Repeat step 2 at 384×384 and measure the actual latency increase against the predicted 8.58×.

Step 4 is worth doing deliberately rather than reading about. The divergence is fast and unmistakable, and it makes the learning-rate table above stop being arbitrary advice.

Summary

ViT splits a 224×224 image into 196 patches of 16×16, flattens each to 768 values, projects them, adds a CLS token for 197 tokens, and adds learned positional embeddings — 151,296 parameters just to say where each patch came from.

Attention scores each token pair by dot product, divides by $\sqrt{d_k}$, softmaxes, and takes a weighted sum of values. With scores 2, 0, 1 and $d_k = 4$, the weights are 0.5065, 0.1863, 0.3072. Skip the scaling and they become 0.6653, 0.0900, 0.2447 — sharper, and at real dimensions sharp enough to kill the gradient.

Two thirds of each block’s 7.1M parameters live in the MLP, not attention. But attention costs O(n²), so 224 → 384 pixels multiplies token pairs by 8.58 and 512 pixels by 27.

The decision is about data. Below roughly 10M pretraining images a ResNet wins; above it ViT does. And ViT needs its own recipe: AdamW at 1e-5, warmup, weight decay 0.05, gradient clipping.

What comes next

You now have a model that produces attention maps, which look like an explanation. They are not, quite. Model explainability covers Grad-CAM, what attention maps do and do not tell you, and how to run a failure analysis that finds the bias you did not know your dataset had.