Vision Transformer

Problem

Convolutional networks dominated vision for a decade because they build in two assumptions: nearby pixels matter together (locality), and a cat is a cat wherever it appears (translation equivariance). Meanwhile the transformer took over language with almost no built-in assumptions, and it kept getting better as data and models grew. Can the same architecture, unchanged, learn vision?

The obvious move fails. Treat every pixel as a token and a 224 × 224 image becomes 50,176 tokens. Self-attention compares every token with every other, so one head in one layer would compute 2,517,630,976 scores, around 2.5 billion. The question is how to turn an image into a sequence short enough for attention while keeping the architecture generic.

Intuition

Cut the picture into tiles, the way you might cut a photo into a 14 × 14 grid of squares. Each tile is small enough to describe with one vector, and there are only 196 of them. Treat each tile as a word: flatten its pixels, multiply by one learned matrix so every tile lands in the same 768-dimensional space, and you have a sentence of 196 tokens. Attention then lets any tile look at any other from the very first layer; a CNN needs many layers before a pixel in one corner can influence a pixel in the other.

Two pieces are borrowed from BERT's encoder-only design. Attention ignores order, so each token gets a learned position embedding that says which tile it came from. And one extra token, the class token, starts as a learned vector with no image content; it attends to all the patches in every layer, and its final state is what the classifier reads. It is an empty seat at the table whose job is to summarize the conversation.

Mechanics

Patches

An image x of height H, width W and C channels is cut into non-overlapping P × P patches. There are N = HW / P² of them, and each flattens to a vector of P²C numbers. For ViT-B/16 on a 224 × 224 RGB image: N = (224 / 16)² = 14² = 196, and each patch holds 16 · 16 · 3 = 768 numbers. (That 768 matching the model width is a coincidence of the B/16 choice; ViT-L/16 has the same patches and width 1024.) The image side must be a multiple of P.

Patch embedding as a strided convolution

Each flattened patch is multiplied by one shared matrix E of shape [P²C, D] and a bias is added. That is a linear layer applied to every patch. It is also exactly a convolution with D output channels, kernel size P and stride P: with stride equal to the kernel, the windows tile the image without overlap, each output position sees exactly one patch, and every position uses the same weights. The conv weight [D, C, P, P] reshaped to [D, C·P·P] is Eᵀ, as long as the patch is flattened in the same channel, row, column order. Implementations use the convolution because it does the cutting, flattening and multiplying in one fused kernel.

The convolution returns [b, D, 14, 14]: a 14 × 14 grid of 768-vectors. flatten(2) merges the two grid axes into [b, D, 196] in row-major order (patch 1 is the top-left, patch 14 the top-right, patch 15 starts the second row), and transpose(1, 2) gives the sequence layout [b, 196, D] the encoder expects.

Class token and position embeddings

z₀ = [x_class; x_p¹E; x_p²E; …; x_p^NE] + E_pos,   E ∈ ℝ^(P²C)×D,   E_pos ∈ ℝ^(N+1)×D

