Warmup Guide — Attention & Transformers

Zero-to-expert primer for Phase 04: the transformer derived piece by piece — why each component exists, what breaks without it — culminating in the decoder-only GPT architecture you implement in the mini-transformer lab.

Table of Contents

Chapter 1: The Problem Attention Solves

Phase 03 ended with the RNN's twin constraints: information from token 5 reaches token 500 only through 495 sequential state updates (a lossy game of telephone), and training can't parallelize over time. Attention dissolves both with one move: let every position directly query every other position, with content-dependent weights, all positions computed simultaneously. Path length between any two tokens: 1. Training parallelism: total. The price — O(n²) pairs — was accepted in 2017 and has been the field's central engineering battle ever since (this track's Phase 09; the model-accuracy track's FlashAttention). Phase 04 is where you internalize what that price buys.

Chapter 2: Scaled Dot-Product Attention, Derived

Build it from the retrieval metaphor: each token asks a question and offers an answer. From each input $x_i$, three learned projections:

  • query $q_i = W_Q x_i$ — what am I looking for?
  • key $k_i = W_K x_i$ — what can I be found by?
  • value $v_i = W_V x_i$ — what do I contribute if selected?

Relevance of $j$ to $i$: the dot product $q_i \cdot k_j$. Softmax the scores into weights; output the weighted sum of values:

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

