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.
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.