Lesson 92: Vision Transformer (ViT)

1. Vision Transformer ViT from Scratch

Title

USAAIO · Lesson 92 · Phase 3

Patches as tokens: convert an image into a sequence, pass it through a standard transformer encoder, classify on the CLS token. Every component built and verified in PyTorch — from patch embedding to attention maps.

2. By the end of this lesson you can

Objectives

1. Compute N patches and sequence length for any image size and patch size P
2. Implement patch embedding using nn.Conv2d with stride=P and explain why it is equivalent to flattening + linear projection
3. Describe the role of the CLS token and learned positional embeddings in a ViT
4. Trace the full ViT forward pass (patch embed → prepend CLS → add pos embed → transformer encoder → CLS head) with concrete tensor shapes
5. State when ViT outperforms a CNN and when it does not, and explain the data-scaling reason

3. What survived from Flash Attention & Memory-Efficient Transformers?

Warm-up

4. Patches as tokens — the core idea

Section

Part 1 of 4

5. An image is worth N×P² words

Concept

A standard transformer (Lesson 88) expects a sequence of vectors. ViT (Dosovitskiy 2021) creates that sequence by chopping the image into non-overlapping P×P patches and treating each flattened patch as a token.

\[ N = \left(\frac{H}{P}\right)^2 \quad\text{(square image } H{\times}H\text{, patch size }P\text{)} \]

6. Fill in: P for An image is worth N×P² words

Comparison

7. Patch embedding: flatten → project

Concept

Each P×P×C patch is flattened to a vector of dimension P²·C, then linearly projected to the model dimension d_model. Two equivalent implementations:

• Naive: slice N patches, call nn.Linear(P²·C, d_model) — easy to understand, O(N) Python loops
• Efficient: nn.Conv2d(C, d_model, kernel_size=P, stride=P) — one pass over the image, all patches in parallel; output shape (B, d_model, H/P, W/P) which you then flatten and transpose

The Conv2d approach is mathematically identical to flattening + nn.Linear (same matrix multiply, just batched spatially). stride=P ensures non-overlapping patches.

8. Break it if you can: Patch embedding: flatten → project

Counterexample

9. Guess the shape of the answer: Patch embedding forward pass — tensor shapes

Estimation

10. Patch embedding forward pass — tensor shapes

Worked example

Trace the patch embed layer on a batch of 4 images of size (1, 8, 8) with P=4, d_model=32. Predict the output shape at each step before reading on.

import torch, torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size, patch_size, in_chans, d_model):
        super().__init__()
        self.proj = nn.Conv2d(
            in_chans, d_model,
            kernel_size=patch_size, stride=patch_size
        )
    def forward(self, x):
        x = self.proj(x)         # (B, d_model, H/P, W/P)
        x = x.flatten(2)         # (B, d_model, N)
        x = x.transpose(1, 2)    # (B, N, d_model)
        return x

torch.manual_seed(42)
patch_embed = PatchEmbedding(img_size=8, patch_size=4, in_chans=1, d_model=32)
x = torch.randn(4, 1, 8, 8)
out = patch_embed(x)
print(out.shape)  # torch.Size([4, 4, 32])

11. What each one costs: Patch embedding forward pass — tensor shapes

Trade off

12. CLS token & positional embeddings

Section

Part 2 of 4

13. The CLS token

Concept

A ViT prepends a learnable [CLS] token to the patch sequence before the transformer encoder. It has no corresponding image patch — its role is to accumulate global information through attention.

After the encoder, only the CLS token's output vector is passed to the classification head. This is the same design as BERT (Lesson 90): the sequence is [CLS, patch_1, patch_2, ..., patch_N].

\[ \mathbf{z}_0 = [\mathbf{x}_{\text{cls}};\; E\mathbf{p}_1;\; E\mathbf{p}_2;\; \ldots;\; E\mathbf{p}_N] + \mathbf{E}_{\text{pos}} \]

14. By analogy: The CLS token

Analogy

15. Learned positional embeddings

Concept

Unlike the sinusoidal scheme, ViT uses learned positional embeddings — a parameter matrix of shape (N+1, d_model) added to the token sequence. They are trained end-to-end and encode 2D position implicitly.

• Shape: (1, N+1, d_model) — broadcast over the batch
• Initialized near zero (trunc_normal_(std=0.02)), learned during training
• Visualizing them after training reveals a 2D spatial structure — rows/columns of the patch grid

16. Guess the shape of the answer: CLS prepend + positional embedding

Estimation

17. CLS prepend + positional embedding

Worked example

Continue from the patch embed output (4, 4, 32). Add a CLS token and positional embedding. Trace shapes at each step.

import torch, torch.nn as nn

B, N, d = 4, 4, 32
patch_out = torch.randn(B, N, d)          # simulated patch embed output

