Vision & Multimodal Transformers

Images are sliced into a grid of non-overlapping patches, projected into linear embeddings, and processed through standard transformer blocks alongside text tokens.

Vision transformers split images into patch embeddings that feed standard multi-head self-attention encoders alongside multimodal text tokens.

Why Does This Exist?

For nearly a decade, Convolutional Neural Networks (CNNs) dominated computer vision through hard-coded inductive biases: translation equivariance (a feature detector produces the same response regardless of where an object appears) and two-dimensional spatial locality (pixels close together matter much more than pixels far apart). While these biases allow CNNs to learn sample-efficiently on modest datasets like ImageNet with 1.2 million images, they impose rigid computational constraints. A $3 \times 3$ kernel cannot connect pixels on opposite edges of an image without stacking dozens of layers, diluting semantic information across repeated downsampling operations.

When industrial pretraining reached tens or hundreds of millions of images, these spatial inductive biases transformed from helpful guardrails into performance ceilings. Natural language processing had already converged on the standard Transformer architecture, which possesses almost no locality bias and scales compute and parameters with empirical power-law predictability.

Vision Transformers (ViT) eliminated custom vision architectures by discarding convolutional filters entirely. An image is cut into a fixed grid of square patches, each patch is flattened and linearly projected into a vector embedding, and the resulting sequence is handed directly to a vanilla Transformer encoder. By converting spatial pixels into token sequences, the exact same attention engine can simultaneously process image patches and text tokens, unlocking modern unified vision-language foundation models. Prerequisite understanding relies on backpropagation across residual layers.

Think of It Like This

A mosaic artist reading tiles as a sentence

Imagine a stained-glass window depicting a harbor scene. A convolutional inspector examines the window through a tiny magnifying glass, sliding it two centimeters at a time across the glass, noting local edges and color gradients, and then repeating the scan at lower resolution.

A Vision Transformer does something completely different. It cuts the window into a clean $14 \times 14$ grid of square glass tiles. It numbers each tile from 1 to 196 to preserve its original position on the wall, and lays them out in a straight line on a long examination table.

Beside the tiles, the artist places a blank notebook labeled [CLS]. Because the artist possesses global vision across the entire table, tile 1 (the morning sun in the top-left corner) can immediately exchange light and color information with tile 190 (the ship's reflection in the bottom water) in a single glance. No tile is forced to communicate solely with its immediate physical neighbor. When reading text captions alongside the window, words and glass tiles sit on the exact same table, allowing the word "mast" to attend directly to tile 84 containing the ship's rigging.

How It Actually Works

Patch Extraction and Linear Projection

Given an input RGB image $x \in \mathbb{R}^{H \times W \times C}$ where $H$ is image height, $W$ is width, and $C = 3$ is color channels, the model divides the spatial canvas into non-overlapping square patches of size $P \times P$ . The total number of patches, which forms the effective sequence length $N$ , is:

$N = \frac{H \cdot W}{P^2}$

Each patch $x_p^i \in \mathbb{R}^{P^2 C}$ is flattened into a one-dimensional vector of dimension $P^2 C$ . A learnable linear projection matrix $E \in \mathbb{R}^{(P^2 C) \times D}$ maps each flattened patch into the transformer's latent model dimension $D$ :

$z_0 = \left[ x_{\text{class}};\, x_p^1 E;\, x_p^2 E;\, \dots;\, x_p^N E \right] + E_{\text{pos}}$

where:

$x_{\text{class}} \in \mathbb{R}^{1 \times D}$ is a prepended learnable classification token whose state at the output layer serves as the aggregate image representation.
$E_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}$ is a learnable 1D position embedding matrix added element-wise to retain spatial ordering, since standard self-attention is permutation-invariant.

The combined token sequence $z_0 \in \mathbb{R}^{(N+1) \times D}$ passes through $L$ identical Transformer encoder layers. Each layer consists of Multi-Head Self-Attention ( $\text{MHSA}$ ) and a Multi-Layer Perceptron ( $\text{MLP}$ ) with GeLU activations, wrapped in pre-LayerNorm ( $\text{LN}$ ) operations and residual connections:

$z_\ell' = \text{MHSA}(\text{LN}(z_{\ell-1})) + z_{\ell-1}, \quad \ell = 1, \dots, L$

