KB KEDBYTE TECHNOLOGIES PRIVATE LIMITED
CHAPTER
48

The Transformer

Part H · Games and Machine Intelligence|32,603 words|about 142 min read|Volume 5
Fast-moving material. Figures, model names, prices and version numbers in this chapter were verified in August 2026. Claims are separated into established fact, active research and marketing claim. Re-check anything you intend to rely on.

48.0 What this chapter gives you#

  1. You will be able to say what came before the transformer, why it worked, and the exact reason the field abandoned it in about eighteen months.
  2. You will be able to explain attention in ordinary words to somebody who has never seen a formula, using a sentence whose meaning flips on one adjective.
  3. You will be able to write the scaled dot-product attention formula from memory and say what every symbol in it does.
  4. You will be able to compute a full attention pass by hand on a four-token sentence and get numbers that match a real implementation.
  5. You will be able to explain the causal mask, draw it as a matrix, and say what changes when a model does not use one.
  6. You will be able to explain multi-head attention, work out the tensor shapes for a real model configuration, and say honestly what heads do and do not do.
  7. You will be able to explain why attention has no sense of order at all, and describe sinusoidal, learned, ALiBi and rotary position encoding.
  8. You will be able to draw a whole transformer block from memory, with residual connections and normalization in the right places, and say why.
  9. You will be able to trace one prompt through a real model as a series of tensor shapes, from token ids to a probability over the vocabulary.
  10. You will be able to explain the quadratic cost of attention, compute a real key-value cache size, and name the fixes that are in production today.

48.1 What came before, and why it was not enough#

PLAIN48.1.1 in simple words#

  1. Before 2017, the standard way to handle a sentence was to read it one word at a time, left to right, the way a person reads aloud.
  2. The machine that did this is called a recurrent neural network, usually shortened to RNN. Recurrent just means it feeds its own output back into itself.
  3. It kept a single running summary of everything it had read so far. That summary is called the hidden state. It is a list of numbers.
  4. Read word one, update the summary. Read word two, update the summary again. Keep going to the end.
  5. That is the whole design. One summary, updated once per word, in order.
  6. It worked. It was the best thing anybody had for sequences for about twenty-five years.
  7. It had two problems, and the second one killed it.
  8. Problem one: information from early words faded. By word two hundred, word three had almost no influence left on the summary.
  9. Problem two, and this is the one that mattered: you cannot do the work for word ten until you have finished the work for word nine.
  10. That is a rule you cannot break, because word ten literally needs the summary that word nine produced.
  11. A modern graphics chip has tens of thousands of small calculators that all want work at the same time.
  12. A design that insists on doing one word, then the next, then the next, cannot feed those calculators. They sit idle.
  13. So training was slow, and it stayed slow no matter how much hardware you bought.
  14. The transformer’s central idea is to throw away the left-to-right rule and let every word look at every other word at once.
  15. Once you do that, all the words can be processed in parallel, and the machine finally has enough work to do.

PLAIN48.1.2 a picture in your head#

  1. Imagine you are told a long story down a telephone line, one word at a time, and you are not allowed to write anything down.
  2. You are allowed to keep exactly one small notepad in your head with, say, twenty facts on it. That is the hidden state.
  3. Every new word arrives and you must decide what to erase to make room.
  4. At word ten you still remember the beginning clearly.
  5. At word two hundred the beginning has been overwritten many times. It is not gone in principle, but almost nothing of it survives.
  6. Now imagine the person telling the story is speaking as fast as you can process. You cannot skip ahead. You cannot ask a friend to handle the second half while you do the first.
  7. The story is strictly serial. That is the cost.
  8. The transformer is the opposite arrangement. The whole story is printed on a sheet of paper in front of you, all at once, and you may look at any word whenever you want.

Where this comparison breaks: a real hidden state is not a list of twenty readable facts. It is a vector of a few hundred to a few thousand numbers with no individually meaningful entries. And nothing “decides what to erase” in a deliberate way. A learned matrix multiplication and a gate value between 0 and 1 produce that effect, and the effect looks like forgetting from outside.

PLAIN48.1.3 a worked example#

  1. Here is why the fading happens, with real arithmetic.
  2. In training, the correction signal for word 3 has to travel backwards from word 200 through 197 update steps.
  3. At each step it gets multiplied by a number. Call that number the shrink factor. If it is less than 1, the signal shrinks.
  4. Here is what repeated multiplication does:
Shrink factor After 10 steps After 100 steps
0.5 0.00098 7.9e-31
0.8 0.107 2.0e-10
0.9 0.349 2.7e-05
0.99 0.904 0.366
  1. At a shrink factor of 0.9, after 100 steps the signal is 0.0000266 of its original size. It has effectively vanished.
  2. That is the vanishing gradient problem, and over a sequence it is vanishing over time steps rather than over layers.
  3. Now the opposite. If the factor is 1.1 instead of 0.9, after 100 steps the signal is 13,780 times bigger. Training explodes and produces nonsense.
  4. So you need the factor to sit very near 1.0 for hundreds of steps in a row, which is a very hard thing to arrange by accident.
  5. Now the parallel problem, also with real arithmetic.
  6. Take a sequence of 1,000 tokens.
  7. An RNN needs 1,000 steps that must happen one after another. Nothing about the sequence can be done out of order.
  8. A transformer needs 1 step for the whole sequence, because the operation is one big matrix multiplication over all 1,000 positions at once.
  9. That is not a small speed-up. It is the difference between a chip running at a few per cent of its capacity and a chip running near its limit.

PLAIN48.1.4 what is really happening inside#

  1. The RNN update is one line: new summary equals a function of the old summary and the new word.
  2. Written plainly: h_t = tanh(W_h h_{t-1} + W_x x_t + b).
  3. h_t is the summary after word t. x_t is the word. W_h and W_x are learned matrices. tanh squashes everything into the range -1 to +1.
  4. Notice that h_{t-1} appears on the right-hand side. That is the whole dependency chain, and it is unbreakable.
  5. To learn, you unroll this into a chain 200 links long and push the error backwards along it, which is the repeated multiplication just shown.
  6. The LSTM, invented by Sepp Hochreiter and Jürgen Schmidhuber in 1997, was the fix for the fading.
  7. It added a second track, called the cell state, that flows forward with addition rather than repeated multiplication.
  8. Addition does not shrink a signal, so the correction can travel a long way without dying.
  9. It also added gates: small learned circuits that output a number between 0 and 1 for each position, deciding how much to keep, how much to let in and how much to output.
  10. The GRU, from Kyunghyun Cho and colleagues in 2014, is a simpler version with two gates instead of three, and usually performs about the same.
  11. These fixes worked well enough for real products. Google Translate switched to an eight-layer LSTM system in September 2016.
  12. But gates fix the fading. They do not fix the serial rule. Step t still needs step t-1. That part is structural.
  13. Meanwhile a second problem appeared in translation. The standard design was an encoder-decoder: one RNN reads the whole source sentence into a single vector, and a second RNN writes the translation out of it.
  14. Everything in a fifty-word sentence had to fit through one fixed-size vector. That vector is a bottleneck, and it is the same size whether the sentence has five words or fifty.
  15. In September 2014, Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio published a fix. Instead of one vector, keep all the encoder’s per-word states, and let the decoder look back and pick out the relevant ones.
  16. That “look back and weight what matters” step is attention. It was invented as an accessory bolted onto an RNN.
  17. The 2017 transformer paper’s move was to delete the RNN and keep only the accessory.

TECHNICAL48.1.5 the engineer’s version#

  1. The vanishing and exploding gradient problem in recurrent networks was identified by Sepp Hochreiter in his 1991 diploma thesis at the Technical University of Munich, and analysed formally by Yoshua Bengio, Patrice Simard and Paolo Frasconi in 1994.
  2. The mechanism is the chain rule applied through time. The gradient of the loss at step T with respect to the state at step t contains the product of Jacobians from t to T.
  3. If the largest singular value of the recurrent Jacobian is consistently below 1, that product decays geometrically. Above 1, it grows geometrically.
  4. With tanh activation the derivative is at most 1.0 and typically well below it, so decay is the common case. With sigmoid gates the derivative peaks at 0.25.
  5. Exploding gradients are cheaply fixed by gradient clipping, which rescales the gradient vector when its norm passes a threshold. Vanishing gradients are not fixable that way, because you cannot amplify information that is gone.
  6. LSTM, from Hochreiter and Schmidhuber, Neural Computation volume 9 issue 8, 1997, introduced the constant error carousel: an additive cell state path whose local derivative is close to 1.
  7. The forget gate was not in the original LSTM. Felix Gers, Jürgen Schmidhuber and Fred Cummins added it in 1999 and 2000, and the version with the forget gate is what everybody now means by “LSTM”.
  8. Parameter counts per unit differ: a vanilla RNN cell has one weight matrix pair, an LSTM has four, and a GRU has three.
Model Year Gates Sequential steps
Elman RNN 1990 0 n
LSTM 1997 3 n
GRU 2014 2 n
Transformer 2017 none 1
  1. Sequence-to-sequence learning with a fixed context vector was published by Ilya Sutskever, Oriol Vinyals and Quoc Le in September 2014, using a 4-layer LSTM with 1,000 units per layer and a 1,000-dimensional context vector.
  2. Additive attention, also called Bahdanau attention, was published the same month by Bahdanau, Cho and Bengio in the paper “Neural Machine Translation by Jointly Learning to Align and Translate”. It scores each encoder state with a small feed-forward network and softmaxes the scores.
  3. Multiplicative attention, from Minh-Thang Luong, Hieu Pham and Christopher Manning in 2015, replaced that small network with a dot product. That is the direct ancestor of the 2017 formula.
  4. Google Neural Machine Translation, deployed in September 2016, used 8 encoder and 8 decoder LSTM layers with attention, residual connections and model parallelism across 8 GPUs. It cut translation errors substantially and was still fundamentally serial.
  5. The training-time argument is about arithmetic intensity. An RNN step at batch size B computes a matrix times a matrix of shape [B, d] by [d, d]. With B = 32 and d = 1024 that is 67 MFLOP, then a synchronization, then the next step.
  6. A transformer layer at sequence length 1,000 computes [1000, d] by [d, d], which at d = 4096 is 33.6 GFLOP in a single call, with no dependency inside the call.
  7. Modern accelerators need large matrix-matrix operations to reach peak. This is covered in Chapter 22, which explains why a GPU is a wide parallel machine rather than a fast serial one. Long thin operations leave most of the chip idle.
  8. This is established fact and was stated openly in the abstract of the 2017 paper: the models are “more parallelizable and requiring significantly less time to train”.

WORDS48.1.6 remember these#

  1. Recurrent neural network — a network that reads one item at a time and keeps a running summary — a network with a cyclic dependency h_t = f(h_{t-1}, x_t), unrolled over time for training.
  2. Hidden state — the running summary — the vector carried from one time step to the next, of size equal to the layer width.
  3. Vanishing gradient — the correction signal dying out over distance — the geometric decay of a product of Jacobians whose singular values are below 1.
  4. LSTM — a recurrent cell with gates that remembers longer — long short-term memory, with input, forget and output gates and an additive cell state.
  5. GRU — a simpler gated cell — gated recurrent unit, with update and reset gates and no separate cell state.
  6. Encoder-decoder — one network reads, another writes — a sequence-to-sequence architecture with a context representation passed between them.
  7. Bottleneck vector — squeezing a whole sentence into one fixed list of numbers — the fixed-dimension context vector in pre-attention seq2seq.
  8. Attention — looking back and weighting what matters — a learned weighted average over a set of value vectors, with weights from a scoring function.

48.2 Attention in plain language, before any mathematics#

PLAIN48.2.1 in simple words#

  1. Read this sentence carefully: “the trophy did not fit in the suitcase because it was too big”.
  2. What does “it” mean here? The trophy.
  3. You did not need a rule book. You worked it out.
  4. Now read the same sentence with one word changed: “the trophy did not fit in the suitcase because it was too small”.
  5. What does “it” mean now? The suitcase.
  6. The sentence is identical except for the last word, and the meaning of “it” flipped completely.
  7. Stop and notice what you just did. To understand the word “it”, you looked back at other words in the sentence.
  8. You did not look back at all of them equally. “the”, “did” and “in” contributed almost nothing.
  9. You looked hard at “trophy”, “suitcase”, “fit” and the final adjective.
  10. Then you combined what you found and produced a meaning for “it”.
  11. That is attention. Look back, weight some things more than others, combine.
  12. The model does exactly this, for every word, at every layer, with the weights worked out by learned matrices rather than by a person.
  13. There is one important extra bit. The model does this for every word at once, not just for the confusing ones.
  14. So while it is deciding what “it” points to, it is simultaneously deciding what “big” applies to and what “fit” takes as its subject.
  15. And it does not know in advance which words are the confusing ones. Every position gets the same treatment.

PLAIN48.2.2 a picture in your head#

  1. Picture a meeting where a decision has to be made, and eight people are in the room.
  2. You are the chairperson. You ask a question.
  3. Everybody in the room has something to say. Some of it is highly relevant. Most of it is not.
  4. You do not take a plain average of all eight opinions, because that would drown the useful answer in noise.
  5. Instead you listen to each person, decide how much their answer bears on your question, and give each one a share of your final decision.
  6. Person 3 gets 60 per cent of your attention because they clearly know. Person 7 gets 25 per cent. The other six split the remaining 15 per cent.
  7. Your conclusion is a blend, weighted by relevance.
  8. That is one attention operation. Your question is the query. Each person’s relevance is a weight. Each person’s opinion is a value.
  9. Now the part that makes it a transformer: every person in the room is doing this at the same time, each with their own question, including you.

Where this comparison breaks: in a real meeting, relevance is judged by understanding. In attention, relevance is a dot product between two learned vectors, which is a geometric measure of alignment. It correlates with what we would call relevance because training pushed it that way, and it has no independent notion of what the question is about. Also, the weights always add up to exactly 1, so a position cannot decide that nothing in the sentence is relevant. It must spread its 1.0 somewhere, and it often dumps it on a token that carries no meaning, which is a real and measured effect.

PLAIN48.2.3 a worked example#

  1. Take the sentence with “big”, and consider the position holding “it”.
  2. Here is roughly how a trained model distributes attention from “it”, based on what published attention maps for coreference look like. These are illustrative figures to show the shape, not a measurement of one named model.
Sentence: the trophy did not fit in the suitcase
          because it was too big

Attention from the position holding "it":

  trophy    ############################  0.41
  suitcase  ############                  0.18
  big       ##########                    0.15
  fit       ######                        0.09
  it        #####                         0.08
  the       ##                            0.03
  other     ####                          0.06
                                          -----
                                          1.00
  1. Now swap “big” for “small”. The same position, in the same model, shifts:
Sentence: the trophy did not fit in the suitcase
          because it was too small

  suitcase  ###########################   0.39
  trophy    ############                  0.17
  small     ###########                   0.16
  fit       ######                        0.10
  it        #####                         0.08
  the       ##                            0.03
  other     ####                          0.07
                                          -----
                                          1.00
  1. One word changed at the end of the sentence, and the top weight moved from “trophy” to “suitcase”.
  2. This is the point. The weights are not fixed by the identity of the word “it”. They are computed fresh from the actual contents of this sentence.
  3. That is why the same word gets a different representation in different sentences, which is the thing older methods could not do.
  4. Chapter 47 covered embeddings, which are the fixed starting vectors for tokens. The embedding for “it” is one single vector, the same every time.
  5. Attention is what turns that one fixed vector into a different vector in every different context. That is the entire job.
  6. This sentence pair is not invented for this book. It is a Winograd schema, a test format proposed by Hector Levesque, Ernest Davis and Leora Morgenstern in 2011, named after Terry Winograd, who used an example of this shape in 1972.
  7. The point of the format is that no amount of word-frequency statistics solves it. You have to model the situation being described.

PLAIN48.2.4 what is really happening inside#

  1. Every token position holds a vector of numbers. Call it the position’s current understanding of itself.
  2. Attention runs in three moves, and every position does all three.
  3. Move one: each position emits a query, meaning “here is what I am looking for”. It is a vector, computed from the position’s own current vector.
  4. Move two: each position also emits a key, meaning “here is what I am, if anybody is looking”. Also a vector, also computed from the same input.
  5. Move three: each position emits a value, meaning “here is what I will contribute if you pick me”.
  6. Now the matching. Position 10’s query is compared against every key in the sentence, including its own.
  7. The comparison is a dot product: multiply the two vectors element by element and add up. A big result means the two vectors point in a similar direction.
  8. That gives one raw score per position. For a 12-token sentence, position 10 gets 12 scores.
  9. Those 12 scores go through softmax, which turns any list of numbers into positive numbers that add to exactly 1.
  10. Now they are weights. 0.41 on “trophy”, 0.18 on “suitcase”, and so on.
  11. Finally, take every position’s value vector, multiply it by that position’s weight, and add them all up.
  12. The result is position 10’s new vector. It now contains a blend of the whole sentence, weighted by what position 10 was looking for.
  13. That new vector replaces the old one, and the next layer does the whole thing again with different learned matrices.
  14. Repeat 32 times for a typical 8-billion-parameter model, and information from anywhere can reach anywhere, by many different routes.

TECHNICAL48.2.5 the engineer’s version#

  1. The operation just described is self-attention: queries, keys and values are all derived from the same input sequence.
  2. Cross-attention is the same operation where queries come from one sequence and keys and values from another. It is what connects the decoder to the encoder in a translation model.
  3. Self-attention is a set operation, not a sequence operation. It is permutation-equivariant, which is proved numerically in section 48.7.
  4. The pairwise score matrix has shape [n, n] for a sequence of length n. Every position produces one row.
  5. The distance between two tokens does not appear anywhere in the computation. Position 1 and position 4,000 are exactly as easy to connect as position 1 and position 2. That is the structural difference from an RNN, where the path length between positions i and j is |i - j| update steps.
  6. Attention weight maps are widely published and easy to produce with the output_attentions=True flag in Hugging Face transformers, or with the BertViz tool from Jesse Vig, 2019.
  7. Two honest warnings about reading those maps.
  8. The honest version: an attention weight is not an explanation. Sarthak Jain and Byron Wallace published “Attention is not Explanation” in 2019, showing that alternative attention distributions can produce identical predictions. Sarah Wiegreffe and Yuval Pinter replied with “Attention is not not Explanation” the same year, arguing the test was too strict. The dispute is not settled. Attention weights show where information flowed, not why.
  9. Second warning: a large fraction of attention mass in trained models lands on the first token or on punctuation. This is called an attention sink, and was analysed by Guangxuan Xiao and colleagues in 2023 in the StreamingLLM work. The softmax must sum to 1, so a head with nothing to say has to put its weight somewhere, and it learns to park it on a fixed harmless position.
  10. Established fact: attention computes content-dependent weighted averages, and this is measurable. Active research: what individual heads compute and whether their weights can be read as explanations. Marketing claim: any statement that the model “focuses on” or “pays attention to” something in the human sense.

WORDS48.2.6 remember these#

  1. Self-attention — every word looking at every word in the same sentence — attention where Q, K and V all derive from one input sequence.
  2. Cross-attention — one sequence looking at a different sequence — attention where Q comes from the target and K, V from the source.
  3. Query — what a position is looking for — a learned linear projection of the position’s hidden state, of dimension d_k.
  4. Key — what a position advertises about itself — a learned projection compared against queries by dot product, of dimension d_k.
  5. Value — what a position contributes if selected — a learned projection that is averaged with the attention weights, of dimension d_v.
  6. Softmax — turning any scores into shares that add to 1 — exp(x_i) divided by the sum of exp over all j, applied along the key axis.
  7. Attention weight — how much one word listened to another — one entry of the row-stochastic [n, n] matrix produced by softmax over scaled scores.
  8. Attention sink — a position that soaks up leftover weight — a token, typically the first, that receives large attention mass with no semantic role, arising from the sum-to-one constraint.
  9. Winograd schema — a sentence pair whose meaning flips on one word — a coreference test designed to be unsolvable by surface statistics.

48.3 Query, key and value, taught slowly#

PLAIN48.3.1 in simple words#

  1. Three words do all the work in attention: query, key and value.
  2. They sound abstract. They are not. They are three different views of the same token, made by three different learned recipes.
  3. The query is the question this position is asking. “I am a pronoun. I need to know what noun I stand for.”
  4. The key is the label this position wears. “I am a singular concrete noun, and I am the subject of the sentence.”
  5. The value is the content this position offers up if selected. “Here is my actual meaning, ready to be copied into whoever picks me.”
  6. Every single position produces all three, every time.
  7. So a position is simultaneously asking a question, advertising an answer, and holding content to hand over.
  8. The matching happens between queries and keys only. Values take no part in deciding who wins.
  9. Once the weights are decided, the values are what actually get blended and passed on.
  10. That separation is the useful part of the design. What makes something findable is allowed to be different from what it contains.
  11. And the crucial fact people miss: all three come from the same starting vector. There is no separate store of keys and values somewhere.
  12. Take one token’s vector, multiply it by matrix number 1 to get the query, matrix number 2 to get the key, matrix number 3 to get the value.
  13. Those three matrices are learned during training and frozen afterwards. They are ordinary parameters, in the sense of Chapter 47.

