Hitchhiker's Guide — The Transformer From Scratch
The compressed practitioner tour. If
WARMUP.mdis the professor, this is the senior who leans over and says "here's what you actually need to remember to build and debug a transformer."
The 30-second mental model
A transformer is an embedding lookup, then N identical blocks, then a norm + LM head. A
block is two residual sublayers: X += Attn(norm(X)) and X += FFN(norm(X)). Attention is
softmax(Q·Kᵀ/√d_k + mask)·V — content-based weighted retrieval; the causal mask (-∞ above the
diagonal) makes it autoregressive. Multi-head runs that in h parallel subspaces. RoPE
injects relative position by rotating Q/K. RMSNorm keeps magnitudes sane. The FFN holds ~2/3
of the parameters. At decode time the KV-cache stores past K/V so you never recompute — and
it's provably identical to the full recompute because the last query, under the mask, sees exactly
the cache. That's the whole machine.
The numbers to tattoo on your arm
| Thing | Number |
|---|---|
| Model width | d_model (768 GPT-2-small, 4096 Llama-2-7B) |
| Heads | n_heads; d_k = d_head = d_model / n_heads |
| Attention scale | 1/√d_k (variance argument) |
| Causal mask | 0 on/below diagonal, -∞ strictly above |
| FFN width | d_ff ≈ 4·d_model (GELU) or ~8/3·d_model (SwiGLU, 3 matrices) |
| Parameter split | ~1/3 attention (4d²), ~2/3 FFN (8d²) per layer |
| RoPE base | θ_i = base^(-2i/d), base 10000 (or 1e6 long-context) |
| RMSNorm | x / sqrt(mean(x²) + ε) · weight; no mean, no bias |
| KV-cache | 2 · n_layers · T · d_model · bytes · batch (Phase 00) |
| GQA cache shrink | × n_kv/n_head (8/32 ⇒ 4× smaller) |
| Decode complexity | O(T) with cache vs O(T²) without |
Back-of-envelope one-liners
# "Head dimension for d_model=4096, 32 heads?"
d_k = 4096 / 32 = 128 → scale attention by 1/sqrt(128) ≈ 0.088
# "Where are the params in a layer (d=4096, SwiGLU d_ff≈11008)?"
attn = 4 * 4096^2 ≈ 67M
ffn = 3 * 4096 * 11008 ≈ 135M → FFN ≈ 2/3. Always.
# "Is my causal mask right? Token 0's attention output should equal V[0]."
SDPA(Q,K,V, causal_mask(T))[0] == V[0] # row 0 sees only itself
# "RoPE sanity: position 0 is identity, and the norm never changes."
rope(v, 0) == v ; ||rope(v, p)|| == ||v|| for all p
The framework one-liners (where these live in real tools)
import torch.nn.functional as F
# scaled dot-product attention, fused (FlashAttention under the hood):
out = F.scaled_dot_product_attention(Q, K, V, is_causal=True) # the mask, for free
torch.nn.MultiheadAttention(embed_dim, num_heads) # the whole MHA module
torch.nn.RMSNorm(d_model) # RMSNorm (HF: LlamaRMSNorm)
# RoPE: HF LlamaRotaryEmbedding + apply_rotary_pos_emb(q, k, cos, sin)
# KV-cache: model.generate(..., use_cache=True); past_key_values / DynamicCache
# config knobs that change the math: num_attention_heads, num_key_value_heads (GQA),
# rope_theta, rms_norm_eps, intermediate_size (d_ff)
# nanoGPT model.py is this entire lab in ~150 readable lines.
War stories
- The mask that leaked the future. Off-by-one in the causal mask let each token see one future
position. Training loss dropped beautifully (the model was peeking at the answer); inference was
gibberish (the future it leaned on wasn't there). Symptom: great train loss, broken generation.
Always test
out[0] == V[0]— token 0 must see only itself. - The RoPE base mismatch. Fine-tuned a
rope_theta=1e6long-context base model but loaded it with the default10000. Short prompts were fine; anything past a few thousand tokens degraded into nonsense because the positional angles no longer matched what the weights expected. Position config must match the checkpoint exactly — RoPE base andmax_position_embeddings. - The forgotten
√d_k. Someone "simplified" attention by dropping the scale. Withd_k=128the scores blew up, softmax went one-hot, gradients died, and the model wouldn't learn. The fix was one/ math.sqrt(d_k). - The KV-cache that didn't match. A custom decode loop diverged from the prefill path after a few tokens — they'd applied RoPE to the cached K again on each step. The cache stores already-rotated keys; rotate once. Golden test: cached decode must equal full-recompute's last row to tolerance.
- The pre-norm/post-norm swap. Ported a model assuming post-norm; the deep stack wouldn't train (activations blew up). Modern models are pre-norm — normalize the sublayer input, keep the residual stream clean.
Vocabulary (rapid-fire)
- Q / K / V — query/key/value, three learned linear projections of the input.
- SDPA — scaled dot-product attention,
softmax(QKᵀ/√d_k + mask)·V. - Causal mask —
-∞above the diagonal; makes attention autoregressive. - Head /
d_k— a subspace of widthd_model/n_heads; the per-head key dimension. - MQA / GQA — share K/V across query heads (1 / a few) to shrink the KV-cache.
- KV-cache — stored past K/V so decode is
O(T); the inference memory wall (Phase 00). - RoPE — rotary position embedding; rotates Q/K so the score is relative.
- RMSNorm — mean-free, bias-free normalization; the modern default.
- SwiGLU — gated FFN
down(silu(gate) ⊙ up); where ~2/3 of params live. - Residual stream — the running
+Xhighway every block reads and writes. - Pre-norm — normalize the sublayer input; stable for deep stacks.
- FlashAttention — IO-aware tiled attention; exact, never materializes
T². - Logits — the per-vocab scores; sampling turns a row into the next token.
Beginner mistakes
- Forgetting the
1/√d_kscale (softmax saturates, training dies). - An off-by-one causal mask (future leak: low train loss, broken inference).
- Applying RoPE to the residual stream or re-rotating cached keys instead of rotating Q/K once.
- Thinking "the model outputs a token" — it outputs logits; sampling is a separate step.
- Quoting attention as the parameter-heavy part — it's the FFN (~2/3).
- Using
==on floats in tests instead of a tolerance (softmax/RoPE are floating-point). - Confusing GQA's win (memory / bandwidth) with a FLOP speedup.
- Treating FlashAttention as an approximation — it's exact, just IO-aware.
The one thing to take away
A transformer is embed → N×(residual attention + residual FFN) → norm → logits, and attention is
softmax(Q·Kᵀ/√d_k + mask)·V. If you can write those two lines from memory, explain the causal
mask and √d_k, and prove the KV-cache equals the full recompute, you can build, read, and debug any
modern decoder model. Everything else in the curriculum is a change to this machine.