Vision Transformers (ViT)

Module: Computer Vision | Difficulty: Advanced

Patch Embedding

Vision Transformers adapt the NLP transformer architecture for images by treating image patches as tokens. Unlike CNNs that process pixels through local convolutions, ViT processes non-overlapping patches through global self-attention, enabling direct modeling of long-range dependencies across the entire image.

Given an input image , we reshape it into patches of size , where is the patch size (typically 16). Each patch is flattened and linearly projected to embedding dimension :

Where each parameter means:

• — flattened patch
• — patch embedding matrix
• — learnable classification token
• — positional embeddings
• Intuition: Each patch becomes a "word" in the transformer; the class token aggregates information from all patches for final classification

Positional Encoding

Since transformers have no built-in spatial bias, positional information must be explicitly added:

Where each parameter means:

• — token embeddings at layer
• — learnable positional embeddings
• Intuition: Positional embeddings tell the model where each patch came from in the image; without them, the model cannot distinguish spatial relationships

Self-Attention for Images

Multi-Head Self-Attention

Each transformer block applies multi-head self-attention, allowing the model to attend to different spatial relationships:

Where each parameter means:

• — number of attention heads (12 for ViT-Base)
• — projection matrices for head
• — output projection matrix
• — dimension per head (64 for ViT-Base)
• Intuition: Different heads learn different types of relationships (local vs. global, horizontal vs. vertical), providing richer representations than single-head attention

Self-Attention Complexity

The computational cost of self-attention is:

Where each parameter means:

• — number of tokens (patches + class token)
• — embedding dimension
• Intuition: Self-attention compares every pair of tokens, leading to quadratic complexity in the number of patches; this is why ViT uses relatively large patches (16×16) to keep N manageable

Feed-Forward Network

Each transformer block includes a position-wise feed-forward network:

Where each parameter means:

• — first linear layer (expansion)
• — second linear layer (contraction)
• — Gaussian Error Linear Unit activation
• Intuition: The FFN processes each token independently with a 4× expansion ratio, providing non-linear transformation capacity

ViT Model Comparison

Complete ViT Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F


class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_channels, embed_dim,
                              kernel_size=patch_size, stride=patch_size)
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches + 1, embed_dim))

    def forward(self, x):
        B = x.shape[0]
        x = self.proj(x).flatten(2).transpose(1, 2)
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)
        return x + self.pos_embed


class TransformerBlock(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, int(embed_dim * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(embed_dim * mlp_ratio), embed_dim)
        )

    def forward(self, x):
        h = self.norm1(x)
        h, _ = self.attn(h, h, h)
        x = x + h
        x = x + self.mlp(self.norm2(x))
        return x


class VisionTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3,
                 num_classes=1000, embed_dim=768, depth=12, num_heads=12):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size,
                                          in_channels, embed_dim)
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

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


model = VisionTransformer(
    img_size=224, patch_size=16, embed_dim=768,
    depth=12, num_heads=12, num_classes=1000
)
params = sum(p.numel() for p in model.parameters())
print(f"ViT-Base parameters: {params:,}")

Vision Transformer Transfer Learning

Linear Probing

Linear probing freezes the pretrained backbone and trains only a linear classifier:

Where each parameter means:

• — frozen ViT backbone
• — classification token output
• — learnable linear layer
• Intuition: By keeping the backbone fixed, linear probing evaluates the quality of learned representations without adaptation; typically achieves 1-2% lower accuracy than fine-tuning

Full Fine-Tuning

Full fine-tuning updates all parameters with a lower learning rate:

Where each parameter means:

• — all model parameters
• — learning rate (typically 1e-4 to 1e-5)
• Intuition: Fine-tuning adapts the entire model to the target task, typically achieving the best accuracy but requiring more data and computation

Adapter Tuning

Adapter tuning inserts small trainable modules while keeping the backbone frozen:

Where each parameter means:

• — hidden state
• — small MLP (typically 64-dim bottleneck)
• — layer normalization
• Intuition: Adapters add less than 1% parameters while achieving competitive fine-tuning accuracy; this enables efficient multi-task adaptation

Vision Transformer Applications

Dense Prediction

ViT can be adapted for dense prediction tasks through hierarchical feature extraction:

Where each parameter means:

• — feature map from transformer layer
• Features at different layers capture different semantic levels
• Intuition: By extracting features from intermediate layers, ViT provides multi-scale representations for segmentation and detection

Video Understanding

ViT extends to video through spatiotemporal patch embedding:

