← The Arithmetic of Intelligence
Part II · The Classical Era · Chapter 6

Attention

Recurrence reaches the past by multiplying through every intervening step. Attention reaches it in one. The consequences — for scale, for cost, and for a famous old trick — are all arithmetic.

A recurrent network reaches a token nine hundred positions back by passing through all nine hundred intervening steps. Attention reaches it in one. That difference sounds like a mere efficiency, but it changes what the network can learn, what it costs, and — as an old preprocessing trick will reveal — what we mean when we say a technique “helps.”

An attention layer compares every position to every other by a dot product, turns those comparisons into weights with a softmax, and mixes the values accordingly. Three quantitative facts govern the design: the dot products must be rescaled by the square root of the dimension, or the softmax saturates; the cost is quadratic in sequence length but crosses over with the feed-forward cost only at a specific length; and the maximum distance information must travel drops from linear to constant. Each falls out of tools already in hand.

6.1 The block, and what it holds

Before pricing attention we have to say what it is. An attention layer does not receive queries and keys; it manufactures them. From the sequence of activations \( X \), with one row of width \( d \) per position, it forms three projections using learned matrices:

\[ Q = XW_Q, \qquad K = XW_K, \qquad V = XW_V, \]

and then compares every position to every other, normalizes the comparisons, and mixes the values accordingly:

\[ \operatorname{Attn}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_{\text{head}}}}\right)V. \]

The softmax is applied along each row, so each position's output is a weighted average of all positions' value vectors, the weights summing to one. Section 2.2 defined it; Section 6.2 explains the divisor.

A layer does not do this once but \( H \) times in parallel, with \( H \) separate triples of projections. These are the heads, and each works in a narrower space of width \( d_{\text{head}} = d/H \). Their outputs are concatenated back to width \( d \) and passed through a fourth learned matrix \( W_O \). The point of splitting is that different heads can attend to different things; the price of splitting is nothing at all, and that is worth seeing rather than being told.

Count the projections. Each head's \( W_Q, W_K, W_V \) is \( d \times d_{\text{head}} \), and there are \( H \) of them, so the three together hold \( 3 \times H \times d \times (d/H) = 3d^2 \) weights — the \( H \) cancels. Add \( W_O \) at \( d \times d \) and attention holds \( 4d^2 \) weights per layer, whatever \( H \) is. Sixteen heads and one head cost the same. Splitting is free because it partitions a fixed width rather than adding to it, and forgetting this is how people come to believe that more heads means a bigger model.

The rest of the block is the position-wise feed-forward network of Section 2.1, widening \( d \) to \( 4d \) and back, which Worked Example 1.7 already counted at \( 8d^2 \). Each of the two parts is wrapped in a residual connection and a normalization layer — Chapter 4's construction, applied here — which together add only \( O(d) \) parameters and are negligible in any count. So a transformer layer holds

\[ \underbrace{4d^2}_{\text{attention}} + \underbrace{8d^2}_{\text{feed-forward}} = 12d^2 \ \text{weights,} \]

and by Chapter 1's corollary it costs about that many multiply–accumulates per position. One number, and it prices every transformer in the rest of this book.

Worked Example 6.1

Seven billion, counted from four numbers

Part III costs everything against a model described only as “7 B parameters, 32 layers, 32 heads of dimension 128”. Check that those numbers are consistent, and recover the vocabulary's share.

The head dimension fixes the width: \( d = H\,d_{\text{head}} = 32 \times 128 = 4096 \). Then each layer holds \( 12d^2 = 12 \times 4096^2 = 201{,}326{,}592 \) weights, and thirty-two layers hold \( 32 \times 201{,}326{,}592 = 6.44 \times 10^9 \).

The tables are extra. At a vocabulary of \( 32{,}000 \) with untied embedding and readout, they add \( 2 \times 32{,}000 \times 4096 = 0.26 \times 10^9 \), for a total of \( 6.70 \times 10^9 \) — the seven billion, to two figures, from four numbers and one formula.

Notice what the tables are here: four percent of the model, against the forty-five percent they occupied in the recurrent model of Worked Example 5.1. Nothing about vocabularies changed. The layers grew.

