A Vision Transformer (ViT) applies the standard Transformer encoder — the same block used by BERT — to images. Instead of words, the input tokens are fixed-size image patches. Introduced by Dosovitskiy et al. in "An Image is Worth 16x16 Words" (2020), ViT showed that with enough pretraining data, a model with almost no image-specific inductive bias can match or beat convolutional networks on classification.
This page covers patch embedding, the encoder block, why ViTs need more data than CNNs, the main variants (DeiT, Swin, hierarchical ViTs), and a fine-tuning example with timm. For the attention mechanism itself see Self-Attention; for the parent topic see Transformer Architecture.
An image of size H × W × C is cut into non-overlapping patches of size P × P. Each patch is flattened to a vector of length P·P·C (768 for a 16×16 RGB patch) and projected by a single learned linear layer to the model width D. In practice this is implemented as a Conv2d with kernel_size=P, stride=P, which is mathematically identical and faster.
A learnable [CLS] token is prepended, and a learnable position embedding (one vector per token position) is added. The sequence z0 ∈ R(N+1)×D then goes through L encoder blocks. The final [CLS] vector, after a LayerNorm, feeds a linear classification head. Many later models replace [CLS] with global average pooling over patch tokens, which works as well or slightly better.
Patch size is the main compute knob: halving P quadruples the token count N = HW/P2, and attention cost grows with N2. ViT-B/16 at 224px uses 196 patches; ViT-B/32 uses 49 and is much cheaper but less accurate.
Each block is pre-norm: z' = z + MHSA(LN(z)), then z'' = z' + MLP(LN(z')). Multi-head self-attention computes softmax(QKT/√dk)V per head over all 197 tokens, so every patch can exchange information with every other patch in the very first layer. The MLP expands to 4D with GELU and projects back.
| Variant | Layers L | Width D | Heads | Params (approx.) |
|---|---|---|---|---|
| ViT-B (Base) | 12 | 768 | 12 | 86M |
| ViT-L (Large) | 24 | 1024 | 16 | 307M |
| ViT-H (Huge) | 32 | 1280 | 16 | 632M |
Attention maps from the [CLS] token are often visualised (attention rollout) to see which patches the model relied on — useful for sanity checks, though not a rigorous explanation.
Convolutions hard-code locality and translation equivariance. ViT only has the patch grid and must learn spatial structure from data. The original paper found that trained only on ImageNet-1k, ViT underperformed comparable ResNets; pretrained on ImageNet-21k or the larger JFT-300M, it overtook them. Learned position embeddings end up encoding 2D row/column structure on their own.
Two practical consequences: (1) do not train a ViT from scratch on a few thousand images — start from a pretrained checkpoint; (2) strong augmentation and regularisation (RandAugment, Mixup, CutMix, stochastic depth) matter more for ViTs than for CNNs.
The timm library provides pretrained ViTs; replacing the head and fine-tuning with a small learning rate is the standard recipe.
import torch, timm
from torch import nn
from torchvision import datasets, transforms
model = timm.create_model("vit_base_patch16_224", pretrained=True, num_classes=10)
cfg = timm.data.resolve_data_config({}, model=model)
train_tf = transforms.Compose([
transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(),
transforms.ToTensor(), transforms.Normalize(cfg["mean"], cfg["std"]),
])
ds = datasets.CIFAR10("data", train=True, download=True, transform=train_tf)
dl = torch.utils.data.DataLoader(ds, batch_size=64, shuffle=True, num_workers=4)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
# lower LR for the backbone, higher for the new head
opt = torch.optim.AdamW([
{"params": [p for n, p in model.named_parameters() if not n.startswith("head")], "lr": 1e-5},
{"params": model.head.parameters(), "lr": 1e-3},
], weight_decay=0.05)
loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)
model.train()
for step, (x, y) in enumerate(dl):
x, y = x.to(device), y.to(device)
with torch.autocast(device_type=device, dtype=torch.bfloat16, enabled=device == "cuda"):
loss = loss_fn(model(x), y)
opt.zero_grad(); loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
if step % 100 == 0:
print(step, loss.item())
if step == 500:
breakInputs must be resized to the resolution the checkpoint expects. To fine-tune at a higher resolution, the position embeddings are interpolated in 2D (timm does this automatically when you pass img_size).
Working through ViT-B/16 by hand shows where the compute and parameters go. A 224 × 224 RGB image split into 16 × 16 patches gives (224/16)2 = 196 patches, plus the [CLS] token, for a sequence of 197. Each patch flattens to 16·16·3 = 768 values, and the patch-embedding projection to width D = 768 has 768 × 768 + 768 ≈ 0.59M parameters.
Inside each encoder block, the Q, K, V and output projections hold about 4D2 ≈ 2.36M weights, and the MLP with hidden width 4D = 3072 holds about 2 · D · 4D ≈ 4.72M. That is roughly 7.1M per block and about 85M over 12 blocks, which together with embeddings and LayerNorms matches the commonly quoted 86M. Two-thirds of the parameters live in the MLPs, not in attention.
Now raise the resolution to 384 × 384 with the same patch size: (384/16)2 = 576 patches, 577 tokens. The parameter count is unchanged (only the position embeddings are interpolated), but each attention matrix grows from 1972 ≈ 38.8K to 5772 ≈ 333K entries, about 8.6 times larger, while the MLP cost grows linearly, about 2.9 times. This is why higher-resolution fine-tuning often helps accuracy but is expensive, and why windowed designs like Swin exist.
Labels are the bottleneck for supervised ViT pretraining, so most strong vision backbones are now pretrained without them. MAE hides a large random fraction of patches (around 75% in the paper), runs the encoder only on the visible ones, and asks a small decoder to reconstruct the hidden pixels. Because the encoder sees a quarter of the tokens, pretraining is much cheaper than running the full sequence. DINO and DINOv2 instead train a student network to match a teacher's output on different crops of the same image; their features work well frozen, with only a linear probe or nearest-neighbour lookup on top. Contrastive image-text pretraining (CLIP, SigLIP) aligns images with captions and enables zero-shot classification by text prompts.
For dense tasks, a plain ViT produces a single-scale grid of patch features (14 × 14 for 224px at P = 16). Detection and segmentation heads usually want multiple scales, so practitioners either use a hierarchical model (Swin, MViT) or build a simple feature pyramid from the plain ViT's last layer, as in ViTDet. Segmentation models such as SegFormer and the Segment Anything image encoder are transformer-based. A practical choice table:
| Task | Reasonable starting backbone |
|---|---|
| Image classification, medium data | Pretrained ViT-B or DeiT, fine-tuned |
| Retrieval / similarity, no labels | Frozen DINOv2 or CLIP features + vector index |
| Detection / segmentation | Swin or ViTDet-style backbone; see Object Detection |
| Edge device, tight latency | Small CNN or mobile hybrid; benchmark on the target hardware |