Self-attention

The problem fixed context cannot solve

Your MLP language model flattens the last k tokens into one vector. Token at position 99 cannot influence the representation at position 100 if it fell outside the window. Recurrent networks pass hidden state forward, but sequential computation is hard to parallelize and long-range gradients vanish.

Self-attention lets every position directly look at every other position (within a block) in parallel. The model learns which past tokens matter for predicting the next one — not a fixed window chosen by you.

This is the core inductive bias of transformers. You will implement it from equations, not only from diagrams, so when you read GPT source code or debug shape errors, you recognize Q, K, and V immediately.

Callout — attention is weighted retrieval: Each position queries the sequence: "what information here is relevant to me?" Keys say what each position offers; values are what get retrieved.

Scaled dot-product attention

Given input sequence X with shape (batch, seq_len, d_model), project into three matrices:

Q = X @ W_q    # queries:  (batch, seq_len, d_k)
K = X @ W_k    # keys:     (batch, seq_len, d_k)
V = X @ W_v    # values:   (batch, seq_len, d_v)

Attention scores combine queries and keys:

scores = Q @ K^T / sqrt(d_k)     # (batch, seq_len, seq_len)
weights = softmax(scores, dim=-1)
output = weights @ V              # (batch, seq_len, d_v)

Each row of weights is a probability distribution over positions — how much position i attends to every other position.

Scaling by √d_k prevents dot products from growing large when d_k is big, which would push softmax into extreme values (near one-hot) and shrink gradients.

Causal (decoder) mask for language modeling

GPT-style models must not peek at future tokens during training. Apply a causal mask before softmax: set scores for positions j > i to −∞ so weights for future tokens are zero.

# mask shape (seq_len, seq_len), lower triangular allowed
scores = scores.masked_fill(mask == 0, float('-inf'))

Without this mask, the model cheats by reading the answer during next-token prediction.

Shape checklist — verify on paper first

For batch B, sequence T, model dimension D, head dimension d_k (often D / num_heads):

Shape bugs are the #1 failure mode when implementing transformers. Write shapes in comments on every line until the block passes a unit test.

Multi-head attention

Multi-head attention runs h parallel attention operations with smaller d_k each, then concatenates:

head_i = Attention(X @ W_q_i, X @ W_k_i, X @ W_v_i)
MultiHead = Concat(head_1, ..., head_h) @ W_o

Why multiple heads? Different heads can specialize — one tracks local syntax, another long-range coreference, another punctuation patterns. You will not see clean human labels for heads in a tiny model, but the capacity helps even at small scale.

Implementation tip: reshape to (B, h, T, d_k) so attention is batched over heads in one matmul.

Residual connection and layer norm

A transformer block wraps attention (and later FFN) with:

x = x + Attention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

Residual connections let gradients flow around sublayers — critical for deep stacks.

Layer normalization stabilizes activations. GPT-2 uses pre-norm (norm before sublayer); original transformer used post-norm. Match whichever reference implementation you follow for lesson 3.5.

You can implement attention alone in this lesson; the full block lands in the tiny GPT assembly.

Intuition: what do weights look like?

On a short sentence like "the cat sat on the mat", a causal LM predicting token after "sat" might attend strongly to "cat" (subject) and "on" (preposition heading next phrase). Visualize the attention matrix as a heatmap: rows = query position, columns = key position.

In tiny models trained on small data, weights are often noisy — do not expect GPT-4-quality interpretability. The exercise is verifying your implementation produces sensible softmax rows (non-negative, sum to 1, upper triangle zero when masked).

Tools: matplotlib imshow, or seaborn heatmap on a (T, T) slice for one head.

Callout — attention weights ≠ explanation: Pretty heatmaps are suggestive, not proof of reasoning. Use visualization to debug shapes and sanity-check training, not to claim the model "understands" grammar.

Complexity and why it still wins

Naive attention is O(T²) in sequence length for the score matrix. Long contexts get expensive — this is why FlashAttention, sparse patterns, and sliding windows exist in production systems.

For your tiny GPT with T ≤ 128 on CPU, naive implementation is fine. Optimize when you feel the pain, not before.

Engineering problem (staff framing)

Need parallel long-range dependency modeling. Self-attention is weighted retrieval over the sequence.

Diagram — Scaled dot-product attention

flowchart LR
  X[X] --> Q[Q=XW_q]
  X --> K[K=XW_k]
  X --> V[V=XW_v]
  Q --> S["scores = QK^T / sqrt(d_k)"]
  K --> S
  S --> W[softmax]
  W --> O["out = W V"]
  V --> O

Precise definitions & mental model

• Q queries, K keys, V values; scale by √d_k for stable softmax.
• Multi-head — multiple subspaces.
• Causal mask — block future positions for LMs.
• Complexity O(n²) in sequence length for dense attention.

Tradeoffs — when to use what

Failure modes (interview + on-call)

Wrong mask → cheat; d_k scale omitted → soft saturates; shape bugs on heads.

Production & OSS practices

KV-cache grows with heads×layers×seq; GQA is a serving optimization born here.

Interview cue card

Derive shapes for multi-head attention given B,T,H,D.

Deep dive (FAANG / OSS bar)

Numerical stability and the √d_k term

Dot products of d_k-dimensional vectors have variance that grows with d_k if components are roughly unit scale. Softmax then saturates to one-hots, gradients die. Scaling by √d_k keeps logits in a healthy range — implement it; do not treat it as optional ornament.

Causal masking as a software contract

For decoder-only LMs, illegal attention to future positions is a correctness bug, not a quality issue. In code:

• Build mask shape (T, T) with -inf above diagonal (or boolean mask consumed by the kernel).
• Test with a unit that asserts position i cannot change when future tokens flip.

KV-cache (why serving engineers care)

At decode step t, keys/values for tokens 0…t-1 are reused. Memory ≈ layers × kv_heads × t × d_k × dtype_bytes × batch. This is why:

• Long context is expensive even when FLOPs per new token look fine.
• GQA/MQA reduce kv_heads.
• Continuous batching packs sequences with different t.

flowchart LR
  Prefill[Prefill: compute K/V for prompt] --> Cache[(KV cache)]
  Cache --> Decode[Decode: attend new Q to cache]
  Decode --> Cache

Multi-head intuition without mysticism

Heads are not guaranteed to be "syntax" vs "coreference." Empirically they specialize sometimes; product engineering should not depend on interpreting a single head. Use attention maps as debugging aids, not legal explanations.

Failure table

Micro-project: Implement attention block

In m3/attention/:

1. Implement scaled_dot_product_attention(Q, K, V, mask=None) returning output and weights.
2. Implement MultiHeadAttention with configurable d_model, num_heads, causal mask for LM use.
3. Unit test shapes: random (B=2, T=8, D=64) input → output same shape as input (after output projection).
4. Forward a real token sequence (10–20 tokens) through your block; save attention_heatmap.png for one head at the final position.
5. Write SHAPES.md: table of every tensor shape in forward pass; note where you had bugs.

Do not depend on nn.MultiheadAttention for the core exercise — hand-roll at least the scaled dot-product step. You may compare against PyTorch's implementation for numerical sanity.

Optional: verify causal property — weights for future columns are exactly zero.

Checklist

• Scaled dot-product attention passes shape tests
• Causal mask prevents attending to future tokens
• Attention heatmap saved for a short sequence
• SHAPES.md documents the full forward pass