Scaled Dot-Product Attention

The one-sentence idea

Before any matrices, hold the intuition. Attention is differentiable dictionary lookup. In an ordinary hash map you supply a key, find the one slot that matches exactly, and read its value. Attention softens every part of that: instead of one exact match you compute a similarity between your query and every key, and instead of reading one value you read a weighted blend of all values, where the weights are how well each key matched. Nothing is discrete, so gradients flow through the whole operation and the projections that produce queries, keys, and values can be learned.

In self-attention every token plays all three roles at once. Token i emits a query — ‘what information am I looking for?’ — every token emits a key — ‘what do I offer?’ — and a value — ‘here is the content to pass along if you pick me.’ The output for token i is a mixture of everyone’s values, dominated by the tokens whose keys best answered token i’s query. That single mechanism, stacked and repeated, is how a transformer moves information between positions.

Where Q, K, and V come from

The queries, keys, and values are not given — they are learned linear projections of the input. Start with the input sequence as a matrix X of shape [N, d_model]: N tokens down the rows, each a d_model-dimensional embedding. Three weight matrices, learned during training, turn that one matrix into three:

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

Row i of Q is the query vector for token i, and likewise for K and V. The projections matter: W_Q and W_K learn what to compare (they carve out the subspace in which similarity is measured), while W_V learns what to carry. Because keys and queries must be dotted together they share the width d_k; values live in their own width d_v, which is often equal to d_k but need not be. In a single-head layer you will commonly see d_k = d_v = d_model.

The score matrix: Q times K-transpose

The first real computation is comparing every query with every key. The natural similarity for vectors is the dot product, and computing all of them at once is exactly a matrix multiply:

scores = Q K^T          [N, d_k] × [d_k, N]  =  [N, N]
scores[i, j] = Σ_r Q[i, r] · K[j, r]  =  ⟨Q[i], K[j]⟩

The result is an N × N matrix. Entry scores[i, j] is the raw affinity of token i’s query for token j’s key — large and positive when the two vectors point the same way, near zero when they are orthogonal, negative when they oppose. Reading across row i gives token i’s affinity for every token in the sequence, itself included. This matrix is the object everything else acts on, and it is also the source of attention’s cost: it has N^2 entries, one for every ordered pair of positions. For a 2048-token context that is over four million scores per attention layer, per head — the quadratic term that dominates long-context compute.

Why divide by sqrt(d_k): the variance argument

The ‘scaled’ in the name is the division by √d_k, and it is not cosmetic. Model the entries of a query and a key as independent random numbers with mean 0 and variance 1 — a reasonable approximation just after initialization. Their dot product is a sum of d_k independent product terms:

s = Σ_{r=1..d_k} q_r k_r
E[q_r k_r] = 0
Var(q_r k_r) = 1        (for unit-variance, independent q_r, k_r)
Var(s) = Σ Var(q_r k_r) = d_k      →   std(s) = √d_k

So the raw scores grow with the width: their standard deviation is √d_k. For d_k = 64 a typical logit swings by ±8, and some pairs land far larger. Feed logits that big into a softmax and it saturates — almost all the weight collapses onto a single key and the distribution becomes nearly one-hot. In that regime the softmax’s gradient is tiny (the function is flat where it is saturated), so learning stalls. Dividing by √d_k rescales the scores back to roughly unit variance regardless of head width, keeping the softmax in its responsive, well-conditioned range where gradients stay healthy. It is the cleanest possible fix: the exact factor that cancels the √d_k the dot product introduced.

Softmax: turning scores into weights over keys

Scaled scores are still arbitrary real numbers. To average values we need non-negative weights that sum to one, and softmax is the standard map from scores to such a distribution. It is applied independently to each row — each query gets its own distribution over the keys:

attn[i, j] = exp(scores[i, j] / √d_k)
             ----------------------------------
             Σ_{j'} exp(scores[i, j'] / √d_k)

Σ_j attn[i, j] = 1   for every row i

