{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Lesson 58: Decoder-Only Language Models\n", "\n", "Every Transformer this course has built so far — Lesson {{lesson:vision_transformers}}'s ViT, Lesson {{lesson:dino_self_distillation}}'s DINO, Lesson {{lesson:segment_anything}}'s image encoder, Lesson {{lesson:vision_language_models}}'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." ] }, { "cell_type": "code", "execution_count": null, "id": "2c9b7f3d", "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 time" ] }, { "cell_type": "markdown", "id": "a9257f94", "metadata": {}, "source": [ "## A tiny language, with no image anywhere\n", "\n", "Reuse Lesson {{lesson:vision_language_models}}'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 \"` or `\"a 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." ] }, { "cell_type": "code", "execution_count": null, "id": "81703c35", "metadata": {}, "outputs": [], "source": [ "SPECIAL = ['', '', '']\n", "WORDS = ['a', 'photo', 'of', 'in', 'the', 'large', 'small', 'plus', 'circle']\n", "VOCAB = SPECIAL + WORDS\n", "W2ID = {w: i for i, w in enumerate(VOCAB)}\n", "\n", "SIZES = ['large', 'small']\n", "SHAPES = ['plus', 'circle']\n", "MAX_LEN = 9 # + up to 6 words + , with a little slack\n", "\n", "def caption_words(style, size, shape):\n", " if style == 0:\n", " return ['a', 'photo', 'of', 'a', size, shape]\n", " else:\n", " return ['a', size, shape, 'in', 'the', 'photo']\n", "\n", "def tokenize(words):\n", " ids = [W2ID['']] + [W2ID[w] for w in words] + [W2ID['']]\n", " ids = ids + [W2ID['']] * (MAX_LEN - len(ids))\n", " return ids[:MAX_LEN]\n", "\n", "def make_dataset(rng, n):\n", " seqs = []\n", " for _ in range(n):\n", " style = rng.integers(0, 2)\n", " size, shape = rng.choice(SIZES), rng.choice(SHAPES)\n", " seqs.append(tokenize(caption_words(style, size, shape)))\n", " return np.array(seqs, dtype=np.int64)\n", "\n", "rng = np.random.default_rng(0)\n", "N = 800\n", "data = make_dataset(rng, N)\n", "split = int(0.85 * N)\n", "train, test = data[:split], data[split:]\n", "\n", "for i in range(4):\n", " words = [VOCAB[t] for t in train[i] if t != W2ID['']]\n", " print(' '.join(words))" ] }, { "cell_type": "markdown", "id": "712de3bb", "metadata": {}, "source": [ "## A decoder-only block: an encoder block plus one mask\n", "\n", "Lesson {{lesson:transformer_architecture}}'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.\"\n", "\n", "The explicit query/key/value projections below (rather than `nn.MultiheadAttention`) exist for the same reason Lesson {{lesson:visual_question_answering}} 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." ] }, { "cell_type": "code", "execution_count": null, "id": "70485448", "metadata": {}, "outputs": [], "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", "\n", "\n", "class GPTBlock(nn.Module):\n", " \"\"\"Causal self-attention only -- no cross-attention, no separate encoder anywhere.\n", " Explicit Q/K/V projections so generation can cache keys/values (see 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.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)\n", " self.lm_head = nn.Linear(dim, vocab_size)\n", "\n", " def _causal_self_attn(self, q, k, v, causal_offset):\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", " # query position i (offset by however many earlier tokens are already cached) may\n", " # attend to key position j only if j <= i -- the causal mask, and the ONLY thing\n", " # distinguishing this block from Lesson 46's encoder block.\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, 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 from an earlier call appends\n", " only the new token's key/value instead of recomputing the whole prefix -- the\n", " 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._causal_self_attn(self.q_proj(x), k, v, causal_offset=start_pos)\n", " x = self.ln1(x + self.out_proj(attn_out))\n", " x = self.ln2(x + self.ff(x))\n", " return self.lm_head(x), (k, v)\n", "\n", "\n", "torch.manual_seed(0)\n", "model = GPTBlock(len(VOCAB), dim=32)\n", "opt = torch.optim.Adam(model.parameters(), lr=0.003)\n", "Tt = torch.tensor(train)\n", "n = len(train)\n", "for epoch in range(800):\n", " idx = np.random.default_rng(epoch).permutation(n)[:128]\n", " batch = Tt[idx]\n", " dec_in, target = batch[:, :-1], batch[:, 1:] # teacher forcing: predict each next word from the true prefix\n", " logits, _ = model(dec_in)\n", " loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target.reshape(-1), ignore_index=W2ID[''])\n", " opt.zero_grad()\n", " loss.backward()\n", " opt.step()\n", "\n", "print(f'final training loss: {loss.item():.4f}')" ] }, { "cell_type": "markdown", "id": "c4c3e3ea", "metadata": {}, "source": [ "## Generating, one token at a time\n", "\n", "Start from just ``, 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." ] }, { "cell_type": "code", "execution_count": null, "id": "078931b9", "metadata": {}, "outputs": [], "source": [ "def generate(B, use_cache):\n", " tokens = torch.full((B, 1), W2ID[''], dtype=torch.long)\n", " past_kv = None\n", " for _ in range(MAX_LEN - 1):\n", " if use_cache:\n", " logits, past_kv = model(tokens[:, -1:], past_kv=past_kv)\n", " else:\n", " logits, _ = model(tokens, 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", "\n", "def trim_at_eos(seq):\n", " seq = seq.tolist()\n", " return seq[:seq.index(W2ID['']) + 1] if W2ID[''] in seq else seq\n", "\n", "def is_valid_caption(words):\n", " s = ' '.join(words)\n", " return any(' '.join(caption_words(style, size, shape)) == s\n", " for style in (0, 1) for size in SIZES for shape in SHAPES)\n", "\n", "with torch.no_grad():\n", " gen_cached = generate(20, use_cache=True)\n", " gen_uncached = generate(20, use_cache=False)\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", " print('greedy decoding from alone, 3 runs:')\n", " for i in range(3):\n", " words = [VOCAB[t] for t in trim_at_eos(gen_cached[i]) if t not in (W2ID[''], W2ID[''], W2ID[''])]\n", " print(' ', ' '.join(words))" ] }, { "cell_type": "markdown", "id": "cf33a0cb", "metadata": {}, "source": [ "## Greedy decoding hides what the model actually learned\n", "\n", "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 ``, 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." ] }, { "cell_type": "code", "execution_count": null, "id": "1dc9e9eb", "metadata": {}, "outputs": [], "source": [ "def generate_sampled(B, temperature=0.8, seed=0):\n", " g = torch.Generator().manual_seed(seed)\n", " tokens = torch.full((B, 1), W2ID[''], dtype=torch.long)\n", " past_kv = None\n", " for _ in range(MAX_LEN - 1):\n", " logits, past_kv = model(tokens[:, -1:], past_kv=past_kv)\n", " probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)\n", " next_tok = torch.multinomial(probs, num_samples=1, generator=g)\n", " tokens = torch.cat([tokens, next_tok], dim=1)\n", " return tokens\n", "\n", "with torch.no_grad():\n", " sampled = generate_sampled(20, temperature=0.8, seed=5)\n", " valid, unique = 0, set()\n", " for i in range(20):\n", " words = [VOCAB[t] for t in trim_at_eos(sampled[i]) if t not in (W2ID[''], W2ID[''], W2ID[''])]\n", " unique.add(' '.join(words))\n", " valid += is_valid_caption(words)\n", " if i < 6:\n", " print(' '.join(words))\n", "\n", "print(f'valid captions: {valid}/20, unique sentences seen: {len(unique)}/8 possible')" ] }, { "cell_type": "markdown", "id": "719a0d80", "metadata": {}, "source": [ "## Conditioning, without any cross-attention\n", "\n", "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 `` ``...\") 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." ] }, { "cell_type": "code", "execution_count": null, "id": "3d4c281d", "metadata": {}, "outputs": [], "source": [ "def generate_from_prompt(prompt_words):\n", " prompt_ids = [W2ID['']] + [W2ID[w] for w in prompt_words]\n", " tokens = torch.tensor([prompt_ids], dtype=torch.long)\n", " logits, past_kv = model(tokens[:, :-1], past_kv=None) # prime the cache with the prompt\n", " for _ in range(MAX_LEN - tokens.shape[1]):\n", " logits, past_kv = model(tokens[:, -1:], past_kv=past_kv)\n", " next_tok = logits[:, -1, :].argmax(-1, keepdim=True)\n", " tokens = torch.cat([tokens, next_tok], dim=1)\n", " return tokens[0]\n", "\n", "with torch.no_grad():\n", " for prompt in [['a', 'photo'], ['a', 'small'], ['a', 'large']]:\n", " gen = generate_from_prompt(prompt)\n", " words = [VOCAB[t] for t in trim_at_eos(gen) if t not in (W2ID[''], W2ID[''], W2ID[''])]\n", " print(f'prompt {prompt!r:<22} -> {\" \".join(words)}')" ] }, { "cell_type": "markdown", "id": "0f7aca7b", "metadata": {}, "source": [ "## Measuring the KV cache's actual speedup\n", "\n", "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.\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": null, "id": "ef8ddaaa", "metadata": {}, "outputs": [], "source": [ "BATCH, REPEATS = 512, 6\n", "\n", "with torch.no_grad():\n", " generate(BATCH, use_cache=True) # warm up (lazy initialization, etc.)\n", " generate(BATCH, use_cache=False)\n", "\n", " t0 = time.perf_counter()\n", " for _ in range(REPEATS):\n", " generate(BATCH, use_cache=False)\n", " t_uncached = time.perf_counter() - t0\n", "\n", " t0 = time.perf_counter()\n", " for _ in range(REPEATS):\n", " generate(BATCH, use_cache=True)\n", " t_cached = time.perf_counter() - t0\n", "\n", "print(f'uncached (full recompute every step): {t_uncached * 1000:.0f} ms total, {REPEATS} generations of a batch of {BATCH}')\n", "print(f'cached (KV cache): {t_cached * 1000:.0f} ms total')\n", "print(f'speedup: {t_uncached / t_cached:.2f}x')" ] }, { "cell_type": "markdown", "id": "2045db7b", "metadata": {}, "source": [ "## Self-attention vs. cross-attention\n", "\n", "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 {{lesson:self_supervised_learning}} (and \"cross-attention\" once, in Lesson {{lesson:open_vocabulary_segmentation}}'s aside about SAM 3's real decoder), but never actually defined either term.\n", "\n", "**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.\n", "\n", "**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).\n", "\n", "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 {{lesson:vision_language_models}}'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.\n", "\n", "Lesson {{lesson:visual_question_answering}}, 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." ] }, { "cell_type": "markdown", "id": "dfab8b1b", "metadata": {}, "source": [ "## Encoder-only vs. decoder-only\n", "\n", "| Model family | What it does | Generates autoregressively? | Built in this course |\n", "|---|---|---|---|\n", "| **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 {{lesson:vision_transformers}}'s ViT, Lesson {{lesson:segment_anything}}'s image encoder, Lesson {{lesson:vision_language_models}}'s CLIP, Lesson {{lesson:dino_self_distillation}}'s DINO, and more. |\n", "| **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`. |\n", "\n", "A third family, **encoder-decoder** — a separate encoder plus a decoder that cross-attends back into it — is Lesson {{lesson:visual_question_answering}}'s job next, once a genuine cross-attention mechanism actually exists to put in the table." ] }, { "cell_type": "markdown", "id": "7e64efa9", "metadata": {}, "source": [ "## Real decoder-only LLMs and VLMs\n", "\n", "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 {{lesson:visual_question_answering}} builds next — both exist in production systems today, and now there's a concrete example of each to compare." ] }, { "cell_type": "markdown", "id": "c581f4b7", "metadata": {}, "source": [ "### Exercises\n", "\n", "1. Drop the causal mask entirely in `_causal_self_attn` (let every query attend to every key, matching Lesson {{lesson:transformer_architecture}}'s encoder exactly) and retrain. Can this version still be used for autoregressive generation at all — and if you try anyway, what goes wrong?\n", "2. Raise `temperature` in `generate_sampled` from 0.8 to 2.0, and separately lower it to 0.3. How does `valid`/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?\n", "3. The KV-cache benchmark above used `MAX_LEN=9` and a batch of 512. Predict, then check, how the measured speedup changes if you (a) grow `MAX_LEN` substantially (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?" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.x" } }, "nbformat": 4, "nbformat_minor": 5 }