Vision Transformer — Module 10.2

How This Lesson Fits the Module & Volume

This is the capstone of Volume 10: the same encoder stack you built from Module 10.1 primitives and Module 10.2 blocks now reads images as sequences of patches. The Vol. 07 ViT lecture introduced the vision angle; here we wire it explicitly to MHSA, positional embeddings, skips, and LayerNorm.

After this, Volume 11 begins with Language Models—decoder-style Transformers at scale. ViT proves the architecture is modality-agnostic: tokens can be words or patches.

Learning Objectives

By the end of this lesson, students should be able to:

Describe how an image is split into patches and linearly embedded.
Explain the role of the class ([CLS]) token and positional embeddings in ViT.
Recognize ViT as a standard Transformer encoder stack on patch tokens.
Implement a minimal ViT encoder forward pass in PyTorch.
Compare CNN inductive bias with ViT’s attention-over-patches approach.
Bridge from Vol. 10 architecture fluency to Vol. 11 language models.

Definition

A Vision Transformer (ViT) treats an image as a sequence of flattened patches, maps each patch to a d_model vector, prepends a learnable class token, adds positional embeddings, and runs a Transformer encoder. The class token’s final state (or pooled patch states) feeds an MLP classification head.

From Pixels to Tokens

Patchify

Split image into P×P cells.

Embed

Linear map each patch.

+ CLS & PE

Class token + positions.

Encoder

MHSA + FFN + skips.

Piece	Role	Vol. 10 link
Patch embedding	Pixels → token vectors	Like token embedding, different input
Positional embedding	2D layout as 1D indices	Positional Embedding
MHSA encoder blocks	Global patch mixing	Encoder Block
Skip + LayerNorm	Deep stable stack	Skip, LayerNorm
MLP head	Class logits from CLS	Task head (cf. full stack)

CNN vs ViT Inductive Bias

CNNs (Vol. 07)

Local receptive fields + weight sharing.
Strong bias for natural images.
Data-efficient on smaller sets.

ViT

Global attention from layer one.
Flexible but data-hungry.
Scales superbly with pretraining.

Shared lesson

Representation learning + heads.
Transfer learning remains key.
Hybrids (Conv stem + ViT) exist.

Code: Minimal ViT Encoder Sketch

import torch
from torch import nn

class MiniViT(nn.Module):
    def __init__(self, img_size=32, patch=8, d_model=64, nhead=4, depth=2, num_classes=10):
        super().__init__()
        assert img_size % patch == 0
        self.n_patches = (img_size // patch) ** 2
        self.patch_embed = nn.Conv2d(3, d_model, kernel_size=patch, stride=patch)
        self.cls = nn.Parameter(torch.zeros(1, 1, d_model))
        self.pos = nn.Parameter(torch.zeros(1, 1 + self.n_patches, d_model))
        layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=4*d_model, batch_first=True
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=depth)
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, num_classes)

    def forward(self, x):
        # x: (B, 3, H, W)
        x = self.patch_embed(x)              # (B, d_model, H/P, W/P)
        x = x.flatten(2).transpose(1, 2)     # (B, N, d_model)
        cls = self.cls.expand(x.size(0), -1, -1)
        x = torch.cat([cls, x], dim=1) + self.pos
        x = self.encoder(x)
        return self.head(self.norm(x[:, 0]))  # classify from CLS

model = MiniViT()
imgs = torch.randn(4, 3, 32, 32)
print(model(imgs).shape)  # torch.Size([4, 10])
print(model.n_patches)    # 16 patches for 32x32 with patch=8

Strengths and Tradeoffs

Strengths

Same Transformer recipe as NLP—shared tooling and intuition.
Global context without deep CNN stacks.
Excellent scaling with data and compute.

Tradeoffs

Weaker locality bias; needs more data or augmentation.
Quadratic cost in number of patches (resolution).
Patch size trades detail vs sequence length.

Common Misconception

“ViT replaces attention with convolutions.” ViT replaces the CNN backbone with a Transformer encoder; the core operation remains self-attention over patch tokens. Convolutions may still appear in patch embedding (as a strided conv) or hybrid stems, but the mixer is attention.

Bridge to Volume 11

You now have the full Vol. 10 story: attention primitives (10.1) → Transformer architecture (10.2) → vision as one modality. Volume 11 specializes the decoder-only path: tokenization, next-token prediction, sampling, and inference tricks like the KV cache. The blocks are the ones you already know—masked MHSA, FFN, skips, LayerNorm—scaled into language models.

Knowledge Check

Short Answer: How does ViT turn an image into a sequence? Answer: Split into patches, flatten/embed each patch as a token.
True/False: ViT typically uses a Transformer encoder stack (bidirectional self-attention over patches). Answer: True.
Multiple Choice: The class token is: (a) a JPEG marker, (b) a learnable vector prepended to patch tokens, (c) the softmax temperature. Answer: (b).
Short Answer: Why add positional embeddings to patches? Answer: So the model knows where each patch sits in the image grid.
True/False: ViT and BERT use unrelated attention mathematics. Answer: False—same MHSA family.
Multiple Choice: Compared with CNNs, vanilla ViTs often need: (a) less data, (b) more data/pretraining, (c) no positions. Answer: (b).
Short Answer: Which Module 10.2 lecture covers the encoder block ViT stacks? Answer: Encoder Block.
Short Answer: Name one Vol. 07 resource for the vision-first ViT view. Answer: volumes/vol-07/module-7-3/vision-transformer.html.
Multiple Choice: Volume 11 primarily builds on: (a) only CNNs, (b) language modeling with Transformer decoders, (c) k-means. Answer: (b).
True/False: This lecture caps Volume 10’s Attention & Transformers arc. Answer: True.

Key Takeaways

ViT = patch embed + CLS + positional embeddings + Transformer encoder + head.
It reuses Module 10.2 blocks; only the tokenizer of the image changes.
Global attention trades CNN locality bias for scalability.
Vol. 07 covers ViT in the vision track; Vol. 10 ties it to the attention stack.
Next volume: Language Model (Vol. 11).

Trainer’s Guide

Hands-on idea: On CIFAR-scale tensors, print patch count vs patch size; discuss the O(N²) attention cost as resolution grows.

Discussion prompt: Capstone synthesis—draw one diagram that shows QKV attention serving both a translation decoder and a ViT encoder.

Recap: Vision Transformers prove the Module 10.2 stack is modality-agnostic. Continue to Vol. 11 Language Model.