PLAIN48.3.2 a picture in your head#

  1. Picture a library, but a strange one where you may only ask one question and must get one answer back immediately.
  2. You walk in with a question written on a card. That is your query.
  3. Every book in the library has a small index card clipped to its spine. The card says what the book is about, in a compressed form. That is the key.
  4. The book’s actual pages are the value.
  5. You do not read any books. You hold your question card next to every index card in turn and score how well they line up.
  6. Book 3’s card scores highly. Book 7’s card scores moderately. Most cards score near zero.
  7. Then you do something no real library allows. You do not pick one book. You take a weighted blend of every book’s contents: 60 per cent of book 3, 25 per cent of book 7, and small slivers of the rest.
  8. That blend is your answer.
  9. The index card is deliberately not the book. A card can say “this is a cookbook about bread” while the book contains actual recipes.
  10. That is the whole reason keys and values are separate. Findability and content are different jobs.

Where this comparison breaks: in a real library the index card was written by a librarian who read the book. Here both the card and the book come from the same raw vector through two different learned matrices, so they are not independent. Also, a library has a fixed catalogue. Here the keys are recomputed from scratch for every input, so the same word gets a different key in a different sentence. And a real search returns one book. Attention always returns a blend and can never return nothing, because the weights are forced to sum to 1.

PLAIN48.3.3 a worked example#

  1. Now the projection, with numbers small enough to check on paper.
  2. We will use a four-token sentence: “the cat was hungry”.
  3. Each token arrives as a vector of 4 numbers. In a real model this would be 4,096 numbers, from the embedding table of Chapter 47. Here it is 4.
  4. Stack the four token vectors as rows. That is the input matrix X, shape [4 tokens, 4 dimensions]:
X (4 tokens x 4 dimensions)

           d0  d1  d2  d3
the         1   0   1   0
cat         0   1   0   1
was         1   1   0   0
hungry      0   1   1   0
  1. The honest version: real embedding values are floating-point numbers like -0.0113 or 0.0392, never neat integers. Integers are used here only so you can verify every multiplication by hand.
  2. Now the three learned matrices. Each is shape [4, 3], because we are choosing to project from 4 dimensions down to 3.
W_Q (4x3)        W_K (4x3)        W_V (4x3)
 1  0  0          1  0  0          2  0  1
 1  1  2          1  0  2          0  2  0
 0  1  0          1  1  1          1  1  2
 1  1  0          0  1  2          0  1  1
  1. To get the query for “hungry”, take the row for “hungry”, which is [0, 1, 1, 0], and multiply it by W_Q.
  2. Doing it column by column:
  3. Column 0: (0 x 1) + (1 x 1) + (1 x 0) + (0 x 1) = 1.
  4. Column 1: (0 x 0) + (1 x 1) + (1 x 1) + (0 x 1) = 2.
  5. Column 2: (0 x 0) + (1 x 2) + (1 x 0) + (0 x 0) = 2.
  6. So the query for “hungry” is [1, 2, 2].
  7. Do the same for every token and every matrix. The three results:
Q = X W_Q          K = X W_K          V = X W_V
the     1  1  0    the     2  1  1    the     3  1  3
cat     2  2  2    cat     1  1  4    cat     0  3  1
was     2  1  2    was     2  0  2    was     2  2  1
hungry  1  2  2    hungry  2  1  3    hungry  1  3  2
  1. Check one more by hand. The key for “cat” is row [0, 1, 0, 1] times W_K.
  2. Column 0: 0 + 1 + 0 + 0 = 1. Column 1: 0 + 0 + 0 + 1 = 1. Column 2: 0 + 2 + 0 + 2 = 4. So [1, 1, 4]. That matches the table.
  3. Notice the shapes. X is [4, 4]. Each W is [4, 3]. Each of Q, K, V is [4, 3].
  4. The number of tokens never changes. Only the width changes.

PLAIN48.3.4 what is really happening inside#

  1. There is no lookup, no dictionary, no search index. There are three matrix multiplications.
  2. A matrix multiplication of a [4, 4] by a [4, 3] is just 12 dot products, one per output cell, which is exactly what was done above by hand.
  3. The words “query”, “key” and “value” are borrowed from database retrieval. They are a naming choice, not a mechanism.
  4. Nothing enforces that W_Q learns “questions” and W_K learns “labels”. Those are the roles they end up playing because of where they sit in the formula.
  5. W_Q and W_K only ever appear together, inside a dot product. Their product Q K-transpose is what the model actually uses.
  6. That means the pair behaves as a single learned bilinear form. You could fold them into one matrix and get identical scores.
  7. They are kept separate because it costs far fewer parameters. Two matrices of [d, d_k] cost 2 d d_k parameters, while one full [d, d] matrix costs d squared. With d = 4096 and d_k = 128 that is 1,048,576 against 16,777,216.
  8. W_V is genuinely separate, because its output does not go into any score. It goes into the answer.
  9. In code, all three projections for a whole batch are usually done as one matrix multiplication with the three weight matrices concatenated, then split afterwards. That is an implementation detail for speed, not part of the idea.
  10. On the hardware side, these are exactly the dense matrix multiplications Chapter 22 described as the operation a GPU is built to do. The whole sequence goes through in one call.

TECHNICAL48.3.5 the engineer’s version#

  1. The projections are three linear maps with no bias in most modern models: Q = X W_Q, K = X W_K, V = X W_V.
  2. Shapes for one attention layer, with input X of shape [B, n, d_model]:
Tensor Shape Notes
X [B, n, d_model] layer input
W_Q [d_model, h x d_k] learned
W_K [d_model, h_kv x d_k] learned
W_V [d_model, h_kv x d_v] learned
  1. In the 2017 base model, d_model = 512, h = 8, d_k = d_v = 64, so each of the four projection matrices is 512 by 512 and holds 262,144 parameters.
  2. In Llama 3.1 8B, d_model = 4096, 32 query heads, 8 key-value heads, head dimension 128. So q_proj is [4096, 4096], while k_proj and v_proj are [1024, 4096]. That asymmetry is grouped-query attention, covered in 48.13.
  3. Bias terms: GPT-2 uses biases on these projections. The Llama family, Mistral and Qwen do not. This is an implementation detail with negligible effect on quality and a small effect on parameter count.
  4. Real parameter names in a checkpoint, which you can list without loading the model using the safetensors header:
model.layers.0.self_attn.q_proj.weight   [4096, 4096]
model.layers.0.self_attn.k_proj.weight   [1024, 4096]
model.layers.0.self_attn.v_proj.weight   [1024, 4096]
model.layers.0.self_attn.o_proj.weight   [4096, 4096]
  1. Note the convention: PyTorch stores a linear layer’s weight as [out_features, in_features] and computes x times W-transpose. So the printed shape is the transpose of the mathematical W in the formula. This trips up almost everybody once.
  2. The low-rank factorization argument is exact. Scores depend only on the product W_Q W_K-transpose, which has rank at most d_k. Choosing d_k = 128 with d_model = 4096 imposes a rank-128 constraint on the score matrix, and that constraint is the parameter saving.
  3. Whether that rank limit hurts has been studied directly. Srinadh Bhojanapalli and colleagues at Google Research published “Low-Rank Bottleneck in Multi-head Attention Models” in 2020, showing that tying d_k to d_model divided by the number of heads does constrain what a head can express, and that decoupling the two can help.

WORDS48.3.6 remember these#

  1. Projection — turning a vector into a different-shaped vector — multiplication by a learned matrix, usually without a bias term.
  2. Linear layer — a matrix multiply, optionally plus a bias — nn.Linear in PyTorch, storing weight as [out_features, in_features].
  3. d_model — the width of the main pipe through the model — the residual stream dimension, 512 in the 2017 base model, 4096 in Llama 3.1 8B.
  4. d_k — the width of one head’s queries and keys — the per-head key dimension, 64 in the 2017 paper, 128 in most current models.
  5. Bilinear form — one matrix sandwiched between two vectors — the effective scoring matrix W_Q W_K-transpose, of rank at most d_k.
  6. Low-rank factorization — writing a big matrix as two thin ones — replacing a [d, d] matrix with [d, r] times [r, d], costing 2dr instead of d squared.

48.4 Scaled dot-product attention, with the actual mathematics#

PLAIN48.4.1 in simple words#

  1. Here is the whole formula. It is one line, and it is the most quoted line in modern machine learning.
Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V
  1. Read it right to left in four steps.
  2. Step one: Q K^T. Compare every query with every key by dot product. Output is a square grid of scores, one row per token, one column per token.
  3. Step two: divide by the square root of d_k, the width of a key. This is a fixed number, not learned. It stops the scores from getting too large.
  4. Step three: softmax each row. Each row of scores becomes a row of weights that are all positive and add up to exactly 1.
  5. Step four: multiply by V. Each row of the output is a weighted average of all the value vectors, using that row’s weights.
  6. That is it. Four steps, three of which are matrix multiplications and one of which is a division.
  7. The result has the same number of rows as the input. One output vector per token, exactly as many as went in.
  8. Everything else in this chapter is either a variation on this line, a way to make it faster, or a way to fix something it cannot do on its own.

PLAIN48.4.2 a picture in your head#

  1. Picture a school report card, but sideways.
  2. Down the left are the students asking questions. Across the top are the students offering answers. Same list of students on both sides.
  3. In each cell you write a number: how useful is that answer for that question. This grid is Q K-transpose.
  4. Now the trick. Nobody trusts raw scores, so each row is converted into percentages that add to 100. That is softmax.
  5. Row 4 might read: 3 per cent, 60 per cent, 3 per cent, 34 per cent. That is student 4’s opinion of who to listen to.
  6. Finally each student writes down their answer on a card. Student 4 makes a mixture: 3 per cent of card 1, 60 per cent of card 2, and so on.
  7. What student 4 walks away with is a blend of the whole class, weighted by their own judgement of relevance.
  8. And every student does this at the same time, so the entire grid is filled in one operation.

Where this comparison breaks: the “percentages” from softmax are never exactly zero. Every student always gets a nonzero share, even an irrelevant one, because the exponential function never reaches zero. And “how useful is that answer” is not a judgement. It is the cosine-like alignment of two vectors, scaled by their lengths. Two vectors can align strongly for reasons that have nothing to do with meaning.

PLAIN48.4.3 a worked example#

  1. This is the full computation, using the Q, K and V produced in section 48.3 for the sentence “the cat was hungry”. Every number below was computed and checked with numpy.
  2. Recall the matrices:
Q                  K                  V
the     1  1  0    the     2  1  1    the     3  1  3
cat     2  2  2    cat     1  1  4    cat     0  3  1
was     2  1  2    was     2  0  2    was     2  2  1
hungry  1  2  2    hungry  2  1  3    hungry  1  3  2
  1. Step one: Q K-transpose. Every entry is one dot product.
  2. Row “hungry”, column “cat”: q_hungry is [1, 2, 2], k_cat is [1, 1, 4].
  3. (1 x 1) + (2 x 1) + (2 x 4) = 1 + 2 + 8 = 11.
  4. Row “hungry”, column “the”: [1, 2, 2] against [2, 1, 1] gives 2 + 2 + 2 = 6.
  5. Row “hungry”, column “was”: [1, 2, 2] against [2, 0, 2] gives 2 + 0 + 4 = 6.
  6. Row “hungry”, column “hungry”: [1, 2, 2] against [2, 1, 3] gives 2 + 2 + 6 =
  7. Doing all sixteen gives the raw score matrix:
S = Q K^T          the   cat   was  hungry
     the             3     2     2     3
     cat             8    12     8    12
     was             7    11     8    11
     hungry          6    11     6    10
  1. Step two: divide by sqrt(d_k). Here d_k = 3, and sqrt(3) = 1.7320508.
S / sqrt(3)        the     cat     was  hungry
     the        1.7321  1.1547  1.1547  1.7321
     cat        4.6188  6.9282  4.6188  6.9282
     was        4.0415  6.3509  4.6188  6.3509
     hungry     3.4641  6.3509  3.4641  5.7735
  1. Step three: softmax each row. Take the row for “hungry”: [3.4641, 6.3509, 3.4641, 5.7735].
  2. Exponentiate each: e to the 3.4641 = 31.9477, e to the 6.3509 = 572.9812, e to the 3.4641 = 31.9477, e to the 5.7735 = 321.6624.
  3. Add them: 31.9477 + 572.9812 + 31.9477 + 321.6624 = 958.5392.
  4. Divide each by the total: 31.9477 / 958.5392 = 0.0333, 572.9812 / 958.5392 = 0.5978, 31.9477 / 958.5392 = 0.0333, 321.6624 / 958.5392 = 0.3356.
  5. Check: 0.0333 + 0.5978 + 0.0333 + 0.3356 = 1.0000. Good.
  6. The full attention weight matrix, every row summing to 1:
A                  the     cat     was  hungry
     the        0.3202  0.1798  0.1798  0.3202
     cat        0.0452  0.4548  0.0452  0.4548
     was        0.0436  0.4393  0.0777  0.4393
     hungry     0.0333  0.5978  0.0333  0.3356
  1. Read that matrix as language. “hungry” puts 0.5978 of its weight on “cat”. Whatever “hungry” describes, it is mostly getting it from “cat”.
  2. “was” splits its weight between “cat” at 0.4393 and “hungry” at 0.4393, which is what a linking verb joining a subject to a predicate would do.
  3. “the” is nearly flat, which is what a determiner with no strong question would do.
  4. This is a toy with hand-chosen matrices. Do not conclude that this pattern is what a trained model produces. It shows the shape of the computation.
  5. Step four: multiply A by V. Take the “hungry” row again:
0.0333 x [3, 1, 3]  =  [0.1000, 0.0333, 0.1000]   from "the"
0.5978 x [0, 3, 1]  =  [0.0000, 1.7933, 0.5978]   from "cat"
0.0333 x [2, 2, 1]  =  [0.0667, 0.0667, 0.0333]   from "was"
0.3356 x [1, 3, 2]  =  [0.3356, 1.0067, 0.6712]   from "hungry"
                       -----------------------
                sum =  [0.5022, 2.9000, 1.4022]
  1. The complete output matrix, one row per token:
Output = A V           dim0    dim1    dim2
     the             1.6405  2.1798  1.9607
     cat             0.6807  2.8645  1.5452
     was             0.7257  2.8350  1.5266
     hungry          0.5022  2.9000  1.4022
  1. Shape check: A is [4, 4], V is [4, 3], so the output is [4, 3]. Four tokens in, four tokens out, width 3.
  2. Notice that “hungry” came in as [0, 1, 1, 0] and leaves as [0.5022, 2.9000, 1.4022]. It is no longer a context-free word. It now carries mostly “cat”.

PLAIN48.4.4 what is really happening inside#

  1. Now the piece everybody skips: why divide by the square root of d_k.
  2. A dot product adds up d_k separate products. The more terms you add, the larger the total tends to be, in both directions.
  3. If the entries of q and k are independent, have mean 0 and variance 1, then each product q_i k_i has mean 0 and variance 1.
  4. Adding d_k independent things with variance 1 gives something with variance d_k, so the typical size grows like the square root of d_k.
  5. Measured with 20,000 random pairs per dimension:
d_k Variance of q.k Std dev sqrt(d_k)
4 4.00 2.00 2.00
64 63.54 7.97 8.00
512 505.98 22.49 22.63
4096 4120.58 64.19 64.00
  1. The measured standard deviation tracks sqrt(d_k) exactly. That is the whole argument.
  2. Now why large scores are bad. Softmax exponentiates. A gap of 10 between two scores becomes a ratio of about 22,026 to 1 after exponentiating.
  3. Here is a real case with d_k = 64 and four random keys:
raw scores       2.153  -7.355  -9.911  -4.582
softmax(raw)     0.9987  0.0001  0.0000  0.0012

after dividing by sqrt(64) = 8
scaled scores    0.2691  -0.9194  -1.2389  -0.5728
softmax(scaled)  0.5110  0.1557  0.1131  0.2202
  1. Without scaling, one key gets 99.87 per cent of the weight and the rest get essentially nothing. The attention has become a hard selection.
  2. That is bad for two reasons. It throws away information, and worse, it kills learning.
  3. Softmax’s gradient for an output p is proportional to p times (1 - p).
  4. At p = 0.5 that is 0.25. At p = 0.999 it is 0.000999. At p = 0.99999 it is 0.00001.
  5. So once softmax saturates, the correction signal flowing back through it is almost zero, and the model stops being able to adjust those scores.
  6. Dividing by sqrt(d_k) rescales the scores to variance about 1 regardless of head width, keeping softmax in its responsive range at the start of training.
  7. The 2017 paper states this reason directly in a footnote, and it is the only place in the paper where the word “suspect” appears about a design choice.
  8. One more subtlety. The scaling matters most at initialization, when weights really are near-random and the variance argument holds exactly. Later in training the weights are not random and the argument is approximate. It still works, and nobody has found a reason to change it.

TECHNICAL48.4.5 the engineer’s version#

  1. The canonical formula, from equation 1 of “Attention Is All You Need”: Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V.
  2. Shapes, single head, batch B, sequence n:
Step Operation Output shape
scores Q times K^T [B, n, n]
scaled divide by sqrt(d_k) [B, n, n]
weights softmax over last dim [B, n, n]
output weights times V [B, n, d_v]
  1. Floating-point operations, forward pass, single head: 2 n^2 d_k for the score matrix and 2 n^2 d_v for the weighted sum, so 4 n^2 d in total when d_k = d_v = d.
  2. Softmax is implemented in the numerically stable form: subtract the row maximum before exponentiating. Without that, exp of a score above about 88 overflows in float32 and above about 11 in float16.
  3. That max-subtraction is mathematically a no-op, since softmax is invariant to adding a constant to all inputs. It is purely a numerical safety measure and it is universal.
  4. Attention scores are almost always computed in float32 even when the model runs in bfloat16, because the softmax denominator accumulates n terms and bfloat16 has only 8 bits of mantissa.
  5. In PyTorch 2.0 and later, the whole thing is one call: torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True), which dispatches to a fused kernel such as FlashAttention when the shapes and dtypes allow it.
  6. The reference implementation, which is worth being able to write from memory:
import torch, math

def attention(q, k, v, mask=None):
    # q: [B, h, n, dk]  k,v: [B, h, n, dk]
    dk = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(dk)
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    w = torch.softmax(scores.float(), dim=-1).to(v.dtype)
    return w @ v
  1. Alternatives to the dot product exist and lost. Additive attention (Bahdanau) uses a one-hidden-layer network and is theoretically stronger at large d_k but far slower, because it cannot be expressed as a single matrix multiply. The 2017 paper reports it performs comparably at small d_k.
  2. Scaling variants exist. Some papers use 1/d_k, some use a learned temperature. The 1/sqrt(d_k) convention is universal in practice, and it is a convention, not a standard. Nothing enforces it.
  3. Query-key normalization, which normalizes Q and K to unit length before the dot product, appeared around 2020 and is used in some large-scale training runs to prevent logit growth. It changes the scaling argument but not the structure.

WORDS48.4.6 remember these#

  1. Dot product — multiply matching entries and add them up — the inner product of two vectors, equal to |a||b|cos(theta).
  2. Score matrix — the grid of how well every query matched every key — the [n, n] tensor Q K^T before scaling and softmax.
  3. Scaling factor — the fixed divisor that keeps scores small — 1/sqrt(d_k), chosen so the score variance is about 1 at initialization.
  4. Saturation — when softmax has already picked a winner and stops responding — the regime where one probability approaches 1 and the Jacobian term p(1-p) approaches 0.
  5. Numerically stable softmax — subtracting the biggest score first — the standard implementation exp(x - max) / sum exp(x - max), avoiding overflow.
  6. Attention output — the blended result for each position — the [n, d_v] matrix of weighted sums of value vectors.
  7. Additive attention — the older, slower scoring method — Bahdanau’s v^T tanh(W_1 q + W_2 k), a small network instead of a dot product.

48.5 Masking: stopping the model reading ahead#

