Attention & the Transformer

The RNN's weakness in Module 6 was structural: to relate two distant words it had to carry information across every step between them, and both the gradient and the memory leaked away. And it was sequential — step t could not start until step t−1 finished.Attention deletes both problems at once. Any token can look directly at any other in a single step, no matter how far apart, and every one of those look-ups happens in parallel. The 2017 Transformer built an entire architecture out of this one operation — no recurrence, no convolution — and swept the field. Every LLM you have used is a stack of the block we build in this chapter.

7.1 Query, key, value

Think of each token as posting a small ad. Its query says "here is what I'm looking for." Every token also publishes a key — "here is what I am" — and a value — "here is what I'll hand over if you pick me." A token decides how much to listen to each other token by comparing its query against that token's key with adot product: aligned vectors score high, unrelated ones score low. Those scores go through asoftmax to become weights that sum to 1, and the output is the weighted average of everyone's values.

The 1/√d is not decoration. Dot products of d-dimensional vectors grow with d; without the scaling the softmax would see huge numbers, saturate, and collapse onto a single token — killing the gradient. Dividing by √dkeeps the scores in a sane range. Below, seven tokens sit in a tiny 2-D key space. Drag the query dot — or click a token to attend as it — and watch the attention redistribute. The thickness of each thread is how much that token is attended to; the hollow star is the output, the weighted blend of the attended values.

7.2 Self-attention as a matrix

In a real layer every token is a query and a key at the same time — the sequence attends to itself, which is why it's called self-attention. Doing all the queries at once turns the operation into a single matrix: row i is the attention that token i pays to every token j, softmax-normalized so each row sums to 1. Below is exactly that matrix for our sentence — brighter cells mean more attention. Read across a row to see what one token looks at; notice how "it" lands mostly on"cat", the coreference resolving itself out of pure vector similarity.

There is one more essential piece. When a model generates text it must predict token i using only tokens 1…i — it cannot peek at words it hasn't written yet. So decoder self-attention applies acausal mask: every score above the diagonal is set to −∞ before the softmax, zeroing it out. Toggle the mask and watch the upper triangle go dark. And note the shape: the matrix is T×T for a sequence of length T — its cost grows with the square of the context length, the single fact that dominates the endgame below.

7.3 Many heads, many patterns

A single attention pass is a head, and one head can only track one kind of relationship at a time. So a Transformer runs many in parallel — multi-head attention — each with its own learned query/key/value projections, so each head slices the sequence differently. Their outputs are concatenated and mixed by one more linear layer. When researchers open up a trained model (Transformer Circuits) they find heads that specialize: some attend to the previous token, some to nearby words, some resolve what a pronoun refers to. Below, three hand-built heads over the same sentence — flip between them and watch the arcs change.

Each head is cheap and narrow; together they give the layer a rich, multi-relational view of the sentence in a single parallel pass. A head that attends to the previous token, stacked with a head in the next layer that copies from wherever the first one pointed, is the seed of an induction head — the circuit that lets a model continue a pattern it has seen once. This is where in-context learning begins.

7.4 The Transformer block

Attention is the new part; the rest of the block is machinery you already met. Wrap the multi-head attention in a residual connection and a layer normalization, then follow it with a small position-wise feed-forward MLP — the same token, run through a two-layer network that expands to ~4× width and back — with its own residual and norm. That pairing is the whole block: attention mixes information across tokens; the MLP transforms each token on its own. Modern models put the norm before each sublayer (pre-norm), which trains more stably. Follow a tensor through one block below.

Stack a few dozen of these identical blocks and you have the body of a modern LLM. Because attention itself is order-blind — it is just a weighted set, with no notion of first or last — the model must also be toldwhere each token sits. That's positional encoding: the original Transformer added fixed sinusoids to the embeddings; today's models rotate the query and key vectors by a position-dependent angle (RoPE), which we assemble in Module 9.