14. Vision Transformers

Vision Transformers (ViT)

In the Transformers class we built an architecture for sequences of text tokens. In 2021, Dosovitskiy et al.¹ asked a provocative question: what if we feed an image to that exact same encoder? Their answer — "An Image is Worth 16×16 Words" — showed that, with enough data, a near-vanilla Transformer can beat Convolutional Neural Networks at image classification, without any convolutions at all.

The ViT is the bridge between the convolutional world (class 11) and the Transformer world (class 13). Once you understand it, the image encoders inside CLIP, Stable Diffusion and Diffusion Transformers stop being black boxes.

The inductive bias trade-off

A CNN comes with two strong inductive biases baked into its architecture:

• Locality — a convolutional kernel only looks at a small neighborhood of pixels.
• Translation equivariance — the same filter slides across the whole image, so a feature is detected regardless of where it appears.

These priors are exactly why CNNs learn so efficiently from small datasets: the architecture already "knows" that nearby pixels matter and that objects can move around.

A pure self-attention layer has none of these priors. Every patch can attend to every other patch from layer 1, so the model must learn spatial relationships from data. This is a double-edged sword:

The ViT pipeline, step by step

The trick is to turn an image into a sequence of tokens so the Transformer encoder we already know can consume it. Step through the pipeline below:

Patch embedding — the only genuinely new piece

Everything after the first step is the encoder you already know. The one new mechanism is how the image becomes tokens.

An image \(x \in \mathbb{R}^{H \times W \times C}\) is reshaped into a sequence of \(N\) flattened patches, then each patch is projected by a single shared linear layer:

where \(P\) is the patch size, \(N = HW/P^2\) is the number of patches, and \(E \in \mathbb{R}^{(P^2 C) \times d}\) is the patch embedding matrix. For a \(224 \times 224\) image with \(P = 16\), that is \(N = 196\) tokens.

A learnable [CLS] token \(z_{\text{cls}}\) is prepended, and learnable positional embeddings \(E_{\text{pos}}\) are added (attention alone is permutation-invariant, so without positions the model could not tell a patch in the top-left from one in the bottom-right):

Note. A linear projection over non-overlapping \(P \times P\) patches is mathematically identical to a Conv2d with kernel_size = stride = P. That is exactly how it is implemented in practice — one convolution does the patchifying and the projection in a single op.

The encoder and the classification head

The sequence \(z_0\) goes through \(L\) identical Transformer encoder blocks⁴ — the same MHSA + Add&Norm + FFN blocks from the Transformers class (ViT uses the pre-norm variant and GELU):

For classification, only the final state of the [CLS] token is read out and passed through a small MLP head:

That is the whole model. No convolutions, no pooling pyramids — just patchify, then a standard Transformer encoder.

Data hunger: why ViT needs pretraining

Because it lacks the convolutional priors, ViT only shines at scale. The original paper made this concrete:

This is precisely the pretrain-then-finetune recipe of the Transfer Learning class: pretrain the encoder on a huge dataset, then fine-tune the cheap MLP head (or the whole model with a low learning rate) on your task. It is also why CLIP could train a ViT image encoder on 400M image-text pairs — at that scale, the weak inductive bias becomes an advantage.

Two lines of work soften this data hunger. Careful augmentation and regularization recipes let ViTs train competitively on ImageNet-1k alone, without a giant private dataset⁵. And self-supervised pretraining learns strong ViT features without any labels — by reconstructing masked patches (MAE⁶) or by self-distillation (DINO⁷, whose attention maps segment objects for free).

CNN vs. ViT at a glance

Hybrids and successors. Several variants reintroduce some spatial bias to get the best of both worlds: DeiT² (data-efficient training with distillation, no JFT needed), and Swin Transformer³ (windowed attention with a hierarchical, pyramid-like structure that brings back locality and makes ViT practical for detection and segmentation). Conversely, ConvNeXt⁸ modernized a pure CNN to match ViTs head-to-head — evidence that the training recipe and scale matter as much as the convolution-vs-attention choice itself.

Implementation reference

import torch
import torch.nn as nn

class PatchEmbed(nn.Module):
    """Image -> sequence of patch tokens (Conv2d does patchify + projection)."""
    def __init__(self, img_size=224, patch=16, in_ch=3, dim=768):
        super().__init__()
        self.n_patches = (img_size // patch) ** 2
        self.proj = nn.Conv2d(in_ch, dim, kernel_size=patch, stride=patch)

    def forward(self, x):                  # x: (B, C, H, W)
        x = self.proj(x)                   # (B, dim, H/p, W/p)
        return x.flatten(2).transpose(1, 2)  # (B, N, dim)

class ViT(nn.Module):
    def __init__(self, dim=768, depth=12, heads=12, n_classes=1000):
        super().__init__()
        self.patch_embed = PatchEmbed(dim=dim)
        n = self.patch_embed.n_patches
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos = nn.Parameter(torch.zeros(1, n + 1, dim))
        layer = nn.TransformerEncoderLayer(dim, heads, dim * 4,
                                           activation='gelu', norm_first=True,
                                           batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, depth)
        self.head = nn.Linear(dim, n_classes)

    def forward(self, x):
        B = x.size(0)
        x = self.patch_embed(x)                       # (B, N, dim)
        cls = self.cls.expand(B, -1, -1)
        x = torch.cat([cls, x], dim=1) + self.pos     # prepend CLS + add positions
        x = self.encoder(x)
        return self.head(x[:, 0])                      # classify from CLS token

import timm
# load ImageNet-21k -> 1k pretrained weights and fine-tune the head
model = timm.create_model('vit_base_patch16_224', pretrained=True,
                          num_classes=10)