Skip to content

12.5.1 — Attention and the Transformer

Translating a sentence with a recurrent network in 2015 worked like this: read the source one word at a time into a hidden state, then generate the translation from that final state.

The entire source sentence had to fit into one fixed-size vector. For a short sentence, fine. For a paragraph, the beginning was gone by the end. And because each step depended on the previous one, a 50-word sentence was 50 sequential operations that no amount of hardware could parallelise.

Attention was introduced to fix the first problem: instead of one summary vector, let the decoder look back at all the source positions and weight them for each word it produces. In 2017 Attention Is All You Need removed the recurrence entirely and kept only that mechanism — which fixed the second problem too, and that is why the architecture took over.

1. Attention as retrieval

Every position produces three vectors, and the names come from a database analogy that is worth taking literally.

  • Query — what this position is looking for.
  • Key — what this position offers, as a label.
  • Value — the content this position will hand over.

A query is compared against every key. Where they match, that position's value is pulled in. Unlike a database lookup, the match is not exact — it is a similarity score, and the result is a weighted blend of all values rather than one row.

Concretely, in "The animal didn't cross the street because it was too tired", the query from it scores highly against the key from animal, so it's new representation absorbs a large share of animal's value. That is coreference resolution falling out of a dot product, with nobody writing a rule for it.

2. The computation, step by step

X (tokens)Q = XWqK = XWkV = XWvQKᵀ / √dsoftmaxweights × Voutevery position attends to every position — n² scores, all computed at once
Scaled dot-product attention in one formula. Every arrow is a matrix multiply, and none of them depends on the previous position — which is the whole reason this replaced recurrence.

\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V

Read it in five steps.

1. Project. Each token's embedding is multiplied by three learned matrices to produce its query, key and value. These matrices are the learned part of attention — the mechanism itself has no parameters.

2. Score. QK^{\top} computes the dot product of every query with every key, giving an n \times n matrix of raw scores. Chapter 12.2's point that a dot product measures alignment is doing all the work here.

