Attention Math — Query, Key, Value and Weighted Sum

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.

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.

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.

Masking: editing scores, never weights

A causal decoder must not let position i read position j > i. The implementation is to add a mask to the scores before the softmax, using -∞ (in practice a large negative constant) for forbidden pairs:

S'_ij = S_ij / √d_k + M_ij,    M_ij = 0 if j ≤ i else -∞

Because exp(-∞) = 0, masked entries get exactly zero weight and the survivors renormalise, so each row still sums to one. That renormalisation is the whole reason the mask goes before the softmax: zeroing weights afterwards leaves a row summing to less than one, quietly shrinking the output and breaking the convex-combination property. Padding masks, sliding windows, and block-sparse patterns are all just entries of M.

What the shapes cost, especially on a CPU

Two matrix multiplies dominate. QK^T costs O(N · M · d_k) and AV costs O(N · M · d_v), so self-attention is O(N^2 d) in time and, if the score matrix is materialised, O(N^2) in memory. Note that the projections themselves cost O(N d^2), which actually dominates while N < d — the quadratic term only dominates at longer contexts.

On a CPU the binding constraint is memory traffic, not arithmetic. During single-token decoding N = 1, so both operations degrade to matrix-vector products with almost no arithmetic intensity: the machine spends its time streaming the cached K and V tensors out of RAM. That is why KV-cache quantisation and grouped-query attention buy far more on a small CPU model than any reduction in FLOPs.

Misreadings worth avoiding

Four recur. “Attention weights are explanations.” They show where information was mixed from, not why a prediction happened. “Q, K, V are different data.” In self-attention they are three linear views of the same tokens; the asymmetry lives entirely in the learned matrices. “Softmax normalises over queries.” It normalises along the key axis, within each row. And “attention is where features are built.” Attention only mixes existing values — it is the strictly linear-in-V part of the block, and every new nonlinear feature comes from the MLP that follows.