The exponential makes every weight positive and amplifies differences: a key that scores a little higher than its neighbours gets a disproportionately larger share. Softmax is also shift-invariant — adding a constant to a whole row leaves the result unchanged — which is why real implementations subtract each row’s maximum before exponentiating, a numerical-stability trick that prevents exp of a large number from overflowing. The output attn is still an N × N matrix, but now every row is a proper probability distribution: attn[i, j] is the fraction of token i’s attention spent on token j.

The weighted sum with V

The final step spends those weights. Multiply the attention matrix by the value matrix and each output row becomes a convex combination of value vectors:

out = attn V            [N, N] × [N, d_v]  =  [N, d_v]
out[i] = Σ_j attn[i, j] · V[j]

Output row i is a blend of the value vectors of the tokens that token i attended to, weighted by how strongly it attended to each. Because the weights are non-negative and sum to one, the output always lies inside the convex hull of the value vectors — attention interpolates, it never extrapolates past the values it is given. Notice the clean division of labour: Q and K decided how much weight each token gets, and V decides what content that weight pulls in. The output shape is [N, d_v] — one refreshed vector per input token — ready to be projected and passed up to the next sublayer.

Assembling the full formula

Chaining the four steps gives the canonical single line, and reading it right-to-left recovers the pipeline exactly:

Attention(Q, K, V) = softmax( Q K^T / √d_k ) V

  Q K^T          scores        [N, N]
  / √d_k        scale         [N, N]
  softmax(·)     weights       [N, N], rows sum to 1
  · V            weighted sum  [N, d_v]

Every symbol now has a job. Q K^T measures all-pairs similarity; √d_k keeps that similarity numerically tame; the softmax turns it into per-query distributions; and the multiply by V reads out content. The operation is permutation-equivariant — it has no built-in notion of order, which is precisely why transformers add positional information to the input separately. It is also fully parallel across positions: unlike a recurrent network there is no sequential dependency in computing the scores, which is a large part of why attention trains so efficiently on modern hardware.

Causal masking: keeping the future out

In a decoder — the setting for text generation — token i must not peek at tokens that come after it, or the model would be trained to predict a word using words it has not been allowed to see yet. Attention enforces this with a causal mask applied to the scores before the softmax:

masked[i, j] = scores[i, j] / √d_k   if j ≤ i
             = -∞                     if j > i

Setting the future entries to negative infinity means exp(-∞) = 0, so those keys receive exactly zero weight after the softmax — token i can attend only to itself and to earlier positions. The result is a lower-triangular attention matrix. In practice -∞ is a large negative constant, and the mask is applied post-scaling so the softmax normalizes over only the visible keys. This is what lets a decoder be trained on an entire sequence in parallel while still respecting the left-to-right generation order — each row learns from its own valid prefix. Encoders, which are allowed to see the whole input, simply omit the mask (bidirectional attention).

Complexity: why attention is O(N^2 d)

Count the arithmetic. Two matrix multiplies dominate the operation, and both scale the same way:

Q K^T :  [N, d_k] × [d_k, N]  →  N*N*d_k  = O(N^2 d_k)
attn V:  [N, N]  × [N, d_v]   →  N*N*d_v  = O(N^2 d_v)

compute:  O(N^2 d)      memory (scores):  O(N^2)

With d = d_k = d_v, the total compute is O(N^2 d) and the score matrix alone needs O(N^2) memory. The quadratic in N is the headline: double the sequence length and both the compute and the attention-matrix memory quadruple. The softmax itself is only O(N^2) element-wise work, cheaper than the multiplies, and the linear-in-d projections that build Q, K, and V cost O(N d^2) — which actually dominates when d > N (short sequences), so attention is only quadratically expensive once the context is long. This single N^2 term is the reason FlashAttention, sliding windows, and every other long-context trick exists: they attack the quadratic without materializing the full N × N matrix.

A worked example on tiny matrices

Make it concrete with N = 3 tokens and d_k = d_v = 2. Take one query and the three keys and values (already projected):