$z_\ell = \text{MLP}(\text{LN}(z_\ell')) + z_\ell', \quad \ell = 1, \dots, L$

In multimodal architectures such as CLIP, ALIGN, or Flamingo, visual patch tokens and textual subword tokens are projected into a shared semantic space. Cross-attention layers allow text tokens to attend directly to visual patch representations:

$\text{Attention}(Q_v, K_t, V_t) = \text{softmax}\left(\frac{Q_v K_t^T}{\sqrt{d_k}}\right) V_t$

where visual queries $Q_v$ retrieve semantic attributes from textual keys $K_t$ and values $V_t$ .

Worked Example

Consider a standard ViT-Base processing a single image with resolution $H = 224$ , $W = 224$ , $C = 3$ , patch size $P = 16$ , and embedding dimension $D = 768$ :

Patch count and input vector sizing: $N = \frac{224 \times 224}{16 \times 16} = \frac{50176}{256} = 196 \text{ patches}$ Each individual patch contains $16 \times 16 \times 3 = 768$ raw pixel values.
Patch projection: Projection matrix $E$ has shape $(768, 768)$ . Multiplying 196 flattened vectors of length 768 by $E$ yields a patch embedding tensor of shape $(196, 768)$ .
Prepend class token and position embeddings: Adding the single learnable vector $x_{\text{class}}$ produces sequence length $196 + 1 = 197$ . The resulting tensor has shape $(197, 768)$ . The 1D position embedding matrix $E_{\text{pos}}$ of shape $(197, 768)$ is added element-wise.
Self-attention matrix computation: With 12 attention heads, each head operates on key-query dimension $d_k = 768 / 12 = 64$ . For each head, the attention map $A = \text{softmax}(Q K^T / \sqrt{64})$ has dimensions $(197, 197)$ . Every forward pass computes $197 \times 197 = 38,809$ attention coefficients per head, allowing every patch to observe all 195 other patches across the canvas.

Code

1import torch2import torch.nn as nn3
4class PatchEmbedding(nn.Module):5    """Slices an image into non-overlapping patches and projects to embedding dim."""6    def __init__(self, img_size: int = 224, patch_size: int = 16, in_channels: int = 3, embed_dim: int = 768):7        super().__init__()8        self.num_patches = (img_size // patch_size) ** 29        # A 2D convolution with kernel_size=patch_size and stride=patch_size performs10        # patch extraction and linear projection in a single optimized GPU operation11        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)12
13    def forward(self, x: torch.Tensor) -> torch.Tensor:14        # x: (batch_size, 3, 224, 224)15        x = self.proj(x)  # -> (batch_size, 768, 14, 14)16        x = x.flatten(2)  # -> (batch_size, 768, 196)17        return x.transpose(1, 2)  # -> (batch_size, 196, 768)18
19class MinimalViT(nn.Module):20    def __init__(self, img_size: int = 224, patch_size: int = 16, embed_dim: int = 768, num_classes: int = 1000):21        super().__init__()22        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels=3, embed_dim=embed_dim)23        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))24        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + self.patch_embed.num_patches, embed_dim))25        26        encoder_layer = nn.TransformerEncoderLayer(27            d_model=embed_dim, nhead=12, dim_feedforward=3072, activation="gelu", batch_first=True28        )29        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)30        self.head = nn.Linear(embed_dim, num_classes)31
32    def forward(self, x: torch.Tensor) -> torch.Tensor:33        b = x.shape[0]34        tokens = self.patch_embed(x)  # (b, 196, 768)35        cls_tokens = self.cls_token.expand(b, -1, -1)  # (b, 1, 768)36        x = torch.cat((cls_tokens, tokens), dim=1)  # (b, 197, 768)37        x = x + self.pos_embed  # (b, 197, 768)38        encoded = self.encoder(x)  # (b, 197, 768)39        logits = self.head(encoded[:, 0])  # Take [CLS] output -> (b, 1000)40        return logits41
42# Verification43model = MinimalViT()44dummy_image = torch.randn(2, 3, 224, 224)45output = model(dummy_image)46print("Output shape:", output.shape)47# -> Output shape: torch.Size([2, 1000])

Watch Out For

Interpolating position embeddings across input resolutions

Practitioners commonly pretrain a Vision Transformer on $224 \times 224$ images and then fine-tune on higher resolutions like $384 \times 384$ or $448 \times 448$ to boost fine-grained classification accuracy. However, increasing image dimensions from $224$ to $384$ increases patch count from $196$ to $(384/16)^2 = 576$ .

Directly loading pretrained weights throws a runtime shape mismatch error because $E_{\text{pos}}$ was instantiated with length $197$ , while the model now requires $577$ positions. The standard fix is bicubic interpolation: separate the $[CLS]$ embedding from the spatial grid embeddings, reshape the remaining $196$ vectors into a $14 \times 14 \times D$ grid, interpolate spatially to a $24 \times 24 \times D$ grid, flatten back to $576$ vectors, and re-attach the $[CLS]$ token. Skipping spatial 2D grid reshaping and performing naive 1D linear interpolation destroys the two-dimensional spatial geometric relationships learned during pretraining.

The Quick Version

Vision Transformers bypass convolutional kernels by slicing an image into a regular 2D grid of non-overlapping patches ( $16 \times 16$ ), projecting each into an embedding vector.
A learnable $[CLS]$ token and 1D position embeddings are added to the patch sequence before passing through standard Multi-Head Self-Attention layers.
Self-attention enables immediate global context across any pair of patches regardless of spatial distance, scaling effectively to multimodal architectures where vision and text share a unified sequence representation.