--- title: "Attention Mechanism" description: "Self-attention and transformers" source: https://vizlearn.app/ai/nlp/attention-mechanism --- # Attention Mechanism The attention mechanism is how a Transformer makes each token's vector depend on the tokens around it. Every token projects a Query, a Key, and a Value; scaled dot-product attention scores each Query against every Key, turns the scores into weights with a softmax, and blends the Values in those proportions — producing a context-aware vector per token. The load-bearing picture is a lookup that refuses to commit. A database lookup matches your query against the keys and returns exactly one value. Attention runs the same query-key-value machinery but returns a weighted mixture of all the values, with the weights set by how well the query matches each key. *Bank* next to *river* is a slope of earth; *bank* next to *savings* holds money — and the same `bank` vector retrieves a different mixture in each sentence, because the keys on offer changed. ## How the Attention Mechanism Works Start with a row of token vectors `X`. Three learned matrices give each token three roles in the lookup: - **Query** `q` — the question this token asks: what am I looking for? - **Key** `k` — the label this token shows: what do I offer to match against? - **Value** `v` — the content this token hands over if attended to. For one token, attention is three moves. **Score:** dot its Query with every Key — a larger `q·k` means a better match — then divide each score by `√dₖ`, the square root of the key dimension, which controls score variance under common scale assumptions and reduces the risk of softmax saturation; it does not guarantee unsaturated scores. **Softmax:** convert the scores into attention weights, positive numbers summing to 1, so they read as a budget: spend 69% of my attention here, 14% there. **Blend:** sum every token's Value scaled by its weight. The result replaces the token's vector with a weighted summary of what it looked up. Run all three moves for every token simultaneously — that is what the matrix form `softmax(QKᵀ/√dₖ)V` computes. It has two parts: the attention matrix `W = softmax(QKᵀ/√dₖ)`, whose row *i* is token *i*'s budget over all the tokens, and the output `Z = WV`, whose row *i* is token *i*'s new vector. ## Attention Step by Step The default sentence is `by the river bank`, and the animation walks the lookup for **bank** in full. The vectors are 4-dimensional so you can read them as numbers, which makes `√dₖ = 2` a clean divisor. 1. **Tokens.** Four tokens, each a small embedding vector. 2. **Project.** Every token sprouts its Query, Key, and Value. Bank's query comes out as `[2.3, 2.3, 0, 0]` — an ambiguous word asking loudly for content words of either sense. 3. **Score each key.** Dot bank's Query with every Key and divide by 2. Against `by`: `0.92 / 2 = 0.46`. Against `the`: `0`. Against `river`: `4.6 / 2 = 2.3` — bank's question lines up squarely with river's label. Against itself: `1.38 / 2 = 0.69`. 4. **Softmax.** The scores `[0.46, 0, 2.3, 0.69]` become the weights `[0.11, 0.07, 0.69, 0.14]`. Bank spends 69% of its budget on *river*, 14% on itself, and the change on the function words. 5. **Blend.** Each Value is scaled by its weight, then summed. River's Value dominates the pile, and bank's new vector lands at about `[1.43, 0.06, 0.14, 0]` — pointing hard down the geographic dimension. Bank now carries river inside it. 6. **Parallel.** The same three moves run for all four tokens at once, filling the full 4×4 attention matrix — one budget row per token. Switch the context to `at the savings bank` and the identical `bank` embedding attends to *savings* instead, blending in the financial Value. Nothing about the word changed; the keys available to match against did. On the canvas, watch bank's matrix row fill left to right — raw numbers during the score sweep, heat after the softmax — and the thick violet ray on the attend beat point at the token bank chose to listen to. ## Why Attention Is a Soft Lookup A hard lookup is brittle in two ways that matter for learning. It returns one row, so near-matches contribute nothing; and it is a discontinuous choice, so no gradient flows through it. Softmax fixes both at once: every Value contributes in proportion to its match, and the score, normalize, and blend operations are differentiable. The model can *learn* what to ask for, because being slightly wrong produces a slightly different answer rather than no answer. The Normalization selector lets you break the lookup on purpose, which is the fastest way to feel why each part exists. **Scaled** is the real formula. **No √dₖ** feeds the raw dot products to the softmax; they become larger in this example, making the distribution more concentrated — a soft lookup collapsing back into a hard one, which is precisely the failure the `√dₖ` term prevents (and saturated softmaxes pass almost no gradient during training). **Uniform** ignores the scores and forces every weight to `1/4`; every token's output becomes the same plain average of all four Values. That average still carries the sentence's content, but no token selects what it needs: query-specific selection is gone, the baseline of a head that learned nothing. Both are labelled diagnostics; Scaled is the standard scaled dot-product rule being taught, while unscaled dot-product and uniform weighting are alternative computations. The price of looking everything up against everything: with `n` tokens of dimension `d`, the score matrix alone costs O(n²·d). That quadratic in `n` is exactly why long context windows are expensive and why so much Transformer research is attention made cheaper. ## Edge Cases - **A token attends to itself.** The diagonal is part of the lookup — bank keeps 14% on itself here. Normal, not a bug. - **Weights always sum to 1.** Softmax runs across the keys, one row at a time, so each token spends exactly 100% of its attention. - **The matrix is not generally symmetric.** Each row is its own budget: in the Scaled run, `bank → river` and `river → bank` are different numbers, and the parallel beat highlights one such pair. Forced Uniform is the symmetric exception, with every entry `1/4`. - **Attention is blind to word order.** Nothing in score → softmax → blend reads a token's position; shuffle the tokens with their embeddings and the outputs shuffle identically. Transformers can encode order with additive positions, relative-position terms, or rotations of queries and keys — the hand-set embeddings here skip that, which is a simplification of the input, not a property of the architecture. ## Common Mistakes - **Treating Q, K, and V as one vector.** They are three different projections of the same token, and a token's query need not match its own key — bank's barely does. - **Softmaxing the wrong axis.** Normalize across the keys, so each query's row sums to 1. The wrong axis silently breaks the layer while producing plausible-looking numbers. - **Dropping the `/√dₖ`.** It is not a rounding nicety; omitting it can sharpen the distribution and weaken the softmax Jacobian, depending on score scale. The No √dₖ diagnostic shows the collapse directly. - **Reading the matrix backwards.** Row *i*, column *j* means token *i* (the query) attends to token *j* (the key) — never the reverse. ## A Note on Simplification This lesson runs the real scaled dot-product formula on a toy head: dimension 4, a single head, embeddings and projections hand-set so the `bank → river` / `bank → savings` pattern is legible on screen. Production model and per-head dimensions vary; many larger models run this computation at wider dimensions, across many heads, in dozens of stacked layers, with every matrix learned rather than hand-set — and they use a positional mechanism, which the fixed inputs here skip. The arithmetic is the genuine article; what's missing is scale and learned parameters.