Vision Transformers (ViT) & Attention in Computer Vision

1 The Paradigm Shift: From Local Convolutions to Global Transformers

For nearly a decade following the breakthrough of AlexNet in 2012, Convolutional Neural Networks (CNNs) stood as the undisputed architectural foundation of computer vision. Whether processing handwritten digits (LeNet), deep residual identity shortcuts (ResNet), real-time object bounding boxes (YOLO), or biomedical pixel segmentation (U-Net), the primary computational engine remained the same: discrete $2\text{D}$ spatial convolution.

Convolutional networks owe their historical dominance to strong built-in assumptions known as inductive biases:

In late 2020, researchers from Google Research, Brain Team published a landmark paper that disrupted the field: "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" (Dosovitskiy et al., ICLR 2021). They demonstrated that a pure Transformer architecture—originally developed for Natural Language Processing (Vaswani et al., 2017)—could be applied directly to sequences of image patches, matching and outperforming state-of-the-art CNNs while requiring substantially fewer computational resources to pretrain.

2 Why Do We Need Attention? The Anatomy of Relational Context

To comprehend a visual scene, an intelligent agent cannot evaluate visual features in isolation. Consider an image depicting a cyclist navigating an intersection:

• The cyclist's hands must be interpreted in relationship to the handlebars.
• The cyclist's feet must be contextualized relative to the pedals and crankset.
• The trajectory of the bicycle must be evaluated in relation to the traffic signal and the pedestrian crossing 50 meters ahead.

In a standard CNN, the interaction between these spatially disparate regions can only occur after multiple pooling stages have compressed spatial resolution down to tiny feature maps ($7 \times 7$), discarding fine-grained edge alignments.

Attention fundamentally solves this by treating every visual region as an interactive entity capable of querying any other region directly, regardless of spatial distance.

3 The Transformer Encoder Architecture

The core processing engine of a Vision Transformer is a stack of identical Transformer Encoder blocks (typically $L=12$ in ViT-Base and $L=24$ in ViT-Large). Each block executes a two-stage transformation using Pre-Layer Normalization (Pre-LN) and residual identity connections:

4 How Does a Vision Transformer Process an Image?

Transformers require a 1D sequence of token embeddings as input. Because treating individual pixels as tokens is computationally intractable ($224 \times 224 = 50,176$ tokens yields a staggering $2.5 \times 10^9$ attention score matrix), ViT breaks the 2D image into a sequence of small non-overlapping patches.

# Patch embedding layer in PyTorch
patch_embed = nn.Conv2d(
    in_channels=3,
    out_channels=768,      # Embedding dimension D
    kernel_size=16,        # Patch size P
    stride=16              # Non-overlapping stride
)
# Input: [B, 3, 224, 224] -> Output: [B, 768, 14, 14] -> Flatten: [B, 768, 196] -> Transpose: [B, 196, 768]

5 Positional Embeddings: Overcoming Permutation Invariance

A critical mathematical property of pure self-attention is permutation equivariance. For any permutation matrix $\mathbf{P} \in \{0, 1\}^{N \times N}$:

In plain English: if you randomly shuffle the 196 image patches like pieces of a jigsaw puzzle, raw self-attention computes the exact same representations, merely permuted in order! A picture of a cat with its head below its paws would produce identical predictions to a normal upright cat.

To restore spatial geometry, ViT injects positional embeddings $\mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}$ by element-wise addition before the first encoder block:

6 The Mathematics of Scaled Dot-Product Attention

At the very heart of the Transformer lies the scaled dot-product attention formula:

7 Multi-Head Self-Attention (MHSA): Projecting Multiple Subspaces

A single attention function averages all relationships together, potentially diluting specialized visual features. Multi-Head Self-Attention (MHSA) executes $h$ separate attention functions in parallel, allowing the model to jointly attend to information from different representation subspaces at different positions:

8 The Vision Transformer Architectural Family

The original Dosovitskiy et al. paper established standard naming conventions based on model size (Base, Large, Huge) and patch resolution ($16 \times 16$ vs $14 \times 14$):

9 Computational Complexity: The Quadratic Scaling Bottleneck

