🧲 Queries, keys, values: a differentiable lookup

A dictionary lookup retrieves the single value whose key exactly matches the query. Self-attention softens both halves of that sentence: every key matches a little, and the result is a weighted blend of all values. Because the blend weights come from a softmax — smooth in its inputs — the whole lookup is differentiable, so gradient descent can learn what to ask for and what to advertise. In matrix form, for a sequence packed into rows of X:

The "self" in self-attention: queries, keys, and values are all learned linear projections of the same input sequence,

so the sequence is querying itself. Written per token, the weight that token i places on token j, and the resulting output, are

A concrete sequence carries the arithmetic for the rest of the note: "the robot picked up the ball because it was heavy". The token "it" is the interesting one — resolving what "it" refers to requires looking at other tokens, which is exactly what a lookup provides. Suppose the trained projections give "it" a query whose dot products (already scaled) with the other tokens' keys are: ball 3.4, heavy 2.8, robot 2.1, and 0 for the remaining seven tokens.

Nothing in the formula privileges nearby tokens, earlier tokens, or the token itself: relevance is decided entirely by the learned q·k geometry. That single fact explains both attention's power (any token can reach any other in one step) and the two complications handled below — the need for masking in §2 and for positional information in §7.

🗺️ Reading an attention pattern

One row of the attention matrix — one token's outgoing weights — is a probability distribution over the sequence. Below, the arcs and bars show that row for the selected token, using a hand-built pattern of the kind a trained coreference-tracking head produces.

Causal masking. A decoder generating left-to-right must not let token i read tokens that come after it — at generation time they do not exist yet. The implementation is not "zero out the weights afterwards": the scores of future positions are set to −∞ before the softmax,

so the surviving weights renormalize over the visible prefix and still sum to 1. The demo makes the difference concrete: under the mask, "it" (position 8 of 10) can no longer see "heavy" — and the weight "heavy" used to receive does not vanish into nothing, it flows proportionally to "ball", "robot", and the rest.

Click (or hover) a token to see where its query lands. Then toggle the causal mask with "it" selected and watch the renormalization: "heavy" is cut off, and its share flows to the remaining tokens — "ball" rises above its unmasked 49%.

🌡️ Why divide by √d: the scaling is a temperature

The one non-obvious constant in the formula is the √d_k divisor. The problem it solves is statistical: a dot product is a sum of d_k terms, and sums grow with the number of terms. If the components of q and k are independent with mean 0 and variance 1, then

At d_k = 64 the raw scores are typically 8× larger than at d_k = 1. Feed scores that large into a softmax and it saturates: one weight pins to ≈1, the rest to ≈0, and the gradient through the softmax — which is proportional to α(1−α) terms — collapses toward zero. Early in training that means the attention pattern freezes into an arbitrary hard lookup before it has learned anything. Dividing by √d_k restores the scores to unit variance regardless of dimension.

Seen through the lens of the temperature note, this is not a new trick at all:

The divisor is a fixed softmax temperature, chosen so that the effective score distribution neither freezes into an argmax (cold) nor flattens into a uniform average (hot). The same vocabulary applies to reading trained heads: the Shannon entropy of a weight row measures how focused the head is — a near-one-hot coreference row has low entropy, a broad averaging row has high entropy. The next demo puts a hand on this dial.

🧭 Attention weight is vector alignment

Everything upstream of the softmax is a dot product, and a dot product is geometry:

A key scores highly by pointing the same way as the query (cosθ → 1) or by being long (a learned way for a token to shout). And because the softmax normalizes across keys, weights are relative — a key wins not by scoring well but by scoring better than the others. Keys compete.

🎭 Many heads, many relations

One softmax produces one distribution — one pattern of relevance. But a token often needs several relations at once: "it" wants its referent (a coreference pattern), its local context (an adjacency pattern), and the salient content words (a salience pattern). A single head forced to serve all three can only average them into mush. Multi-head attention runs h independent copies of the lookup, each in its own learned d_model/h-dimensional projected subspace, then concatenates the results and mixes them with one more linear map:

Because each head works at width d_model/h, the parameter count and compute are about the same as one full-width head — the heads are not extra capacity, they are the same capacity partitioned so that separate relations get separate softmaxes. (The original transformer uses h = 8 heads of d_k = 64 at d_model = 512.)

Click a token in the strip; its row is outlined in each head's matrix. With "it" selected, head A points at "ball" (coreference), head B at "was" (previous token), head C at the three content words (salience) — three simultaneous, mutually incompatible patterns. Flip to 1 head: the averaged scores hedge every row into a blur — visibly higher entropy, and no single relation survives cleanly. The strip underneath shows the mechanics: each head writes a d_model/h-slice of the output, and W^O mixes the concatenation back to width d_model.

🕰️ Where attention came from, and why recurrence was dropped

Attention predates the transformer by three years, and it was invented to fix a bottleneck. The RNN encoder–decoders of 2014 translated by reading the whole source sentence into one fixed-size vector, then generating from it — every fact about a 50-word sentence squeezed through the same few hundred numbers. Bahdanau et al. (2015) let the decoder instead look back at all encoder hidden states at every step, scoring each with a small learned network (additive attention):

Luong et al. (2015) simplified the score to a plain dot product (multiplicative attention),

which is the direct ancestor of QK^T. In both, attention was an add-on: the RNN still did the representation work, attention just fetched. The 2017 move was to delete the RNN and keep only the fetching — attention is all you need — for two structural reasons:

The last row is the honest price: every token attends to every token, so compute and the attention matrix itself grow quadratically in sequence length. For d > n the trade is favorable; for long contexts it is the central cost of the architecture, and §9 is about paying it down.

