Vision Transformer from Scratch

Implementing ViT end-to-end and walking through the pieces that matter.

A from-scratch walkthrough of the Vision Transformer architecture in a notebook.

Building a Vision Transformer From Scratch

After implementing a small GPT-style Transformer from scratch, I wanted to understand how the same architecture could be applied to images.

Vision Transformers, or ViTs, approach image understanding by treating an image as a sequence of patches. Instead of processing the entire image through a convolutional network, the model divides it into smaller regions, converts those regions into vectors, and processes them using Transformer blocks.

In this notebook, I build a small Vision Transformer from scratch using PyTorch.

The goal is to understand how an image becomes a sequence of tokens, how positional information is added, and how the Transformer architecture can be adapted for image classification.

This is a deliberately small first implementation. It prioritizes understanding the architecture and its tensor shapes over achieving state-of-the-art accuracy.

What we’re building

A Vision Transformer consists of two major parts:

1. Converting an image into a sequence of patch embeddings.
2. Processing that sequence using a Transformer encoder.

The implementation follows this structure:

Image
  ↓
Patch Embedding
  ↓
Class Token + Positional Embeddings
  ↓
Transformer Blocks × N
  ├── LayerNorm
  ├── Multi-Head Self-Attention
  └── MLP
  ↓
Class Token Representation
  ↓
Classification Head
  ↓
Class Logits

We will implement the components in the following order:

1. Configuration
2. Patch Embedding
3. Class Token and Positional Embeddings
4. LayerNorm
5. Single-Head Attention
6. Multi-Head Attention
7. MLP
8. Transformer Block
9. Full Vision Transformer
10. Dataset and Training
11. Evaluation

Setup

import math
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Device:", device)

SEED = 42
random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed(SEED)

Device: cpu

Configuration

Before implementing the model, we define the architectural hyperparameters in one place.

The important image-related parameters are:

• img_size: height and width of the input image
• patch_size: height and width of each image patch
• in_channels: number of image channels
• d_model: size of each patch representation
• n_heads: number of attention heads
• d_head: dimension of each attention head
• d_mlp: hidden dimension of the MLP
• n_layers: number of Transformer blocks
• n_classes: number of output classes

For this first implementation, we use small images and a small model so that the complete architecture can be trained and inspected without requiring a large computational budget.

