7  Transformers and Attention

Learning objectives. After this chapter you can:

7.1 Self-attention from first principles

By Chapter 6 we could predict the next word, but our recurrent models had a bottleneck: everything the model knew about the sentence had to be squeezed through a fixed-size state, one step at a time. That was literal, not metaphorical — a sequence-to-sequence translator encoded the entire source sentence into a single vector before emitting a word of the target (Sutskever et al. 2014). Long-range relationships — a pronoun and the noun it refers to, a verb and its distant subject — kept getting lost, and longer sentences translated worse.

Attention arose as an answer to this bottleneck. Bahdanau et al. (2015) let the decoder look back at all of the encoder’s states and learn which ones mattered at each output step, and on the translation tasks they studied the length-related collapse largely went away. Note the date: this is 2015, and there is still a recurrent network underneath.

The Transformer’s contribution two years later (Vaswani et al. 2017) was to notice that this mechanism was doing most of the work, and to drop the recurrent network altogether; the rest of the block keeps a feed-forward layer, residual connections and normalisation. Its answer is strikingly direct: let every token look at every token it is allowed to see, and learn how much to take from each. If you have ever computed a weighted average, you already own the core intuition. Attention is a weighted average of token vectors where the weights themselves are predicted from the data.

Each token’s vector is projected three ways — a query (“what am I looking for?”), a key (“what do I contain?”), and a value (“what do I pass along if selected?”). The score between token \(i\) and token \(j\) is the dot product \(q_i \cdot k_j\), scaled and pushed through a softmax so each row becomes a probability distribution:

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

That softmax row is exactly the set of weights in the weighted average. A statistician’s translation: for every token, the model takes a data-dependent weighted average over the whole sentence — closer to Nadaraya–Watson kernel regression than to fitting coefficients — deciding which other tokens are informative right now. Those weights are the mixing coefficients of the value vectors; on their own they do not answer “which word matters?”, because the value transform, the output projection and the later layers shape the result too.

The shape of the whole model

Before opening anything up, it is worth seeing where attention actually sits. A transformer is a short pipeline with one part repeated:

Where attention sits in the whole model
TOKENS kedi sütü içti Embedding table lookup n × d + Position added here, or inside every block × L BLOCKS Transformer block Multi-head attention Feed- forward n × d Final LayerNorm Unembedding to vocab n × |V| logits p(next) The token count n and the width d are identical at every arrow: a block rewrites the vectors, it never reshapes them. This chapter opens the navy box up, one part at a time.

The whole architecture. Most of what makes a transformer a transformer is inside the repeated block; the pieces on either side of it are a lookup table and a matrix multiply. But what the model can do also depends on how position is supplied and on the mask that decides which token may see which; both come later.

Two properties of that picture are worth fixing in mind now, because everything later leans on them.

The shape never changes. \(n\) tokens enter a block as an \(n \times d\) matrix and leave as an \(n \times d\) matrix. A block is a rewriting step, not a reshaping step, which is exactly why stacking is trivial: block 12 presents the same interface as block 1, so depth is a free parameter.

Only one component ever looks sideways. The feed-forward sublayer, the normalisations, and the unembedding all treat each token in isolation. Attention is the only place in the entire model where position \(i\) can read position \(j\). Every claim in this book about a model “connecting” two distant words cashes out, mechanically, in the numbers of attention matrices — usually not one number but several heads and layers, together with the value transforms and the residual path.

Why three projections, and not one?

Here is a fair objection to the equation above. If attention needs to know which tokens are related, why build \(Q\), \(K\) and \(V\) at all? Why not compare the token vectors directly, \(x_i \cdot x_j\), and save three matrix multiplies?

It is worth taking seriously, because trying it shows the logic of the design. Using the four-dimensional embeddings from the worked example below — kedi = [0.9, 0.1, 0.0, 0.5], sütü = [0.1, 0.9, 0.0, 0.5], içti = [0.0, 0.1, 0.9, 0.5] — the raw dot products are:

\[ XX^{\top} = \begin{array}{r|rrr} & \texttt{kedi} & \texttt{sütü} & \texttt{içti}\\\hline \texttt{kedi} & \mathbf{1.07} & 0.43 & 0.26\\ \texttt{sütü} & 0.43 & \mathbf{1.07} & 0.34\\ \texttt{içti} & 0.26 & 0.34 & \mathbf{1.07} \end{array} \]

Two limitations show up in this example, and each one seriously weakens it.

Similarity is symmetric. \(x_i \cdot x_j = x_j \cdot x_i\), necessarily. So “the verb wants its object” and “the object wants its verb” are forced to be the same score. Be precise about what that does and does not force: the attention weights can still differ, because softmax normalises each row by its own denominator, so \(\alpha_{ij} \neq \alpha_{ji}\) in general. What is lost is control over the score table itself: no head can give verb-wants-object a higher score than object-wants-verb, because the two are the same entry. Row normalisation still separates the weights — with \(d_k = 1\) and \(Q = K = [1, 2]^{\top}\) the two off-diagonal weights come out 0.731 and 0.119 — so directional attention does not disappear. What disappears is the freedom to choose it: the head can only express the directions that row-normalising a symmetric table happens to produce. Grammar is full of relations that run one way and not the other, and separate projections are how you get to score them that way.

Self-similarity, in this example, swamps everything else. A token’s score against itself is \(x \cdot x\), its squared length, and here that is the largest entry in every row. Push the matrix through a softmax and içti puts its largest weight, 0.4235, on itself, and only 0.2825 on kedi. The verb reaches for nothing. The layer computes an expensive approximation to leave every token alone.

Strictly, \(x_i \cdot x_i \ge x_i \cdot x_j\) is guaranteed whenever \(\|x_j\| \le \|x_i\|\), by Cauchy–Schwarz. The converse does not hold — with \(x_i = (1,0)\) and \(x_j = (0,2)\) the neighbour is longer and the diagonal still wins, 1 against 0 — so a longer neighbour can outscore you but need not. Our three toy vectors all have \(\|x\|^2 = 1.07\), so here the diagonal wins outright. The general claim is weaker but still damaging: an unprojected score inflates self-similarity by construction, whether or not the two tokens are related.

Now the same sentence with the projections restored, which is the head you will build in a moment. Both columns are softmax rows, each scaled by the square root of its own width — \(\sqrt{4}\) on the left, \(\sqrt{2}\) on the right — so they are comparable:

içti attends to no projections with \(W^Q\), \(W^K\)
kedi (its agent) 0.2825 0.4230
sütü (its object) 0.2940 0.5012
içti (itself) 0.4235 0.0759

The verb now spends 92.41% of its attention on its two arguments instead of 42.35% on itself. And the score matrix is no longer symmetric: içti scores sütü at 2.97, while sütü scores içti at 0.27 — an 11× asymmetry, which is the whole point. No unprojected score could have produced it.

So the three matrices are not decoration. They exist to separate jobs that a single vector cannot hold at once:

  • \(W^Q\) builds the query — what am I looking for? Read from the token doing the looking.
  • \(W^K\) builds the key — what do I advertise to others? Read from the token being looked at.
  • \(W^V\) builds the value — what do I hand over if I am chosen?

The division of labour is the design. \(Q\) and \(K\) together decide where information moves; \(V\) supplies what moves. (The head’s output is then transformed by \(W^O\) and added to the residual stream; \(V\) alone does not settle the final effect.) Keeping them apart lets a token advertise one thing and deliver another — and lets two tokens that are not remotely similar still be strongly connected, which is what a verb and its subject need.

Two matrices, or one?

This subsection is a deeper aside. On a first reading you can skip to the next figure and come back after Section 7.2.

Be careful about what that argument proves. Write the score out and the projections collapse:

\[ q_i \cdot k_j \;=\; (x_i W^Q)(x_j W^K)^{\top} \;=\; x_i\,\underbrace{W^Q W^{K\top}}_{M,\;\; d \times d}\, x_j^{\top} \]

Only the product \(M\) ever touches the score. A single learned \(d \times d\) matrix would give the same scores — still asymmetric, still free of the self-similarity trap, because nothing forces \(M\) to be symmetric. So the honest claim is not “two matrices are necessary”; it is two matrices are a choice you make once you also want the model to be cheap. Three concrete reasons:

  • Parameters. For a 7B-class model (Llama 2 7B: \(d = 4096\), 32 heads of \(d_k = 128\)): one head’s \(W^Q, W^K\) pair costs \(2 \times 4096 \times 128 = 1{,}048{,}576\) numbers, while a full \(M\) costs \(4096^2 = 16{,}777{,}216\) — 16× more. The ratio is the same for a whole layer, since you would need one \(M\) per head: \(32\) heads cost \(2d^2 = 33.5\)M factored against \(32d^2 = 537\)M unfactored. The factorisation is a deliberate low-rank constraint: \(M\) is forced to rank \(\le d_k\). A full \(M\) would offer a wider family of scores; factoring gives some of that expressiveness up in exchange for cost.
  • Arithmetic. With \(Q\) and \(K\) already computed, pairwise scoring costs \(O(n^2 d_k)\). Scoring through \(M\) costs \(O(n^2 d)\) — you would pay \(d/d_k = 32\times\) more for the same numbers. (The projections themselves come on top: \(O(n d\,d_k)\) for \(Q\) and \(K\), \(O(n d^2)\) for \(XM\).)
  • Cache width. This is the one that decides the design. With \(M\) you can still cache — you just store \(x_j\) itself, since \(x_i^{\top} M x_j = (M^{\top} x_i) \cdot x_j\) — but what you store is stuck at the full width \(d\). Factoring puts keys in a \(d_k\)-wide space of their own, and once they live there you can shrink that space: letting each key head serve four query heads (8 key heads instead of 32) takes the keys in the cache from \(4096\) numbers per token per layer to \(1024\). That trick is called GQA, it is the subject of
    1. A separate \(W^K\) makes that narrow, shared cache natural and efficient; a single full \(M\) gives no direct way to build the same layout.

The same non-identifiability applies on the value path: a head’s output enters the residual stream as \(x W^V W^O_{\text{head}}\), another product of two matrices that only ever appears together. Factoring is an engineering decision about rank and cost, taken twice, for the same reasons.

The mechanism this produces is a soft version of something you already know:

A dictionary looks up one key; attention looks up all of them
A DICTIONARY — EXACT MATCH, ONE VALUE KEY VALUE "sütü" kedi sütü içti cat milk drank milk Two keys out of three are ignored completely. Miss the key and there is no answer at all. ATTENTION — PARTIAL MATCH, BLENDED VALUE KEY WEIGHT VALUE q(içti) k(kedi) k(sütü) k(içti) 0.4230 0.5012 0.0759 v(kedi) v(sütü) v(içti) Every key is consulted, and every value contributes something. Nothing is ignored; it is only weighted. Same three parts — query, key, value — but the match is graded instead of exact, so the answer is a blend of every value. That is the one change, and it is what makes attention differentiable: you cannot take a gradient through "which key matched", but you can take one through a weight of 0.5012.

Attention is a differentiable dictionary. The weights on the right are the real ones for içti from the head built in Section 7.2.

That last point deserves its own sentence, because it explains why the softmax is there at all rather than some simpler normalisation. A hard lookup — take the single best-matching key — is a step function, and step functions have zero gradient almost everywhere, so with ordinary backpropagation nothing upstream could learn (there are other ways to train discrete choices, all more roundabout). Softmax is the price of being able to train the thing: it keeps the “pick the best match” behaviour in spirit, while leaving a usable derivative with respect to every score.

Explore: what is a head looking at?

Trained models are known to have heads with different, sometimes interpretable, patterns. The widget below is not taken from a trained model: it is a stylised attention matrix for one sentence, built by hand to compare three possible patterns. The matrix is unmasked (bidirectional): in the pronoun head the it row gives a small weight to ready, which comes after it; in a GPT-style causal model that cell would be zero (Section 7.5). Each row is a token doing the looking; each column is a token being looked at; darker means more weight. Switch heads and watch the pattern change — especially the row for it.

Attention head explorer — “The robot poured coffee because it was ready.”

The ambiguity in the pronoun head is real, and it is the classic Winograd setup. Only the ending changes between these two sentences, and the referent of it moves from one noun to the other:

  • “The robot poured coffee because it was ready.” → it leans to the coffee.
  • “The robot poured coffee because it was instructed to.” → it is the robot.

The first of those is a lean and not a certainty — a robot can be ready too, which is why the head above splits 0.46/0.34 rather than committing. The canonical version of the schema removes that slack: the trophy does not fit in the suitcase because it is too big (the trophy) against too small (the suitcase) (Levesque et al. 2012). Either way the point is the same: nothing in the syntax distinguishes the pair; the grammar is identical. The expected reading is decided by world knowledge, which is one of the things large-scale pretraining tries to buy (in real text even that does not always settle a single reading).

What we will look at. Attention first, derived rather than announced, then run once by hand on numbers small enough to check with a pencil. Then the block it sits inside, and why one head is never enough. Then the three things the equation leaves out and every real model has to add: position, a choice of which tokens may see which, and the arithmetic that decides whether you can afford to run it at all.

7.2 One layer, three tokens, every number

That equation is four symbols wide and hides five separate steps. The fastest way to stop being afraid of it is to run it once, completely, on numbers small enough to check with a pencil.

The sentence is Turkish: kedi sütü içti, “the cat drank the milk”. It is chosen because Turkish puts the verb last, so içti (“drank”) has to reach backwards for both of the things it needs — who did the drinking, and what got drunk. That makes the pattern we are looking for easy to state in advance: if attention is doing anything sensible, the row for içti should lean on the other two tokens and not on itself. We will pick the matrices to make exactly that happen: the arithmetic below shows how a head can represent the pattern, not that a model learns the subject and object relation from training.

