← Vision Transformer (ViT) Series
Deep Learning · Architectures · Vision Transformer (ViT) Chapter 01 / 2

1. Vision Transformer (ViT) from First Principles: Patchify, Stride-$P$ Conv2d Projection, [CLS] Token, 2D Positional Interpolation & Inductive Biases

Before 2020, computer vision was synonymous with Convolutional Neural Networks (CNNs) because 2D convolutions bake in two powerful spatial priors—local connectivity and translation equivariance—while keeping compute linear in pixel count. When Dosovitskiy et al. (2020) introduced the Vision Transformer (ViT) in An Image is Worth 16x16 Words, their core design principle was radical minimalism: change as little of Vaswani et al.'s (2017) standard NLP Transformer Encoder as possible. By partitioning an image of size $H \times W \times C$ into $N = (H/P)(W/P)$ non-overlapping patches of $P \times P$ pixels and linearly projecting each patch into a $d$-dimensional vector, an image becomes a short 1D sequence of $N$ visual "words."

Vision Transformer end-to-end pipeline: 224x224x3 image split into 196 16x16 patches, projected via stride-16 Conv2d, prepended with [CLS], added with 197x768 positional embeddings, and passed through L Pre-Norm Transformer blocks.
Figure 1: Complete Vision Transformer (ViT) pipeline from raw image tensor $\mathbf{I} \in \mathbb{R}^{224 \times 224 \times 3}$ to $N+1 = 197$ patch+[CLS] tokens and $L$ Pre-Norm Transformer Encoder layers.

1. Why Patches Instead of Pixels? The Quadratic Bottleneck

In standard scaled dot-product self-attention over a sequence of $n$ tokens, every token queries every other token, forming an $n \times n$ attention matrix per head per layer. If we treat each pixel of an RGB image $\mathbf{I} \in \mathbb{R}^{H \times W \times C}$ as an individual token, then for a modest $224 \times 224$ image:

$$n_{\text{pixels}} = H \cdot W = 224 \times 224 = 50{,}176 \implies n_{\text{pixels}}^2 = 2{,}517{,}630{,}976 \approx 2.52 \times 10^9\text{ pairs per head}$$

Even in `float16` ($2\text{ bytes}$ per element), storing a single layer's attention scores across $H = 12$ heads for batch size $B = 1$ would require $12 \times 2.52 \times 10^9 \times 2 \approx 60.4\text{ GB}$ of VRAM—before accounting for gradients or activations.

By grouping pixels into non-overlapping square patches of side length $P$, ViT reduces the sequence length by a factor of $P^2$ and reduces the quadratic attention matrix size by a factor of $P^4$:

$$N = \frac{H}{P} \cdot \frac{W}{P} = \frac{H W}{P^2}, \qquad N^2 = \frac{(H W)^2}{P^4}$$

For $H = W = 224$ and $P = 16$, $P^2 = 256$ and $P^4 = 65{,}536$:

2. Patchify Math & Exact Equivalence to Stride-$P$ Conv2d

Let the input image be $\mathbf{I} \in \mathbb{R}^{H \times W \times C}$ where $H$ and $W$ are divisible by the patch size $P$. We partition $\mathbf{I}$ into $N = (H/P)(W/P)$ non-overlapping spatial blocks and flatten each $P \times P \times C$ block into a row vector $\mathbf{x}_p^{(i)} \in \mathbb{R}^{P^2 C}$ for $i \in \{1, \dots, N\}$:

$$\mathbf{I} \in \mathbb{R}^{H \times W \times C} \xrightarrow{\text{reshape \& flatten}} \mathbf{x}_p = \begin{bmatrix} \text{vec}(\mathbf{I}_{1,1})^\top \\ \text{vec}(\mathbf{I}_{1,2})^\top \\ \vdots \\ \text{vec}(\mathbf{I}_{H/P, W/P})^\top \end{bmatrix} \in \mathbb{R}^{N \times (P^2 C)}$$

