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.