chapter 06 / transformer · the single most important chapter of the course · estimated study time 180-240 min

Transformer
Attention Is All You Need

Three modifications: from Bahdanau to Self-Attention

At the end of Chapter 5 we recast attention as "differentiable soft retrieval": a query finds keys and pulls values weighted by relevance. The 2017 paper Attention Is All You Need made three deceptively simple modifications and then threw the RNN out entirely:

• ① Query yourself: no longer the decoder querying the encoder—instead, a single sequence queries itself internally. The representation of "mat" is shaped jointly by "cat," "sat," and "on." Contextualization no longer relies on recurrent propagation.
• ② Separation of powers: Q/K/V: in Bahdanau, $\mathbf{h}_j$ served as both key and value; now each token's vector $\mathbf{x}$ passes through three independently learnable projections that play three roles: $\mathbf{q} = W_Q\mathbf{x}$ (what am I looking for), $\mathbf{k} = W_K\mathbf{x}$ (what can find me), $\mathbf{v} = W_V\mathbf{x}$ (what do I hand over once found). "How I am retrieved" is decoupled from "what content I provide."
• ③ Scoring simplified from an MLP to a dot product: $e_{ij} = \mathbf{q}_i^\top \mathbf{k}_j$—all the scores across the entire sequence become a single matrix multiplication $QK^\top$, the GPU's favorite.

