Vision Transformer (ViT) Model and Its Implementation Using Python
For nearly a decade, Convolutional Neural Networks (CNNs) stood as the undisputed standard for computer vision tasks. By leveraging sliding kernels, local receptive fields, and translation invariance, CNNs excelled at extracting hierarchical spatial representations from image pixel grids.
However, in 2020, Dosovitskiy et al. introduced “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale,” presenting the Vision Transformer (ViT). By completely discarding convolutional operations and adapting the self-attention mechanism originally designed for natural language processing, ViTs demonstrated that treating an image as a sequence of visual patches allows a network to capture global, long-range contextual relationships right from the very first layer.
In this article, we will explore the core architecture and mathematical principles of the Vision Transformer, unpack the concepts of patch embeddings and positional encodings, and build a complete, end-to-end ViT model in Python using PyTorch.
An effective visual diagram helps clarify how an image is split into non-overlapping patches, linearly projected into embeddings, and processed through Multi-Head Self-Attention layers:
[Image: Vision Transformer (ViT) Architecture Diagram showing Image Patch Extraction, Linear Projection, Class Token Insertion, Positional Encoding, Transformer Encoder Blocks, and MLP Classification Head]
1. Architectural Foundations: How Transformers Process Images
Traditional Transformers expect a 1D sequence of token embeddings as input, such as word vectors in a text document. Images, however, are 3D arrays of shape (Height x Width x Channels).
To bridge this structural gap without using convolutions, the Vision Transformer applies a sequence of spatial transformations:
A. Patch Extraction and Flattening
An input image of spatial dimensions (H x W x C) is divided into a grid of non-overlapping square patches of size (P x P). The total number of patches, N, is calculated as:
N = (H * W) / (P²)
For example, a 224x224 RGB image split into 16x16 patches yields N = (224 * 224) / (16 * 16) = 196 patches. Each 3D patch is then flattened into a 1D vector of length (P² * C) = 16 * 16 * 3 = 768 features.
B. Linear Projection and Patch Embeddings
Each flattened patch vector is passed through a learnable Linear Projection matrix to project the raw pixel values into a D-dimensional vector space (known as the embedding dimension). This creates an initial patch embedding sequence of shape (N, D).
C. The Classification Token (CLS Token)
Following the design of BERT in NLP, a learnable vector parameter called the Classification Token ([CLS]) is prepended to the sequence of patch embeddings. The final state of this token after passing through the Transformer encoder serves as the global image representation fed into the classification head.
Get Iqra Mahmood’s stories in your inbox
Join Medium for free to get updates from this writer.
The sequence length becomes (N + 1) tokens.
D. Positional Encodings
Because self-attention operations are inherently permutation-invariant (they process all tokens simultaneously without awareness of order), 1D learnable Positional Encodings of dimension D are added to the patch embeddings. This retains spatial structure, allowing the model to know where each patch originated within the original 2D grid.
2. Under the Hood: The Transformer Encoder Block
Once patch embeddings and positional encodings are combined, the sequence passes through a stack of L Transformer Encoder blocks. Each block consists of two main sub-layers:
- Multi-Head Self-Attention (MHSA): Allows every patch token to interact with and attend to every other patch token across the entire image simultaneously. This captures both local details and global relationships regardless of spatial distance.
- Multi-Layer Perceptron (MLP): A feedforward block consisting of two dense linear layers with a GeLU activation function in between, typically expanding feature dimensions by a factor of 4 before projecting back down.
Residual skip connections and Layer Normalization are applied before each sub-layer (Pre-LN architecture) to ensure numerical stability during gradient updates.
3. Complete Python Implementation Using PyTorch
Below is an end-to-end, modular Python script that builds a Vision Transformer from scratch using PyTorch. It defines Patch Embeddings, Multi-Head Self-Attention, Transformer Encoder blocks, and the full ViT classifier.
import torch
import torch.nn as nn
import torch.optim as optim# Ensure reproducible runs
torch.manual_seed(42)# -------------------------------------------------------------------
# Step 1: Patch Embedding Layer
# -------------------------------------------------------------------
class PatchEmbedding(nn.Module):
"""
Splits a 2D image into patches and projects them into a D-dimensional space
using a 2D Convolution with kernel_size equal to stride size.
"""
def __init__(self, in_channels=3, patch_size=16, embed_dim=768, img_size=224):
super(PatchEmbedding, self).__init__()
self.patch_size = patch_size
self.num_patches = (img_size // patch_size) ** 2
# Efficient projection using Conv2d with stride equal to patch size
self.projection = nn.Conv2d(
in_channels,
embed_dim,
kernel_size=patch_size,
stride=patch_size
) def forward(self, x):
# Input shape: (Batch, Channels, Height, Width)
# Output shape: (Batch, Embed_Dim, Grid_H, Grid_W)
x = self.projection(x)
# Flatten spatial dimensions into sequence length: (Batch, Embed_Dim, Num_Patches)
x = x.flatten(2)
# Transpose to sequence format: (Batch, Num_Patches, Embed_Dim)
x = x.transpose(1, 2)
return x
# -------------------------------------------------------------------
# Step 2: Multi-Head Self-Attention Block
# -------------------------------------------------------------------
class MultiHeadSelfAttention(nn.Module):
def __init__(self, embed_dim=768, num_heads=12, dropout=0.1):
super(MultiHeadSelfAttention, self).__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.qkv = nn.Linear(embed_dim, embed_dim * 3)
self.proj = nn.Linear(embed_dim, embed_dim)
self.dropout = nn.Dropout(dropout)
self.scale = self.head_dim ** -0.5 def forward(self, x):
batch_size, seq_len, embed_dim = x.shape
# Calculate Queries, Keys, and Values in a single matrix multiplication
qkv = self.qkv(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4) # (3, Batch, Heads, Seq_Len, Head_Dim)
q, k, v = qkv[0], qkv[1], qkv[2]
# Scaled Dot-Product Attention: (q @ k.T) / sqrt(d_k)
attn_scores = (q @ k.transpose(-2, -1)) * self.scale
attn_weights = torch.softmax(attn_scores, dim=-1)
attn_weights = self.dropout(attn_weights)
# Combine attention with Values
out = (attn_weights @ v).transpose(1, 2).reshape(batch_size, seq_len, embed_dim)
out = self.proj(out)
return out
# -------------------------------------------------------------------
# Step 3: Transformer Encoder Layer
# -------------------------------------------------------------------
class TransformerEncoderBlock(nn.Module):
def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout=0.1):
super(TransformerEncoderBlock, self).__init__()
self.norm1 = nn.LayerNorm(embed_dim)
self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
self.norm2 = nn.LayerNorm(embed_dim)
hidden_dim = int(embed_dim * mlp_ratio)
self.mlp = nn.Sequential(
nn.Linear(embed_dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, embed_dim),
nn.Dropout(dropout)
) def forward(self, x):
# Pre-Layer Normalization and Residual Skip Connections
x = x + self.attn(self.norm1(x))
x = x + self.mlp(self.norm2(x))
return x
# -------------------------------------------------------------------
# Step 4: Full Vision Transformer Architecture
# -------------------------------------------------------------------
class VisionTransformer(nn.Module):
def __init__(
self,
img_size=224,
patch_size=16,
in_channels=3,
num_classes=10,
embed_dim=768,
depth=6,
num_heads=12,
mlp_ratio=4.0,
dropout=0.1
):
super(VisionTransformer, self).__init__()
self.patch_embed = PatchEmbedding(in_channels, patch_size, embed_dim, img_size)
num_patches = self.patch_embed.num_patches
# Learnable [CLS] token and Positional Encodings
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
self.pos_drop = nn.Dropout(p=dropout)
# Stacked Transformer Encoder Blocks
self.blocks = nn.ModuleList([
TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, dropout)
for _ in range(depth)
])
# Final Layer Normalization and Classification Head
self.norm = nn.LayerNorm(embed_dim)
self.head = nn.Linear(embed_dim, num_classes) # Initialize parameters
nn.init.trunc_normal_(self.pos_embed, std=0.02)
nn.init.trunc_normal_(self.cls_token, std=0.02) def forward(self, x):
batch_size = x.shape[0]
# Step A: Convert input images into patch embeddings
x = self.patch_embed(x) # (Batch, Num_Patches, Embed_Dim)
# Step B: Prepend [CLS] token to input sequence
cls_tokens = self.cls_token.expand(batch_size, -1, -1)
x = torch.cat((cls_tokens, x), dim=1) # (Batch, Num_Patches + 1, Embed_Dim)
# Step C: Add Positional Encodings
x = x + self.pos_embed
x = self.pos_drop(x)
# Step D: Pass through stacked Transformer Blocks
for block in self.blocks:
x = block(x)
x = self.norm(x)
# Step E: Extract [CLS] token output for classification
cls_token_final = x[:, 0]
logits = self.head(cls_token_final)
return logits
# -------------------------------------------------------------------
# Step 5: Model Instantiate and Verification
# -------------------------------------------------------------------
if __name__ == "__main__":
# Create synthetic input tensor representing a batch of 2 RGB images (224x224)
dummy_input = torch.randn(2, 3, 224, 224)
# Instantiate Vision Transformer model
model = VisionTransformer(
img_size=224,
patch_size=16,
in_channels=3,
num_classes=10,
embed_dim=384, # ViT-Small configuration
depth=6,
num_heads=6
)
# Execute forward pass
outputs = model(dummy_input)
print("--- Vision Transformer Execution Summary ---")
print(f"Input Shape: {list(dummy_input.shape)} (Batch, Channels, Height, Width)")
print(f"Output Shape: {list(outputs.shape)} (Batch, Num_Classes)")
print(f"Total Trainable Parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
4. CNNs vs. Vision Transformers: Key Trade-Offs
While Vision Transformers deliver state-of-the-art performance across computer vision benchmarks, choosing between CNNs and ViTs involves understanding distinct trade-offs:
- Inductive Bias vs. Data Scale: CNNs feature strong built-in inductive biases (translation invariance and local pixel proximity). Consequently, CNNs train effectively on small or medium datasets. ViTs lack these structural assumptions, requiring large-scale pre-training (e.g., ImageNet-21k or JFT-300M) or heavy data augmentation to surpass CNN performance.
- Global Context from Layer 1: Convolutions process visual information within small, local receptive fields that only expand in deeper layers. Self-attention in ViTs computes relationships across all image patches globally from the very first layer.
- Computational Complexity: In standard self-attention, computational complexity scales quadratically with respect to sequence length (number of patches). Fine-tuning ViTs on higher-resolution images requires efficient attention mechanisms or windowing strategies (such as Swin Transformers).
