{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": "# Lesson 46: Vision Transformers (ViT)\n\nLesson 45's Transformer encoder operates on a sequence of vectors. An image isn't a sequence — so the **Vision Transformer** (\"An Image is Worth 16x16 Words\", Dosovitskiy et al., 2020★) turns it into one: chop the image into fixed-size patches, flatten and linearly embed each patch into a vector, and feed the resulting sequence straight into an ordinary Transformer encoder, exactly as built in Lesson 45. No convolution anywhere. This lesson builds that pipeline, and then runs the exact translation-generalization test from Lesson 33 to show precisely what a Transformer gives up by dropping convolution's built-in inductive bias."
},
{
"cell_type": "code",
"id": "80542b81",
"source": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "2c42e54a",
"source": "## Patchify: turning an image into a sequence\n\nSplit a `16x16` image into `4x4` patches, giving a sequence of 16 patches, each flattened to a 16-number vector. This is exactly equivalent to a strided convolution with a kernel the same size as the stride (`nn.Conv2d(..., kernel_size=P, stride=P)`) — both slide a non-overlapping window and apply the same linear map to each — which is how ViT's patch embedding is usually implemented in practice.",
"metadata": {}
},
{
"cell_type": "code",
"id": "0b2a28a0",
"source": "def patchify(img, patch_size):\n B, C, H, W = img.shape\n P = patch_size\n patches = img.unfold(2, P, P).unfold(3, P, P) # (B, C, H/P, W/P, P, P)\n patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)\n return patches.reshape(B, -1, C * P * P) # (B, num_patches, C*P*P)\n\nimg = torch.randn(2, 1, 16, 16)\npatches = patchify(img, patch_size=4)\nprint(f'image {tuple(img.shape)} -> patch sequence {tuple(patches.shape)} '\n f'({(16 // 4) ** 2} patches of {1 * 4 * 4} values each)')\n\npatch_dim, embed_dim = 16, 8\nlinear_embed = nn.Linear(patch_dim, embed_dim)\nconv_embed = nn.Conv2d(1, embed_dim, kernel_size=4, stride=4)\nwith torch.no_grad():\n conv_embed.weight.copy_(linear_embed.weight.view(embed_dim, 1, 4, 4))\n conv_embed.bias.copy_(linear_embed.bias)\n\nout_linear = linear_embed(patches)\nout_conv = conv_embed(img).flatten(2).transpose(1, 2)\nprint(f'linear-on-patches vs. strided-conv max abs diff: {(out_linear - out_conv).abs().max().item():.2e}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "b2abe808",
"source": "## The [CLS] token, and assembling a full ViT\n\nTo turn a sequence of patch embeddings into a single whole-image prediction, ViT prepends one extra, learned \"classification\" token to the sequence before running the Transformer encoder (Lesson 45) — after attention has let it gather information from every patch, that one token's final output is what the classification head reads. Positional encoding (Lesson 45) is added so patch order (i.e. patch *location*) isn't invisible to attention.",
"metadata": {}
},
{
"cell_type": "code",
"id": "d269aebd",
"source": "def positional_encoding(T, D):\n pos = torch.arange(T).unsqueeze(1).float()\n i = torch.arange(D).unsqueeze(0).float()\n angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))\n angles = pos * angle_rates\n pe = torch.zeros(T, D)\n pe[:, 0::2] = torch.sin(angles[:, 0::2])\n pe[:, 1::2] = torch.cos(angles[:, 1::2])\n return pe\n\nclass TinyViT(nn.Module):\n def __init__(self, img_size=16, patch_size=4, in_ch=1, embed_dim=32, n_heads=4, n_layers=2):\n super().__init__()\n self.patch_size = patch_size\n n_patches = (img_size // patch_size) ** 2\n patch_dim = in_ch * patch_size * patch_size\n self.embed = nn.Linear(patch_dim, embed_dim)\n self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)\n self.register_buffer('pos_enc', positional_encoding(n_patches + 1, embed_dim))\n layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,\n batch_first=True, dropout=0.0)\n self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)\n self.head = nn.Linear(embed_dim, 1)\n\n def forward(self, x):\n p = patchify(x, self.patch_size)\n tok = self.embed(p)\n cls = self.cls_token.expand(x.shape[0], -1, -1)\n tok = torch.cat([cls, tok], dim=1) + self.pos_enc\n out = self.encoder(tok)\n return self.head(out[:, 0]).squeeze(-1) # classify from the [CLS] token's output\n\nmodel = TinyViT()\nx = torch.randn(3, 1, 16, 16)\nout = model(x)\nprint(f'TinyViT: input {tuple(x.shape)} -> output {tuple(out.shape)} (one logit per image)')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "a3bf1a12",
"source": "## What convolution's inductive bias buys you\n\nLesson 33 trained a CNN and a flatten-based MLP on plus-vs-circle shapes at one range of positions, then tested both at *unseen* positions — the CNN generalized (global max pooling makes it exactly translation-invariant); the flatten-based MLP didn't (its first layer's weights are tied to absolute pixel coordinates). Rerun that exact experiment, swapping in `TinyViT` for the MLP.\n\nA ViT has no built-in translation invariance either — patch embedding and positional encoding are both position-specific, same as the MLP's flattened input. The chief practical mechanism ViTs use to compensate is scale: pretraining on enormous datasets lets attention *learn* something like translation invariance from data, rather than getting it for free from the architecture the way a CNN does. On the tiny dataset in this course, expect that learning to fail.",
"metadata": {}
},
{
"cell_type": "code",
"id": "a0f686af",
"source": "def make_image(shape_type, cx, cy, size=16):\n img = np.zeros((size, size), dtype=np.float32)\n if shape_type == 'plus':\n img[cy-1:cy+2, cx-3:cx+4] = 1.0\n img[cy-3:cy+4, cx-1:cx+2] = 1.0\n else:\n yy, xx = np.mgrid[0:size, 0:size]\n img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0\n return img\n\ndef make_dataset(rng_local, n, position_range):\n imgs, labels = [], []\n for _ in range(n):\n shape_type = rng_local.choice(['plus', 'circle'])\n cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)\n imgs.append(make_image(shape_type, cx, cy))\n labels.append(0.0 if shape_type == 'plus' else 1.0)\n return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.float32)\n\ndata_rng = np.random.default_rng(1)\nX_train, y_train = make_dataset(data_rng, 300, (5, 11)) # training positions\nX_test, y_test = make_dataset(data_rng, 150, (3, 5)) # UNSEEN positions\n\nclass CNNClassifier(nn.Module):\n def __init__(self):\n super().__init__()\n self.conv = nn.Sequential(\n nn.Conv2d(1, 8, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),\n nn.Conv2d(8, 16, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1))\n self.fc = nn.Linear(16, 1)\n\n def forward(self, x):\n return self.fc(self.conv(x).flatten(1)).squeeze(-1)\n\ndef train_and_eval(model_cls, Xtr, ytr, Xte, yte, seed, epochs=200, lr=0.001):\n torch.manual_seed(seed)\n model = model_cls()\n opt = torch.optim.Adam(model.parameters(), lr=lr)\n for _ in range(epochs):\n opt.zero_grad()\n loss = F.binary_cross_entropy_with_logits(model(Xtr), ytr)\n loss.backward()\n opt.step()\n with torch.no_grad():\n train_acc = ((model(Xtr) > 0).float() == ytr).float().mean().item()\n test_acc = ((model(Xte) > 0).float() == yte).float().mean().item()\n return train_acc, test_acc\n\nXtr_t = torch.tensor(X_train).unsqueeze(1); ytr_t = torch.tensor(y_train)\nXte_t = torch.tensor(X_test).unsqueeze(1); yte_t = torch.tensor(y_test)\n\nfor seed in range(3):\n vit_train, vit_test = train_and_eval(TinyViT, Xtr_t, ytr_t, Xte_t, yte_t, seed=seed)\n cnn_train, cnn_test = train_and_eval(CNNClassifier, Xtr_t, ytr_t, Xte_t, yte_t, seed=seed)\n print(f'seed={seed}: ViT train={vit_train:.1%} test(unseen positions)={vit_test:.1%} '\n f'| CNN train={cnn_train:.1%} test(unseen positions)={cnn_test:.1%}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "dbc06358",
"source": "Both models fit the training positions perfectly, but the ViT collapses to chance-level (or worse) on unseen positions every time, while the CNN generalizes perfectly — reproducing Lesson 33's MLP-vs-CNN gap almost exactly, with the ViT standing in for the MLP. This is the well-documented empirical finding behind real ViTs: they need either far more training data, or heavy data augmentation, or a hybrid architecture that reintroduces some convolutional structure, to match a CNN's data efficiency on small-to-medium datasets — because a CNN's translation invariance is a hard architectural guarantee, while a ViT's has to be learned from examples. At the scale of hundreds of millions of images, ViTs match or beat CNNs handily; at the scale of hundreds of images, they don't.\n\n### Exercise\n\n1. Increase the training set size in `make_dataset(data_rng, 300, ...)` from 300 to 3000 images at the same training positions. Does the ViT's unseen-position accuracy improve substantially, partially, or not at all — and what does that suggest about *how much* more data a ViT needs to compensate for its missing inductive bias?\n2. Reduce `patch_size` from 4 to 2 (finer patches, longer sequence). Does finer patching help the ViT generalize better to unseen positions, or is the failure mode unrelated to patch resolution?\n3. A \"hybrid\" architecture feeds a small CNN's feature map (instead of raw pixel patches) into a Transformer encoder, gaining some of the CNN's spatial inductive bias back. Sketch how you'd modify `TinyViT.forward` to patchify the *output* of a small `nn.Conv2d` stack instead of the raw image.",
"metadata": {}
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}