Vision Transformer from scratch — ViT vs CNN on Fashion-MNIST

# Description

Vision Transformer (Dosovitskiy et al., 2020) implemented from scratch in PyTorch on Fashion-MNIST — no transformer libraries: PatchEmbed (Conv2d with stride=patch: the 28×28 image is split into 49 4×4 patches), learned [CLS] token, learned positional embeddings (50×64, trunc_normal init), Multi-Head Self-Attention (4 heads, head_dim 16, dropout 0.1) and Pre-LN Transformer blocks (norm → attention → residual; norm → GELU MLP → residual), 105,546 params. Stratified split 3,200 train / 800 val / 10,000 test (seed 42), train-only augmentation (horizontal flip + affine 5°/6%), AdamW lr 1e-3 wd 1e-4 with cosine annealing, early stopping patience 5, batch 128. Honest baseline: classic CNN with comparable params (117,498 params, Conv-BN-ReLU 32→80→128 + GlobalAvgPooling, no dense layers). Test results: CNN 0.8031 vs ViT 0.7180 (Δ −8.51 pp) — with only 3,200 samples the convolutional inductive bias (locality + translation equivariance) beats global attention: transformers are data-hungry (classic finding: ImageNet-21k → fine-tune). The ViT still learns with no locality prior: attention rollout (Abnar & Zuidema 2020, R = ∏(0.5·A + 0.5·I)) shows the [CLS] token focusing on discriminative regions (trouser waist, sweater hem, boot upper); positional embeddings learn implicit 2D structure (cosine similarity decays with distance); the 4 heads of the middle block specialize on different regions; t-SNE of CLS embeddings clusters by class. Both architectures fail on Shirt (0.184 vs 0.315 — overlaps with T-shirt, Pullover and Coat), the hardest Fashion-MNIST class. 8 visualizations.

$ ./deploy.sh $ git clone

# Key features

$ Complete ViT from scratch in PyTorch: PatchEmbed (Conv2d stride=patch), [CLS] token, learned positional embeddings and Pre-LN Transformer blocks — no transformer libraries

$ CPU-first architecture: 4×4 patches → 49 tokens, dim 64, 3 blocks, 4 heads, dropout 0.1 — 105,546 params

$ Honest baseline: classic CNN with comparable params (117,498 params, Conv-BN-ReLU 32→80→128 + GlobalAvgPooling)

$ Fashion-MNIST with stratified split (3,200/800/10,000), train-only augmentation, AdamW + cosine annealing + early stopping

$ CNN wins on small data (test 80.3% vs 71.8%, Δ −8.5 pp): convolutional inductive bias beats global attention with 3,200 samples — ViTs are data-hungry

$ Attention Rollout (Abnar & Zuidema 2020): the [CLS] token attends to discriminative regions — trouser waist, sweater hem, boot upper

$ Positional embeddings learn implicit 2D structure (cosine similarity) and middle-block heads specialize on different regions

$ 8 visualizations: patches, pos-embeddings, ViT vs CNN curves, attention rollout, CLS t-SNE, attention heads, summary and confusion matrices

# Gallery

Project terminal