Vision Transformer (ViT): Patch Embedding, Self-Attention, and PyTorch Architecture

Inductive Biases vs. Universal Sequence Scalability

Convolutional Neural Networks (CNNs; ResNet, ConvNeXt) enforce two explicit architectural inductive biases:

1. Spatial Locality: Filters process small local pixel neighborhoods ($3\times 3,7\times 7$), requiring deep cascades of $L$ layers to achieve a global receptive field.
2. Translation Equivariance: Convolution kernels share identical parameters across spatial coordinates: $f(\mathrm{Shift}(\mathbf{x}))=\mathrm{Shift}(f(\mathbf{x}))$.

While these priors provide exceptional sample efficiency on small datasets ($\le 10^6$ images), they impose structural constraints on asymptotic scaling.

The Vision Transformer (ViT) (Dosovitskiy et al., ICLR 2021) eliminates convolutional inductive biases entirely. An input image is partitioned into non-overlapping spatial patches ${\mathbf{x}}_p\in {\mathbb{R}}^{P\times P\times C}$, linearly projected into visual tokens, prepended with a learnable class token (${\mathbf{x}}_{\text{class}}$), augmented with 1D/2D positional embeddings, and processed by a standard isotropic Transformer encoder:

${\mathbf{z}}_0=[{\mathbf{x}}_{\text{class}};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\dots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{\text{pos}},N=\frac{HW}{P^2}$

ViT replaces local receptive fields with instant global self-attention from layer 1, enabling predictable power-law scaling across compute budgets from $5\text{M}$ to $22\text{B}$ parameters.

Visual Architecture Paradigms Comparison

Mathematical Foundations

Figure 1: Vision Transformer pipeline decomposing 2D images into linear patch projections, adding positional embeddings, and processing via multi-head self-attention.

1. Patch Projection and Linear Embedding

Let an input image be represented as $\mathbf{x}\in {\mathbb{R}}^{H\times W\times C}$. Given a square patch resolution $P\times P$, the image is partitioned into $N=\frac{HW}{P^2}$ patches:

${\mathbf{x}}_p\in {\mathbb{R}}^{N\times (P^{2}C)}$

Each flattened patch vector ${\mathbf{x}}_p^i\in {\mathbb{R}}^{P^{2}C}$ is linearly mapped to latent embedding dimension $D$ via matrix $\mathbf{E}\in {\mathbb{R}}^{(P^{2}C)\times D}$:

${\mathbf{z}}_0=\left[{\mathbf{x}}_{\text{class}};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\dots ;{\mathbf{x}}_p^N\mathbf{E}\right]+{\mathbf{E}}_{\text{pos}}$

where ${\mathbf{x}}_{\text{class}}\in {\mathbb{R}}^{1\times D}$ is the learnable classification token, and ${\mathbf{E}}_{\text{pos}}\in {\mathbb{R}}^{(N+1)\times D}$ represents the learnable 1D spatial position embeddings.

Equivalent 2D Strided Convolutional Formulation:

${\mathbf{z}}_0[1:]=\mathrm{Reshape}\left(Conv2D\left(\mathbf{x},{\mathbf{W}}_{\text{conv}},\text{kernel}=P,\text{stride}=P\right)\right)$

where ${\mathbf{W}}_{\text{conv}}\in {\mathbb{R}}^{D\times C\times P\times P}$. This equivalence enables optimized Tensor Core SIMD execution.

2. Multi-Head Self-Attention over Visual Tokens

The token sequence $\mathbf{z}\in {\mathbb{R}}^{(N+1)\times D}$ propagates through $L$ identical Transformer encoder blocks:

${\mathbf{z}}_{\ell }^{\prime }=\mathrm{MSA}\left(\mathrm{LayerNorm}({\mathbf{z}}_{\ell -1})\right)+{\mathbf{z}}_{\ell -1},\ell =1,\dots ,L$

${\mathbf{z}}_{\ell }=\mathrm{MLP}\left(\mathrm{LayerNorm}({\mathbf{z}}_{\ell }^{\prime })\right)+{\mathbf{z}}_{\ell }^{\prime },\ell =1,\dots ,L$

Multi-Head Self-Attention (MSA):

$\mathrm{MSA}(\mathbf{X})=\left[{\mathrm{head}}_1;{\mathrm{head}}_2;\dots ;{\mathrm{head}}_h\right]{\mathbf{W}}^O$

${\mathrm{head}}_i=\mathrm{softmax}\left(\frac{(\mathbf{X}{\mathbf{W}}_i^Q)(\mathbf{X}{\mathbf{W}}_i^K{)}^{\top }}{\sqrt{D_h}}\right)(\mathbf{X}{\mathbf{W}}_i^V)$

where ${\mathbf{W}}_i^Q,{\mathbf{W}}_i^K,{\mathbf{W}}_i^V\in {\mathbb{R}}^{D\times D_h}$, $D_h=D/h$, and ${\mathbf{W}}^O\in {\mathbb{R}}^{D\times D}$.

3. Positional Embedding 2D Bicubic Interpolation

When fine-tuning a model on higher-resolution images ($H^{\prime }\times W^{\prime }>H\times W$), the token sequence length expands from $N=\frac{HW}{P^2}$ to $N^{\prime }=\frac{H^{\prime }{W}^{\prime }}{P^2}$.

Because positional embeddings ${\mathbf{E}}_{\text{pos}}\in {\mathbb{R}}^{1\times (N+1)\times D}$ are learned at fixed resolution, the patch coordinates must be interpolated:

1. Separate the classification token embedding ${\mathbf{E}}_{\text{pos}}^{\text{class}}\in {\mathbb{R}}^{1\times 1\times D}$ from spatial patch embeddings ${\mathbf{E}}_{\text{pos}}^{\text{patch}}\in {\mathbb{R}}^{1\times N\times D}$.
2. Reshape spatial embeddings to a 2D grid: ${\mathbb{R}}^{1\times \sqrt{N}\times \sqrt{N}\times D}$.
3. Apply 2D bicubic spline interpolation to target grid size: $\sqrt{N}\times \sqrt{N}\to \sqrt{N^{\prime }}\times \sqrt{N^{\prime }}$.
4. Flatten back to ${\mathbb{R}}^{1\times N^{\prime }\times D}$ and concatenate with ${\mathbf{E}}_{\text{pos}}^{\text{class}}$.

4. Computational Complexity: ViT vs. CNN

Production PyTorch Implementation

Below is a complete, modular PyTorch implementation of the VisionTransformer architecture featuring nn.Conv2d patch projections, Pre-LN LayerNorm, DropPath stochastic depth, and dynamic 2D bicubic positional interpolation.

Step 1: Patch Embedding and Stochastic Depth (DropPath)

from __future__ import annotations

import torch
import torch.nn as nn


class DropPath(nn.Module):
    """Stochastic Depth per sample (residual branch drop)."""

    def __init__(self, drop_prob: float = 0.0) -> None:
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.drop_prob == 0.0 or not self.training:
            return x
        keep_prob = 1.0 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
        random_tensor.floor_()
        return x.div(keep_prob) * random_tensor


class PatchEmbed(nn.Module):
    """2D Image to Patch Embedding via 2D Strided Convolution."""

    def __init__(self, img_size: int = 224, patch_size: int = 16, in_chans: int = 3, embed_dim: int = 768) -> None:
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size // patch_size, img_size // patch_size)
        self.num_patches = self.grid_size[0] * self.grid_size[1]

        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, C, H, W] -> proj: [B, D, H/P, W/P] -> flatten: [B, N, D]
        x = self.proj(x)
        x = x.flatten(2).transpose(1, 2)
        return x