The payoff is structural: in an RNN, information at position 1 must travel $n$ steps to affect position $n$ (Chapter 5's exponential decay happens along the way); in self-attention any two positions are one step apart. An RNN must wait serially for $\mathbf{h}_{t-1}$; self-attention computes all positions simultaneously. The price is $O(n^2)$—the answer to Chapter 5's exercise, and the starting point for the KV cache in Chapter 9.

Scaled dot-product attention and the derivation of √d

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

That $\sqrt{d_k}$ is the paper's only "magic constant," and it is worth working the derivation through. Assume the components of $\mathbf{q}, \mathbf{k}$ are independent, with mean 0 and variance 1. The dot product $\mathbf{q}^\top\mathbf{k} = \sum_{i=1}^{d_k} q_i k_i$ is a sum of $d_k$ independent terms:

$$\mathbb{E}[\mathbf{q}^\top\mathbf{k}] = 0, \qquad \text{Var}(\mathbf{q}^\top\mathbf{k}) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = d_k$$

That is, the typical magnitude of the dot product is $\sqrt{d_k}$. When $d_k = 128$, scores easily reach ±11, and inside the softmax $e^{11}$ and $e^{-11}$ differ by a factor of $10^9$—the output saturates into a one-hot. Recall Chapter 2: where the softmax saturates, gradients vanish and attention can no longer learn. Dividing by $\sqrt{d_k}$ normalizes the variance back to 1, keeping the softmax in a region where gradients stay healthy. This is a microcosm of the entire Transformer's design philosophy: every component serves the goal of "letting gradients flow smoothly."

The causal mask: what makes GPT a GPT

The training objective of a language model (the star of Chapter 7) is to predict the next word, so position $i$ must never see position $j > i$—otherwise it is just copying the answer. The implementation is crude but efficient: before the softmax, set the entire upper triangle to $-\infty$, and after the softmax those positions have weight exactly 0.

Interactive lab: a step-by-step self-attention calculator

A toy size of 5 tokens with $d{=}4$, but the computation is real (the matrices really are $QK^\top$). Required experiments: ① walk through all five stages with "Next step," checking each step against the formula in §2; ② toggle the "causal mask" switch back and forth at stages ③④—watch the upper-right triangle go from −∞ to having values, and how the softmax row distributions reshuffle (turning the mask off gives BERT-style bidirectional attention); ③ count at stage ④: why is the first row always 1.00? ("cat" can only see itself).

Multi-head attention: parallel perspectives

A single attention can only learn one kind of "relevance"—but "mat" needs to attend simultaneously to the grammatical subject (cat), the spatial relation (on), and the modifying structure. The solution: split the $d$-dimensional space into $h$ parts, do attention independently in each, then concatenate and project:

$$\text{head}_i = \text{Attention}(XW_Q^{(i)},\, XW_K^{(i)},\, XW_V^{(i)}), \qquad \text{MHA}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\, W_O$$

Each head works in a $d_k = d/h$ subspace (GPT-2: $d{=}768, h{=}12, d_k{=}64$), and the total computation is the same as a single head—multiple perspectives for free. Interpretability research (Anthropic's work, linked in the appendix's sources) has indeed found heads with clear divisions of labor: ones that fixate on the previous word, ones that track bracket matching, "induction heads" that carry object information over to the verb… and, of course, plenty of redundant heads (which can be pruned).

Positional encoding: from sinusoids to RoPE

Pure attention has a fatal blind spot: it is a set operation—shuffle the input order and the output shuffles the same way (permutation-equivariant); "cat sat mat" and "mat sat cat" look identical to it. Positional information must be injected by hand:

• Sinusoidal positional encoding (the original paper): for position $p$, add $\sin(p / 10000^{2i/d})$ to dimension $2i$ and $\cos(\cdot)$ to dimension $2i{+}1$—different dimensions are "rulers" of different wavelengths, from the word level to the paragraph level. Parameter-free and, in theory, extrapolatable.
• Learned absolute positional embeddings (GPT-2/BERT): one trainable vector per position, simple and effective, but a model that only saw 1024 during training dies at 1025.
• RoPE rotary positional encoding (Jianlin Su, 2021, the signature work of "Scientific Spaces" in the appendix's sources): standard equipment in modern LLMs (LLaMA/Qwen/DeepSeek…). The idea is elegant: treat each pair of adjacent dimensions of $\mathbf{q}, \mathbf{k}$ as a point on the complex plane, and rotate it by angle $m\theta_i$ according to position $m$. The key property—after rotation, the dot product depends only on relative position: $$\langle R_m\mathbf{q},\, R_n\mathbf{k} \rangle = \langle \mathbf{q},\, R_{n-m}\mathbf{k} \rangle$$ "Five words apart" is the same geometric relationship anywhere in the sequence. The form of absolute position, the substance of relative position, and not a single added parameter. Long-context extension techniques (NTK interpolation, YaRN) are all built on adjusting RoPE's frequencies—more on this in Chapter 7.

FFN, normalization, and residuals: assembling a GPT block

Attention is responsible for transporting information (between tokens), but we still need a component to process information (within a token): the feed-forward network (FFN), which independently applies a two-layer MLP to each position, expanding to $4d$ in the middle:

$$\text{FFN}(\mathbf{x}) = W_2\, \text{GELU}(W_1 \mathbf{x}), \qquad W_1: d \to 4d, \quad W_2: 4d \to d$$

It accounts for about 2/3 of the Transformer's parameters, and interpretability research views it as the model's "key-value store"—factual knowledge is stored mainly here. Modern LLMs use the SwiGLU variant (the last row of Chapter 3's activation-function table). Assemble all the parts into a single block (the Pre-LN structure, normalizing before entering the sublayer—far more stable to train than the original paper's Post-LN, with no warmup magic needed even at great depth):

x = x + MHA(LayerNorm(x))     # sublayer 1: attention + residual
x = x + FFN(LayerNorm(x))     # sublayer 2: feed-forward + residual

Those two residual connections are exactly Chapter 4's $y = F(x) + x$—GPT is just a stack of residual blocks (with the gradient highway paved), plus the token embedding at the bottom and the lm_head at the top (in GPT-2 the two share weights, saving half the embedding parameters). The full architecture is now complete before your eyes:

tokens → Embedding(+position) → [GPT Block] × L → LayerNorm → lm_head → next-word probability

Counting parameters: verifying GPT-2's 124M by hand

The basic skill for reading any LLM configuration table. Per block: attention $W_Q, W_K, W_V, W_O$ are each $d^2$ → $4d^2$; the FFN's two matrices are $d \times 4d$ → $8d^2$. That totals $12d^2$ per block, giving the famous quick-estimate formula:

$$\text{non-embedding parameters} \approx 12\, L\, d^2$$

Verify GPT-2 small ($L{=}12, d{=}768$, vocabulary 50257): block parameters $12 \times 12 \times 768^2 \approx 85\text{M}$, embedding $50257 \times 768 \approx 38.6\text{M}$ (shared with lm_head), totaling $\approx 124\text{M}$ ✓. Apply the same formula to LLaMA-2 7B ($L{=}32, d{=}4096$): $12 \times 32 \times 4096^2 \approx 6.4\text{B}$, plus embedding gives exactly ~7B. From now on, whenever you see an "xB model," you can work backward to its skeleton.

Code in practice: writing the heart of GPT in ~60 lines

python · minimal_gpt_block.py (nanoGPT-style condensed version)

import torch, torch.nn as nn, torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d, h, n_max=1024):
        super().__init__()
        self.h, self.dk = h, d // h
        self.qkv = nn.Linear(d, 3 * d)          # Q,K,V in one projection
        self.proj = nn.Linear(d, d)             # W_O
        mask = torch.tril(torch.ones(n_max, n_max))
        self.register_buffer('mask', mask)      # lower triangle = visible

    def forward(self, x):
        B, N, d = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        # split into heads: (B, N, d) → (B, h, N, dk)
        q, k, v = (t.view(B, N, self.h, self.dk).transpose(1, 2) for t in (q, k, v))
        att = (q @ k.transpose(-2, -1)) / self.dk ** 0.5          # ② scaled dot product
        att = att.masked_fill(self.mask[:N, :N] == 0, float('-inf'))  # ③ causal mask
        att = F.softmax(att, dim=-1)                              # ④ attention weights
        y = att @ v                                               # ⑤ weighted transport
        y = y.transpose(1, 2).contiguous().view(B, N, d)          # reassemble to (B, N, d)
        return self.proj(y)

class Block(nn.Module):
    """A complete GPT block: every LLM you're reading about is this, repeated L times"""
    def __init__(self, d=768, h=12):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = CausalSelfAttention(d, h)
        self.ffn = nn.Sequential(
            nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # Pre-LN + residual (Chapter 4's highway)
        x = x + self.ffn(self.ln2(x))
        return x

x = torch.randn(2, 10, 768)              # batch=2, 10 tokens
print(Block()(x).shape)                  # → (2, 10, 768)
print(sum(p.numel() for p in Block().parameters()) / 1e6, 'M')  # ≈ 7.1M ≈ 12d²/1e6 ✓

The grand finale: Karpathy writes a complete, trainable GPT from an empty file in 3.5 hours (including data loading, the training loop, and sampling/generation). This is the only video in the entire course marked "you must code along"—after you finish it, your understanding of LLMs will shift from "I've read the derivation" to "I've built one with my own hands":

Chapter quiz