Hard Primitives

Implement Scaled Dot-Product Attention

Attention builds an output for each query by taking a weighted average of value vectors. A query's dot products with the keys determine those weights. This exercise implements one attention head without batching, masking or dropout.

Scores, weights and values

Compute these three arrays in order:

scores = (Q @ K.T) / sqrt(d_k)
weights = softmax(scores, axis=-1)
output = weights @ V

Each score row belongs to one query; each column belongs to one key. Softmax normalizes across the keys, so each row's weights sum to one. The matching rows of V supply the vectors to average. Query and key vectors have width d_k, but the values can have a different width d_v.

The scale uses the query/key width, not the number of tokens or value width. For independent zero-mean, unit-variance components, a dot product's variance grows with d_k; dividing by its square root keeps that scale from growing in the same way. It is not a guarantee against numerical overflow.

For Q = [[1, 0]], identity rows in K, and V = [[10, 0], [0, 10]], the scores are approximately [0.707107, 0]. The weights are about [0.669762, 0.330238], giving output [[6.697615, 3.302385]].

Your task

Implement scaled_dot_product_attention(Q, K, V) using matrix products and a row-wise softmax. Do not call a built-in attention function. For a stable softmax, subtract each row's maximum before exponentiating.

Q: a floating-point tensor of shape (seq_len_q, d_k).
K: a floating-point tensor of shape (seq_len_k, d_k).
V: a floating-point tensor of shape (seq_len_k, d_v).

Dimensions are positive. The number of queries need not equal the number of keys. Inputs use compatible dtypes and, in PyTorch, the same device.

Output

Return a tensor of shape (seq_len_q, d_v), leaving the inputs unchanged. Return the weighted values, not the score matrix or attention weights.

The grid below shows which query/key pairs a mask would allow, not their learned weights. Full attention corresponds to this task; the other patterns are comparisons for masked variants.

Loading visualization…

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Attention combines values using learned weights

Attention lets each output position combine information from source positions. It uses three sets of vectors:

A query for each output position.
A key for each source position, compared with the query to get a score.
A value for each source position, containing the information to combine.

Scaled dot-product attention has three steps: compute query-key scores, turn those scores into weights with softmax, then take a weighted sum of the values.

$$A=\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}+M\right),\qquad O=AV.$$

Here, d_k is the width of each query and key vector. The mask M is zero for allowed positions and negative infinity for blocked positions.

For a single head, Q has shape (query_length, key_width), K has shape (source_length, key_width), and V has shape (source_length, value_width). The weights A have shape (query_length, source_length). Softmax runs over source positions, so each query gets its own set of weights that sum to one. The output has shape (query_length, value_width).

For weights [0.25, 0.75] and scalar values [2, 6], the output is 0.25 * 2 + 0.75 * 6 = 5. Both values contribute; attention does not simply choose the value with the largest score.

Mask before softmax

A causal mask blocks future positions. A padding mask blocks positions added to make sequences the same length. Setting a blocked score to negative infinity makes its exponential zero, so it receives zero weight. Setting it to zero does not work: exp(0) = 1.

Each row needs at least one allowed, finite score. Applying ordinary softmax to an all-negative-infinity row produces undefined values. If you instead zero weights after softmax, the remaining weights no longer sum to one unless you normalise them again.

Why divide by the square root of the key width?

If query and key components are independent, with mean zero and variance one, the variance of their dot product is d_k. Dividing by sqrt(d_k) brings it back to one. This helps keep softmax from concentrating nearly all the weight on one source position, where its derivatives become small. Scaling helps under these assumptions; it does not prevent extreme weights for every possible input.

In self-attention, queries, keys, and values come from the same sequence. In cross-attention, queries come from one sequence and keys and values come from another.

Detecting runtime…