Vision Transformers

Images are not obviously sequences. They’re 2D grids of pixels, not tokens in a sentence. CNNs were designed for exactly this structure — convolutional filters exploit translation invariance, pooling builds hierarchical representations, local receptive fields match the local structure of natural images. ViT discards all of that by asking a simpler question: what if we just flop the image into a sequence of patches and feed it to a transformer? The answer turned out to be: it works, badly at small data scale, extraordinarily well at large data scale. Knowing which regime you’re in is the only thing that matters for deciding which architecture to use.

---

1. Vision Transformer — Core Idea

Paper: Dosovitskiy et al. (2020), arXiv:2010.11929 — “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale”

The Patch Embedding Step

Take an image of shape (H, W, C) — say (224, 224, 3). Split it into a grid of non-overlapping P×P patches, where P=16:

Number of patches: (224/16) × (224/16) = 14 × 14 = 196 patches
Each patch: (16, 16, 3) = 768-dimensional vector when flattened

Linearly project each flattened patch to d_model dimensions. Prepend a learnable [CLS] token. Add learnable positional embeddings. Feed the resulting sequence of 197 tokens to a standard transformer encoder.

Input sequence length: 197 tokens (196 patches + 1 [CLS])
Each token: d_model dimensions (768 for ViT-B/16)

The classification head: a linear layer on top of the [CLS] token output after the final encoder layer.

# Conceptual ViT patch embedding
import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, d_model=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2
        # Conv2d with kernel=patch_size, stride=patch_size = non-overlapping patch extraction + linear projection
        self.proj = nn.Conv2d(in_channels, d_model, kernel_size=patch_size, stride=patch_size)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, d_model))
    
    def forward(self, x):
        B = x.size(0)
        # (B, C, H, W) → (B, d_model, H/P, W/P) → (B, num_patches, d_model)
        x = self.proj(x).flatten(2).transpose(1, 2)
        # Prepend [CLS] token
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)
        # Add positional embedding
        x = x + self.pos_embed
        return x

---

2. ViT Variants

Model	Layers	Hidden	Heads	Params	ImageNet-1K Top-1 (fine-tuned from JFT)
ViT-S/16	12	384	6	22M	81.4%
ViT-B/16	12	768	12	86M	84.0% (from ImageNet-21K)
ViT-L/16	24	1024	16	307M	87.76% (from JFT-300M)
ViT-H/14	32	1280	16	632M	88.55% (from JFT-300M)

For comparison: ResNet-152 trained on ImageNet-1K from scratch: 78.3%.

---

3. The Honest Data Tradeoff

The ViT paper’s most important finding is one that most tutorials omit: ViT underperforms ResNets when trained on ImageNet-1K (1.28M images) from scratch.

From the paper (Table 5):

• ViT-B/16 trained on ImageNet-1K from scratch: 77.9% top-1

• ResNet-50 trained on ImageNet-1K from scratch: 76.9% (comparable)

• ViT-L/16 trained on ImageNet-1K from scratch: 76.5% — worse than ViT-B/16

ViT only starts winning when pretrained on much larger datasets:

• ImageNet-21K (~14M images): ViT-B/16 reaches 84.0%

• JFT-300M (300M images): ViT-L/16 reaches 87.76%

Why? CNNs have inductive biases baked in: translation invariance (convolutions), local structure (small kernels), spatial hierarchy (pooling). These biases match natural image statistics. ViT has no such biases — it must learn them from data. With enough data, ViT surpasses CNNs because it can learn more flexible representations without the constraint of locality. With insufficient data, those CNN biases are an advantage.

Practical implication for an applied engineer:

• <10K labeled images: use a pretrained CNN (EfficientNet, ConvNeXt) or pretrained ViT features (DINOv2) — do not train from scratch

• 10K–100K labeled images: pretrained ViT from HuggingFace (google/vit-base-patch16-224) fine-tuned beats from-scratch ViT and is competitive with pretrained CNNs

• 1M labeled images, or when pretraining on unlabeled domain data: ViT dominates

---

4. CLIP — Contrastive Language-Image Pretraining

Paper: Radford et al. (2021), arXiv:2103.00020
Training data: 400 million (image, text) pairs scraped from the web

CLIP trains two encoders jointly:

• Image encoder: ViT (or ResNet) that maps an image to a 512-dimensional embedding

• Text encoder: Transformer that maps a text string to a 512-dimensional embedding

Training objective: for a batch of N (image, text) pairs, maximize cosine similarity of the N correct pairs and minimize similarity of the N² - N incorrect pairs. This is the InfoNCE contrastive loss.

Correct pair:    cos_sim(image_i_embedding, text_i_embedding) → 1
Incorrect pair:  cos_sim(image_i_embedding, text_j_embedding) → 0  (for i ≠ j)

Zero-Shot Classification

Once trained, CLIP can classify images into any categories without labeled training data:

import torch
import clip
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

image = preprocess(Image.open("dog.jpg")).unsqueeze(0).to(device)
text_prompts = clip.tokenize(["a dog", "a cat", "a car", "a building"]).to(device)

with torch.no_grad():
    image_features = model.encode_image(image)
    text_features = model.encode_text(text_prompts)
    
    # Cosine similarity
    logits_per_image, _ = model(image, text_prompts)
    probs = logits_per_image.softmax(dim=-1)

