Katkat
Back

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

ByteDanceSoftware Engineer · Onsite
MediumSoftware EngineerOnsiteMachine Learning

Last updated

Write a multi-head attention layer from scratch: project queries, keys and values, split them into heads, apply scaled dot-product attention with optional padding or causal masks, and merge the heads. It tests tensor shape discipline, numerical stability, and the cost of attention as sequence length grows.

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.

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?

Hints

Hint 1

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 2

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.

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?