PLAIN48.5.1 in simple words#

  1. A language model is trained by hiding the next word and asking it to guess.
  2. That only works if the model genuinely cannot see the next word.
  3. But attention, as built so far, lets every position see every position, including the ones after it.
  4. So during training on the sentence “the cat was hungry”, the position for “was” can see “hungry”, and predicting “hungry” becomes trivial copying.
  5. The model would score perfectly during training and be useless at generation time, when there is nothing after the current word to copy.
  6. The fix is called a mask. Before the softmax, you go into the score grid and destroy every score that looks forward.
  7. Destroy means set to negative infinity, because e to the power of negative infinity is 0, so softmax gives that position a weight of exactly zero.
  8. You cannot just set the weights to zero after the softmax, because then the remaining weights would not add to 1. Doing it before the softmax means the surviving weights renormalize automatically.
  9. The shape of the destroyed region is the triangle above the diagonal, so this is called a causal mask, or a look-ahead mask.
  10. Causal here means “the past can affect the future, and the future cannot affect the past”, the ordinary meaning of cause.
  11. There is a second, different kind of mask. When you run many sentences of different lengths together, you pad the short ones with filler tokens.
  12. A padding mask blocks attention to that filler, so real words never average in meaningless positions.
  13. And there are models that use no causal mask at all. They read the whole text both ways at once. Those are excellent at understanding a fixed piece of text and cannot generate.

PLAIN48.5.2 a picture in your head#

  1. Imagine an exam where the questions and answers are printed on the same long scroll, in order.
  2. Question 1, answer 1, question 2, answer 2, all the way down.
  3. You are given the scroll rolled up, and a rule: you may unroll it only as far as the question you are currently answering, never further.
  4. So when you are on question 3, you can see questions 1, 2 and 3 and answers 1 and 2, and nothing beyond.
  5. That is exactly the causal mask. Position t sees positions 1 through t.
  6. Now here is the part that surprises people. The examiner does not hand you one question at a time. The examiner hands the whole class the whole scroll at once, with each student allowed a different unrolling limit.
  7. Student 1 may see 1 word. Student 2 may see 2 words. Student 500 may see 500.
  8. All 500 students work at the same time, in one operation, and each is graded on a genuinely fair prediction.
  9. That is why masked training is fast. One pass over a 4,000-token document produces 4,000 separate prediction problems, all solved in parallel.

Where this comparison breaks: the students are not separate. They are 500 rows of one matrix, and the “unrolling limit” is not enforced by discipline but by arithmetic that makes forbidden information contribute exactly zero. Also, a real exam has one right answer per question. Here the target is the actual next token in the training text, which is one sample from a distribution of acceptable continuations, and the model is graded against that one sample.

PLAIN48.5.3 a worked example#

  1. Take the scaled score matrix computed in section 48.4 for “the cat was hungry”, and apply a causal mask.
  2. Before masking, the scaled scores were:
                   the     cat     was  hungry
     the        1.7321  1.1547  1.1547  1.7321
     cat        4.6188  6.9282  4.6188  6.9282
     was        4.0415  6.3509  4.6188  6.3509
     hungry     3.4641  6.3509  3.4641  5.7735
  1. Now set everything strictly above the diagonal to negative infinity:
                   the     cat     was  hungry
     the        1.7321    -inf    -inf    -inf
     cat        4.6188  6.9282    -inf    -inf
     was        4.0415  6.3509  4.6188    -inf
     hungry     3.4641  6.3509  3.4641  5.7735
  1. The pattern is a lower triangle of real numbers and an upper triangle of negative infinity. This is the causal mask, and it is always this shape.
  2. Now softmax each row. Row “the” has exactly one surviving score, so it gets weight 1.0 and everything else 0.
  3. Row “cat” has two surviving scores, 4.6188 and 6.9282. Their difference is 2.3094, so e to the 2.3094 = 10.07, and the weights are 1 / 11.07 = 0.0903 and 10.07 / 11.07 = 0.9097.
  4. The full masked attention matrix:
A causal           the     cat     was  hungry
     the        1.0000  0.0000  0.0000  0.0000
     cat        0.0903  0.9097  0.0000  0.0000
     was        0.0778  0.7836  0.1386  0.0000
     hungry     0.0333  0.5978  0.0333  0.3356
  1. Every row still sums to exactly 1. The zeros are true zeros, not small numbers.
  2. Compare the “was” row with the unmasked version. Unmasked it was [0.0436, 0.4393, 0.0777, 0.4393]. Masked it is [0.0778, 0.7836, 0.1386, 0].
  3. The 0.4393 that was going to “hungry” did not disappear. It was redistributed proportionally across the three legal positions. That is what the renormalization in softmax does for free.
  4. Multiplying by V gives the masked output:
Output (causal)        dim0    dim1    dim2
     the             3.0000  1.0000  3.0000
     cat             0.2710  2.8193  1.1807
     was             0.5107  2.7057  1.1556
     hungry          0.5022  2.9000  1.4022
  1. Note the first row. Position 1 can only see itself, so its output is exactly its own value vector [3, 1, 3]. That is always true of the first token in a causal model, in every layer.
  2. Note the last row. It is identical to the unmasked last row, because the last position had nothing ahead of it to mask anyway.

PLAIN48.5.4 what is really happening inside#

  1. In code, the mask is not really negative infinity. It is a large negative number, or a genuine negative infinity that the softmax handles.
  2. In float32, a common choice is -1e9 or the smallest representable value. In float16 it must be about -65504 or smaller in magnitude than that, because float16 cannot hold -1e9.
  3. Using a finite value works because e to the power of -1e9 underflows to exactly 0.0 in floating point. The result is identical to true zero.
  4. The honest version: there is a real failure mode here. If an entire row is masked, every score becomes negative infinity, and softmax computes 0 divided by 0, giving NaN, which then poisons the whole model. Implementations guard against this. It happens in practice with badly constructed padding masks.
  5. The mask is not a parameter. Nothing about it is learned. It is a fixed boolean pattern computed from the sequence length.
  6. It costs no memory to store, in good implementations, because the triangular pattern is generated on the fly inside the kernel rather than materialized as an [n, n] tensor.
  7. Now the padding mask. Suppose you batch three prompts of lengths 5, 9 and 3.
  8. To make one rectangular tensor, you pad all three to length 9 with a filler token id.
  9. Without a mask, position 2 of the third prompt would average in six filler positions, and its output would be contaminated.
  10. The padding mask sets the columns of the padded positions to negative infinity for every row.
  11. Padding masks and causal masks combine by taking whichever is more restrictive at each cell, which in boolean terms is a logical OR of the two blocked patterns.
  12. A neat trick used by almost every training framework avoids padding entirely: concatenate many documents into one long stream, cut it into fixed-length chunks, and never pad at all. This is called sequence packing.

TECHNICAL48.5.5 the engineer’s version#

  1. The causal mask M is defined by M[i][j] = 0 if j is less than or equal to i, and negative infinity otherwise. It is added to the scaled scores before softmax.
  2. The masked formula: softmax(Q K^T / sqrt(d_k) + M) V.
  3. In PyTorch, torch.triu(torch.ones(n, n), diagonal=1).bool() produces the blocked pattern, and scores.masked_fill(mask, float("-inf")) applies it.
  4. Since PyTorch 2.0, scaled_dot_product_attention(..., is_causal=True) is preferred, because fused kernels skip the masked blocks entirely rather than computing and discarding them, roughly halving the work.
  5. That saving is real and worth naming: a causal attention over n tokens does about n^2/2 useful score computations, not n^2.
  6. Encoder models use no causal mask. BERT, from Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova at Google in October 2018, attends both directions at every layer.
  7. That is why BERT cannot be trained on next-token prediction. If every position sees every position, the answer is already visible.
  8. BERT is trained instead on masked language modelling: replace 15 per cent of tokens with a special [MASK] token and predict the originals. The masking there is of the input tokens, not of the attention matrix. Two different meanings of the same word, and this confuses people constantly.
  9. The consequences of the choice:
Property Causal (decoder) Bidirectional (encoder)
Sees future tokens No Yes
Can generate text Yes No
Training signal per token Every token 15 per cent of tokens
Best at Generation, chat Classification, embedding
  1. The training-efficiency line in that table is important and often missed. A causal model gets a prediction target at every one of n positions. A masked language model gets one at roughly 0.15n positions, so it extracts less signal per pass over the data.
  2. Prefix language modelling is the hybrid: bidirectional attention over a prompt prefix, causal attention over the generated part. Used in UniLM in 2019 and in some of the T5 variants.
  3. Sliding window attention adds a second constraint: block anything further back than w positions. Mistral 7B, released October 2023, used w = 4096. That is covered in 48.13.
  4. Observation tools: printing the mask is one line, but for real debugging the useful check is that every attention row sums to 1 and contains no NaN. A single NaN in one row propagates to the entire output within one layer.

WORDS48.5.6 remember these#

  1. Mask — blocking some positions from being looked at — an additive [n, n] tensor of 0 and negative infinity applied before softmax.
  2. Causal mask — you may not read ahead — the strictly upper-triangular block pattern, also called a look-ahead or autoregressive mask.
  3. Padding mask — ignore the filler used to square up a batch — a per-sequence column mask blocking positions beyond the true length.
  4. Autoregressive — predicting the next thing from all the previous things — a factorization of the joint probability as a product of conditionals.
  5. Masked language modelling — hiding some input words and guessing them — the BERT objective, corrupting about 15 per cent of input tokens.
  6. Sequence packing — gluing documents together to avoid padding — concatenating a corpus into a token stream and slicing fixed-length windows.
  7. Bidirectional — able to see both left and right — an unmasked encoder attention pattern.

48.6 Multi-head attention#

PLAIN48.6.1 in simple words#

  1. One attention operation produces one set of weights per position. One opinion about what matters.
  2. But a word usually needs several different things at once.
  3. Take the word “hungry” in “the cat was hungry”. It needs to know what it describes, which is “cat”. It also needs to know its grammatical role. It also needs to know where it sits in the sentence.
  4. Those are different questions, and one set of weights cannot answer them all, because the weights must add up to 1 and get spent once.
  5. So the model runs several attention operations side by side, each with its own W_Q, W_K and W_V. Each one is called a head.
  6. Head 1 might mostly look at the previous word. Head 2 might look at the subject of the sentence. Head 3 might look at matching brackets.
  7. They all run at the same time on the same input, and they do not talk to each other while running.
  8. When they finish, their outputs are laid side by side in one long row, which is called concatenation.
  9. Then one final learned matrix mixes that long row back down to the original width. That matrix is W_O, the output projection.
  10. The clever part is that heads cost nothing extra. If the model is 512 wide and you want 8 heads, each head is 64 wide, and 8 times 64 is 512.
  11. So the total work is the same as one 512-wide head. You have bought variety for free by splitting rather than adding.
  12. That is the trade: eight narrow views instead of one wide view, at identical cost.

PLAIN48.6.2 a picture in your head#

  1. Picture a panel of eight specialist doctors examining the same patient at the same time.
  2. The cardiologist looks at the heart. The neurologist looks at the nerves. The radiologist looks at the scans.
  3. Each one asks different questions, of the same patient, and gets a different answer.
  4. None of them sees the whole picture. Each is deliberately narrow.
  5. When they are done, all eight write their findings on one long sheet.
  6. A general physician reads the whole sheet and produces one combined assessment. That is W_O, the output projection.
  7. The panel is faster than eight sequential consultations, because they work simultaneously.
  8. And the panel is better than one generalist, because a generalist would have to spend the same total effort spread across all the specialities.

Where this comparison breaks: real specialists chose their speciality and know what it is. Heads are not assigned roles. They are eight identical copies with different random starting values, and whatever specialization appears is an accident of training that nobody designed. Many heads end up doing nothing useful at all, which no hospital would tolerate. And a doctor examines the actual patient. A head sees only a 64-dimensional slice of a vector that has already passed through many previous layers.

PLAIN48.6.3 a worked example#

  1. Work out every shape for the 2017 base model, with a 10-token sentence.
  2. The configuration: d_model = 512, 8 heads, d_k = d_v = 64, since 512 / 8 = 64.
input X                       [1, 10, 512]

project with W_Q [512, 512]   [1, 10, 512]
reshape into heads            [1, 10, 8, 64]
transpose head axis forward   [1, 8, 10, 64]

same for K and V              [1, 8, 10, 64]

scores per head Q K^T         [1, 8, 10, 10]
softmax over last axis        [1, 8, 10, 10]
weights times V               [1, 8, 10, 64]

transpose back                [1, 10, 8, 64]
concatenate heads             [1, 10, 512]
project with W_O [512, 512]   [1, 10, 512]
  1. Look at the first and last lines. The input was [1, 10, 512] and the output is [1, 10, 512]. The shape is unchanged, which is why blocks can be stacked.
  2. Look at the score tensor: [1, 8, 10, 10]. That is 8 separate 10-by-10 grids, one per head. 800 numbers.
  3. There is no reshaping magic here. “Reshape into heads” just means reading the same 512 numbers as 8 groups of 64, without moving any data.
  4. Head 0 gets dimensions 0 to 63. Head 1 gets 64 to 127. And so on to head 7, which gets 448 to 511.
  5. Parameters: W_Q, W_K, W_V and W_O are each 512 by 512, so 262,144 each, and 1,048,576 for the whole multi-head attention block.
  6. That is exactly 4 times d_model squared, which is a formula worth memorizing: attention costs 4 d^2 parameters per layer.
  7. Now do the same for a modern model, Llama 3.1 8B: d_model = 4096, 32 query heads, 8 key-value heads, head dimension 128.
  8. Note that 32 times 128 = 4096, so the query side still adds up to d_model.
  9. But the key and value side uses only 8 heads, so 8 times 128 = 1024.
q_proj  [4096, 4096]  =  16,777,216 parameters
k_proj  [1024, 4096]  =   4,194,304 parameters
v_proj  [1024, 4096]  =   4,194,304 parameters
o_proj  [4096, 4096]  =  16,777,216 parameters
                         ----------
total attention        =  41,943,040 per layer
  1. Each of the 8 key-value heads is shared by 4 query heads. That is grouped-query attention, and its purpose is to shrink the memory cost of generation, explained fully in 48.13.
  2. Across 32 layers that is 1,342,177,280 attention parameters, which is 16.7 per cent of the model’s 8,030,261,248 total.

PLAIN48.6.4 what is really happening inside#

  1. Heads are not separate networks. They are one matrix multiplication whose output is interpreted as several groups.
  2. W_Q is a single [512, 512] matrix. Nothing in it says “this column belongs to head 3”. The head boundaries exist only in how the result is reshaped.
  3. That is why the cost is identical to single-head attention with the full width. The reshape is free.
  4. The heads become genuinely independent at the score step, because a query from head 3 is only ever dotted with keys from head 3.
  5. This makes the attention computation block diagonal across heads. Eight small independent problems rather than one large one.
  6. The only place heads mix is W_O at the end. Every output dimension of W_O is a weighted combination of all 512 concatenated numbers, so every head can influence every output.
  7. That means W_O is not just a resize. It is the thing that decides how much each head’s finding is worth, and it is learned.
  8. A useful way to see it: W_O can be split into 8 blocks of [64, 512], one per head. Then the block’s output is added to the residual stream. Each head writes its own contribution independently, and they sum.
  9. This decomposition is exactly how the interpretability literature analyses attention, and it is mathematically exact, not an approximation.
  10. Heads within one layer cannot see each other’s results. Head 3 cannot use what head 5 found. That only becomes possible in the next layer.

TECHNICAL48.6.5 the engineer’s version#

  1. Multi-head attention, from the 2017 paper: concat(head_1, …, head_h) W_O, where head_i = Attention(Q W_Q_i, K W_K_i, V W_V_i).
  2. Standard configurations:
Model d_model Heads Head dim
Transformer base 2017 512 8 64
BERT-base 2018 768 12 64
GPT-2 small 2019 768 12 64
Llama 3.1 8B 2024 4096 32 128
  1. The convention d_k = d_model / h keeps total cost constant. It is a convention, not a requirement, and models that break it exist.
  2. The 2017 paper ran an ablation: with d_model fixed at 512, 1 head scored 0.9 BLEU worse than 8 heads, and 32 heads was also slightly worse than 8. So more heads is not monotonically better once each head becomes too narrow.
  3. What heads have been observed to learn, from published analysis:
    1. Positional heads that attend almost entirely to the previous token or the next token. These are the most reliable finding and appear in nearly every model.
    2. Syntactic heads that attend from a verb to its direct object, or from a determiner to its noun. Kevin Clark, Urvashi Khandelwal, Omer Levy and Christopher Manning documented these in BERT in 2019.
    3. Coreference heads that attend from a pronoun to its antecedent, also in the same 2019 analysis, at above-chance but well below perfect accuracy.
    4. Induction heads, which complete a repeated pattern by finding an earlier occurrence of the current token and copying what followed it. Described by Catherine Olsson and colleagues at Anthropic in 2022 and connected to in-context learning.
  4. The honest version: this is not settled science. Head interpretability is active research, not established fact. The named roles above are real published observations on specific models, and they do not mean that heads in general have clean human-nameable jobs.
  5. Redundancy is well established and cuts the other way. Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich and Ivan Titov showed in 2019 that pruning 38 of 48 encoder heads in a translation model cost only 0.15 BLEU on English-Russian WMT data.
  6. Paul Michel, Omer Levy and Graham Neubig published the same finding independently in 2019, in a paper titled “Are Sixteen Heads Really Better Than One?”. Many heads can be deleted at test time with little loss.
  7. So the correct summary is: some heads do identifiable jobs, most heads are partly redundant, and the population as a whole is not a tidy division of labour.
  8. Marketing claim to reject: any product page saying the model “has heads dedicated to reasoning” or similar. No such assignment is designed or verified.
  9. Reproducing head analysis is straightforward. Load a model with output_attentions=True, take the [B, h, n, n] tensor, and for a previous-token head you will see nearly all mass on the first subdiagonal.

WORDS48.6.6 remember these#

  1. Head — one independent attention view — one of h parallel attention computations, each of width d_model / h by convention.
  2. Multi-head attention — several attention views run side by side — h parallel scaled dot-product attentions, concatenated and projected by W_O.
  3. Concatenation — laying results side by side into one long row — joining the h per-head outputs along the feature axis to restore width d_model.
  4. Output projection — the matrix that mixes the heads back together — W_O, of shape [d_model, d_model], applied after concatenation.
  5. Head dimension — how wide one head is — d_k, typically 64 in 2017-era models and 128 in current ones.
  6. Induction head — a head that continues a repeated pattern — a two-head circuit that matches the current token to a previous occurrence and copies the following token.
  7. Head pruning — deleting heads that do not earn their keep — zeroing whole heads at inference, shown to cost little accuracy in several 2019 studies.

48.7 Position: attention has no idea what order things are in#

PLAIN48.7.1 in simple words#

  1. Here is a fact about attention that catches everybody by surprise.
  2. Attention has no concept of order. None. It does not know which word came first.
  3. To attention, a sentence is a bag of items, not a sequence.
  4. Prove it to yourself. Look back at the formula: softmax(Q K^T / sqrt(d_k)) V.
  5. Find the position number in it. There is none. There is no i, no t, no index of any kind.
  6. Every operation is a dot product between two token vectors, and a dot product does not care where either token sits.
  7. So “dog bites man” and “man bites dog” would produce exactly the same set of output vectors, just shuffled the same way as the input.
  8. That would be a catastrophe for language, where order is most of the meaning.
  9. So order has to be added by hand. The model is given, for each position, some extra numbers that say where that position is.
  10. The original method was to add a fixed pattern of sine and cosine waves to each token’s vector before the first layer.
  11. A second method is to learn a separate vector for position 1, position 2, and so on, exactly like an embedding table for positions.
  12. The modern method, used by almost every current model, is different and cleverer. It rotates the query and key vectors by an angle proportional to their position.
  13. It is called rotary position embedding, or RoPE, and the reason it won is that it makes the score between two tokens depend only on how far apart they are, not on where they are in absolute terms.

PLAIN48.7.2 a picture in your head#

  1. Imagine a set of magnetic word tiles scattered on a table, one word per tile.
  2. Attention reads all the tiles at once and works out which ones relate to which others.
  3. Now shake the table. The tiles land in a different arrangement. Attention gives exactly the same answer, because it never looked at the arrangement.
  4. To fix this, you write a number on the back of each tile: 1, 2, 3, 4.
  5. That is learned absolute position embedding. Simple, and it fails as soon as you meet tile 5,000 when you only ever wrote numbers up to 4,096.
  6. Now a better fix. Instead of writing numbers, you tilt each tile by an angle. Tile 1 is tilted 10 degrees, tile 2 is tilted 20, tile 3 is tilted 30.
  7. When two tiles are compared, what matters is the difference in their tilts. Tiles 1 and 2 differ by 10 degrees. Tiles 500 and 501 also differ by 10.
  8. So the comparison naturally depends on the gap between positions, not the absolute positions.
  9. That is rotary position embedding. Tilt by angle, compare by difference.

