Scaled dot-product attention with a mask
Medium15 min11 hidden tests
A search model reads a query and a list of passages with one attention layer. Write attention(Q, K, V, mask=None) that returns scaled dot-product attention: score every query against every key, divide by sqrt(d) where d is the key width, softmax over the keys, and average the value rows with those weights. An optional boolean mask says which keys each query may use (True means allowed). A blocked key must get a weight of exactly 0. In the first example, the query matches the first key better, so the output, about 13.30, sits closer to that key's value of 10 than to 20.
Tap a cell to block or allow a key
softmax(Q Kᵀ / √d), blocked keys get 0Query 1: scores q · k / √2 = 1.414, 2.828. Softmax over the allowed keys gives weights 0.196, 0.804, and a blocked key gets exactly 0. Output = weights @ V = (0.1956, 0.8044).
Query 1 output 0.196, 0.804.
Examples
Weights about 0.67 and 0.33
attention([[1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]], [[10.0], [20.0]])
returns [[13.3024]]
A causal mask: the first query sees only the first key
attention([[1.0, 1.0], [2.0, 0.0]], [[1.0, 1.0], [2.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]], [[True, False], [True, True]])
returns [[1, 0], [0.19557, 0.80443]]
Requirements
Qhas shape (n_q, d),Kshape (n_k, d) andVshape (n_k, d_v). Lists or arrays.- Weights are
softmax(Q @ K.T / sqrt(d))along each row, over the n_k keys. Stay finite when the scores are in the thousands. maskis None or an (n_q, n_k) array or nested list of booleans. False entries get weight 0, and every row has at least one True.- Return the (n_q, d_v) float array
weights @ V.
The idea behind it
Attention lets every query pull information from the most relevant keys. Each query is compared with each key by a dot product. The scores go through a softmax, so each query's weights are positive and sum to 1, and the output is the weighted average of the value rows: softmax(Q K^T / sqrt(d)) V. The division by sqrt(d) matters. With d independent unit-variance entries, a dot product has variance d, so wide heads would produce huge scores and a softmax that puts almost all weight on one key, with gradients near 0. Dividing by sqrt(d) brings the variance back to 1. Masks handle causality and padding. Setting a blocked score to -inf before the softmax gives it a weight of exactly 0 while the rest still sum to 1. This is what torch.nn.functional.scaled_dot_product_attention computes.
Tested in interview case Fix the attention layer, then add a KV cache.