For ViT-Base/16 (`ViT-B/16`) on a $224 \times 224 \times 3$ image:

Why Every Production Codebase Uses nn.Conv2d(C, d, kernel_size=P, stride=P)

Explicitly slicing an image tensor into patches via `einops.rearrange` or `unfold` and multiplying by $E \in \mathbb{R}^{(P^2 C) \times d}$ requires materializing an intermediate patch buffer in memory. However, consider what a 2D convolution with kernel size $P \times P$ and stride $P$ computes at spatial grid coordinates $(r, c) \in \{0, \dots, H/P - 1\} \times \{0, \dots, W/P - 1\}$ for output channel $k \in \{1, \dots, d\}$:

$$[\text{Conv2d}(\mathbf{I})]_{r, c, k} = \sum_{u=0}^{P-1} \sum_{v=0}^{P-1} \sum_{ch=0}^{C-1} W_{k, ch, u, v} \, \mathbf{I}_{rP + u,\; cP + v,\; ch} + b_k = \mathbf{x}_p^{(r \cdot \frac{W}{P} + c + 1)} E_{:, k} + b_k$$

Because the stride $P$ equals the kernel width $P$, the convolutional receptive fields are strictly non-overlapping and tile the image identically to the patch grid! Reshaping the 4D convolutional output $[B, d, H/P, W/P] \to [B, d, N] \to [B, N, d]$ yields the exact same tensor as flattening patches and applying `nn.Linear(P*P*C, d)`, while leveraging highly optimized cuDNN implicit-GEMM kernels without allocating an intermediate `unfold` buffer.

Worked Numerical Example (Toy $4 \times 4$ Single-Channel Image, $P = 2$, $d = 3$)

Suppose $H = W = 4$, $C = 1$, and patch size $P = 2$, giving $N = (4/2)(4/2) = 4$ patches, each of flattened length $P^2 C = 4$:

$$\mathbf{I} = \begin{bmatrix} 1 & 2 & \big| & 0 & 1 \\ 3 & 4 & \big| & 1 & 2 \\ \hline 2 & 0 & \big| & 3 & 1 \\ 1 & 1 & \big| & 0 & 2 \end{bmatrix} \implies \mathbf{x}_p = \begin{bmatrix} 1 & 2 & 3 & 4 \\ 0 & 1 & 1 & 2 \\ 2 & 0 & 1 & 1 \\ 3 & 1 & 0 & 2 \end{bmatrix} \in \mathbb{R}^{4 \times 4}$$

With projection matrix $E \in \mathbb{R}^{4 \times 3}$ and zero bias:

$$E = \begin{bmatrix} 1 & 0 & -1 \\ 0 & 1 & 1 \\ 1 & -1 & 0 \\ 0 & 1 & 1 \end{bmatrix} \implies \mathbf{x}_p E = \begin{bmatrix} 1(1)+3(1) & 2(1)+3(-1)+4(1) & 1(-1)+2(1)+4(1) \\ 0+1 & 1-1+2 & 0+1+2 \\ 2+1 & 0-1+1 & -2+0+1 \\ 3+0 & 1-0+2 & -3+1+2 \end{bmatrix} = \begin{bmatrix} 4 & 3 & 5 \\ 1 & 2 & 3 \\ 3 & 0 & -1 \\ 3 & 3 & 0 \end{bmatrix} \in \mathbb{R}^{4 \times 3}$$

3. The Learnable [CLS] Token & 1D/2D Positional Embeddings

Following BERT (Devlin et al., 2019), Dosovitskiy et al. prepend a learnable class token $\mathbf{x}_{\text{cls}} \in \mathbb{R}^{1 \times d}$ to the sequence of $N$ projected patch embeddings. Why prepend a dedicated synthetic token instead of just picking one of the $N$ patch tokens at the output?