Everything is deliberately tiny. Three tokens, a model width of \(d_{\text{model}} = 4\), and a head size of \(d_k = d_v = 2\).

These matrices are set by hand, not learned. They are chosen so the arithmetic stays legible and so each dimension means something you can name. A trained model finds matrices that are wider, denser, and not interpretable one column at a time. What transfers from this example is the procedure and the shapes — never these particular values.

We give the four embedding dimensions meanings, so that “why is this number big?” always has an answer:

token meaning animate liquid action always-on
kedi cat 0.9 0.1 0.0 0.5
sütü milk (object) 0.1 0.9 0.0 0.5
içti drank 0.0 0.1 0.9 0.5

The fourth dimension is the same for all three tokens. Watch what happens to it: because it never varies, it cannot help one token advertise anything that distinguishes it, and it will drop out of every comparison below. We will come back at the end of the section and ask whether “a constant feature is worth nothing” is really true.

Step 1 — project each token three ways

The same input vector is multiplied by three different learned matrices to produce a query, a key, and a value. This is the only place this head’s learned parameters enter the attention calculation (the rest of the block has its own: the output projection, the normalisation, the feed-forward layer).

\[ Q = XW^Q, \qquad K = XW^K, \qquad V = XW^V \]

Keep the shapes in view. \(X\) has one row per token, so it is \(n \times d_{\text{model}}\) (here \(3 \times 4\)). \(W^Q\) and \(W^K\) are \(d_{\text{model}} \times d_k\) and \(W^V\) is \(d_{\text{model}} \times d_v\); in general \(d_v\) can differ from \(d_k\). Here both are 2 (so \(4 \times 2\)), and \(Q\), \(K\) and \(V\) are \(n \times 2\): still one row per token, now only two numbers wide.

The three matrices encode three different questions. \(W^Q\) builds what am I looking for, \(W^K\) builds what do I advertise, and \(W^V\) builds what do I hand over if chosen:

\[ W^Q=\begin{bmatrix}2&0\\0&2\\2&2\\0&0\end{bmatrix}\; W^K=\begin{bmatrix}1.5&0\\0&1.5\\0&0\\0&0\end{bmatrix}\; W^V=\begin{bmatrix}1&0\\0&1\\0&0\\0.2&0.2\end{bmatrix} \]

Read \(W^Q\) one row at a time. The action row is [2, 2], so anything verb-like issues a query that looks for an agent and an object. In \(W^K\) the action row is [0, 0]: a verb advertises neither, so nothing will want to attend to it. Multiplying through:

token query \(q_i\) key \(k_i\) value \(v_i\)
kedi [1.80, 0.20] [1.35, 0.15] [1.00, 0.20]
sütü [0.20, 1.80] [0.15, 1.35] [0.20, 1.00]
içti [1.80, 2.00] [0.00, 0.15] [0.10, 0.20]

Two rows are worth pausing on. içti has the only query that is large in both coordinates — it is shopping for two things at once. And içti has a value of [0.10, 0.20], which is nearly nothing: the verb carries no agent-content and no object-content of its own. Hold on to that; it is the punchline.

Step 2 — every query meets every key

Score token \(i\) against token \(j\) with a dot product, \(q_i \cdot k_j\). Nine token pairs, nine numbers, arranged so that each row is one token doing the looking and each column is a token being looked at:

\[ QK^{\top} = \begin{array}{r|rrr} & \texttt{kedi} & \texttt{sütü} & \texttt{içti}\\\hline \texttt{kedi} & 2.46 & 0.54 & 0.03\\ \texttt{sütü} & 0.54 & 2.46 & 0.27\\ \texttt{içti} & 2.73 & 2.97 & 0.30 \end{array} \]

The dot product is large when two vectors point the same way, so these numbers are similarity scores and nothing more exotic. The largest entry in the table is \(2.97\), and it is içti looking at sütü: the hand-picked projections give the verb–object pair the highest score. The smallest entry in the içti row, \(0.30\), is the verb looking at itself — exactly as designed, since verbs advertise nothing.

Step 3 — divide by the square root of the head size

Why divide at all? Because dot products grow with dimension, and a softmax handed large numbers becomes almost one-hot — one token takes nearly all the weight and the gradient through the rest becomes very small. If the components of \(q\) and \(k\) are independent with mean zero and unit variance, dividing makes the scale of the scores independent of the head width and reduces that saturation; a trained head’s vectors need not satisfy that assumption, so the division does not prevent saturation in every case.

\[ \frac{QK^{\top}}{\sqrt{d_k}} = \frac{QK^{\top}}{\sqrt{2}} = \frac{QK^{\top}}{1.4142} \]

\[ \begin{array}{r|rrr} & \texttt{kedi} & \texttt{sütü} & \texttt{içti}\\\hline \texttt{kedi} & 1.7395 & 0.3818 & 0.0212\\ \texttt{sütü} & 0.3818 & 1.7395 & 0.1909\\ \texttt{içti} & 1.9304 & 2.1001 & 0.2121 \end{array} \]

This step looks like fussy bookkeeping and is not. A dot product of two random \(d_k\)-dimensional vectors (independent components with zero mean and unit variance) has a standard deviation that grows like \(\sqrt{d_k}\), so without the division the scores get systematically larger as heads get wider — and softmax reacts to absolute size. Measured over 20,000 random pairs with independent standard-normal components (not a trained head; this tests the assumption itself):

\(d_k\) std. dev. of \(q\cdot k\) \(\sqrt{d_k}\) after dividing by \(\sqrt{d_k}\)
2 1.414 1.414 1.000
8 2.859 2.828 1.011
64 8.029 8.000 1.004
512 22.746 22.627 1.005

The middle two columns agree to within about 1% — 1.10% at \(d_k = 8\) and under 0.6% for the wider heads. The residual is sampling noise from a 20,000-pair estimate, not a systematic gap. That is the whole point: the measured spread of raw scores really does grow as \(\sqrt{d_k}\), so dividing by \(\sqrt{d_k}\) really does flatten it to a constant, which the last column confirms across a 256× range of widths. The consequence shows up when 64 random keys compete for one query — here is the average weight the winning key receives:

\(d_k\) with \(1/\sqrt{d_k}\) without it
2 0.104 0.165
8 0.105 0.420
64 0.108 0.786
512 0.107 0.928

Without the scaling, a 512-dimensional head hands 93% of its weight to a single token — before it has learned anything. That is a softmax in saturation, and a saturated softmax has very small derivatives with respect to its scores, so teaching the head where to look becomes very hard. With the scaling, the competition stays at roughly the same temperature no matter how wide the head is. The scaled column barely moves: 0.104, 0.105, 0.108, 0.107.

Step 4 — softmax turns scores into weights

Softmax exponentiates each row and divides by the row’s total, which forces every row to be non-negative and to sum to one. Spelled out completely for the içti row:

\[ \begin{aligned} \text{scaled scores} &= [\,1.9304,\; 2.1001,\; 0.2121\,]\\ \exp(\cdot) &= [\,6.8923,\; 8.1670,\; 1.2363\,]\\ \textstyle\sum \exp(\cdot) &= 16.2956\\ \text{weights} &= [\,6.8923,\; 8.1670,\; 1.2363\,]/16.2956\\ &= [\,0.4230,\; 0.5012,\; 0.0759\,] \end{aligned} \]

Doing that to all three rows gives the attention matrix:

\[ A = \operatorname{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right) = \begin{array}{r|rrr} & \texttt{kedi} & \texttt{sütü} & \texttt{içti}\\\hline \texttt{kedi} & \mathbf{0.6961} & 0.1791 & 0.1249\\ \texttt{sütü} & 0.1750 & \mathbf{0.6804} & 0.1446\\ \texttt{içti} & 0.4230 & \mathbf{0.5012} & 0.0759 \end{array} \]

Every row sums to 1.000. The two content words spend about 70% of their attention on themselves. The verb does something different: it splits 42% / 50% across its subject and its object and keeps only 7.6% for itself. That is the pattern we aimed for when we picked the matrices; it does not show that nouns always attend to themselves in a trained model.

Step 5 — the weighted average

Finally, multiply the weights by the values. Nothing new happens here; this is an ordinary weighted average, with the weights supplied by step 4.

\[ \operatorname{Attention}(Q,K,V) = AV \]

The shapes close the loop: \(A\) is \(n \times n\) (one row of weights per token), \(V\) is \(n \times d_v\), so \(AV\) is \(n \times d_v\) — one output vector per token, the same number of rows we started with.

token value before attention output after attention
kedi [1.00, 0.20] [0.7444, 0.3433]
sütü [0.20, 1.00] [0.3256, 0.7443]
içti [0.10, 0.20] [0.5308, 0.6009]

Look at the last row, because it is the reason the whole mechanism exists. içti arrived carrying [0.10, 0.20] — almost no agent-content, almost no object-content. It leaves carrying [0.5308, 0.6009]. The verb has pulled the content of both of its arguments into its own position: it asked, and the two tokens with matching keys answered. In this example we built that match by picking the matrices; a trained model does the same only if it learns it from data. The output is then transformed by \(W^O\) and added to the residual stream; what the next layer sees at position 3 is that sum.

The two nouns moved less, and in the direction you would expect: kedi went from [1.00, 0.20] to [0.74, 0.34], keeping most of its own identity while picking up a little of the object it is now associated with.

Every step above is either a matrix multiply or a softmax. There is no loop over positions, no recurrence, and nothing that has to happen in order — which is exactly why this replaced the recurrent models of Chapter 6. All nine scores can be computed at the same time, on hardware that is very good at computing many things at the same time. That parallelism holds for a text whose positions are known: in training and when reading a prompt. When generating text, each new token still waits for the one before.

Do it yourself

The widget runs the arithmetic above live. Step through the pipeline one stage at a time; the two switches let you break the mechanism on purpose and watch what happens.

One attention layer, stage by stage — “kedi sütü içti”

Two things to try. Turn causal mask on and watch the top-right of the matrix go to zero — kedi can then only see itself, which is the constraint that makes a model able to generate text left to right (Section 7.5). Turn divide by √dₖ off and watch every weight drift toward the extremes. At \(d_k=2\) the effect is mild; it becomes serious only for wide heads, which this widget does not show — see the random-vector experiment in step 3 for that — and which is why Vaswani et al. (2017), working with \(d_k = 64\), added the scaling in the first place.

The same picture, geometrically

Because we chose \(d_k = 2\), every query and key is a point on a plane and the dot product becomes something you can see. Since \(q \cdot k = \|q\|\,\|k\|\cos\theta\), it depends on both angle and length: a query scores highly against a key when the two arrows point in a similar direction and both are long; a key pointing the same way but short can stay low.

Queries and keys live in the same 2-D space

Select içti. Its query (red) points diagonally, up and to the right, so it has a large dot product with both kedi’s key and sütü’s key. Select kedi and the query swings towards the horizontal axis, lining up with its own key and nothing else. The attention pattern is not a lookup table — it is this geometry, repeated for every token. (We can name the axes “agent-ness” and “object-ness” only because we picked the matrices by hand; the axes of trained heads usually do not split into one-word readings like that.)

7.3 Multi-head attention and the transformer block

The head we just built has exactly one idea about language — find me an agent and an object — and it applies that idea to every token whether or not it makes sense. A model that could only do this would be a very narrow instrument.

The fix is simple: run several heads side by side on the same input, each with its own \(W^Q\), \(W^K\), \(W^V\), and let each represent different relations. Do not assume a trained head reduces to a single job: its pattern can change with context (see the box below).

A second head, same sentence

Keep everything about kedi sütü içti and add one more head, with a different job (an unmasked, bidirectional one: in a GPT-style causal model kedi and sütü could not see içti, which comes after them). Head 1 asks who are my arguments. Head 2 asks the opposite question — which verb governs me, which in a verb-final language means the nouns should both point forward to içti:

\[ W^{Q}_{2}=\begin{bmatrix}2.5&0\\2.5&0\\0&2.5\\0&0\end{bmatrix}\; W^{K}_{2}=\begin{bmatrix}0&1.6\\0&1.6\\2&0\\0&0\end{bmatrix}\; W^{V}_{2}=\begin{bmatrix}0&0.6\\0&0.6\\1&0\\0&0\end{bmatrix} \]

The query dimensions now mean looking for an action and looking for a noun; the key dimensions mean I am an action and I am a noun. Running the same five steps gives a completely different matrix:

Head 1 — argument structure

kedi sütü içti
kedi 0.6961 0.1791 0.1249
sütü 0.1750 0.6804 0.1446
içti 0.4230 0.5012 0.0759

Head 2 — which verb governs me

kedi sütü içti
kedi 0.0383 0.0383 0.9234
sütü 0.0383 0.0383 0.9234
içti 0.4675 0.4675 0.0650

The two heads disagree about almost every row, and that disagreement is the entire point. In head 1, kedi mostly looks at itself. In head 2, kedi spends 92% of its attention on the verb. We built both heads, each designed to read a different relation off the same three words. What they show is that several heads can express different flows of information; which one is grammatically right is a separate question.

Real heads are not this tidy. Published analyses of trained models do find heads with recognisable jobs — attending to the previous token, to the matching bracket, to the syntactic head of a phrase (Clark et al. 2019) — but most heads resist any one-sentence description, many are redundant, and a large fraction can be pruned with little loss (Michel et al. 2019). Treat “head 2 finds the verb” as a description of this hand-built example, not as a claim about what trained heads do.

Concatenate, then project

Each head produces a \(d_v\)-wide vector per token. Stack them side by side and push the result through one more matrix \(W^O\):