cls_token = nn.Parameter(torch.zeros(1, 1, d))
pos_embed  = nn.Parameter(torch.zeros(1, N + 1, d))  # N+1 for CLS
nn.init.trunc_normal_(cls_token, std=0.02)
nn.init.trunc_normal_(pos_embed, std=0.02)

cls = cls_token.expand(B, -1, -1)         # (4, 1, 32)
x   = torch.cat([cls, patch_out], dim=1)  # (4, 5, 32)
x   = x + pos_embed                       # broadcast add
print(x.shape)   # torch.Size([4, 5, 32])

18. Which is which, by shape

Discrimination

19. Something is wrong here: CLS token at position 0 vs appended at the end

Anomaly

20. Trap: CLS token at position 0 vs appended at the end

Trap

21. Break it on purpose: CLS token at position 0 vs appended at the…

Break the constraint

22. Transformer encoder on patch sequences

Section

Part 3 of 4

23. Standard transformer encoder (recap Lesson 88)

Concept

ViT reuses the standard transformer encoder block unchanged from Lesson 88 — no modifications needed. Each block: LayerNorm → MultiHeadAttention → residual → LayerNorm → MLP → residual.

• Pre-norm variant: ViT uses LN before attention, not after (unlike the original 'Attention is All You Need')
• MLP: two linear layers with GELU activation, hidden dim = 4 × d_model
• No causal mask: vision is not autoregressive — every patch attends to every other patch freely
• Depth: ViT-Base uses 12 blocks, ViT-Large uses 24; our toy uses 2

24. Teach it back: Standard transformer encoder (recap Lesson 88)

Explain it

25. Guess the shape of the answer: TinyViT8 — full forward pass with shapes

Estimation

26. TinyViT8 — full forward pass with shapes

Worked example

Build TinyViT8 (8×8 images, P=4, d_model=32, 2 heads, 2 encoder blocks, 10 classes). This is a runnable ViT — no pretrained weights, no downloads.

import torch, torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d, n_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(d)
        self.attn  = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d)
        mlp_d = d * 4
        self.mlp = nn.Sequential(
            nn.Linear(d, mlp_d), nn.GELU(), nn.Linear(mlp_d, d)
        )
    def forward(self, x):
        a, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
        x = x + a
        x = x + self.mlp(self.norm2(x))
        return x

class TinyViT8(nn.Module):
    def __init__(self):
        super().__init__()
        d, P, n_heads = 32, 4, 2
        self.patch_proj = nn.Conv2d(1, d, kernel_size=P, stride=P)
        n_p = (8 // P) ** 2                       # 4 patches
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d))
        self.pos_embed  = nn.Parameter(torch.zeros(1, n_p + 1, d))
        self.blocks = nn.ModuleList([TransformerBlock(d, n_heads) for _ in range(2)])
        self.norm = nn.LayerNorm(d)
        self.head = nn.Linear(d, 10)
    def forward(self, x):
        B = x.shape[0]
        x = self.patch_proj(x).flatten(2).transpose(1, 2)  # (B,4,32)
        x = torch.cat([self.cls_token.expand(B,-1,-1), x], 1) + self.pos_embed
        for blk in self.blocks:
            x = blk(x)
        return self.head(self.norm(x)[:, 0])

torch.manual_seed(42)
m = TinyViT8()
print(sum(p.numel() for p in m.parameters()), 'params')
print(m(torch.randn(4, 1, 8, 8)).shape)

27. Work backwards from the answer: TinyViT8 — full forward pass with shapes

Reverse engineer

28. Something is wrong here: forgetting to index CLS after the encoder

Anomaly

29. Trap: forgetting to index CLS after the encoder

Trap

30. Which of these survive contact with Lesson 92: Vision Transformer (ViT)?

Two truths and a lie

31. ViT vs CNN — when to use which

Section

Part 4 of 4

32. ViT needs more data — but scales better

Concept

CNNs have strong inductive biases baked in: translation equivariance (Lesson 55) and locality. A CNN achieves good accuracy on small datasets because the architecture already knows 'nearby pixels matter more.'

ViT has no such inductive bias — it treats all patch pairs equally and must learn spatial structure from scratch. This means ViT needs large datasets (JFT-300M, ImageNet-21k) to match CNN performance.

33. Fill in: 100M+ images for ViT needs more data — but scales better

Comparison

34. Fine-tuning vs training from scratch

Concept

In practice, ViT models are pretrained on massive datasets (ImageNet-21k or JFT) and then fine-tuned on the target task — the same transfer-learning workflow as BERT (Lesson 90).

• Fine-tuning: freeze early layers or lower lr; replace the classification head; converges in ~10-30 epochs
• Training from scratch: requires 300+ epochs, large datasets, strong augmentation (RandAugment, MixUp, CutMix)
• Key difference from CNNs: positional embeddings must be interpolated when fine-tuning at a different resolution than pretraining

