Multi-Head Attention: splitting into independent heads
The problem, from first principles
[01-scaled-dot-product-attention] computes one single attention pattern per call. But a single pattern can only capture one type of relationship per layer (e.g. "attend to the previous token"). Multi-Head Attention (MHA) splits a model-dimension input into n_heads independent attention heads, runs scaled dot-product attention in each head in parallel, then concatenates the head outputs and projects them back to the model dimension — so each head can learn a different attention pattern cheaply.
From theory to code
Implement multi_head_attention(X, W_Q, W_K, W_V, W_O, n_heads, mask=None):
Q_h = X @ W_Q[h], K_h = X @ W_K[h], V_h = X @ W_V[h] for h = 1..n_heads
head_h = softmax(Q_h K_h^T / sqrt(d_head)) @ V_h
output = concat(head_1, ..., head_H) @ W_O
Constraints
d_modelmust be evenly divisible byn_heads;d_head = d_model // n_heads.- Each head gets its own Q/K/V weight slice, all packed into single
(d_model, d_model)matricesW_Q,W_K,W_V. - Apply scaled dot-product attention (scale by
sqrt(d_head)) independently per head. - Concatenate heads in head order, then apply the output projection
W_Oof shape(d_model, d_model).
Hints
Hint 1: Splitting into heads without a loop
Reshape (seq_len, d_model) into (seq_len, n_heads, d_head) by splitting the last axis, then transpose to (n_heads, seq_len, d_head). This is exactly equivalent to slicing n_heads separate weight matrices out of one big projection, and lets you do the whole thing as one batched matmul instead of a Python loop over heads.
Hint 2: Reassembling after attention
After computing (n_heads, seq_len, d_head) head outputs, transpose back to (seq_len, n_heads, d_head) and reshape to (seq_len, d_model) before the final W_O projection — this must invert exactly the split you did going in, in head order.
Click "Run Tests" to test your implementation