\[ \operatorname{MultiHead}(X) = \operatorname{Concat}(\text{head}_1, \dots, \text{head}_h)\, W^O \]

Split into heads, run them in parallel, glue the results back together
X n × d HEAD 1 — ARGUMENT STRUCTURE its own W_Q, W_K, W_V içti → 0.4230 / 0.5012 / 0.0759 n × d/h HEAD 2 — WHICH VERB GOVERNS ME different W_Q, W_K, W_V içti → 0.4675 / 0.4675 / 0.0650 n × d/h Concat n × d W_O d × d output n × d Because each head is only d/h wide, h heads cost the same weights as one full-width head: 3d² for the projections plus d² for W_O, whatever h is. What grows is the score tensor, h·n².

Two heads on kedi sütü içti. The rows shown are each head’s real attention weights for içti, computed below — the same token, read two different ways at the same time.

With two heads of width 2, the concatenation is exactly 4 wide — back to \(d_{\text{model}}\), which is not a coincidence. In the classic layout the head size is chosen as \(d_k = d_{\text{model}} / h\) precisely so that concatenating them lands you where you started. That is a convention, not a necessity: in general \(W^O\) projects the concatenated \(h\,d_v\)-wide vector back to \(d_{\text{model}}\) anyway:

token h1 agent h1 object h2 action h2 noun
kedi 0.7444 0.3433 0.8310 0.1014
sütü 0.3256 0.7443 0.8310 0.1014
içti 0.5308 0.6009 0.0585 0.5649

\(W^O\) then decides how much of each head’s report to write back into the model’s working memory. Ours routes head 1’s agent and object content into the animate and liquid dimensions, head 2’s action content into the action dimension, and a fifth (0.2) of head 2’s noun content into both animate and liquid — which is why kedi’s animate value below is 0.7646 rather than head 1’s 0.7444. Its last column is all zeros, so nothing is written to the fourth dimension:

\[ W^O=\begin{bmatrix}1&0&0&0\\0&1&0&0\\0&0&1&0\\0.2&0.2&0&0\end{bmatrix} \qquad\Longrightarrow\qquad \begin{array}{r|rrrr} & \text{animate} & \text{liquid} & \text{action} & \text{bias}\\\hline \texttt{kedi} & 0.7646 & 0.3635 & 0.8310 & 0.0000\\ \texttt{sütü} & 0.3458 & 0.7646 & 0.8310 & 0.0000\\ \texttt{içti} & 0.6438 & 0.7139 & 0.0585 & 0.0000 \end{array} \]

Multi-head attention costs almost nothing extra. \(h\) heads of width \(d_{\text{model}}/h\) have the same total parameter count and the same matrix-multiply FLOP count (FLOP: one floating-point operation, a single multiply or add) as one head of width \(d_{\text{model}}\): the \(W^Q\), \(W^K\), \(W^V\) projections are \(3d^2\) either way, \(W^O\) adds \(d^2\), and scoring costs \(h \cdot n^2 \cdot (d/h) = n^2 d\) regardless of \(h\). Splitting into heads buys several independent attention patterns for very close to nothing in parameters and FLOPs, which is why the design stuck. Real latency is another matter: \(h\) separate score and softmax matrices can change memory traffic and the kernel used.

Two costs do move, and they matter later in this chapter:

  • The score tensor grows. In a naive implementation that materialises the scores, one head stores \(n^2\) of them and \(h\) heads store \(h \cdot n^2\) (blocked implementations such as FlashAttention never build the full tensor). At \(n = 1024\), \(h = 32\) that is 33.5M entries instead of 1.05M — a 32× increase in the memory that 1 is about, and in the softmax work, which is not a matrix multiply.
  • Each head gets narrower. A head of width \(d/h\) constrains the pre-softmax score matrix \(QK^{\top}\) to rank \(\le d/h\). At \(d = 4096\) that is 128 per head instead of 4096. Note where the bound stops: row-wise softmax does not preserve it. With \(d_k = 1\) and \(Q = K = [1, 2]^{\top}\), \(QK^{\top}\) has rank 1, while the softmaxed matrix has determinant 0.1497 and therefore rank 2. So the trade is one high-rank score relation for 32 low-rank ones — usually a good trade, and always a trade.

The payoff, in numbers

Now read the attention output as a change to each token. Every token started with a vector saying what it is in isolation; the attention sublayer hands back what it learned from the rest of the sentence:

token started with attention added reading
içti animate 0.0, liquid 0.1 animate +0.6438, liquid +0.7139 the verb has acquired an agent and an object
kedi action 0.0 action +0.8310 the noun has found the verb it belongs to
sütü action 0.0 action +0.8310 so has the other one

This is the first step of what “contextual embedding” means, stated as arithmetic. (In a real model this output also passes through the residual, the normalisation and the feed-forward layer, and the numbers come from training, not from us.) The verb içti entered the layer knowing only that it was an action; it leaves carrying numeric traces of who drank and what got drunk. That is the difference between the static vectors of Chapter 3, where içti has one vector forever, and a transformer, where its vector is rebuilt from its neighbours at every layer.

The bias column gained exactly zero, in every row — an artefact of this \(W^O\), whose last column is all zeros and so sends nothing back to that dimension.

Now for the promised return to the constant dimension. The tidy story about it is only two-thirds true. Split the score into the pieces the constant \(c\) contributes, writing each token as \(x_i = [z_i, c]\):

\[ q_i \cdot k_j = \underbrace{z_i A \cdot z_j B}_{\text{the real content}} \;+\; \underbrace{c\,(z_i A \cdot b)}_{\text{same for all } j} \;+\; \underbrace{c\,(a \cdot z_j B)}_{\textbf{varies with } j} \;+\; \underbrace{c^2 (a \cdot b)}_{\text{same for all } j} \]

where \(a\) and \(b\) are the rows of \(W^Q\) and \(W^K\) that read the constant feature. Softmax is invariant to adding the same number to a whole row, so the second and fourth terms genuinely vanish. Three conclusions, and only the first two are what you would guess:

  • On the key side, for one head’s scores, it is useless. A constant cannot advertise. Whatever \(b\) is, its contribution is identical across the \(j\) of a query row and cancels in the row softmax.
  • On the value side it is a bias term. Our \(W^V\) gives it a weight of 0.2, so it adds a constant \(+0.1\) to both components of every value vector.
  • On the query side it is not useless at all. The third term depends on \(j\), and it does not cancel. It is a fixed preference over keys that every token asks for regardless of what it is — a learned global prior. In our example it disappears only because this \(W^Q\) has a zero row for that feature; set that row to \([0.9, -0.7]\) instead and the attention weights move by as much as 0.1592.

A constant feature can never answer a question, but it can ask one. That third term is not a curiosity: it is one plausible route by which a model builds an attention sink, a position that every token dumps leftover weight onto when it has nothing in particular to look at. It is not enough on its own; key geometry, the mask and the course of training play a part too.

Residual connections: adding, not replacing

Notice the word added. The sublayer output is not the new representation; it is summed onto the old one:

\[ x \;\leftarrow\; x + \operatorname{MultiHead}(x) \]

This is the residual connection, and it is doing more work than its one-line definition suggests. Because every sublayer adds to a running vector rather than overwriting it, the sequence of vectors running through the depth of the model is often called the residual stream: a shared workspace that each layer reads from and writes to. A layer that has nothing useful to contribute can output near-zero and leave the stream intact, which means a layer added to a transformer has a “do no harm” option ready to hand. Training is not guaranteed to take it, but it is easy to find.

residual connection — a shortcut that adds a sublayer’s input to its output, so the layer only has to learn a correction rather than a whole new representation. ↩︎

