Lesson 59: Visual Question Answering (VQA)¶
Lesson 55'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.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Scenes with (question, answer) pairs¶
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.
SIZE = 32
QUESTION_VOCAB = ['<pad>', 'what', 'shape', 'is', 'this', 'how', 'many', 'shapes', 'are', 'there',
'a', 'circle', 'plus', 'in', 'the', 'image']
Q2ID = {w: i for i, w in enumerate(QUESTION_VOCAB)}
ANSWER_VOCAB = ['plus', 'circle', '0', '1', '2', '3', 'yes', 'no']
A2ID = {w: i for i, w in enumerate(ANSWER_VOCAB)}
def make_shape(img, shape_type, cx, cy, r=4):
yy, xx = np.mgrid[0:img.shape[0], 0:img.shape[1]]
if shape_type == 'plus':
img[cy-1:cy+2, cx-r:cx+r+1] = 1.0
img[cy-r:cy+r+1, cx-1:cx+2] = 1.0
else:
img[((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2] = 1.0
return img
def make_scene(rng, size=SIZE, max_shapes=3):
n = rng.integers(0, max_shapes + 1)
img = np.zeros((size, size), dtype=np.float32)
shapes = []
tries = 0
while len(shapes) < n and tries < 30:
tries += 1
cx, cy = rng.integers(6, size - 6), rng.integers(6, size - 6)
if any(abs(cx - sx) < 9 and abs(cy - sy) < 9 for sx, sy, _ in shapes):
continue
st = rng.choice(['plus', 'circle'])
make_shape(img, st, cx, cy)
shapes.append((cx, cy, st))
img = np.clip(img + rng.normal(0, 0.03, img.shape), 0, 1).astype(np.float32)
return img, shapes
def make_qa(rng, shapes):
n_circle = sum(1 for _, _, t in shapes if t == 'circle')
n_plus = sum(1 for _, _, t in shapes if t == 'plus')
qtype = rng.choice(['what_shape', 'count', 'is_there_circle', 'is_there_plus'])
if qtype == 'what_shape':
words = ['what', 'shape', 'is', 'in', 'the', 'image']
answer = shapes[0][2] if len(shapes) == 1 else str(min(len(shapes), 3))
elif qtype == 'count':
words = ['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image']
answer = str(len(shapes))
elif qtype == 'is_there_circle':
words = ['is', 'there', 'a', 'circle', 'in', 'the', 'image']
answer = 'yes' if n_circle > 0 else 'no'
else:
words = ['is', 'there', 'a', 'plus', 'in', 'the', 'image']
answer = 'yes' if n_plus > 0 else 'no'
# qtype and the raw shapes are returned too (unused by the closed-vocabulary classifier
# below) because the generative decoder later in this lesson needs them to build full
# phrase answers like "there are two shapes" instead of a single class id.
return qtype, words, [Q2ID[w] for w in words], A2ID[answer]
def pad_question(ids, max_len=8):
return ids + [0] * (max_len - len(ids)) # 0 = <pad>, distinct from every real word
rng = np.random.default_rng(4)
N = 1000
imgs, questions, answers, question_words, scene_shapes, question_types = [], [], [], [], [], []
for _ in range(N):
img, shapes = make_scene(rng)
qtype, words, q_ids, a_id = make_qa(rng, shapes)
imgs.append(img)
questions.append(pad_question(q_ids))
answers.append(a_id)
question_words.append(words)
scene_shapes.append(shapes)
question_types.append(qtype)
imgs = np.array(imgs, dtype=np.float32)
questions = np.array(questions, dtype=np.int64)
answers = np.array(answers, dtype=np.int64)
split = int(0.85 * N)
Xtr, Qtr, Atr = imgs[:split], questions[:split], answers[:split]
Xte, Qte, Ate = imgs[split:], questions[split:], answers[split:]
shapes_tr, shapes_te = scene_shapes[:split], scene_shapes[split:]
qtypes_tr, qtypes_te = question_types[:split], question_types[split:]
fig, axes = plt.subplots(1, 4, figsize=(9, 2.5))
for ax, im, words, a in zip(axes, imgs[:4], question_words[:4], answers[:4]):
ax.imshow(im, cmap='gray'); ax.axis('off')
ax.set_title(f'{" ".join(words)}?\n-> {ANSWER_VOCAB[a]}', fontsize=7)
plt.show()
Fusing an image encoder and a question encoder¶
Both encoders are small Transformers (Lesson 47'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 55'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.
def patchify(img, patch_size):
B, C, H, W = img.shape
P = patch_size
patches = img.unfold(2, P, P).unfold(3, P, P)
patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)
return patches.reshape(B, -1, C * P * P)
def positional_encoding(T, D):
pos = torch.arange(T).unsqueeze(1).float()
i = torch.arange(D).unsqueeze(0).float()
angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))
angles = pos * angle_rates
pe = torch.zeros(T, D)
pe[:, 0::2] = torch.sin(angles[:, 0::2])
pe[:, 1::2] = torch.cos(angles[:, 1::2])
return pe
class ImageEncoder(nn.Module):
"""A small ViT (Lesson 47's pattern): patchify, embed,
prepend a [CLS] token, run a Transformer encoder, read the [CLS] token's output."""
def __init__(self, img_size=SIZE, patch_size=4, out_dim=32, embed_dim=32, n_heads=4, n_layers=2):
super().__init__()
self.patch_size = patch_size
n_patches = (img_size // patch_size) ** 2
patch_dim = patch_size * patch_size
self.embed = nn.Linear(patch_dim, embed_dim)
self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
self.register_buffer('pos_enc', positional_encoding(n_patches + 1, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,
batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, x):
p = patchify(x, self.patch_size)
tok = self.embed(p)
cls = self.cls_token.expand(x.shape[0], -1, -1)
tok = torch.cat([cls, tok], dim=1) + self.pos_enc
out = self.encoder(tok)
return self.proj(out[:, 0])
class QuestionEncoder(nn.Module):
"""A small bidirectional Transformer encoder over the question's word tokens, with
padding masked out of both attention and the final pooling."""
def __init__(self, vocab_size, out_dim=32, embed_dim=32, n_heads=4, n_layers=1, max_len=8):
super().__init__()
self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.register_buffer('pos_enc', positional_encoding(max_len, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,
batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, ids):
pad_mask = (ids == 0) # True where padding
x = self.embed(ids) + self.pos_enc[:ids.shape[1]]
x = self.encoder(x, src_key_padding_mask=pad_mask)
valid = (~pad_mask).float().unsqueeze(-1)
pooled = (x * valid).sum(1) / valid.sum(1).clamp(min=1)
return self.proj(pooled)
class VQAModel(nn.Module):
def __init__(self, vocab_size, n_answers, dim=32):
super().__init__()
self.img_enc = ImageEncoder(out_dim=dim)
self.q_enc = QuestionEncoder(vocab_size, out_dim=dim)
self.head = nn.Sequential(nn.Linear(dim, 32), nn.ReLU(), nn.Linear(32, n_answers))
def forward(self, img, q_ids):
fused = self.img_enc(img) * self.q_enc(q_ids)
return self.head(fused)
torch.manual_seed(0)
model = VQAModel(len(QUESTION_VOCAB), len(ANSWER_VOCAB))
opt = torch.optim.Adam(model.parameters(), lr=0.003)
Xt = torch.tensor(Xtr).unsqueeze(1); Qt = torch.tensor(Qtr); At = torch.tensor(Atr)
n = len(Xtr)
for epoch in range(300):
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
opt.zero_grad()
loss = F.cross_entropy(model(Xt[idx], Qt[idx]), At[idx])
loss.backward()
opt.step()
with torch.no_grad():
preds = model(torch.tensor(Xte).unsqueeze(1), torch.tensor(Qte)).argmax(1).numpy()
acc = (preds == Ate).mean()
most_common = np.bincount(Atr).argmax()
baseline_acc = (Ate == most_common).mean()
print(f'VQA test accuracy: {acc:.1%}')
print(f'baseline (always predict the most common training answer, "{ANSWER_VOCAB[most_common]}"): {baseline_acc:.1%}')
The real test: does the same image get different answers to different questions?¶
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.
test_rng = np.random.default_rng(123)
demo_img, demo_shapes = make_scene(test_rng)
while len(demo_shapes) != 2 or demo_shapes[0][2] == demo_shapes[1][2]:
demo_img, demo_shapes = make_scene(test_rng) # find a scene with one of each shape
demo_questions = [
['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image'],
['is', 'there', 'a', 'circle', 'in', 'the', 'image'],
['is', 'there', 'a', 'plus', 'in', 'the', 'image'],
]
img_t = torch.tensor(demo_img[None, None])
print(f'true scene contents: {[s[2] for s in demo_shapes]}')
with torch.no_grad():
for words in demo_questions:
q_ids = torch.tensor([pad_question([Q2ID[w] for w in words])])
pred = model(img_t, q_ids).argmax(1).item()
print(f' "{" ".join(words)}?" -> {ANSWER_VOCAB[pred]}')
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.
From closed-vocabulary classification to real generative VQA¶
Treating 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 46 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 58 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.
Phrase answers instead of single labels¶
Instead 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 <bos> ("beginning of sequence") token and ends with <eos> ("end of sequence"), exactly the convention real language models use to know where to start generating and when to stop.
NUM_WORDS = {0: 'zero', 1: 'one', 2: 'two', 3: 'three'}
def answer_phrase(qtype, shapes):
n = len(shapes)
n_circle = sum(1 for _, _, t in shapes if t == 'circle')
n_plus = sum(1 for _, _, t in shapes if t == 'plus')
if qtype == 'what_shape':
if n == 1:
return ['it', 'is', 'a', shapes[0][2]]
elif n == 0:
return ['there', 'is', 'no', 'shape']
else:
return ['there', 'are', NUM_WORDS[min(n, 3)], 'shapes']
elif qtype == 'count':
if n == 0:
return ['there', 'are', 'no', 'shapes']
elif n == 1:
return ['there', 'is', 'one', 'shape']
else:
return ['there', 'are', NUM_WORDS[min(n, 3)], 'shapes']
elif qtype == 'is_there_circle':
return ['yes', 'there', 'is', 'a', 'circle'] if n_circle > 0 else ['no', 'there', 'is', 'no', 'circle']
else:
return ['yes', 'there', 'is', 'a', 'plus'] if n_plus > 0 else ['no', 'there', 'is', 'no', 'plus']
ANSWER_WORDS = ['<pad>', '<bos>', '<eos>', 'it', 'is', 'a', 'there', 'are', 'no', 'shape', 'shapes',
'yes', 'circle', 'plus', 'zero', 'one', 'two', 'three']
AW2ID = {w: i for i, w in enumerate(ANSWER_WORDS)}
MAX_LEN = 20 # bos + up to 5 words + eos, padded well past the longest real phrase -- see the KV-cache timing section below for why
def tokenize_answer(words):
ids = [AW2ID['<bos>']] + [AW2ID[w] for w in words] + [AW2ID['<eos>']]
return ids + [AW2ID['<pad>']] * (MAX_LEN - len(ids))
AtrP = np.array([tokenize_answer(answer_phrase(qt, sh)) for qt, sh in zip(qtypes_tr, shapes_tr)], dtype=np.int64)
AteP = np.array([tokenize_answer(answer_phrase(qt, sh)) for qt, sh in zip(qtypes_te, shapes_te)], dtype=np.int64)
for i in range(3):
words = [ANSWER_WORDS[t] for t in AtrP[i] if t not in (AW2ID['<pad>'], AW2ID['<bos>'], AW2ID['<eos>'])]
print(f'"{" ".join(question_words[i])}?" -> "{" ".join(words)}"')
A causal, cross-attending decoder — built explicitly, so it can be cached later¶
This decoder combines the two attention patterns Lesson 58 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 58'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 58 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.
Lesson 58 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.
class ImageEncoder2(nn.Module):
"""Same ViT pattern as this lesson's earlier ImageEncoder -- a separate copy so it gets
its own trained weights, independent of the closed-vocabulary classifier's encoder."""
def __init__(self, out_dim=32):
super().__init__()
self.patch_size = 4
n_patches = (SIZE // self.patch_size) ** 2
patch_dim = self.patch_size * self.patch_size
embed_dim = 32
self.embed = nn.Linear(patch_dim, embed_dim)
self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
self.register_buffer('pos_enc', positional_encoding(n_patches + 1, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, 4, dim_feedforward=embed_dim * 2, batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=2)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, x):
p = patchify(x, self.patch_size)
tok = self.embed(p)
cls = self.cls_token.expand(x.shape[0], -1, -1)
tok = torch.cat([cls, tok], dim=1) + self.pos_enc
out = self.encoder(tok)
return self.proj(out[:, 0])
class QuestionEncoder2(nn.Module):
def __init__(self, vocab_size, out_dim=32):
super().__init__()
embed_dim = 32
self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.register_buffer('pos_enc', positional_encoding(8, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, 4, dim_feedforward=embed_dim * 2, batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=1)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, ids):
pad_mask = (ids == 0)
x = self.embed(ids) + self.pos_enc[:ids.shape[1]]
x = self.encoder(x, src_key_padding_mask=pad_mask)
valid = (~pad_mask).float().unsqueeze(-1)
pooled = (x * valid).sum(1) / valid.sum(1).clamp(min=1)
return self.proj(pooled)
class CachedDecoder(nn.Module):
"""A from-scratch causal self-attention + cross-attention decoder block with explicit
K/V projections, so generation can optionally reuse cached keys/values instead of
recomputing them -- see the KV-cache section below."""
def __init__(self, vocab_size, dim=32, n_heads=4, max_len=MAX_LEN):
super().__init__()
self.dim, self.n_heads = dim, n_heads
self.tok_embed = nn.Embedding(vocab_size, dim, padding_idx=0)
self.pos_embed = nn.Embedding(max_len, dim)
self.q_proj = nn.Linear(dim, dim); self.k_proj = nn.Linear(dim, dim); self.v_proj = nn.Linear(dim, dim)
self.out_proj = nn.Linear(dim, dim)
self.cross_q = nn.Linear(dim, dim); self.cross_k = nn.Linear(dim, dim); self.cross_v = nn.Linear(dim, dim)
self.cross_out_proj = nn.Linear(dim, dim)
self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))
self.ln1 = nn.LayerNorm(dim); self.ln2 = nn.LayerNorm(dim); self.ln3 = nn.LayerNorm(dim)
self.answer_head = nn.Linear(dim, vocab_size)
def _mha(self, q, k, v, causal_offset=None):
B, Tq, D = q.shape
Tk = k.shape[1]
H, hd = self.n_heads, D // self.n_heads
q = q.view(B, Tq, H, hd).transpose(1, 2)
k = k.view(B, Tk, H, hd).transpose(1, 2)
v = v.view(B, Tk, H, hd).transpose(1, 2)
scores = (q @ k.transpose(-2, -1)) / hd ** 0.5
if causal_offset is not None:
# position i (in the new query chunk) may attend to key position j only if
# j <= i + causal_offset -- the usual causal mask, shifted to account for
# however many earlier tokens' keys are already sitting in the cache.
i = torch.arange(Tq).unsqueeze(1) + causal_offset
j = torch.arange(Tk).unsqueeze(0)
scores = scores.masked_fill(j > i, float('-inf'))
weights = scores.softmax(dim=-1)
return (weights @ v).transpose(1, 2).reshape(B, Tq, D)
def forward(self, tokens_in, memory, past_kv=None):
"""past_kv=None recomputes everything from scratch (used for training, and as the
'naive' generation baseline). Passing a (k, v) tuple returned by an earlier call
appends only the new token's key/value instead of recomputing the whole prefix --
the actual KV cache."""
B, T = tokens_in.shape
start_pos = 0 if past_kv is None else past_kv[0].shape[1]
pos = torch.arange(start_pos, start_pos + T)
x = self.tok_embed(tokens_in) + self.pos_embed(pos)[None]
k_new, v_new = self.k_proj(x), self.v_proj(x)
if past_kv is not None:
k = torch.cat([past_kv[0], k_new], dim=1)
v = torch.cat([past_kv[1], v_new], dim=1)
else:
k, v = k_new, v_new
attn_out = self._mha(self.q_proj(x), k, v, causal_offset=start_pos)
x = self.ln1(x + self.out_proj(attn_out))
cross_out = self._mha(self.cross_q(x), self.cross_k(memory), self.cross_v(memory), causal_offset=None)
x = self.ln2(x + self.cross_out_proj(cross_out))
x = self.ln3(x + self.ff(x))
return self.answer_head(x), (k, v)
torch.manual_seed(0)
img_enc2 = ImageEncoder2(32)
q_enc2 = QuestionEncoder2(len(QUESTION_VOCAB), 32)
decoder = CachedDecoder(len(ANSWER_WORDS), dim=32)
opt2 = torch.optim.Adam(list(img_enc2.parameters()) + list(q_enc2.parameters()) + list(decoder.parameters()), lr=0.002)
Xt2, Qt2, At2 = torch.tensor(Xtr).unsqueeze(1), torch.tensor(Qtr), torch.tensor(AtrP)
n2 = len(Xtr)
for epoch in range(800):
idx = np.random.default_rng(epoch).permutation(n2)[:256] # mini-batches, same reason as the classifier above
memory = torch.stack([img_enc2(Xt2[idx]), q_enc2(Qt2[idx])], dim=1) # (B, 2, dim): the image features + question features
dec_in, target = At2[idx][:, :-1], At2[idx][:, 1:] # teacher forcing: predict each next word from the true prefix
logits, _ = decoder(dec_in, memory)
loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target.reshape(-1), ignore_index=AW2ID['<pad>'])
opt2.zero_grad(); loss.backward(); opt2.step()
print(f'final training loss: {loss.item():.4f}')
Generating an answer, one token at a time¶
Start from just <bos>, 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.
def generate(imgs_t, q_t, use_cache):
B = imgs_t.shape[0]
memory = torch.stack([img_enc2(imgs_t), q_enc2(q_t)], dim=1)
tokens = torch.full((B, 1), AW2ID['<bos>'], dtype=torch.long)
past_kv = None
for _ in range(MAX_LEN - 1):
if use_cache:
logits, past_kv = decoder(tokens[:, -1:], memory, past_kv=past_kv)
else:
logits, _ = decoder(tokens, memory, past_kv=None)
logits = logits[:, -1:]
next_tok = logits[:, -1, :].argmax(-1, keepdim=True)
tokens = torch.cat([tokens, next_tok], dim=1)
return tokens
def trim_at_eos(seq):
seq = seq.tolist()
return seq[:seq.index(AW2ID['<eos>']) + 1] if AW2ID['<eos>'] in seq else seq
with torch.no_grad():
Xte2, Qte2 = torch.tensor(Xte).unsqueeze(1), torch.tensor(Qte)
gen_cached = generate(Xte2, Qte2, use_cache=True)
gen_uncached = generate(Xte2, Qte2, use_cache=False)
match_rate = (gen_cached == gen_uncached).float().mean().item()
print(f'cached vs. uncached generation, token-for-token agreement: {match_rate:.4f} (should be 1.0000)')
correct = sum(trim_at_eos(g) == trim_at_eos(a) for g, a in zip(gen_cached, torch.tensor(AteP)))
print(f'exact-phrase accuracy on held-out test scenes: {correct}/{len(AteP)} = {correct / len(AteP):.1%}')
for i in range(4):
words = [ANSWER_WORDS[t] for t in trim_at_eos(gen_cached[i]) if t not in (AW2ID['<bos>'], AW2ID['<eos>'])]
print(f' "{" ".join(question_words[split + i])}?" -> "{" ".join(words)}"')
Measuring the KV cache's actual speedup¶
Recomputing 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 58 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.
At 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.
import time
BATCH, REPEATS = 512, 6
Xb = torch.tensor(np.tile(Xte, (BATCH // len(Xte) + 1, 1, 1))[:BATCH]).unsqueeze(1)
Qb = torch.tensor(np.tile(Qte, (BATCH // len(Qte) + 1, 1))[:BATCH])
with torch.no_grad():
generate(Xb, Qb, use_cache=True) # warm up (lazy initialization, etc.)
generate(Xb, Qb, use_cache=False)
t0 = time.perf_counter()
for _ in range(REPEATS):
generate(Xb, Qb, use_cache=False)
t_uncached = time.perf_counter() - t0
t0 = time.perf_counter()
for _ in range(REPEATS):
generate(Xb, Qb, use_cache=True)
t_cached = time.perf_counter() - t0
print(f'uncached (full recompute every step): {t_uncached * 1000:.0f} ms total, {REPEATS} generations of a batch of {BATCH}')
print(f'cached (KV cache): {t_cached * 1000:.0f} ms total')
print(f'speedup: {t_uncached / t_cached:.2f}x')
The third family: encoder-decoder¶
Lesson 58 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.
The wrinkle worth naming directly: Lesson 53's "mask decoder" and Lesson 57'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.
Exercises¶
- 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 34'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? - Replace elementwise multiplication fusion (
img_enc(img) * q_enc(q_ids)) with concatenation followed by a linear layer (nn.Linear(2 * dim, dim)ontorch.cat([img_feat, q_feat], dim=-1)). Does accuracy change noticeably, and which fusion method converges faster during training? - Unlike a bag-of-words mean, this lesson's
QuestionEncoderhas both positional encoding and self-attention, so"is there a circle"and"circle there is"are not guaranteed to produce the same embedding (Lesson 55'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? - The KV-cache benchmark above used
MAX_LEN=20and a batch of 512 to make the speedup visible. Predict, then check, how the measured speedup changes if you (a) doubleMAX_LENwith 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?