ML coding · Deep learning layers
Scaled dot-product attention with a causal mask
Given queries Q (n × d), keys K (m × d) and values
V (m × dv), return softmax(QKᵀ/√d)V as an n × dv matrix, with the
softmax over each row. If causal is true (then n = m), query i may only attend to
keys j ≤ i.
Examples
Input Q = [[1, 0]], K = [[1, 0], [0, 1]], V = [[10, 0], [0, 10]], causal = False
Output [[6.6976, 3.3024]]
The query is closer to the first key, so the output leans towards the first value.
Input Q = [[1, 0], [0, 1]], K = [[1, 0], [0, 1]], V = [[1, 2], [3, 4]], causal = True
Output [[1, 2], [2.3395, 3.3395]]
Causal: the first query sees only the first key, so its output is exactly V[0].
Input Q = [[0, 0]], K = [[1, 1], [2, 2], [3, 3]], V = [[1], [2], [3]], causal = False
Output [[2]]
A zero query scores every key equally: the output is the mean of V.
Constraints
- 1 ≤ n, m ≤ 256, 1 ≤ d, d_v ≤ 128
- If causal, n = m
- Checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
Scores are Q @ K.T, shape (n, m), divided by √d where d is the key width.
Hint 2
For the mask, set scores above the diagonal to −∞ before the softmax. np.tril builds the mask.
Hint 3
Shift each row by its maximum before exponentiating. With a causal mask every row keeps its diagonal, so no row is all −∞.
Where candidates go wrong
- Masking with 0 instead of −∞. A zero score still gets weight e⁰ = 1.
- Scaling by d instead of √d, or by the value width d_v.
- Taking the softmax over axis 0, so that each key's weights sum to 1 instead of each query's.
- Masking after the softmax and not renormalising.
The idea
S = QKᵀ/√d gives every query-key score at once. For the causal case, np.where(np.tril(mask), S, −∞) removes the future. A row-wise stable softmax turns scores into weights that sum to one, and the weights times V average the values.
The scaling keeps the variance of the scores near 1 regardless of d, so the softmax stays in the range where its gradient is useful. The masked entries contribute e−∞ = 0 exactly, so the result is identical to running the softmax on the visible keys alone.
Time O(nmd + nmd_v) · Space O(nm)
Reference solution
import numpy as np
def attention(Q, K, V, causal):
Q, K, V = (np.asarray(a, dtype=float) for a in (Q, K, V))
S = Q @ K.T / np.sqrt(Q.shape[1]) # (n, m) scaled scores
if causal:
future = ~np.tril(np.ones(S.shape, dtype=bool)) # True where j > i
S = np.where(future, -np.inf, S) # weight e^-inf = 0
S = S - S.max(axis=1, keepdims=True) # each row keeps its diagonal, so the max is finite
W = np.exp(S)
W /= W.sum(axis=1, keepdims=True)
return W @ VFollow-ups an interviewer asks
- Multi-head attention: split d into h heads, attend in each, concatenate, project.
- A padding mask for variable-length sequences. What should a fully masked row return?
- Compute the softmax in blocks with a running maximum and sum, so the n × m matrix is never stored. That is the idea behind FlashAttention.
- Add a KV cache for autoregressive decoding: what is the cost of generating one more token?
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
Since transformers took over, this has become the most asked deep learning implementation question. A complete answer has the scaling, a stable softmax and a correct mask, and the candidate can say why each is there.
The theory behind it is in the lesson Attention, normalisation and convolution from scratch.
Keep going
Layer normalisation
Normalise each row over its features, then scale by gamma and shift by beta.
Open the problem → Medium15 minBatch normalisation in training mode
Normalise each feature over the batch and update the running mean and variance the way PyTorch does: biased variance to normalise, unbiased to track..
Open the problem → Medium20 minScaled dot-product attention with a causal mask
Implement softmax(QKᵀ/√d)V with an optional causal mask in NumPy.
Open the problem → Medium20 min2-D convolution with stride and padding
Implement a 2-D convolution (cross-correlation) with zero padding and stride, first with loops, then with sliding_window_view and einsum..
Open the problem →