Where this comparison breaks: a tile has one tilt angle. A real RoPE vector is split into 64 pairs of numbers, and each pair is rotated at its own speed, from one full turn every 6.3 positions down to one full turn every two and a half million positions. So it is 64 clocks running at 64 different speeds, not one tilt. And the tilt is applied to the query and key vectors only, never to the values.

PLAIN48.7.3 a worked example#

  1. First, the proof that attention ignores order. Take the four-token example from 48.4, which produced this output:
Output          dim0    dim1    dim2
  the         1.6405  2.1798  1.9607
  cat         0.6807  2.8645  1.5452
  was         0.7257  2.8350  1.5266
  hungry      0.5022  2.9000  1.4022
  1. Now shuffle the input rows into the order hungry, cat, the, was, and run the identical computation with the identical matrices:
Output of shuffled input     dim0    dim1    dim2
  hungry                   0.5022  2.9000  1.4022
  cat                      0.6807  2.8645  1.5452
  the                      1.6405  2.1798  1.9607
  was                      0.7257  2.8350  1.5266
  1. Every row is identical to before. Only their order changed, in exactly the same way the input changed.
  2. That property has a name: permutation equivariance. Shuffle the input, and the output is shuffled the same way, with no other change.
  3. This was verified numerically, not asserted. The two matrices agree to floating-point precision.
  4. Now sinusoidal position encoding. The formula from the 2017 paper, for position pos and dimension index i:
PE(pos, 2i)   = sin( pos / 10000^(2i/d_model) )
PE(pos, 2i+1) = cos( pos / 10000^(2i/d_model) )
  1. With d_model = 8 there are 4 sine-cosine pairs. Their frequencies:
Pair Frequency Wavelength
0 1.0 6.28 positions
1 0.1 62.8 positions
2 0.01 628 positions
3 0.001 6283 positions
  1. The actual values for the first five positions:
pos 0:  0.0000  1.0000  0.0000  1.0000  0.0000  1.0000  0.0000  1.0000
pos 1:  0.8415  0.5403  0.0998  0.9950  0.0100  1.0000  0.0010  1.0000
pos 2:  0.9093 -0.4161  0.1987  0.9801  0.0200  0.9998  0.0020  1.0000
pos 3:  0.1411 -0.9900  0.2955  0.9553  0.0300  0.9996  0.0030  1.0000
pos 4: -0.7568 -0.6536  0.3894  0.9211  0.0400  0.9992  0.0040  1.0000
  1. Read the columns. The first pair changes fast and wraps around quickly. The last pair barely moves at all across five positions.
  2. That is the design. Fast waves distinguish neighbours precisely, slow waves distinguish far-apart positions coarsely. It is a positional binary counter made of continuous waves.
  3. Now RoPE, with real numbers. Take a query pair q = [1.0, 0.0] and a key pair k = [0.6, 0.8]. Their plain dot product is 0.6.
  4. RoPE rotates q by m times theta and k by n times theta, where m and n are the two positions. With theta = 1.0 radian:
m n m - n Dot after rotation
1 0 1 0.997358
5 4 1 0.997358
100 99 1 0.997358
3 0 3 -0.481099
7 4 3 -0.481099
  1. Look at the last column. Every pair with the same gap gives the same result, to six decimal places, whether the positions are 1 and 0 or 100 and
  2. That is the whole selling point of RoPE, demonstrated numerically. The score depends on relative distance only.

PLAIN48.7.4 what is really happening inside#

  1. Sinusoidal encoding is added to the token embedding, once, before layer 1. Token vector plus position vector, element by element.
  2. That means position information and word information are mixed together in the same numbers from the start, and every layer sees the sum.
  3. Learned absolute embeddings work identically, except the position vectors are parameters found by training rather than computed from a formula.
  4. GPT-2 uses learned absolute embeddings: a table of 1,024 rows by 768 columns, which is 786,432 parameters spent purely on saying where things are.
  5. Both of these have the same weakness. They are absolute. Position 3,000 has its own vector, unrelated to position 2,999.
  6. And if the table has 1,024 rows, position 1,024 does not exist. The model cannot be run on longer text at all.
  7. Sinusoids can at least be computed for any position, but a model trained only on positions 0 to 511 has never seen what position 800 looks like, and in practice quality collapses beyond the trained length.
  8. RoPE takes a completely different route. It never adds anything. It rotates.
  9. Split the query vector into consecutive pairs: dimensions (0,1), (2,3), (4,5), and so on. Treat each pair as a point on a flat plane.
  10. Rotate pair j by an angle of position times theta_j, where theta_j gets smaller for later pairs, exactly like the sinusoid frequencies.
  11. Do the same to the key vector, using the key’s own position.
  12. Now take the dot product. Rotating both vectors by the same amount leaves the dot product unchanged, so only the difference in rotation survives.
  13. That is basic geometry: a dot product is invariant under rotation. Rotating q by m and k by n is the same as rotating q by (m - n) and leaving k alone.
  14. So the score naturally becomes a function of m - n. No extra terms, no lookup table, no learned parameters at all.
  15. RoPE is applied inside every attention layer, to Q and K only, every time. It is not applied to V, and it is not added to the residual stream.

TECHNICAL48.7.5 the engineer’s version#

  1. Permutation equivariance is exact: for any permutation matrix P, Attention(PX) = P Attention(X). This is why position injection is mandatory rather than helpful.
  2. Sinusoidal, 2017. Fixed, no parameters, defined for all positions. The base 10,000 sets the slowest wavelength. The paper’s stated reason for choosing sinusoids is that PE(pos + k) is a linear function of PE(pos) for any fixed offset k, so relative offsets are in principle learnable as a linear map. This was verified numerically for this chapter: for the fastest pair with theta = 1, rotating PE(5) by 3 radians reproduces PE(8) exactly.
  3. The paper also reports that learned and sinusoidal encodings performed almost identically, and chose sinusoidal for its extrapolation potential.
  4. Sinusoid dot products fall with distance, which is a useful inductive bias. Measured at d_model = 64: the dot product of PE(0) with itself is 32.0, with PE(1) is 30.92, with PE(10) is 21.05, and with PE(200) is 11.44.
  5. Learned absolute, used by BERT (2018), GPT-1, GPT-2 and GPT-3. A table of shape [max_positions, d_model]. Hard cap on length, no extrapolation at all.
  6. Relative position representations, from Peter Shaw, Jakob Uszkoreit and Ashish Vaswani in 2018, add a learned vector for each relative offset into the key and value computation.
  7. T5 relative position bias, from Colin Raffel and colleagues in 2019, adds a learned scalar to the attention logit per bucket of relative distance, shared across layers, with logarithmically spaced buckets. Cheap and effective.
  8. ALiBi, from Ofir Press, Noah Smith and Mike Lewis in 2021, adds a linear penalty to the attention logit: minus m times the distance, where m is a fixed per-head slope, typically a geometric sequence such as 1/2, 1/4, 1/8. No parameters at all, and it extrapolates well beyond the training length. The paper’s title is “Train Short, Test Long”.
  9. RoPE, from Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen and Yunfeng Liu, in the RoFormer paper first posted in April 2021. It is the dominant choice as of 2026. Llama, Mistral, Qwen, Gemma, DeepSeek and Phi all use it.
  10. RoPE frequencies use theta_j = base^(-2j/d_head). With base 10,000 and head dimension 128, the fastest pair has wavelength 6.3 positions and the slowest has wavelength 54,410 positions.
  11. Raising the base stretches everything. With base 500,000, as used in Llama 3.1, the slowest pair has wavelength 2,559,196 positions. That is how a 128k-token context is supported without the slowest frequency wrapping.
Method Year Parameters Extrapolates
Sinusoidal 2017 0 Poorly
Learned absolute 2018 max_pos x d Not at all
T5 bias 2019 ~ 32 per head Moderately
ALiBi 2021 0 Well
RoPE 2021 0 With extension
  1. Context extension. A model trained at 4,096 tokens fails at 32,768 because the rotation angles at position 30,000 are outside anything it saw.
  2. Position interpolation, from Shouyuan Chen, Sherman Wong, Liangjian Chen and Yuandong Tian at Meta in June 2023, divides every position index by a scale factor s. To go from 4,096 to 32,768, use s = 8, so position 32,767 is presented to the model as 4,095.9.
  3. That keeps every angle inside the trained range, at the cost of squashing fine positional detail. It requires a short fine-tune, typically 1,000 steps, and works remarkably well.
  4. NTK-aware scaling, which appeared as a community proposal in mid-2023, changes the RoPE base instead of the positions, stretching low frequencies more than high ones so nearby tokens keep their resolution.
  5. YaRN, from Bowen Peng, Jeffrey Quesnelle, Honglu Fan and Enrico Shippole in 2023, combines the two: interpolate the low-frequency dimensions, leave the high-frequency ones alone, apply a ramp in between, and add a temperature adjustment to the attention logits. It reaches long contexts with about ten times less fine-tuning data than plain interpolation.
  6. Llama 3.1, released July 2024, used a variant of this family to reach a 131,072-token context from a shorter pretraining length, together with base 500,000.
  7. Established fact: RoPE gives exact relative-distance dependence, and interpolation methods extend context in practice. Active research: how much real long-range capability extension gives, as opposed to merely not crashing. Benchmarks such as RULER, from NVIDIA in 2024, found many models claiming 32,000-token contexts performed well at much shorter lengths.

WORDS48.7.6 remember these#

  1. Permutation equivariance — shuffle the input, get the same output shuffled — Attention(PX) = P Attention(X) for any permutation matrix P.
  2. Positional encoding — extra numbers telling the model where each token sits — any scheme injecting order information into an order-blind operation.
  3. Sinusoidal encoding — fixed sine and cosine waves at many speeds — the 2017 scheme with frequencies base^(-2i/d), base 10,000.
  4. Absolute position — where a token is in the whole sequence — an index from 0 to n-1, encoded directly.
  5. Relative position — how far apart two tokens are — the offset m - n, which is what a score should depend on.
  6. RoPE — rotating queries and keys by an angle set by position — rotary position embedding, applying a block-diagonal rotation to Q and K.
  7. RoPE base (theta) — the number controlling how slow the slowest wave is — 10,000 originally, raised to 500,000 or higher for long contexts.
  8. Position interpolation — squeezing long positions into the trained range — dividing position indices by a scale factor s, then fine-tuning.
  9. YaRN — a smarter context extension that treats fast and slow waves differently — wavelength-dependent interpolation with a logit temperature correction.
  10. ALiBi — penalizing distant tokens with a straight-line penalty — adding minus m times |i - j| to the attention logit, with fixed per-head slopes.

48.8 The rest of the block: feed-forward, residuals and normalization#

PLAIN48.8.1 in simple words#

  1. Attention is only half of a transformer block. The other half is a small ordinary network applied to each position separately.
  2. It is called the feed-forward network, or FFN, or the MLP. All three names mean the same thing.
  3. It has exactly two steps: widen, then narrow. Multiply by a matrix that makes the vector about four times wider, apply a non-linear function, multiply by a matrix that brings it back to the original width.
  4. The important word is “separately”. Every position goes through the same FFN, on its own, with no reference to any other position.
  5. So attention is where positions talk to each other, and the FFN is where each position thinks about what it just heard.
  6. Now the surprising part: this simple two-matrix block holds most of the model’s parameters. About 70 per cent for an 8-billion model, and about 80 per cent for larger ones.
  7. Chapter 47 counted this exactly. For Llama 3.1 8B, the feed-forward matrices hold 176,160,768 parameters per layer against 41,943,040 for attention.
  8. So attention gets all the attention, and the FFN holds the weights.
  9. Two more pieces complete the block, and both are about making deep stacks trainable rather than about capability.
  10. A residual connection means: whatever a sub-layer computes, add the original input back on. Output equals input plus change, not just change.
  11. Without residuals, a 32-layer stack cannot be trained. The correction signal dies on the way back, exactly as it did in recurrent networks.
  12. Layer normalization rescales each vector so its numbers have a predictable size, which stops values drifting to extremes as they pass through dozens of layers.
  13. That is the whole block: attention, add the input back, normalize, feed-forward, add the input back, normalize. Then repeat.

PLAIN48.8.2 a picture in your head#

  1. Picture a long assembly line with 32 stations, each one modifying a product.
  2. In a bad design, each station takes the previous station’s output and rebuilds it from scratch. If station 7 does something wrong, everything downstream is corrupted, and the factory manager cannot tell which station caused it.
  3. In the residual design, each station is given the product and told: do not rebuild it. Just add your modification on top and pass it along.
  4. Now the product carries a continuous history, and every station’s change is a small addition rather than a replacement.
  5. If the manager finds a fault at the end, the complaint can travel back down the line unchanged, because the additions form a straight unbroken path.
  6. That straight path is exactly what a residual connection creates for gradients. It is often called the residual stream, and it runs the whole depth of the model.
  7. Normalization is the quality inspector between stations, resetting the product to a standard size and shape so the next station receives something in the range it expects.

Where this comparison breaks: nothing is physically added to a product, and the “modification” is a vector added element by element, which can cancel earlier contributions rather than accumulate. Also, real assembly stations have fixed jobs. These are all trained together, so a station’s job is defined only by what the others ended up doing.

PLAIN48.8.3 a worked example#

  1. First, why residual connections matter, with arithmetic.
  2. During backpropagation, the correction signal is multiplied by something at each layer. Call it the per-layer factor.
  3. Without residuals, that factor depends entirely on the layer’s weights and there is no reason for it to be near 1.
Per-layer factor After 32 layers After 96 layers
0.90 0.034 0.00004
0.95 0.194 0.0073
1.00 1.000 1.000
1.05 4.765 108.2
  1. At 0.9 per layer, a 96-layer model receives 0.004 per cent of the signal at the bottom. Training the early layers becomes impossible.
  2. With a residual connection, the layer computes x + f(x), so the derivative is 1 + f’(x). The 1 is there whatever f does.
  3. So there is always a path with factor exactly 1, and the product along that path stays exactly 1 no matter how deep the stack is.
  4. That is the entire trick, and it came from image models: Kaiming He, Xiangyu Zhang, Shaoqing Ren and Jian Sun published ResNet in December 2015, trained 152 layers, and won the ImageNet competition that year.
  5. Now layer normalization, worked on a real vector of 8 numbers:
x        =  2.0  -1.0   0.5   4.0  -2.5   1.0   0.0   3.0
mean     =  0.875
variance =  3.921875
std dev  =  1.980372
  1. Subtract the mean and divide by the standard deviation:
normed   =  0.5681 -0.9468 -0.1894  1.5780
           -1.7042  0.0631 -0.4418  1.0730
  1. Check: the mean of the normalized vector is 0.0 and its standard deviation is 1.0. That is guaranteed by construction.
  2. Then multiply by a learned scale and add a learned shift, one per dimension, so the layer can undo the normalization if that helps.
  3. Crucially, the mean and variance were computed across the 8 features of this one token. Not across the batch. Not across the sequence.
  4. That is why it works with any batch size, including batch size 1, which is the normal case when you chat with a model.
  5. RMSNorm is the cheaper modern variant. It skips the mean subtraction and divides by the root mean square:
root mean square of x = 2.165064
normed = 0.9238 -0.4619  0.2309  1.8475
        -1.1547  0.4619  0.0000  1.3856
  1. It has one learned parameter per dimension instead of two, and in practice performs the same. Llama, Mistral, Qwen and Gemma all use it.

PLAIN48.8.4 what is really happening inside#

  1. The feed-forward network, written out: FFN(x) = W_2 (activation(W_1 x + b_1))
    • b_2.
  2. W_1 has shape [d_model, d_ff] and W_2 has shape [d_ff, d_model], with d_ff usually 4 times d_model.
  3. The 4x ratio is a convention from the 2017 paper, where d_model was 512 and d_ff was 2048. It is not a law and modern models vary it.
  4. The activation function was ReLU in 2017, then GELU in BERT and GPT-2, and now SwiGLU in most current models. Chapter 46 covered what these do.
  5. SwiGLU uses three matrices rather than two: a gate, an up-projection and a down-projection, with the gate multiplied elementwise into the up path.
  6. Because it has three matrices, models using SwiGLU shrink d_ff to about 8/3 of d_model to keep the parameter count similar. Llama 3.1 8B uses 14,336 with d_model 4096, which is 3.5x rather than the naive 4x.
  7. Now pre-norm versus post-norm, which is a real and consequential difference.
  8. Post-norm, the 2017 original: x -> attention -> add x -> normalize. The normalization sits after the addition, on the main path.
  9. Pre-norm, used by almost everything since about 2019: x -> normalize -> attention -> add x. The normalization sits inside the branch, and the main path from input to output has nothing on it but additions.
  10. That difference is everything. In pre-norm, the residual stream is a clean sum with no normalization interrupting it, so the gradient path really is unobstructed all the way down.
  11. In post-norm, every residual addition is immediately rescaled, which shrinks the accumulated signal and makes deep stacks unstable.
  12. The practical consequence: post-norm transformers need a learning-rate warmup or they diverge. Pre-norm ones train stably without it.
  13. The 2017 paper used post-norm with 4,000 warmup steps, and that warmup was not optional.

TECHNICAL48.8.5 the engineer’s version#

  1. The complete pre-norm decoder block, drawn:
        x  (residual stream, width d_model)
        |
        +---------------------------+
        |                           |
     [RMSNorm]                      |
        |                           |
   [Multi-head attention]           |
   [ +RoPE on Q and K   ]           |
   [ +causal mask       ]           |
        |                           |
        +----------> (add) <--------+
                       |
                       x1
                       |
        +--------------+------------+
        |                           |
     [RMSNorm]                      |
        |                           |
   [FFN: gate/up -> SwiGLU -> down] |
        |                           |
        +----------> (add) <--------+
                       |
                    output  (width d_model)
  1. Parameter accounting per layer, d = d_model:
Component Parameters GPT-2 small
Attention (4 matrices) 4 d^2 2,359,296
FFN (2 matrices, 4x) 8 d^2 4,718,592
Two LayerNorms 4 d 3,072
Block total 12 d^2 + 4d 7,080,960
  1. So the FFN is exactly two thirds of a classic block’s parameters, from the ratio 8 d^2 to 12 d^2.
  2. For Llama 3.1 8B with SwiGLU: attention is 41,943,040 and the FFN is 3 x 4096 x 14336 = 176,160,768, so the FFN is 80.8 per cent of the layer.
  3. Across the whole model, including the two 525,336,576-parameter vocabulary tables, the FFN share is 70.2 per cent and attention is 16.7 per cent. These figures match Chapter 47 exactly.
  4. Layer normalization is from Jimmy Lei Ba, Jamie Ryan Kiros and Geoffrey Hinton, July 2016. It normalizes over the feature axis of a single example, unlike batch normalization from Sergey Ioffe and Christian Szegedy in 2015, which normalizes over the batch axis.
  5. That distinction is why layer norm won in language models: batch statistics are useless when sequences have different lengths and batch size can be 1.
  6. RMSNorm is from Biao Zhang and Rico Sennrich, 2019. It drops the mean subtraction and the shift term, saving one reduction and d parameters per norm, for roughly the same quality.
  7. Epsilon inside the square root prevents division by zero. Typical values are 1e-5 for LayerNorm and 1e-5 or 1e-6 for RMSNorm. This is an implementation detail that occasionally matters for numerical reproducibility.
  8. Pre-norm was analysed formally by Ruibin Xiong and colleagues in the 2020 paper “On Layer Normalization in the Transformer Architecture”, which showed that post-norm gradients at initialization scale with the square root of depth while pre-norm gradients do not, explaining the warmup requirement.
  9. Baevski and Auli in 2018 and Nguyen and Salazar in 2019 had already reported the practical benefit. GPT-2, in February 2019, used pre-norm.
  10. A third variant exists: sandwich norm, or normalization both before and after the sub-layer, used in some very large training runs including Gemma 2 in 2024, to control activation growth.
  11. The residual stream view: because every block writes x + f(x), the model’s activations at layer L are the embedding plus the sum of every sub-layer’s output so far. This decomposition is exact and is the foundation of mechanistic interpretability work, notably “A Mathematical Framework for Transformer Circuits” from Anthropic in December 2021.
  12. Observation: activation norms typically grow with depth in trained transformers, often by an order of magnitude from layer 1 to the last layer. This is measurable with a forward hook and is a routine sanity check when training goes wrong.

WORDS48.8.6 remember these#

  1. Feed-forward network — the widen-then-narrow part applied to each position — FFN or MLP, two or three linear layers with a non-linearity, d_ff typically about 4 x d_model.
  2. Position-wise — applied to each token independently — the same weights applied to every position with no mixing across the sequence axis.
  3. Residual connection — adding the input back onto the output — the identity shortcut x + f(x), from ResNet in 2015.
  4. Residual stream — the main pipe running the depth of the model — the d_model-wide activation that every block reads from and adds to.
  5. Layer normalization — rescaling one token’s numbers to a standard size — normalizing over the feature axis per example, with learned scale and shift.
  6. RMSNorm — cheaper layer normalization — dividing by the root mean square with no mean subtraction and no shift term.
  7. Pre-norm — normalize before the sub-layer, add after — the arrangement that keeps the residual path free of normalization, standard since about 2019.
  8. Post-norm — normalize after adding — the 2017 arrangement, which needs learning-rate warmup to train stably.
  9. SwiGLU — a gated activation using three matrices — swish-gated linear unit, standard in the Llama, Mistral and Qwen families.

