Building Vision Transformers (ViT) for Image Tasks: A Complete From-Scratch Guide

Vision Transformers treat images as sequences of patch tokens and apply transformer architectures—complete with multi-head self-attention and feed-forward networks—to perform image classification and representation learning.

This guide walks through the complete implementation of a Vision Transformer (ViT) in the rohitg00/ai-engineering-from-scratch repository. Instead of relying on pre-built libraries, the code assembles the full pipeline using modular PyTorch components, making it ideal for understanding how modern computer vision models actually work under the hood.

How Vision Transformers Process Images

The ViT pipeline converts spatial image data into a sequence format that standard transformer encoders can process. According to the source code in phases/19-capstone-projects/58-vision-encoder-patches/code/main.py, this involves three distinct operations performed by the VisionFrontEnd class: patch embedding, CLS token prepending, and positional encoding injection.

Patch Tokenization with Conv2d Projection

Rather than processing individual pixels, the model divides the image into non-overlapping patches using a convolutional projection. The PatchEmbed class in phases/19-capstone-projects/58-vision-encoder-patches/code/main.py implements this via a Conv2d layer where kernel size equals stride size equals patch size (typically 16×16 pixels).

This projection reduces the quadratic attention complexity from pixels² to patches². For a standard 224×224 image with 16×16 patches, this yields 196 patch tokens, each mapped to a hidden dimension vector (default 768 dimensions).

CLS Token and Positional Embeddings

After patchification, the model prepends a learnable [CLS] token—implemented as an nn.Parameter in VisionFrontEnd.__init__—to the sequence. This token serves as the aggregate image representation used for downstream tasks.

The sinusoidal_2d routine then generates fixed 2-D sinusoidal position tables that encode spatial relationships using row and column sine/cosine frequencies. These positional encodings are added to every token (patches + CLS) to preserve spatial information that would otherwise be lost during the flattening process.

Transformer Encoder Architecture

The core computation occurs in the transformer stack defined in phases/19-capstone-projects/59-vit-transformer/code/main.py. Here, the VisionEncoder combines the front-end with a deep transformer network composed of Block layers, each containing multi-head self-attention (MHSA) and a feed-forward network (FFN).

Multi-Head Self-Attention and Feed-Forward Networks

Each Block contains a MultiHeadSelfAttention module that computes scaled dot-product attention across multiple heads, allowing the model to focus on different spatial relationships simultaneously. The FeedForward module applies a 4× expansion of the hidden dimension (e.g., 768 → 3072 → 768) using GELU activation, providing the network capacity to model complex visual features.

As implemented in the Block class, the architecture uses a pre-LayerNorm configuration where layer normalization is applied before the attention and FFN sub-layers rather than after. This stabilizes training without requiring elaborate learning-rate warm-up schedules.

Pre-LayerNorm and Residual Connections

The forward pass through each Block follows a strict pattern: pre-LayerNorm → MHSA → residual addition → pre-LayerNorm → FFN → residual addition. This wiring ensures gradient flow remains strong through deep stacks, which is critical when training vision transformers from scratch.

After the final transformer block and a concluding LayerNorm, the VisionEncoder extracts the image representation by taking the first token (tokens[:, 0]), which corresponds to the [CLS] token.

End-to-End Implementation

The repository provides runnable examples demonstrating how to instantiate and use the complete ViT pipeline.

Encoding a Synthetic Image

This example shows how to compose the front-end and transformer to encode an image:

from pathlib import Path
import torch
from phases.19_capstone_projects.58_vision_encoder_patches.code.main import (
    VisionFrontEnd, FrontEndConfig, synthesize_image,
)
from phases.19_capstone_projects.59_vit_transformer.code.main import (
    VisionEncoder, ViTConfig,
)

# Build the encoder (front-end + transformer)

cfg = ViTConfig()                         # default 224×224, 16-px patches, 768-dim

encoder = VisionEncoder(cfg).eval()

# Create a deterministic fixture image (seed = 0)

img = synthesize_image(seed=0)            # shape (1, 3, 224, 224)

# Forward pass → token sequence + CLS vector

