10 Attention
A recurrent network reads a sequence the way a clerk with one notepad reads a ledger: each entry is folded into the running summary, and by the end the summary is all that remains. Chapter 8 built models of exactly this shape and Chapter 12 will confront their limit directly: a fixed-size state must eventually discard something, and the model cannot know at reading time what it will be asked later. Attention removes the constraint by refusing to summarize. It keeps every entry, and when a question arrives it looks the answer up: compare the question against a key describing each entry, convert the comparison scores into weights, and return the weighted average of the entries’ values. The three names — query, key, value — come from databases, and the analogy is exact except on one point: a database returns the single best match, while attention returns a soft blend of all of them. That one change makes lookup differentiable, and a differentiable lookup can be learned.
This chapter is about that mechanism, on its own terms. The reason it deserves a chapter is what became of it. Attention entered deep learning in 2014 as a patch: machine-translation models were forcing whole sentences through a fixed-size vector, and letting the decoder look back at the source at every step repaired the bottleneck (Bahdanau et al. 2015). Three years later the patch swallowed the architecture — Vaswani et al. (2017) discarded the recurrence entirely and kept only attention, and the resulting transformer is the subject of Chapter 11. The mechanism then outgrew its birthplace: the models that read and write text, classify and generate images, transcribe speech, and predict protein structures are, at their core, stacks of the same lookup. Its weights are a quantity the field measures, engineers around, and argues about; its quadratic cost sets the economics of long-context models; and its learned circuits are the best-understood fragments of what large models actually compute. A mechanism with that reach is worth understanding precisely, before any architecture is wrapped around it.
The six sections build it up in order. Section 10.1 states the lookup abstraction and shows that classical kernel regression is attention with hand-picked weights — the case for learning them. Section 10.2 turns similarity into scores: scaled dot products, why the \(1/\sqrt{d}\) factor is a matter of softmax saturation rather than convention, and masking, the small bookkeeping device that lets one implementation serve padded batches and causal language models alike. Section 10.3 runs several attention functions in parallel and proves, on the smallest task that defeats it, that a single head must average what a query asks for separately while two heads recover it, then separates the two wirings — self-attention, where a sequence queries itself, and cross-attention, where one sequence queries another. Section 10.4 restores what the lookup deliberately ignores: order. Attention is permutation-equivariant, so position must be injected, and we follow the idea from sinusoidal encodings to the rotary embeddings used by essentially every current model, ending with an experiment on what happens beyond the training length — whose outcome is not what the folklore promises. Section 10.5 faces the price: computing every query–key pair is quadratic in sequence length, and we measure it, then implement the three escapes: computing exact attention without ever materializing the score matrix, restricting it to a window, and linearizing it, which turns attention back into a recurrence and hands the story to Chapter 12. Section 10.6 closes by asking what trained attention layers actually do, and answers with running code: a two-layer attention-only model visibly learns an induction circuit — find the previous occurrence of the current token, copy what followed it — and that small mechanism is the clearest laboratory example of how in-context learning arises.
Along the way we train one model architecture: a character-level, attention-only language model a few blocks deep, introduced in Section 10.4 and trained afresh on a different task in Section 10.6, with no run longer than about a minute. Everything else executes in seconds on data chosen so that the phenomenon under study is unmistakable. A word on what this chapter is not. It builds no transformer (no feed-forward layers, no normalization, no training recipes); that assembly, and the architecture zoo around it, is Chapter 11. It says nothing about optimizers, which have Chapter 9 to themselves. And it treats the efficient-attention literature with deliberate economy: a decade of approximations is summarized in a paragraph, because the variants that survived are the ones we implement. Where attention touches state-space models, Chapter 12 owns the recurrence side of the correspondence.
Resources and Further Reading
The references follow the chapter’s arc: intuition for the mechanism, the founding papers, the cost of attention, and the circuits view. All are freely accessible online.
Visual introductions
- Attention in transformers, step-by-step — 3Blue1Brown (2024) — the geometric reading of Section 10.1: queries as questions, keys as advertisements, values as contributions, animated in embedding space; deliberately code-free, which this chapter’s notebooks complement.
- The Illustrated Transformer — Jay Alammar (2018) — the field’s shared picture vocabulary for Section 10.2 and Section 10.3; the “it”-coreference attention pattern reproduced in half the talks you will ever see started here.
- Transformers from scratch — Peter Bloem (2019) — derives self-attention as a permutation-equivariant set operation before any architecture appears, the same order of ideas as Section 10.4’s opening.
Founding and load-bearing papers
- Neural Machine Translation by Jointly Learning to Align and Translate — Bahdanau et al. (2014) — where the mechanism began, as the alignment story Section 10.2 retells; the attention-weight heatmap over a translation originates here.
- Attention Is All You Need — Vaswani et al. (2017) — scaled dot-product attention and multi-head attention as this chapter teaches them; the architecture half of the paper belongs to Chapter 11.
- RoFormer: Enhanced Transformer with Rotary Position Embedding — Su et al. (2021) — rotary embeddings as derived in Section 10.4, now the default positional scheme of open models.
- Train Short, Test Long — Press et al. (2022) — ALiBi, and the extrapolation question that Section 10.4 puts to experiment.
The cost of attention
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al. (2022) — the industrial form of Section 10.5’s chunked online-softmax cell: exact attention, quadratic time, linear memory, organized around memory traffic rather than arithmetic.
- Transformers are RNNs — Katharopoulos et al. (2020) — the linear-attention equivalence that Section 10.5 verifies numerically, and the cleanest bridge between this chapter and Chapter 12.
- The Transformer Family v2.0 — Lilian Weng (2023) — a taxonomy of the attention-variant design space; the map for which Section 10.5 implements the surviving territory.
What attention computes
- A Mathematical Framework for Transformer Circuits — Elhage et al. (2021) — the QK/OV factorization and residual-stream view that Section 10.6 teaches in executable form.
- In-context Learning and Induction Heads — Olsson et al. (2022) — the evidence, at model scale, for the phase change our two-block model reproduces in minutes.
- Thinking Like Transformers (Transformer Puzzles) — Sasha Rush — RASP puzzles on what a fixed-depth attention machine can compute at all — a complement to this chapter’s exercises that requires no training.
Exercises with teeth
- Stanford CS224n, self-attention and transformers assignment — the graded counterpart of this chapter: constructions for what single heads cannot do, the permutation-equivariance proof, and a RoPE relative-position derivation, all of which appear here as section material or exercises.