48.9 The whole model, assembled#

PLAIN48.9.1 in simple words#

  1. Everything so far has been parts. Here is the whole machine, in order.
  2. Step 1: your text is chopped into tokens and each token becomes an integer id. Chapter 47 covered this.
  3. Step 2: each id is used to look up a row in the embedding table. Now every token is a vector.
  4. Step 3: position information is injected. In modern models this happens inside attention, as a rotation, rather than as an addition here.
  5. Step 4: the stack. N identical blocks, one after another. Each block is attention then feed-forward, both with residual connections and normalization.
  6. Identical means identical in shape, not in values. Every block has its own separate parameters.
  7. Step 5: one final normalization at the top of the stack.
  8. Step 6: multiply by one last matrix that turns each d_model-wide vector into one number per token in the vocabulary. Those numbers are called logits.
  9. Step 7: softmax over the vocabulary turns the logits into probabilities that add up to 1.
  10. The output is one probability distribution per input position. For a 10-token prompt, you get 10 distributions.
  11. During generation you only use the last one, because that is the prediction for what comes next.
  12. During training you use all of them, because each position is a separate prediction problem, which is why masked parallel training is so efficient.

PLAIN48.9.2 a picture in your head#

  1. Picture a tall narrow building with 32 identical floors.
  2. On the ground floor, a crowd of people walks in. One person per token. Each is handed a card with numbers on it.
  3. On every floor, two things happen. First everybody talks to everybody and updates their card based on what they heard. Then each person sits alone and thinks, updating their card again.
  4. Nobody joins and nobody leaves. The same number of people walk out of floor 32 as walked into floor 1.
  5. The cards get richer on every floor. A person who arrived holding only the word “it” leaves floor 32 holding something that encodes what “it” refers to, what grammatical role it plays, and what is likely to come next.
  6. At the top, only the last person in the queue is asked a question: given everything you now know, what word comes next?
  7. That person’s card is compared against a list of 128,256 possible words, and a score is produced for each.

Where this comparison breaks: the people do not know anything and are not thinking. Each is a vector being multiplied by matrices. And the last person’s answer is not derived from a conversation, but from a single dot product of their card against every row of a very large table. Also, at training time every person is asked a question, not just the last one.

PLAIN48.9.3 a worked example#

  1. Trace one prompt through Llama 3.1 8B, showing the tensor shape at every stage. Configuration: vocabulary 128,256, d_model 4,096, 32 layers, 32 query heads, 8 key-value heads, head dimension 128, d_ff 14,336.
  2. The prompt is six tokens long, batch size 1.
stage                          shape                numbers
-----------------------------  -------------------  ----------
text                           "the cat was hungry"  -
token ids                      [1, 6]                6
embedding lookup               [1, 6, 4096]          24,576
block 1 input                  [1, 6, 4096]          24,576
  norm                         [1, 6, 4096]          24,576
  q after projection           [1, 32, 6, 128]       24,576
  k after projection           [1, 8, 6, 128]        6,144
  v after projection           [1, 8, 6, 128]        6,144
  k,v repeated to 32 heads     [1, 32, 6, 128]       24,576
  scores                       [1, 32, 6, 6]         1,152
  weights after softmax        [1, 32, 6, 6]         1,152
  attention output             [1, 32, 6, 128]       24,576
  after o_proj                 [1, 6, 4096]          24,576
  after residual add           [1, 6, 4096]          24,576
  norm                         [1, 6, 4096]          24,576
  ffn hidden (gate and up)     [1, 6, 14336]         86,016
  after down_proj              [1, 6, 4096]          24,576
  after residual add           [1, 6, 4096]          24,576
block 2 ... block 32           [1, 6, 4096]          24,576
final RMSNorm                  [1, 6, 4096]          24,576
lm_head projection             [1, 6, 128256]        769,536
softmax over vocabulary        [1, 6, 128256]        769,536
take last position             [1, 128256]           128,256
  1. Read down the shape column. From the embedding to the final norm, the shape never changes: [1, 6, 4096]. That is the residual stream, and it is constant width for the entire depth.
  2. The only places the shape changes are inside a block, temporarily, and at the very end.
  3. The last projection is the expensive one in terms of output size. It turns 24,576 numbers into 769,536 numbers, because the vocabulary is large.
  4. That final matrix, lm_head.weight, is [128256, 4096], holding 525,336,576 parameters. It is 6.5 per cent of the model on its own.
  5. Some models tie this matrix to the input embedding table, using the same numbers for both. GPT-2 does this. Llama 3.1 8B does not.
  6. Tying halves the vocabulary cost. Untying gives slightly better quality at larger scale. This is a real trade-off and both choices are in current use.

PLAIN48.9.4 what is really happening inside#

  1. A logit is a raw score before it becomes a probability. That is its only definition. It can be any real number, positive or negative.
  2. The name comes from statistics, where the logit function is the logarithm of the odds. In this context it just means “the thing softmax eats”.
  3. There is exactly one logit per vocabulary entry. For Llama 3.1, 128,256 of them, every single time a token is produced.
  4. Each logit is one dot product: the final hidden vector of length 4,096 dotted with one row of the output matrix, also length 4,096.
  5. So the model is asking, for every possible next token: how well does my current state line up with the direction that means this token?
  6. Softmax then converts. Exponentiate every logit, add them all up, divide.
  7. A logit of 4.0 against a logit of 3.2 becomes probabilities 0.57 and 0.26, because what matters is the difference, not the absolute values.
  8. Adding 100 to every logit changes nothing at all. Softmax is invariant to a constant shift. That is why logits are not comparable across models or even across steps.
  9. During training, the loss is the negative log of the probability the model gave to the token that actually came next. That is cross-entropy loss, from Chapter 46.
  10. If the model gave the right token probability 0.5, the loss is 0.693. If it gave 0.01, the loss is 4.605. The worse the guess, the larger the penalty, growing without bound.
  11. Nothing else is optimized. There is no separate objective for grammar, for facts, or for helpfulness during pretraining. One number, minimized.

TECHNICAL48.9.5 the engineer’s version#

  1. Full forward pass, in the order the code runs:
h = embed_tokens(ids)                    # [B, n, d]
for layer in layers:                     # N identical blocks
    r = h
    h = rmsnorm(h, layer.attn_norm)
    h = attention(h, rope, causal_mask)  # +RoPE inside
    h = r + h
    r = h
    h = rmsnorm(h, layer.ffn_norm)
    h = swiglu_ffn(h)
    h = r + h
h = rmsnorm(h, final_norm)               # [B, n, d]
logits = h @ lm_head.T                   # [B, n, vocab]
  1. Reference configurations, all published:
Model Layers d_model Vocab
Transformer base 2017 6 + 6 512 37,000
GPT-2 small 2019 12 768 50,257
Llama 3.1 8B 2024 32 4096 128,256
Llama 3.1 70B 2024 80 8192 128,256
  1. Total parameter count for Llama 3.1 8B, worked exactly: 32 layers of 218,112,000, plus 525,336,576 for the input embedding, plus 525,336,576 for the output projection, plus 4,096 for the final norm, giving 8,030,261,248.
  2. Activation memory for a forward pass is separate from weights and is dominated by the FFN hidden state. At n = 4096 and bfloat16, one layer’s [1, 4096, 14336] gate tensor is 117 MiB, and there are two of them.
  3. Weight tying is controlled by a flag, tie_word_embeddings, in the Hugging Face config. It is true for GPT-2 and for Gemma, false for Llama 3.1 8B.
  4. The final logit tensor is genuinely large. At n = 4,096 and vocabulary 128,256 in float32, the logits alone are 2.1 GiB. Training frameworks compute the loss in chunks along the sequence axis to avoid materializing it.
  5. At inference only the last position’s logits are needed, so implementations slice the hidden state to [B, 1, d] before the lm_head multiply. Skipping this is a common and expensive beginner mistake.
  6. Some models scale the logits. Gemma multiplies the embedding output by the square root of d_model; some models divide final logits by a constant, called logit softcapping, which Gemma 2 uses with a cap of 30.0 on the final layer.
  7. Numerical note: logits are computed in float32 even in bfloat16 models, because cross-entropy over 128,256 classes is sensitive to precision in the log-sum-exp reduction.

WORDS48.9.6 remember these#

  1. Logit — a raw score before it is turned into a probability — the pre-softmax output, one per vocabulary entry, unbounded in both directions.
  2. lm_head — the final matrix that scores every possible next token — the [vocab, d_model] output projection, sometimes tied to the input embedding.
  3. Weight tying — reusing the input embedding table as the output matrix — sharing one [vocab, d] tensor for both, halving vocabulary parameters.
  4. Residual stream — the constant-width pipe running the depth of the model — the [B, n, d_model] activation that all blocks read and write.
  5. Cross-entropy loss — how wrong the guess was, in one number — the negative log probability assigned to the correct next token, averaged over positions.
  6. Forward pass — running the model once from input to output — one evaluation of the network with no gradient computation.
  7. Hidden state — the working vector for one token at one depth — a d_model-wide slice of the residual stream at a given layer and position.

48.10 The 2017 paper#

PLAIN48.10.1 in simple words#

  1. The transformer was introduced in one paper: “Attention Is All You Need”.
  2. It was posted publicly on 12 June 2017, and presented at the NeurIPS conference in December 2017 in Long Beach, California.
  3. There were eight authors: Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser and Illia Polosukhin.
  4. Most were at Google Brain or Google Research. Aidan Gomez was at the University of Toronto, working with Google Brain.
  5. The paper has an unusual footnote saying the authors contributed equally and the listing order is random. That was deliberate and unusual for the time.
  6. It was a machine translation paper. Nothing in it is about chatbots, reasoning or general-purpose assistants.
  7. The claim was narrow: you can throw away recurrence entirely and get better translation quality in a fraction of the training time.
  8. The headline results were 28.4 BLEU on English-to-German and 41.8 BLEU on English-to-French, both on the WMT 2014 test sets.
  9. BLEU is a machine translation score from 0 to 100 that measures overlap with human reference translations. Higher is better.
  10. The 41.8 was a new single-model best, reached after 3.5 days on eight GPUs.
  11. That last number is the point of the whole paper. Competing systems needed far more compute for worse results.
  12. What was new was not attention. Attention had existed since 2014.
  13. What was new was deleting the recurrent network and keeping only attention, which made the whole sequence processable in parallel.

PLAIN48.10.2 a picture in your head#

  1. Imagine a road crew resurfacing a long road, working one metre at a time, because each metre must set before the next can be laid.
  2. Adding more workers does not help. The job is serial by construction.
  3. Now imagine somebody invents a material that sets independently everywhere, so a thousand workers can lay a thousand metres simultaneously.
  4. The material is not better. One metre of road is one metre of road. What changed is that the work is now divisible.
  5. That is what removing recurrence did. It converted a serial job into a parallel one, and hardware that was already sitting idle could finally be used.

Where this comparison breaks: attention is not merely a parallel version of the same computation. It also gives every position direct access to every other position, which recurrence never did. So the change was both a speed change and a capability change, and the paper’s own results reflect both.

PLAIN48.10.3 a worked example#

  1. The two model sizes reported:
Setting Base Big
Parameters 65 million 213 million
d_model 512 1024
Heads 8 16
Training time 12 hours 3.5 days
  1. Both used 6 encoder layers and 6 decoder layers, and were trained on 8 NVIDIA P100 GPUs in one machine.
  2. The base model ran 100,000 steps at about 0.4 seconds each, which is 12 hours.
  3. The big model ran 300,000 steps at about 1.0 second each, which is 3.5 days.
  4. Now the parallelism argument in concrete terms. Take a 1,000-token training sequence.
  5. A recurrent model must execute 1,000 dependent steps. Each step is a small matrix multiplication that cannot fill a GPU, and between each pair of steps the chip waits.
  6. A transformer executes one matrix multiplication over all 1,000 positions. With d_model = 512, that single call is 2 x 1000 x 512 x 512, which is 524 MFLOP of dense work handed to the chip in one piece.
  7. The chip’s utilization goes from a few per cent to near its limit, and the speed-up is far larger than the reduction in arithmetic.
  8. That is why the paper could report state-of-the-art translation from a single 8-GPU machine in under four days.

PLAIN48.10.4 what is really happening inside#

  1. The architecture in the paper is an encoder-decoder, not the decoder-only design used by chat models today.
  2. The encoder has 6 layers, each with bidirectional self-attention and a feed-forward network. It reads the source sentence.
  3. The decoder has 6 layers, each with three sub-layers: causal self-attention over what has been produced so far, cross-attention into the encoder’s output, and a feed-forward network.
  4. Post-norm was used throughout, with a warmup schedule that raised the learning rate for 4,000 steps then decayed it as the inverse square root of the step number.
  5. Dropout of 0.1 was applied to every sub-layer output and to the sum of the embedding and positional encoding.
  6. Label smoothing of 0.1 was used, which hurts perplexity and helps BLEU, and the paper says so plainly.
  7. Decoding used beam search with beam size 4 and length penalty 0.6.
  8. Every one of those settings is an ordinary engineering choice. None of them is the idea. The idea is one sentence: remove recurrence.

TECHNICAL48.10.5 the engineer’s version#

  1. Reported results, WMT 2014 test sets, from Table 2 of the paper:
System EN-DE BLEU EN-FR BLEU
Previous best single 25.16 41.29
Previous best ensemble 26.36 41.16
Transformer base 27.3 38.1
Transformer big 28.4 41.8
  1. Training cost in floating-point operations: 3.3e18 for the base model and 2.3e19 for the big model, both reported in the paper, and both one to two orders of magnitude below the competing systems in the same table.
  2. The paper also reported English constituency parsing results, to argue the architecture generalizes beyond translation. Those results were good but not record-setting, and they are rarely cited.
  3. The complexity table in the paper is worth knowing. Self-attention costs O(n^2 d) per layer with O(1) sequential operations and a maximum path length of O(1) between any two positions.
  4. A recurrent layer costs O(n d^2) with O(n) sequential operations and a maximum path length of O(n).
  5. So self-attention is cheaper than recurrence whenever n is less than d, which was true for translation sentences in 2017, where n was around 30 and d was
    1. The quadratic cost only became the dominant problem years later.
  6. Established fact: the architecture, the results and the training times are all in the published paper and have been reproduced many times.
  7. The honest version: the paper did not predict what the transformer became. It contains no mention of language modelling at scale, no scaling laws, no in-context learning and no suggestion that this design would generalize to images, audio, protein structure or code.
  8. The final sentence of the paper’s conclusion is about applying attention to images, audio and video, and about making generation less sequential. Those were reasonable next steps, not a prediction of GPT-3.
  9. The name “transformer” was chosen by Jakob Uszkoreit, and by several published accounts the team considered other names first. That is anecdote, not documented in the paper.
  10. Reproduction is easy today. The tensor2tensor library released alongside the paper contained the original code, and Alexander Rush’s annotated walkthrough from Harvard in 2018 reimplements the paper in a few hundred lines.

WORDS48.10.6 remember these#

  1. BLEU — a translation score based on word overlap — bilingual evaluation understudy, from Papineni and colleagues at IBM in 2002, scored 0 to 100.
  2. WMT — the standard machine translation benchmark — the Workshop on Machine Translation shared task, with the 2014 English-German and English-French sets still used as reference points.
  3. NeurIPS — the main machine learning conference — the Conference on Neural Information Processing Systems, called NIPS until 2018.
  4. Encoder-decoder — one stack reads, another writes — the original transformer layout, with cross-attention linking the two.
  5. Warmup — raising the learning rate slowly at the start — a schedule required by post-norm transformers, 4,000 steps in the 2017 paper.
  6. Label smoothing — never demanding 100 per cent confidence — replacing the one-hot target with 0.9 on the correct class and the rest spread out.

48.11 The three architectures#

PLAIN48.11.1 in simple words#

  1. The 2017 paper had two stacks: an encoder that reads and a decoder that writes. The field then split them apart and used them separately.
  2. Encoder-only: keep the reading stack, throw away the writing stack. No causal mask, so every word sees every other word in both directions.
  3. That is BERT, from Google in October 2018. It is superb at understanding a fixed piece of text and cannot generate text at all.
  4. Use it for classification, for search, for named-entity extraction, and for producing embedding vectors of documents.
  5. Decoder-only: keep the writing stack, throw away the reading stack. Causal mask, so every word sees only what came before.
  6. That is the GPT line, and essentially every chat model you have ever used.
  7. Encoder-decoder: keep both, as in the original. The encoder reads the input, the decoder writes the output, and cross-attention connects them.
  8. That is T5, and most dedicated translation and speech-recognition models.
  9. Decoder-only won for general-purpose models, and the reason is worth stating plainly.
  10. A decoder-only model can do everything by writing. Ask it to classify, and it writes the class name. Ask it to translate, and it writes the translation.
  11. Every task becomes text in, text out, so one model and one training objective covers all of them.
  12. It also trains on more signal per token, needs no separate cross-attention machinery, and gets a prediction target at every position.

PLAIN48.11.2 a picture in your head#

  1. Think of three kinds of worker.
  2. The encoder-only worker is a proofreader. Given a finished page, they read it forwards and backwards as many times as they like and tell you what it is about. They never write anything.
  3. The decoder-only worker is a novelist writing live, one word at a time, never allowed to see the words they have not written yet.
  4. The encoder-decoder worker is a translator. They read the whole source document freely, then write the translation one word at a time, glancing back at the source whenever they need to.
  5. The novelist turned out to be the most useful, because you can ask a novelist to write a summary, a classification or a translation, and they will.

Where this comparison breaks: a proofreader could write if asked. BERT genuinely cannot, because it was never trained to predict a next token and its attention has no causal structure to build on. The limitation is architectural and training-based, not a matter of instruction.

PLAIN48.11.3 a worked example#

  1. Real models, sorted by architecture:
Model Year Architecture
BERT 2018 Encoder-only
RoBERTa 2019 Encoder-only
DeBERTa 2020 Encoder-only
T5 2019 Encoder-decoder
BART 2019 Encoder-decoder
Whisper 2022 Encoder-decoder
GPT-2, GPT-3 2019, 2020 Decoder-only
Llama, Mistral 2023 on Decoder-only
Qwen, Gemma 2023 on Decoder-only
DeepSeek-V3 2024 Decoder-only
  1. Sizes for the encoder models: BERT-base has 110 million parameters, 12 layers and d_model 768. BERT-large has 340 million, 24 layers and d_model 1024.
  2. Notice that encoder-only models stopped getting bigger. Nobody built a 100-billion-parameter BERT, because the returns went to generative models.
  3. Encoder-only models are still heavily used, and not as a legacy. Almost every semantic search system in production embeds text with a small encoder, because it is cheap, bidirectional and fast.
  4. A concrete size comparison for a search system: an encoder producing a 384- dimensional embedding might have 22 million parameters and run in under a millisecond per short document on a CPU.

PLAIN48.11.4 what is really happening inside#

  1. The architectural differences reduce to two switches: is attention masked, and is there a second stack with cross-attention.
  2. Encoder-only: no mask, one stack, no cross-attention.
  3. Decoder-only: causal mask, one stack, no cross-attention.
  4. Encoder-decoder: no mask in the encoder, causal mask in the decoder, plus cross-attention layers in the decoder.
  5. Everything else, the blocks, the residuals, the normalization, the feed-forward networks, is identical in all three.
  6. So they are not three architectures in any deep sense. They are one architecture with different masking and different wiring.
  7. The training objective differs too, and that matters more than the wiring.
  8. Encoder-only models are trained with masked language modelling, which corrupts about 15 per cent of tokens, giving a learning signal at 15 per cent of positions.
  9. Decoder-only models are trained with next-token prediction, which gives a signal at 100 per cent of positions.
  10. Encoder-decoder models like T5 use span corruption, dropping contiguous runs of tokens and asking the decoder to reproduce them.
  11. That efficiency difference compounds over trillions of tokens, and it is a large part of why the decoder-only line scaled further.

