Overview

A convolution decides in advance which inputs may influence which outputs: neighbours, always, with the same weights. Attention makes the opposite choice. It computes the connection strengths from the data, at run time, for every pair of positions.

That is the whole mechanism, and it is worth holding onto as a one-line definition: attention is a dense layer whose weights are a function of the input. Everything else in this chapter — the \(\sqrt{d_k}\), the masks, the multiple heads, the cache that dominates your inference bill — falls out of that sentence and the cost of taking it literally.

Historically it arrived as a patch. Sequence-to-sequence translation compressed the whole source sentence into a single vector, and long sentences got worse the longer they got; Bahdanau et al. let the decoder look back at every source position instead¹. Three years later the patch had eaten the architecture: "Attention Is All You Need" removed the recurrence and kept only the looking back².

From text to the next word, step by step

A language model never sees letters. The text is cut into tokens, each token becomes a vector, attention lets those vectors exchange information, and at the end the vector of the last token is turned into a probability for every token that could come next. Seven steps, each worked out on its own page:

Steps 1 to 4 are this chapter, 5 to 7 the Transformers chapter. A side page, The cost of attention, covers what attention costs when it runs for real: the KV cache and the \(O(n^2)\).

The running example

One sentence goes through every page: "the cat climbed the", which the model should continue with " roof". The model is a toy built exactly like a real one — GPT-2 small, the model of the labs and of Transformer Explainer — only with smaller numbers, so that every matrix fits on the page. Its weights were picked by hand, mostly 0, 1 and halves, so that you can follow every sum and so that the toy does on this sentence what a trained model does.

Slides

Click the slides and use ← → to move through them; F goes full screen and P opens the presenter view, with the notes. Six of the simulators of these pages run inside the slides. Open the slides in a new tab

What came before: recurrence, and how far the gradient reaches

Before attention, a sequence was read one step at a time. An RNN carries a state \(h_t\) and rewrites it at every token:

Think of the state as a fixed-size notepad you rewrite after every word. To remember the first word at the end of a paragraph, its note has to survive every rewrite along the way.

The LSTM³ adds to that a cell \(c_t\) and three gates the network itself learns to open and close — forget, input and output (the forget gate is a later addition, from 2000⁴; the 1997 cell kept its content with a fixed factor of 1):

The cell is the whole trick. A plain RNN multiplies the state by a matrix at every step; the LSTM multiplies by \(f_t\), a number between 0 and 1 the network picks. With \(f_t\) near 1, information crosses many steps without being rewritten — the same idea as the residual shortcut of chapter 9, nineteen years earlier.

But near 1 is not 1, and the gradient travelling from the last token to the first is still a product of \(T\) factors, and a product of many numbers below 1 shrinks fast: \(0.9^{64} \approx 10^{-3}\). The panel measures all three paths on one scale:

At 64 tokens with weights at scale 1, the gradient reaches the first token at \(5 \times 10^{-4}\) of what left the last one in the plain RNN, and at \(4 \times 10^{-2}\) along the LSTM's cell path. Go to 256 tokens and those become \(10^{-14}\) and \(2 \times 10^{-6}\). Both curves are straight lines on a logarithmic plot, the signature of exponential decay: the LSTM did not solve the problem, it bought a better constant.

Attention has no such product. The output reads every position at once, so the path to any token is one step, and the gradient arrives scaled by one attention weight, not by a product with one factor per step of distance. Add to that recurrence's second cost, which the plot does not show: the \(T\) steps are sequential, and a whole GPU waits on one state at a time. It was that combination — reach and parallelism — that retired recurrence, not an isolated gain in accuracy.

The idea: every word re-reads the others

Compare two sentences:

• "I sat on the bank of the river."
• "I went to the bank to withdraw cash."

The model's first layer (the embedding, step 2) hands "bank" the same vector in both: it looks only at the token, not at its neighbours. Attention is what resolves the ambiguity. It lets every word build a new vector by mixing the vectors of the others, with more weight on the ones that matter: "bank" pulls from "sat" and "river" in the first sentence and from "withdraw" and "cash" in the second. That is what an attention layer outputs — one vector per token, now with context.

How that mixing is computed — and how the model decides which neighbours matter — is the subject of steps 3 and 4, starting at Query, key, value. First the text has to become vectors: step 1, tokens.

Lab: the mechanism, computed and measured

Every claim of this chapter is computed in a Colab notebook. It runs in the browser, on CPU; there is nothing to install, and the only download is a small public language model (SmolLM2-135M, about 270 MB) in section 9. Ten sections, each ending in a checkpoint with a check value:

1. the soft dictionary by hand: the worked example of the Query, key, value page, and what moves when the value or the query changes;
2. every query at once, the shape of every matrix, and our three lines against F.scaled_dot_product_attention;
3. why \(\sqrt{d_k}\): the score variance measured, and the softmax going one-hot without it;
4. the backward pass written by hand and checked against autograd, and the gradient of a saturated row;
5. masks: the causal mask, and padding with and without its mask;
6. permutation equivariance ("dog bites man"), and RoPE turning absolute rotations into offsets;
7. multi-head as a reshape, checked against nn.MultiheadAttention, and GQA;
8. the \(n^2\) bill: memory and time of a naive layer against the fused kernel;
9. the KV cache of a real model, measured against the formula, generation with and without it, and what its heads look at;
10. a one-layer attention model that has to use order to answer, with a table to fill in.

Open the lab in Colab — save a copy to your Drive (File → Save a copy in Drive) to run and edit it.

Key takeaways

1. Text reaches the model as tokens: pieces of a fixed vocabulary, found by BPE (merge the most frequent pair, again and again), each with an integer id. A tokenizer fitted on English writes Portuguese with more tokens — about 1.9 per word in GPT-2 against 1.1 for English.
2. An id is a name, not a number to compute with. The embedding table turns it into a vector by a lookup, and a position is added (or, with RoPE, applied to q and k), because attention on its own does not see order.
3. Attention is a dense layer whose weights are computed from the input — the opposite design choice from a convolution's fixed, local, shared weights. It beat recurrence on reach (one step to any token) and parallelism.
4. It works as a soft dictionary: it compares a question with every key and returns a weighted average of the values. The output is one vector with context per token.
5. Q, K, V are three projections of the same input. The separation is what lets a token ask for something it is not.
6. \(\sqrt{d_k}\) exists because score variance grows as \(d_k\). Without it the softmax saturates and the gradient vanishes. Log attention entropy; it is cheap and diagnostic.
7. The mask is one added matrix, and it is the entire encoder/decoder distinction. Forgetting the document mask when packing is a real and common pretraining bug.
8. Attention does not see order. RoPE encodes absolute rotations that yield relative offsets in the dot product — which is both why it works and why context can be extended by rescaling frequencies.
9. Heads are free parallelism, not extra capacity: \(h\) heads of width \(d/h\) cost what one head of width \(d\) costs.
10. The KV cache dominates inference memory (128 KiB per token in Llama 3 8B). \(h_{kv}\) is the one term you can shrink, which is why GQA is universal and MLA is where the frontier went. Inference is bandwidth-bound, not compute-bound: optimize bytes moved.
11. FlashAttention is exact and free. Sparsity, sliding windows and SSMs trade something; in 2026 the winning trades are hybrid stacks and trainable sparse attention.