Vision Transformer › Worked Example: Tracing a Patch Through ViT

To solidify the mechanics, let's trace a single patch from raw pixels all the way to the transformer encoder output, working through the dimensions explicitly.

Consider an RGB image of size $224 \times 224 \times 3$ and a model configuration of patch size $P = 16$, embedding dimension $D = 768$, 12 transformer layers, and 12 attention heads. We focus on the patch at grid position $(r=0, c=1)$, the second patch in the first row.

Step 1: Extract the patch. The patch at $(r=0, c=1)$ corresponds to pixels $[0:16, 16:32, :]$ in the image tensor. Slicing this region gives a sub-tensor of shape $(16, 16, 3)$.

Step 2: Flatten. We flatten this sub-tensor into a 1D vector: $(16, 16, 3) \to (768,)$. In Python: patch_vec = patch.reshape(-1) yields a vector $\mathbf{x}^{(2)} \in \mathbb{R}^{768}$. (We index from 1, so $(r=0, c=1)$ is patch 2 in the sequence.)

Step 3: Linear projection. The patch embedding matrix $\mathbf{E} \in \mathbb{R}^{768 \times 768}$ transforms the raw patch vector to an embedding: $\mathbf{z}_0^{(2)} = \mathbf{x}^{(2)} \mathbf{E} \in \mathbb{R}^{768}$.

Step 4: Add position embedding. The sequence index of this patch is $i = 0 \times 14 + 1 = 1$ in 0-indexed patch coordinates, which becomes position 2 in the full sequence (after the CLS token at position 0). We add the position embedding: $\mathbf{z}_0^{(2)} \leftarrow \mathbf{z}_0^{(2)} + \mathbf{e}_{\text{pos}, 2}$.

Step 5: Transformer processing. The full sequence $\mathbf{z}_0 \in \mathbb{R}^{197 \times 768}$ (196 patches plus 1 CLS token) passes through 12 transformer layers. In each layer, self-attention is computed across all 197 tokens. Our patch at position 2 can receive information from the CLS token (position 0) and all other 195 patches. After 12 layers, the representation $\mathbf{z}_{12}^{(2)} \in \mathbb{R}^{768}$ reflects a complex mixture of information from all spatial locations, weighted by learned attention patterns.

Step 6: Classification. The final representation of the CLS token, $\mathbf{z}_{12}^{(0)} \in \mathbb{R}^{768}$, is passed to the linear classification head. Our individual patch representation $\mathbf{z}_{12}^{(2)}$ is not directly used for classification but contributed to the CLS token's representation through attention in all 12 layers.

This trace shows that the path from raw pixels to a classification decision passes through a sequence of linear transformations and attention-based mixing operations, all operating in the same $D$-dimensional space. No special-purpose vision operations are required after the initial patching step.