3. Scale by \sqrt{d_k}. For random vectors of dimension d_k, dot products have variance proportional to d_k, so with d_k = 64 the scores are large. Large inputs to softmax produce a nearly one-hot output, and its gradient then vanishes (Chapter 12.2's exponential amplification). Dividing by \sqrt{d_k} keeps the variance around 1 and training stable. This is the whole reason for the square root, and it is a small numerical fix rather than a deep idea.

4. Softmax per row. Each token's scores become a distribution summing to 1 — how much of its attention goes where.

5. Weighted sum of values. Multiply the weights by V. Each output position is a blend of every value, weighted by relevance.

A worked micro-example. Three tokens with d_k = 4, and query 3 producing raw scores [8, 2, 4] against the three keys. Scaled by \sqrt{4} = 2: [4, 1, 2]. Softmax: roughly [0.84, 0.04, 0.11]. So position 3's output is 84% of value 1, 4% of value 2, 11% of value 3. Position 3 is mostly looking at position 1, which is what "attends to" means numerically.

3. Multi-head attention

One attention operation learns one kind of relationship. Multi-head attention runs several in parallel with separate projection matrices, each in a smaller dimension, then concatenates and projects the result.

With a model dimension of 512 and 8 heads, each head works in 64 dimensions, so the total cost is roughly the same as one full-size head — you get diversity for free.

Different heads specialise, and inspection of trained models finds heads tracking syntactic dependencies, heads that attend to the previous token, heads that match up brackets or quotes, and heads that carry positional patterns. Nobody assigns these roles; they emerge because the heads are initialised differently and the model has capacity for several relationship types.

Grouped-query and multi-query attention are the modern efficiency variants: several query heads share one key/value head, which cuts memory during generation for a small quality cost. Section 7 explains why that memory is the binding constraint.

4. Position: the part attention loses

Attention is permutation-invariant. Nothing in \text{softmax}(QK^{\top})V depends on order — shuffle the tokens and each one produces the same output. "Dog bites man" and "man bites dog" would be identical.

So position must be injected explicitly.

Sinusoidal encodings — the original: sine and cosine waves of different frequencies added to the embeddings, chosen so that relative offsets are expressible as a linear function of the encodings.

Learned positional embeddings — a trainable vector per position. Simple, and it cannot extrapolate beyond the longest position seen in training.

RoPE (rotary position embeddings) — what most current models use. Instead of adding a position vector, it rotates the query and key vectors by an angle proportional to their position. The dot product between two rotated vectors then depends on their relative distance, which is what actually matters, and it extends to longer contexts far more gracefully. RoPE is also the mechanism behind context-window extension: scaling the rotation frequencies lets a model trained at 8k tokens work at much longer lengths with modest additional training.

ALiBi takes a simpler route — add a linear penalty to attention scores proportional to distance — and also extrapolates well.

5. The block

A transformer layer is two sublayers, each wrapped in a residual connection and a normalisation (Chapter 12.4.2):

x = x + Attention(LayerNorm(x))        # (1) pre-norm
x = x + FeedForward(LayerNorm(x))      # (2)

(1) Attention mixes information between positions. (2) The feed-forward network processes each position independently — the same small network applied at every position:

\text{FFN}(x) = W_2 \,\text{GELU}(W_1x + b_1) + b_2

It expands to 4× the model dimension and back, and this is where most of the parameters live — roughly two-thirds of a typical transformer's weights. A useful reading: attention decides what to look at, and the feed-forward layer does the thinking about what was gathered. There is evidence that these layers act as key-value memories storing factual associations, which is why editing a model's facts targets them.

Stack this block 12 times for a small model, 32 for a medium one, 100+ for a large one.

6. Encoder, decoder, and the mask

Encoder-only (BERT-style). Every token attends to every other, in both directions. Trained by masking words and predicting them. Right for understanding tasks — classification, entity extraction, and producing embeddings (Chapter 12.5.2). It cannot generate text left to right.

Decoder-only (GPT-style). Causal masking sets the attention scores for future positions to -\infty before the softmax, so they receive zero weight. Each position can only see itself and what came before. Trained to predict the next token. This is nearly every modern large language model.

Encoder-decoder (T5, translation). An encoder reads the input bidirectionally, and a decoder generates while attending both to its own output and to the encoder's — cross-attention.

Why decoder-only won. Next-token prediction is a completely general objective: summarising, translating, answering and classifying can all be phrased as "continue this text". One architecture, one objective, one scaling story — and simplicity wins when scale is the strategy.

7. The quadratic problem

Attention computes n^2 scores for a sequence of length n. Double the context and the work quadruples. That single fact drives most of the engineering around modern models.

  • 1,000 tokens → 1 million scores.
  • 10,000 tokens → 100 million.
  • 100,000 tokens → 10 billion.

Four responses, and the difference between them matters.

FlashAttentionnot an approximation. It computes exact attention while never materialising the full n \times n matrix in slow memory, processing it in tiles that fit in the GPU's fast on-chip memory. Since attention is memory-bandwidth-bound rather than compute-bound, this is a large speed-up and a large memory saving with identical results. It is the single most important systems contribution to long contexts.

Sparse and sliding-window attention — each token attends only to a nearby window plus a few global tokens. Linear rather than quadratic, and it is an approximation.

The KV cache — during generation, each new token would otherwise recompute keys and values for the whole prefix. Cache them instead, and generating token n+1 costs one step rather than n. This is what makes generation practical, and the cache is the memory constraint in serving: it grows with batch size × sequence length × layers × heads, and at long contexts it exceeds the model weights themselves. Grouped-query attention exists to shrink it.

Linear-attention and state-space models (Mamba and relatives) replace the mechanism with one that scales linearly. They are genuinely promising, competitive on several tasks, and have not displaced transformers at the frontier — the honest 2026 position is that hybrids combining both are the active direction.

8. Why it won

Parallel training. Every position's attention is computed simultaneously, so a whole sequence uses the whole GPU. This is the decisive advantage over recurrence, and it is what made training on internet-scale text affordable.

Constant path length between any two positions. In an RNN, information from token 1 reaches token 500 through 499 sequential updates. In a transformer it is one attention step. Long-range dependencies stop being a distance problem.

It scales predictably. More data, more parameters, more compute produced reliably better models over several orders of magnitude — the scaling laws in Chapter 12.5.3 — which is a rare property and the reason enormous investment followed.

It generalises across modalities. Chop an image into patches, or audio into frames, and treat each as a token: the same architecture handles vision and speech. One architecture for everything is a large practical advantage, and it is why the same infrastructure serves text, image and audio models.

Recall

  • Attention is retrieval: a query is compared against every key, and the matching values are blended. The comparison is a dot product, so Chapter 12.2's alignment measure is doing all the work.
  • \text{softmax}(QK^{\top}/\sqrt{d_k})V. The \sqrt{d_k} keeps score variance near 1 — without it, softmax saturates and the gradient vanishes.
  • Multi-head runs several attentions in smaller dimensions for roughly the cost of one, and heads specialise on their own into syntax, previous-token and bracket-matching roles.
  • Attention is permutation-invariant, so position must be injected. RoPE rotates queries and keys by position, making the dot product depend on relative distance — and is also how context windows get extended.
  • A block is attention (mixes between positions) + feed-forward (processes each position alone), each with a residual and pre-norm. The feed-forward layers hold most of the parameters.
  • Causal masking sets future scores to -\infty, making a decoder generate left to right. Decoder-only won because next-token prediction expresses every task.
  • Attention is O(n^2). FlashAttention is exact and tiles the computation to fit fast memory; sparse and sliding-window attention are approximations; the KV cache makes generation practical and becomes the serving memory constraint, which is why grouped-query attention exists.
  • It won on parallel training, constant path length between any two positions, predictable scaling, and working across text, images and audio with one architecture.

Self-test: What do query, key and value each mean? · Why divide by \sqrt{d_k}? · What does a transformer lose that an RNN gets for free, and how is it restored? · What does causal masking actually change in the computation? · Why is FlashAttention not an approximation? · What grows until it exceeds the model weights during serving?