TECHNICAL48.11.5 the engineer’s version#

  1. T5, from Colin Raffel and colleagues at Google in October 2019, cast every task as text-to-text with a task prefix, and scaled from 60 million to 11 billion parameters.
  2. T5 established the text-to-text framing that decoder-only models later adopted without needing the encoder.
  3. The empirical comparison exists. The 2022 paper “What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization?” by Thomas Wang and colleagues found that decoder-only with causal language modelling gave the best zero-shot performance after pretraining alone, while encoder-decoder with masked objectives was stronger after multitask fine-tuning.
  4. So the decoder-only win is not absolute. It is a win under the specific regime that turned out to matter: pretrain once, then prompt.
  5. Practical reasons decoder-only dominates in deployment:
    1. One stack means one set of key-value caches, so serving is simpler.
    2. The prompt and the generation share the same cache, so a long prompt costs one prefill pass and nothing more.
    3. Any input format can be expressed as text, so no schema changes are needed per task.
  6. Encoder-decoder still wins where the input is genuinely a different modality or language from the output. Whisper, released by OpenAI in September 2022, is an audio encoder with a text decoder, and that split is natural.
  7. Encoder-only remains standard for retrieval. The sentence-transformers library, from Nils Reimers and Iryna Gurevych in 2019, is the usual tool, and the MTEB leaderboard is the usual benchmark.
  8. A hybrid worth naming: many modern retrieval systems now use decoder-only models with the causal mask removed for embedding, which works but requires adaptation training.

WORDS48.11.6 remember these#

  1. Encoder-only — reads text both ways, cannot write — a bidirectional transformer stack trained with masked language modelling.
  2. Decoder-only — writes text one token at a time — a causally masked stack trained with next-token prediction.
  3. Encoder-decoder — reads one sequence, writes another — two stacks joined by cross-attention, the original 2017 layout.
  4. Bidirectional — able to use both left and right context — attention with no causal mask.
  5. Span corruption — deleting runs of words and asking for them back — the T5 pretraining objective, replacing spans with sentinel tokens.
  6. Text-to-text — every task expressed as text in and text out — the T5 framing, inherited by all current chat models.

48.12 Generation: producing text one token at a time#

PLAIN48.12.1 in simple words#

  1. The model does not write a sentence. It writes one token, then starts again with that token appended.
  2. The loop is simple. Feed in the prompt. Get a probability for every possible next token. Choose one. Append it. Feed in the whole thing again.
  3. Stop when a special end-of-text token is chosen, or when a length limit is reached.
  4. Every word you have ever seen a language model produce came out of this loop, one token at a time.
  5. The interesting decision is step three: choose one. There are several ways.
  6. Greedy: always take the highest-probability token. Fully deterministic and often dull and repetitive.
  7. Temperature: divide all the logits by a number before softmax. Below 1 sharpens the distribution towards the top choice. Above 1 flattens it.
  8. Top-k: keep only the k most likely tokens, throw the rest away, and sample from what remains after renormalizing.
  9. Top-p, also called nucleus sampling: keep the smallest set of tokens whose probabilities add up to p, typically 0.9 or 0.95.
  10. Min-p: keep tokens whose probability is at least some fraction of the top token’s probability. It adapts to how confident the model is.
  11. Repetition and frequency penalties: reduce the score of tokens that have already appeared, to stop loops.
  12. Beam search: keep several candidate continuations alive at once and pick the best complete sequence at the end.
  13. This is why the same prompt gives different answers. Unless you use greedy decoding, a random draw happens at every single token.

PLAIN48.12.2 a picture in your head#

  1. Picture a weather forecaster who can only forecast one hour ahead.
  2. To forecast a whole day, they forecast hour one, write it down, then forecast hour two assuming hour one happened exactly as written.
  3. Errors compound. If hour three is wrong, hours four to twenty-four are built on a mistake and cannot recover, because the forecaster treats their own written output as fact.
  4. That is autoregressive generation, and it explains why long outputs drift.
  5. Temperature is the forecaster’s willingness to name an unlikely outcome. At temperature 0 they always say the single most likely thing. At high temperature they will occasionally forecast snow in July.

Where this comparison breaks: the forecaster knows their earlier forecast was a guess. The model has no such flag. Once a token is in the context, it is indistinguishable from text you wrote yourself, which is exactly why a model can confidently build on something it invented three sentences ago.

PLAIN48.12.3 a worked example#

  1. Take a real five-way choice. The model has produced “the cat sat on the” and these are the top five candidate logits:
token    logit
mat       4.0
floor     3.2
chair     2.5
roof      1.1
moon      0.4
  1. Softmax at temperature 1.0 gives:
mat    0.5699    floor  0.2561    chair  0.1272
roof   0.0314    moon   0.0156
  1. Temperature divides every logit by T before the softmax. Here is the same distribution at several temperatures, computed exactly:
Token T=0.5 T=1.0 T=2.0
mat 0.7965 0.5699 0.3933
floor 0.1608 0.2561 0.2636
chair 0.0397 0.1272 0.1858
roof 0.0024 0.0314 0.0923
moon 0.0006 0.0156 0.0650
  1. At T = 0.5 the top token has 79.65 per cent, and “moon” has 0.06 per cent. At T = 2.0 the top token is down to 39.33 per cent and “moon” is up to 6.5 per cent, a hundredfold increase.
  2. At T = 0.1 the top token reaches 0.9997, which is greedy decoding in all but name. Temperature 0 is usually implemented as exactly greedy, since dividing by zero is undefined.
  3. Top-k with k = 2: keep “mat” and “floor” only. Renormalize 0.5699 and 0.2561 by their sum 0.8260, giving 0.6900 and 0.3100. “chair” is now impossible.
  4. Top-p with p = 0.9: accumulate from the top. 0.5699, then 0.8260, then 0.9531 which passes 0.9. So three tokens are kept, renormalized to 0.5979, 0.2687 and 0.1334.
  5. Note that top-p with 0.9 and 0.95 give the same set here, because the cumulative sum jumps from 0.826 to 0.953 in one step.
  6. Min-p with 0.1: the threshold is 0.1 times the top probability, so 0.1 x 0.5699 = 0.057. Tokens above that are “mat”, “floor” and “chair”. Same set as top-p here.
  7. The methods differ on a peaked distribution. If the top token has 0.9947 and the rest share 0.0053, min-p keeps only “mat”, because the threshold is 0.0995 and nothing else clears it. Top-p at 0.9 also keeps only “mat”.
  8. On a flat distribution where the top is 0.2419, min-p’s threshold is 0.0242 and all five survive, which is the point: the cut adapts to confidence.
  9. Repetition penalty of 1.2 on a token already generated divides its logit by 1.2 when positive. “mat” goes from 4.0 to 3.3333, and its probability falls from 0.5699 to 0.4048.
  10. Frequency penalty of 0.5 subtracts 0.5 for each previous occurrence. If “mat” appeared twice, its logit drops from 4.0 to 3.0 and “floor” becomes the most likely token at 0.4002.

PLAIN48.12.4 what is really happening inside#

  1. Generation has two phases, and they behave completely differently.
  2. Prefill: the whole prompt goes through the model in one pass. All prompt positions are computed in parallel. This is compute-bound and fast per token.
  3. Decode: each new token requires a full forward pass over one position. The model reads every weight in the file to produce one token.
  4. That is why decode is memory-bandwidth-bound. For an 8-billion model at 16-bit, each token requires reading 16 GB of weights from memory.
  5. On a card with 2,000 GB/s of memory bandwidth, that sets a hard ceiling of about 125 tokens per second for a single sequence, before any other cost.
  6. Beam search keeps b partial sequences alive, extends each by every candidate token, keeps the b best by total log probability, and repeats.
  7. It optimizes for the highest-probability complete sequence, which is exactly what you want in translation, where there is one correct output.
  8. It is wrong for chat for two reasons. The highest-probability continuation of open-ended text is bland and repetitive, a result documented by Ari Holtzman and colleagues in 2019. And beam search costs b times the compute and b times the key-value cache.
  9. So translation systems use beam 4 or 5. Chat systems use sampling with temperature and top-p, and never beam search.
  10. Why does the same prompt give different answers? Three separate reasons.
  11. First, sampling. Every token is a random draw, so different random seeds give different text.
  12. Second, even at temperature 0, floating-point addition is not associative. Changing the batch size changes the order of reductions inside the matrix multiplications, which changes the last bits of the logits, which can flip a near-tie between two tokens.
  13. Third, providers change models, system prompts and serving configurations without notice. Determinism across dates is not offered by any major API.

TECHNICAL48.12.5 the engineer’s version#

  1. The autoregressive factorization: P(x_1 … x_n) = product over t of P(x_t given x_1 … x_{t-1}). Every generation method is a search or sampling strategy over this factorization.
  2. Temperature applies as softmax(z / T). It is a monotonic transform, so it never changes the ranking, only the gaps.
  3. Method origins:
Method Year Source
Beam search pre-1980s speech recognition
Top-k 2018 Fan, Lewis, Dauphin
Top-p (nucleus) 2019 Holtzman and colleagues
Repetition penalty 2019 Keskar and colleagues, CTRL
Min-p 2024 Nguyen and colleagues
  1. Typical production defaults as of 2026: temperature 0.7 to 1.0, top_p 0.9 to 1.0, top_k either 40 or disabled, repetition penalty 1.0 to 1.1. These are conventions, and every provider picks slightly different ones.
  2. OpenAI’s API exposes frequency_penalty and presence_penalty, both in the range -2.0 to 2.0, applied additively to logits. Presence penalty applies once if a token has appeared at all; frequency penalty scales with the count.
  3. The Hugging Face generate method applies processors in a fixed order: penalties first, then temperature, then top-k, then top-p, then min-p, then the sample. Order matters and is worth checking when results surprise you.
  4. Greedy decoding is do_sample=False. Setting temperature to 0 in most APIs is mapped internally to greedy.
  5. Speculative decoding, from Yaniv Leviathan, Matan Kalman and Yossi Matias at Google in 2022 and independently from DeepMind, uses a small draft model to propose several tokens and the large model to verify them in one pass. It is exact, meaning the output distribution is unchanged, and gives two to three times speedup in practice.
  6. Constrained decoding masks logits to force valid JSON or a regular expression. Libraries such as Outlines and llama.cpp’s GBNF grammars do this by zeroing the probability of any token that would break the grammar.
  7. Established fact: sampling parameters change output distribution in exactly the ways computed above. Active research: which defaults are best for which task. Marketing claim: any statement that a specific temperature makes a model “more creative” in a measurable, general sense.

WORDS48.12.6 remember these#

  1. Autoregressive generation — writing one token at a time, feeding each back in — sampling from P(x_t given all previous tokens) in a loop.
  2. Greedy decoding — always take the most likely token — argmax over logits, fully deterministic given identical arithmetic.
  3. Temperature — a dial that flattens or sharpens the choice — the divisor T in softmax(z / T), with T below 1 sharpening.
  4. Top-k sampling — only consider the k best candidates — truncate to the k highest logits, renormalize, sample.
  5. Top-p sampling — consider just enough candidates to cover p of the mass — nucleus sampling, keeping the smallest prefix of the sorted distribution summing to at least p.
  6. Min-p sampling — keep anything within a fraction of the top choice — a threshold of p times max probability, adapting to model confidence.
  7. Beam search — keep several drafts alive and pick the best — breadth-limited search over sequences, standard for translation, unsuitable for chat.
  8. Prefill — processing the whole prompt at once — the parallel phase of inference, compute-bound.
  9. Decode — producing one token per forward pass — the sequential phase, memory-bandwidth-bound.
  10. Speculative decoding — a small model guesses, a big model checks — a draft-and-verify scheme that preserves the output distribution exactly.

48.13 The cost of attention: the quadratic problem#

PLAIN48.13.1 in simple words#

  1. Attention compares every token with every token. For n tokens that is n times n comparisons.
  2. Double the length and the work quadruples. Ten times the length and the work goes up a hundredfold.
  3. That is what “quadratic” means, and it is the single biggest engineering problem in the architecture.
  4. At 1,000 tokens it is a million comparisons, which is nothing.
  5. At 1,000,000 tokens it is a trillion comparisons, which is an enormous amount for a single layer of a single head.
  6. There is a second cost, separate and often larger in practice: the key-value cache.
  7. When generating, the model has already computed the keys and values for every earlier token. Recomputing them every step would be wasteful, so they are stored.
  8. That store grows with every token generated, and it must live in fast memory next to the chip.
  9. For a 70-billion-parameter model, every token of context costs 320 kilobytes of cache. A 128,000-token conversation costs 40 gigabytes.
  10. That is more than the model itself in some configurations, and it is why long-context serving is expensive.
  11. There are four families of fix, and three of them are in production today.
  12. Make the exact computation memory-smarter (FlashAttention). Share keys and values across heads (grouped-query attention). Restrict which tokens can be looked at (sliding window, sparse). Or replace attention with something cheaper (linear attention, state-space models).

PLAIN48.13.2 a picture in your head#

  1. Picture a party where every guest must shake hands with every other guest.
  2. With 10 guests that is 45 handshakes. Manageable.
  3. With 1,000 guests it is 499,500. With 100,000 guests it is about 5 billion.
  4. The number of guests grew by a factor of 100 and the handshakes grew by a factor of 10,000.
  5. Nobody solves this by shaking hands faster. You solve it by changing the rules: only shake hands with the nearest twenty people, or shake hands in small groups, or appoint representatives.
  6. Every one of those is a real technique in the list of fixes.

Where this comparison breaks: a handshake is symmetric and attention is not. Token 5 attending to token 2 is a different computation from token 2 attending to token 5, and in a causal model only one of the two happens at all. So the true count is n(n+1)/2, not n(n-1)/2, and the arithmetic differs by a factor of about two rather than the shape being different.

PLAIN48.13.3 a worked example#

  1. The score matrix has n squared entries. Here is what that means, with the memory it would take at 16 bits per number, for a single head in a single layer:
Tokens Entries Memory, one head
1,024 1,048,576 2 MB
32,768 1,073,741,824 2,048 MB
131,072 17,179,869,184 32,768 MB
1,048,576 1,099,511,627,776 2,097,152 MB
  1. Read the last row again. Two terabytes, for one head, in one layer, for one sequence.
  2. A model like Llama 3.1 8B has 32 heads in each of 32 layers. If you tried to store the score matrices for a 131,072-token sequence in one layer, that is 32 times 32,768 MB, which is 1,024 GiB.
  3. No accelerator has that memory. This is why the naive implementation simply cannot run at long context, and why FlashAttention exists.
  4. Now the cost of attention relative to the rest of the model. The score computation costs about 4 n squared d operations per layer, and the feed-forward network costs about 16 n d squared.
  5. Setting those equal gives a crossover at n = 4d. For Llama 3.1 8B with d = 4096, that is 16,384 tokens.
  6. Below 16,384 tokens the feed-forward network dominates the compute. Above it, attention does. Measured for that configuration:
Sequence length Attention share
512 2 per cent of FFN
4,096 19 per cent of FFN
24,576 114 per cent of FFN
131,072 609 per cent of FFN
  1. Now the key-value cache, computed exactly. The formula is 2 x layers x kv_heads x head_dim x bytes, per token.
  2. For Llama 3.1 8B: 2 x 32 x 8 x 128 x 2 = 131,072 bytes per token, which is exactly 128 KiB.
Context Llama 3.1 8B Llama 3.1 70B
2,048 tokens 0.25 GiB 0.62 GiB
32,768 tokens 4.00 GiB 10.00 GiB
131,072 tokens 16.00 GiB 40.00 GiB
  1. Without grouped-query attention, using 32 key-value heads instead of 8, the 8B model’s cache would be 512 KiB per token, so 64 GiB at full context instead of 16 GiB. That single design choice saves four times the memory.
  2. Batching multiplies this. Serving 32 concurrent conversations of 8,192 tokens each on Llama 3.1 8B costs 32 GiB of cache, on top of 16 GB of weights.
  3. That is the actual reason long context costs more per token in commercial APIs. It is memory, not arithmetic.

PLAIN48.13.4 what is really happening inside#

  1. FlashAttention is the most important fix, and it is important to understand what it is not. It is not an approximation. Its output is exactly the same as the naive implementation, bit-for-bit equivalent up to floating-point reassociation.
  2. The insight is that the n-by-n score matrix never needs to exist all at once.
  3. The naive implementation writes the whole score matrix to the chip’s main memory, reads it back for the softmax, writes it again, and reads it again for the multiply by V.
  4. Those trips to main memory are the bottleneck, not the arithmetic. Chapter 22 covered why: a GPU’s arithmetic units are far faster than its memory.
  5. FlashAttention splits Q, K and V into tiles that fit in the chip’s small fast on-chip memory, and computes the output tile by tile, keeping a running softmax normalizer as it goes.
  6. The running softmax is possible because of an algebraic trick published by Maxim Milakov and Natalia Gimelshein in 2018, called online softmax, which lets you update a softmax result as new terms arrive.
  7. The score matrix is therefore never stored. Memory use drops from order n squared to order n, and the speed rises because far fewer bytes move.
  8. Grouped-query attention attacks the cache instead. Instead of 32 separate key and value heads, use 8, and let four query heads share each one.
  9. The cache shrinks by exactly that ratio, and the quality loss is small.
  10. Sliding window attention simply forbids looking back more than w tokens. Cost becomes linear in n, and information travels further only by hopping layer by layer, at w tokens per layer.
  11. Linear attention rewrites the formula so the value sum can be accumulated without ever forming the score matrix, giving linear cost at the price of a weaker, approximate similarity function.

TECHNICAL48.13.5 the engineer’s version#

  1. FlashAttention, from Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra and Christopher Re, May 2022. FlashAttention-2 from Tri Dao, July 2023, improved work partitioning. FlashAttention-3, July 2024, targets Hopper-class hardware with asynchrony and low-precision paths.
  2. Complexity: standard attention uses O(n^2) memory. FlashAttention uses O(n) and reduces high-bandwidth-memory accesses from O(n^2 d) to about O(n^2 d^2 / M), where M is on-chip SRAM size. Reported speedups on training are roughly 2x to 4x end to end.
  3. It recomputes the attention matrix during the backward pass rather than storing it, trading a little arithmetic for a large memory saving. This is gradient checkpointing applied at kernel level.
  4. Multi-query attention, from Noam Shazeer in November 2019, uses one key-value head for all query heads. Grouped-query attention, from Joshua Ainslie and colleagues at Google in May 2023, interpolates between the two.
Scheme KV heads (32 q heads) Cache size
Multi-head (MHA) 32 1.0x
Grouped-query (GQA) 8 0.25x
Multi-query (MQA) 1 0.031x
  1. Multi-head latent attention, introduced in DeepSeek-V2 in 2024 and used in DeepSeek-V3, compresses keys and values into a low-rank latent vector and reconstructs them, cutting cache further than GQA at comparable quality.
  2. Sliding window attention: Longformer, from Iz Beltagy, Matthew Peters and Arman Cohan in April 2020, combined a local window with a few global tokens. Mistral 7B, October 2023, used a 4,096-token window with 32 layers, giving a theoretical receptive field of 131,072 tokens.
  3. Sparse attention: the Sparse Transformer, from Rewon Child, Scott Gray, Alec Radford and Ilya Sutskever in April 2019, used strided and fixed patterns to reach O(n sqrt(n)). BigBird, from Manzil Zaheer and colleagues in July 2020, combined window, global and random attention and proved it is a universal approximator.
  4. Linear attention: Angelos Katharopoulos and colleagues in June 2020 showed that replacing softmax with a kernel feature map lets you reorder the multiplications to get O(n) cost, and that the resulting model is a recurrent network at inference. The paper is titled “Transformers are RNNs”.
  5. Performer, from Krzysztof Choromanski and colleagues in September 2020, approximates the softmax kernel with random features. Reformer, from Nikita Kitaev, Lukasz Kaiser and Anselm Levskaya in January 2020, used locality-sensitive hashing.
  6. The honest version: the approximate methods have largely not displaced exact attention in frontier models. Exact attention plus FlashAttention plus GQA is what production systems actually run as of 2026. The approximations remain valuable in specific niches.
  7. State-space models are the serious current alternative. Mamba, from Albert Gu and Tri Dao in December 2023, uses a selective state-space layer with linear cost in sequence length and a constant-size state, so there is no growing key-value cache at all.
  8. Mamba-2 followed in 2024 and showed a formal duality between state-space models and a restricted form of attention. Jamba, from AI21 Labs in March 2024, is a production hybrid interleaving Mamba and attention layers.
  9. Active research, stated plainly: hybrid models are being trained at scale and show competitive quality with much better long-context economics. Whether they displace attention entirely is not known. Pure state-space models have measured weaknesses on tasks requiring precise recall of arbitrary earlier tokens, which is exactly what attention is good at.
  10. Observation tooling: nvidia-smi shows memory use, but for cache accounting the useful figures come from the serving framework. vLLM, which implements PagedAttention from Woosuk Kwon and colleagues in 2023, reports cache block usage directly and manages the cache in fixed-size pages like an operating system manages virtual memory.