Step 2: Transformer Encoder Block with Fused Attention

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from patch_embed import DropPath


class Attention(nn.Module):
    """Multi-Head Self-Attention with PyTorch Scaled Dot-Product Attention."""

    def __init__(self, dim: int, num_heads: int = 12, qkv_bias: bool = True) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = 1.0 / (self.head_dim ** 0.5)

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, n, c = x.shape
        qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0) # [B, num_heads, N, head_dim]

        # Hardware-accelerated FlashAttention / SDPA
        out = F.scaled_dot_product_attention(q, k, v, scale=self.scale)
        out = out.transpose(1, 2).reshape(b, n, c)
        return self.proj(out)


class MLP(nn.Module):
    """Feed-Forward Network with GELU non-linearity."""

    def __init__(self, in_features: int, hidden_features: int, out_features: int) -> None:
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_features, out_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(self.act(self.fc1(x)))


class Block(nn.Module):
    """Isotropic Transformer Encoder Block with Pre-LN."""

    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, drop_path: float = 0.0) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
        self.attn = Attention(dim, num_heads=num_heads)
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
        self.mlp = MLP(in_features=dim, hidden_features=int(dim * mlp_ratio), out_features=dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.drop_path(self.attn(self.norm1(x)))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

Step 3: Complete VisionTransformer Model with Dynamic Positional Interpolation

from __future__ import annotations

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from patch_embed import PatchEmbed
from vit_block import Block


class VisionTransformer(nn.Module):
    """Vision Transformer (ViT) Architecture."""

    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 16,
        in_chans: int = 3,
        num_classes: int = 1000,
        embed_dim: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        mlp_ratio: float = 4.0,
        drop_path_rate: float = 0.1,
    ) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.embed_dim = embed_dim
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

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

        # Stochastic depth decay
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
        self.blocks = nn.ModuleList([
            Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, drop_path=dpr[i])
            for i in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim, eps=1e-6)
        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()

        # Weight initialization
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m: nn.Module) -> None:
        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 interpolate_pos_encoding(self, x: torch.Tensor, w: int, h: int) -> torch.Tensor:
        npatch = x.shape[1] - 1
        n = self.pos_embed.shape[1] - 1
        if npatch == n and w == h:
            return self.pos_embed
        
        cls_pos = self.pos_embed[:, 0]
        patch_pos = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        
        orig_size = int(math.sqrt(n))
        patch_pos = patch_pos.reshape(1, orig_size, orig_size, dim).permute(0, 3, 1, 2)
        
        target_w, target_h = w // self.patch_embed.patch_size, h // self.patch_embed.patch_size
        patch_pos = F.interpolate(patch_pos, size=(target_h, target_w), mode="bicubic", align_corners=False)
        patch_pos = patch_pos.permute(0, 2, 3, 1).flatten(1, 2)
        return torch.cat((cls_pos.unsqueeze(1), patch_pos), dim=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, _, h, w = x.shape
        x = self.patch_embed(x)
        cls_tokens = self.cls_token.expand(b, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        
        x = x + self.interpolate_pos_encoding(x, w, h)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x)
        return self.head(x[:, 0])


def vit_base_patch16_224(num_classes: int = 1000) -> VisionTransformer:
    return VisionTransformer(img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12, num_classes=num_classes)


if __name__ == "__main__":
    model = vit_base_patch16_224(num_classes=1000)
    img = torch.randn(2, 3, 224, 224)
    logits = model(img)
    print(f"Logits output shape: {logits.shape} (Expected: [2, 1000])")

Empirical Benchmark Evaluation

Classification fidelity, pre-training compute efficiency, and transfer metrics across visual backbones:

Troubleshooting Common Vision Transformer Faults

1. Gradient Divergence / NaN Loss in Early Training Iterations

• Symptom: Training loss diverges to NaN or explodes within the first 100 iterations on AdamW.
• Root Cause: Lacking translation-equivariant inductive biases, random initialization creates chaotic cross-patch attention gradients.
• Remedy: Enforce linear learning rate warmup over at least $10\%$ of total epochs (or 5 epochs), initialize classification head weights to exact zeros (nn.init.zeros_(head.weight)), and apply gradient clipping at norm $1.0$.

2. Catastrophic Underfitting on Small Target Datasets

• Symptom: ViT-B/16 trained from scratch on CIFAR-10 or a domain-specific dataset ($<50,000$ samples) plateaus at $<65\%$ accuracy.
• Root Cause: Transformers require immense data density to learn spatial patch adjacency that CNNs possess a priori.
• Remedy: Use a pre-trained self-supervised backbone (DINOv2 / MAE), or apply aggressive data augmentations (RandAugment, Mixup $\alpha =0.8$, CutMix $\alpha =1.0$) with DropPath stochastic depth $0.1$.

3. Accuracy Collapse After Fine-Tuning on Variable Image Resolutions

• Symptom: Deploying a $224\times 224$ trained ViT on $384\times 384$ or $448\times 448$ images causes immediate classification collapse.
• Root Cause: Flattening or nearest-neighbor resizing positional embeddings introduces high-frequency phase artifacts in 2D coordinate embeddings.
• Remedy: Apply 2D bicubic spline interpolation to spatial patch embeddings while explicitly detaching and preserving the 1D classification token.

References

1. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. International Conference on Learning Representations (ICLR 2021).
2. Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., & Jégou, H. (2021). Training data-efficient image transformers & distillation through attention (DeiT). ICML 2021.
3. Liu, Z., et al. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. ICCV 2021.
4. He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. (2022). Masked Autoencoders Are Scalable Vision Learners (MAE). CVPR 2022.
5. Oquab, M., et al. (2024). DINOv2: Learning Robust Visual Features without Supervision. Transactions on Machine Learning Research (TMLR 2024).