Drill 6.1

(a) A model has \( d = 5120 \) and 40 heads. Give \( d_{\text{head}} \), and the parameters in one layer.   (b) The same model has 40 layers. Give the non-embedding parameter count.   (c) Its designers double the head count to 80, holding \( d \) fixed. By how much does the parameter count change?

Show answers

(a) \( d_{\text{head}} = 5120/40 = 128 \); one layer holds \( 12 \times 5120^2 = 314{,}572{,}800 \), about \( 315 \) M.   (b) \( 40 \times 3.15\times10^8 = 1.26\times10^{10} \), about \( 12.6 \) B.   (c) Not at all. The heads partition \( d \); doubling \( H \) halves \( d_{\text{head}} \) and the \( H \) cancels out of \( 3d^2 \) exactly as before.

6.2 Why divide by the square root of the dimension

An attention score is a dot product of a query vector and a key vector, each of dimension \( d_{\text{head}} \) — written \( d \) throughout this section, since only the dimension of the vectors being compared matters to the argument. If their components are independent with unit variance, Chapter 1 already told us the variance of that dot product.

Derivation 6.1

The score variance, and the scale that tames it

For \( q, k \in \mathbb{R}^d \) with independent, zero-mean, unit-variance components, the score is \( q \cdot k = \sum_{i=1}^{d} q_i k_i \). Each term is a product of independent unit-variance variables, hence has variance 1; the \( d \) terms are independent, so their variances add:

\[ \operatorname{Var}(q \cdot k) = d, \qquad \text{standard deviation} = \sqrt{d}. \]

For \( d = 512 \), the raw scores have standard deviation \( \sqrt{512} = 16\sqrt{2} \approx 22.6 \). Feeding logits of that magnitude into a softmax pushes almost all its mass onto a single entry: the output is nearly one-hot, and its gradient — which is largest when the distribution is spread out — collapses toward zero. This is the sigmoid saturation of the recurrent era reborn inside attention, and it is solved the same way: keep the pre-activation at unit scale. Dividing each score by \( \sqrt{d} \) restores the variance to 1 and keeps the softmax responsive.

The lesson generalizes a theme running through the book. In Chapter 3 the fix for exploding activations was a \( 1/\sqrt{m} \) on the weights; here the fix for saturating scores is a \( 1/\sqrt{d} \) on the logits. Both are the same instinct — hold the variance at one as signal crosses an operation — and both are one line of variance algebra.

Drill 6.2

Attention operates at dimension \( d = 128 \). (a) What is the standard deviation of an unscaled score? (b) By what factor must scores be divided to restore unit variance?

Show answers

(a) \( \sqrt{128} = 8\sqrt{2} \approx 11.3 \). (b) Divide by \( \sqrt{128} \approx 11.3 \); this restores variance to 1.

6.3 What attention costs

A transformer layer has two expensive parts: the attention itself, and the position-wise feed-forward network that follows it. Both costs come from the matrix-multiply primitive of Chapter 1 — multiplying an \( a\times b \) matrix by a \( b\times c \) matrix costs \( abc \) multiply–accumulates — and comparing them tells you which regime a model lives in.

Derivation 6.2

The attention–feedforward crossover

For a sequence of length \( n \) and model width \( d \): attention forms an \( n\times n \) score matrix from \( n\times d \) queries and keys (\( n^2 d \) operations), then mixes an \( n\times n \) weight matrix with \( n\times d \) values (another \( n^2 d \)), totaling \( 2n^2 d \). The feed-forward network maps width \( d \) up to \( 4d \) and back, at each of \( n \) positions: \( n \cdot d \cdot 4d \) up and the same down, totaling \( 8nd^2 \). Setting them equal,

\[ 2n^2 d = 8nd^2 \quad\Longrightarrow\quad n = 4d. \]

Below \( n = 4d \), the feed-forward network dominates and attention is a minor cost. Only past that crossover does the quadratic term take over. This is why quadratic attention was tolerable for the sentence-length sequences of early machine translation and became the binding constraint only when context windows grew to thousands of tokens — a story Chapter 11 resumes in the modern setting.