35. By analogy: Fine-tuning vs training from scratch

Analogy

36. Without one step: The ViT architecture recipe

Constraint

37. The ViT architecture recipe

Pattern

1. Patch count: N = (H/P)²; sequence length = N+1 (CLS prepended). For 224×224, P=16 → N=196.
2. Patch embed: nn.Conv2d(C, d_model, kernel_size=P, stride=P) → flatten(2).transpose(1,2) → shape (B, N, d_model)
3. CLS prepend: torch.cat([cls_token.expand(B,-1,-1), patches], dim=1) → (B, N+1, d_model)
4. Positional embed: add learned (1, N+1, d_model) parameter — broadcast over batch
5. Encoder: L × [PreNorm → MHA → residual → PreNorm → MLP(4d) → residual]; no causal mask
6. Head: norm(x)[:, 0] → Linear(d_model, n_classes) — index 0 = CLS token output
7. Data rule: ViT lags CNNs on < ~10k images; matches/beats with 100M+ samples or pretrained fine-tuning

38. Where does it stop working: The ViT architecture recipe

Edge cases

39. Rule out three: Check yourself — patch sequence length

Elimination

40. Check yourself — patch sequence length

Check

Work it out before clicking.

41. Answer it before you see the options: Check yourself — patch embed equivalence

Prediction

42. Check yourself — patch embed equivalence

Check

Think about what Conv2d with stride=P computes.

43. Rule out three: Check yourself — ViT vs CNN data scaling

Elimination

44. Check yourself — ViT vs CNN data scaling

Check

Apply the inductive bias reasoning.

45. Your turn: TinyViT8 from scratch

Section

Project

46. Project: train TinyViT8 on load_digits

Concept

Build TinyViT8 — a minimal but structurally correct ViT — on load_digits (1,797 samples, 8×8 grayscale, 10 classes). Three milestones: patch shape tracing → full forward pass → training loop.

Build rules: type every line; print tensor shapes after each operation; verify the CLS index is 0 before classification; use torch.manual_seed(42) for reproducibility.

47. Break it if you can: Project: train TinyViT8 on load_digits

Counterexample

48. Milestone 1 — patch embed shape trace

Worked example

Your turn: implement PatchEmbedding for 8×8 images with P=4, d=32. Predict the output shape before running. What is out_size = (8 - 4) // 4 + 1?

Hint: Conv2d with kernel_size=P, stride=P produces (B, d_model, H/P, W/P). Then flatten(2) merges spatial dims, transpose(1,2) swaps channels and sequence.

import torch, torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size, patch_size, in_chans, d_model):
        super().__init__()
        self.proj = nn.Conv2d(
            in_chans, d_model, kernel_size=patch_size, stride=patch_size
        )
    def forward(self, x):
        x = self.proj(x)          # (B, d_model, H/P, W/P)
        x = x.flatten(2)          # (B, d_model, N)
        return x.transpose(1, 2)  # (B, N, d_model)

torch.manual_seed(42)
pe = PatchEmbedding(8, 4, 1, 32)
x = torch.randn(4, 1, 8, 8)
out = pe(x)
print(out.shape)   # torch.Size([4, 4, 32])

49. Milestone 2 — CLS + positional embedding

Worked example

Your turn: prepend the CLS token and add positional embeddings. Predict the final sequence shape before running. Why must pos_embed have shape (1, N+1, d_model) and not (1, N, d_model)?

Hint: cls_token.expand(B, -1, -1) broadcasts the (1,1,d) parameter to (B,1,d) without copying. torch.cat([cls, patches], dim=1) then gives (B, N+1, d).

import torch, torch.nn as nn

B, N, d = 4, 4, 32
patches = torch.randn(B, N, d)

cls_token = nn.Parameter(torch.zeros(1, 1, d))
pos_embed  = nn.Parameter(torch.zeros(1, N + 1, d))
nn.init.trunc_normal_(cls_token, std=0.02)
nn.init.trunc_normal_(pos_embed, std=0.02)

cls  = cls_token.expand(B, -1, -1)        # (4, 1, 32)
seq  = torch.cat([cls, patches], dim=1)   # (4, 5, 32)
seq  = seq + pos_embed                    # (4, 5, 32)
print(seq.shape, '  N+1 =', N+1)

50. Which is which, by shape

Discrimination

51. Milestone 3 — train to >95% accuracy

Worked example

Your turn: complete the training loop for 150 epochs with Adam lr=3e-4. Predict how quickly ViT converges compared to the CNN from Lesson 55.