WORDS48.13.6 remember these#

  1. Quadratic complexity — work grows with the square of the length — O(n^2) in sequence length for the score matrix and the weighted sum.
  2. KV cache — stored keys and values so earlier tokens are not recomputed — 2 x layers x kv_heads x head_dim x bytes per token of context.
  3. FlashAttention — computing attention without ever storing the score matrix — an exact, IO-aware tiled kernel with online softmax and backward recomputation.
  4. Online softmax — updating a softmax as new terms arrive — a rescaling recurrence that avoids a second pass over the scores.
  5. Multi-query attention — one key-value head shared by all query heads — MQA, cutting cache by the number of heads at some quality cost.
  6. Grouped-query attention — a few key-value heads shared in groups — GQA, the standard compromise, 8 KV heads for 32 query heads in Llama 3.
  7. Sliding window attention — only look back a fixed distance — local attention with window w, giving linear cost and a layer-multiplied receptive field.
  8. State-space model — a linear recurrence with a fixed-size state — the Mamba family, linear in sequence length with no growing cache.
  9. PagedAttention — managing the cache in fixed-size blocks — the vLLM technique that removes cache fragmentation, modelled on virtual memory paging.

48.14 Mixture of experts: paying for a big model, running a small one#

PLAIN48.14.1 in simple words#

  1. Every transformer block has two halves. Attention mixes information between tokens. The feed-forward network then thinks about each token on its own.
  2. In an ordinary model there is exactly one feed-forward network per block, and every token goes through it. That block holds most of the model’s numbers.
  3. A mixture of experts model replaces that one feed-forward network with many copies of it, called experts, sitting side by side.
  4. A tiny extra network called the router looks at each token and picks a small number of experts for it. Usually two, four or eight out of dozens.
  5. Only the chosen experts run. The rest sit still and cost nothing in arithmetic for that token.
  6. So the model is enormous when you count its numbers, and small when you count the work done for any one token.
  7. This gives two different parameter counts, and the difference between them is the whole point of the idea.
  8. Total parameters is every number in the file. It sets how much memory you need and how much a copy costs to store and to move.
  9. Active parameters is how many numbers actually take part in producing one token. It sets the speed, the electricity bill and the training cost.
  10. A real example: DeepSeek-V3, released December 2024, holds 671 billion parameters in total and uses 37 billion for each token. That is 5.5 per cent.
  11. So it costs as much arithmetic per token as a 37-billion model, while holding the knowledge capacity of something eighteen times larger.
  12. Nothing is free. You must keep all 671 billion numbers loaded, because you cannot know in advance which experts the next token will want.
  13. That is the trade in one line: you pay in memory, and you are paid back in speed.
  14. There is one nasty problem. Left alone, the router learns to send almost every token to the same two or three favourite experts, and the rest die.
  15. Fixing that needs a deliberate extra pressure during training, and most of the engineering in this area is about exactly that.

PLAIN48.14.2 a picture in your head#

  1. Picture a large hospital with 64 specialist doctors and one triage nurse at the door.
  2. You walk in. The nurse spends two seconds on you and sends you to two of the 64 doctors. You see those two. The other 62 never learn you existed.
  3. The hospital still has to employ all 64, pay all 64, and keep 64 consulting rooms heated. That is the memory cost.
  4. But your visit only occupies two doctors’ time. That is the compute cost.
  5. Now the failure. Suppose the nurse sends nearly everyone to Doctor Smith because Doctor Smith seemed good on the first day.
  6. Doctor Smith sees thousands of patients and gets very good. The other 63 see nobody and stay bad. So the nurse sends even more people to Smith.
  7. You have not built a hospital with 64 specialists. You have built an expensive building containing one overworked doctor.
  8. To stop this the hospital management adds a rule: the nurse is penalized when the workload is uneven, whatever the nurse thinks is best.
  9. That rule is real, and it has a name in the papers: a load balancing loss.

Where this comparison breaks: real experts are not topic specialists, and this is the single most over-sold part of the idea. The Mixtral paper of January 2024 looked for topic-based routing and did not find it. Tokens about medicine do not reliably go to a “medicine expert”. What routing does correlate with is much duller: token identity, position, and surface patterns such as indentation in code. Also, a patient sees one doctor and gets one opinion. A token gets a weighted blend of two experts’ outputs, added together in proportion to the router’s confidence.

PLAIN48.14.3 a worked example#

  1. We will build a mixture-of-experts model from a normal one, count the parameters exactly, and land on a real published model.
  2. Start with the dense shape used earlier in this chapter: d_model 4096, feed forward width 14,336, SwiGLU with three matrices, 32 layers, 8 key-value heads, vocabulary 32,000.
  3. One feed-forward network is 3 x 4096 x 14336 = 176,160,768 parameters.
  4. Attention in one layer is 41,943,040 parameters. Two norms are 8,192.
  5. Now replace the single feed-forward network with 8 copies, and add a router that maps 4096 numbers to 8 scores: 4096 x 8 = 32,768 parameters.
  6. The expert bank in one layer is 8 x 176,160,768 = 1,409,286,144.
  7. One layer therefore holds:
attention        41,943,040
experts (8)   1,409,286,144
router               32,768
two norms             8,192
layer total   1,451,270,144
  1. Thirty-two layers give 46,440,644,608. Add the two vocabulary tables, 2 x 32,000 x 4,096 = 262,144,000, and a final norm of 4,096.
  2. Total: 46,702,792,704 parameters, which is 46.7 billion.
  3. That is Mixtral 8x7B, published by Mistral AI in December 2023, whose stated size is 46.7B. The arithmetic reproduces the published number exactly.
  4. Now count what one token actually touches, with top-2 routing.
attention             41,943,040
2 experts of 8       352,321,536
router                    32,768
two norms                  8,192
per layer active     394,305,536
x 32 layers       12,617,777,152
+ two vocab tables   262,144,000
active total      12,879,921,152
  1. That is 12.9 billion active parameters, which is also exactly the published figure. The ratio is 46.7 to 12.9, or 3.6 to 1.
  2. Note what did not change: attention is dense. Every token runs the full attention block. Only the feed-forward half is sparse.
  3. Now push the idea harder. Cut each expert into 8 narrower ones, so 64 experts, and route to 8 instead of 2.
  4. The arithmetic per token is identical, because 8 narrow experts equal 2 wide ones. But the number of possible expert combinations explodes.
  5. Choosing 2 from 8 gives 28 combinations. Choosing 8 from 64 gives 4,426,165,368. That is the argument for fine-grained experts.
  6. Here is what real models do, as published by each laboratory:
Model Total params Active per token
Mixtral 8x7B, 2023 46.7B 12.9B
DBRX, 2024 132B 36B
Mixtral 8x22B, 2024 141B 39B
Grok-1, 2024 314B about 79B
DeepSeek-V3, 2024 671B 37B
Llama 4 Scout, 2025 109B 17B
Qwen3-235B-A22B, 2025 235B 22B
Llama 4 Maverick, 2025 400B 17B
gpt-oss-120b, 2025 117B 5.1B
Kimi K2, 2025 1,000B 32B
  1. Now the same table as one number, the fraction of the model used per token, sorted by date. The trend is the story:
Model Year Active fraction
Mixtral 8x7B 2023 27.6 per cent
Grok-1 2024 25.0 per cent
DeepSeek-V3 2024 5.5 per cent
Llama 4 Scout 2025 15.6 per cent
Qwen3-235B-A22B 2025 9.4 per cent
gpt-oss-120b 2025 4.4 per cent
Kimi K2 2025 3.2 per cent
  1. Models are getting sparser. In 2023 a mixture-of-experts model used about a quarter of itself per token. By 2025, three per cent was normal.
  2. Finally the memory bill. Mixtral 8x7B at 16 bits is 93.4 gigabytes of weights, which does not fit on one 80-gigabyte accelerator.
  3. A dense model of the same speed, 12.9 billion parameters, would be 25.8 gigabytes and would fit on one card with room for a cache.
  4. So you pay roughly 3.6 times the memory for roughly 1 times the arithmetic. Whether that is a good deal depends entirely on how many users you serve.

PLAIN48.14.4 what is really happening inside#

  1. Here is the whole mechanism for one token at one layer.
      token vector x  (width d_model)
             |
      [ Router: x times W_r -> one score per expert ]
             |
      softmax, then keep the top 2
             |
      e5 (0.71)          e2 (0.29)
        |                   |
   [Expert 5]          [Expert 2]
        |                   |
     y5 x 0.71           y2 x 0.29
         \                 /
          +---- (add) ----+
                  |
               output  (width d_model)
  1. The router is one small matrix, of shape d_model by number of experts. For Mixtral that is 4096 by 8. It is trained by ordinary gradient descent.
  2. It produces one score per expert. A softmax turns those into probabilities.
  3. Top-k gating keeps the k highest, throws the rest away, and renormalizes the survivors so they add to 1.
  4. Only the surviving experts are run. Their outputs are multiplied by their renormalized weights and added together.
  5. The gate weight is what makes this trainable. Without it, the choice of expert would be a hard, non-differentiable decision with no gradient.
  6. With it, the gradient can flow back through the weight and tell the router “this expert helped, raise its score” or “this one hurt, lower it”.
  7. Now the collapse. The router and the experts learn at the same time, which creates a feedback loop.
  8. An expert that gets slightly more tokens early gets more gradient updates, improves faster, produces better outputs, and so earns higher router scores.
  9. That earns it more tokens still. Within a few thousand steps a handful of experts absorb nearly everything and the others are effectively dead weight.
  10. The standard fix is a load balancing loss: an extra number added to the training loss that is smallest when every expert gets an equal share.
  11. It pushes against the router’s preference. The model is being told, in effect, “spread the work, even if you would rather not”.
  12. There is a second, purely practical constraint. On real hardware every expert is given a fixed-size buffer, called its capacity.
  13. If more tokens are routed to an expert than its capacity allows, the extra tokens are dropped: they skip the feed-forward network entirely and pass through on the residual connection only.
  14. Dropped tokens are not an error. They are a deliberate trade that keeps the tensor shapes fixed so the hardware stays efficient.
  15. Finally, where the experts physically live. A 671-billion-parameter model does not fit on one accelerator, so the experts are split across machines.
  16. That is expert parallelism. Each device holds a few experts. Every token must be shipped to whichever devices hold its chosen experts, and the results shipped back.
  17. So each mixture-of-experts layer contains two network exchanges: one to dispatch tokens, one to collect answers. On a slow network this, and not the arithmetic, becomes the limit.

TECHNICAL48.14.5 the engineer’s version#

  1. History. The idea is from Robert Jacobs, Michael Jordan, Steven Nowlan and Geoffrey Hinton, “Adaptive Mixtures of Local Experts”, Neural Computation,
    1. It predates the transformer by twenty-six years.
  2. Noam Shazeer and colleagues, including Hinton and Jeff Dean, revived it for deep learning in January 2017 with “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”, reaching 137 billion parameters in an LSTM, with noisy top-k gating.
  3. GShard, from Dmitry Lepikhin and colleagues at Google in June 2020, put MoE layers inside a transformer and trained a 600-billion-parameter translation model with automatic sharding.
  4. Switch Transformer, from William Fedus, Barret Zoph and Noam Shazeer, January 2021, simplified routing to top-1 and scaled to 1.6 trillion parameters. Switch-C matched T5-Large’s compute per sequence at 2,000 times the size.
  5. GLaM, from Nan Du and colleagues at Google, December 2021: 1.2 trillion total parameters, about 97 billion active per token, 64 experts per MoE layer, top-2. It reported better zero-shot quality than GPT-3 at about one third of the training energy.
  6. ST-MoE, from Barret Zoph and colleagues, February 2022, introduced the router z-loss and produced ST-MoE-32B at 269 billion parameters.
  7. Routing mathematics. With router matrix W_r, gate g = softmax(x W_r), top-k index set T, and renormalized weights g_i / sum_{j in T} g_j:
scores   = x @ W_r                    # [n_tok, n_exp]
probs    = softmax(scores, dim=-1)
val, idx = topk(probs, k)             # weights, expert ids
val      = val / val.sum(-1, keepdim=True)
y = zeros_like(x)
for slot in range(k):
    for e in range(n_exp):
        m = idx[:, slot] == e
        y[m] += val[m, slot, None] * expert[e](x[m])
  1. The Switch Transformer auxiliary loss, which is still the default in most frameworks:
f_i = fraction of tokens in the batch routed to expert i
P_i = mean router probability assigned to expert i
L_aux = alpha * N * sum over i of ( f_i * P_i )
alpha = 0.01,  N = number of experts
minimum value 1.0, reached when all f_i = P_i = 1/N
  1. The product f_i P_i is the trick. f_i is a count and has no gradient. P_i is differentiable. Multiplying them gives the router a gradient proportional to how overloaded each expert already is.
  2. Router z-loss, from ST-MoE: the mean of the squared log-sum-exp of the router logits, with coefficient 0.001. It keeps logits small, which improves numerical stability in bfloat16 and reduces round-off in the softmax.
  3. Expert capacity = (tokens per batch / number of experts) x capacity factor. Typical capacity factors are 1.0 to 1.25 in training and higher at evaluation. Overflow tokens are dropped to the residual.
  4. Expert Choice routing, from Yanqi Zhou and colleagues at Google in 2022, inverts the problem: each expert picks its top tokens instead of each token picking its experts. Balance is then perfect by construction, but the method is not causal-safe for autoregressive decoding without care.
  5. MegaBlocks, from Trevor Gale and colleagues in 2022, reformulates MoE as block-sparse matrix multiplication so no token is ever dropped and no padding is needed. It is used in several production stacks.
  6. DeepSeek-V3, December 2024, made two notable changes. Fine-grained experts: 256 routed experts of intermediate width 2048, top-8, plus one shared expert that every token always uses, to hold common knowledge.
  7. And auxiliary-loss-free load balancing: instead of an extra loss term, a per-expert bias is added to the routing scores and adjusted between steps, raised for underloaded experts and lowered for overloaded ones. The bias affects selection only, never the gate weight, so it adds no gradient interference.
  8. Its reported training cost was about 2.788 million H800 GPU-hours on 14.8 trillion tokens. The compute rule C is about 6ND, covered in Chapter 49, uses the ACTIVE parameter count. That substitution is the entire economic case for MoE.
  9. Expert parallelism. Experts are sharded across devices and each MoE layer performs two all-to-all collectives, dispatch and combine. Communication per token per layer is about 2 x k x d_model x bytes across the fabric.
  10. DeepSeek-V3 bounded this with node-limited routing, capping each token at 4 nodes, and overlapped communication with computation using a scheme it calls DualPipe. Frameworks that implement expert parallelism include DeepSpeed-MoE from Samyam Rajbhandari and colleagues in 2022, Tutel from Microsoft, and Megatron-Core.
  11. Serving arithmetic intensity. At batch size 1 you read only k experts. At batch size 256 the tokens collectively select nearly every expert, so you read the whole expert bank anyway. MoE therefore looks cheap at small batch and behaves closer to dense at large batch.
Setting Weights read per step
Dense 13B, any batch 26 GB at bf16
MoE 47B, batch 1 about 27 GB
MoE 47B, batch 256 up to 93 GB
  1. Established fact. Sparse models reach a lower pretraining loss than dense models given equal training compute; this was measured by GShard, Switch and GLaM and has been reproduced many times. Active parameters set FLOPs and training cost. Total parameters set memory. All experts must be resident. Unbalanced routing occurs reliably without a balancing mechanism.
  2. Active research. How to compare an MoE against a dense model fairly, and what the right ratio of total to active parameters is. Whether experts specialize in any human-meaningful way; the current evidence is weak. Routing without auxiliary losses. Expert merging and pruning after training. Whether the shared-expert design is genuinely better or merely convenient.
  3. Vendor claim. Headline parameter counts. A “1 trillion parameter model” that activates 32 billion is not equivalent to a dense trillion-parameter model, and no laboratory has demonstrated that it is. A common informal rule of thumb places MoE quality near the geometric mean of active and total parameters, which for Mixtral 8x7B would be about 24.5 billion; this is a heuristic, not a measured law, and should not be quoted as one.
  4. Also vendor claim, by omission: GPT-4 has been widely reported to use a mixture-of-experts design since 2023, but OpenAI has never confirmed its architecture or parameter count. Treat any specific number you see for it as unverified.
  5. Observation. For an open-weights model, safetensors index files list every expert tensor by name, so total parameters are directly countable. Active parameters must be derived from the config: num_local_experts and num_experts_per_tok in a Hugging Face config.json give k and N. Serving frameworks such as vLLM and SGLang report expert-parallel placement and per-expert token counts, which is how you detect a collapsed router in production.

WORDS48.14.6 remember these#

  1. Mixture of experts — many small networks with a chooser in front — MoE, replacing a dense feed-forward layer with N experts and a learned router.
  2. Expert — one of the interchangeable feed-forward networks — a standard FFN or SwiGLU block, one of N per MoE layer.
  3. Router, or gate — the small network that picks experts — a d_model by N matrix followed by softmax and top-k selection.
  4. Top-k gating — take the k best experts and ignore the rest — hard sparse selection with renormalized gate weights, k typically 1, 2 or 8.
  5. Total parameters — every number in the model file — the figure that sets memory, storage and download size.
  6. Active parameters — the numbers used for one token — the figure that sets FLOPs per token, latency and training compute in C is about 6ND.
  7. Load balancing loss — an extra penalty for using experts unevenly — the Switch auxiliary loss alpha N sum f_i P_i, alpha typically 0.01.
  8. Router z-loss — a penalty that keeps router scores small — the mean squared log-sum-exp of router logits, coefficient about 0.001, for stability.
  9. Expert capacity — the fixed number of tokens an expert will accept — buffer size = tokens/experts x capacity factor; overflow tokens are dropped.
  10. Token dropping — skipping the feed-forward step for overflow tokens — those tokens pass through on the residual connection unchanged.
  11. Shared expert — an expert every token always uses — the DeepSeek design for holding knowledge common to all inputs, alongside the routed experts.
  12. Fine-grained experts — many narrow experts instead of few wide ones — same FLOPs, far more possible combinations, 64-choose-8 versus 8-choose-2.
  13. Expert parallelism — experts split across devices — sharding along the expert axis, requiring two all-to-all collectives per MoE layer.
  14. Sparse upcycling — turning a trained dense model into an MoE — copying the trained FFN into N experts and continuing training from there.

48.15 Why one architecture does so many tasks#

PLAIN48.15.1 in simple words#

  1. This is the question most people actually want answered. One machine writes code, translates Hindi, summarizes a contract and holds a conversation. Why?
  2. The short answer is that there is only one training task, and doing it well happens to require almost everything else.
  3. The task is next-token prediction. Show the model some text, hide what comes next, make it guess, and correct it. Repeat trillions of times.
  4. There is no list of skills. There is no separate summarizing module and no translation module. There is one number being pushed down: the prediction error.
  5. Now the important argument. To predict text well, you must model whatever produced that text.
  6. To finish “The capital of France is” you need a fact. Guessing needs knowledge, so knowledge lowers the error, so training installs knowledge.
  7. To finish “for i in range(10):” you need to know what an indented block looks like. So the training installs the shape of Python.
  8. To finish the last page of a detective novel you must have worked out who did it from the clues. So training rewards keeping track of a plot.
  9. To finish a polite email you need to know how polite emails go. So training installs style and register.
  10. Grammar, facts, arithmetic, tone, the structure of arguments, the layout of a table: every one of them lowers the prediction error, so every one of them gets learned.
  11. Then the second half of the answer, which is easy to miss. Text on the internet already contains every task, written out in full.
  12. An article is followed by its summary. An English sentence is followed by its French translation. A question is followed by an answer. A bug report is followed by a fix.
  13. So when the text is arranged that way, predicting the next token IS doing the task. The task was never separate. It was always just more text.
  14. That covers translation, summarizing and code. Then there is one behaviour that genuinely surprised the people who built it.
  15. In-context learning. Put three worked examples in the prompt and the model does the fourth, for a task nobody trained it on and that may not exist anywhere in its training data.
  16. Nothing inside the model changes. No number is updated. The examples sit in the input and the model behaves as if it had been taught. That was not designed. It was noticed.
  17. Then instruction following: you write “Summarize this in three bullet points” and it obeys. Here we must be honest.
  18. That is mostly not from pretraining. A raw pretrained model given that instruction is quite likely to continue with another instruction, because that is what usually follows an instruction on a web page.
  19. Obedience is added afterwards, in a separate stage, using human-written examples and human preference data. Chapter 49 covers exactly how.
  20. Finally, emergence: the claim that some abilities are absent in small models and appear suddenly in large ones. This one is genuinely disputed, and we will give both sides.

