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

> Learn to build Vision Transformers ViT from scratch for image tasks. This guide explains treating images as sequences and applying transformer architectures for classification.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: tutorial
- Published: 2026-06-15

---

**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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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:

```python
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:

```python
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:

```python
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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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