📍 Attention is order-blind: positional encoding

The formula in §1 never mentions position: shuffle the input tokens and the (unmasked) outputs shuffle identically — self-attention is permutation-equivariant, a set operation. To the mechanism, "the robot picked up the ball" and any reordering of it are the same bag of tokens. Word order has to be injected into the representations themselves. The original transformer adds a deterministic sinusoidal positional encoding to each token embedding: dimension pair 2i oscillates with position at a wavelength that grows geometrically with i, from 2π up to 10000·2π,

Fast dimensions distinguish neighbors; slow dimensions distinguish regions — together the dimensions form a smooth binary-counter-like code, unique per position. The construction has two useful properties: for any fixed offset Δ, PE(pos+Δ) is a linear function (a rotation) of PE(pos), so relative position is easy for the model to compute; and the code is defined for every position, including ones longer than any training sequence.

Modern practice moved the same idea inside the attention operation: learned absolute embeddings (GPT-style) simply train a vector per position, while RoPE (rotary position embedding, Su et al. 2021) rotates each query and key by an angle proportional to its position, so that q_i·k_j depends only on the relative offset i−j. RoPE is the standard in Llama-class LLMs; the underlying trick — encode position as rotation, so dot products see only differences — is the sinusoidal idea finished properly.

🧱 The full transformer block

Attention only ever mixes information across tokens — a weighted average cannot compute a new feature of a single token. The transformer block pairs it with a position-wise feed-forward network (the same two-layer MLP applied to every token independently, hidden width 4×d_model in the original), which transforms each token but mixes nothing. Mix across, then transform within. Both sublayers sit on a residual stream — each adds a refinement to a representation that flows through unchanged — and each is stabilized by LayerNorm:

(Written here in the pre-LN arrangement of modern stacks — normalize before each sublayer; the 2017 paper normalized after, which trains less stably at depth. The FFN, often treated as an afterthought, holds roughly two-thirds of the block's parameters.)

Stack the block and two architectures fall out. The original machine-translation transformer is an encoder–decoder: the encoder runs unmasked self-attention over the source; the decoder runs causal self-attention over the target plus a cross-attention sublayer whose queries come from the decoder and whose keys and values come from the encoder output —

— the same formula with a different source for K and V. Cross-attention is Bahdanau's decoder-looks-at-encoder attention reborn inside the new architecture. GPT-style LLMs drop the encoder entirely: a decoder-only stack of causal self-attention blocks, where "everything is context" and generation is next-token prediction.

⚡ Attention at inference: cache, share, tile

Generation exposes the O(n²) honestly: producing token t naively re-runs attention over the whole prefix, every step. Three fixes, in escalating order, are why chat-scale inference works at all.

1. The KV cache. Under causal masking, the keys and values of past tokens never change — token t cannot alter what tokens 1…t−1 advertise. So cache them: at step t only the new query q_t is computed, dotted against t stored keys. Per-step cost drops from O(t²) to O(t); the price is memory,

(the 2 is keys and values), which for long contexts rivals the model weights themselves and becomes the binding constraint on batch size.

2. Share the keys and values. The cache is dominated by the factor h — every head stores its own K and V. Multi-query attention (MQA, Shazeer 2019) keeps per-head queries but shares a single K and V across all heads, shrinking the cache h-fold at a small quality cost. Grouped-query attention (GQA, Ainslie et al. 2023) interpolates: heads are divided into g groups, each sharing one K/V pair — near-MQA memory at near-full quality, and the default in current open-weight LLMs.

3. Never materialize the matrix. FlashAttention (Dao et al. 2022) computes exact attention while refusing to ever write the n×n weight matrix to GPU main memory: it tiles Q, K, V through fast on-chip SRAM and maintains running softmax statistics (the "online softmax"), merging partial results block by block. Same math, same output — the speedup comes entirely from doing fewer slow-memory round-trips. It is the cleanest illustration of the section's theme: the quadratic lives in the pattern, so the wins come from how the pattern is scheduled, not from changing the formula.

📚 Where this lives in the literature

• Bahdanau, Cho & Bengio (2015), Neural Machine Translation by Jointly Learning to Align and Translate — additive attention over RNN encoder states; the fixed-vector bottleneck argument.
• Luong, Pham & Manning (2015), Effective Approaches to Attention-based NMT — dot-product (multiplicative) attention.
• Vaswani et al. (2017), Attention Is All You Need — scaled dot-product and multi-head attention, sinusoidal positions, the encoder–decoder transformer.
• Shazeer (2019), Fast Transformer Decoding: One Write-Head Is All You Need — multi-query attention.
• Su et al. (2021), RoFormer — rotary position embedding (RoPE).
• Dao et al. (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.
• Ainslie et al. (2023), GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.

🎯 Consequences worth remembering

• Self-attention is a differentiable soft lookup: softmax(q·k) weights over value vectors, with Q, K, V all projections of the same sequence.
• The √d divisor is a softmax temperature chosen to keep score variance at 1 — not a heuristic, a variance calculation.
• Masking is renormalization, not post-hoc zeroing: −∞ before the softmax, so the visible prefix reclaims the masked weight.
• Heads are partitioned, not extra, capacity — h parallel lookups in d/h-dimensional subspaces at the cost of one full-width lookup.
• Attention itself is order-blind; position must be injected (added sinusoids/embeddings, or RoPE rotating it directly into the dot product).
• The O(n²) lives in the pattern, not the parameters — which is why the fixes are systems fixes: cache the K/V, share them across heads, tile the computation.