{ "cells": [ { "cell_type": "markdown", "id": "ea9aeaf3", "metadata": {}, "source": "# Lesson 59: Visual Question Answering (VQA)\n\nLesson {{lesson:vision_language_models}}'s CLIP matched a *fixed* set of text prompts against an image. **Visual Question Answering (VQA)** asks for something more flexible: given an image and an arbitrary natural-language question about it, produce an answer. This is the mechanism behind modern multimodal assistants that can look at a picture and answer questions about it — a generative-VLM capability, distinct from CLIP's contrastive matching. This lesson builds the classic (and still widely used) simplification: treat VQA as classification over a small, closed vocabulary of possible answers, fusing an image encoder and a question encoder before the final prediction." }, { "cell_type": "code", "execution_count": null, "id": "f80c1b03", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import matplotlib.pyplot as plt" ] }, { "cell_type": "markdown", "id": "0c4fd57c", "metadata": {}, "source": [ "## Scenes with (question, answer) pairs\n", "\n", "Scenes contain 0-3 plus/circle shapes. Each scene gets one of four question templates: \"what shape is in the image\" (only meaningful when exactly one shape is present), \"how many shapes are there\", \"is there a circle in the image\", \"is there a plus in the image\" — with the true answer computed directly from the scene's ground truth. The answer space is a small, fixed vocabulary: shape names, counts, and yes/no." ] }, { "cell_type": "code", "execution_count": null, "id": "aaa2ce35", "metadata": {}, "outputs": [], "source": "SIZE = 32\n\nQUESTION_VOCAB = ['', 'what', 'shape', 'is', 'this', 'how', 'many', 'shapes', 'are', 'there',\n 'a', 'circle', 'plus', 'in', 'the', 'image']\nQ2ID = {w: i for i, w in enumerate(QUESTION_VOCAB)}\nANSWER_VOCAB = ['plus', 'circle', '0', '1', '2', '3', 'yes', 'no']\nA2ID = {w: i for i, w in enumerate(ANSWER_VOCAB)}\n\ndef make_shape(img, shape_type, cx, cy, r=4):\n yy, xx = np.mgrid[0:img.shape[0], 0:img.shape[1]]\n if shape_type == 'plus':\n img[cy-1:cy+2, cx-r:cx+r+1] = 1.0\n img[cy-r:cy+r+1, cx-1:cx+2] = 1.0\n else:\n img[((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2] = 1.0\n return img\n\ndef make_scene(rng, size=SIZE, max_shapes=3):\n n = rng.integers(0, max_shapes + 1)\n img = np.zeros((size, size), dtype=np.float32)\n shapes = []\n tries = 0\n while len(shapes) < n and tries < 30:\n tries += 1\n cx, cy = rng.integers(6, size - 6), rng.integers(6, size - 6)\n if any(abs(cx - sx) < 9 and abs(cy - sy) < 9 for sx, sy, _ in shapes):\n continue\n st = rng.choice(['plus', 'circle'])\n make_shape(img, st, cx, cy)\n shapes.append((cx, cy, st))\n img = np.clip(img + rng.normal(0, 0.03, img.shape), 0, 1).astype(np.float32)\n return img, shapes\n\ndef make_qa(rng, shapes):\n n_circle = sum(1 for _, _, t in shapes if t == 'circle')\n n_plus = sum(1 for _, _, t in shapes if t == 'plus')\n qtype = rng.choice(['what_shape', 'count', 'is_there_circle', 'is_there_plus'])\n if qtype == 'what_shape':\n words = ['what', 'shape', 'is', 'in', 'the', 'image']\n answer = shapes[0][2] if len(shapes) == 1 else str(min(len(shapes), 3))\n elif qtype == 'count':\n words = ['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image']\n answer = str(len(shapes))\n elif qtype == 'is_there_circle':\n words = ['is', 'there', 'a', 'circle', 'in', 'the', 'image']\n answer = 'yes' if n_circle > 0 else 'no'\n else:\n words = ['is', 'there', 'a', 'plus', 'in', 'the', 'image']\n answer = 'yes' if n_plus > 0 else 'no'\n # qtype and the raw shapes are returned too (unused by the closed-vocabulary classifier\n # below) because the generative decoder later in this lesson needs them to build full\n # phrase answers like \"there are two shapes\" instead of a single class id.\n return qtype, words, [Q2ID[w] for w in words], A2ID[answer]\n\ndef pad_question(ids, max_len=8):\n return ids + [0] * (max_len - len(ids)) # 0 = , distinct from every real word\n\nrng = np.random.default_rng(4)\nN = 1000\nimgs, questions, answers, question_words, scene_shapes, question_types = [], [], [], [], [], []\nfor _ in range(N):\n img, shapes = make_scene(rng)\n qtype, words, q_ids, a_id = make_qa(rng, shapes)\n imgs.append(img)\n questions.append(pad_question(q_ids))\n answers.append(a_id)\n question_words.append(words)\n scene_shapes.append(shapes)\n question_types.append(qtype)\nimgs = np.array(imgs, dtype=np.float32)\nquestions = np.array(questions, dtype=np.int64)\nanswers = np.array(answers, dtype=np.int64)\n\nsplit = int(0.85 * N)\nXtr, Qtr, Atr = imgs[:split], questions[:split], answers[:split]\nXte, Qte, Ate = imgs[split:], questions[split:], answers[split:]\nshapes_tr, shapes_te = scene_shapes[:split], scene_shapes[split:]\nqtypes_tr, qtypes_te = question_types[:split], question_types[split:]\n\nfig, axes = plt.subplots(1, 4, figsize=(9, 2.5))\nfor ax, im, words, a in zip(axes, imgs[:4], question_words[:4], answers[:4]):\n ax.imshow(im, cmap='gray'); ax.axis('off')\n ax.set_title(f'{\" \".join(words)}?\\n-> {ANSWER_VOCAB[a]}', fontsize=7)\nplt.show()" }, { "cell_type": "markdown", "id": "304afe47", "metadata": {}, "source": "## Fusing an image encoder and a question encoder\n\nBoth encoders are small Transformers (Lesson {{lesson:vision_transformers}}'s ViT pattern for the image, a small bidirectional encoder over the question's tokens) mapping to the same-size vector, same pattern as Lesson {{lesson:vision_language_models}}'s CLIP. The difference is what happens next: CLIP compared image and text embeddings with a dot product to check *if they match*. VQA needs to actually *combine* them into one joint representation and classify — a standard, simple fusion is elementwise multiplication (the image vector and the question vector gate each other), followed by a small classifier head over the fixed answer vocabulary." }, { "cell_type": "code", "execution_count": null, "id": "ecddd06d", "metadata": {}, "outputs": [], "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)\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)\n\ndef 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 ImageEncoder(nn.Module):\n \"\"\"A small ViT (Lesson {{lesson:vision_transformers}}'s pattern): patchify, embed,\n prepend a [CLS] token, run a Transformer encoder, read the [CLS] token's output.\"\"\"\n def __init__(self, img_size=SIZE, patch_size=4, out_dim=32, 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 = 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.proj = nn.Linear(embed_dim, out_dim)\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.proj(out[:, 0])\n\nclass QuestionEncoder(nn.Module):\n \"\"\"A small bidirectional Transformer encoder over the question's word tokens, with\n padding masked out of both attention and the final pooling.\"\"\"\n def __init__(self, vocab_size, out_dim=32, embed_dim=32, n_heads=4, n_layers=1, max_len=8):\n super().__init__()\n self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)\n self.register_buffer('pos_enc', positional_encoding(max_len, 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.proj = nn.Linear(embed_dim, out_dim)\n\n def forward(self, ids):\n pad_mask = (ids == 0) # True where padding\n x = self.embed(ids) + self.pos_enc[:ids.shape[1]]\n x = self.encoder(x, src_key_padding_mask=pad_mask)\n valid = (~pad_mask).float().unsqueeze(-1)\n pooled = (x * valid).sum(1) / valid.sum(1).clamp(min=1)\n return self.proj(pooled)\n\nclass VQAModel(nn.Module):\n def __init__(self, vocab_size, n_answers, dim=32):\n super().__init__()\n self.img_enc = ImageEncoder(out_dim=dim)\n self.q_enc = QuestionEncoder(vocab_size, out_dim=dim)\n self.head = nn.Sequential(nn.Linear(dim, 32), nn.ReLU(), nn.Linear(32, n_answers))\n\n def forward(self, img, q_ids):\n fused = self.img_enc(img) * self.q_enc(q_ids)\n return self.head(fused)\n\ntorch.manual_seed(0)\nmodel = VQAModel(len(QUESTION_VOCAB), len(ANSWER_VOCAB))\nopt = torch.optim.Adam(model.parameters(), lr=0.003)\nXt = torch.tensor(Xtr).unsqueeze(1); Qt = torch.tensor(Qtr); At = torch.tensor(Atr)\nn = len(Xtr)\nfor epoch in range(300):\n idx = np.random.default_rng(epoch).permutation(n)[:128] # mini-batches: a Transformer encoder makes a full-batch pass too slow to repeat hundreds of times\n opt.zero_grad()\n loss = F.cross_entropy(model(Xt[idx], Qt[idx]), At[idx])\n loss.backward()\n opt.step()\n\nwith torch.no_grad():\n preds = model(torch.tensor(Xte).unsqueeze(1), torch.tensor(Qte)).argmax(1).numpy()\nacc = (preds == Ate).mean()\n\nmost_common = np.bincount(Atr).argmax()\nbaseline_acc = (Ate == most_common).mean()\n\nprint(f'VQA test accuracy: {acc:.1%}')\nprint(f'baseline (always predict the most common training answer, \"{ANSWER_VOCAB[most_common]}\"): {baseline_acc:.1%}')" }, { "cell_type": "markdown", "id": "d27a4875", "metadata": {}, "source": [ "## The real test: does the same image get different answers to different questions?\n", "\n", "A model that's secretly ignoring the question and just pattern-matching the image would give the same answer regardless of what's asked. Pose all four question types against one fixed test image and check." ] }, { "cell_type": "code", "execution_count": null, "id": "06acc6c9", "metadata": {}, "outputs": [], "source": [ "test_rng = np.random.default_rng(123)\n", "demo_img, demo_shapes = make_scene(test_rng)\n", "while len(demo_shapes) != 2 or demo_shapes[0][2] == demo_shapes[1][2]:\n", " demo_img, demo_shapes = make_scene(test_rng) # find a scene with one of each shape\n", "\n", "demo_questions = [\n", " ['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image'],\n", " ['is', 'there', 'a', 'circle', 'in', 'the', 'image'],\n", " ['is', 'there', 'a', 'plus', 'in', 'the', 'image'],\n", "]\n", "img_t = torch.tensor(demo_img[None, None])\n", "print(f'true scene contents: {[s[2] for s in demo_shapes]}')\n", "with torch.no_grad():\n", " for words in demo_questions:\n", " q_ids = torch.tensor([pad_question([Q2ID[w] for w in words])])\n", " pred = model(img_t, q_ids).argmax(1).item()\n", " print(f' \"{\" \".join(words)}?\" -> {ANSWER_VOCAB[pred]}')" ] }, { "cell_type": "markdown", "id": "e1b89946", "metadata": {}, "source": "The same fixed image produces three different, individually correct answers depending on what's asked — the fusion step is genuinely reading both inputs, not just memorizing image-to-answer shortcuts.\n\n## From closed-vocabulary classification to real generative VQA\n\nTreating VQA as classification over a small fixed answer set (as above) was the dominant approach for years and is still used when the answer space genuinely is small and known in advance. It breaks down the moment an answer needs to be open-ended text — \"describe what's unusual about this image,\" or an answer word that never appeared in training. Modern VQA and image-captioning systems solve this the way Lesson {{lesson:transformer_architecture}} promised they would: replace the classification head with an **autoregressive text decoder** — causal self-attention over previously generated words, cross-attention into the image features — that generates the answer one token at a time, exactly the mechanism behind large multimodal assistants that can describe, question, and reason about an image in free-form language rather than picking from a fixed list. Lesson {{lesson:decoder_only_language_models}} already built the causal-self-attention half of that recipe and defined self-attention vs. cross-attention precisely; the rest of this lesson adds the cross-attention half, completing the picture with the third and last Transformer model family." }, { "cell_type": "markdown", "id": "ebc44dce", "source": "### Phrase answers instead of single labels\n\nInstead of one label per answer, generate a short word sequence: `\"it is a circle\"`, `\"no there is no plus\"`, `\"there are two shapes\"`. Every phrase starts with a `` (\"beginning of sequence\") token and ends with `` (\"end of sequence\"), exactly the convention real language models use to know where to start generating and when to stop.", "metadata": {} }, { "cell_type": "code", "id": "0d54c01c", "source": "NUM_WORDS = {0: 'zero', 1: 'one', 2: 'two', 3: 'three'}\n\ndef answer_phrase(qtype, shapes):\n n = len(shapes)\n n_circle = sum(1 for _, _, t in shapes if t == 'circle')\n n_plus = sum(1 for _, _, t in shapes if t == 'plus')\n if qtype == 'what_shape':\n if n == 1:\n return ['it', 'is', 'a', shapes[0][2]]\n elif n == 0:\n return ['there', 'is', 'no', 'shape']\n else:\n return ['there', 'are', NUM_WORDS[min(n, 3)], 'shapes']\n elif qtype == 'count':\n if n == 0:\n return ['there', 'are', 'no', 'shapes']\n elif n == 1:\n return ['there', 'is', 'one', 'shape']\n else:\n return ['there', 'are', NUM_WORDS[min(n, 3)], 'shapes']\n elif qtype == 'is_there_circle':\n return ['yes', 'there', 'is', 'a', 'circle'] if n_circle > 0 else ['no', 'there', 'is', 'no', 'circle']\n else:\n return ['yes', 'there', 'is', 'a', 'plus'] if n_plus > 0 else ['no', 'there', 'is', 'no', 'plus']\n\nANSWER_WORDS = ['', '', '', 'it', 'is', 'a', 'there', 'are', 'no', 'shape', 'shapes',\n 'yes', 'circle', 'plus', 'zero', 'one', 'two', 'three']\nAW2ID = {w: i for i, w in enumerate(ANSWER_WORDS)}\nMAX_LEN = 20 # bos + up to 5 words + eos, padded well past the longest real phrase -- see the KV-cache timing section below for why\n\ndef tokenize_answer(words):\n ids = [AW2ID['']] + [AW2ID[w] for w in words] + [AW2ID['']]\n return ids + [AW2ID['']] * (MAX_LEN - len(ids))\n\nAtrP = np.array([tokenize_answer(answer_phrase(qt, sh)) for qt, sh in zip(qtypes_tr, shapes_tr)], dtype=np.int64)\nAteP = np.array([tokenize_answer(answer_phrase(qt, sh)) for qt, sh in zip(qtypes_te, shapes_te)], dtype=np.int64)\n\nfor i in range(3):\n words = [ANSWER_WORDS[t] for t in AtrP[i] if t not in (AW2ID[''], AW2ID[''], AW2ID[''])]\n print(f'\"{\" \".join(question_words[i])}?\" -> \"{\" \".join(words)}\"')", "metadata": {}, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "ac843df8", "source": "### A causal, cross-attending decoder — built explicitly, so it can be cached later\n\nThis decoder combines the two attention patterns Lesson {{lesson:decoder_only_language_models}} distinguished: **causal self-attention** over the answer generated so far (each word can only look at itself and earlier words — it can't peek at a word it hasn't generated yet, exactly Lesson {{lesson:decoder_only_language_models}}'s decoder-only block), followed by **cross-attention into the image features** (the fused image-and-question embeddings from above, treated as a two-token \"memory\" — queries from the answer tokens, keys and values from this separate memory, the defining feature Lesson {{lesson:decoder_only_language_models}} deliberately left out). This is the classic encoder-decoder captioning pattern: a separate encoder (the image/question encoders above) produces a fixed representation, and the decoder repeatedly attends back into it while generating.\n\nLesson {{lesson:decoder_only_language_models}} already used explicit Q/K/V projections instead of `nn.MultiheadAttention` for its own decoder, for the same reason this one does: generating a sequence one token at a time can either recompute attention over the whole growing sequence at every step, or **cache** the keys and values already computed for earlier tokens and only do the work for the new one, which needs direct access to the keys and values that `nn.MultiheadAttention` doesn't expose. This decoder now has two attention sub-layers that could each, in principle, be cached — but only the self-attention one actually needs to be: the cross-attention memory is the same fixed two tokens at every generation step, so its keys/values can simply be computed once and reused, with nothing to \"grow\" the way the self-attention cache does.", "metadata": {} }, { "cell_type": "code", "id": "b5262e1d", "source": "class ImageEncoder2(nn.Module):\n \"\"\"Same ViT pattern as this lesson's earlier ImageEncoder -- a separate copy so it gets\n its own trained weights, independent of the closed-vocabulary classifier's encoder.\"\"\"\n def __init__(self, out_dim=32):\n super().__init__()\n self.patch_size = 4\n n_patches = (SIZE // self.patch_size) ** 2\n patch_dim = self.patch_size * self.patch_size\n embed_dim = 32\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, 4, dim_feedforward=embed_dim * 2, batch_first=True, dropout=0.0)\n self.encoder = nn.TransformerEncoder(layer, num_layers=2)\n self.proj = nn.Linear(embed_dim, out_dim)\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.proj(out[:, 0])\n\nclass QuestionEncoder2(nn.Module):\n def __init__(self, vocab_size, out_dim=32):\n super().__init__()\n embed_dim = 32\n self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)\n self.register_buffer('pos_enc', positional_encoding(8, embed_dim))\n layer = nn.TransformerEncoderLayer(embed_dim, 4, dim_feedforward=embed_dim * 2, batch_first=True, dropout=0.0)\n self.encoder = nn.TransformerEncoder(layer, num_layers=1)\n self.proj = nn.Linear(embed_dim, out_dim)\n\n def forward(self, ids):\n pad_mask = (ids == 0)\n x = self.embed(ids) + self.pos_enc[:ids.shape[1]]\n x = self.encoder(x, src_key_padding_mask=pad_mask)\n valid = (~pad_mask).float().unsqueeze(-1)\n pooled = (x * valid).sum(1) / valid.sum(1).clamp(min=1)\n return self.proj(pooled)\n\nclass CachedDecoder(nn.Module):\n \"\"\"A from-scratch causal self-attention + cross-attention decoder block with explicit\n K/V projections, so generation can optionally reuse cached keys/values instead of\n recomputing them -- see the KV-cache section below.\"\"\"\n def __init__(self, vocab_size, dim=32, n_heads=4, max_len=MAX_LEN):\n super().__init__()\n self.dim, self.n_heads = dim, n_heads\n self.tok_embed = nn.Embedding(vocab_size, dim, padding_idx=0)\n self.pos_embed = nn.Embedding(max_len, dim)\n self.q_proj = nn.Linear(dim, dim); self.k_proj = nn.Linear(dim, dim); self.v_proj = nn.Linear(dim, dim)\n self.out_proj = nn.Linear(dim, dim)\n self.cross_q = nn.Linear(dim, dim); self.cross_k = nn.Linear(dim, dim); self.cross_v = nn.Linear(dim, dim)\n self.cross_out_proj = nn.Linear(dim, dim)\n self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))\n self.ln1 = nn.LayerNorm(dim); self.ln2 = nn.LayerNorm(dim); self.ln3 = nn.LayerNorm(dim)\n self.answer_head = nn.Linear(dim, vocab_size)\n\n def _mha(self, q, k, v, causal_offset=None):\n B, Tq, D = q.shape\n Tk = k.shape[1]\n H, hd = self.n_heads, D // self.n_heads\n q = q.view(B, Tq, H, hd).transpose(1, 2)\n k = k.view(B, Tk, H, hd).transpose(1, 2)\n v = v.view(B, Tk, H, hd).transpose(1, 2)\n scores = (q @ k.transpose(-2, -1)) / hd ** 0.5\n if causal_offset is not None:\n # position i (in the new query chunk) may attend to key position j only if\n # j <= i + causal_offset -- the usual causal mask, shifted to account for\n # however many earlier tokens' keys are already sitting in the cache.\n i = torch.arange(Tq).unsqueeze(1) + causal_offset\n j = torch.arange(Tk).unsqueeze(0)\n scores = scores.masked_fill(j > i, float('-inf'))\n weights = scores.softmax(dim=-1)\n return (weights @ v).transpose(1, 2).reshape(B, Tq, D)\n\n def forward(self, tokens_in, memory, past_kv=None):\n \"\"\"past_kv=None recomputes everything from scratch (used for training, and as the\n 'naive' generation baseline). Passing a (k, v) tuple returned by an earlier call\n appends only the new token's key/value instead of recomputing the whole prefix --\n the actual KV cache.\"\"\"\n B, T = tokens_in.shape\n start_pos = 0 if past_kv is None else past_kv[0].shape[1]\n pos = torch.arange(start_pos, start_pos + T)\n x = self.tok_embed(tokens_in) + self.pos_embed(pos)[None]\n k_new, v_new = self.k_proj(x), self.v_proj(x)\n if past_kv is not None:\n k = torch.cat([past_kv[0], k_new], dim=1)\n v = torch.cat([past_kv[1], v_new], dim=1)\n else:\n k, v = k_new, v_new\n attn_out = self._mha(self.q_proj(x), k, v, causal_offset=start_pos)\n x = self.ln1(x + self.out_proj(attn_out))\n cross_out = self._mha(self.cross_q(x), self.cross_k(memory), self.cross_v(memory), causal_offset=None)\n x = self.ln2(x + self.cross_out_proj(cross_out))\n x = self.ln3(x + self.ff(x))\n return self.answer_head(x), (k, v)\n\ntorch.manual_seed(0)\nimg_enc2 = ImageEncoder2(32)\nq_enc2 = QuestionEncoder2(len(QUESTION_VOCAB), 32)\ndecoder = CachedDecoder(len(ANSWER_WORDS), dim=32)\nopt2 = torch.optim.Adam(list(img_enc2.parameters()) + list(q_enc2.parameters()) + list(decoder.parameters()), lr=0.002)\n\nXt2, Qt2, At2 = torch.tensor(Xtr).unsqueeze(1), torch.tensor(Qtr), torch.tensor(AtrP)\nn2 = len(Xtr)\nfor epoch in range(800):\n idx = np.random.default_rng(epoch).permutation(n2)[:256] # mini-batches, same reason as the classifier above\n memory = torch.stack([img_enc2(Xt2[idx]), q_enc2(Qt2[idx])], dim=1) # (B, 2, dim): the image features + question features\n dec_in, target = At2[idx][:, :-1], At2[idx][:, 1:] # teacher forcing: predict each next word from the true prefix\n logits, _ = decoder(dec_in, memory)\n loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target.reshape(-1), ignore_index=AW2ID[''])\n opt2.zero_grad(); loss.backward(); opt2.step()\n\nprint(f'final training loss: {loss.item():.4f}')", "metadata": {}, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "40494530", "source": "### Generating an answer, one token at a time\n\nStart from just ``, feed it through the decoder, take the most likely next word, append it, and repeat — this is **greedy decoding**. The `use_cache` flag switches between the two generation strategies: with it off, every step reruns the decoder over the *entire* sequence generated so far; with it on, every step only processes the single newest token and reuses the keys/values already cached from earlier steps. Both must produce identical output, since caching changes nothing about *what* gets computed, only how much of it gets recomputed — that's checked explicitly below before trusting the speed comparison that follows.", "metadata": {} }, { "cell_type": "code", "id": "934ba61e", "source": "def generate(imgs_t, q_t, use_cache):\n B = imgs_t.shape[0]\n memory = torch.stack([img_enc2(imgs_t), q_enc2(q_t)], dim=1)\n tokens = torch.full((B, 1), AW2ID[''], dtype=torch.long)\n past_kv = None\n for _ in range(MAX_LEN - 1):\n if use_cache:\n logits, past_kv = decoder(tokens[:, -1:], memory, past_kv=past_kv)\n else:\n logits, _ = decoder(tokens, memory, past_kv=None)\n logits = logits[:, -1:]\n next_tok = logits[:, -1, :].argmax(-1, keepdim=True)\n tokens = torch.cat([tokens, next_tok], dim=1)\n return tokens\n\ndef trim_at_eos(seq):\n seq = seq.tolist()\n return seq[:seq.index(AW2ID['']) + 1] if AW2ID[''] in seq else seq\n\nwith torch.no_grad():\n Xte2, Qte2 = torch.tensor(Xte).unsqueeze(1), torch.tensor(Qte)\n gen_cached = generate(Xte2, Qte2, use_cache=True)\n gen_uncached = generate(Xte2, Qte2, use_cache=False)\n\n match_rate = (gen_cached == gen_uncached).float().mean().item()\n print(f'cached vs. uncached generation, token-for-token agreement: {match_rate:.4f} (should be 1.0000)')\n\n correct = sum(trim_at_eos(g) == trim_at_eos(a) for g, a in zip(gen_cached, torch.tensor(AteP)))\n print(f'exact-phrase accuracy on held-out test scenes: {correct}/{len(AteP)} = {correct / len(AteP):.1%}')\n\n for i in range(4):\n words = [ANSWER_WORDS[t] for t in trim_at_eos(gen_cached[i]) if t not in (AW2ID[''], AW2ID[''])]\n print(f' \"{\" \".join(question_words[split + i])}?\" -> \"{\" \".join(words)}\"')", "metadata": {}, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "7abdc939", "source": "### Measuring the KV cache's actual speedup\n\nRecomputing self-attention from scratch at generation step $t$ costs $O(t)$ work (attending over all $t$ tokens so far); doing that at every step from $1$ to $T$ costs $O(T^2)$ in total. Caching makes every step cost $O(1)$ new work (just the newest token's attention against the growing cache), for $O(T)$ total — the whole reason real LLMs can generate long outputs without generation cost exploding quadratically with length. This is the identical idea Lesson {{lesson:decoder_only_language_models}} measured in its simpler, pure-self-attention setting; the only difference here is the extra (uncached) cross-attention step riding along at every generation step.\n\nAt this lesson's toy scale (a few hundred examples, `dim=32`, `MAX_LEN=20` tokens) that asymptotic gap is small in absolute terms, and most of the wall-clock time is Python/PyTorch call overhead rather than attention FLOPs — so to actually *see* the effect, this benchmark generates a large batch at once and repeats it several times, exactly the regime (long sequences, large batches) where KV caching earns its keep in practice.", "metadata": {} }, { "cell_type": "code", "id": "a2f9ebd1", "source": "import time\n\nBATCH, REPEATS = 512, 6\nXb = torch.tensor(np.tile(Xte, (BATCH // len(Xte) + 1, 1, 1))[:BATCH]).unsqueeze(1)\nQb = torch.tensor(np.tile(Qte, (BATCH // len(Qte) + 1, 1))[:BATCH])\n\nwith torch.no_grad():\n generate(Xb, Qb, use_cache=True) # warm up (lazy initialization, etc.)\n generate(Xb, Qb, use_cache=False)\n\n t0 = time.perf_counter()\n for _ in range(REPEATS):\n generate(Xb, Qb, use_cache=False)\n t_uncached = time.perf_counter() - t0\n\n t0 = time.perf_counter()\n for _ in range(REPEATS):\n generate(Xb, Qb, use_cache=True)\n t_cached = time.perf_counter() - t0\n\nprint(f'uncached (full recompute every step): {t_uncached * 1000:.0f} ms total, {REPEATS} generations of a batch of {BATCH}')\nprint(f'cached (KV cache): {t_cached * 1000:.0f} ms total')\nprint(f'speedup: {t_uncached / t_cached:.2f}x')", "metadata": {}, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "733ff88e", "source": "## The third family: encoder-decoder\n\nLesson {{lesson:decoder_only_language_models}} laid out two of the three Transformer model families: **encoder-only** (every prior Transformer in this course — maps an input to a representation and stops, bidirectional self-attention, no generation) and **decoder-only** (a single causal stream, no separate encoder, no cross-attention). This lesson's generative decoder is the third: **encoder-decoder** — a separate encoder (the image/question encoders above) produces a representation, and a decoder generates an output step by step, cross-attending back into that representation at every step. Unlike decoder-only, there are now two distinct token sets in play (the growing answer, and the fixed image+question memory), and unlike encoder-only, generation happens autoregressively, one token at a time.\n\n**The wrinkle worth naming directly:** Lesson {{lesson:segment_anything}}'s \"mask decoder\" and Lesson {{lesson:open_vocabulary_segmentation}}'s detector are *also* decoders in the generic Transformer-block sense — they cross-attend into an encoder's output and produce something. But they're **non-autoregressive**: one forward pass produces a fixed-size result (a set of masks) all at once, nothing is generated token-by-token, and there's no causal self-attention anywhere in them. That's architecturally close to DETR, not to GPT, and not to this lesson's decoder either. So \"decoder\" alone is still ambiguous even now that all three families are on the table — the useful question to always ask is: *does it generate a sequence one step at a time, or does it just produce a fixed-size output in a single pass?* Only the former is what \"a decoder\" usually means in modern LLM/VLM conversation, and it's exactly the ingredient — and the only ingredient — that makes a KV cache meaningful at all: caching speeds up *repeated* forward passes during step-by-step generation, so it has nothing to offer a model that was never going to take more than one forward pass in the first place.", "metadata": {} }, { "cell_type": "markdown", "id": "d8310c9b", "source": "### Exercises\n\n1. Add a fifth question template asking about *position* (e.g. \"is there a shape in the top half of the image\", answer yes/no based on `cy < SIZE // 2`). Retrain and check whether accuracy on this new question type matches the others. Unlike Lesson {{lesson:convolutional_neural_networks}}'s CNN with its position-destroying global max pooling, this lesson's ViT-style image encoder keeps an explicit positional encoding on every patch token, all the way into the `[CLS]` token's output — does that mean position-based questions come out *easier* here than they would with a CNN encoder, or does the encoder still have to learn to actually use that positional information, which isn't guaranteed just because it's technically present in the input?\n2. Replace elementwise multiplication fusion (`img_enc(img) * q_enc(q_ids)`) with concatenation followed by a linear layer (`nn.Linear(2 * dim, dim)` on `torch.cat([img_feat, q_feat], dim=-1)`). Does accuracy change noticeably, and which fusion method converges faster during training?\n3. Unlike a bag-of-words mean, this lesson's `QuestionEncoder` has both positional encoding and self-attention, so `\"is there a circle\"` and `\"circle there is\"` are *not* guaranteed to produce the same embedding (Lesson {{lesson:vision_language_models}}'s exercise made the same point about captions). Verify this directly: feed the encoder the same question's words in their normal order and in a scrambled order, and compare the two embeddings. Then check whether this order-sensitivity actually changes the model's predicted answer on this dataset — since every question is built from one of four fixed templates, does the model ever have a real reason to *use* word order, even though it technically could?\n4. The KV-cache benchmark above used `MAX_LEN=20` and a batch of 512 to make the speedup visible. Predict, then check, how the measured speedup changes if you (a) double `MAX_LEN` with everything else fixed, and (b) shrink the batch back down to 8. Does the direction of each change match the $O(T^2)$ vs. $O(T)$ argument above, and why would a tiny batch make the *Python-loop overhead* dominate the picture rather than the attention cost itself?", "metadata": {} } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 5 }