DL S030 · Vision Transformer — Attention All the Way Down

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 30 / 80 · Module M05

The story

Story hook · the heresy that worked

In late 2020, a small team at Google Brain — Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Neil Houlsby and colleagues — quietly uploaded "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" to arXiv. The paper's central claim was heretical to the vision community:

"We show that this reliance on CNNs is not necessary and a pure transformer applied directly to sequences of image patches can perform very well on image classification tasks."

To appreciate how heretical that was, you have to understand what the vision field believed in 2020. Every serious classification, detection, or segmentation architecture used convolutions. Not by convention — because everyone was convinced they had to. The convolutional inductive bias (locality, translation equivariance, hierarchy) was considered a necessary prior for making vision work with feasible amounts of data and compute. Convnets had 30 years of theoretical justification (Fukushima 1980, LeCun 1989, Hubel & Wiesel 1962) and a decade of empirical dominance (AlexNet, VGG, ResNet, EfficientNet). "Just use attention" was the kind of proposal that a senior reviewer would reject in one sentence.

The ViT paper's headline results were split. On ImageNet-1k alone, ViT lost to a well-tuned ResNet. That was the expected outcome. But when the same ViT was pretrained on JFT-300M — Google's internal, never-released, 300-million-image dataset — and then fine-tuned on ImageNet-1k, it won by ~2% top-1. And more dramatically, ViT's accuracy kept climbing as they scaled data and model size, while ResNet's plateaued.

That scaling behaviour — the fact that ViT got hungrier and hungrier for data but rewarded you if you fed it — turned out to be the transformative insight. The 2020s in vision have been about removing inductive biases and replacing them with more data. The move from CNN to ViT is the same move as the move from RNN to Transformer in NLP, and for the same reason: attention scales; hand-crafted priors don't.

Within two years, the entire vision community had at least one ViT variant in production. DINOv2 (2023) made ViT the go-to backbone for dense prediction. Swin (2021) re-added a bit of the convolutional locality prior via windowed attention. ConvNeXt (2022) hit back with a modernised pure convnet and closed most of the gap. And Perceiver / Perceiver-IO (Jaegle et al. 2021) generalised the ViT idea to arbitrary modalities — audio, video, point clouds.

And there is a deeper unifying insight lurking here: ViT is not really "non-convolutional". The patch-embed layer at the front is literally nn.Conv2d(3, embed_dim, kernel_size=patch_size, stride=patch_size). It's a stride-16 convolution. The rest of the network is attention, but the first thing ViT does to a raw image is a very traditional convolutional projection. The border between CNN and Transformer is less a wall and more a spectrum. That is section 2 of this session.

Today we build the front end of a ViT from scratch, understand the three critical components (patch embedding, class token, positional embedding), and finish with the 2025 vision-language backbones that descended from this one paper.

The result that surprised everyone

In 2020, Google Brain published a paper titled "An Image is Worth 16x16 Words". Its claim: you don't need convolutions to do image classification. Take an image, chop it into 16×16 patches, flatten each patch into a vector, feed the sequence of vectors into a plain Transformer encoder — the same architecture used for NLP — and you get results competitive with ResNet.

At the time this was heretical. The entire field of computer vision had been built on the convolutional inductive bias: locality + translation equivariance were assumed to be necessary for images. ViT threw all of that out and worked anyway, provided you gave it enough data.