After prepending $\mathbf{x}_{\text{cls}}$, the sequence has length $N + 1$. Because self-attention is permutation-equivariant—permuting the $N$ patches in the input simply permutes the output rows without changing any attention weights—spatial geometry must be injected explicitly via a learnable positional embedding table $\mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times d}$:

$$\mathbf{z}_0 = \begin{bmatrix} \mathbf{x}_{\text{cls}} \\ \mathbf{x}_p^{(1)} E \\ \mathbf{x}_p^{(2)} E \\ \vdots \\ \mathbf{x}_p^{(N)} E \end{bmatrix} + \mathbf{E}_{\text{pos}}, \qquad E \in \mathbb{R}^{(P^2 C) \times d}, \quad \mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times d}$$

Why Does 1D Learnable Position Embedding Match 2D-Aware Embeddings in ViT?

Dosovitskiy et al. (Appendix D.4) ablated four positional encoding strategies on ImageNet:

  1. No positional embedding: Treating the image as an unordered "bag of $16 \times 16$ patches." Accuracy drops by more than $3\%$, proving spatial arrangement matters beyond patch-internal structure.
  2. 1D learnable positional embedding ($\mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times d}$): Raster-scan indexing $0, 1, \dots, N$.
  3. 2D-aware learnable positional embedding: Separate learnable row embeddings $\mathbf{E}_{\text{row}} \in \mathbb{R}^{(H/P) \times (d/2)}$ and column embeddings $\mathbf{E}_{\text{col}} \in \mathbb{R}^{(W/P) \times (d/2)}$ concatenated into $\mathbb{R}^d$ at each grid cell $(r, c)$.
  4. Relative 2D positional bias: Encoding pairwise grid distances $(\Delta r, \Delta c)$ directly inside attention logits.

Surprisingly, 1D learnable, 2D-aware, and relative 2D embeddings all achieved virtually identical accuracy (within $0.2\%$). Why? At $14 \times 14 = 196$ patches, there are only $14$ rows and $14$ columns—so learning the 2D grid topology from scratch requires discovering that index $i + 1$ is the horizontal neighbor and index $i + 14$ is the vertical neighbor. Visualizing the cosine similarity matrix $\mathbf{E}_{\text{pos}}[1:] \mathbf{E}_{\text{pos}}[1:]^\top \in \mathbb{R}^{196 \times 196}$ of a trained 1D ViT reveals unmistakable 2D periodic stripes: each patch's learned 1D position vector has highest cosine similarity with patches in the same row and the same column!

4. The $L$-Layer Pre-Norm Transformer Encoder Stack

Unlike Vaswani et al. (2017), which placed LayerNorm after the residual addition (Post-Norm), ViT uses the Pre-Norm formulation (Wang et al., 2019; Radford et al., 2019) where LayerNorm is applied to the input of each sub-layer while the identity residual highway remains completely unobstructed from $\mathbf{z}_0$ to $\mathbf{z}_L$:

