Part I — Basics

The Transformer, from scratch

Tokenization (01) handed you a stream of integer ids and a number — the compression ratio that turns bytes into ≈ bytes / 4 tokens. That is all it produced: a sequence of ids and the certainty that something must now turn each prefix into a prediction of the next id. This lesson builds that something — the decoder-only Transformer — shape by shape, deriving every constant (why √d_head, why 4d, why pre-norm) rather than asserting it, and ends by counting its parameters and FLOPs so the budget chain from 00 has its first concrete number to spend.

1 · Embeddings and the residual stream

Notation, fixed once: B = batch, T = sequence length in tokens, V = vocab size (50,257 for GPT-2), d = model width (the residual stream), h = number of heads, d_head = d / h per-head width, L = number of blocks.

An autoregressive LM factors the joint distribution as a product of conditionals, p(x₁…x_T) = ∏_t p(x_t | x_<t), and trains by minimizing the summed per-token negative log-likelihood. That is the whole objective. Every shape below exists to compute all T of those conditionals in one differentiable pass.

The input is idx of shape (B, T), int64 ids. Two lookup tables turn ids into vectors: a token table W_te of shape (V, d) and an absolute-position table W_pe of shape (T_max, d). We add them, not concatenate:

tok = W_te[idx]          # (B, T, d)  — what token is here
pos = W_pe[:T]           #    (T, d)  — where in the sequence (broadcast over B)
x   = tok + pos          # (B, T, d)  — the residual stream is born

Addition keeps the width at d (concatenation would force every downstream matrix wider for no measured gain), and any linear layer can still read off either component since a projection of a sum is the sum of projections. That tensor x of shape (B, T, d) is the residual stream: every block reads it, computes a correction, and adds the correction back. Its shape never changes from here to the final norm — the entire model is a sequence of (B,T,d) → (B,T,d) refinements. (GPT-2 uses a learned position table, which caps usable length at T_max and cannot extrapolate; lesson 3 replaces it with RoPE.)

2 · Self-attention, shape by shape