The other cost that matters is not arithmetic but distance. In a recurrence, information from position \( i \) reaches position \( j \) only by passing through all \( |i-j| \) steps between them: the path length is linear in the gap, and it is along that long path that gradients decay. In attention, any position reaches any other in a single operation: the maximum path length is constant, independent of the gap. A token nine hundred positions back is no harder to consult than its neighbor. That collapse from linear to constant path length — not the arithmetic — is the deepest reason attention learns long-range structure where recurrence struggled.

Drill 6.3

A model has width \( d = 512 \) and feed-forward width \( 2048 \). At what sequence length do attention and feed-forward costs balance? Below that length, which dominates?

Show answer

\( n = 4d = 2048 \). Below \( n = 2048 \), the feed-forward network dominates.

6.4 Temperature: reshaping confidence without reordering it

When a model samples its next token, it divides the logits by a temperature \( T \) before the softmax. This is worth a moment because a simple monotonicity argument settles what temperature can and cannot do. Dividing by a positive \( T \) is a monotone transformation of the logits: it can sharpen the distribution (small \( T \), toward the argmax) or flatten it (large \( T \), toward uniform), but it can never change which logit is largest. Temperature reshapes confidence; it never reorders the candidates.

Worked Example 6.2

Two temperatures on a two-token vocabulary

Logits are \( z_1 = 2, z_2 = 0 \). At \( T = 1 \), the probability of token 1 is \( \sigma(2) = 1/(1+e^{-2}) = 1/(1.135) = 0.881 \). At \( T = 0.5 \), the logits double to \( 4, 0 \), so \( \sigma(4) = 1/(1+e^{-4}) = 1/(1.018) = 0.982 \): colder sampling has sharpened the lead. As \( T \to 0 \) the distribution tends to the argmax; as \( T \to \infty \) it tends to uniform. At no finite positive \( T \) does token 2 overtake token 1 — monotonicity forbids it.

Worked Example 6.3

A perplexity drop that sounds enormous and is half a bit

A language model improves from perplexity 114.5 to perplexity 78.4 — a 31% relative reduction. How large is that gain in bits per token?

Perplexity is \( 2^H \) where \( H \) is cross-entropy in bits, so \( H = \log_2(\text{perplexity}) \). Before: \( \log_2(114.5) = 6.84 \) bits/token. After: \( \log_2(78.4) = 6.29 \) bits/token. The improvement is \( 6.84 - 6.29 = 0.55 \) bits per token — barely half a bit.

The lesson: perplexity is exponential in bits. A gain that sounds dramatic in the exponential unit can be modest in the linear unit that file sizes, bandwidth budgets, and the compression arguments of Chapter 10 are denominated in. Whenever someone reports a perplexity improvement, convert it to bits before deciding whether to be impressed.

6.5 The reversal trick, and what it reveals

Before attention, sequence-to-sequence models used a recurrent encoder to read the source and a recurrent decoder to write the target. A curious preprocessing trick helped: reversing the source sequence before encoding. Why should reading a sentence backward help translate it? The answer is one of the most instructive small calculations in the book, because it distinguishes two things students routinely conflate — optimization and representation.

Work out what you think reversal should do before reading on. The natural expectation is that it shortens something — that reading the sentence backwards must somehow bring related words closer together, since otherwise why would it help? Write down what you would predict for the average distance between a source word and its translation. Then compare it with what follows. The gap between the two is the lesson.

Derivation 6.3

The lag of the reversal trick, and its unchanged mean

Suppose source and target both have length \( n \), with a strictly monotonic alignment: target position \( j \) corresponds to source position \( j \). Define the lag of position \( j \) as the number of recurrent steps between reading source token \( j \) and emitting target token \( j \).

Unreversed. Source token \( j \) is read at encoder step \( j \); target token \( j \) is emitted at step \( n + j \) (after the whole source is read). The lag is \( (n+j) - j = n \), the same for every position.

Reversed. The source is fed backward, so source token \( j \) is now read at step \( n - j + 1 \). Target token \( j \) is still emitted at step \( n + j \). The lag is \( (n+j) - (n-j+1) = 2j - 1 \).