$$\begin{aligned} \mathbf{z}'_\ell &= \mathbf{z}_{\ell-1} + \operatorname{MSA}\!\left(\operatorname{LN}(\mathbf{z}_{\ell-1})\right), &&\ell = 1, \dots, L \\[4pt] \mathbf{z}_\ell &= \mathbf{z}'_\ell + \operatorname{MLP}\!\left(\operatorname{LN}(\mathbf{z}'_\ell)\right), &&\ell = 1, \dots, L \\[4pt] \mathbf{y} &= \operatorname{LN}\!\left(\mathbf{z}_L^0\right) \in \mathbb{R}^d \end{aligned}$$

Inside each Multi-Head Self-Attention ($\operatorname{MSA}$) block with $H$ heads of dimension $d_h = d / H$, let $\tilde{\mathbf{z}} = \operatorname{LN}(\mathbf{z}_{\ell-1}) \in \mathbb{R}^{(N+1) \times d}$:

$$[\mathbf{Q}_h, \mathbf{K}_h, \mathbf{V}_h] = \tilde{\mathbf{z}} W_{qkv}^{(h)} \in \mathbb{R}^{(N+1) \times 3 d_h}, \qquad \mathbf{A}_h = \operatorname{softmax}\!\left(\frac{\mathbf{Q}_h \mathbf{K}_h^\top}{\sqrt{d_h}}\right) \in \mathbb{R}^{(N+1) \times (N+1)}$$
$$\operatorname{MSA}(\tilde{\mathbf{z}}) = \operatorname{Concat}\!\left(\mathbf{A}_1 \mathbf{V}_1, \dots, \mathbf{A}_H \mathbf{V}_H\right) W_{\text{out}}, \qquad \operatorname{MLP}(\mathbf{u}) = \operatorname{GELU}(\mathbf{u} W_1 + \mathbf{b}_1) W_2 + \mathbf{b}_2$$

Crucially, unlike autoregressive language models (GPT), ViT applies bidirectional (unmasked) self-attention: every patch token can attend to every other patch token (top-left corner to bottom-right corner) starting in Layer 1.

5. Patch Size Scaling ($P = 32, 16, 14$) & 2D Bicubic Positional Interpolation

Notice that the parameter count of the Transformer Encoder ($L$ blocks of $\operatorname{MSA}$ and $\operatorname{MLP}$) depends only on the hidden dimension $d$ and layer count $L$, not on the patch size $P$ or image resolution $H \times W$. Shrinking $P$ slightly decreases the patch stem parameters ($P^2 C d$), yet explodes the compute FLOPs because token length $N = (H/P)(W/P)$ grows quadratically as $P$ shrinks:

Impact of patch size $P$ and image resolution $H \times W$ on sequence length $N+1$ and attention matrix complexity for a ViT-Base ($d=768, L=12$) backbone.
Configuration Input Resolution Patch Grid $(H/P) \times (W/P)$ Sequence Length $N+1$ Attention Pairs $(N+1)^2$ Relative Compute
ViT-B/32 $224 \times 224$ $7 \times 7 = 49$ $50$ $2{,}500$ $1.0\times$ (fastest)
ViT-B/16 $224 \times 224$ $14 \times 14 = 196$ $197$ $38{,}809$ $\sim 4.4\times$
ViT-B/14 $224 \times 224$ $16 \times 16 = 256$ $257$ $66{,}049$ $\sim 6.1\times$
ViT-B/16 @ $384\text{px}$ $384 \times 384$ $24 \times 24 = 576$ $577$ $332{,}929$ $\sim 16.5\times$

Fine-Tuning at Higher Resolution via 2D Bicubic Interpolation of $\mathbf{E}_{\text{pos}}$

In practice, pretraining happens at $224 \times 224$ for throughput, followed by fine-tuning at a higher resolution $H' \times W'$ (such as $384 \times 384$ in supervised ViT, or $336 \times 336$ in CLIP ViT-L/14@336px). When we keep patch size $P$ fixed and increase image resolution from $(H, W)$ to $(H', W')$:

  1. The patch projection `Conv2d(3, d, kernel_size=P, stride=P)` and all $L$ Transformer blocks operate on per-token dimension $d$, so their weights require zero shape changes.
  2. The patch count increases from $N = (H/P)(W/P)$ to $N' = (H'/P)(W'/P)$ (e.g., for $P=16$ from $224^2 \to 384^2$, the grid grows from $14 \times 14 = 196$ to $24 \times 24 = 576$ patches).
  3. The pretrained positional embedding $\mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times d}$ only has $197$ rows, whereas we now need $577$ rows. To preserve the learned 2D geometry:
    • Separate the [CLS] position vector $\mathbf{e}_{\text{pos}}^{\text{cls}} = \mathbf{E}_{\text{pos}}[0:1, :] \in \mathbb{R}^{1 \times d}$.
    • Reshape the remaining $N$ patch position vectors $\mathbf{E}_{\text{pos}}[1:, :] \in \mathbb{R}^{N \times d}$ back into a 2D spatial grid of shape $\left[\frac{H}{P}, \frac{W}{P}, d\right]$ (e.g., $14 \times 14 \times 768$).
    • Apply 2D bicubic interpolation along the two spatial axes to resize $\left[\frac{H}{P}, \frac{W}{P}, d\right] \to \left[\frac{H'}{P}, \frac{W'}{P}, d\right]$ (e.g., $24 \times 24 \times 768$).
    • Flatten back to $\mathbb{R}^{N' \times d}$ and concatenate $\mathbf{e}_{\text{pos}}^{\text{cls}}$ at index $0$ to obtain $\tilde{\mathbf{E}}_{\text{pos}} \in \mathbb{R}^{(N'+1) \times d}$.

