Katkat
Back

Implement the Forward Pass of Multi-Head Attention From Scratch

MetaMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

An ML coding question that asks for the forward pass of multi-head attention from scratch. It tests query, key and value projections, splitting and merging heads, scaled dot-product attention with a stable softmax, causal and padding masks, complexity in sequence length, and self-made correctness checks.

Implement the forward pass of a multi-head attention layer from scratch. The interviewer provides no test cases, so correctness has to be argued from the shapes and from small checks you write yourself.

Unless told otherwise, assume self-attention over an input x of shape (batch, seq_len, d_model), a head count num_heads that divides d_model, learned projection matrices for queries, keys, values and the output, and an optional mask. Only the forward pass is required: no training and no backward pass.

Clarifying Questions

  • Should the implementation use NumPy, or a deep-learning framework's tensor operations without its built-in attention module?
  • Is this self-attention, or cross-attention in which keys and values come from a different sequence?
  • What mask format is expected: a causal mask, a padding mask or a general boolean mask, and does True mean "may attend" or "blocked"?
  • Do the projections include bias terms, and should dropout on the attention weights be modeled?
  • Should the function also return the attention weights?

Hints

Hint 1

Track the shapes: Write the shape of every intermediate tensor next to the line that produces it. Most bugs in this exercise are a reshape or transpose that produces the right size but the wrong layout.

Hint 2

Guard the softmax: Consider what happens to the softmax when scores are large, and when every position in a row is masked out.

Follow-up questions

  • How would you add a key-value cache for autoregressive decoding, and how does the cost per generated token change?
  • How do grouped-query and multi-query attention change the shapes and the size of the cache?
  • Why is the score tensor a memory bottleneck for long sequences, and how do fused, tiled attention kernels avoid storing it?
  • Where do layer normalization and the residual connection go around this layer, and why does their placement matter?