Multi-Head Attention: splitting into heads, per-head attention
The problem, from first principles
[01-scaled-dot-product-attention]'s single attention computation produces exactly ONE set of attention weights per query, meaning it can only really capture ONE notion of "relevance" at a time. But natural language relationships are genuinely MULTI-FACETED: a word might need to attend to a nearby word for SYNTACTIC reasons (subject-verb agreement) and to a distant word for SEMANTIC reasons (what does "it" refer to) simultaneously, and a single shared attention pattern has to somehow blend both needs into one compromise. Multi-Head Attention's fix: instead of computing ONE attention pattern over the FULL d_model-dimensional vectors, split each of query/key/value into several smaller, INDEPENDENT "heads," each with its own SLICE of the full dimensionality, and let each head learn to specialize in capturing a DIFFERENT kind of relationship, entirely in parallel.
This question implements exactly the SPLITTING and PER-HEAD ATTENTION half of Multi-Head Attention; Multi-Head Attention: concatenating heads plus output projection, immediately following this question, implements the second half, recombining every head's separate output back into one unified representation.
From theory to code
Implement split_heads(x, num_heads), reshaping x (shape (batch_size, seq_len, d_model)) into (batch_size, num_heads, seq_len, d_k), where d_k = d_model // num_heads. Implement multi_head_attention_per_head(query, key, value, num_heads, mask), splitting all three inputs via split_heads, then calling [01-scaled-dot-product-attention]'s scaled_dot_product_attention ONCE on the split, head-batched inputs, letting it process every head in parallel since it's already shape-agnostic to leading batch dimensions.
Constraints
split_headsmust reshape so thatd_modelsplits intonum_headsCONSECUTIVE, non-overlapping chunks of sized_keach (the firstd_kdimensions become head 0, the nextd_kbecome head 1, and so on), then move the head dimension to position1(right after batch).d_modelmust be evenly divisible bynum_heads(assume this holds; no remainder-handling needed).multi_head_attention_per_headmust callscaled_dot_product_attentiona SINGLE time on the already-split,(batch_size, num_heads, seq_len, d_k)-shaped inputs, not loop over heads manually calling it once per head.- The result's shape must be
(batch_size, num_heads, seq_len, d_k), the head dimension still separate (not yet recombined, that's[05-mha-concat-output-projection]'s job).
Hints
Hint 1: split_heads's two-step reshape
First, x.reshape(batch_size, seq_len, num_heads, d_k) splits the LAST dimension (d_model) into two: num_heads and d_k, without changing the total number of elements or the order of any existing dimension. Then, .transpose(0, 2, 1, 3) swaps the seq_len and num_heads axes, moving the head dimension to sit right after batch: (batch_size, num_heads, seq_len, d_k).
Hint 2: Why this specific reshape, not a different one
Reshaping (batch, seq_len, d_model) directly to (batch, seq_len, num_heads, d_k) (rather than some other splitting order) preserves the natural CONSECUTIVE-chunk interpretation: dimensions 0 through d_k-1 of the original d_model axis become head 0's slice, dimensions d_k through 2*d_k-1 become head 1's slice, and so on, matching how the learned projection weights that PRODUCE query/key/value in a real Transformer are structured.
Hint 3: multi_head_attention_per_head
query_heads = split_heads(query, num_heads) (and the same for key/value), then return scaled_dot_product_attention(query_heads, key_heads, value_heads, mask=mask) directly, letting that function's existing shape-agnostic implementation handle the (batch_size, num_heads, ...) leading dimensions in one call.
Click "Run Tests" to test your implementation