Katkat
Back

Implement Multi-Head Attention from Scratch

ByteDanceApplied Scientist · Technical Screen
MediumApplied ScientistTechnical ScreenMachine Learning

Last updated

An ML coding question asking you to implement multi-head attention from scratch in NumPy or PyTorch, from the query, key and value projections to the output projection. It tests tensor shape handling, scaled dot-product attention, masking, numerical stability and complexity analysis.

Implement multi-head attention, the core layer of the Transformer, from scratch. Your implementation takes query, key and value inputs and returns the attention output for every query position.

  • query has shape (batch, L_q, d_model); key and value have shape (batch, L_k, d_model). For self-attention, all three are the same tensor.
  • The layer has num_heads heads. It projects the queries, keys and values with learned matrices, runs attention in every head in parallel, concatenates the heads, and applies a learned output projection.
  • The output has shape (batch, L_q, d_model).
  • An optional boolean mask marks which key positions each query may attend to.

Use NumPy or PyTorch tensor operations. Assume that a framework's ready-made attention layer or fused attention function is off limits, since the point is to write the layer yourself.

Constraints and Clarifications

  • d_model is divisible by num_heads; each head works in dimension d_k = d_model / num_heads.
  • The implementation is batched, with no Python loop over heads or positions.
  • Masked key positions must receive zero attention weight.

Clarifying Questions

  • A NumPy forward pass with explicit weight arguments, or a trainable PyTorch module?
  • Self-attention only, or separate query and key/value inputs as in cross-attention?
  • What is the mask convention: does True mean "may attend" or "blocked"? Is a causal mask needed?
  • Should the layer also return the attention weights, and should it include dropout and bias terms?

Hints

Hint 1

Shapes first: Write down the shape of every tensor after every line. Most bugs in this task are a reshape or transpose that silently mixes heads with positions.

Hint 2

Keep the softmax healthy: Consider how the size of the dot products changes as the per-head dimension grows, and what that does to the softmax and its gradients.

Follow-up questions

  • Add a causal mask for decoder self-attention. What changes, and what happens to a query row whose keys are all masked?
  • During autoregressive generation, how would you change the layer to process one new token at a time without recomputing the keys and values of earlier tokens?
  • Why use several heads of size `d_model / num_heads` instead of one head of size `d_model`? What do multi-query and grouped-query attention change?
  • The sequence length grows to 32,000 tokens. Where does the layer run out of memory, and what can you do about it?