Q[1] = [1, 0]
K = [ [1,0],   [0,1],   [1,1] ]       (K1, K2, K3)
V = [ [10,0],  [0,10],  [5,5] ]       (V1, V2, V3)

Step 1 — scores for query 1: ⟨Q1,K1⟩ = 1, ⟨Q1,K2⟩ = 0, ⟨Q1,K3⟩ = 1, so scores = [1, 0, 1]. Step 2 — scale by √d_k = √2 ≈ 1.414: [0.707, 0, 0.707]. Step 3 — softmax: exp = [2.028, 1.000, 2.028], sum = 5.056, so attn = [0.401, 0.198, 0.401] (sums to 1, and the two matching keys share most of the weight).

out = 0.401·[10,0] + 0.198·[0,10] + 0.401·[5,5]
    = [4.01, 0] + [0, 1.98] + [2.005, 2.005]
    = [6.02, 3.98]

The output leans toward V1 and V3 (the keys that matched) and sits inside the triangle of value vectors, exactly as the convex-combination property promised. The other two output rows follow by repeating the same three steps with Q[2] and Q[3].

Reading the same example with a causal mask

Suppose this is a decoder and the query above is token 2 (0-indexed: it may see positions 0, 1, 2 — here we index tokens 1..3, so token 2 sees keys 1 and 2 only, not the future key 3). The mask sets the score for key 3 to -∞ before the softmax:

masked scores = [0.707, 0, -∞]
exp           = [2.028, 1.000, 0]
sum           = 3.028
attn          = [0.670, 0.330, 0]

Key 3 now receives exactly zero weight, and the remaining weight renormalizes over the visible keys so the row still sums to one. The output becomes 0.670·[10,0] + 0.330·[0,10] = [6.70, 3.30] — computed without a single number from the future leaking in. Do this for every row and the attention matrix comes out lower-triangular: row 1 sees only key 1, row 2 sees keys 1–2, row 3 sees all three. That triangular structure is the entire mechanical content of ‘the model cannot look ahead.’

Common pitfalls and misconceptions

A few errors recur. Transposing the wrong matrix: it is Q K^T, giving an [N, N] score matrix; writing Q^T K produces a [d_k, d_k] matrix and is simply wrong. Scaling by d_k instead of √d_k: the variance argument is specific — the dot product’s standard deviation is √d_k, so that is the factor that restores unit variance; dividing by d_k over-shrinks the logits and flattens the attention toward uniform. Softmax along the wrong axis: it must normalize over keys (within each row), not over queries; normalizing down the columns mixes independent queries and breaks the ‘each row is a distribution’ invariant.

Two more: forgetting the mask is pre-softmax — masking after the softmax leaves nonzero weights that no longer sum to one; and conflating d_k with d_model — in multi-head attention each head uses a smaller d_k = d_model / h, and the scaling uses the per-head width, not the full model width. Single-head attention hides that distinction because the two coincide, which is exactly why it is the right place to learn the mechanism first.

What this means for CPU and small models

On a CPU-hosted small language model the constants behind the big-O matter as much as the exponents. The O(N^2) score matrix is not just compute — it is memory traffic, and CPUs are far more bandwidth-limited than GPUs. Materializing a full N × N matrix per layer per head thrashes the cache once N grows, so the same FlashAttention-style tiling that helps GPUs — computing attention in blocks and never storing the whole matrix — is doubly valuable on a CPU, where it keeps the working set inside L2/L3.

The projection cost O(N d^2) is a reminder that for the short prompts a small on-device model typically handles, attention may not even be the bottleneck — the feed-forward and projection matmuls, linear in N but quadratic in d, often dominate. Knowing which term rules your regime tells you what to optimize: quantize and tile the projections for short-context chat, and attack the quadratic score matrix only once you push toward long documents. Either way, the exact operation is the same four steps derived above — project, score, scale-and-softmax, blend — and every performance decision is really a decision about how to compute those four steps without wasting memory bandwidth.