Multi-Head Self-Attention

Attention is the function that lets a token's representation pull information from other tokens in the same sequence. Self-attention means queries, keys, and values all come from the same input. Multi-head means the projection is split into H parallel attention problems, whose outputs are concatenated and projected back.

The efficient pattern is one linear layer that projects from D to 3 * D. You slice the result into three views and reshape each into H heads of size D // H. The matmul, softmax, and weighted sum then run as batched tensor operations, so the heads run in parallel on the accelerator.

This lesson builds that block and adds the causal mask, so the same code works as the attention layer in a decoder-only language model. The next lesson stacks the block into a full transformer, and the lesson after that trains it.

The block takes (B, T, D) and returns (B, T, D). The mask is (T, T) or broadcastable to it. Inside, the intermediate tensors are (B, H, T, d_head) with d_head = D // H, which requires D % H == 0.

All the parameters live in two linear layers: the QKV projection and the output projection. The mask, the softmax, the matmuls, and the reshapes have no parameters.

The naive implementation uses three separate linear layers, one each for Q, K, and V. The efficient one uses a single layer that outputs 3 * D features and splits the result. The two are mathematically equivalent: three matrix multiplications by (D, D) weights are one multiplication by a (3D, D) weight stacked from them.

The single layer is faster because the accelerator launches one matmul instead of three. It is also easier to initialize, since the three sub-matrices live in one parameter tensor and get initialized together.

After the split, Q, K, and V are each (B, T, D). To get H parallel attention problems, reshape to (B, T, H, d_head) and transpose to (B, H, T, d_head). That puts the head dimension next to the batch dimension, so PyTorch treats per-head attention as a batched operation over B * H independent instances.

Keep d_head last. The score matmul Q @ K.transpose(-2, -1) contracts it and returns (B, H, T, T) per-head attention scores.

Scores are divided by sqrt(d_head) before the softmax. Without that, dot products grow with d_head and push the softmax into a regime where one entry holds almost all the mass and the rest are vanishingly small. Gradients there are tiny and learning stalls. Dividing by sqrt(d_head) keeps the variance of the scores roughly constant across head sizes.

The causal mask blocks each position from attending to the positions after it. A decoder-only language model can only condition on the past when predicting the next token, and the mask enforces that. Before the softmax, every entry above the diagonal of the (T, T) score matrix is replaced by negative infinity, and after the softmax those positions get weight zero.

The mask is registered as a buffer at construction, so it lives on the same device as the model and stays out of the gradient graph. It covers the maximum context length the block will see, and at forward time we slice the top-left (T, T) corner.

After attention you have per-head context vectors of shape (B, H, T, d_head). Transpose back to (B, T, H, d_head), reshape to (B, T, D), and apply a final (D, D) linear projection. That projection is where the model mixes the heads. Without it, the H heads would only recombine in later layers, which constrains the block artificially.

Pass return_weights=True to the forward pass and the block returns the per-head attention weights, shape (B, H, T, T), alongside the output. The demo prints a heatmap of one head's weights on a short input (one random token id repeated across the context), so you can see the causal-triangle structure and the per-position focus.

In a trained model, different heads learn different patterns. Some attend to the immediately previous token, some to the start of the sequence, and some spread attention almost uniformly. This hook is the entry point for that interpretability work.

The demo at the bottom of main.py wires the attention block to a tiny LM head and trains the whole thing on a repeat task. Each input row is a single random id replicated across the context, and the target is the input shifted by one, so the model has to learn that the next token equals the previous one. The loss is cross-entropy. With H=4, D=32, T=12, and a vocabulary of 64, the loss starts near random (around log(64) ~ 4.16) and the demo prints the average loss for each of three epochs on CPU. The script asserts only that the final epoch's average loss is below the first epoch's.

The demo isn't meant to produce a useful model. It confirms that gradients flow through every piece of the block and that the heads learn something on a problem where the answer is obvious.

No feed-forward block. A real transformer layer is attention followed by a two-layer MLP, with a residual connection and layer norm around each. The next lesson adds those.

No rotary or AliBi positional encoding. Both apply at the QKV projection step of this same block, but they are a separate teaching unit. Transforming Q and K before the matmul makes the block compatible with either.

No KV cache for inference. Caching keys and values across forward passes is the optimization that makes autoregressive decoding fast. It changes the shape contract on K and V but not on Q, and it belongs in the inference lesson.

main.py defines MultiHeadSelfAttention. The class holds two linear layers and a registered mask buffer. Its forward pass projects, reshapes, scores, masks, applies the softmax, weights the values, reshapes back, and projects again. The demo at the bottom builds a small model that wraps the attention with token and positional embeddings and an LM head, trains it on a copy task for three epochs, and prints the loss curve and a per-head attention heatmap. The tests in code/tests/test_attention.py pin the shape contract, the causality property, the softmax property, the head-split property, and gradient flow.

Run the demo. Then increase n_heads from 4 to 8 (keeping d_model=32, so d_head=4) and watch the heatmap change.