Hand-Write a Multi-Head Attention Layer: Projections, Head Splitting and Masking

Implement a multi-head attention layer from scratch, as you would in a live coding round. The layer receives query, key and value inputs, projects each of them with learned weights, splits the projections into several heads, runs scaled dot-product attention independently in every head, then concatenates the heads and applies an output projection. It should accept an optional attention mask. ```hint Write down every shape Annotate the tensor shape after each line. Most bugs in this exercise are a reshape or a transpose that silently mixes heads with sequence positions. ``` ```hint Masks and the softmax Decide at which step the mask must be applied for the attention weights to remain a valid distribution, and check what your code returns for a query that is not allowed to attend to anything. ``` ### Constraints and Clarifications - Assume NumPy or PyTorch. Basic tensor operations (matrix multiplication, reshape, transpose, softmax, linear layers) are allowed; a library attention module such as `torch.nn.MultiheadAttention` or `torch.nn.functional.scaled_dot_product_attention` is not. - Inputs are batched: queries have shape `(batch, len_q, d_model)`, and keys and values have shape `(batch, len_k, d_model)`. - Assume `d_model` is divisible by the number of heads. ### Clarifying Questions - Is this self-attention only, or must the same layer handle cross-attention, where queries and keys have different lengths? - Which masks are needed: key padding, causal, or both? Does `True` in a boolean mask mean "may attend" or "blocked"? - Should the layer apply dropout to the attention weights, and should it return the weights along with the output? - Is a forward pass enough, or should the implementation also be checked for gradients and for numerical behavior in half precision? ### What a Strong Answer Covers - Correct projections, head splitting and head merging, with consistent shapes for self- and cross-attention - The right scaling factor, a softmax over the correct axis, and masking applied at the correct step - Numerical robustness: large scores, fully masked rows, reduced precision - Time and memory cost as functions of sequence length, model width and number of heads - How the implementation would be tested against a reference and for mask correctness ### Follow-up Questions - Why are the scores divided by $\sqrt{d_k}$, and what happens during training if that step is left out? - What are the time and memory costs in sequence length, and how would you reduce them for long sequences? - With the same `d_model`, what does multi-head attention gain over a single head, given that the number of parameters is the same? - How would you extend the layer for step-by-step autoregressive decoding without recomputing attention over the whole prefix at every step?