Attention is the only operation in the block that moves information between positions. Each position emits three vectors: a query (what am I looking for?), a key (what do I offer?), and a value (what I'll hand over if matched). One fused linear produces all three at once — one GEMM, one read of x from memory, instead of three:

qkv = x @ W_qkv                              # (B, T, 3d)
q, k, v = qkv.chunk(3, dim=-1)               # each (B, T, d)
# reshape (B,T,d) -> (B,T,h,d_head) -> transpose -> (B,h,T,d_head)
q, k, v = [t.view(B,T,h,d_head).transpose(1,2) for t in (q,k,v)]
scores  = (q @ k.transpose(-2,-1)) / d_head**0.5   # (B, h, T, T)
scores  = scores.masked_fill(causal_mask[:T,:T]==0, float('-inf'))
attn    = softmax(scores, dim=-1)            # (B, h, T, T)
out     = attn @ v                           # (B, h, T, d_head)
out     = out.transpose(1,2).reshape(B,T,d)  # merge heads -> (B, T, d)
out     = out @ W_o                          # output projection (B, T, d)

Why divide by √d_head

Take the entries of q and k to be roughly i.i.d. with mean 0 and variance 1. A single score is a dot product of two d_head-vectors: a sum of d_head independent products, each of variance ≈ 1. Variance of the sum is d_head; its standard deviation is √d_head. So unscaled logits grow like √d_head. Feed those into softmax and the largest logit dominates exponentially — softmax saturates toward one-hot, and a one-hot softmax has a near-zero Jacobian, so the gradient vanishes and the layer stops learning. Dividing by √d_head pins the score variance back to ≈ 1 regardless of head width, keeping softmax in its informative regime.

Worked: GPT-2 small has d = 768, h = 12 ⇒ d_head = 64, so unscaled std ≈ √64 = 8. A logit 8 above its neighbors gives e⁸ ≈ 3,000× the mass — essentially one-hot before training even starts. With /√64 the std is 1 and attention can actually learn where to look. This is one of the rare constants the architecture derives rather than tunes.

The causal mask

The mask is forced, not stylistic. The autoregressive factorization p(x₁,…,x_T) = ∏_t p(x_t | x_<t) says the predicted distribution at position t is conditioned only on tokens before t. To train all T−1 of those predictions in one parallel pass, position t's output must depend only on positions < t: if it could attend to t+1,…,T, the next-token target would leak straight into the input and the model would learn the trivial copy. Zeroing attention from t to every j > t — the causal mask — is therefore the exact, forced mechanism that makes the parallel forward pass compute the autoregressive factorization honestly.

Concretely the mask is a lower-triangular T × T matrix; entries above the diagonal are set to −∞ before softmax, so their probability is exactly 0. Position t attends only to keys j ≤ t. This is precisely what makes "one forward pass = T−1 honest supervised targets" true: the prediction at t never sees the answer at t+1. Drop the mask and train loss crashes to zero (the model just copies the next token) while inference — where future tokens do not exist — is impossible.

Multi-head split and merge

Why slice d into h heads of width d_head = d/h instead of one big head? A single head computes one similarity geometry — one notion of "relevant." Splitting gives h independent attention patterns (interpretability work finds previous-token heads, induction heads, copy heads emerging on their own), recombined by the output projection W_o. The reshape (B,T,d) → (B,T,h,d_head) → (B,h,T,d_head) is a free view + transpose: the matmul then batches over the leading (B,h) dims and multiplies the trailing (T, d_head) matrices. Total FLOPs are unchanged because the work scales with the full width d; splitting only rearranges where the bytes sit. So you get h attention patterns for free.

One cost the diagram makes visible: the score tensor is (B, h, T, T) — quadratic in T. Materializing it is fine at T = 1,024 but becomes the memory wall at long context, which is the entire reason FlashAttention exists (see the kernel view linked below).

3 · The MLP: d → 4d → d

Attention moved information between positions; it did not transform it much (it is mostly weighted averaging plus a linear projection). The per-position MLP does the heavy nonlinear feature work, applied identically and independently at every t:

h1 = GELU(x @ W_1)   # W_1: (d, 4d)   ->  (B, T, 4d)   expand
y  = h1 @ W_2        # W_2: (4d, d)   ->  (B, T,  d)   contract

It expands the width 4× into a hidden space, applies a pointwise nonlinearity, and contracts back. GELU is a smooth, non-monotonic gate (x · Φ(x), close to SiLU/Swish) — not simply a rounded ReLU (Φ = standard-normal CDF; it actually dips slightly negative for small negative inputs). The smoothness gives a small nonzero gradient for slightly-negative inputs and trains marginally better than ReLU at negligible cost. The 4× ratio is the historical Transformer width: wide enough to give the nonlinearity room, narrow enough not to swamp the budget. Conceptually the MLP is the block's memory — the place that stores "if the features look like X, write feature Y into the stream" associations. It also dominates the parameters: per block attention holds 4d² (QKV is 3d², output proj d²) while the MLP holds 8d² (two d × 4d matrices) — about 2/3 of every block's weights live in the MLP, a fact people routinely underestimate.

4 · Pre-norm residuals make depth trainable

A block is two sublayers, each wrapped in a residual add with the normalization placed inside the branch:

x = x + Attn(LN(x))      # pre-norm
x = x + MLP (LN(x))      # pre-norm

Two mechanisms are doing the work, and you need both to train a deep stack.

Concretely: post-norm routes the backward gradient through L LayerNorm Jacobians in series, so small errors compound and the run spikes unless you baby it with warmup. Pre-norm leaves a clean additive skip path open, so the gradient at layer 0 is the gradient at layer L plus well-behaved corrections. The cost of switching is exactly zero parameters and zero FLOPs, and it removes a hyperparameter — the kind of free win you always take. After the last block a final LN normalizes the stream before the head.

5 · Weight tying, the LM head, and the loss

The head turns each d-vector into a distribution over the vocabulary by multiplying by a (d, V) matrix. Instead of a fresh matrix, tie it to the embedding: the head is W_teᵀ.

x      = LN_final(x)        # (B, T, d)
logits = x @ W_te.T         # (B, T, V)   — head shares storage with the embedding
loss   = cross_entropy(logits[:, :-1], idx[:, 1:])   # scalar

Tying is sound because input and output spaces are the same vocabulary: "the vector for token v" should mean the same thing whether you read v in or predict v out. It saves V·d parameters (for GPT-2: 50,257 × 768 ≈ 38.6M, a large share of the small model), gives the embedding rows far more gradient updates, and empirically improves perplexity.

The loss is cross-entropy of logits at positions [0..T−2] against targets at [1..T−1] — i.e. each position predicts the next token. Cross-entropy of a one-hot target is exactly the negative log-likelihood of the true next token, −log p_θ(x_t+1 | x_≤t); minimizing the mean over all B·T positions is maximum-likelihood training of the joint distribution. There is no separate objective — the entire model exists to make this one number small.

6 · Parameter count and the FLOP preview

Count per block, then scale by L and add the embeddings. The O(d) bias/norm terms are negligible at scale, so:

The block term is quadratic in d and linear in L; the embeddings are only linear in d. So at any real scale the model is "all blocks." Check it on GPT-2 small (L=12, d=768, V=50,257, T_max=1,024): blocks = 12 · 12 · 768² ≈ 85M; embeddings ≈ (50,257 + 1,024) · 768 ≈ 39M; total ≈ 124M — exactly GPT-2 small's published count.

The FLOP preview, which the whole track turns on: a matmul of shapes (·, in) × (in, out) costs 2·in·out FLOPs per row (one multiply + one add per element). Summed over the model's weight matrices, a forward pass costs ≈ 2N FLOPs per token; the backward pass costs about twice the forward, so training is ≈ 6N per token. Over D training tokens that is the budget identity from lesson 0: C ≈ 6·N·D. (This counts the dense linears and ignores attention's O(T²) term, which is a rounding error until T approaches d — lesson 5 makes the full accounting precise.)