Hint: use the same five-step loop as Lesson 40 (zero_grad → forward → loss → backward → step). A lower learning rate (3e-4 vs 1e-3) is typical for transformers — they are sensitive to lr.

import torch, torch.nn as nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import numpy as np

digits = load_digits()
X = digits.data.reshape(-1,1,8,8).astype(np.float32)/16.0
y = digits.target.astype(np.int64)
Xtr,Xte,ytr,yte = train_test_split(X,y,test_size=0.2,random_state=42)
Xt=torch.tensor(Xtr); yt=torch.tensor(ytr)
Xtt=torch.tensor(Xte); ytt=torch.tensor(yte)

torch.manual_seed(42)
model = TinyViT8()           # defined in earlier milestone
opt   = torch.optim.Adam(model.parameters(), lr=3e-4)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(150):
    opt.zero_grad()
    loss = loss_fn(model(Xt), yt)
    loss.backward(); opt.step()
with torch.no_grad():
    acc=(model(Xtt).argmax(1)==ytt).float().mean()
print(f'test acc: {acc.item()*100:.1f}%')  # 96.4%

52. What each one costs: Milestone 3 — train to >95% accuracy

Trade off

53. The full program

Concept

import torch, torch.nn as nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import numpy as np

class PatchEmbedding(nn.Module):
    def __init__(self, img_size, patch_size, in_chans, d):
        super().__init__()
        self.proj = nn.Conv2d(in_chans, d, kernel_size=patch_size, stride=patch_size)
    def forward(self, x):
        return self.proj(x).flatten(2).transpose(1, 2)

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.n1 = nn.LayerNorm(d); self.n2 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, h, batch_first=True)
        self.mlp  = nn.Sequential(nn.Linear(d, d*4), nn.GELU(), nn.Linear(d*4, d))
    def forward(self, x):
        a,_ = self.attn(self.n1(x), self.n1(x), self.n1(x))
        x = x + a; x = x + self.mlp(self.n2(x)); return x

class TinyViT8(nn.Module):
    def __init__(self, d=32, n_heads=2, n_layers=2, n_cls=10):
        super().__init__()
        self.pe = PatchEmbedding(8, 4, 1, d)
        self.cls = nn.Parameter(torch.zeros(1,1,d))
        self.pos = nn.Parameter(torch.zeros(1,5,d))   # 4 patches + CLS
        self.blocks = nn.ModuleList([Block(d, n_heads) for _ in range(n_layers)])
        self.norm = nn.LayerNorm(d); self.head = nn.Linear(d, n_cls)
        nn.init.trunc_normal_(self.cls, std=0.02); nn.init.trunc_normal_(self.pos, std=0.02)
    def forward(self, x):
        B=x.shape[0]; x=self.pe(x)
        x=torch.cat([self.cls.expand(B,-1,-1),x],1)+self.pos
        for b in self.blocks: x=b(x)
        return self.head(self.norm(x)[:,0])

digits=load_digits(); X=digits.data.reshape(-1,1,8,8).astype(np.float32)/16.0; y=digits.target.astype(np.int64)
Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=0.2,random_state=42)
Xt=torch.tensor(Xtr); yt=torch.tensor(ytr)
torch.manual_seed(42); m=TinyViT8()
opt=torch.optim.Adam(m.parameters(),lr=3e-4); L=nn.CrossEntropyLoss()
for _ in range(150):
    opt.zero_grad(); L(m(Xt),yt).backward(); opt.step()
with torch.no_grad():
    acc=(m(torch.tensor(Xte)).argmax(1)==torch.tensor(yte)).float().mean()
print(f'{acc.item()*100:.1f}%')  # 96.4%

54. Fill in: value for The full program

Comparison

55. Show it off

Concept

Out loud, slides closed: (1) explain the full ViT forward pass from raw image pixels to logits, naming every tensor shape at each stage; (2) explain why the CLS token is at index 0 and not the end; (3) state the inductive bias argument for why ViT underperforms CNN on small datasets.

Stretch (homework from the lesson plan): visualize attention maps — after training, extract attn_weights from nn.MultiheadAttention using need_weights=True, reshape to (N+1, N+1), and inspect which patches the CLS token attends to. Compare ViT vs the SmallCNN from Lesson 55 at 1k vs 10k training examples.

56. Connect it up: Lesson 92: Vision Transformer (ViT)

Connect it up

57. What you can do now

Recap

• Compute N = (H/P)² patches and sequence length N+1 for any ViT configuration
• Implement patch embedding as nn.Conv2d(C, d, k=P, s=P) → flatten(2).transpose(1,2)
• Prepend the CLS token and add learned positional embeddings before the encoder
• Build a full ViT forward pass using nn.MultiheadAttention + pre-norm MLP blocks
• State the data-scaling argument: ViT needs large data to match CNNs; fine-tuning from pretrained weights closes the gap