That last clause is the key. On ImageNet-1k alone, ViT loses to a well-tuned ResNet. On JFT-300M (Google's internal 300M-image dataset), ViT wins by ~2%. The lesson: the convolutional inductive bias saves data, and if you have enough data you don't need the bias.

Today we build the front end of a ViT — the part that turns an image into a sequence — and understand its three components: patch embedding, class token, positional embedding. The Transformer encoder that sits on top is the same one we'll build from scratch in Module M07.

Prerequisites

• S025–S029 (all of M05 — you need the CNN baseline to appreciate what ViT gives up).
• We'll use nn.TransformerEncoderLayer as a black box today; the internals get built from scratch in S036–S039.

1 · The recipe in one picture

That's the whole model. No convolution. No pooling. Everything is matmul.

2 · Patch embedding — two ways to write the same thing

The patchify step reshapes (B, 3, 224, 224) into (B, 196, 768). There are two idiomatic implementations, and it turns out they are identical.

2.1 The obvious way: unfold + linear

import torch
import torch.nn as nn
 
class PatchEmbeddingUnfold(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3, embed_dim=768):
        super().__init__()
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2   # 196
        self.proj = nn.Linear(patch_size * patch_size * in_ch, embed_dim)
 
    def forward(self, x):
        B, C, H, W = x.shape
        p = self.patch_size
        # unfold spatial dims into patches
        x = x.unfold(2, p, p).unfold(3, p, p)   # (B, C, H/p, W/p, p, p)
        x = x.permute(0, 2, 3, 1, 4, 5)          # (B, H/p, W/p, C, p, p)
        x = x.reshape(B, self.num_patches, -1)    # (B, 196, C·p·p)
        return self.proj(x)                       # (B, 196, embed_dim)

2.2 The clever way: a Conv2d with kernel = stride = patch_size

class PatchEmbeddingConv(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_ch, embed_dim,
                              kernel_size=patch_size, stride=patch_size)
 
    def forward(self, x):
        # (B, C, H, W) -> (B, embed_dim, H/p, W/p)
        x = self.proj(x)
        # (B, embed_dim, 14, 14) -> (B, 196, embed_dim)
        return x.flatten(2).transpose(1, 2)

Why are they equivalent? A Conv2d(3, 768, kernel=16, stride=16) applies 768 different 16×16×3 kernels at each of the 14×14 non-overlapping positions. That's exactly what "flatten each 16×16×3 patch and multiply by a 768×768 matrix" does — the 3072 weights per output channel are the 768 rows of the linear projection.

Use the Conv2d version in real code. It's faster (dispatched to cuDNN) and simpler to read.

3 · The class token — where does the prediction come from?

You now have a sequence of 196 patch embeddings, each (768,). What single vector do you feed to the classifier at the end?

Three options historically considered:

1. Average all 196. Works okay, ~0.5% worse than the alternative.
2. Take the first patch's embedding. Bad — that patch is a fixed spatial location.
3. Prepend a learnable "[CLS]" token. The winner.

The class token is a single trainable (1, D) vector prepended to every sequence:

self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
nn.init.trunc_normal_(self.cls_token, std=0.02)
 
def forward(self, x):
    B = x.size(0)
    cls = self.cls_token.expand(B, -1, -1)   # (B, 1, D)
    x = torch.cat([cls, x], dim=1)            # (B, 197, D)
    ...

The idea: through self-attention, the CLS token pulls information from every patch token. Its final embedding is a learned "summary" of the whole image, purpose-built for classification. The last layer feeds only output[:, 0] (the CLS position) to the classification head, ignoring the other 196 patch outputs.

4 · Positional embedding — because attention is permutation-invariant

Self-attention is order-agnostic — permute the input tokens and you get the same permuted output. That's fine for a set, but an image is a grid: patch (0, 0) is in the top-left corner and patch (13, 13) is in the bottom-right. The network needs to know this.

ViT's fix: add a learnable position embedding to every token.

self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))   # +1 for CLS
nn.init.trunc_normal_(self.pos_embed, std=0.02)
 
def forward(self, x):
    # ... after cls prepend, x is (B, 197, D)
    x = x + self.pos_embed
    return self.dropout(x)

Note the design choices:

• Learnable, not fixed. The original Transformer (Attention Is All You Need) used sinusoidal positional encodings. ViT went learnable because empirically it works about the same, and lets the network specialize per position. We'll compare both in S038.
• Added, not concatenated. Additive positional info integrates naturally with the token embeddings.
• Same dropout everywhere. Positional embeddings are dropout'd along with tokens.

4.1 The interpolation trick (for fine-tuning to different resolutions)

If you pretrain at 224×224 (196 patches + CLS = 197 positional embeddings) and want to fine-tune at 384×384 (576 patches + CLS = 577), the positional embedding table has the wrong number of rows. Standard fix: bilinear-interpolate the 14×14 grid of position embeddings to a 24×24 grid before initializing the new model. Practically all ViT fine-tuning pipelines do this transparently.

5 · Putting it together — a minimal ViT

