Scaled Dot-Product Attention

A recurrent network passes information from the first token to the last through every step in between, and loses some of it at each one. Attention lets every position look directly at every other position in a single operation. After this lesson your program computes the attention weights and the attention output for a sequence, with an optional causal mask.

A soft dictionary lookup

Think of a dictionary: you have a query, you compare it to every key, and you return the value of the best match. Attention does the same, but softly: it compares the query to all keys, turns the similarities into weights that sum to 1, and returns the weighted average of all the values.

• Query q_i: what position i is looking for.
• Key k_j: what position j offers to be matched on.
• Value v_j: what position j hands over if it is chosen.

output_i = sum over j of  weight[i][j] * v_j
weight[i][.] = softmax over j of ( q_i . k_j / sqrt(d_k) )

In matrix form, softmax(Q K^T / sqrt(d_k)) V, where Q and K have d_k columns and V has d_v columns. The output has the same number of rows as the input and d_v columns.

A worked trace

Take three tokens with Q = K = V = [[1, 1], [1, -1], [0, 2]] (so d_k = 2).

scores for token 0 :  q0.k0 = 2,  q0.k1 = 0,  q0.k2 = 2     divide by sqrt(2): 1.414, 0, 1.414
softmax            :  [0.4458, 0.1084, 0.4458]
output 0 = 0.4458*[1,1] + 0.1084*[1,-1] + 0.4458*[0,2] = [0.5542, 1.2290]

Token 0 is attracted equally to itself and to token 2 because both give a dot product of 2.

Why divide by the square root of d_k

If the components of q and k have variance 1, their dot product has variance d_k. With d_k = 64 typical scores are about +-8, and softmax of numbers that far apart puts almost all the weight on one position, so its gradient everywhere else vanishes. Dividing by sqrt(d_k) brings the scores back to variance 1 regardless of width.

Causal masking

A language model predicting the next token must not see future tokens. Before the softmax, set every score where j > i to minus infinity: e^(-inf) = 0, so those positions get exactly zero weight and the rest renormalize. Continuing the trace with the mask, token 0 can only see itself, so its weights are [1, 0, 0] and its output is its own value. Token 1 sees tokens 0 and 1: its scores are 0 and 1.414, which give the weights [0.1956, 0.8044].

The mask has to be applied before the softmax. After it, the future positions already hold probability that has leaked into the other weights. Equivalently, implement it by running the softmax over only the first i + 1 scores of row i and filling the rest with 0.

Cost

The score matrix is T by T, so attention costs time and memory proportional to the square of the sequence length. Doubling the context quadruples the work, which is why long-context models use tricks such as FlashAttention (compute in blocks, never storing the full matrix) or sparse patterns. Real transformers also run several attention heads in parallel, each with its own learned projections of the input into Q, K and V; here they are given directly.

Common mistakes

• Dividing by d_k instead of sqrt(d_k), or forgetting to scale at all.
• Applying the softmax down the columns instead of along each row.
• Multiplying by K instead of V in the last step, which only works when they happen to be equal.
• Masking after the softmax, or masking the past instead of the future.
• Forgetting the max-subtraction in the softmax with large scores.

Your exercise

Commands load Q, K and V and switch the mask. The starter parses and prints. Implement softmax, attention_weights (the T by T weights, mask included) and attend (the weighted sum of the values). Compute a score as the dot product times 1 / sqrt(d_k).