print(probs)  # Distribution over ["dog", "cat", "car", "building"]

Zero-shot ImageNet accuracy: 76.2% — comparable to a ResNet-50 trained with full supervision on ImageNet. This is genuinely remarkable. The model has never seen an ImageNet training example; it generalizes from the structure of 400M internet (image, caption) pairs.

Practical uses:

• Semantic image search: embed query text, retrieve images with highest cosine similarity

• Zero-shot classification on proprietary domain without labeling data

• Image-text retrieval: “find images of broken circuit boards”

• Feature extraction for downstream tasks (CLIP embeddings are strong baselines)

---

5. DINOv2

Paper: Oquab et al. (2023), arXiv:2304.07193
Self-supervised ViT trained with a knowledge distillation objective (no labels). DINOv2 features are state-of-the-art for linear probing — fitting a linear classifier on frozen features — across many vision tasks.

When to use DINOv2: you have limited labeled data but unlabeled images are available. Freeze the DINOv2 encoder, train a small head on your labeled examples. With 1% of ImageNet labels, DINOv2 ViT-L achieves 86.3% top-1.

---

6. ViT vs. CNN Decision Flowchart

Do you have a pretrained ViT checkpoint for your domain?
├── Yes → Use it. Fine-tune with a small learning rate (1e-5 to 3e-5).
└── No  → How much labeled training data do you have?
          ├── < 10K → Use pretrained DINOv2 or CLIP features + linear head/SVM
          ├── 10K–100K → Fine-tune ViT-B/16 (pretrained ImageNet-21K)
          │              vs. fine-tune EfficientNet-B4 — benchmark both
          └── > 1M → ViT-L/16 or larger. Train properly with warmup,
                     cosine decay, mixup/cutmix augmentation.
                     
Is inference latency critical (< 10ms, mobile/edge)?
└── Yes → ResNet or EfficientNet. ViT-B/16 at 224px: ~3ms on A100, 
          ~80ms on CPU. ConvNeXt-Tiny: ~15ms on CPU.

---

7. Code — Fine-Tune ViT-B/16 on Custom Classification

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from transformers import ViTForImageClassification, ViTImageProcessor
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
import numpy as np

# ── Setup ─────────────────────────────────────────────────────────────────────
NUM_CLASSES = 5  # adjust for your dataset
MODEL_NAME = "google/vit-base-patch16-224"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

processor = ViTImageProcessor.from_pretrained(MODEL_NAME)
model = ViTForImageClassification.from_pretrained(
    MODEL_NAME,
    num_labels=NUM_CLASSES,
    ignore_mismatched_sizes=True  # replaces the pretrained head
).to(DEVICE)

# ── Data transforms ───────────────────────────────────────────────────────────
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=processor.image_mean, std=processor.image_std),
])
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=processor.image_mean, std=processor.image_std),
])

# Assumes ImageFolder structure: data/train/{class_0, class_1, ...}
train_ds = datasets.ImageFolder("data/train", transform=train_transform)
val_ds   = datasets.ImageFolder("data/val",   transform=val_transform)
train_dl = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4)
val_dl   = DataLoader(val_ds,   batch_size=64, shuffle=False, num_workers=4)

# ── Training loop ─────────────────────────────────────────────────────────────
optimizer = AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)
scheduler = CosineAnnealingLR(optimizer, T_max=10)
criterion = nn.CrossEntropyLoss()

for epoch in range(10):
    model.train()
    for imgs, labels in train_dl:
        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
        outputs = model(pixel_values=imgs)
        loss = criterion(outputs.logits, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    scheduler.step()
    
    # Validation
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for imgs, labels in val_dl:
            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
            preds = model(pixel_values=imgs).logits.argmax(dim=-1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)
    print(f"Epoch {epoch+1}: Val Accuracy = {correct/total:.4f}")

ResNet-18 baseline to compare against: Use torchvision.models.resnet18(pretrained=True), replace model.fc with nn.Linear(512, NUM_CLASSES). Same training loop. Typical result on a 5-class dataset with ~5K images per class: ResNet-18 ≈ 88-92%, ViT-B/16 fine-tuned ≈ 91-94%.

---

What Most People Get Wrong

ViT is not universally better than CNNs. This is stated directly in the ViT paper, but gets glossed over in every tutorial that was written after ViT went viral. The paper’s conclusion: “ViT attains excellent results when pre-trained on large amounts of data and transferred to tasks with fewer data points.”

The key phrase is pre-trained on large amounts of data. If you’re fine-tuning from google/vit-base-patch16-224 (which was pretrained on ImageNet-21K), you have that large pretraining baked in. If you’re training from scratch on your 500-sample dataset, you’re in the “ResNets win” regime and you should know it.

The other failure mode: ViT with 16×16 patches is not good for high-resolution images where fine-grained local features matter (e.g., pathology slides, satellite imagery). For those, either use smaller patches (8×8 — 4x more tokens, quadratic attention cost), hierarchical ViTs (Swin Transformer, arXiv:2103.14030), or just use a CNN backbone.

---

Return to [README.md] · Next: [04_diffusion_models_foundations.md]