6. Inductive Biases: Why CNNs Win on Small Data and ViT Wins at Scale

Why did Dosovitskiy et al. find that ViT slightly underperformed ResNets (BiT) when trained from scratch on mid-sized ImageNet-1k ($1.28\text{M}$ images), matched ResNets on ImageNet-21k ($14\text{M}$ images), and decisively surpassed ResNets when pretrained on JFT-300M ($300\text{M}$ images) or WIT-400M (CLIP)?

7. Clean PyTorch Implementation of Standard ViT + 2D Pos Interpolation

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

class VisionTransformer(nn.Module):
    """
    Standard Vision Transformer (Dosovitskiy et al., ICLR 2021)
    with dynamic 2D bicubic positional embedding interpolation.
    """
    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_rate: float = 0.1,
    ):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = img_size // patch_size
        self.num_patches = self.grid_size * self.grid_size

        # 1. Patchify + Linear Projection fused as a stride-P Conv2d
        self.patch_embed = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True
        )

        # 2. Learnable [CLS] token and 1D Positional Embeddings
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)

        # 3. L Pre-Norm Transformer Encoder Blocks
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=int(embed_dim * mlp_ratio),
            dropout=drop_rate,
            activation="gelu",
            batch_first=True,
            norm_first=True,  # Pre-Norm: LN before MSA and MLP
        )
        self.blocks = nn.TransformerEncoder(encoder_layer, num_layers=depth)

        # 4. Final LayerNorm on [CLS] and Classification Head
        self.norm = nn.LayerNorm(embed_dim, eps=1e-6)
        self.head = nn.Linear(embed_dim, num_classes)

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

    def interpolate_pos_encoding(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
        """2D bicubic interpolation of patch position embeddings when H, W != img_size."""
        N = x.shape[1] - 1
        if N == self.num_patches and H == W:
            return self.pos_embed

        cls_pos = self.pos_embed[:, :1]       # [1, 1, d]
        patch_pos = self.pos_embed[:, 1:]     # [1, N_orig, d]
        d = x.shape[-1]
        h_new, w_new = H // self.patch_size, W // self.patch_size

        # Reshape [1, g, g, d] -> permute [1, d, g, g] -> bicubic resize -> [1, h_new*w_new, d]
        patch_pos = patch_pos.reshape(1, self.grid_size, self.grid_size, d).permute(0, 3, 1, 2)
        patch_pos = F.interpolate(
            patch_pos, size=(h_new, w_new), mode="bicubic", align_corners=False
        )
        patch_pos = patch_pos.permute(0, 2, 3, 1).flatten(1, 2)
        return torch.cat((cls_pos, patch_pos), dim=1)

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

        # Prepend [CLS] token: [B, 1, d] + [B, N, d] -> [B, N+1, d]
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # Add (interpolated) positional embeddings
        x = self.pos_drop(x + self.interpolate_pos_encoding(x, H, W))

        # L Pre-Norm Transformer blocks
        x = self.blocks(x)

        # Readout from [CLS] at index 0
        cls_out = self.norm(x[:, 0])
        return self.head(cls_out)