Vision Transformers

Build the mental model

From the Deep Learning course, you already know that transformers operate on a sequence of tokens — each one a vector — attending to every other token to build contextual representations. Text has an obvious tokenization: words or subwords, in order. Images don't have an inherent sequence structure; pixels are arranged on a 2D grid, and treating each pixel as a token would produce sequences of tens of thousands of elements, making self-attention (which is quadratic in sequence length) computationally infeasible. The Vision Transformer's key insight is to tokenize at the level of patches instead of pixels: chop the image into a grid of, say, 16x16 pixel squares, and treat each square as one token. A 224x224 image becomes a 14x14 grid of patches — 196 tokens, a completely manageable sequence length for a standard transformer encoder.

The clever implementation trick is that 'cut into patches and linearly project each patch to an embedding vector' is exactly what a convolution with kernel_size equal to the patch size and matching stride does in a single operation. A Conv2d with kernel_size=16, stride=16 slides a 16x16 window across the image without overlap, and each window position produces one output vector of length embed_dim — that vector is the patch's token embedding, and the convolution's learned weights are the linear projection. Reshape the resulting (batch, embed_dim, grid_h, grid_w) feature map into (batch, num_patches, embed_dim) and you have exactly the token sequence a transformer encoder expects, ready for the same multi-head self-attention and feed-forward blocks you already know from text models — just with a positional embedding added to preserve 2D spatial information that a plain sequence would otherwise lose.

Connect it to a real scenario

The Tutorial Platform could use a ViT-style patch-embedding step to auto-tag lesson screenshots by content type — code editor, terminal output, architecture diagram, or browser UI — because splitting each screenshot into a 16x16 patch sequence and feeding it to a lightweight transformer classifier captures layout patterns (a terminal's monospace grid, a diagram's sparse whitespace and connecting lines) that a single global CNN feature vector tends to blur together.

• Self-attention over raw pixels is computationally infeasible, so ViT tokenizes images at the patch level instead, drastically shrinking the sequence length.
• A single Conv2d with kernel_size and stride both set to the patch size performs 'cut into patches + linearly project' in one operation.
• The resulting (batch, num_patches, embed_dim) tensor is a token sequence that a standard transformer encoder can consume directly, just like text embeddings.

Try the working example

import torch
import torch.nn as nn

patch_size = 16
embed_dim = 768
image_size = 224

patch_embed = nn.Conv2d(
    in_channels=3,
    out_channels=embed_dim,
    kernel_size=patch_size,
    stride=patch_size,
)

images = torch.randn(2, 3, image_size, image_size)  # fake batch of 2 RGB images
patches = patch_embed(images)                # (batch, embed_dim, H/patch, W/patch)
print(patches.shape)

tokens = patches.flatten(2).transpose(1, 2)  # (batch, num_patches, embed_dim)
print(tokens.shape)

torch.Size([2, 768, 14, 14])
torch.Size([2, 196, 768]) — the convolution turns each 16x16 pixel patch into one 768-dimensional vector, producing a 14x14 grid of patches (since 224/16=14), which is then flattened and transposed into a sequence of 196 tokens (14*14) each of dimension 768 — exactly the (batch, num_patches, embed_dim) shape a transformer encoder expects.

Wikipedia — Vision transformer — Computer Vision