Lesson 58: Decoder-Only Language Models¶
Every Transformer this course has built so far — Lesson 47's ViT, Lesson 51's DINO, Lesson 53's image encoder, Lesson 55's CLIP — is encoder-only: it maps an input to a representation and stops. This lesson builds the opposite: a decoder-only model, the GPT-style design behind modern large language models — a single causal stream, no cross-attention, no separate encoder. This lesson also defines, precisely, what "self-attention" and "cross-attention" actually mean — a distinction this course has used informally but never defined.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import time
A tiny language, with no image anywhere¶
Reuse Lesson 55's small synthetic caption vocabulary, but drop the image entirely — this lesson is pure sequence modeling. Every sentence describes a shape two different ways, chosen at random: "a photo of a <size> <shape>" or "a <size> <shape> in the photo". The two styles start identically ("a") but diverge at the very next word ("photo" vs. a size word) — so predicting the next token genuinely requires having inferred, from the tokens generated so far, which style this particular sentence is.
SPECIAL = ['<pad>', '<bos>', '<eos>']
WORDS = ['a', 'photo', 'of', 'in', 'the', 'large', 'small', 'plus', 'circle']
VOCAB = SPECIAL + WORDS
W2ID = {w: i for i, w in enumerate(VOCAB)}
SIZES = ['large', 'small']
SHAPES = ['plus', 'circle']
MAX_LEN = 9 # <bos> + up to 6 words + <eos>, with a little slack
def caption_words(style, size, shape):
if style == 0:
return ['a', 'photo', 'of', 'a', size, shape]
else:
return ['a', size, shape, 'in', 'the', 'photo']
def tokenize(words):
ids = [W2ID['<bos>']] + [W2ID[w] for w in words] + [W2ID['<eos>']]
ids = ids + [W2ID['<pad>']] * (MAX_LEN - len(ids))
return ids[:MAX_LEN]
def make_dataset(rng, n):
seqs = []
for _ in range(n):
style = rng.integers(0, 2)
size, shape = rng.choice(SIZES), rng.choice(SHAPES)
seqs.append(tokenize(caption_words(style, size, shape)))
return np.array(seqs, dtype=np.int64)
rng = np.random.default_rng(0)
N = 800
data = make_dataset(rng, N)
split = int(0.85 * N)
train, test = data[:split], data[split:]
for i in range(4):
words = [VOCAB[t] for t in train[i] if t != W2ID['<pad>']]
print(' '.join(words))
A decoder-only block: an encoder block plus one mask¶
Lesson 46's encoder block applies multi-head self-attention, where every token can see every other token — including ones "after" it in the sequence. A decoder-only block changes exactly one thing: a causal mask on that same self-attention, so token $t$ can only attend to tokens $1, \ldots, t$, never anything later. Nothing else about the block changes — no cross-attention, no second set of tokens, no separate encoder feeding it anything. That one mask is the entire architectural difference between "this block looks at a whole sequence at once" and "this block can only generate a sequence one token at a time, each new token conditioned on everything generated so far and nothing else."
The explicit query/key/value projections below (rather than nn.MultiheadAttention) exist for the same reason Lesson 59 will use them: generating one token at a time can either recompute self-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 — explored below in its simplest possible setting, since there's no cross-attention memory to complicate it here.
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 GPTBlock(nn.Module):
"""Causal self-attention only -- no cross-attention, no separate encoder anywhere.
Explicit Q/K/V projections so generation can cache keys/values (see 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.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.lm_head = nn.Linear(dim, vocab_size)
def _causal_self_attn(self, q, k, v, causal_offset):
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
# query position i (offset by however many earlier tokens are already cached) may
# attend to key position j only if j <= i -- the causal mask, and the ONLY thing
# distinguishing this block from Lesson 46's encoder block.
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, past_kv=None):
"""past_kv=None recomputes everything from scratch (used for training, and as the
'naive' generation baseline). Passing a (k, v) tuple from 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._causal_self_attn(self.q_proj(x), k, v, causal_offset=start_pos)
x = self.ln1(x + self.out_proj(attn_out))
x = self.ln2(x + self.ff(x))
return self.lm_head(x), (k, v)
torch.manual_seed(0)
model = GPTBlock(len(VOCAB), dim=32)
opt = torch.optim.Adam(model.parameters(), lr=0.003)
Tt = torch.tensor(train)
n = len(train)
for epoch in range(800):
idx = np.random.default_rng(epoch).permutation(n)[:128]
batch = Tt[idx]
dec_in, target = batch[:, :-1], batch[:, 1:] # teacher forcing: predict each next word from the true prefix
logits, _ = model(dec_in)
loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target.reshape(-1), ignore_index=W2ID['<pad>'])
opt.zero_grad()
loss.backward()
opt.step()
print(f'final training loss: {loss.item():.4f}')
Generating, one token at a time¶
Start from just <bos>, feed it through the block, take the most likely next word, append it, and repeat — greedy decoding. The use_cache flag switches between recomputing the whole sequence from scratch at every step versus reusing cached keys/values and only processing the newest token. Both must produce identical output, since caching only changes how much gets recomputed, never what gets computed — checked explicitly below before trusting anything built on top of it.
def generate(B, use_cache):
tokens = torch.full((B, 1), W2ID['<bos>'], dtype=torch.long)
past_kv = None
for _ in range(MAX_LEN - 1):
if use_cache:
logits, past_kv = model(tokens[:, -1:], past_kv=past_kv)
else:
logits, _ = model(tokens, 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(W2ID['<eos>']) + 1] if W2ID['<eos>'] in seq else seq
def is_valid_caption(words):
s = ' '.join(words)
return any(' '.join(caption_words(style, size, shape)) == s
for style in (0, 1) for size in SIZES for shape in SHAPES)
with torch.no_grad():
gen_cached = generate(20, use_cache=True)
gen_uncached = generate(20, 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)')
print('greedy decoding from <bos> alone, 3 runs:')
for i in range(3):
words = [VOCAB[t] for t in trim_at_eos(gen_cached[i]) if t not in (W2ID['<bos>'], W2ID['<eos>'], W2ID['<pad>'])]
print(' ', ' '.join(words))
Greedy decoding hides what the model actually learned¶
Every greedy run above produces the exact same sentence — not a bug, but a direct consequence of always picking the single most-likely next word: starting from the same <bos>, that choice is deterministic at every step. That makes greedy decoding a poor way to check whether the model learned the whole grammar (both styles, both sizes, both shapes) rather than just memorizing whichever single sequence happens to be most probable overall. Sampling — drawing the next token from the predicted probability distribution instead of always taking its argmax — reveals the difference.
def generate_sampled(B, temperature=0.8, seed=0):
g = torch.Generator().manual_seed(seed)
tokens = torch.full((B, 1), W2ID['<bos>'], dtype=torch.long)
past_kv = None
for _ in range(MAX_LEN - 1):
logits, past_kv = model(tokens[:, -1:], past_kv=past_kv)
probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)
next_tok = torch.multinomial(probs, num_samples=1, generator=g)
tokens = torch.cat([tokens, next_tok], dim=1)
return tokens
with torch.no_grad():
sampled = generate_sampled(20, temperature=0.8, seed=5)
valid, unique = 0, set()
for i in range(20):
words = [VOCAB[t] for t in trim_at_eos(sampled[i]) if t not in (W2ID['<bos>'], W2ID['<eos>'], W2ID['<pad>'])]
unique.add(' '.join(words))
valid += is_valid_caption(words)
if i < 6:
print(' '.join(words))
print(f'valid captions: {valid}/20, unique sentences seen: {len(unique)}/8 possible')
Conditioning, without any cross-attention¶
How does a decoder-only model get "told" what to talk about, if it has no cross-attention and no separate encoder to feed it anything? The only mechanism available is the one already built: feed it a different prompt — a different starting sequence of tokens — and let the same causal self-attention continue from there. Since style 0 ("a photo of...") and style 1 ("a <size> <shape>...") diverge at the second token, prompting with "a photo" versus "a small" should steer generation into the matching style, purely because the model has learned which continuations are likely after each prefix — exactly the mechanism real chatbots use to follow a prompt, with no architecture beyond what's already here.
def generate_from_prompt(prompt_words):
prompt_ids = [W2ID['<bos>']] + [W2ID[w] for w in prompt_words]
tokens = torch.tensor([prompt_ids], dtype=torch.long)
logits, past_kv = model(tokens[:, :-1], past_kv=None) # prime the cache with the prompt
for _ in range(MAX_LEN - tokens.shape[1]):
logits, past_kv = model(tokens[:, -1:], past_kv=past_kv)
next_tok = logits[:, -1, :].argmax(-1, keepdim=True)
tokens = torch.cat([tokens, next_tok], dim=1)
return tokens[0]
with torch.no_grad():
for prompt in [['a', 'photo'], ['a', 'small'], ['a', 'large']]:
gen = generate_from_prompt(prompt)
words = [VOCAB[t] for t in trim_at_eos(gen) if t not in (W2ID['<bos>'], W2ID['<eos>'], W2ID['<pad>'])]
print(f'prompt {prompt!r:<22} -> {" ".join(words)}')
Measuring the KV cache's actual speedup¶
Recomputing self-attention from scratch at generation step $t$ costs $O(t)$ work; 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 cleanest possible setting to measure that in: no cross-attention memory, no image encoder, nothing but the causal self-attention itself.
At this lesson's toy scale, most of the wall-clock time is Python/PyTorch call overhead rather than attention FLOPs, so a large batch and several repeats are needed to see the effect above that noise floor — exactly the regime (long sequences, large batches) where KV caching earns its keep in real deployments too.
BATCH, REPEATS = 512, 6
with torch.no_grad():
generate(BATCH, use_cache=True) # warm up (lazy initialization, etc.)
generate(BATCH, use_cache=False)
t0 = time.perf_counter()
for _ in range(REPEATS):
generate(BATCH, use_cache=False)
t_uncached = time.perf_counter() - t0
t0 = time.perf_counter()
for _ in range(REPEATS):
generate(BATCH, 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')
Self-attention vs. cross-attention¶
Every attention computation this lesson has built and run so far is self-attention — look back at GPTBlock: one set of tokens, used as the queries, the keys, and the values, all at once. This course has used the word "self-attention" informally in passing since Lesson 50 (and "cross-attention" once, in Lesson 57's aside about SAM 3's real decoder), but never actually defined either term.
Self-attention: a block's queries, keys, and values are all computed from the same set of tokens. The sequence attends to itself — exactly what GPTBlock does above.
Cross-attention: the queries come from one set of tokens; the keys and values come from a different set, usually produced by something else entirely (a separate encoder).
That's the whole distinction — and notice it says nothing about how many modalities are involved. It's tempting to assume "self-attention is what happens within one image, cross-attention is what happens across modalities or frames," but that's not the rule, and one of this course's own lessons is proof: Lesson 55's CLIP is thoroughly multimodal — an image tower and a text tower — yet the two towers never interact inside any attention layer at all; each self-attends only within its own modality, and they're compared only afterward, with a plain dot product. There is no cross-attention anywhere in CLIP.
Lesson 59, next, builds this course's first genuine example of true cross-attention: a separate image+question encoder produces a representation, and a decoder's queries attend into that separate representation at every generation step.
Encoder-only vs. decoder-only¶
| Model family | What it does | Generates autoregressively? | Built in this course |
|---|---|---|---|
| Encoder-only | Maps an input to a fixed representation (an embedding, or per-token features) and stops. Bidirectional self-attention — every token can see every other token, including ones "after" it. | No | Every Transformer block built in this course before this lesson: Lesson 47's ViT, Lesson 53's image encoder, Lesson 55's CLIP, Lesson 51's DINO, and more. |
| Decoder-only | A single causal stream, no separate encoder, no cross-attention anywhere. Everything — a prompt and whatever gets generated after it — is just more tokens in one sequence, each attending only backward. | Yes | This lesson's GPTBlock. |
A third family, encoder-decoder — a separate encoder plus a decoder that cross-attends back into it — is Lesson 59's job next, once a genuine cross-attention mechanism actually exists to put in the table.
Real decoder-only LLMs and VLMs¶
This exact recipe — causal self-attention, a growing context, next-token prediction, scaled up enormously — is the GPT family (Radford et al., 2019) and most of today's large language models. It's not just a text-only design, either: many modern multimodal LLMs (LLaVA-style) are decoder-only too, prepending image patch embeddings as extra tokens onto the front of the causal sequence and running ordinary self-attention over the whole mixed sequence jointly — no cross-attention block anywhere, just a longer sequence with two kinds of tokens in it. That's a real, different architectural choice from the encoder-decoder captioning design Lesson 59 builds next — both exist in production systems today, and now there's a concrete example of each to compare.
Exercises¶
- Drop the causal mask entirely in
_causal_self_attn(let every query attend to every key, matching Lesson 46's encoder exactly) and retrain. Can this version still be used for autoregressive generation at all — and if you try anyway, what goes wrong? - Raise
temperatureingenerate_sampledfrom 0.8 to 2.0, and separately lower it to 0.3. How doesvalid/20 change in each direction, and why does a very high temperature eventually break the model's grasp of the grammar, while a very low one pushes sampling back toward greedy's single deterministic answer? - The KV-cache benchmark above used
MAX_LEN=9and a batch of 512. Predict, then check, how the measured speedup changes if you (a) growMAX_LENsubstantially (e.g. lengthen the toy sentences), 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?