Now compute the mean lag over all positions:

\[ \frac{1}{n}\sum_{j=1}^{n} (2j - 1) = \frac{1}{n}\cdot n^2 = n. \]

The mean lag is exactly unchanged at \( n \). And the maximum lag has grown, from \( n \) to \( 2n - 1 \). By every aggregate measure of distance, reversal made things no better — and in the worst case, worse.

So why does it help? Because it changes the distribution of lags without changing their mean. Unreversed, every alignment is equally distant, \( n \) steps away; the network has no easy foothold anywhere. Reversed, the early alignments become very short — the first target position now lags its source by just one step — at the cost of longer lags later. Gradient descent can seize those early, short-lag alignments immediately, and once the beginning of the sentence is aligned, the rest follows. The trick works because optimization needs an early handhold, not because the representational problem got any smaller.

Trap · T3

The reversal trick is the canonical mean-versus-maximum confusion. It is natural to assume a trick that helps must reduce the average distance information travels — but here the mean is provably unchanged and the maximum is worse. Whenever a technique helps without improving the aggregate you first reach for, suspect that it is reshaping a distribution to aid optimization, not shrinking a quantity to aid representation. Attention, incidentally, removes the trade-off entirely by making every lag constant.

6.6 What the chapter bought

Attention's \( \sqrt{d} \) scaling is the same variance-control instinct as sensible initialization, applied to logits. Its cost is quadratic in length but only overtakes the feed-forward network past \( n = 4d \), which is why quadratic attention was affordable for years. Its constant path length is the true source of its reach. And the reversal trick, properly analyzed, is a lesson in reading a technique correctly: some interventions change what a network can represent, and others merely change how easily gradient descent finds a representation — and telling the two apart is a skill the rest of the book will keep demanding. Chapter 7 takes up what becomes, once architecture stops being the bottleneck, the field's central question: given a fixed budget, how much model and how much data — answered by the derivation that turned compute-optimal training from a matter of taste into a calculation.

Exercises

A · Drills

  1. Attention runs at \( d = 256 \). Give the standard deviation of an unscaled score and the required scaling factor.
  2. A model has width \( d = 768 \). Give the sequence length at which attention and feed-forward costs balance.
  3. Logits are \( 3, 1 \). Give the probability of the first token at temperature \( T = 1 \) and \( T = 0.5 \).
  4. For a source of length \( n = 6 \) with the reversal trick, give the lag of positions 1 and 6.

B · Problems

  1. Which regime. A model has width \( d = 1024 \) and is run at sequence length \( n = 8192 \). Compute the attention and feed-forward MAC counts (in units of \( d^2 \) and \( n \)), give their ratio, and state which dominates. Then find the sequence length at which they would balance.
  2. The reversal trick in full. For a monotonic length-\( n \) alignment, derive the reversed lag \( 2j-1 \) from scratch, prove the mean lag is exactly \( n \), and give the maximum lag. Then explain in two sentences why the trick aids optimization despite leaving the mean unchanged.
  3. Temperature cannot reorder. Prove that for any positive temperature \( T \) and any two logits \( z_a > z_b \), the softmax probability of \( a \) exceeds that of \( b \). Conclude that sampling temperature reshapes confidence but never changes the ranking of candidates.

C · Challenge

  1. One instinct, three chapters. The \( 1/\sqrt{m} \) of Chapter 3, the normalization-inside-the-branch of Chapter 4, and the \( 1/\sqrt{d} \) of this chapter are all the same principle. State that principle in one sentence, then show how each of the three is an instance of it, identifying in each case what quantity is being held at what value and what would go wrong otherwise.
Gate 6 · Pass before Chapter 7

Reproduce Derivations 5.1 (the \( \sqrt{d} \) scaling), 5.2 (the \( n = 4d \) crossover), and 5.3 (the reversal lag) on blank paper. Attempt B-1 and B-2 closed book. You pass when the crossover and the mean-lag proof come without hesitation — the latter is the one most often botched, because the sum \( \sum(2j-1) = n^2 \) must be seen instantly.

Readings for Chapter 6