It also helps the gradient. Differentiating \(x + f(x)\) with respect to \(x\) gives \(1 + f'(x)\) (for vectors, the Jacobian \(I + J_f\)): there is a direct path from the loss back to every earlier layer whose derivative is exactly 1. That does not mean the total gradient is preserved — the other branch can shrink it — but it makes vanishing much harder.

Residual connections are often introduced as a cure for vanishing gradients. He et al. (2016) are explicit that this is not what they were solving. Their networks already used batch normalisation, so the gradients were verified healthy. The problem they named was degradation: a 56-layer plain network scored worse than a 20-layer one on training error. It could not learn to imitate the shallower net, because making a stack of layers compute the identity is hard — and the residual path hands it that identity for free. Better gradient flow is a real side effect, not the motivating disease.

Layer normalisation, and why it moved to the front

Residual streams grow. Every sublayer adds something, so the numbers drift upward in scale as depth increases, and both attention (through the softmax) and the feed-forward nonlinearity care about scale. Layer normalisation rescales each token’s vector on its own (Ba et al. 2016) — across its \(d_{\text{model}}\) features, not across the batch, and not across the sequence:

\[ \operatorname{LN}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta, \qquad \mu = \frac{1}{d}\sum_j x_j, \qquad \sigma^2 = \frac{1}{d}\sum_j (x_j - \mu)^2 \]

layer normalization — rescaling each token’s feature vector to zero mean and unit variance, then applying a learned scale \(\gamma\) and shift \(\beta\); computed across that token’s own features rather than across the batch. ↩︎

Apply it in the block’s real order. In a pre-norm block the normalisation before the feed-forward sublayer sees \(x + \operatorname{Attn}(\cdot)\), so normalise the residual sum, not the sublayer output. First do the addition of the previous section:

One simplification to declare. A real pre-norm block normalises before attention too — the worked example fed raw embeddings straight to \(W^Q\), \(W^K\), \(W^V\) so the numbers stayed legible. Everything downstream inherits that choice. The normalisation you are watching here is the block’s second one, in its correct position.

token animate liquid action bias
kedi 1.6646 0.4635 0.8310 0.5000
sütü 0.4458 1.6646 0.8310 0.5000
içti 0.6438 0.8139 0.9585 0.5000

and only then normalise those rows, with \(\gamma = 1\) and \(\beta = 0\):

token animate liquid action bias
kedi 1.6543 −0.8299 −0.0699 −0.7545
sütü −0.8509 1.6508 −0.0602 −0.7397
içti −0.4932 0.4908 1.3268 −1.3244

Every row now has mean 0 and standard deviation 1. Note what that costs: after normalisation a feature is only meaningful relative to the token’s own other features. içti gained a lot of animate content above, but it gained even more liquid and action content, so its animate value comes out negative. In this table (\(\gamma = 1\), \(\beta = 0\)) absolute magnitudes are gone; only the profile survives. In a real model the learned \(\gamma\) and \(\beta\), and the residual path that flows on un-normalised, can still carry some scale information.

That independence from other examples is why layer norm and not batch norm. Batch normalisation computes statistics across the batch, which needs extra machinery such as padding and masks when sequences have different lengths, makes training depend on batch composition, and requires storing running statistics for inference. Layer norm needs one token’s own numbers and nothing else, so it behaves identically at batch size 512 and batch size 1.

Where it goes matters. The original Transformer put the normalisation after the residual add, \(\operatorname{LN}(x + f(x))\) — post-norm. Since GPT-2, the standard has been to put it before the sublayer, \(x + f(\operatorname{LN}(x))\) — pre-norm. The difference is that pre-norm leaves the residual path completely clean: there is an unbroken identity connection from the input embedding to the final layer, with nothing sitting on it. Post-norm rescales the stream at every block, so the identity path is no longer an identity; Xiong et al. (2020) showed that the deep post-norm models they studied could not train stably without a learning-rate warmup, because of large gradients at the start, and that pre-norm largely removes the need. Everything below is pre-norm.

And in practice, often not even layer norm. GPT-2 and BERT use layer norm; the recent open-weight families this chapter names — Llama, Mistral, Qwen, Gemma — use RMSNorm (Zhang and Sennrich 2019):

\[\operatorname{RMSNorm}(x) = \frac{x}{\sqrt{\tfrac{1}{d}\sum_j x_j^2 + \epsilon}} \odot \gamma\]

Two things are gone: the mean subtraction and the shift \(\beta\). The claim — supported by ablation rather than by theory — is that re-centring was never the part that helped; only the re-scaling was. Its normalisation layer has half the parameters of layer norm’s (no \(\beta\); a negligible difference for the model’s total), and skipping the mean makes it cheaper to run — its authors report 7–64% shorter running times on the models they studied. The gain in your own model depends on the architecture and the implementation. It is also why a Llama config file has an rms_norm_eps field and no layer_norm_eps. Everything above about where the normalisation goes applies unchanged; only the function inside the box gets simpler.

“Standard” is not “universal”. Gemma 2 uses sandwich norm, normalising before and after each sublayer, \(x + \operatorname{N}_2(f(\operatorname{N}_1(x)))\). OLMo 2 goes the other way and drops the input norm, keeping only an output one, \(x + \operatorname{N}(f(x))\). Both were adopted for training stability at scale — and in both the normalisation still sits on the branch, so the identity path stays unbroken. That is the property nobody gives up. Check the config before assuming pre-norm of a specific checkpoint.

The feed-forward sublayer

Attention moves information between tokens, and that is close to all it can do.

Be precise about why, because the usual shorthand — “attention is linear” — is wrong. Attention is not a linear function of \(x\); the softmax of a quadratic form is thoroughly nonlinear.

The real limit is elsewhere. Given the weights, the output is a convex combination of the value vectors (a weighted average with non-negative weights that sum to 1), so it can only ever land inside the convex hull of the values it was handed — the region “between” them. It can reweight and blend. It cannot compute a feature that none of its inputs carried. (That bound belongs to this intermediate step: \(W^O\) and the residual add move the representation into other coordinates; a step that processes content nonlinearly is still needed.)

So something has to supply a nonlinearity acting on the content of a vector, and something has to give each token a chance to process what it just received, privately.

That is the second sublayer — a two-layer MLP applied to each position separately and identically:

\[ \operatorname{FFN}(x) = \sigma(x W_1 + b_1)\, W_2 + b_2, \qquad W_1 \in \mathbb{R}^{d \times d_{\text{ff}}}, \qquad W_2 \in \mathbb{R}^{d_{\text{ff}} \times d} \]

\(\sigma\) is the nonlinearity, and which one it is has changed twice. Vaswani et al. (2017) used ReLU, \(\max(0, z)\). GPT and BERT switched to GELU (Hendrycks and Gimpel 2016), \(z\,\Phi(z)\) — a smooth ReLU that weights an input by the probability a standard normal falls below it, so small negative values leak through instead of being clipped to exactly zero. Since Llama the standard has been SwiGLU (Shazeer 2020), which is not a pointwise function at all but a gate:

\[ \operatorname{FFN}_{\text{SwiGLU}}(x) = \big(\operatorname{Swish}(xW_1) \odot xW_3\big)W_2, \qquad \operatorname{Swish}(z) = z\,\sigma_{\text{logistic}}(z) \]

One branch decides how much to let through and the other decides what — which is where the third matrix in the table below comes from, and why its \(d_{\text{ff}}\) ratios are smaller than 4 while its parameter counts are larger. The worked example further down uses ReLU, because a threshold you can read off a table is easier to watch fire than a smooth curve.

The words separately and identically carry the weight: kedi’s FFN computation does not see sütü at all. Every cross-token interaction in a transformer happens in attention, and nowhere else. This is worth remembering when you are debugging a model that mixes up two entities in a long document — the mixing must have gone through attention, though usually not through one attention map but carried across several heads and layers.

Why \(d_{\text{ff}} = 4d\)?

The short answer is that Vaswani et al. (2017) picked 4 and it worked, and the best-known models of the next few years — BERT, GPT-2, GPT-3 — kept it. There is no derivation. But there is a reading that makes the shape make sense: the FFN behaves like a key–value memory (Geva et al. 2021). Each of the \(d_{\text{ff}}\) columns of \(W_1\) — the rows of the \(W_1^{\top}\) printed below — is a pattern detector, and the matching row of \(W_2\) is what gets written to the residual stream when that detector fires. Under that reading \(d_{\text{ff}}\) roughly sets how many distinct patterns this layer can hold (no unit is one separate “fact”; this is a way of reading the network, not a count), and the expand-then-contract shape is the point: you need room wider than \(d\) to hold more distinct patterns than you have dimensions.

Watch what happens when the surrounding design changes, because the “4” turns out to be entirely negotiable. Modern models replace the single \(W_1\) with a gated unit — SwiGLU — which uses three matrices instead of two, so the sublayer costs \(3\,d\,d_{\text{ff}}\) rather than \(2\,d\,d_{\text{ff}}\). What did the field do with that?

model \(d\) \(d_{\text{ff}}\) ratio FFN parameters
GPT-2, GELU 768 3,072 4.00 \(8d^2\)
Llama 2 7B, SwiGLU 4,096 11,008 2.69 \(8.1d^2\)
Llama 3 8B, SwiGLU 4,096 14,336 3.50 \(10.5d^2\)
Llama 3 70B, SwiGLU 8,192 28,672 3.50 \(10.5d^2\)
Mistral 7B, SwiGLU 4,096 14,336 3.50 \(10.5d^2\)

Llama 2 picked \(\tfrac{8}{3}d \approx 2.67d\) (rounded up to 11,008, a multiple of 256), which is the width that holds the parameter count at \(8d^2\) despite the third matrix. Llama 3 and Mistral chose a different balance of parameters and quality and spent 31% more on the FFN. So neither the 4 nor the \(8d^2\) is a law. What survived is the shape — expand well past \(d\), apply a nonlinearity, contract back — and the fact that this sublayer is where most of a model’s parameters go. Treat any specific constant as a hyperparameter someone tuned, not a result someone derived.

Watching it fire

Take \(d_{\text{ff}} = 8\) (a real model at \(d=4\) would use 16; this is shrunk to fit on a page) and hand-set the columns of \(W_1\) so each neuron is a recognisable detector, with the bias as its firing threshold. Reading the four input dimensions in the order animate, liquid, action, bias:

\[ W_1^{\top}=\begin{bmatrix} 1&0&1&0\\ 0&1&1&0\\ 1&1&0&0\\ 1&0&0&0\\ 0&1&0&0\\ 0&0&1&0\\ 0.5&0.5&0&0\\ 0&0&0&1 \end{bmatrix}, \qquad b_1=\begin{bmatrix}-0.9\\-0.9\\-0.9\\-0.5\\-0.5\\-0.5\\-0.5\\-0.5\end{bmatrix} \]

Each row is a question and each bias is how convinced the neuron has to be before it answers. Feed it the normalised vectors from the previous section:

detector its column of \(W_1\) asks kedi sütü içti
animate + action is this an animate thing attached to a verb? 0.6845 0 0
liquid + action is this a liquid attached to a verb? 0 0.6906 0.9176
animate + liquid is this animate and liquid at once? 0 0 0
animate only is this animate, verb or no verb? 1.1543 0 0
liquid only is this a liquid, verb or no verb? 0 1.1508 0
action only is this an action? 0 0 0.8268
any noun is this animate or liquid, weakly? 0 0 0
constant is the always-on dimension large? 0 0 0

Read the columns. kedi fires animate + action — after attention it is no longer just an animate thing, it is an animate thing attached to a verb. (The reading comes from the labels we gave the columns; it does not show that a trained neuron matches one grammatical concept.) içti fires liquid + action, having absorbed its object. The ReLU zeroes everything below threshold, which is how a linear layer becomes a set of switches.

Now notice which rows are doing the interesting work, because it is not all of them. The three single-feature detectors would have fired on the raw embeddings too — they tell you nothing attention did not already know.

The composite detectors are the ones that matter, and specifically for the two nouns. Before attention ran, kedi and sütü both had exactly 0.0 action content. Nothing could have answered “animate and attached to a verb”, because half the question did not exist yet.

(içti did start with a little liquid content of its own, but not enough: pushed through the same normalisation, its raw embedding gives the liquid + action detector a pre-activation of about −0.20, below threshold. After attention it is 0.9176.)

Which is the point of the whole layer. The FFN is reading features that attention created one sublayer earlier, and that is the whole reason the two sublayers alternate.

Two honest observations. First, içti misses the animate + action detector by 0.066 — it has the agent content, but layer normalisation pushed that feature below the token’s own average, and a threshold is a threshold. Second, the constant neuron never fires for anything: after normalisation the always-on dimension is below every token’s mean. Hand-designed detectors are brittle in exactly this way: sensitive to the thresholds and the normalisation we chose. In trained networks features are usually spread across several neurons, but how much cannot be read off this table.

Writing the answer back

Detecting is only half of it. \(W_2\) maps those eight neurons back down to the four dimensions of the residual stream, and by the reading above each row of \(W_2\) is the message its detector broadcasts when it fires. Keeping the same row order, and taking \(b_2 = 0\):

\[ W_2=\begin{bmatrix} 0.4&0&0.5&0\\ 0&0.4&0.5&0\\ 0.4&0.4&0&0\\ 0.3&0&0&0\\ 0&0.3&0&0\\ 0&0&0.3&0\\ 0.15&0.15&0&0\\ 0&0&0&0 \end{bmatrix} \]

Each detector writes into exactly the features it asked about. The composite ones write hardest into action (\(0.5\)), because “this noun is attached to a verb” is the fact worth broadcasting; the single-feature ones write a milder \(0.3\), because they are only restating what the token already advertised. Multiply out and add to the residual, \(x + \operatorname{FFN}(\operatorname{N}(x))\):

token animate liquid action bias
kedi 2.2847 0.4635 1.1733 0.5
sütü 0.4458 2.2860 1.1763 0.5
içti 0.6438 1.1810 1.6654 0.5

That is the block output — stage 4 in the widget below. Two things to notice. The bias column is untouched, because no row of \(W_2\) writes to it. And rows 3, 7 and 8 never got used: those detectors did not fire for any of the three tokens, so whatever their rows contain made no difference here. Under ReLU, exact zeros like these are the normal case rather than the exception. Two caveats before you generalise it: the GELU and SiLU activations most current models use do not zero the whole negative half-axis the way ReLU does, so “switched off” becomes “contributing very little”, and either way the dense layer still computes every row. That second point is what Section 7.7 is about — skipping part of the arithmetic outright: MoE skips the multiplications of the expert networks not chosen for a token, while every unit inside the chosen experts, and the router, is still computed.

The whole block

One pre-norm transformer block
RESIDUAL STREAM x x' LayerNorm per token Multi-head attention 4d² params + LayerNorm Feed- forward 8d² params + MIXES ACROSS TOKENS EACH TOKEN ALONE The only place information crosses between positions is the accent-coloured box. TWO THIRDS OF THE PARAMETERS ARE IN THE FEED-FORWARD HALF.

The pre-norm arrangement used in this chapter, common since GPT-2: the normalisation sits on the branch, not on the stream (real models may differ in norm type and placement; check the config), so there is an unbroken identity path from the input embedding to the output. Stack this block \(L\) times — 12 for GPT-2 small, 80 for a 70B-class model — and you have a transformer.

Every part, and the problem it solves

Nothing in that diagram is decorative. Each piece is a response to a specific failure, and the fastest way to understand a component is to ask what goes wrong when you delete it. This table is the chapter in one page — every row is argued for somewhere above or below.

component the problem it exists to solve what breaks without it
any learned score matrix raw similarity \(x_i \cdot x_j\) is symmetric and inflates the diagonal the two directions of a relation share one score; every token attends mostly to itself
factoring it into \(W^Q\), \(W^K\) one \(d \times d\) matrix per head is 16× the parameters, and pins the cache at width \(d\) training cost balloons and GQA becomes inexpressible
separate \(W^V\) what a token advertises is not what it should deliver routing and content are forced to be the same thing
dot product need a cheap comparable score for all \(n^2\) pairs, computed as a dense matrix multiply other score functions exist (additive attention, say), but this efficient matrix-multiply route is lost
\(\div\sqrt{d_k}\) with independent unit-variance components, score spread grows like \(\sqrt{d_k}\), and softmax reacts to absolute size wide heads risk saturating the softmax, and the gradient becomes very small
softmax weights must be positive, sum to 1, and be differentiable a hard “best match” has no usable gradient under ordinary backpropagation; other normalisations change how weights and gradients behave
multiple heads one weighted average can track several relations at once only in a limited way the layer expresses one blend per position instead of several independent patterns
\(W^O\) head outputs live in separate subspaces and must re-enter the stream no way to combine or route what the heads found
residual add a layer must be able to do nothing, and deep stacks are prone to degradation without that option deep networks become much harder to train; adding depth can stop helping and start to hurt
layer norm (RMSNorm in practice) the stream’s magnitude grows with depth; softmax and ReLU are scale-sensitive activations drift and training destabilises
pre-norm placement post-norm rescales the stream, so the identity path is not an identity the deep post-norm models studied could not train stably without learning-rate warmup
feed-forward attention only mixes: its output is a convex combination of values no private per-token nonlinear processing, and most capacity disappears
positional signal unmasked attention is permutation-equivariant — it cannot see order an unmasked encoder reads every sentence as a bag of words; a causally masked decoder can recover some order from the mask (NoPE, below)

Read the middle column on its own and you have the design rationale for the entire architecture. Every one of these is an answer to a concrete defect, not an aesthetic choice (though not the only possible answer) — which is why the shape has barely moved since 2017 while almost everything around it has.

Watch the representation change

Everything above happens to three vectors of four numbers each. Step through the block and watch them move.

The residual stream, one stage at a time

Compare stage 1 with stage 4. Every token leaves the block with the same number of dimensions but a larger norm (length) than it arrived with: norms go from 1.03 to roughly 2.2–2.7 (a result of the matrices chosen here, not a rule for every block). More importantly, every token leaves with content it did not have. kedi entered as animate and leaves as animate, attached to an action. That is one block. Stack twelve and each one gets to condition on what the previous eleven worked out.

Where the parameters actually live

The block has a fixed shape, so its size is a formula. With \(d_{\text{ff}} = 4d\):

part matrices parameters
attention \(W^Q, W^K, W^V, W^O\), each \(d \times d\) \(4d^2\)
feed-forward \(d \times 4d\) and \(4d \times d\) \(8d^2\)
block total \(\mathbf{12d^2}\)

Biases and layer-norm parameters are \(O(d)\) and round to nothing. So:

\(d_{\text{model}}\) attention feed-forward one block
768 2,359,296 4,718,592 7,077,888
1,024 4,194,304 8,388,608 12,582,912
4,096 67,108,864 134,217,728 201,326,592
8,192 268,435,456 536,870,912 805,306,368

With \(d_{\text{ff}} = 4d\) and full-width attention, two thirds of a transformer block is the feed-forward half. People say “attention is all you need” and then build blocks in which attention is a third of the weights.

Say block, not model. The ratio holds inside a block; across a whole model the embeddings dilute it, and how much depends on the size. In GPT-2 small below, the 12 blocks are only 68% of the parameters, so the feed-forward share of the model is about 46%. At \(d = 8192\) the embeddings are a rounding error and the model ratio converges back on two thirds. Small models are mostly vocabulary; large models are mostly feed-forward.

And two thirds is not a ceiling. It assumes the GPT-2 shape: \(8d^2\) of feed-forward against \(4d^2\) of attention. A Llama-3 block moves in both directions at once. Its SwiGLU feed-forward costs \(10.5d^2\), and its attention costs only \(2.5d^2\) — grouped-query attention shrinks \(W^K\) and \(W^V\) from \(d \times d\) to \(d \times 1024\) (1), so that pair costs \(0.5d^2\) instead of \(2d^2\). The feed-forward share is therefore \(10.5/13 = \mathbf{81\%}\), not the 72% you get by assuming full-width attention. That is this configuration’s in-block share; in the Llama family two changes of recent years (GQA and SwiGLU) moved weight from attention into the feed-forward layer, but other families may choose different intermediate widths, shared experts or other attention types. Before carrying a ratio over to a model, compute it from that model’s config.

A rule of thumb is only worth having if it reproduces a number you can look up. GPT-2 small is public: vocabulary 50,257, \(d = 768\), 12 layers, context 1,024.

part arithmetic parameters
token embeddings \(50{,}257 \times 768\) 38,597,376
position embeddings \(1{,}024 \times 768\) 786,432
12 blocks \(12 \times 12 \times 768^2\) 84,934,656
total 124,318,464

The published size is ~124M, and the rule lands within 0.098% of it. The residual is not mysterious — it is exactly the \(O(d)\) terms the rule throws away. Restore them and the count closes completely:

dropped term arithmetic parameters
attention biases \(12 \times (3d + d)\) 36,864
feed-forward biases \(12 \times (4d + d)\) 46,080
layer-norm \(\gamma, \beta\) \(12 \times 2 \times 2d + 2d\) 38,400
exact total 124,318,464 + 121,344 124,439,808

So the rule is not an approximation with unexplained slop; it is an exact count with three named terms deleted. (The exact count belongs to this configuration: 12 layers, \(d = 768\), a 50,257-token vocabulary, 1,024 learned positions, tied input and output weights; another checkpoint will not show the same gap.) Trust it for capacity planning, and remember which terms you dropped when someone asks why your number and the model card differ in the last six digits.

Notice what the table does not contain: an output projection back to the vocabulary. GPT-2 ties the output weights to the input embedding matrix, so those 38.6M parameters are counted once and used twice.

7.4 Positional information: absolute, learned, RoPE, ALiBi

Because attention itself is order-blind, position must be injected explicitly — the model literally cannot tell “dog bites man” from “man bites dog” without it. This is the same failure the averaged embeddings hit in Chapter 4, where kalite iyi fiyat kötü and kalite kötü fiyat iyi came out at cosine 1.0000.

First, prove the problem is real

That claim is easy to state and easy to believe without checking, so check it. Take the worked example, shuffle the three tokens into the order içti kedi sütü, run attention again, then un-shuffle the outputs and compare:

token outputs, original order outputs, shuffled then un-shuffled
kedi [0.7444, 0.3433] [0.7444, 0.3433]
sütü [0.3256, 0.7443] [0.3256, 0.7443]
içti [0.5308, 0.6009] [0.5308, 0.6009]

Largest disagreement anywhere in the table: \(1.11 \times 10^{-16}\), which is floating-point noise and not a difference. Unmasked self-attention is permutation equivariant: shuffle the inputs and the outputs shuffle with them, otherwise untouched.

This is not a bug that someone forgot to fix; it follows from the definition. Every token’s output is a weighted sum over all tokens, and a sum does not care what order you add things in. Nothing anywhere in \(QK^{\top}\), softmax, or \(AV\) can see an index. Each token still gets its own, context-dependent output, but those outputs do not depend on order. Pool them in a position-independent way (an average, say) and two sentences with the same tokens in a different order get identical representations: a bag of words with a very clever averaging scheme is still a bag of words.

“Unmasked” is doing real work in that sentence. Add the causal mask of Section 7.5 and the symmetry breaks: token 1 may attend to one position, token 5 to five, so the number of visible tokens already encodes an index. Decoder-only models can therefore learn position with no positional encoding at all — the so-called NoPE result (Kazemnejad et al. 2023), which at small scale even extrapolates to longer sequences better than some explicit schemes on the tasks studied. Most widely used generative models still add an explicit signal; a count can be a weak and indirect way to learn order, and not every decoder learns position to the same degree. But the honest statement is that the mask leaks position, not that a decoder is blind without an encoding.

Adding a position signal

The original Transformer’s answer is to build a vector for each slot and add it to whatever word lands there (Vaswani et al. 2017). The vector is made of sines and cosines at geometrically spaced frequencies:

\[ \text{PE}_{(\text{pos},\, 2i)} = \sin\!\left(\frac{\text{pos}}{10000^{2i/d}}\right), \qquad \text{PE}_{(\text{pos},\, 2i+1)} = \cos\!\left(\frac{\text{pos}}{10000^{2i/d}}\right) \]

At \(d = 4\) there are two frequencies — one fast, one a hundred times slower — and the whole encoding fits in three rows:

position \(\sin(\text{pos})\) \(\cos(\text{pos})\) \(\sin(\text{pos}/100)\) \(\cos(\text{pos}/100)\)
0 0.0000 1.0000 0.0000 1.0000
1 0.8415 0.5403 0.0100 1.0000
2 0.9093 −0.4161 0.0200 0.9998

Add \(0.3 \times \text{PE}\) to the embeddings (0.3 is just a scale chosen to keep this illustration’s numbers legible, not part of the method) and repeat the shuffle test. The subtlety that makes it work: when the words move, the position vectors do not move with them. Position 0 is still position 0, so a word that was first and is now third gets a different vector than before:

token outputs, original order outputs, shuffled then un-shuffled
kedi [0.8029, 0.7345] [1.1319, 0.5094]
sütü [0.6385, 1.0023] [0.7489, 0.6805]
içti [0.7865, 0.8342] [0.8767, 0.6607]

Largest disagreement: 0.3289. The order-blindness is gone.

positional encoding — the signal that lets a model use the position of tokens or the distance between them; it can be added to the embedding, applied to the query and key (RoPE) or added to the attention score as a bias (ALiBi). Needed because unmasked attention alone is order-blind. ↩︎

Why sinusoids in particular

Any distinct vector per slot would break the tie. Sinusoids were chosen because of a property that is easy to state and, unusually for this literature, exactly true. For a positional encoding built this way,

\[ \text{PE}_{p} \cdot \text{PE}_{p+k} \;=\; \sum_{i} \cos(k\,\omega_i) \]

The \(p\) terms cancel by the angle-difference identity, and what survives depends only on \(k\) — the offset. So the similarity between two positions is a function of how far apart they are and not of where they sit in the sequence.

If that were only approximately true it would be worth much less, so measure it. Sweeping \(p\) across 256 positions at \(d = 64\) and holding the offset fixed:

offset \(k\) mean \(\text{PE}_p \cdot \text{PE}_{p+k}\) std over \(p\) range over \(p\)
0 32.000000 1.6e−15 1.1e−14
1 30.916832 7.8e−15 4.6e−14
2 28.303862 1.1e−14 7.8e−14
4 23.934362 7.4e−15 5.0e−14
8 22.413912 5.8e−15 3.9e−14
16 19.372820 5.8e−15 3.9e−14
32 19.553972 1.0e−14 7.1e−14
64 13.906089 1.2e−14 9.1e−14

The identity is exact in mathematics; the standard deviation of \(10^{-14}\) left across 256 different absolute positions comes only from finite-precision floating point. Every position vector also has the same length, \(\sqrt{d/2} = 5.656854\), to ten digits. That is a property of the position vector, not of the sum: adding it changes the embedding’s norm, since \(\lVert x + p \rVert^2 = \lVert x \rVert^2 + \lVert p \rVert^2 + 2\,x \cdot p\) depends on the angle between them. RoPE, further down, is the scheme that leaves the norm alone.

Read the middle column honestly. The mean is not monotonic in \(k\): offset 32 scores higher than offset 16. It is a sum of cosines, so it wobbles. The correct claim is “this quantity depends only on the offset”, which the last two columns establish beyond argument. The tempting claim — “nearby positions look more similar than distant ones” — is not what the table shows, and a textbook that says it anyway is asking you to trust a picture over the numbers.

Sinusoidal position encodings, as a picture

Low dimensions (top) flip sign every few positions; high dimensions (bottom) barely change across the whole window. That spread of wavelengths is the point — the fast dimensions are suited to telling nearby positions apart and the slow ones distant positions. Whether a model actually uses them for “which neighbour” or “where in the document” cannot be read off this picture. Push d_model up and watch more slow dimensions appear at the bottom.

Learned absolute positions

Sinusoids are elegant, but not the only option. BERT and GPT-2 both used an ordinary embedding table indexed by position instead: row 0 is the vector for slot 0, learned by gradient descent like any other parameter (Devlin et al. 2019; Radford et al. 2019). In those models it worked about as well (which does not show sinusoids are unnecessary for every task), and you saw the cost in the parameter table above — GPT-2 small spends 786,432 parameters on it.

The real cost is elsewhere. A learned table has exactly as many rows as you trained it with. GPT-2’s table has 1,024 rows, so token 1,025 has no vector at all — not a degraded one, none. The context limit is not a tuning choice you can relax at inference time; it is baked into the shape of a weight matrix. Handling longer input in this family means extending the table (initialising and training new rows, or interpolating the old ones) or redefining position and re-adapting the model; changing an inference setting on the existing checkpoint is not enough.

RoPE: rotate instead of add

The dominant choice today is rotary position embedding, and its idea is a genuine change of kind rather than a better table (Su et al. 2024). RoPE adds nothing to the embedding. It takes the query and key vectors — after projection, inside each head — splits them into consecutive pairs of coordinates, and rotates each pair by an angle proportional to the position:

\[ \begin{bmatrix} q'_{2i} \\ q'_{2i+1} \end{bmatrix} = \begin{bmatrix} \cos m\theta_i & -\sin m\theta_i \\ \sin m\theta_i & \phantom{-}\cos m\theta_i \end{bmatrix} \begin{bmatrix} q_{2i} \\ q_{2i+1} \end{bmatrix}, \qquad \theta_i = 10000^{-2i/d_k} \]

Here \(m\) is the token’s position and \(d_k\) is the head dimension, since the rotation happens inside each head. Rotating a vector by \(m\theta\) and another by \(n\theta\) and then taking their dot product leaves an angle of \((m-n)\theta\). So for the same query and key contents the attention score itself becomes a function of the offset (the score still depends on the contents too; different query–key pairs at the same distance do not get the same score) — not a side channel added to the embedding and hoped to survive the projections, but the score.

That is a strong claim, so here it is measured. One fixed random query and key, \(d = 64\), scored at a fixed offset from four wildly different absolute positions (\(n = 0, 100, 1000, 5000\)):

offset \(m-n\) score spread across the four positions
0 2.152941 2.2e−15
1 3.606170 1.5e−13
2 5.453257 3.1e−13
5 6.457170 5.0e−13
10 8.827710 1.4e−13
50 −0.061180 1.5e−13
200 −13.740668 4.1e−14

Now the same measurement on the additive scheme, where the position vector is summed into the embedding before projection:

offset \(m-n\) score at \(n=0\) at \(n=100\) at \(n=1000\) spread over all four
0 31.19 26.98 48.15 26.3
1 31.82 27.54 47.52 26.6
10 13.74 20.55 38.85 25.1
200 1.86 6.43 25.00 23.1

Three of the four positions are printed; the spread column covers all four, which is why the first two rows are wider than the visible columns alone would give.

A spread of \(10^{-13}\) against a spread of 26. With addition, “five tokens back” means something different at position 0 than at position 5,000, and the model has to learn to cope with that. With rotation it means the same thing everywhere, for free. That shows relative position entering the score geometry directly; long-context quality still has to be measured separately. RoPE also preserves vector norms exactly — a rotation is a rotation — so it cannot distort the magnitude of a query the way addition can.

This is why RoPE is in Llama, Qwen, Mistral, Gemma and most open-weight models released since 2023 (ALiBi, learned and hybrid schemes are in use too; check a model’s position scheme in its config), and why the main context-extension methods for those models (position interpolation, NTK-aware scaling, YaRN) are rescalings of \(\theta_i\) or of the position index rather than new tables, often followed by extra long-context training.

ALiBi: skip the encoding, penalise distance

The most economical answer is to encode nothing and simply subtract a penalty proportional to distance, directly on the attention scores (Press et al. 2022):

\[ \text{score}(i, j) = \begin{cases} \dfrac{q_i \cdot k_j}{\sqrt{d_k}} - m\,(i - j), & j \le i \\[4pt] -\infty, & j > i \end{cases} \]

The penalty is subtracted after the \(\sqrt{d_k}\) scaling, not before. That is our reading of why it matters rather than a claim the paper makes: added before the division, the effective slope would be \(m/\sqrt{d_k}\) and would change with head width; added after, the same slopes \(m\) mean the same thing in models of different width.

Each head gets its own fixed slope \(m\), on a geometric schedule. With eight heads the slopes are \(\tfrac12, \tfrac14, \tfrac18, \dots, \tfrac{1}{256}\), so the head with \(m = 0.5\) has accumulated a penalty of 1 after 2 tokens (a score lower by 1, which, other scores being equal, divides the attention weight by \(e \approx 2.7\)) and leans local, while the head with \(m = 1/256\) needs 256 tokens to accumulate the same penalty and leans global. Different heads lean towards different distance scales, by construction; how local a head actually behaves also depends on the content scores — a distant token with a high content score can still get a high weight.

The bias a head with \(m = 0.5\) adds, over six tokens:

0 1 2 3 4 5
pos 0 0.0 −∞ −∞ −∞ −∞ −∞
pos 1 −0.5 0.0 −∞ −∞ −∞ −∞
pos 2 −1.0 −0.5 0.0 −∞ −∞ −∞
pos 3 −1.5 −1.0 −0.5 0.0 −∞ −∞
pos 4 −2.0 −1.5 −1.0 −0.5 0.0 −∞
pos 5 −2.5 −2.0 −1.5 −1.0 −0.5 0.0

That single matrix is the causal mask and the position signal at once, and it contains no learned parameters whatsoever. Because it is defined by a formula over \(i - j\), it is defined for any distance. The original paper reported a 1.3B-parameter ALiBi model trained at length 1,024 doing well at length 2,048 — the result the paper’s title advertises. A formula defined at every length does not mean every ALiBi model runs at every length without degrading.

Choosing

scheme how position enters extra parameters beyond training length
sinusoidal added to embeddings none formula defined; quality must be measured, often degrades
learned absolute added to embeddings \(n_{\text{ctx}} \times d\) undefined — no row in the existing table
RoPE rotates Q and K per head none formula defined; typically degrades, usable with extension methods and often extra training
ALiBi subtracted from scores none formula defined; did well at twice the training length in the original paper

RoPE does not usually extrapolate well on its own. Because it is a formula over \(i - j\), RoPE is defined at any distance — but a model trained at 4k often falls apart well before 8k, because the rotation angles it sees past the training length are ones it never encountered. That failure is exactly why the long-context toolkit exists: Position Interpolation squeezes unseen positions back into the trained range (Chen et al. 2023), and NTK-aware scaling and YaRN (Peng et al. 2024) rescale \(\theta\) frequency-by-frequency instead of uniformly. An “extended to 128k” announcement about a RoPE model usually describes one of these plus continued training (some methods are also tried at inference time alone) — not a free property of RoPE. Measure the quality you expect on the target task.

The column that matters in practice is the last one. When a vendor announces a context window, they are describing a property of the position scheme as much as of the training data — and when a model becomes strangely unreliable past some length, the position scheme is the first place to look. Advertised context length and usable context length are different numbers, a gap we will measure rather than assume in Chapter 25.

7.5 Encoder-only, decoder-only, encoder-decoder

Here is a fact that sounds like a simplification and mostly is not: the structural difference between BERT-style and GPT-style self-attention is which entries of the attention matrix are allowed to be non-zero. Two qualifications travel with it, and the rest of this section makes both: an encoder–decoder is two stacks joined by cross-attention, which no mask over one stack produces, and the names BERT and GPT also carry particular embeddings, normalisation placements and training objectives.

Same blocks, same multi-head attention, same feed-forward sublayer, same residual stream. What differs, at the level of self-attention’s access, is a causal mask applied to the scores before the softmax — and, following from that, what the model can be trained to do. The behaviour of concrete models also depends on the training objective, the data, the output head and the other architectural choices.

causal mask — blocks every position from attending to later ones, so the model can be trained to generate left to right. ↩︎

The mask is the architecture

Look at the three patterns.

Which token may look at which

Select the causal mask and the top-right triangle goes grey. kedi, in position 0, can then attend only to itself. That constraint is what lets a decoder be trained on raw text efficiently, in a single parallel pass: with it, the model can be asked to predict every next token in a sentence simultaneously: position \(i\) sees itself and what came before, and its target is the next token, \(i+1\), so the target never appears in its own input. That turns one document into thousands of training signals. (When generating, new tokens are still chosen one at a time.)

Mechanically it is one line — set the masked scores to \(-\infty\) before the softmax, so \(e^{-\infty} = 0\) and those tokens receive exactly zero weight. Setting the weights to zero after the softmax would not be equivalent, because the remaining weights would no longer sum to 1.

The three families

Three families, one difference: which cells are allowed
ENCODER-ONLY every token sees every token trained to fill in blanks BERT — classify, extract, embed DECODER-ONLY a token sees only the past trained to predict the next token GPT, Llama, Qwen — generate ENCODER-DECODER cross source target full input, causal output trained to emit a target sequence T5, BART — translate, summarize A filled cell means "this query row may read this key column". Nothing else about the architecture changes between the three — same block, same maths, different mask.

Rows are queries, columns are keys. In terms of self-attention access, the families differ in which cells are switched on; defining the encoder–decoder family also takes the cross-attention that joins its two stacks.

family mask trained to representative good at
encoder-only bidirectional fill in masked tokens BERT (Devlin et al. 2019) classification, extraction, retrieval embeddings
decoder-only causal predict the next token GPT, Llama, Qwen generation, and — since scale — many tasks that can be solved by prompting
encoder-decoder bidirectional in, causal out produce a target sequence T5 (Raffel et al. 2020), BART translation, summarization, any input→output rewrite

An encoder-decoder needs one extra mechanism. Its decoder runs self-attention over what it has generated so far, and then a second attention sublayer called cross-attention in which the queries come from the decoder while the keys and values come from the encoder. It is the same arithmetic you worked through above — the only change is that \(Q\) and \(K, V\) are computed from two different sequences. With \(m\) target tokens and \(n\) source tokens, the score matrix is \(m \times n\) — no longer square. That is how “attend to the source sentence while writing the translation” is expressed.

Some multimodal models attach an image this way — Flamingo (Alayrac et al. 2022) lets the language model look at the image through cross-attention, with the image encoder producing keys and values. Most open vision-language models instead project image tokens straight into the language model’s input sequence (Liu et al. 2023). We come back to both in Chapter 21.

Why decoder-only won

In 2019 the sensible prediction was that each family would keep its niche. For large general-purpose models that is not what happened; most of what you will use in this book is decoder-only. Three plausible reasons (data, infrastructure and investment are tangled in too; the causes of this history have not been measured):

One objective, all the data. Next-token prediction needs no labels, no corruption scheme, and no paired data. Every document on the internet is training data for it, and every position in that document is a training example.

Task-as-text. Once a model is good enough at continuation, classification and translation and summarization can all be written down as prompts rather than built as separate heads. The architectural specialisation of an encoder-decoder stops paying for itself when a general model can be asked politely. This is the story of Chapter 17.

Cheap incremental decoding. With a causal mask, the representation of token \(i\) never depends on anything after it — so when you append a token, all the earlier work is still valid. That is what makes the KV cache possible, and the KV cache is why generation costs what it costs.

Encoder-only models are not obsolete. For classification and retrieval at volume, a 110M-parameter BERT-style encoder can be far cheaper per document than a generative model and is often close in accuracy, because it emits one vector in one forward pass instead of decoding token by token. How much cheaper and how accurate depends on the models, the hardware and the task; measure it on yours. Most high-volume retrieval systems (Chapter 19) still run encoder embedding models, although the top of the embedding leaderboards is increasingly held by decoders adapted to the job. Choosing a 7B decoder for sentiment analysis is a common and expensive mistake — and Chapter 4 already showed how often the cheaper representation wins.

7.6 Efficiency: FlashAttention, GQA/MQA, KV cache

Everything so far has been about what attention computes. This section is about what it costs: how the computation grows with sequence length, memory traffic and the way text is generated. Those costs weigh heavily on which hardware, and how many concurrent requests, a model can run on.

The quadratic term, in bytes

\(QK^{\top}\) is \(n \times n\). That is the whole problem in one line. Written naively, a single head at a single layer must hold that matrix in memory. At fp16 (16-bit floating point) each number takes 2 bytes; a GiB is \(2^{30}\) bytes, about 1.07 GB:

sequence length entries in \(QK^{\top}\) one score matrix at fp16
512 262,144 0.0005 GiB
4,096 16,777,216 0.03 GiB
32,768 1,073,741,824 2.00 GiB
100,000 10,000,000,000 18.63 GiB

Eighteen gibibytes for one head at one layer. Multiply by 32 heads and 32 layers and storing all the scores at once stops being a memory problem and becomes a fantasy. That is the arithmetic of a naive implementation that writes the matrix out; blocked implementations such as FlashAttention, below, compute the same result without ever building it.

When does attention actually dominate?

The quadratic term is not always the term that matters, and people routinely overestimate it. Assuming full multi-head attention and \(d_{\text{ff}} = 4d\), per layer in a full-sequence forward pass, roughly (as FLOP counts):

\[ \underbrace{24\,n\,d^2}_{\text{projections and FFN}} \qquad\text{versus}\qquad \underbrace{4\,n^2 d}_{\text{scores and weighting}} \]

Setting them equal gives \(n = 6d\). That is a rough threshold for operation counts; real latency also depends on memory access, the kernel and the hardware. Below it, the transformer spends most of its arithmetic on matrix multiplications that are perfectly ordinary and linear in sequence length:

\(d_{\text{model}}\) sequence length \(n\) attention’s share of FLOPs
768 1,024 18%
768 4,608 crossover 50%
768 32,768 88%
4,096 4,096 14%
4,096 24,576 crossover 50%
4,096 100,000 80%

At contexts of a few thousand tokens, attention is a minority of the arithmetic. A 4,096-wide model running 4,096 tokens spends 86% of its arithmetic on the feed-forward and projection matrices. The quadratic term is a wall you hit at long context, not a tax you pay at every length. Knowing which regime you are in is a good first hint as to whether optimising attention is worth your time; but even with a small FLOP share, attention can be a bottleneck through memory traffic (next section), so decide on measured latency.

FlashAttention: don’t build the matrix

Here is the observation that made long context practical: on a modern GPU, at long sequence lengths, attention is often not limited by arithmetic. It is limited by moving that \(n \times n\) matrix between fast on-chip memory and slow high-bandwidth memory. The FLOPs are cheap; the round trips are not (Dao et al. 2022).

FlashAttention therefore never materialises the score matrix. It walks through the sequence in tiles, keeping a running softmax normaliser so that a partially-computed softmax can be corrected as new tiles arrive, and accumulates the output directly.

FlashAttention is exact. (We mean the original, dense algorithm; the block-sparse variant in the same paper is approximate.) It is not a sparse approximation, a low-rank factorisation, or a windowed heuristic. It computes the same function as the naive implementation and differs only in what it stores. Exact is not the same as bit-identical: the online softmax re-associates the sums, and floating-point addition is not associative, so the last digits can move — the same way any two matrix-multiply kernels disagree in the last digits. The dense algorithm trades away no quality, which is why it went from a paper to a default in about eighteen months, and why memory use in attention is now linear in \(n\) even though the computation is still quadratic.

The KV cache: the number that decides deployment

Generation is incremental. Having produced 500 tokens, the model produces the 501st — and because the mask is causal, the keys and values of the first 500 have not changed. Recomputing them would be pure waste, so they are kept. That store is the KV cache. Instead of running the whole prefix through the model at every step, only the new token is processed and its query looks at the \(n\) keys in the past: per-token attention cost goes from quadratic to linear. Generation still is not constant-time; the new token’s projections and feed-forward, and reading the weights, remain.

KV cache — stored keys and values from earlier tokens, so generation does not recompute them. Its size, not the model’s, usually decides how many users a GPU can serve. ↩︎

This splits a request into two phases with very different characters. The prompt is processed all at once, in parallel, exactly as in training, and its keys and values fill the cache — this is called prefill. After that, each new token is one small step (decode): little arithmetic per step, but in a classic full-attention decoder all of the model’s weights and every layer’s cache are read from memory for every new token (sliding windows, compressed caches and state-space layers change how much; below). That is why generation speed is usually limited by memory bandwidth rather than by arithmetic, a point Chapter 31 builds on.

What the cache saves, and what it costs
GENERATING TOKEN 6 — NO CACHE recompute all 6 tok 1 k, v tok 2 k, v tok 3 k, v tok 4 k, v tok 5 k, v tok 6 k, v 6 × work GENERATING TOKEN 6 — WITH A CACHE reuse 5, compute 1 tok 1 cached tok 2 cached tok 3 cached tok 4 cached tok 5 cached tok 6 NEW 1 × work The causal mask is what makes this legal: token 3's key can never depend on token 6, so it never needs recomputing. The arithmetic you save comes back as memory you must budget.

The cache is the reason generation is affordable and the reason long context is not. It trades re-processing the whole prefix at every step (per layer, the projections and feed-forward of the prefix’s \(n\) tokens and the attention among them) for \(O(n)\) storage, held separately for every concurrent request. The new query’s look at \(n\) keys still grows linearly with length.

It also eats your GPU:

\[ \text{bytes} = 2 \times L \times h_{\text{kv}} \times d_{\text{head}} \times n_{\text{tokens}} \times \text{bytes per value} \]

The leading 2 is for K and V. Note what is absent: the number of query heads, and the model’s parameter count. Note what is present: the number of tokens, linearly, and per concurrent request.

model shape layers KV heads head dim context KV cache
7B-class 32 32 128 4,096 2.1 GB
7B-class 32 32 128 32,768 17.2 GB
70B-class, hypothetical MHA 80 64 128 32,768 85.9 GB
7B-class 32 32 128 100,000 52.4 GB

The table gives the raw bytes of the K/V tensors with 16-bit values, every layer full-attention and each request’s cache separate (GB decimal; 52.4 GB is about 48.8 GiB). Quantisation, GQA, sliding windows, paging and working space change what is actually allocated. A 7B model in fp16 is 14 GB of weights. At 100,000 tokens of context its KV cache is 52 GB — nearly four times the model itself, for one conversation. This is the arithmetic behind every “why is long context so expensive?” question, and it is why serving systems care so much about how that memory is partitioned and reused (Chapter 31).

The third row is marked hypothetical on purpose: today’s 70B-class models do not have 64 KV heads. The original LLaMA 65B still used full multi-head attention, but Llama 2 70B, Llama 3, Mistral’s models and their descendants use 8, and the next section explains why. Redo the arithmetic with \(h_{\text{kv}} = 8\) and the 85.9 GB becomes 10.7 GB. That one substitution is the difference between a single 32k-context request whose cache is larger than an entire 80 GB card (85.9 GB) and one that fits in an eighth of it (10.7 GB). The weights still need a node — 140 GB in fp16 — but on the four-card node of exercise 9 the 180 GB left over holds, counting only those two items, 16 requests at the same context instead of 2 — eight times as many. In a real service, working space and the deployment layout lower that ceiling. It is the single most consequential number in this chapter.

KV cache calculator

GQA and MQA: fewer keys, same queries

The formula has an obvious lever. Nothing requires every query head to have its own key and value head — several query heads can share one set.

  • Multi-query attention takes this to the limit: all query heads share a single K/V head (Shazeer 2019).
  • Grouped-query attention is the compromise that won: query heads are divided into groups, one K/V head per group (Ainslie et al. 2023).

On the 32-layer, 32-head, 100,000-token model above:

scheme KV heads KV cache reduction
MHA 32 52.4 GB 1×
GQA 8 13.1 GB 4×
MQA 1 1.6 GB 32×

MQA’s 32× is tempting and often slightly too aggressive — in the GQA paper’s experiments MQA lost quality and trained less stably. GQA with 4 or 8 groups kept most of the saving and recovered nearly all the quality in those experiments, which is why Llama 2 70B, Llama 3, Mistral, and many models since ship with it. It is a rare case where the efficiency trick costs very little accuracy, which is why it became so common (still measure the quality–speed balance for your own model and task).

Notice, finally, that this is a decision you inherit rather than tune. You cannot flip an MHA checkpoint to GQA in your serving config. You can convert one — the GQA paper’s own recipe mean-pools the key and value heads of each group and then uptrains on roughly 5% of the original pretraining compute (Ainslie et al. 2023) — but 5% of a pretraining run is still a training run, with a cluster, a budget and days to weeks depending on the model’s scale. For every practical purpose, “how big is the KV cache” belongs in your model-selection checklist, not in your optimisation backlog.

Other ways out of the quadratic term and the growing cache. The techniques in this section fall into three branches. Kernel optimisations that speed up full attention — FlashAttention, say — leave the result unchanged and only reduce memory traffic. Sparse or windowed attention lets each token look at a subset only: in sliding-window attention a token sees just the last \(w\) tokens, and if the implementation drops older K/V entries that layer’s cache stays bounded by \(w\) (the first Mistral 7B used it in every layer, Gemma 2 and 3 in some of them). Methods that compress the cache or replace attention go further: DeepSeek’s MLA (multi-head latent attention) compresses keys and values into a low-dimensional latent vector and caches that; linear attention and state-space layers — Mamba, Gated DeltaNet — summarise the past in a fixed-size state and never build the \(n \times n\) matrix. Some current models mix these: of the 40 layers of qwen3.6:35b-a3b, measured in this book, 30 are Gated DeltaNet and only 10 are full attention (the layer_types field of its official config, checked on 15 October 2026), so the KV-cache formula above applies to those 10 layers only, while the other 30 keep a fixed-size state.

7.7 Mixture of Experts

Two thirds of every transformer block is the feed-forward sublayer, and every token goes through all of it. A mixture of experts asks whether it has to.

Replace the single feed-forward network in a block with \(E\) copies — the experts — plus a small learned router that scores the experts for each token and sends it to the best \(k\) — \(k = 2\) in the example below, \(k = 8\) in many recent models. Only those \(k\) run. The other \(E - k\) sit idle for that token, and a different token in the same sentence may take a different pair. This is the simplest MoE; in real models the width of the experts, shared experts that every token visits, and the number of layers that use MoE can all vary (below).

mixture of experts — replaces one feed-forward network with many, routing each token to a few of them, so total parameters and per-token compute stop being the same number. ↩︎

One token, eight experts, two of them run
token Router scores 8 expert 1 expert 2 0.62 expert 3 expert 4 expert 5 expert 6 0.38 expert 7 expert 8 + output n × d 8 EXPERTS, TOP-2 ROUTING, d = 4096 TOTAL 36.5B — all of it must be resident ACTIVE 10.7B — this sets per-token compute You rent the GPU for the first number. Every token picks its own pair, so a different word in the same sentence takes a different route.

The router is a small learned layer that scores the experts and normalises the top \(k\). Six of the eight boxes do no arithmetic for this token — but all eight still occupy memory.

This severs a link that had held since the beginning: total parameters and per-token compute stop being the same number. Take \(d = 4096\), 32 layers, 8 experts, top-2 routing:

total parameters active per token
dense 6.4B 6.4B
MoE, 8 experts, top-2 36.5B 10.7B

In this hypothetical table (embeddings excluded) the MoE model holds 5.7× the parameters of the dense one and uses 1.67× the active parameters per token. Its memory bill follows the first number; the second is a first-order estimate of the per-token arithmetic, not a latency figure. What you actually wait for also depends on how evenly the routing spreads across experts, on the communication that spreading costs, and on batch size and bandwidth.

You provision memory for the total and do arithmetic for the active. All 36.5B parameters must be reachable — the router’s choice is per token, so any expert may be needed at any moment. Those weights can be split across several GPUs, quantised, or partly moved to CPU memory; each choice changes memory, communication and latency together. A “small active parameter count” in a launch post tells you about speed, not about what you need to buy. This is the single most misread number in open-weight model announcements, and it matters directly when you cost out self-hosting in Chapter 31.

The idea goes back to the early 1990s; its sparse, gated form at large scale is due to Shazeer et al. (2017), and the Switch Transformer simplified the routing to one expert per token and pushed the scale further, an important step in the approach spreading (Fedus et al. 2022). The hard part in practice is not the routing rule but load balancing: nothing stops the router from sending almost every token to the same few experts and starving the rest, which wastes the parameters you paid for and can collapse training.

The router itself is a single matrix. For token \(x\) it produces \(s = xW_r \in \mathbb{R}^{E}\), keeps the top \(k\) scores, renormalises them, and mixes only those experts:

\[ \operatorname{MoE}(x) = \sum_{e \,\in\, \operatorname{top-}k(s)} \tilde g_e\, E_e(x), \qquad \tilde g_e = \frac{\exp s_e}{\sum_{e' \in \operatorname{top-}k(s)} \exp s_{e'}} \]

The classical fix for load balancing is an auxiliary loss penalising uneven usage — an explicit correction for a failure mode the main objective is happy to walk into. But it is a second objective fighting the first, and some large open MoEs demoted it: DeepSeek-V3 balances mainly with a per-expert bias added to the routing score and adjusted online, carrying no gradient of its own. “Mainly” is the operative word — it keeps a small complementary sequence-wise balancing loss alongside that mechanism. Check which mechanism a model actually uses, and in what proportion, rather than assuming the auxiliary loss is either present or gone.

The 29% active fraction in the table above is a property of this configuration, not of MoE. The direction the field went is finer granularity: instead of 8 large experts with top-2, use 64–256 narrow experts with top-8, plus one or two shared experts that every token visits. Same total, far smaller active fraction — which is how a 35B-total model ends up near 3B active rather than 10.7B, as qwen3.6:35b-a3b does in Chapter 12. When you read a launch post, the ratio to look for is active over total (it depends on the number of experts, top-\(k\), expert width and shared parts), and across well-known shipped models it has ranged from about 30% down to 3%.

7.8 Scaling laws and what they predict

Transformers turned out to show, unusually cleanly and over an unusually wide range, a property that is not unique to them — LSTMs follow power laws too, only with a worse constant (Kaplan et al. 2020): make them bigger and their loss falls, over the range studied, in a way you can predict in advance. Loss follows a power law in parameters, in data, and in compute — a straight line on log-log axes over many orders of magnitude (Kaplan et al. 2020). These are empirical fits: beyond the measured range, or with different data, architecture or training setup, the coefficients and exponents can change. That predictability is the reason anyone was willing to spend nine figures on a training run.

The first version of the analysis suggested growing the model much faster than the data. It was wrong, in a way that matters and that is worth understanding because the error was methodological rather than arithmetical. The Chinchilla authors’ explanation is that every run used the same learning-rate schedule, whatever its number of training tokens, so the loss at smaller token counts was read off runs whose schedule had not finished. Later re-analyses point to further causes — how the last layer’s compute is counted, counting parameters without the embeddings and analysing at small scale, warmup length, and optimiser settings not re-tuned for scale — so treat the schedule as one reason, not the whole story. Redoing it properly gave a different answer — for a fixed compute budget, parameters and tokens should grow at roughly the same rate, about 20 tokens per parameter (Hoffmann et al. 2022). That is an approximate ratio that minimises loss under a fixed training compute budget; it takes no account of data quality or of inference cost (the second qualification below):

model parameters training tokens tokens per parameter
Gopher 280B 300B 1.1
Chinchilla 70B 1.4T 20

Chinchilla is a quarter of Gopher’s size, used the same compute budget, and beat it broadly. A generation of models had been badly undertrained, and nobody knew because everyone was reading the same scaling law.

Two honest qualifications, both of which you will need.

The laws predict loss, not capability. Loss is average negative log-likelihood on held-out text. It is smooth and it is not what you care about. Whether a model can follow a multi-step instruction, or refuse a bad request, or do arithmetic reliably, is not a quantity these curves forecast — and the relationship between a smooth loss curve and a jagged capability curve remains genuinely disputed. Lower loss helps on many tasks, but do not let a scaling plot tell you a model will be able to do your task; evaluate the target task separately.

Compute-optimal is not deployment-optimal. Chinchilla optimises training compute for a given loss. If you are going to serve a model to millions of users, inference cost can dominate its lifetime bill, and it is rational to train a smaller model on far more tokens than the law recommends — paying more up front to pay less forever. This is exactly what the Llama-family models do, and why “trained past Chinchilla-optimal” is a design decision rather than a mistake. Which point is right depends on the expected request volume, the training and serving hardware, and the quality target together. We return to the full training pipeline in Chapter 8.

7.9 Lab: implement attention, then profile it

The point of this lab is to feel the quadratic term rather than read about it. Work in NumPy first; reach for a framework only in step 3.

  1. Implement scaled dot-product attention from the equation in Section 7.2 — projections, scores, scaling, optional mask, softmax, weighted sum. Write your own softmax and subtract the row maximum before exponentiating. No attention library.
  2. Reproduce this chapter’s numbers exactly. Use the same \(E\), \(W^Q\), \(W^K\), \(W^V\), and check that içti’s attention row comes out [0.4230, 0.5012, 0.0759]. If it does not, your bug is in the transpose or in the axis you are normalising over — those are the only two candidates.
  3. Verify against a reference. Compare to torch.nn.functional.scaled_dot_product_attention on random inputs at several shapes, with the same mask, scale and dropout_p=0 (in this API a boolean attn_mask uses True for allowed positions), and report the maximum absolute difference along with the backend used. In fp32, anything above about \(10^{-5}\) is very likely a real disagreement; choose a looser tolerance for fp16 and bfloat16.
  4. Confirm permutation equivariance. With no mask, shuffle the tokens, run again, un-shuffle the outputs, and report the maximum difference. It should stay within your data type’s numerical tolerance (this chapter got machine-epsilon level). Then add a positional encoding and report the number again. Finally repeat the whole test with a causal mask and no positional encoding: equivariance breaks even though you added no position signal. Explain why to someone who has not read this chapter.
  5. Measure the exponent. Time attention at \(n = 128, 256, \dots, 4096\) with \(d = 64\), holding everything else fixed, after warm-up runs, with repeats and (on a GPU) synchronisation. Plot on log-log axes and fit local slopes. The theoretical cost is \(O(n^2 d)\), but the wall-clock slope need not be exactly 2: fixed costs intrude at small \(n\), memory and kernel changes at large \(n\); explain which shows up where.
  6. Measure the memory. First compute the expected bytes of the raw \(n \times n\) score tensor and compare against the table in 1; then measure the process’s peak allocated and reserved memory separately (temporaries and the allocator can make them differ). Find the \(n\) at which your machine refuses. Report it.
  7. Add a causal mask, then prove it works behaviourally rather than by inspection: run on a sequence, append one token, run again, and confirm the outputs for the original positions are unchanged. This invariant is what makes the next step legal.
  8. Build a KV cache. With the same starting prefix, weights, data type and hardware, generate 512 tokens one at a time, first recomputing everything each step, then caching K and V. Time prefill separately from the decode steps. Plot per-token latency for both and report the speedup at token 512. Compute the cache’s size in bytes from the formula in 1 and check it against what your process actually allocated.

Deliverable: the log-log plot with your fitted exponent, the two per-token latency curves, and a short note on the largest sequence length your hardware could handle — with the specific resource that ran out named.

7.10 Intuition checks

Answer each before you open it. These are the six places where the standard mental model of a transformer is wrong in a way that costs money.

Usually false, and believing it sends you optimising the wrong thing.

Per layer the two terms are \(24nd^2\) for the projections and feed-forward against \(4n^2d\) for attention. They are equal at \(n = 6d\) — so for a 4,096-wide model the crossover is at 24,576 tokens. At a 4,096-token context attention is 14% of the arithmetic and the feed-forward and projection matrices are 86% (assuming full multi-head attention and \(d_{\text{ff}} = 4d\), as operation counts).

The quadratic term is a wall you hit at long context, not a tax you pay at every length. Before optimising attention, work out which regime you are in; if \(n < 6d\) most of the arithmetic is in the feed-forward half. But \(n = 6d\) is not a wall-clock threshold: prefill and decode behave differently, and attention can still bottleneck on memory traffic. The measured profile has the last word.

Because they answer three different questions. Which “one matrix” you mean matters: using the same vector as both query and key (one shared projection) makes the score table symmetric, while a single learned \(M\) in \(x_i M x_j^{\top}\) can give directed scores but is expensive (see Two matrices, or one?). What follows is the shared-projection case.

\(Q\) is what a token is looking for; \(K\) is what it offers to be found by; \(V\) is what it actually contributes once found. \(Q\) and \(K\) together decide where information moves; \(V\) supplies what moves. Tie them and a token’s advert becomes identical to its query, so the score table is forced to be a symmetric Gram matrix (a table of all pairwise dot products). Row normalisation can still produce different weights in the two directions; what a head loses is the ability to choose them, since both directions read one entry. It also biases every token toward itself, since \(x \cdot x\) is the largest score in a row unless a neighbour is longer; that second effect is a tendency, not a guarantee.

The worked example in Section 7.2 shows this concretely: the score that makes the example work could not have been produced by any unprojected dot product of the inputs.

Under this example’s assumptions (full MHA, fp16, one GPU): memory, and not the model’s. A 7B model in fp16 is 14 GB of weights. Its KV cache at 100,000 tokens is 52.4 GB — nearly four times the model itself, for one conversation — and that figure is per concurrent request, while the weights are shared across all of them.

This is why the KV cache, not the parameter count, usually decides how many users a GPU can serve, and why GQA became universal: dropping from 32 KV heads to 8 takes the same cache from 52.4 GB to 13.1 GB, a 4× saving with very little quality cost in published experiments. Once GQA, KV quantisation or several GPUs come in, the first limit may be bandwidth or another memory item.

If you answered “the \(n \times n\) score matrix”, that was the right answer before FlashAttention, which never materialises it.

Those are two different questions with two different answers, which is the entire point of the architecture.

You must hold all 35B parameters in memory, because any token might route to any expert. But only the top-\(k\) experts run per token — about 3B active — so per-token arithmetic tracks the small number. Total parameters set your memory bill; active parameters set your per-token arithmetic; latency also depends on bandwidth, routing balance and the deployment layout.

Section 12.9 measured exactly this: qwen3.6:35b-a3b classifies a document in 0.353 s against 0.365 s for a dense 9.7B model and 0.592 s for a dense 27.4B (on the same laptop, through ollama, all four quantised to 4 bits, median per document) — in that setup faster than both, while scoring higher than both, from a 23.9 GB quantised file on disk, which is not the same as working memory. On other hardware or another task the order can change.

With learned absolute positions, nothing good: there is no row 1,025 in the table, so the model has no vector to add. GPT-2’s standard forward pass does not accept a longer sequence — the failure is a lookup error, not a degradation (implementations can work around it by truncating or with a sliding window).

This is the concrete reason the field moved to RoPE and ALiBi. RoPE encodes position as a rotation, so the score between two tokens depends only on their offset. Measured in Section 7.4: the same query-key pair at a fixed offset, evaluated at absolute positions 0, 100, 1000 and 5000, gives scores that vary by at most \(5 \times 10^{-13}\) — against a spread of about 26 when the position vector is added to the embedding instead. When position enters only as an offset, there is no row 1,025 to be missing.

Note what this does not buy: extrapolating far beyond the training length still often degrades quality; it has to be measured. It removes a hard failure and replaces it with a soft one.

Because it is the only thing every sublayer reads from and writes to, and its width never changes.

Attention and the feed-forward network both take the residual stream in and add a correction back. That is why stacking blocks is trivial — block 12 looks exactly like block 1. In pre-norm the normalisation sits at each branch’s input: the branch reads the residual stream (all the sums so far) normalised, and adds its own output back un-normalised.

The practical consequence: when you inspect what a transformer “knows” at some depth, you are reading the residual stream, not any one layer’s output. So does every technique that extracts a vector from a mid-network layer — including the frozen-encoder embeddings Section 12.3 compares against.

7.11 Interview questions

Transformer questions in interviews are usually a test of whether you can do arithmetic about a system, not whether you can recite an architecture diagram. The answers below are the ones with numbers attached.

If you get asked to “explain attention”, resist the temptation to start with the softmax formula. Start with what it is for — every token can look at every token it may see and learn how much to take from each — then derive the formula: Q/K scores, softmax weights, the V mixture. That order makes the link between the computation and its job visible.

  1. Explain self-attention to someone who knows regression. It is a learned weighted average. Each token proposes a query, every token advertises a key, the dot products become weights through a softmax, and the output is those weights applied to the values. The unusual part: training learns the parameters of the Q/K/V projections, but the attention weights are not fitted once — they are recomputed for each input at inference time — which is exactly what makes the same word get a different vector in a different sentence. Then name the scaling: divide by \(\sqrt{d_k}\), because dot products of \(d_k\)-dimensional vectors with independent unit-variance components grow like \(\sqrt{d_k}\) and an unscaled softmax drifts towards saturation.

  2. Why is there a feed-forward network after attention? Because attention mixes tokens but, given its weights, does so linearly in the values — the non-linearity it does have lives in how the weights are computed, not in what it does with what it moves. The feed-forward sublayer is where each token gets non-linear processing of its own; on one reading \(W_1\) is a bank of detectors and \(W_2\) is what each detector broadcasts when it fires. Get the proportion right, because most candidates have it backwards: in the standard \(d_{\text{ff}} = 4d\) block used throughout this chapter, two thirds of a transformer block’s parameters are the feed-forward sublayer, not attention. The ratio follows from that choice of \(d_{\text{ff}}\) and from full-width attention; recompute it for a block with GQA, SwiGLU or MoE.

  3. How do you estimate whether a model fits on your GPU? As a first approximation, two numbers, added. Weights: parameters × bytes per parameter (7B at fp16 = 14 GB). KV cache: \(2 \times \text{layers} \times \text{KV heads} \times \text{head dim} \times \text{tokens} \times \text{bytes}\), per concurrent request. Give the worked case: a 7B model at 32,768 tokens needs 14 GB of weights and 17.2 GB of raw cache, so context has already cost more than the model. Then leave room for activations, attention workspace, the allocator and the deployment; the real requirement is higher. The senior version of the answer names the lever — GQA, quantisation, or a shorter context — and says which one you would pull first.

  4. What does FlashAttention actually change? Not the mathematics: it computes the same function, and only the last digits move, because the tiled sums re-associate floating-point addition. It changes what gets written to memory: the \(n \times n\) score matrix is never materialised, computed instead in tiles that stay in fast on-chip memory. At long sequences that moves attention from memory-bandwidth-bound towards compute-bound; measure the wall-clock effect for your sequence length, head size and hardware. The number that motivates it: at 100,000 tokens one score matrix for one head at one layer is 18.63 GiB in fp16. Multiply by 32 heads and 32 layers and the naive version is not slow, it is impossible.

  5. Encoder-only, decoder-only, or encoder-decoder? Inside self-attention the distinction is which entries of the attention matrix are allowed to be non-zero; an encoder-decoder adds cross-attention between two stacks. Encoder-only sees everything both ways and is the right choice when you need a representation of a complete text — classification, retrieval, and every result in Chapter 12. Decoder-only is causally masked and is what you need to generate. Encoder-decoder keeps both and lets the decoder attend to the encoder, which is why it persists in translation. Say the practical version: for classification with labels, an encoder is still a strong starting point (Section 12.9 has this book’s measurement); make the choice on quality, latency and serving load for your own task.

  6. Why did transformers replace RNNs? Not mainly because they are more expressive — because they are parallelisable in training. An RNN must process token \(t\) before token \(t+1\); a transformer computes all positions of a layer at once (layers still come in sequence), which converts a sequential dependency into a matrix multiplication and lets you use a GPU properly. The cost is the quadratic term of dense attention and a finite context window (an RNN’s context is unbounded in principle, but squeezed into a finite hidden state). Be precise that this was an engineering win that unlocked scale, and that scale is what produced the capabilities — the argument Section 7.8 makes quantitatively.

7.12 Where to go deeper

Three outside visualizations pair well with this chapter: the Illustrated Transformer for step-by-step diagrams, Georgia Tech’s Transformer Explainer, which runs a live GPT-2 in your browser, and Brendan Bycroft’s 3D LLM walkthrough.

If you would rather watch, three sources are a good start: 3Blue1Brown’s Transformers, the tech behind LLMs and Attention, step-by-step for the geometry, Karpathy’s Let’s build GPT: from scratch for the same architecture typed out line by line, and — if that pace is too fast — Umar Jamil’s slower Coding a Transformer from scratch in PyTorch, which includes training and inference. Stanford’s CS25: Transformers United is a seminar series where the authors present their own papers.

The primary sources are unusually readable, and after this chapter you have enough to read them directly. Start with Vaswani et al. (2017) — a short paper whose core equations you now know; its training, optimisation and experimental details are worth reading on their own. Then, roughly in the order the chapter needed them:

topic paper
the architecture Vaswani et al. (2017)
residual connections He et al. (2016)
layer norm, and moving it Ba et al. (2016); Xiong et al. (2020)
what heads actually learn Clark et al. (2019); Michel et al. (2019)
rotary positions Su et al. (2024)
position by distance penalty Press et al. (2022)
exact, IO-aware attention Dao et al. (2022)
shrinking the KV cache Shazeer (2019); Ainslie et al. (2023)
sparse experts Shazeer et al. (2017); Fedus et al. (2022)
scaling laws, and their correction Kaplan et al. (2020); Hoffmann et al. (2022)

If you want one book that derives the whole stack with runnable code rather than a reading list, Dive into Deep Learning (Zhang et al. 2023) is free online and covers attention and transformers from the linear algebra up.

7.13 Exercises

  1. In the attention explorer, the pronoun head splits its weight roughly 46/34 between coffee and robot — leaning to the coffee, but only just (and we built that matrix by hand). Write three one-word substitutions for ready that you would expect to push the reading clearly one way, and for each state the assumption about the world, not about grammar, that supports it. Then discuss what this says about the limits of reading meaning off a single attention map.
  2. Attention is order-blind on its own. Take the word-order pair from Chapter 4 that scored cosine 1.0000 under averaging. Show how the token outputs of unmasked, position-free self-attention change under permutation; then explain which symmetry breaks when you add a position signal or a causal mask, and under which read-out or pooling rule the two sentences’ representations come apart.
  3. A model with 32 layers, 32 heads, and head dimension 128 serves a 100,000 token context. Assuming full multi-head attention (32 KV heads) and 16-bit K/V, compute the raw KV cache in GB and GiB. Then compare the same shape with 8-KV-head GQA and explain why grouped-query attention exists.
  4. In the worked example, the fourth embedding dimension is 0.5 for all three tokens and, because the rows of \(W^Q\) and \(W^K\) for it were chosen as zero, it contributed nothing to any attention weight. Work out how far that generalises, in two parts. (a) Prove that the key side always cancels: show that whatever row \(W^K\) assigns to a constant feature, its contribution is the same for every key \(j\) and therefore vanishes in the softmax. (b) Then show by counterexample that the query side does not. Give \(W^Q\) a non-zero row for the constant feature, recompute içti’s attention weights, and report how far they moved. Explain in one sentence why the same feature can matter in one projection and not the other.
  5. Multiply \(W^Q\) and \(W^K\) each by \(1/\sqrt{3}\), so the product of their scales is 1 instead of 3, and recompute içti’s attention row. Report the weights, and explain what the \(1/\sqrt{d_k}\) term is protecting against in terms of what you observe.
  6. The chapter claims \(h\) heads of width \(d/h\) cost the same as one head of width \(d\). Verify this by counting parameters and multiplications for \(d = 512\) at \(h = 1, 8, 64\), assuming \(d_k = d_v = d/h\). Also write down the naive implementation’s \(h n^2\)-entry score tensor and the per-head score-rank bound. Then state what is genuinely lost as \(h\) grows, whether equal matrix-multiply FLOPs mean equal latency, and at what point you would expect it to hurt.
  7. Take the relative-position table for sinusoidal encodings. The mean at offset 32 is higher than at offset 16. Explain why from the formula \(\sum_i \cos(k\omega_i)\), and say what this means for the common claim that sinusoidal encodings make nearby positions more similar.
  8. Reproduce the RoPE experiment: one fixed random query and key, scored at offset 5 from absolute positions 0, 100, 1,000 and 5,000. Report the spread. Then do it with an additive encoding — say exactly which position vectors and projections you use — and report that spread. In one paragraph, separate what the two numbers show mathematically from what they do not show about quality on a real task.
  9. A 70B-class model has 80 layers, 64 query heads, and head dimension 128. Its fp16 weights are 140 GB, so it is served on a node of four 80 GB accelerators — 320 GB total. Using the KV cache calculator, and counting only the weights and the raw KV cache, find the theoretical longest context you can serve to four concurrent users, first with 64 KV heads and then with 8, as current 70B-class models have. Also say why this is not a usable limit once the weights’ split across cards, activations, working space and a safety margin are added. Report both lengths and the ratio between them. Then name every quantity in the cache formula that changed and every one that did not.
  10. An announcement describes a model as “120B total parameters, 5B active”. A colleague concludes it will run on a 24 GB consumer GPU because 5B is small. Write the two-sentence correction, with the arithmetic; add the latency and quality trade-offs that quantisation, several GPUs or offloading weights to CPU memory bring while cutting memory.
  11. Gopher used 1.1 tokens per parameter and Chinchilla used 20. Both claims came from a scaling-law analysis. Explain what changed between them — the compute budget, whether parameters were counted with the embeddings, the learning-rate schedule, the measured range — and state what you would need to see in a new scaling-law paper (its experimental range, its uncertainty) before you reallocated a training budget on the strength of it.