Multi-Head Attention: concatenating heads plus output projection
The problem, from first principles
[04-mha-split-heads] split a single d_model-wide representation into num_heads separate, smaller heads, computed each head's own independent attention pattern, and left the results as num_heads SEPARATE d_k-wide outputs, shape (batch_size, num_heads, seq_len, d_k). Before this can be used by anything downstream ([05-transformer-block]'s residual connections, later in this curriculum, need a d_model-wide vector to add back to the original input), every head's output needs to be recombined back into ONE full-width representation. Simply gluing the heads' outputs side by side (concatenation) is the obvious first step, but the original Transformer paper's design adds one more crucial piece: a LEARNED linear projection applied to that concatenated result, before it's considered "done."
This output projection isn't a rounding-error detail, it's what gives the model a way to learn how to COMBINE different heads' contributions intelligently (some heads might deserve more weight than others for a given context, or their outputs might need to be blended in non-trivial, LEARNED ways), rather than forcing a fixed, naive side-by-side stacking to be the final answer regardless of what the training data actually calls for.
From theory to code
Implement concat_heads(x), the exact inverse of [04-mha-split-heads]'s split_heads: reshape (batch_size, num_heads, seq_len, d_k) back into (batch_size, seq_len, d_model). Implement multi_head_attention(query, key, value, num_heads, weight_o, bias_o, mask), calling [04-mha-split-heads]'s multi_head_attention_per_head (already provided) to get the per-head outputs, concatenating them via concat_heads, then applying a [03-dl-training/02-layers/01-linear-forward]-style output projection (weight_o, bias_o) to the concatenated result.
Constraints
concat_headsmust be the exact mathematical inverse ofsplit_heads:concat_heads(split_heads(x, num_heads)) == xfor any validx.- The head dimension must move BACK to sit next to
d_kbefore flattening (mirroring[04-mha-split-heads]'s move in the opposite direction), so each head's output lands in the correct CONSECUTIVE chunk of the finald_model-wide vector. multi_head_attentionmust apply the output projection AFTER concatenation, not before, and not to each head separately.weight_ohas shape(d_model, d_model), matching[03-dl-training/02-layers/01-linear-forward]'s(out_features, in_features)convention.
Hints
Hint 1: concat_heads's two-step reshape
First, x.transpose(0, 2, 1, 3) swaps the num_heads and seq_len axes back (undoing split_heads's own transpose), giving (batch_size, seq_len, num_heads, d_k). Then, .reshape(batch_size, seq_len, num_heads * d_k) flattens the last two axes together into one d_model-wide dimension.
Hint 2: Sanity-checking against split_heads
If you're unsure concat_heads is correct, check it directly: concat_heads(split_heads(x, num_heads)) should produce something numerically identical to the original x, for any x and any valid num_heads.
Hint 3: multi_head_attention
per_head_output, weights = multi_head_attention_per_head(query, key, value, num_heads, mask=mask), then concatenated = concat_heads(per_head_output), then output = concatenated @ weight_o.T + bias_o, exactly [03-dl-training/02-layers/01-linear-forward]'s linear_forward formula, applied to the concatenated attention output.
Click "Run Tests" to test your implementation