Patch Embeddings Step by Step

The patch arithmetic

With an image of size H x W and a patch size P, a ViT produces (H / P) x (W / P) patches. A 224 x 224 image with 16-pixel patches gives 14 x 14 = 196 patches; with 14-pixel patches it gives 16 x 16 = 256. Each patch has P x P x 3 numbers, which are flattened and multiplied by a learned matrix (the patch embedding, equivalent to a convolution with kernel size and stride P) to produce a vector of the model width, for example 768. Many ViTs add a special [CLS] token whose final state summarises the whole image, and every token gets a position embedding because attention itself has no notion of order. The number of patches grows with the square of the resolution: doubling height and width gives four times as many tokens, and self-attention cost grows roughly with the square of the token count. That is the root of most cost problems with high-resolution images.

Tracing the patch embedding

Follow the shapes. Each step changes what the array means, not just its size.

import numpy as np

rng = np.random.default_rng(0)
img = rng.random((224, 224, 3), dtype=np.float32)
P = 16                                   # patch size
h, w, c = img.shape

patches = (img.reshape(h // P, P, w // P, P, c)
              .transpose(0, 2, 1, 3, 4)
              .reshape(-1, P * P * c))
print(patches.shape)                     # (num_patches, pixels_per_patch)

W = rng.standard_normal((P * P * c, 768), dtype=np.float32) * 0.02
tokens = patches @ W                     # linear patch embedding
cls = np.zeros((1, 768), dtype=np.float32)
seq = np.concatenate([cls, tokens])
print(seq.shape)                         # [CLS] + 196 patch tokens

1. Line 4: A random 224x224 RGB image stands in for a real photo, stored height-width-channels.
2. Line 5: The patch size decides the token count: 224 / 16 = 14 patches per side.
3. Line 8: reshape splits each axis into (14 blocks, 16 pixels): rows and columns of patches appear.
4. Line 9: transpose brings the two block axes together so each patch's pixels are contiguous.
5. Line 10: Flatten: 196 patches, each a vector of 16*16*3 = 768 raw pixel values.
6. Line 14: A learned linear map projects each patch to the model width; here also 768.
7. Line 16: Prepending a [CLS] token gives the final sequence the transformer receives.

Run it

The same code with its real output: 196 patch vectors, then 197 tokens including [CLS].

import numpy as np

rng = np.random.default_rng(0)
img = rng.random((224, 224, 3), dtype=np.float32)
P = 16                                   # patch size
h, w, c = img.shape

patches = (img.reshape(h // P, P, w // P, P, c)
              .transpose(0, 2, 1, 3, 4)
              .reshape(-1, P * P * c))
print(patches.shape)                     # (num_patches, pixels_per_patch)

W = rng.standard_normal((P * P * c, 768), dtype=np.float32) * 0.02
tokens = patches @ W                     # linear patch embedding
cls = np.zeros((1, 768), dtype=np.float32)
seq = np.concatenate([cls, tokens])
print(seq.shape)                         # [CLS] + 196 patch tokens

Output:

(196, 768)
(197, 768)

Patch size 14 vs 16

Model names often encode this: vit-base-patch16-224 means 16-pixel patches at 224 resolution; clip-vit-large-patch14-336 means 14-pixel patches at 336, giving 24 x 24 = 576 patches. Read the name before estimating token counts.

Quick check: A ViT uses 14-pixel patches on a 336 x 336 image. How many patch tokens does it produce (excluding [CLS])?

• 576
• 256
• 196
• 1,024