tokens, cls = encoder(img)

print("tokens shape :", tokens.shape)      # (1, 197, 768)  – 196 patches + CLS

print("CLS norm    :", cls.norm().item())

Inspecting Attention Patterns

You can examine how the model attends to different patches by accessing the attention scores from the first transformer block:

encoder = VisionEncoder().eval()
with torch.no_grad():
    _ = encoder.vit.blocks[0].attn(encoder.front(img), store_attn=True)

attn = encoder.vit.blocks[0].attn.last_attn   # shape (B, heads, N, N)

cls_head0 = attn[0, 0, 0]                     # CLS → all patches

print("CLS → patch attention sums to:", cls_head0.sum().item())

This confirms the attention mechanism properies—each attention head's output sums to 1.0 after the softmax operation.

Verifying Gradient Flow

To ensure gradients propagate correctly through the CLS token and patch projections:

encoder = VisionEncoder()
img = synthesize_image(seed=2)
_, cls = encoder(img)
loss = (cls ** 2).sum()
loss.backward()

print("Grad on patch proj weight :", encoder.front.patch.proj.weight.grad.norm())
print("Grad on CLS token        :", encoder.front.cls_token.grad.norm())

This demonstrates that gradients flow backward from the final CLS representation through the entire stack, including the convolutional patch projection and learned CLS parameter.

Key Design Decisions in the ViT Implementation

The repository follows specific architectural choices validated by the original Vision Transformer research:

• Patch Size Trade-offs: Halving the patch size (e.g., from 16×16 to 8×8) quadruples the token count and computational cost, as highlighted in the lesson documentation at phases/19-capstone-projects/59-vit-transformer/docs/en.md.

• 4× FFN Expansion: The FeedForward class expands hidden dimensions by 4× before projecting back, balancing model capacity with computational efficiency.

• Fixed Sinusoidal vs. Learned Positions: The implementation uses fixed 2-D sinusoidal positional encodings rather than learned embeddings, providing strong inductive biases for spatial relationships without additional parameters.

• Modular Architecture: Separating VisionFrontEnd from ViT allows swapping components—such as replacing sinusoidal encodings with learned positions or adding register tokens—without modifying the core transformer logic in phases/19-capstone-projects/59-vit-transformer/code/main.py.

Summary

• Vision Transformers convert images to sequences using non-overlapping patch projections via Conv2d layers in the PatchEmbed class.
• The VisionFrontEnd prepends a learnable [CLS] token and adds 2-D sinusoidal positional encodings to preserve spatial information.
• The transformer encoder in phases/19-capstone-projects/59-vit-transformer/code/main.py uses pre-LayerNorm blocks with multi-head self-attention and 4× expanded feed-forward networks.
• CLS pooling extracts the final image representation by selecting the first token after processing through the Block stack.
• The modular design enables experimentation with embedding strategies and attention mechanisms while maintaining a stable training pipeline.

Frequently Asked Questions

What is the purpose of the [CLS] token in Vision Transformers?

The [CLS] token acts as a learnable aggregate representation of the entire image. Unlike convolutional networks that use global average pooling, the ViT relies on this special token to collect information from all patches through the self-attention mechanism. After the final transformer block, the CLS token's embedding serves as the fixed-dimensional image representation for classification or downstream tasks.

How does patch size affect Vision Transformer performance and computation?

Patch size directly determines the sequence length fed into the transformer. Smaller patches (e.g., 8×8) create longer sequences (784 tokens for 224×224 images) that capture finer spatial details but increase computational cost quadratically with respect to sequence length. The repository defaults to 16×16 patches (196 tokens) as a balance between resolution and efficiency, as noted in the implementation documentation.

Why use pre-LayerNorm instead of post-LayerNorm in the transformer blocks?

Pre-LayerNorm applies normalization before the attention and feed-forward sub-layers rather than after. This configuration stabilizes training deep transformer networks without requiring complex learning-rate warm-up schedules. As implemented in the Block class, pre-LayerNorm ensures that gradients flow more directly through residual connections during backpropagation.

Can this implementation

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}