Katkat
Back

Masked Multi-Head Attention Forward Pass Without a Framework

MetaMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

Implement batched multi-head self-attention by hand, including projections, head split and merge, scaled dot products, a numerically stable masked softmax, and causal and padding masks, then prove it correct with tests you design yourself. Tests shape discipline and attention fundamentals.

You're on a team that ships a tiny on-device transformer. The device runtime has no ML framework, so you must write the attention layer yourself. Implement the forward pass of multi-head self-attention. Plain Python lists are fine. If you prefer NumPy, explain how each line maps to the list version.

Interface

def multi_head_attention(x, params, num_heads, mask=None, return_weights=False):
    """
    x:       (B, T, D) batch of token embeddings
    params:  Wq, Wk, Wv, Wo of shape (D, D); optional biases bq, bk, bv, bo of shape (D,)
    mask:    optional (B, T, T) booleans; mask[b][i][j] is True when query i MAY attend to key j
    returns: (B, T, D) output, plus (B, H, T, T) attention weights if requested
    """

Also provide helpers causal_mask(B, T) and padding_mask(lengths, T) that build masks in this format.

Requirements

  1. Project the inputs to queries, keys and values. Split them into num_heads heads of width D / num_heads. Raise an error if D isn't divisible.
  2. Compute scaled dot-product attention per head, then concatenate the heads and apply the output projection.
  3. The softmax must not overflow when scores are large, for example 1000.
  4. A query row whose keys are all blocked must produce zeros, not NaN.
  5. No test cases are provided. Write at least three checks of your own that would catch common bugs.

Example check you should be able to pass

stable_softmax([1000.0, 999.0, -5.0], allowed=[True, True, False])
-> [0.7310585786300049, 0.2689414213699951, 0.0]

stable_softmax([3.0, 4.0], allowed=[False, False])
-> [0.0, 0.0]

Discussion points

  • Why divide by sqrt(head_dim)?
  • Time and memory as a function of sequence length T.
  • What goes wrong if you reshape (T, D) to (H, T, hd) without transposing?

Hints

Hint 1

Write the shape of every intermediate next to the line that creates it: x (T,D) -> q (T,D) -> head slice (T,hd) -> scores (T,T) -> weights (T,T) -> context (T,hd) -> concat (T,D) -> out (T,D).

Hint 2

Head h uses columns h*hd to (h+1)*hd of q, k and v. You don't need a physical transpose if you index by slice.

Hint 3

For stability, subtract the maximum over the allowed scores only before calling exp. If no score is allowed, return all zeros.

Hint 4

Good checks: zero Wq and Wk give uniform attention; under a causal mask, changing a later token must not change earlier outputs; one head should match a hand-written single-head formula.

Follow-up questions

  • Add a key-value cache for autoregressive decoding. What is the cost per new token, and how big does the cache get?
  • How do multi-query and grouped-query attention change the shapes and the cache size?
  • Implement the tiled online-softmax trick so memory stays O(T) per query block.
  • Extend this to cross-attention, where the keys and values come from an encoder sequence of a different length.