Attention mechanism

The idea

A fixed weight matrix mixes inputs the same way for every example. Attention computes its mixing weights from the data, per token — the network decides at runtime which positions to read and how much. The metaphor that survives contact with the math is a soft dictionary lookup: each token emits a query ("what I am looking for"), each token offers a key ("what I contain") and a value ("what I will contribute"). Match strength = dot product of query and key; the returned result is not the single best match but a softmax-weighted blend of all values.

$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left( \frac{QK^\top}{\sqrt{d_k}} \right) V$

Shapes (worth reciting in interviews): for $n$ queries against $m$ key-value pairs, $Q \in \mathbb{R}^{n \times d_k}$, $K \in \mathbb{R}^{m \times d_k}$, $V \in \mathbb{R}^{m \times d_v}$; scores $QK^\top$ are $n \times m$; softmax is applied row-wise; output is $n \times d_v$. In self-attention, $Q, K, V$ are three learned linear projections of the same sequence; in cross-attention the queries come from another sequence. Causal masking sets forbidden scores to $-\infty$ pre-softmax so their weights land at exactly 0.

Why the $\sqrt{d_k}$: with unit-variance components, $q \cdot k$ has variance $d_k$ (a sum of $d_k$ independent unit-variance products), so at $d_k = 64$ raw scores have standard deviation 8. Softmax over scores that spread is essentially argmax — near one-hot weights whose gradients $p(1-p)$ round to zero — and the attention pattern stops learning. The scaling holds the score distribution at unit variance for any head size.

Cost: standard self-attention performs $O(n^2 d)$ arithmetic and naively materializes an $O(n^2)$ score matrix per head. FlashAttention avoids that materialization and reduces memory traffic while retaining quadratic arithmetic; KV caching avoids recomputing old keys and values during autoregressive inference. Sparse or linear variants change the computation.

Worked example

One query attending over three positions, $d_k = 2$, so scale $\sqrt{2} \approx 1.414$. Query and keys:

$q = (1, 0), \qquad k_1 = (1, 0), \quad k_2 = (0, 1), \quad k_3 = (-1, 0)$

with (scalar, for clarity) values $v_1 = 10$, $v_2 = 0$, $v_3 = -10$.

(Weights = each exponential over the sum $3.521$.) Output:

$0.576 \cdot 10 + 0.284 \cdot 0 + 0.140 \cdot (-10) = 5.76 - 1.40 = 4.36$

The aligned key dominates but does not monopolize — position 3, pointing opposite the query, still contributes a little negative value. That softness is the point: gradients flow to every projection, so all three roles keep learning.

Now watch saturation, the thing scaling prevents. At $d_k = 64$ unscaled scores would spread with standard deviation 8 — take scores $(8, 0, -8)$: weights $\propto (e^8, 1, e^{-8}) = (2981, 1, 0.000335)$, i.e. $(0.99966, 0.00034, 0.0000001)$. Effectively a hard, frozen argmax with vanishing gradients — after scaling by $\sqrt{64} = 8$ the scores become $(1, 0, -1)$, right back in the soft regime above.

Cheat sheet

Formula: $\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\!\left(QK^\top / \sqrt{d_k}\right) V$, softmax row-wise.

• Scaling derivation in one line: unit-variance components → $\mathrm{Var}(q \cdot k) = d_k$ → divide by $\sqrt{d_k}$ to restore unit variance → softmax stays soft, gradients survive.
• Causal mask: scores at $j > i$ set to $-\infty$ pre-softmax → exact zero weights.
• Self vs cross: same sequence for Q/K/V vs queries from elsewhere (decoder → encoder, query → retrieved docs).
• Complexity: $O(n^2 d)$ time, $O(n^2)$ score memory; doubling context 4x-es attention compute.
• KV cache: at generation, reuse cached K, V; only the new token is projected → $O(nd)$ per step. Cache K and V, never Q. Memory = layers × heads × context × head-dim × 2 × bytes.
• Attention is permutation-equivariant — order information must be injected separately (positional encodings).
• Gotcha: softmax is over the key dimension (each query row sums to 1), not over queries.
• Gotcha: "attention weights = explanation" is contested — treat attention maps as diagnostics, not proof.