PLAIN48.15.2 a picture in your head#

  1. Imagine a person locked in the world’s largest library, playing one game, forever.
  2. The game: open any book at any page, cover the next word, guess it, then uncover it and see if you were right. Score kept. No other instructions.
  3. On day one they guess badly. Soon they learn that “the” is common, that a full stop is usually followed by a capital letter.
  4. Keep going. Now they are playing on a chemistry textbook. To guess the word after “sodium reacts violently with” they have to learn chemistry.
  5. Keep going. Now a French-English phrasebook, where every line is the same sentence twice. To win, they must learn to translate.
  6. Keep going. Now the last page of a murder mystery. To guess the name after “the killer was” they must have followed the clues and solved it.
  7. Nobody ever told them to learn chemistry, French or detective work. They learned all three because each one improved their score at one game.
  8. That is the whole explanation of why one architecture does so many things. The tasks were never taught. They were the only way to win.

Where this comparison breaks: the person in the library knows they are playing a game, wants to win, and could stop. The model has none of that. It has no goal, no awareness that it is guessing, and no sense of being right or wrong. A person who does not know says so; the model always produces a full probability distribution over every possible next token, and something is always on top. Also, the person learns by understanding. The model learns by a mechanical adjustment of billions of numbers in the direction that lowers one error term.

PLAIN48.15.3 a worked example#

  1. Here is how four different-looking jobs become the same job. Each block is just text, and the model only ever continues it.
[ translation ]
English: the sea otter is asleep
French:

[ summarizing ]
<a 900-word article>
TL;DR:

[ code ]
# return the nth Fibonacci number
def fib(n):

[ conversation ]
User: what is the boiling point of water?
Assistant:
  1. In every case the model does the same thing: predict what comes next. The task lives entirely in how the text was arranged.
  2. The “TL;DR:” trick is not made up. The GPT-2 paper of February 2019 used exactly that string to get summaries out of a model that had no summarizing training at all.
  3. Now in-context learning. This is the format from the GPT-3 paper of May 2020, which is where the phrase “few-shot” entered common use:
Translate English to French.
sea otter => loutre de mer
peppermint => menthe poivree
plush giraffe => girafe peluche
cheese =>
  1. The three examples are not training data. They are input. They are thrown away when the request ends. The model’s weights are identical before and after.
  2. It works anyway, and it works better with more examples. Measured on the TriviaQA question-answering set with the 175-billion-parameter GPT-3:
Setting Examples in prompt Accuracy
Zero-shot 0 64.3 per cent
One-shot 1 68.0 per cent
Few-shot 64 71.2 per cent
  1. Seven points of accuracy, bought with nothing but text in the prompt. No training, no weight change, no gradient.
  2. Now the honest correction about instructions. Take a base model, before any fine-tuning, and give it this:
Write a haiku about rain.
  1. A base model may well continue like this, because that is what actually follows such a line on a real web page:
Write a haiku about rain.
Write a limerick about a cat.
Write a short story about a lighthouse.
Submit your entries by Friday.
  1. It is not being unhelpful. It is predicting correctly. It has recognized a list of writing prompts and is continuing the list.
  2. Turning that into a model that answers is a separate stage of work, with its own data and its own cost, described in Chapter 49.

PLAIN48.15.4 what is really happening inside#

  1. One loss function drives everything. For each position, the model outputs a probability for every token in the vocabulary, and the loss is the negative logarithm of the probability it gave to the token that actually came next.
  2. Averaged over the sequence, that single number is the whole training signal. There is nothing else. No task labels, no reward, no human in the loop.
  3. Because the loss is per-position, one 8,192-token document gives 8,192 separate lessons in one pass. That is why the objective is so efficient.
  4. Implicit multitask learning is the mechanism behind the range. The corpus contains question-answer pairs, parallel translations, code with comments, articles with headlines, arguments with rebuttals.
  5. Each of those is a task demonstration that happens to be written down as ordinary text. The model never learns “this is translation”. It learns that after a French colon, French words are likely.
  6. The honest version: it is often said that models were never trained on these tasks. That is too strong. The tasks are present in the data in enormous quantity, just unlabelled. And since about 2023 laboratories deliberately mix instruction-shaped and textbook-shaped data into pretraining itself, so the clean story of pure raw web text is out of date.
  7. In-context learning works differently, and we do have a partial mechanical account of it.
  8. Earlier in this chapter we met induction heads: a pair of attention heads that between them find an earlier occurrence of the current token and copy what followed it.
  9. That is a complete, if crude, pattern-completion machine built out of attention. Given “A B … A” it predicts “B”. Given three translation pairs it can find the pattern and continue it.
  10. So few-shot prompting is not the model learning. It is the model reading. The examples are in the context, and attention can reach them.
  11. That is why in-context learning vanishes the moment the examples fall out of the context window. Nothing was stored. Nothing was learned in the sense of a weight changing.
  12. Instruction following is the part that is genuinely bolted on. The base model has all the capability and none of the disposition to use it on request.
  13. Fine-tuning on instruction-and-answer pairs changes which continuation is likely, not what the model knows. It is a change of habit, not of knowledge.
  14. That distinction matters, because it explains a common observation: fine tuning for helpfulness can make a model slightly worse at raw prediction while making it far more useful.

TECHNICAL48.15.5 the engineer’s version#

  1. The objective, exactly. Maximize the log-likelihood of the corpus under the chain rule of probability, one factor per token:
loss = - (1/T) * sum over t of log P(x_t | x_1 ... x_{t-1})
  1. This is cross-entropy between the model’s distribution and a one-hot target, in nats per token. Divide by ln 2 for bits per token.
  2. Shannon’s source coding theorem makes the connection to compression exact: the cross-entropy in bits per token is the size of the optimally compressed text under that model. Better prediction is literally better compression.
Method Approximate bits per char
gzip on English prose about 2.4
Shannon’s 1951 human study 0.6 to 1.3
Large model, held-out text under 1.0
  1. Claude Shannon’s “Prediction and Entropy of Printed English”, 1951, estimated English at roughly 0.6 to 1.3 bits per character by having humans guess next letters. Current models are inside that band on ordinary prose. Treat the figures above as approximate; they depend heavily on the text.
  2. Zero-shot task transfer was demonstrated by GPT-2, “Language Models are Unsupervised Multitask Learners”, from Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei and Ilya Sutskever, February 2019, at 1.5 billion parameters.
  3. Few-shot in-context learning was the headline result of GPT-3, “Language Models are Few-Shot Learners”, from Tom Brown and colleagues, May 2020, at 175 billion parameters.
  4. Instruction tuning as a named technique: FLAN, “Finetuned Language Models Are Zero-Shot Learners”, Jason Wei and colleagues at Google, September 2021; and T0, “Multitask Prompted Training Enables Zero-Shot Task Generalization”, Victor Sanh and colleagues, October 2021.
  5. Then InstructGPT, “Training language models to follow instructions with human feedback”, Long Ouyang and colleagues at OpenAI, March 2022. Its key measured result: human labellers preferred the outputs of the 1.3-billion-parameter InstructGPT model over the 175-billion-parameter GPT-3, a hundredfold size difference overturned by post-training alone.
  6. That single result is the strongest available evidence that instruction following is added, not emergent. Full detail is in Chapter 49.
  7. Now emergence, stated fairly, because this is contested.
  8. The claim. “Emergent Abilities of Large Language Models”, Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph and others, June 2022, published in TMLR. Their definition: an ability is emergent if it is not present in smaller models but is present in larger models, and cannot be predicted by extrapolating from smaller models.
  9. Reported examples include multi-digit arithmetic, word unscrambling and some of the harder tasks in BIG-bench, the 204-task benchmark released by Aarohi Srivastava and several hundred co-authors in June 2022, which itself reported “breakthrough” behaviour on some tasks.
  10. The critique. “Are Emergent Abilities of Large Language Models a Mirage?”, Rylan Schaeffer, Brando Miranda and Sanmi Koyejo of Stanford, April 2023. Their argument: nonlinear or discontinuous metrics produce apparent emergence, while linear or continuous metrics on the same models produce smooth, predictable curves.
  11. Their strongest evidence is constructive: they induced apparent emergence in simple vision models, where nobody claims it exists, purely by choosing a harsh metric. And they showed most BIG-bench emergence disappears when exact-match accuracy is replaced by token edit distance.
  12. Here is the artefact in arithmetic you can check yourself. Suppose per-token accuracy improves smoothly with scale, and the answer is 10 tokens long, and the metric is exact string match, so all 10 must be right:
Per-token accuracy 10-token exact match
0.50 0.1 per cent
0.70 2.8 per cent
0.80 10.7 per cent
0.90 34.9 per cent
0.95 59.9 per cent
0.99 90.4 per cent
  1. The left column rises steadily. The right column sits near zero, then leaps. Plot the right column against compute on a log axis and you get a cliff, and nothing whatsoever happened inside the model.
  2. The reply to the critique, given fairly. Three points. First, the critique does not show that scaling is safe or fully predictable; it shows that one class of measurement exaggerates sharpness.
  3. Second, many things people actually care about really are all-or-nothing. A program compiles or it does not. A proof is valid or it is not. If the metric that matters to the user is discontinuous, the user’s experience of a sudden capability jump is real, whatever the smooth underlying curve says.
  4. Third, at least one genuine discontinuity has been found mechanistically, not metrically. “In-context Learning and Induction Heads”, Catherine Olsson and colleagues at Anthropic, March 2022, identified a visible bump in the training loss curve where induction heads form, accompanied by a jump in in-context learning ability. That is a phase change in the model, not in the scoring.
  5. Relatedly, “Grokking”, from Alethea Power and colleagues at OpenAI in January 2022, showed genuine sudden generalization long after memorization on small algorithmic tasks. Sharp transitions do exist.
  6. A later position, from Sheng Lu and colleagues at ACL 2024, argued that many claimed emergent abilities reduce to in-context learning plus instruction tuning rather than any new latent reasoning capacity.
  7. Established fact. Next-token prediction is the sole pretraining objective. Pretraining loss falls as a smooth power law in compute, data and parameters, per Jared Kaplan and colleagues in January 2020 and Jordan Hoffmann and colleagues in March 2022. In-context learning is real and measurable. Instruction following comes overwhelmingly from post-training.
  8. Active research. Why in-context learning works at all. One line of work, including Johannes von Oswald and colleagues in 2022, argues the forward pass implements something equivalent to gradient descent on the in-context examples; this is suggestive and not settled. Whether emergent abilities are real, artefactual, or both depending on the ability, is open, and reasonable researchers disagree in print.
  9. Marketing claim. That emergence shows a model has begun to understand, reason or think. That is a philosophical claim dressed as a measurement. The “Sparks of AGI” paper from Microsoft in March 2023 was widely criticized on exactly this ground, including for evaluating an unreleased model version that others could not reproduce.
  10. Also marketing: presenting a capability as spontaneous when it was in the fine-tuning mix. Without disclosure of the post-training data, which is almost never published, that distinction cannot be checked from outside.
  11. Observation. Per-token loss and log-probabilities are the measurable quantities here, exposed as logprobs by most inference APIs and by output_scores in Hugging Face generate. Task-level claims are measured with harnesses such as EleutherAI’s lm-evaluation-harness or Stanford’s HELM, and the first thing to check in any emergence plot is which metric the y-axis uses.

WORDS48.15.6 remember these#

  1. Next-token prediction — guess the next piece of text — autoregressive language modelling with a cross-entropy loss over the vocabulary.
  2. Pretraining objective — the one thing the model is scored on — maximizing log-likelihood of the corpus under the chain rule of probability.
  3. Cross-entropy loss — how surprised the model was by the true next token — minus the log probability assigned to the target, in nats per token.
  4. Implicit multitask learning — tasks learned without being labelled as tasks — task demonstrations occurring naturally in the pretraining corpus.
  5. In-context learning — learning from examples in the prompt, with no weight change — conditioning on demonstrations within a single forward pass.
  6. Few-shot prompting — putting a handful of examples in the prompt — k-shot conditioning, popularized by the GPT-3 paper of May 2020.
  7. Zero-shot — asking with no examples at all — task specification by instruction or format alone, demonstrated by GPT-2 in February 2019.
  8. Base model — a model that has only been pretrained — a pure next-token predictor with no instruction tuning and no preference optimization.
  9. Instruction tuning — teaching the model to do what it is asked — supervised fine-tuning on instruction-and-response pairs, FLAN 2021, InstructGPT 2022.
  10. Emergent ability — a skill said to appear only above some size — an ability absent in smaller models and present in larger ones, per Wei et al. 2022, with the sharpness disputed by Schaeffer et al. 2023.
  11. Metric artefact — a jump caused by the ruler, not the thing measured — apparent discontinuity produced by exact-match or thresholded scoring.
  12. Scaling law — the predictable fall in loss as you spend more — a power law in parameters, data and compute, per Kaplan 2020 and Hoffmann 2022.

48.98 Common wrong ideas#

  1. Wrong: attention means the model is focusing, the way a person concentrates. Right: attention is a weighted average. Every token computes a similarity score against every other token, those scores go through a softmax, and the result is used to mix value vectors. Nothing is ignored: after a softmax every weight is strictly greater than zero, so a “0.001” token still contributes. There is no spotlight, no effort and no choosing. The name is a 1990s metaphor from Bahdanau, Cho and Bengio’s 2014 alignment work, and it has misled more readers than it has helped.
  2. Wrong: the transformer understands grammar, because everything it writes is grammatical. Right: grammatical output is evidence that grammatical continuations had low prediction loss, nothing more. The model holds no parse tree, no rule for subject-verb agreement and no category called “verb”. Probing studies do find syntax-correlated directions in the activations, and that is a genuine finding, but a correlated direction is not a rule. The honest test is behaviour on rare constructions, where models still fail in ways no speaker of the language would.
  3. Wrong: a bigger context window means better memory. Right: the context window is the maximum input size for one request. Nothing persists between requests, and any apparent memory is your application resending the history. Accuracy also dips for material in the middle of a long input, and NVIDIA’s RULER benchmark in 2024 found many models advertising 32,000 tokens performed well only at much shorter lengths. Longer also costs more: at 131,072 tokens the key-value cache for Llama 3.1 70B is 40 GiB, on top of the weights.
  4. Wrong: it predicts words. Right: it predicts tokens, which are subword pieces from a fixed vocabulary, typically 32,000 to 256,000 entries. “Unbelievable” may be three tokens; a leading space is usually part of the token; a number may split in places no human would choose. And it does not predict one token, it produces a probability for every token in the vocabulary at once, and then a separate sampling step picks one.
  5. Wrong: more heads means more intelligence. Right: heads split the same d_model between them. With d_model 512, going from 8 heads to 16 halves the head dimension from 64 to 32. Total parameters and total FLOPs are unchanged. The 2017 paper measured this directly and found 1 head was worse than 8, and 32 heads were also worse than 8. Head count is a shape choice with an interior optimum, not a quality dial.
  6. Wrong: emergent abilities prove something new appeared inside the model. Right: this is disputed in print. Wei and colleagues reported sharp jumps in 2022; Schaeffer, Miranda and Koyejo at Stanford argued in 2023 that the sharpness is often produced by all-or-nothing metrics, and demonstrated the effect by manufacturing fake emergence in vision models. Some genuine phase changes do exist, notably the formation of induction heads found by Anthropic in 2022. The correct position is that the metric must be checked before any claim about the model is made.
  7. Wrong: a 671-billion-parameter mixture-of-experts model is equivalent to a dense 671-billion model. Right: it does the arithmetic of a 37-billion model per token. You get the memory bill of the large model and the speed of the small one. No laboratory has shown the quality matches a dense model of the same total size, and the informal geometric-mean rule of thumb is a heuristic, not a measured law.
  8. Wrong: attention is where the model’s knowledge is stored, since attention is what makes it a transformer. Right: attention holds about 17 per cent of the parameters in Llama 3.1 8B; the feed-forward blocks hold about 70 per cent, rising to 81 per cent at 405B. Attention moves information between positions. The feed-forward blocks are where the storage happens.
  9. Wrong: the transformer processes text left to right, like a person reading. Right: during training every position is computed at once in one matrix multiplication, which is the entire reason the architecture won. The causal mask enforces the ordering of information without enforcing an ordering of computation. Only generation is sequential, and that is a property of writing text one token at a time, not of the architecture.
  10. Wrong: setting temperature to 0 makes the model correct. Right: it makes the model deterministic, which is a different thing. Greedy decoding picks the highest-probability token every time, which removes variation and often produces repetitive text. If the highest-probability continuation is wrong, temperature 0 guarantees you get the wrong answer every time instead of some of the time.

48.99 Chapter summary in 20 lines#

  1. Before 2017 sequences were handled by recurrent networks, which kept one hidden state and updated it once per word, so information from early words faded and gradients vanished or exploded across hundreds of steps.
  2. The fatal flaw was not memory but order: word ten could not be computed until word nine finished, so a chip with tens of thousands of parallel units sat idle, and no amount of hardware fixed it.
  3. Attention removes the order constraint. Every position looks at every other position at once, scores them for relevance, and takes a weighted average of what it finds. The whole sequence becomes one matrix multiplication.
  4. Each token is projected into three roles: a query saying what it is looking for, a key advertising what it offers, and a value carrying what it passes on if selected. Three learned matrices, three different jobs.
  5. The formula is softmax(QK^T / sqrt(d_k)) V. Dot products give similarity, division by sqrt(d_k) keeps the variance of the scores near 1 so the softmax does not saturate, softmax turns scores into weights that sum to 1, and the multiply by V collects the answer.
  6. Without the sqrt(d_k) scaling, scores at d_k = 64 have standard deviation around 8, the softmax becomes nearly one-hot, and gradients through it fall towards zero. The scaling is not cosmetic.
  7. A causal mask sets every score above the diagonal to minus infinity before the softmax, which makes those weights exactly zero. That is how a model is stopped from reading its own answer during training.
  8. Multi-head attention splits d_model into h independent subspaces, runs attention in each, concatenates and projects. Parameters and FLOPs are unchanged; only the shapes differ. The 2017 paper measured 8 heads as better than 1 and better than 32.
  9. Heads are not clean specialists. Some interpretable ones are real, notably induction heads that complete a repeated pattern and underlie in-context learning, but most heads resist a tidy description and many can be pruned.
  10. Attention is permutation equivariant: shuffle the input and the output shuffles identically. It has no idea what order anything is in, so position must be injected deliberately.
  11. Position schemes went sinusoidal in 2017, learned absolute in BERT and GPT, T5 relative bias in 2019, ALiBi in 2021, and rotary embedding in 2021. RoPE rotates queries and keys by an angle set by position, so scores depend on relative distance exactly, and it is the dominant choice as of 2026.
  12. RoPE’s base controls the slowest wavelength. Raising it from 10,000 to 500,000, as Llama 3.1 did in July 2024, is part of how a 131,072-token context is reached without the slowest frequency wrapping around.
  13. After attention, each position passes alone through a feed-forward network that widens by roughly four times, applies a non-linearity such as SwiGLU, and narrows back. This block holds 70 to 81 per cent of a model’s parameters and is where facts appear to be stored.
  14. Residual connections add the input back to every sub-layer’s output, giving gradients a clean path through depth and creating a residual stream that every block reads from and writes to. Normalization, LayerNorm in 2017 and RMSNorm today, keeps activation sizes stable; pre-norm replaced post-norm around 2019 and removed the need for learning-rate warmup.
  15. The assembled model is: embed token ids, add or rotate in position, run N identical blocks, normalize once more, multiply by the unembedding matrix, and read a probability over the whole vocabulary at every position.
  16. The 2017 paper, “Attention Is All You Need” by Ashish Vaswani and seven co-authors, reached 28.4 BLEU on WMT 2014 English-German with the big model, beating every previous system at one to two orders of magnitude less training compute. It never mentioned scaling laws or in-context learning.
  17. Three architectures came out of it: encoder-only for reading, such as BERT; decoder-only for writing, which is every chat model today; and the original encoder-decoder, joined by cross-attention, which survives in translation and in T5-style text-to-text framing.
  18. Generation is a loop: feed everything, get a distribution, choose one token, append, repeat. The choice is greedy, or sampled with temperature, top-k, top-p or min-p. That sampling step, not the model, is why the same prompt gives different answers.
  19. Attention costs O(n^2). At 131,072 tokens the score matrix for one head in one layer would be 32 GiB, so nobody stores it: FlashAttention computes the same result tile by tile without ever writing it out, grouped-query attention shrinks the key-value cache fourfold, and sliding windows, sparse patterns and state-space models attack the problem from other directions.
  20. Mixture of experts splits total from active parameters, letting DeepSeek-V3 hold 671 billion numbers while using 37 billion per token; and the reason one architecture does so many jobs is that it is trained on one objective, next-token prediction, which rewards grammar, facts, style and reasoning alike, with in-context learning arriving unplanned, instruction following added afterwards by fine-tuning, and emergence still argued over.