Vision Transformer Patch Embedding

The problem

A transformer reads a sequence of vectors. To read an image, the ViT (Vision Transformer; Dosovitskiy et al., 2021) cuts it into square patches, flattens each patch into a vector, and treats the patches as tokens, like words in a sentence. You will write the cutting and the embedding.

For images of shape $(B, H, W, C)$ (batch, height, width, channels) and patch size $P$ (dividing $H$ and $W$):

1. Cut each image into $N = (H/P)(W/P)$ patches of $P \times P$ pixels, in row-major order: left to right, then top to bottom.
2. Flatten each patch into $P \cdot P \cdot C$ numbers, in the order patch[row, col, channel].
3. Project with a learned matrix $E$ of shape $(P^2 C, D)$: each patch becomes a vector of $D$ numbers.
4. Prepend the learned class embedding (shape $(D,)$) and add the position embeddings (shape $(N + 1, D)$).

Example. A toy image of $4 \times 4$ pixels, one channel, with the pixels numbered $0, \dots, 15$ row by row, and $P = 2$:

 0  1  2  3
 4  5  6  7
 8  9 10 11
12 13 14 15

There are $N = 4$ patches of $4$ numbers each. In row-major order: [0, 1, 4, 5] (top-left), [2, 3, 6, 7] (top-right), [8, 9, 12, 13] (bottom-left), [10, 11, 14, 15] (bottom-right). The result of patchify has shape $(1, 4, 4)$.

A real image of $224 \times 224$ pixels with 3 channels and $P = 16$ gives $14 \cdot 14 = 196$ patches of $16 \cdot 16 \cdot 3 = 768$ numbers, and 197 tokens with the class token.

Cutting without loops

With one reshape and one transpose:

1. x.reshape(B, H // P, P, W // P, P, C): the axes are now (batch, patch row, pixel row inside the patch, patch column, pixel column inside the patch, channel).
2. transpose(0, 1, 3, 2, 4, 5): bring the two patch-grid axes together, then the two pixel axes: (batch, patch row, patch column, pixel row, pixel column, channel).
3. reshape(B, N, P * P * C): flatten the grid into a list of patches and each patch into a vector.

What to implement

• patchify(images, P): shape $(B, N, P^2 C)$.
• vit_embed(images, P, E, cls, pos): shape $(B, N + 1, D)$.

Requirements

• No Python loops.
• Raise ValueError if $P$ does not divide $H$ and $W$.