@dataclass
class Config:
    # Image / patch related (NEW)
    img_size: int = 32          # e.g. 32 or 64 for toy experiments
    patch_size: int = 4         # must divide img_size evenly
    in_channels: int = 3        # 3 for RGB, 1 for grayscale
    num_classes: int = 10

    # Model dimensions (SAME as language Transformer)
    d_model: int = 128
    n_heads: int = 4
    d_head: int = 32
    d_mlp: int = 512
    n_layers: int = 4
    dropout: float = 0.0

    # Derived
    @property
    def num_patches(self):
        return (self.img_size // self.patch_size) ** 2

    # Sanity-Check
    def __post_init__(self):
      assert self.d_model == self.n_heads*self.d_head
      assert self.img_size % self.patch_size == 0

Patch Embedding

A Transformer expects a sequence of vectors as input.

An image, however, is represented as a grid of pixels:

(B, C, H, W)

where:

• B is the batch size
• C is the number of channels
• H is the image height
• W is the image width

To apply a Transformer to images, we divide each image into smaller patches.

For example, a 32 × 32 image with a patch size of 4 × 4 produces:

\[ \frac{32}{4} \times \frac{32}{4} = 64 \]

patches.

Each patch is then projected into a vector of size d_model.

The result is a sequence of patch embeddings:

(B, num_patches, d_model)

This is the bridge between image processing and the Transformer architecture.

class PatchEmbedding(nn.Module):
  def __init__(self, cfg: Config):
    super().__init__()
    self.cfg = cfg

    self.proj = nn.Conv2d(
        cfg.in_channels,
        cfg.d_model,
        kernel_size=cfg.patch_size,
        stride=cfg.patch_size,
    )

  def forward(self, x):
    x = self.proj(x)
    x = x.flatten(2)
    x = x.transpose(1,2)

    return x

patch.png

Image Credit

Class Token and Positional Embeddings

After patch embedding, an image is represented as a sequence of patch vectors.

However, we still need two things:

1. A way to preserve spatial information.
2. A representation that can be used for image classification.

Positional embeddings

Self-attention does not inherently know the spatial arrangement of the patches.

The patch at the top-left and the patch at the bottom-right are both just vectors. Without positional information, the model would not know where they came from.

We therefore add a learned positional embedding to each patch representation.

Class token

We also prepend a learned class token to the sequence.

This token is not associated with a particular image patch. Instead, it acts as a location where information from the entire image can be aggregated through the Transformer blocks.

After processing the image, we use the final representation of this class token for classification.

class ViTEmbedding(nn.Module):
  def __init__(self, cfg: Config):
    super().__init__()
    self.cfg = cfg
    self.patch_embed = PatchEmbedding(cfg)

    self.cls_token = nn.Parameter(torch.zeros(1,1,cfg.d_model))
    self.pos_embed = nn.Parameter(torch.zeros(1,1+cfg.num_patches,cfg.d_model))

    self.dropout = nn.Dropout(cfg.dropout)
    self._init_weights()

  def _init_weights(self):
    nn.init.trunc_normal_(self.cls_token, std=0.02)
    nn.init.trunc_normal_(self.pos_embed, std=0.02)

  def forward(self, x):
    x = self.patch_embed(x)
    B = x.shape[0]
    cls_token = self.cls_token.expand(x.shape[0], -1, -1)
    x = torch.cat((cls_token, x), dim=1)
    x = x + self.pos_embed

    return self.dropout(x)

Layer Normalization

The Vision Transformer uses the same LayerNorm mechanism as the language Transformer.

For each token representation, we normalize across the feature dimension:

\[ \hat{x} = \frac{x-\mu}{\sqrt{\sigma^2+\epsilon}} \]

The normalized representation is then scaled and shifted using learned parameters.

In this implementation, normalization is applied independently to each patch token and to the class token.

class LayerNorm(nn.Module):
  def __init__(self, d_model, eps=1e-5):
    super().__init__()
    self.gamma = nn.Parameter(torch.ones(d_model))
    self.beta = nn.Parameter(torch.zeros(d_model))
    self.eps = eps

  def forward(self, x):
    mean = torch.mean(x, dim=-1, keepdim=True)
    var = ((x-mean)**2).mean(dim=-1, keepdim=True)

    x_norm = (x - mean) / torch.sqrt(var + self.eps)

    op = self.gamma * x_norm + self.beta

    return op

Self-Attention

Once the image has been converted into a sequence of patch embeddings, the attention mechanism is almost identical to the one used in the language Transformer.

Each patch representation is projected into:

• Query
• Key
• Value

The attention scores are calculated using:

\[ QK^T \]

and scaled by:

\[ \frac{1}{\sqrt{d_{head}}} \]

The key difference is that Vision Transformers do not need causal masking for ordinary image classification.

Each patch should be able to attend to every other patch, because the model is not predicting future tokens. It is processing the image as a whole.

Therefore, the attention matrix is fully visible rather than lower-triangular.

class AttentionHead(nn.Module):
  def __init__(self, cfg):
    super().__init__()
    self.key = nn.Linear(cfg.d_model, cfg.d_head, bias=False)
    self.query = nn.Linear(cfg.d_model, cfg.d_head, bias=False)
    self.value = nn.Linear(cfg.d_model, cfg.d_head, bias=False)

    self.dropout = nn.Dropout(cfg.dropout)

  def forward(self, x):
    B, T, C = x.shape

    k = self.key(x)
    q = self.query(x)
    v = self.value(x)

    attn_scores = q @ k.transpose(-2, -1) / math.sqrt(cfg.d_head)
    attn_probs = F.softmax(attn_scores, dim=-1)

    attn_probs = self.dropout(attn_probs)

    return attn_probs @ v

Multi-Head Attention

Multi-head attention allows the model to learn several relationships between image patches simultaneously.

In this implementation, each head independently processes the patch sequence, and the resulting representations are concatenated and projected back to d_model.

The implementation is intentionally similar to the one used in the language Transformer notebook.

The main architectural difference is that the input sequence now represents image patches rather than text tokens.

class MultiHeadAttention(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.heads = nn.ModuleList(
        [AttentionHead(cfg) for _ in range(cfg.n_heads)]
    )

    self.proj = nn.Linear(cfg.n_heads * cfg.d_head, cfg.d_model)

    self.dropout = nn.Dropout(cfg.dropout)

  def forward(self, x):
    out = torch.cat([h(x) for h in self.heads], dim=-1)
    out = self.dropout(self.proj(out))

    return out

The Feed-Forward Network

The MLP is the same general component used in the language Transformer.

It expands each patch representation to a larger hidden dimension, applies a nonlinear activation, and projects it back to d_model.

The important distinction is that the MLP operates independently on each token.

In the Vision Transformer, this means each patch representation is transformed independently after attention has allowed information to be exchanged between patches.

class MLP(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.linear1 = nn.Linear(
        cfg.d_model,
        cfg.d_mlp
    )

    self.linear2 = nn.Linear(
        cfg.d_mlp,
        cfg.d_model,
    )

  def forward(self, x):

    x = self.linear1(x)

    # Apply activation

    x = F.gelu(x)

    x = self.linear2(x)

    return x

Transformer Block

The Transformer block combines LayerNorm, self-attention, the MLP, and residual connections.

The structure is the same Pre-LayerNorm design used in the previous Transformer implementation:

x
│
├── LayerNorm
│       ↓
│   Self-Attention
│       ↓
└───── (+)
        │
        ↓
    LayerNorm
        │
       MLP
        │
      (+)
        │
      Output

The block processes the entire patch sequence while preserving its shape:

(B, num_patches + 1, d_model)

class TransformerBlock(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.ln1 = LayerNorm(cfg.d_model)
    self.attn = MultiHeadAttention(cfg)

    self.ln2 = LayerNorm(cfg.d_model)
    self.mlp = MLP(cfg)

  def forward(self, x):

    # Attention + residual
    x = x + self.attn(self.ln1(x))

    # MLP + residual
    x = x + self.mlp(self.ln2(x))

    return x

The Complete Vision Transformer

We can now assemble the complete model.

The input image is converted into patch embeddings, a class token is added, and the resulting sequence is processed through several Transformer blocks.

The final class token representation is passed through a classification head to produce logits for each class.

Image
  ↓
Patch Embedding
  ↓
Class Token + Positional Embeddings
  ↓
Transformer Blocks
  ↓
Final LayerNorm
  ↓
Class Token
  ↓
Linear Classification Head
  ↓
Class Logits

The model output is:

(B, n_classes)

where each row contains the classification logits for one image.

class ViT(nn.Module):
  def __init__(self, cfg: Config):
    super().__init__()
    self.cfg = cfg
    self.embed = ViTEmbedding(cfg)
    self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
    self.norm = LayerNorm(cfg.d_model)
    self.head = nn.Linear(cfg.d_model, cfg.num_classes)
    self.apply(self._init_weights)

  def _init_weights(self, m):
    if isinstance(m, nn.Linear):
      nn.init.trunc_normal_(m.weight, std=0.02)
      if m.bias is not None:
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):
      nn.init.zeros_(m.bias)
      nn.init.ones_(m.weight)

  def forward(self, x):
    x = self.embed(x)
    for block in self.blocks:
      x = block(x)
    x = self.norm(x)
    cls_out = x[:, 0]
    return self.head(cls_out)

Dataset

To keep the experiment simple, we use a small dataset for image classification.

The purpose of this dataset is not to evaluate the full capabilities of Vision Transformers. Instead, it provides a manageable environment for checking whether the model can learn useful visual representations.

A small dataset also makes it easier to inspect training behavior and experiment with the architecture.

# Sanity Check

cfg = Config()
model = ViT(cfg).to(device)

# Fake images
x = torch.randn(4, cfg.in_channels, cfg.img_size, cfg.img_size).to(device)
logits = model(x)

print("Logits shape:", logits.shape)          # expected: (4, 10)
print("Number of parameters:", sum(p.numel() for p in model.parameters()))

Logits shape: torch.Size([4, 10])
Number of parameters: 807818

Synthetic Dataset

def get_batch(batch_size, cfg, device):
    x = torch.randn(batch_size, cfg.in_channels, cfg.img_size, cfg.img_size, device=device)
    y = torch.randint(0, cfg.num_classes, (batch_size,), device=device)
    return x, y

Loss Function

def compute_loss(logits, targets):
    return F.cross_entropy(logits, targets)

Training

The training loop follows the same general workflow as the language Transformer.

For each batch:

1. Pass the images through the model.
2. Calculate the classification loss.
3. Backpropagate through the model.
4. Update the parameters using the optimizer.

The main difference is the training objective.

Instead of predicting the next token, the Vision Transformer predicts a class label for each image.

The model is trained using cross-entropy loss over the classification logits.

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

num_steps = 200
for step in range(num_steps):
    model.train()
    x, y = get_batch(batch_size=16, cfg=cfg, device=device)

    logits = model(x)
    loss = compute_loss(logits, y)

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

    if step % 50 == 0:
        print(f"Step {step:3d} | Loss: {loss.item():.4f}")

Step   0 | Loss: 2.2782
Step  50 | Loss: 2.2705
Step 100 | Loss: 2.2685
Step 150 | Loss: 2.3349

Evaluation

Training loss alone does not tell us whether the model generalizes to unseen images.

We therefore evaluate the model on a separate validation or test set.

The main metric for this classification task is accuracy:

\[ \text{Accuracy} = \frac{\text{Correct Predictions}} {\text{Total Predictions}} \]

For a small learning experiment, accuracy and loss provide a useful first indication of whether the model is learning the intended visual patterns.

model.eval()
with torch.no_grad():
    x, y = get_batch(4, cfg, device)
    logits = model(x)
    preds = logits.argmax(dim=-1)
    print("Predictions:", preds.tolist())
    print("Targets:    ", y.tolist())

Predictions: [3, 3, 3, 3]
Targets:     [2, 9, 0, 3]