Asked about Causal multi-head attention in NumPy
Implement causal self-attention in NumPy.
An answer that works
Project the input three times to get q, k and v of shape (T, d_k). Compute scores as q @ k.T divided by the square root of d_k, giving (T, T). Set the strictly upper triangle to negative infinity, softmax each row, then multiply by v to get (T, d_k). For multiple heads, reshape the feature axis into (n_heads, d_head), transpose heads in front of positions, do the same computation batched, transpose back and reshape, then apply the output projection.
Then they push further
Why negative infinity rather than zeroing the weights afterwards?
Because the softmax has to normalise over the visible positions only. Zeroing after the softmax leaves each row summing to less than one, and the deficit is worst at early positions, which quietly shrinks their outputs. Negative infinity makes the exponential exactly zero before normalisation, so the surviving weights still sum to one.
Do you scale by the square root of d_model or of d_head?
d_head. The dot product happens inside a head, between vectors of length d_head, so that is the dimension whose variance needs correcting. Scaling by d_model would under-scale by the number of heads.
How would you convince me it is correct without a reference implementation?
Two assertions. Each row of the attention weights sums to one, which catches a softmax on the wrong axis or a mask applied too late. The strict upper triangle of the weights is exactly zero, which catches a mask built with the wrong offset or transposed. A third is that truncating the sequence leaves the earlier outputs unchanged, which is causality stated as a property rather than as a shape.
What gets marked up
Strong candidates state the shape of each intermediate as they write it and test their own code unprompted. Getting the mask right before the softmax, and scaling by d_head, separate people who have written this from people who have read it.
Reading an answer is not being able to give one. The section that teaches this is Causal multi-head attention in NumPy, 6 hours of Transformers & attention.
Learn this properlyAlso asked about Transformers & attention