Attention is a lookup you can differentiate

Start with something familiar. A Python dictionary holds pairs (key, value). You hand it a query, it finds the key that is equal to your query, and it hands back that key’s value. Written as math, with keys k_1 … k_M and values v_1 … v_M:

lookup(q) = Σ_j  1[q = k_j] · v_j

The indicator is 1 for the matching key and 0 elsewhere, so the sum collapses to a single value. This is exactly what a neural network wants — “given this token, fetch the information associated with that one” — and exactly what gradient descent cannot use. The indicator is flat almost everywhere and undefined at the jump, so ∂lookup/∂q = 0: no learning signal reaches the query. Attention repairs that one defect.

Advertisement

Relaxing the hard match into a soft one

Two things about the hard lookup have to go. First, equality is too brittle: two vectors are almost never exactly equal, and we want partial matches to count. Replace it with a similarity score s_j = score(q, k_j), large when q and k_j point the same way. Second, picking the single best is an argmax, just as non-differentiable as equality. Replace it with softmax, the standard smooth relaxation:

α_j = exp(s_j) / Σ_m exp(s_m)     α_j > 0,  Σ_j α_j = 1
out   = Σ_j α_j · v_j

Nothing else changed: the computation is still “score every entry, pick by score, return the value.” But now every key contributes a little, every contribution is a smooth function of q, and the retrieval is trainable end to end. The hard lookup is the limit of this as the scores are scaled up without bound.

Advertisement

Why three projections and not one vector

Given a token embedding x, why not use x itself as query, key, and value? Because those are three different jobs. The query encodes what this token is looking for; the key encodes what this token offers as a handle; the value encodes what gets copied if the handle is grabbed. Collapsing them has a concrete failure mode: with Q = K = X, the score matrix X X^T is symmetric, so “A attends to B” is forced to equal “B attends to A” — yet an adjective seeking its noun is not the same relation as the noun looking back. Worse, the diagonal is usually the largest entry in its row, so every token would attend mostly to itself. Three learned linear maps fix both:

Q = X W_Q    X:   [N, d_model]    W_Q: [d_model, d_k]  → Q: [N, d_k]
K = X W_K                      W_K: [d_model, d_k]  → K: [M, d_k]
V = X W_V                      W_V: [d_model, d_v]  → V: [M, d_v]

What the dot product actually measures

The similarity function attention uses is the plain dot product, s_ij = q_i · k_j = Σ_c q_ic k_jc. Geometrically q · k = ‖q‖ ‖k‖ cosθ, so the score rises with alignment of direction and with magnitude. That second dependence is deliberate: a key can make itself globally attractive by growing its norm, and a query can sharpen or soften its own attention by scaling itself.

The choice is also computational: all query-key pairs are one matrix multiply, S = Q K^T with shape [N, M]. Additive attention, s = w^T tanh(W_q q + W_k k), is comparably expressive but cannot be written as a single GEMM, which is why the dot-product form won.

Scaling the scores before the softmax

Raw dot products grow with dimension. If the components of q and k are roughly independent with unit variance, q · k is a sum of d_k such products, so its variance is d_k and its typical magnitude is √d_k. At d_k = 64 the scores routinely reach ±8, and softmax of numbers that spread out is nearly one-hot. Dividing by √d_k restores unit variance:

α_i = softmax( q_i K^T / √d_k )

The reason this matters is not aesthetic. A saturated softmax has vanishing gradients (see the Jacobian below), so an unscaled model can stall before it ever learns which keys are worth attending to. The scale factor is a fixed constant, not a parameter — it corrects a known dimensional effect rather than learning anything.

The softmax step and what the weights mean

Softmax converts a row of scores into a probability distribution over the keys. It is shift invariant: adding any constant c to every score leaves α unchanged, because exp(s_j + c) factors exp(c) out of numerator and denominator alike. Every real implementation exploits this by subtracting the row max before exponentiating, which keeps exp from overflowing in float16. It is also scale sensitive: multiplying all scores by β is a temperature change — large β drives the row toward one-hot, small β toward uniform. Row entropy H = -Σ_j α_j log α_j is a useful diagnostic: a head sitting near log M is just mean-pooling.

The weighted sum is a convex combination

The final step, out_i = Σ_j α_ij v_j, is worth dwelling on because it is the step people gloss over. Since the weights are non-negative and sum to one, the output is a convex combination of the value vectors — it lies inside the convex hull of {v_1 … v_M}. Attention can interpolate between the information already present in the sequence, and it can select a single item exactly, but it cannot invent a direction that no value vector points in. All genuinely new features come from the feed-forward block downstream.

It also bounds the output: ‖out‖ ≤ max_j ‖v_j‖ by the triangle inequality, so attention is non-expansive and cannot blow up activations on its own.

The full formula with shapes attached

Assembling the pieces, and carrying dimensions at every line so nothing is ambiguous:

X       : [N, d_model]              input token representations
Q = XW_Q: [N, d_k]                  what each position asks for
K = XW_K: [M, d_k]                  what each position advertises
V = XW_V: [M, d_v]                  what each position hands over

S = QK^T           : [N, M]         raw similarity scores
S'= S / √d_k       : [N, M]         variance-corrected
A = softmax_row(S'): [N, M]         each row sums to 1
O = A V            : [N, d_v]       convex combos of values

N (queries) and M (keys/values) are independent. All three projections reading the same X gives self-attention; Q from a decoder with K, V from an encoder gives cross-attention. The only hard constraint is that Q and K share the width d_k; d_v is free.

A worked lookup over four keys

Take one query, four keys with d_k = 2, and four values with d_v = 3 chosen as near-one-hot tags so the retrieval is easy to read:

q  = [2, 1]
K  = [1, 0]   [0, 2]   [2, 2]   [-1, 1]
V  = [1,0,0]  [0,1,0]  [0,0,1]  [1,1,1]

Scores are [2, 2, 6, -1]. Scale by √2 ≈ 1.414 to get [1.414, 1.414, 4.243, -0.707], exponentiate to [4.113, 4.113, 69.59, 0.493], and divide by the sum 78.31:

α   = [0.053, 0.053, 0.888, 0.006]
out = 0.053[1,0,0] + 0.053[0,1,0] + 0.888[0,0,1] + 0.006[1,1,1]
    = [0.059, 0.059, 0.894]

The output is almost v_3, softened by a trace of the others — a soft lookup that has nearly, but not quite, committed. Skipping the √d_k division gives [0.018, 0.018, 0.964, 0.001]: sharper, and closer to the saturated regime where gradients die.

How gradients flow back through the weights

This is the payoff for making the lookup soft. The softmax Jacobian is

∂α_i / ∂s_j = α_i (δ_ij - α_j)

and three consequences follow. First, ∂out/∂v_j = α_j: gradient reaches a value vector in proportion to how much it was attended to, so attention routes the backward pass as well as the forward one. Second, the -α_i α_j off-diagonal term makes the weights compete — raising one key’s score lowers every other key’s weight, which is what pushes heads to specialise. Third, when a row saturates (α_i → 1 or 0), the Jacobian goes to zero and learning through that row stops. That is the mechanical reason the √d_k scaling exists.