class ViT(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3,
                 embed_dim=768, depth=12, num_heads=12,
                 mlp_ratio=4, num_classes=1000, dropout=0.1):
        super().__init__()
        self.patch_embed = PatchEmbeddingConv(img_size, patch_size, in_ch, embed_dim)
        num_patches = self.patch_embed.num_patches
 
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.dropout   = nn.Dropout(dropout)
 
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=num_heads,
            dim_feedforward=int(embed_dim * mlp_ratio),
            dropout=dropout, activation="gelu",
            batch_first=True, norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth)
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)
 
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
 
    def forward(self, x):
        x = self.patch_embed(x)                         # (B, N, D)
        cls = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.cat([cls, x], dim=1)                  # (B, N+1, D)
        x = self.dropout(x + self.pos_embed)
        x = self.encoder(x)                             # (B, N+1, D)
        x = self.norm(x[:, 0])                          # CLS token only
        return self.head(x)                             # (B, num_classes)

Parameter count for ViT-Base (D=768, depth=12, heads=12): ~86M. Compare with ResNet-50 at 25M. ViT is 3× bigger.

Two things worth flagging:

• norm_first=True. Pre-norm variant (LayerNorm before attention, not after). Trains more stably; standard since ~2020.
• activation="gelu". Not ReLU. GELU is smoother, marginally better in language and vision transformers. Every LLM uses it.

model = ViT(embed_dim=768, depth=12, num_heads=12, num_classes=1000)
for name, module in model.named_children():
    n = sum(p.numel() for p in module.parameters())
    print(f"{name:15s}  {n/1e6:7.2f}M")
total = sum(p.numel() for p in model.parameters())
print(f"total          {total/1e6:7.2f}M")

6 · Why ViT wins with data, loses without

The ViT paper's central figure shows accuracy vs. pretraining dataset size:

• ImageNet-1k (1.3M images): ResNet > ViT
• ImageNet-21k (14M images): ResNet ≈ ViT
• JFT-300M (300M images): ViT > ResNet by ~2%

Why? The convolutional inductive bias (locality, translation equivariance) is like free training data for a CNN — the architecture already assumes what a CNN would otherwise have to learn from examples. ViT has to learn all those properties from data. With enough data, ViT catches up and surpasses; below that threshold, CNNs use their prior more efficiently.

7 · What came after ViT (context, not required today)

• DeiT (2021): trained ViT-Base to competitive ImageNet-1k accuracy using clever distillation + augmentation. Killed the "ViT needs JFT" narrative.
• Swin Transformer (2021): re-introduced locality via windowed attention. Hybrid of ViT and hierarchical CNN. Won multiple 2021 benchmarks.
• DINOv2 (2023): self-supervised ViT that produces features rivaling supervised models. The current default backbone for "I need general-purpose visual features".
• SAM (2023): ViT-based promptable segmentation. Foundation model era for vision.

Every one of these builds on the same architectural skeleton we wrote today.

War stories

8 · The 2025 pulse · vision-language backbones eat everything

By 2023 the ViT story had already outgrown pure image classification. The interesting question stopped being "does ViT beat ResNet on ImageNet?" and became "can a single ViT learn a joint representation of images and text, so that everything downstream — classification, retrieval, captioning, VQA, robotics — falls out for free?"

Here's the current lineage you should know for 2025.

Foundation models that are just ViTs, trained self-supervised at massive scale:

• DINOv2 (Oquab et al., Meta AI 2023) — ViT-g/14 trained self-supervised (teacher-student distillation + iBOT masked-image prediction + Sinkhorn-Knopp centering) on 142M curated images. Frozen features rival supervised backbones on almost every downstream benchmark. Current best default for dense-prediction transfer.
• ViT-22B (Dehghani et al., Google 2023) — 22 billion parameters. Demonstrated that the transformer scaling laws that work for LLMs also work for vision at extreme scales. Not open, but the paper is essential reading for the recipe.
• Hiera (Ryali et al., FAIR 2023) — stripped-down hierarchical ViT (like Swin) pretrained with MAE. Faster and simpler than Swin, matches or beats it on every benchmark.
• EVA-02 (Fang et al. 2023) — combines MAE pretraining with CLIP-based feature distillation. Strong on both classification and segmentation.

Vision-language backbones (image + text jointly trained):

