Vision Transformers and Patch Embeddings

Transformers were built for sequences of discrete tokens — words in a sentence. Images are neither discrete nor naturally sequential; a photo is a dense grid of pixels with no obvious "next" pixel. For years the assumption was that images needed a fundamentally different architecture, convolutional neural networks, built around the idea that nearby pixels matter more than distant ones. The Vision Transformer's central move questions that assumption: instead of designing a vision-specific architecture, force the image into the shape a transformer already understands — a sequence of tokens — by chopping it into patches, and let self-attention figure out spatial relationships from data, the same way it figures out word relationships in a sentence.

The analogy: reading an image like a jigsaw puzzle laid out as a sentence

Imagine cutting a photograph into a 4×4 grid of 16 square pieces, laying them left to right, top to bottom, and treating each as a single word in a sentence. A person can still recognise the photo from these 16 tokens, using the same contextual reasoning as with a scrambled sentence: piece 3 (sky) relates to piece 4 (also sky) more than to piece 12 (grass), even though nothing in the raw pixels says "these are adjacent" once laid out as a flat list. A Vision Transformer relies on exactly this: patches as tokens, with self-attention learning which relate to which, entirely from data rather than a hard-coded notion of neighbourhood.

From pixels to tokens

An image of size $H \times W \times C$ (height, width, channels) is split into a grid of $N$ non-overlapping patches, each $P \times P \times C$ pixels:

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

In words: the number of patches equals total pixel area divided by the area of one patch.

Each flattened patch (a vector of $P^2 \times C$ raw pixel values) is passed through a single learned linear layer — the patch embedding — projecting it into the same $d$-dimensional space every other token lives in:

$$e_i = W_e\, x_i^{\text{flat}} + b_e$$

In words: each patch's raw pixel vector $x_i^{\text{flat}}$ is multiplied by a shared learned weight matrix $W_e$ and shifted by a bias $b_e$ — exactly the same linear-embedding step a word embedding does for a token ID. The model never sees raw pixels again after this, only $d$-dimensional vectors. A learned position embedding is added to each patch so the model knows where in the grid it came from (self-attention has no inherent sense of order), and one extra learnable "classification" token is prepended, whose final output represents the whole image for downstream classification.

Worked example 1: patch count and embedding size for a real image

A 224×224×3 image, patch size 16×16: $N = 224\times224/16^2 = 196$ patches. Each patch flattens to $16\times16\times3 = 768$ raw numbers. With embedding dimension $d=768$, the patch-embedding matrix $W_e$ is $768\times768$ — one learned linear map applied identically to every patch. The transformer then processes a sequence of $196+1=197$ tokens (patches plus the classification token), exactly like a 197-word sentence.

Worked example 2: how attention connects distant patches

Take a 4×4-patch toy image (16 patches, numbered 1–16). Patch 1 (top-left) and patch 16 (bottom-right) are as far apart as two patches can be. In a convolutional network, information from patch 1 needs several stacked local layers to influence patch 16's representation. In a Vision Transformer, patch 1 and patch 16 are just two tokens in one sequence of 16 — the very first self-attention layer lets patch 16's query directly score against patch 1's key, regardless of spatial distance. Whether that link is actually useful is learned from data, but the architectural capacity for direct interaction exists from layer one, unlike a convolution's fixed local receptive field.

Patch embedding is itself just one linear transform applied to every patch — drag the matrix above to see how a linear map stretches and rotates a space; $W_e$ does the same, just in far higher dimensions.

What this means in practice

Vision Transformers now match or beat convolutional networks on large-scale image classification, but need substantially more training data to get there, because convolutions come with a built-in assumption (locality, translation invariance) that a plain transformer has to learn instead of receiving for free. This trade — less built-in structure, more reliance on data and scale — recurs across ML, and the same patch-based sequencing idea underlies PatchTST and Patched Time-Series Transformers, where "patches" are contiguous chunks of a time series instead of pixels.

Practice

1. For a 384×384 image with patch size 32×32, how many patches result, and how does that compare to using patch size 16×16 on the same image?
2. Why does removing the position embedding from a Vision Transformer make the model unable to distinguish a normal image from the same image with its patches randomly shuffled?