Where each parameter means:

• — number of frames
• — spatial dimensions
• — patch size (typically 16)
• Intuition: Video ViT treats each frame as a sequence of patches, with temporal attention modeling motion; this achieves state-of-the-art on Kinetics-400 (88.7% top-1)

Common Challenges

1. Data Hunger: ViT requires massive pretraining data (300M+ images) to outperform CNNs; with only ImageNet-1K, performance lags behind CNNs
2. Quadratic Complexity: Self-attention scales quadratically with patch count, limiting resolution or requiring efficient attention variants
3. No Inductive Bias: Unlike CNNs, ViT has no built-in translation invariance or locality, requiring more data to learn these properties
4. Computational Cost: Training ViT-Large requires significant GPU resources (8 V100s for 30+ hours on ImageNet)
5. Transfer Learning: ViT features transfer differently than CNN features, requiring different fine-tuning strategies

Case Study: DeiT Training Efficiency

Facebook's DeiT (Data-efficient Image Transformers, 2021) achieved 83.8% top-1 on ImageNet with ViT-Base trained only on ImageNet-1K (1.28M images) for 300 epochs on 4 V100 GPUs in 53 hours. Key innovations included: knowledge distillation from RegNetY-16GF teacher, strong augmentation (RandAugment, Mixup, CutMix, Erasing), and stochastic depth. The distillation token learned to mimic the teacher's output, providing an additional 1.5% accuracy boost. DeiT demonstrated that with proper training strategies, ViT can match CNN performance without large-scale pretraining, making transformers accessible to smaller research groups.

Efficient Attention Mechanisms

Linear Attention

Linear attention reduces complexity from quadratic to linear:

Where each parameter means:

• — feature map (e.g., )
• Intuition: By changing the order of matrix multiplication, linear attention avoids computing the full N x N attention matrix, reducing complexity from O(N^2) to O(N)

Window Attention (Swin)

Swin Transformer computes attention within local windows, reducing complexity:

Where each parameter means:

• — queries, keys, values within window
• Window size typically 7x7
• Intuition: By computing attention only within local windows, Swin reduces complexity from O(N^2) to O(N x M^2) where M is window size; shifted windows provide cross-window connections

Shifted Window Mechanism

Swin uses shifted windows to enable cross-window information flow:

Where each parameter means:

• — tensor shift operation
• — shift amount (typically half the window size)
• Intuition: By shifting the window partition between consecutive layers, adjacent windows can share information, providing global connectivity while maintaining local computation

ViT Pretraining Strategies

Masked Image Modeling (MAE)

MAE randomly masks patches and reconstructs them:

Where each parameter means:

• — set of masked patch indices (typically 75%)
• — original patch
• — reconstructed patch
• Intuition: By reconstructing masked patches, the model learns rich visual representations without labels; MAE achieves 87.8% on ImageNet with ViT-Large

BEiT (BERT for Image Transformers)

BEiT uses discrete visual tokens from DALL-E for masked prediction:

Where each parameter means:

• — transformer output for masked positions
• — visual tokens from DALL-E tokenizer
• Intuition: By predicting discrete visual tokens instead of raw pixels, BEiT learns semantic features; the DALL-E tokenizer provides meaningful visual vocabulary

DINO (Self-Distillation)

DINO trains vision transformers through self-distillation without labels:

Where each parameter means:

• — set of augmented views
• — teacher output (EMA of student)
• — student output
• Intuition: The teacher provides consistent targets across augmentations, learning semantic features; DINO produces excellent attention maps for object discovery

Swin Transformer Architecture

Swin Transformer introduces hierarchical feature maps through patch merging:

Where each parameter means:

• — feature maps at 2x2 grid positions
• — linear projection matrix
• Intuition: By merging 2x2 patches into one, Swin doubles channel dimensions while halving spatial resolution, creating a pyramidal feature hierarchy like CNNs

ViT Model Comparison

Key Takeaways

• Vision Transformers treat image patches as tokens, enabling global self-attention across the entire image
• Patch embedding linearly projects flattened patches into the transformer embedding space
• Positional embeddings are essential since transformers have no built-in spatial inductive bias
• Self-attention complexity is quadratic in the number of patches, limiting resolution
• ViT requires large-scale pretraining to match CNN performance, but excels with sufficient data
• Hierarchical variants (Swin) combine local window attention with global receptive fields
• Self-supervised pretraining (MAE, DINO) enables ViT to learn without labels