• CLIP (Radford et al., OpenAI 2021) — the paper that started it all. Contrastive image-text training on 400M pairs. Zero-shot ImageNet from a text prompt: "a photo of a {class}". Every multimodal system since 2021 traces back to CLIP.
• SigLIP (Zhai et al., DeepMind 2023) — replaces CLIP's softmax contrastive loss with a sigmoid loss. Scales to much larger batch sizes and better small-model performance.
• SigLIP-2 (Tschannen et al., DeepMind 2025) — the current SOTA open vision-language encoder. Adds decoder-based captioning objectives + self-distillation + native multilingual support (100+ languages) + a NaViT-style variable-resolution head. Powers the vision encoders inside Gemma 3 (2025), PaliGemma 2 (2024), and the open Molmo (Allen AI 2024) family.
• Perception Encoder (Meta AI, June 2025) — Meta's latest open ViT-based encoder, jointly trained on classification + captioning + dense prediction. Ships as a drop-in replacement for CLIP in Llama 4's vision path.

The 2025 practical decision tree for "what backbone do I pick?":

The Vision Transformer, seven years on: what started as a heretical claim in 2020 — "convolutions aren't necessary" — has evolved into the more nuanced modern truth: convolutions aren't necessary if you have enough data, but even when they're not necessary they're often the right first layer for computational reasons. Every SigLIP-2 encoder, every DINOv2 model, every ViT-22B still starts with Conv2d(3, D, kernel=stride=16) for the patch embed. The line between conv and attention is a technical detail; the real story is that flexible attention layers on top of a small conv stem now underpin nearly every state-of-the-art vision system in 2025.

That's the end of Module M05. You now have the entire modern vision toolkit in your head: convolution (S025), the architectural lineage from LeNet to ConvNeXt (S026), the training recipe that beats the architecture (S027), transfer learning and self-supervised pretraining (S028), the augmentation stack that adds 5% top-1 for free (S029), and the ViT recipe that unified images and language into one representation (this session).

Next up: we swap image patches for word tokens and start building the language side of the story.

Further reading · ViT and beyond in 2025

• Dosovitskiy et al. (2020) — An Image is Worth 16x16 Words. The paper. Read section 4 (experiments) carefully.
• Steiner et al. (2021) — How to train your ViT? Data, Augmentation, and Regularization. The training-recipe reference.
• Touvron et al. (2020) — Training data-efficient image transformers & distillation (DeiT). How to train ViT without JFT-300M.
• Oquab et al. (2023) — DINOv2.
• Tschannen et al. (2025) — SigLIP-2.
• Dehghani et al. (2023) — Scaling Vision Transformers to 22 Billion Parameters.
• Umar Jamil — Coding a Vision Transformer from Scratch (YouTube). The best code-along walkthrough on the internet.

Retention scaffold

One-line summary (write it in your own words): ______________________________________

Spaced review: re-derive the ViT forward pass end-to-end in 24 hours — patchify → embed → CLS → pos → encoder → classify — without opening the file. On day 7, load timm.create_model('vit_base_patch16_224', pretrained=True), inspect .patch_embed, and confirm it's a stride-16 Conv2d.

Next session (S031): Word2Vec — swap image patches for word tokens. The trick that made every modern LLM possible: represent every word as a learnable dense vector.

Sticky note:

Patch → Linear → [CLS + Pos] → Encoder → Classify from CLS. Patch embed = stride-P Conv2d. ViT wins with data, loses without. DINOv2 + SigLIP-2 are the 2025 defaults.

Legacy recall drills

Retention scaffold

Recall

Stretch

Implement PatchEmbeddingUnfold and PatchEmbeddingConv and verify they produce the same output up to weight initialization. Then load a pretrained timm ViT and inspect its patch_embed layer — it will be nn.Conv2d, confirming the equivalence.

In your own words

Explain the inductive-bias / data tradeoff between CNNs and ViTs in 3 sentences.

Spaced review

• S025 (convolution) — ViT's patch embed is one big stride-16 conv. Understanding this is understanding why the two paradigms aren't as far apart as they seem.
• S021 (LayerNorm vs BatchNorm) — ViT uses LayerNorm exclusively; there's no BN. We'll revisit why in S036–S038.

Next session

S031 — Word2Vec. We swap image patches for word tokens and start building up the language side of the story. The core trick — represent every word as a learnable dense vector — is what makes modern NLP possible.

Bring back tomorrow

• ViT recipe: patch → linear project → CLS + pos embed → Transformer encoder → classify from CLS.
• Patch embed = Conv2d(kernel=stride=patch_size).
• ViT wins with data, loses without.

Previous: ← DL S029 · Next: DL S031 →