ViT Patch Embeddings - PyTorch Basics

ViT Patch Embeddings

What you'll build: patch_embed(...), the input layer of a Vision Transformer: cut an image into square patches, flatten and project each one to a token, prepend a learnable CLS token and add positional embeddings.

Background

A transformer consumes a sequence of vectors. ViT (Dosovitskiy et al., 2020) turns an image into one by treating each P x P patch as a "word": a 224x224 image with P = 16 becomes 196 tokens. The extra CLS token gathers information from all patches through attention, and its final state feeds the classifier. Positional embeddings tell the model where each patch came from, since attention itself ignores order.

Algorithm

For an image (C, H, W) with H and W divisible by P:

1. Split it into num_patches = (H / P) * (W / P) patches of shape (C, P, P), ordered row-major: left to right along the top row of patches, then the next row.
2. Flatten each patch in (C, P, P) order (channel first, then row, then column) to a vector of length C * P * P.
3. Project: tokens = patches @ W_proj, shape (num_patches, embed_dim).
4. Prepend the CLS token: [cls_token; tokens], shape (num_patches + 1, embed_dim).
5. Add the positional embeddings: + pos_embed.

Steps 1 to 3 are exactly a Conv2d(C, embed_dim, kernel_size=P, stride=P), which is how real ViT code implements them.

Inputs

• image: float array (C, H, W).
• patch_size: P, which divides H and W.
• embed_dim: D, the token size (equal to W_proj.shape[1]).
• W_proj: (C * P * P, D); cls_token: (D,); pos_embed: (num_patches + 1, D).

Expected Return Format

An array of shape (num_patches + 1, D). Row 0 is cls_token + pos_embed[0]; row 1 + i * (W / P) + j comes from the patch in patch-row i, patch-column j.

Constraints

• Vectorise the patch extraction with reshape and transpose (loops over patches are acceptable in NumPy, but slow).
• PyTorch track: use reshape/permute (or Tensor.unfold); don't use nn.Conv2d.

Follow-up

• How many tokens does a 224x224 image give with P = 14, and how does attention cost scale with P?
• How do you use a model trained at 224x224 on 384x384 images, given pos_embed has a fixed length?

Input: image (3, 8, 8), patch_size=4, embed_dim=16, W_proj (48, 16), cls_token (16,), pos_embed (5, 16)
Output: shape (5, 16)

Input: patch_embed(np.zeros((1, 4, 4)), 2, 3, np.ones((4, 3)), np.array([1.0, 2.0, 3.0]), np.full((5, 3), 0.5))
Output: array([[1.5, 2.5, 3.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]])

import numpy as np

def patch_embed(image, patch_size, embed_dim, W_proj, cls_token, pos_embed):
    """
    Create ViT patch embeddings.

    Args:
        image: Input image (C, H, W)
        patch_size: Size P of each square patch (divides H and W)
        embed_dim: Dimension D of the embeddings
        W_proj: Projection matrix (C*P*P, D)
        cls_token: CLS token embedding (D,)
        pos_embed: Positional embeddings (num_patches+1, D)

    Returns:
        Embeddings (num_patches + 1, D): CLS first, then patches in row-major order
    """
    # Your code here
    pass