6.17 Attention from first principles

What this is and why it exists

Attention is a retrieval operation, and once you can write it from memory the transformer stops being a diagram you have seen and becomes a mechanism you can reason about. This topic derives it from the retrieval analogy, explains why a particular division appears in the middle of it, and then covers masking — which is what makes both padding and generation work. Almost every attention bug traces to confusion about which tensor is playing which role, so the topic insists on that clarity throughout.

The vocabulary

• Query — what a position is looking for.
• Key — what a position advertises about itself.
• Value — what a position contributes when attended to.
• Attention score — how well one query matches one key.
• Attention weight — a score after the softmax, so the weights over all positions add to one.
• Head — one independent attention operation.
• Mask — a rule preventing certain positions from being attended to.
• Causal — masking that hides everything after the current position.

The mental model

Think of it as a soft dictionary lookup. In an ordinary dictionary you present a key and get back exactly the matching value. In attention, every position presents a query, every position advertises a key, the query is compared against all the keys, and the answer is a blend of all the values weighted by how well each key matched. Nothing is chosen; everything is mixed in proportion to relevance. That is the entire operation, and the rest is arithmetic.

Concretely, three separate learned projections turn each position's representation into a query, a key and a value. Every query is compared with every key by a dot product, which is large when two vectors point in similar directions. Those scores are divided by a constant, passed through a softmax so each position's weights over all positions are positive and sum to one, and used to take a weighted average of the values. Say the three roles out loud when you write it: queries come from the positions doing the asking, keys and values come from the positions being asked about. In self-attention all three come from the same sequence; in cross-attention the queries come from one sequence and the keys and values from another, which is exactly how a decoder consults an encoder.

Now the division, because it exists for a concrete reason. The scores are divided by the square root of the key dimension. Without it, the dot product of two random vectors grows with the number of dimensions being summed, so with a large key dimension the scores become large, and a softmax over large scores is nearly one at its maximum and nearly zero everywhere else. That is a hard selection rather than a blend, and — the part that matters — the softmax's gradient in that regime is almost zero, so nothing learns. Dividing by the square root of the dimension keeps the scores at a scale where the softmax stays soft and gradients survive. PyTorch's fused implementation encodes this as its default: if no scale is given, it is set to one over the square root of the embedding dimension.

Multi-head attention runs several of these in parallel. Split the representation into a number of heads, run a complete attention operation within each, concatenate the results and project once more. Each head can specialise — one attending to the previous word, another to a matching bracket, another to the subject of a sentence — because a single averaged blend cannot represent several different relationships at once. It is a small change with a large effect on capacity, and it costs almost nothing because each head works in a proportionally smaller dimension.

Masking is the last piece and it does two different jobs. Padding masks hide positions that are padding rather than data, so a short sequence in a batch of long ones does not contribute. Forget it and the model attends to filler and your results shift with batch composition, which is a genuinely confusing bug. Causal masking hides everything after the current position, so a position can attend to the past and to itself but never to the future. That is what makes generation possible: it lets the model be trained on every position of a sequence at once while still being usable one token at a time, because no position ever saw an answer it was supposed to predict. Get it wrong and training looks superb — the model is reading ahead — and generation is nonsense, which is one of the most instructive bugs in the field.

Masks are applied by making the forbidden scores very large and negative before the softmax, so their weights become effectively zero. PyTorch's fused operation exposes both routes: a flag whose documentation states that when it is set true "the attention masking is a lower triangular matrix when the mask is a square matrix", and an explicit mask parameter accepting either "a boolean mask where a value of True indicates that the element should take part in attention" or "a float mask of the same type as query, key, value that is added to the attention score". Note the sense of the boolean convention — true means attend — because getting it inverted silently trains a model on exactly the wrong half of its input.

Write it once from memory. Project to queries, keys and values; multiply queries by transposed keys; divide by the square root of the key dimension; add the mask; softmax over the key positions; multiply by values. Six lines. Then use the fused implementation in real code, because it is substantially faster and uses far less memory by never forming the full score matrix — but the six lines are what let you read an error message about a mask shape and know immediately which dimension is wrong.

Two costs to carry forward. Every position attends to every other, so the computation and memory grow with the square of the sequence length, which is the constraint behind everything you will read about long contexts. And attention has no inherent notion of order — shuffle the positions and it computes the same blend — so position information has to be added deliberately, which is a topic of its own in the language models module.

What you should now be able to explain or do

Give the retrieval analogy and name what each of the three roles does. Say where queries, keys and values come from in self-attention and in cross-attention. Write scaled dot-product attention from memory in the right order. Explain why the scores are divided by the square root of the key dimension, in terms of the softmax and its gradient. Say what multi-head attention buys and why it is nearly free. Distinguish padding from causal masking and say what each prevents. State the two structural costs of attention.

Check yourself