The class token x_class is one learned vector of size D, copied for every image in the batch and prepended, so the sequence has N + 1 = 197 tokens. Position embeddings are a learned table with one row per token position, added (not concatenated) to the tokens. ViT uses plain 1-D learned positions: the paper found that 2-D-aware and relative variants gave no measurable gain, and the learned rows end up encoding the 2-D layout on their own (each position's embedding is most similar to those in its own row and column). Without them, attention would see a bag of patches and could not tell a face from a scrambled face.

Because the table has a fixed number of rows, a model trained at 224 × 224 cannot directly read a 384 × 384 image, which has 24 × 24 = 576 patches. The standard fix keeps the class-token row, reshapes the 196 patch rows into their 14 × 14 grid, resizes that grid to 24 × 24 with 2-D bicubic interpolation, and flattens it back. Fine-tuning at a higher resolution than pre-training is common, so this interpolation is part of every ViT checkpoint loader.

The encoder

The encoder is the standard transformer encoder in its pre-norm form (LayerNorm before each sublayer, residual around it), with GELU in the MLP. For layers ℓ = 1 … L:

z′_ℓ = MSA(LN(z_ℓ−1)) + z_ℓ−1,    z_ℓ = MLP(LN(z′_ℓ)) + z′_ℓ

MSA is multi-head self-attention with 12 heads of width 64, with no causal mask: every patch may attend to every other patch and to the class token. The MLP is Linear(768, 3072) → GELU → Linear(3072, 768). The shape stays [b, 197, 768] through all 12 layers. In PyTorch terms one layer is nn.TransformerEncoderLayer(768, 12, 3072, activation="gelu", norm_first=True, batch_first=True), and the stack is nn.TransformerEncoder(layer, 12); the paper's models use LayerNorm ε = 10⁻⁶ rather than PyTorch's default 10⁻⁵.

The classification head

y = Linear(LN(z_L⁰))   ∈ ℝ^K

After the last layer a final LayerNorm is applied and only token 0, the class token, is kept: [b, 197, 768] → [b, 768]. A linear head maps it to K class logits, trained with cross-entropy. During pre-training the paper used an MLP head with one hidden layer; for fine-tuning it is replaced by a single zero-initialized linear layer sized for the new task. Averaging all patch tokens (global average pooling) instead of reading the class token works as well once the learning rate is tuned for it, and many later models do that.

Inductive bias and data

The only image-specific choices in ViT are the patch cut and the position-embedding resize. Everything else, including which patches relate to which, must be learned. That is a weaker prior than a CNN's, and it shows: trained on ImageNet-1k (1.3M images) alone, ViT-B lagged comparable ResNets. Pre-trained on ImageNet-21k (14M images) or JFT-300M and then fine-tuned, ViT matched or beat them at lower pre-training compute. DeiT (Touvron et al., 2021) later closed the gap on ImageNet-1k alone with heavy augmentation, regularization and distillation from a CNN teacher through an extra distillation token. Early layers still learn local filters, because the data reward them; the architecture just does not force them.

Worked example

All numbers below come from tested helpers in ml_models_docs.py (vit_patchify, vit_conv, vit_shape_walk, vit_params), written in plain Python so the page needs no ML package. The parameter counts match the published totals for torchvision's ViT models.

Implementation

A complete ViT in PyTorch, built only from the modules in the table above (dropout and initialization details left at their defaults):

import torch
from torch import nn


class ViT(nn.Module):
    def __init__(self, image=224, patch=16, dim=768, depth=12, heads=12, mlp=3072, classes=1000):
        super().__init__()
        tokens = (image // patch) ** 2 + 1
        self.embed = nn.Conv2d(3, dim, kernel_size=patch, stride=patch)
        self.flatten = nn.Flatten(2)
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos = nn.Parameter(torch.randn(1, tokens, dim) * 0.02)
        layer = nn.TransformerEncoderLayer(dim, heads, mlp, dropout=0.0, activation="gelu",
                                           layer_norm_eps=1e-6, batch_first=True, norm_first=True)
        self.encoder = nn.TransformerEncoder(layer, depth, norm=nn.LayerNorm(dim, eps=1e-6),
                                             enable_nested_tensor=False)
        self.head = nn.Linear(dim, classes)

    def forward(self, x):                                # [b, 3, 224, 224]
        x = self.flatten(self.embed(x)).transpose(1, 2)  # [b, 196, 768]
        cls = self.cls.expand(x.shape[0], -1, -1)        # [b, 1, 768]
        x = torch.cat([cls, x], dim=1) + self.pos        # [b, 197, 768]
        x = self.encoder(x)                              # [b, 197, 768], final LayerNorm included
        return self.head(x[:, 0])                        # [b, 1000]


model = ViT()
print(sum(p.numel() for p in model.parameters()))       # 86567656
print(model(torch.randn(8, 3, 224, 224)).shape)          # torch.Size([8, 1000])

The count matches the table because nn.MultiheadAttention packs Q, K and V into one [2304, 768] in_proj_weight, exactly the paper's layout. To load a checkpoint at a new resolution, interpolate the position table before loading:

def resize_positions(pos, new_side):
    cls, grid = pos[:, :1], pos[:, 1:]
    side = int(grid.shape[1] ** 0.5)
    grid = grid.reshape(1, side, side, -1).permute(0, 3, 1, 2)
    grid = nn.functional.interpolate(grid, size=(new_side, new_side), mode="bicubic", align_corners=False)
    return torch.cat([cls, grid.permute(0, 2, 3, 1).reshape(1, new_side * new_side, -1)], dim=1)

In practice, torchvision's vit_b_16 and the timm library provide pretrained weights, the same structure with their own encoder block classes, and the position resize built into their loaders.

Tradeoffs

• Data appetite. With few built-in priors ViT needs large pre-training sets (or strong augmentation and distillation, as in DeiT) to beat a CNN; on small datasets from scratch a ResNet usually wins.
• Quadratic tokens. Halving the patch size quadruples the tokens and multiplies attention cost by about 16. Dense tasks (detection, segmentation) want small patches and high resolution, so they use hierarchical or windowed variants (Swin) or plain ViTs with windowed attention in most layers (ViTDet).
• Global from layer one. Every patch can attend to every other immediately, which helps with long-range structure and makes ViTs more robust to some occlusions and patch shuffles than CNNs, but costs compute that local convolutions would not.
• Fixed positions. Learned absolute positions tie the model to a grid size; interpolation works but is a mismatch until fine-tuned. Relative or rotary positions avoid that at some complexity.
• Uniform architecture. One token width at every depth makes ViT simple to scale, shard and pair with text encoders (CLIP), and to pre-train by masking patches (MAE); a CNN's pyramid gives multi-scale features for free.
• Patch boundaries. Non-overlapping patches cut objects at arbitrary lines; small convolutional stems before patching (hybrid models) smooth this and often help optimization.

Connections

References

• Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021. arXiv:2010.11929
• Vaswani et al. Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
• Devlin et al. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL 2019. arXiv:1810.04805
• Touvron et al. Training data-efficient image transformers & distillation through attention (DeiT). ICML 2021. arXiv:2012.12877
• Liu et al. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. ICCV 2021. arXiv:2103.14030
• He et al. Masked Autoencoders Are Scalable Vision Learners. CVPR 2022. arXiv:2111.06377
• Hendrycks and Gimpel. Gaussian Error Linear Units (GELUs). 2016. arXiv:1606.08415
• PyTorch v2.15.0: Conv2d, Flatten, TransformerEncoderLayer, TransformerEncoder, MultiheadAttention.