Scaled Dot-Product Attention: the core operation every transformer runs
The problem, from first principles
Every transformer block, in every model, runs the same core operation to decide "which other tokens should this token pay attention to": score a query against a set of keys, turn those scores into a probability distribution, and use that distribution to mix the corresponding values. Given query, key and value matrices for a single attention head, compute this scaled dot-product attention output.
From theory to code
Implement scaled_dot_product_attention(Q, K, V, mask=None):
scores = Q @ K.T
scaled = scores / sqrt(d_k)
scaled[mask == 0] = -inf # only if mask is given
weights = softmax(scaled, axis=-1)
output = weights @ V
Constraints
- Scale scores by
1/sqrt(d_k)whered_kis the last dimension ofQandK. - Apply an additive mask of
-infto positions wheremask == 0, before the softmax, never after. - Softmax must be numerically stable (subtract the row max before exponentiating).
- Return both the attention output and the attention weight matrix.
Hints
Hint 1: Why the scale factor
As the head dimension d_k grows, q . k grows roughly like d_k in magnitude for unit-variance inputs. Large-magnitude scores push softmax into a near one-hot regime, which is brittle. Dividing by sqrt(d_k) keeps the pre-softmax variance close to 1 regardless of head size.
Hint 2: Masking order matters
Do the masking on the SCALED logits, before the softmax — never after. Subtract the row-wise max before exponentiating for numerical stability; it doesn't change the softmax result, since softmax is shift-invariant.
Click "Run Tests" to test your implementation