Why each piece:

  • Separate Q and K (not $x_i \cdot x_j$ directly): asymmetry — "looking for" and "findable as" are different roles ("it" queries for nouns; it doesn't offer itself as one).
  • Separate V: what a token contributes differs from what it matches on.
  • $\sqrt{d_k}$: components i.i.d. with unit variance make $q \cdot k$ have variance $d_k$; at $d_k = 64$, unscaled logits of magnitude ~8 saturate softmax — winner-take- all weights and vanishing gradients through the softmax. Dividing by $\sqrt{d_k}$ restores unit variance. (Verify empirically in the lab — it's a two-line experiment.)
  • Softmax: turns arbitrary scores into a convex combination — outputs stay in the values' span, bounded, differentiable.

Chapter 3: Causal Masking — Training on Every Position at Once

A language model must not see the future (Phase 03 Ch. 1's conditional). In attention, enforcement is brutally simple: before softmax, set $S_{ij} = -\infty$ for $j > i$ (an upper-triangular mask). After softmax those weights are exactly 0.

The consequence that makes LLM training affordable: one forward pass over a sequence trains all n positions simultaneously — position i's output predicts token i+1, every position is a training example, and the mask guarantees no leakage. This "teacher forcing in parallel" is the transformer's economic engine; the lab's correctness test (perturb a future token, assert outputs at earlier positions are bit-identical) is the single most important test you'll write in this phase.

Chapter 4: Multi-Head Attention

One attention computes one weighted average per position — one "relation type" at a time. Multi-head splits $d_{model}$ into $h$ subspaces of $d_k = d_{model}/h$, runs attention independently in each, concatenates, and projects ($W_O$):

$$\text{MHA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h),W_O$$

Same total FLOPs as one full-width head — but $h$ different relation patterns (syntactic heads, positional heads, rare-token heads emerge under analysis). Two implementation truths the lab teaches: it's all done as one batched tensor op (reshape to (B, h, n, d_k), never a Python loop over heads), and the (B, h, n, n) attention-weights tensor is the memory hog — materializing it is what FlashAttention later avoids (model-accuracy Phase 08, when you get there).

Chapter 5: The Rest of the Block — FFN, Residuals, LayerNorm

A transformer block is attention plus three pieces, each load-bearing:

  • FFN: per-position MLP, $W_2,\text{GELU}(W_1 x)$ with hidden dim $4\times d_{model}$. Attention moves information between positions; the FFN processes it pointwise — ~⅔ of the model's parameters live here, and mechanistic work reads it as key-value memory over features. No interaction across positions (that's attention's job alone).
  • Residual connections: $x + \text{Sublayer}(x)$ — the gradient highway (the LSTM cell-state idea, Phase 03 Ch. 5, reborn as architecture). Depth-100 stacks train because identity paths exist around every sublayer.
  • LayerNorm: re-center/re-scale each token's vector (learned $\gamma, \beta$) — keeps activation scales stable across depth. Pre-norm placement (norm before each sublayer, used by GPT-2 onward) leaves the residual path untouched and trains stably without warmup heroics; post-norm (the 2017 original) normalizes the sum and destabilizes deep stacks. Your lab uses pre-norm; know why.

Block: x = x + Attn(LN(x)); x = x + FFN(LN(x)). Stack N times. That's the architecture.

Chapter 6: Position Information

Attention is permutation-equivariant — shuffle inputs, get shuffled outputs. Order must be injected:

  • Learned absolute embeddings (GPT-2, and your lab): a trainable vector per position index, added to token embeddings. Simple, effective, but no extrapolation beyond trained length and no explicit relative structure.
  • Sinusoidal (the 2017 original): fixed sin/cos at geometric frequencies — no parameters, theoretical extrapolation (rarely realized in practice).
  • The modern answers — RoPE and ALiBi — rotate Q/K by position-dependent angles (relative position emerges in the dot product) or bias logits by distance. They're covered in depth in the model-accuracy track's Phase 02 warmup (Ch. 5); for this phase, know that learned-absolute is the pedagogical baseline and why its extrapolation fails (untrained rows in the position table).

Chapter 7: The Full GPT — Assembly and Parameter Accounting

tokens → token_emb (V×d) + pos_emb (n_ctx×d)
       → N × [pre-norm attention block + pre-norm FFN block]
       → final LayerNorm → lm_head (d×V) → logits

Two assembly details with outsized importance:

  • Weight tying: lm_head shares the token-embedding matrix (transpose) — saves V×d parameters (significant at small scale) and improves quality; both GPT-2 and your lab do it.
  • Parameter accounting (do this once by hand — the lab asks for it): per block ≈ $12 d^2$ ($4d^2$ attention QKVO + $8d^2$ FFN); total ≈ $12 N d^2$ + embeddings $Vd$. GPT-2-small: N=12, d=768, V=50257 → ~85M block + ~39M embedding ≈ 124M. ✓ Being able to do this arithmetic is how you sanity-check any config — and it's the FLOPs-per-token estimate ($\approx 2 \times$ params) from the model-accuracy track's Phase 07 in embryo.

Chapter 8: Encoder vs Decoder vs Encoder-Decoder

The same blocks, three wirings — know which and why:

  • Decoder-only (GPT, LLaMA, ~everything generative): causal mask everywhere; one stack does both understanding and generation. Won because of training simplicity, the in-context-learning emergent bonus, and KV-cache-friendly inference.
  • Encoder-only (BERT): bidirectional attention (no mask) + masked-token training — better representations per parameter for understanding tasks (classification, retrieval embeddings — your Phase 07 RAG encoders are these). Cannot generate autoregressively.
  • Encoder-decoder (T5, translation, Whisper): bidirectional encoder over input, causal decoder with cross-attention (decoder queries, encoder keys/values) — still the right shape when input and output are genuinely different sequences/ modalities.

Lab Walkthrough Guidance

Lab 04 — Mini-Transformer (decoder-only GPT, trained on the Phase 03 corpus):

  1. Single-head causal attention first; write the no-future-leakage test (perturb token j, assert positions < j unchanged) before training anything.
  2. Verify the $\sqrt{d_k}$ claim empirically: log attention-logit variance with and without scaling at d_k = 64.
  3. Multi-head via the batched reshape — test: output matches a loop-over-heads reference within 1e-6.
  4. Assemble pre-norm blocks + weight tying; do the Chapter 7 parameter count and assert it matches sum(p.numel()) exactly — off-by-anything means a wiring bug.
  5. Train on the same data as your char-RNN with the same budget; compare loss curves and samples — the gap is the phase's thesis. Then overfit a tiny batch to near-zero loss (the standard can-it-learn sanity check) before any long run.

Success Criteria

You are ready for Phase 05 when you can, from memory:

  1. Derive attention from the Q/K/V retrieval story, justifying all three projections and the $\sqrt{d_k}$.
  2. Explain how the causal mask enables all-positions-at-once training and write its correctness test.
  3. State what FFN does that attention doesn't (and vice versa), and why pre-norm.
  4. Count GPT-2-small's parameters on paper to within a few percent.
  5. Choose decoder-only vs encoder-only vs enc-dec for: chat model, embedding model, speech-to-text — with reasons.
  6. Name the O(n²) cost's two manifestations (training compute, inference KV/attention) — the bridge to Phases 05 and 09.

Interview Q&A

Q: Why three separate projections Q, K, V instead of using the embeddings directly? Roles differ: what a token searches for (Q), what it's discoverable by (K), and what it contributes when found (V) are three different functions of its content. Collapsing them forces symmetric attention ($x_i \cdot x_j$) and couples routing to content transport. Empirically and mechanistically, the asymmetry is where attention's expressiveness lives — heads implement little programs like "pronouns query for recent nouns," which need Q ≠ K.

Q: Walk me through why training a transformer LM is one parallel pass but generating is sequential. Training: all tokens are known, so all n next-token predictions compute simultaneously under the causal mask — the mask, not time, enforces order. Generation: token t+1's identity depends on sampling from position t's output — an inherent data dependency no parallelism removes. That asymmetry creates the prefill/decode split and the KV cache (Phase 09): cache K/V of the fixed prefix; each new token costs one query against cached history instead of recomputing the past.

Q: Remove the FFN entirely — what happens? You lose all per-position nonlinear processing; the network becomes (normed, gated) weighted averaging of value vectors — outputs confined near the span of linear transforms of inputs, and stacking attention alone is known to collapse toward rank-deficient mixing. Practically: most of the parameter budget and the "memory/ feature computation" capacity is in FFNs; models degrade catastrophically. The clean division — attention mixes across positions, FFN computes within them — is the architecture's actual design principle.

References

  • Vaswani et al., Attention Is All You Need (2017) — arXiv:1706.03762
  • Radford et al., GPT-2 report (2019) — the decoder-only recipe your lab follows
  • Karpathy, Let's build GPT from scratch — youtube.com/watch?v=kCc8FmEb1nY — pairs exactly with this lab
  • The Illustrated Transformer and The Annotated Transformer
  • Xiong et al., On Layer Normalization in the Transformer Architecture (2020) — pre-norm vs post-norm, formally
  • Elhage et al., A Mathematical Framework for Transformer Circuits (2021) — transformer-circuits.pub — heads-as-programs, for the curious
  • Press & Wolf, Using the Output Embedding to Improve Language Models (2017) — weight tying