7 · Drive the shapes and the budget

Move the sliders. The left column prints every intermediate tensor shape for one forward pass; the right column breaks the parameter count into emb / attn-per-block / mlp-per-block, multiplies the block by L, ties the head to 0, and reports total N plus the 2N forward and 6N train FLOPs per token. Push d up and watch the 12Ld² block term swamp the embeddings; push L up and watch N — and the FLOPs you'll pay per token — climb linearly. Dragging T resizes every activation (watch the score tensor (B,h,T,T) balloon) but leaves N untouched: the position table W_pe is (T_max, d), fixed at the model's max context (1,024 here), so it never grows with the runtime sequence length. The constraint d mod h = 0 is enforced (changing d snaps h to a divisor).

Failure modes & checklist

Checkpoint

Where this points next

You now have a Transformer that trains and a parameter/FLOP count to plan with. But this is the 2019 block. Push it to modern depth and scale and five cracks open at once: learned absolute positions cap and cannot extrapolate context; LayerNorm's mean-centering and bias are dead weight; plain GELU MLPs are wasteful per parameter; the full multi-head KV is a fat cache you'll pay for at inference; and the residual stream's scale grows with depth so big runs spike and lean on warmup. Each is a real efficiency or stability leak under a fixed budget. Lesson 03, The modern recipe — what changed since GPT-2, swaps in RMSNorm, RoPE, SwiGLU, GQA, and bias-free linears with QK-norm — five changes that each buy stability or efficiency at essentially zero cost.

Interview prompts

• Why divide attention scores by √d_head? (§2 — dot products of unit-variance vectors have variance d_head, so unscaled logits grow like √d_head, saturating softmax to one-hot with near-zero gradient; the scale pins score variance to ≈1 independent of head width.)
• What exactly does the causal mask buy you, and what breaks without it? (§2 — it lets one forward pass produce T−1 honest next-token targets in parallel by forbidding position t from seeing >t; without it train loss collapses to ≈0 by copying the future and inference is impossible.)
• Why pre-norm over the original post-norm? (§4 — pre-norm leaves a clean additive skip on the backward path and bounds residual variance ~linearly in L, so it trains stably at depth without warmup; post-norm stacks L norm Jacobians in series and spikes.)
• Derive the parameter count of a decoder block. (§6 — attention is QKV 3d² + output proj d² = 4d²; MLP is two d×4d matrices = 8d²; total 12d² per block, so N ≈ 12Ld² + V·d + T_max·d.)
• Where do most of a block's parameters live, and why does it matter? (§3/§6 — ~2/3 in the MLP (8d² vs attention's 4d²); it dictates which part dominates FLOPs and memory and why MoE later targets the MLP specifically.)
• Why is training ≈ 6N FLOPs/token? (§6 — a matmul costs 2·in·out, summing over weights gives 2N forward; backward ≈ 2× forward, so total ≈ 6N/token ⇒ C ≈ 6ND, the budget identity.)
• What does weight tying do and when can't you use it? (§5 — shares W_te as the head, saving V·d params and improving PPL via shared semantics + more gradient updates; you can't tie when input and output vocabularies differ.)

Companion reads: gpt_mini · 01 Architecture walks the same block line-by-line as runnable code; gpu_kernels · 11 Attention kernel pipeline shows how that (B,h,T,T) score tensor is actually computed on the GPU.