Hitchhiker's Guide — The Transformer From Scratch

The compressed practitioner tour. If WARMUP.md is 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

ThingNumber
Model widthd_model (768 GPT-2-small, 4096 Llama-2-7B)
Headsn_heads; d_k = d_head = d_model / n_heads
Attention scale1/√d_k (variance argument)
Causal mask0 on/below diagonal, -∞ strictly above
FFN widthd_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)
RMSNormx / sqrt(mean(x²) + ε) · weight; no mean, no bias
KV-cache2 · n_layers · T · d_model · bytes · batch (Phase 00)
GQA cache shrink× n_kv/n_head (8/32 ⇒ 4× smaller)
Decode complexityO(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=1e6 long-context base model but loaded it with the default 10000. 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 and max_position_embeddings.
  • The forgotten √d_k. Someone "simplified" attention by dropping the scale. With d_k=128 the 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 width d_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 +X highway every block reads and writes.
  • Pre-norm — normalize the sublayer input; stable for deep stacks.
  • FlashAttention — IO-aware tiled attention; exact, never materializes .
  • Logits — the per-vocab scores; sampling turns a row into the next token.

Beginner mistakes

  • Forgetting the 1/√d_k scale (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.