Why can't we simply set patch size $P=1 \times 1$ to capture pixel-perfect detail? The answer lies in the quadratic time and memory complexity of the attention matrix:

10 PyTorch Implementation: Pretrained ViT-B/16 Inference

Torchvision provides direct access to modern Vision Transformers pretrained on ImageNet-1k and ImageNet-21k. Let's inspect the canonical end-to-end evaluation pipeline:

import torch
import torchvision.transforms as T
from torchvision.models import vit_b_16, ViT_B_16_Weights
from PIL import Image

# 1. Load pretrained model with modern Weights API
weights = ViT_B_16_Weights.DEFAULT
model = vit_b_16(weights=weights)
model.eval()

# 2. Inspect the official preprocessing transform
# Standard ViT preprocessing: Resize(256) -> CenterCrop(224) -> ToTensor() -> Normalize
preprocess = weights.transforms()
print("Transforms applied:", preprocess)

# 3. Load input image and construct 4D batch tensor [B, C, H, W]
image = Image.open("sample_image.jpg").convert("RGB")
input_tensor = preprocess(image).unsqueeze(0)  # Shape: [1, 3, 224, 224]

# 4. Execute forward inference under no_grad context
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
input_tensor = input_tensor.to(device)

with torch.no_grad():
    logits = model(input_tensor)                 # [1, 1000] ImageNet class logits
    probs = torch.softmax(logits, dim=-1)        # Calibrate probabilities

# 5. Extract Top-5 predictions with confidence percentages
top5_probs, top5_indices = torch.topk(probs, k=5, dim=-1)
categories = weights.meta["categories"]

print("\n--- Top-5 ImageNet Predictions ---")
for rank in range(5):
    class_idx = top5_indices[0, rank].item()
    conf = top5_probs[0, rank].item() * 100.0
    label = categories[class_idx]
    print(f"Rank {rank+1}: {label:<25} | Confidence: {conf:6.2f}%")

11 Transfer Learning & Fine-Tuning Pipeline

Adapting a pretrained Vision Transformer to a custom dataset requires replacing the 1,000-class ImageNet classification head and executing a two-stage training strategy:

import torch
import torch.nn as nn
from torchvision.models import vit_b_16, ViT_B_16_Weights

# Step 1: Instantiate pretrained model
num_classes = 3  # Custom classes: Cats, Dogs, Horses
model = vit_b_16(weights=ViT_B_16_Weights.DEFAULT)

# Step 2: Replace classification projection head
in_features = model.heads.head.in_features  # 768 in ViT-Base
model.heads.head = nn.Linear(in_features, num_classes)

# Step 3: Stage 1 — Freeze Transformer Backbone (Linear Probing)
for param in model.parameters():
    param.requires_grad = False
for param in model.heads.parameters():
    param.requires_grad = True

# Stage 1 Optimizer: Train only classification head
optimizer_head = torch.optim.AdamW(
    model.heads.parameters(),
    lr=1e-3,
    weight_decay=0.01
)
criterion = nn.CrossEntropyLoss()

# Step 4: Stage 2 — Unfreeze Backbone with Differential Learning Rates
# After 3-5 epochs, unfreeze full network but apply delicate LR to pretrained weights
for param in model.parameters():
    param.requires_grad = True

optimizer_full = torch.optim.AdamW([
    {"params": model.conv_proj.parameters(), "lr": 1e-5},
    {"params": model.encoder.parameters(),   "lr": 1e-5},
    {"params": model.heads.parameters(),     "lr": 1e-4}
], weight_decay=0.05)

12 Data Augmentation & Regularization Strategies

Because Vision Transformers lack the hard-coded inductive biases of CNNs, they are prone to severe overfitting when trained or fine-tuned without heavy regularization. The seminal DeiT paper (Touvron et al., ICML 2021) showed that aggressive data augmentation is the single most critical factor in unlocking Transformer accuracy:

13 Beyond Classification: DETR & SegFormer

The power of self-attention extends far beyond single-label image classification. Modern vision systems employ Transformers for dense spatial prediction:

14 Top 10 Vision Transformer Pitfalls & Anti-Patterns

15 Hands-on Practical Assignment: Build a 3-Class Custom ViT Classifier

16 Key Architectural Takeaways