# Transformers from First Principles: From Counting Words to a Memory That Computes
Read this sentence and guess the last word.
> Flo gave Ann a pen and yes well then Ann thanked ___
A person says "Flo" without effort. The word is not in any dictionary of likely continuations of "thanked". It is in the sentence, eleven words earlier, and the person found it by asking one question: *who gave the pen?* This note is about the machines that learned to ask that question, and about what each design had to give up or invent to be able to.
The plan is to build the transformer by pressure, not by decree. We start with the simplest honest model of text, a bag of words, and ask of it what we will ask of every model in this note: **what story does it tell about how the text was generated, what does that story explain, and what does it have to remember to predict the next word?** Each model class fails at something specific and measurable. The failure forces the next design. Counting words fails because the name it needs is in this sentence, and a table of counts keeps no record of this sentence. Carrying a fixed summary fails because the summary is too small. Looking back over everything that was written, with a lookup by content, is attention. And once the lookup is read as a *memory that the forward pass writes*, a transformer becomes something more than a pattern matcher: a small machine with a memory it fills while it reads, and programs, in its weights, that decide what to write and what to read.
## The ruler: one tiny language
A discussion of language models needs a ruler, a language small enough that every probability is known exactly, yet hard in the one way that matters. Here it is. A sentence of **ToyLang** has the form
> X gave Y a T $f_1\ldots f_L$ then Y thanked X .
where X and Y are two different names out of $N$ available, T is one of 3 things (book, pen, cup), and the $L$ fillers $f_1,\dots,f_L$ are words drawn at random from a list of $F=6$ (and, so, well, now, yes, ok). Four sentences from the generator, with $N=8$ and $L=3$:
```
Flo gave Ann a pen and yes well then Ann thanked Flo .
Di gave Ben a book now so and then Ben thanked Di .
Flo gave Di a cup ok so yes then Di thanked Flo .
Cy gave Flo a pen yes now and then Flo thanked Cy .
```
The sentence makes only four random choices (the giver, the receiver, the thing, and the fillers) and then *repeats* two of them: the receiver after "then", and the giver after "thanked". Those two repeated names are the entire difficulty. They are free for anyone who can look back, and expensive for anyone who cannot. The gap length $L$ sets how far back the model must look: the second name is $L+4$ tokens behind its source, and the final name is $L+8$ behind.
**What a perfect model pays.** The generator's information per sentence is $\ln N+\ln(N-1)+\ln 3+L\ln F$ nats: one name, a different name, a thing, and $L$ fillers; every copy is free. For $N=8$ and $L=3$ that is $2.079+1.946+1.099+5.375=10.50$ nats. This is the ruler's zero: a model that has understood the generating process pays 10.50 nats per sentence and not one nat more.
**What a model that cannot look back pays.** A bag of words knows only how often each name occurs, so it pays $\ln 26=3.26$ nats on the final name instead of 0. A model that sees only the previous word pays $\ln 8=2.08$ nats. One that sees the previous two words, which include the second name, knows the answer is *some other* name and pays $\ln 7=1.95$. Only a model that sees all of the previous eleven tokens pays 0, and a table that stores every such context would need $8\cdot7\cdot3\cdot6^3=36{,}288$ of them. Move the gap to $L=10$ and the table needs more than ten billion. These numbers are exact; the first part of the note derives them.
**The question to carry through the note** is therefore: what must a model hold, and where does it keep it, so that the two copied names cost nothing? Every part answers it differently.
## The route
Five parts, each ending where the next begins. Hover a part to follow it to its limit.
The turn in the story is at C. Until then every model carries the past forward *in a summary* (a table, a state); attention keeps the past as it was written and asks it questions. That leaves one more idea, the thread of D and E: **a transformer has two memories on two clocks**.
Two memories, two clocks. Weights and context are read the same way; they differ in who writes them and for how long.
## How to read the notes
**Three colours.** One colour code runs through the text, the equations and the figures. Blue is what is observed: tokens, and the keys that are matched against. Copper is the question or the hidden quantity: the query, a latent pointer, the true process behind the data. Violet is what is stored and returned, or derived: values, and the laws we derive.
In the block diagrams the same roles hold: blue boxes are tokens, keys and the context, copper the query, violet what is stored and returned, green a result or the new state, red a limit or a cost.
Boxes marked **Check** state a numerical confirmation, and boxes marked **Question** hold things that are not understood yet.
**Derivations are folded.** Every derivation longer than a line sits in a closed block after the claim and the figure. Open it when you want the proof; skipping it costs nothing.
**One notation.**
| Symbol | Meaning |
| --- | --- |
| $w_t$, $\mathcal V$ | the token at position $t$ ($t=1,\dots,n$); the vocabulary, of size $\lvert\mathcal V\rvert$ |
| $\Pr$, $P^\star$, $p$ | probability of an event; the true generating process; a model |
| $N$, $L$, $F$ | names, gap length and filler words in ToyLang |
| $n$, $r$ | sequence length; the window of an $r$-gram model |
| $d$, $d_k$, $d_v$ | model, key and value dimensions; $d_h$ hidden units of a layer |
| $e(w)$, $x_t$ | the embedding of a token; the vector at position $t$ in the residual stream |
| $k_i$, $v_i$, $q$ | key and value of item $i$; the query. $K$, $V$ stack keys and values |
| $W_Q,W_K,W_V,W_O$ | learned projections; $W_h,W_x,W_{\text{out}}$ in the recurrent net of Part B |
| $\sigma$, $\beta$ | softmax; inverse temperature (scores are multiplied by $\beta$) |
| $a_j$, $A$ | the attention weight of item $j$; the attention matrix, whose row $t$ is how position $t$ spreads its attention |
| $S_t$ | the fast-weight state, a $d_v\times d_k$ matrix, after $t$ writes (the same object as the linear memory of Part B) |
| $\eta$, $\gamma$, $\psi$ | write strength; decay factor; kernel feature map (Part D) |
| $m$ | the number of stored pairs |
| $h_t$ | the recurrent hidden state (Part B) |
| $n_h$ | the number of attention heads |
| $\ell$ | a loss, or a count of layers |
| $\theta$ | the parameters of a model |
A few letters are reused locally inside one chapter (a Lagrange multiplier, a count, a rotation); each chapter defines them where it uses them.
The older note [Transformers as memory](transformer) goes deeper on the linear algebra of memories, with proofs in appendices. This note is the on-ramp to it: where a result there is used here, it is derived again from the start.
## A. Counting words: models that explain the data
**The question of this part.** A machine reads some text and must say what comes next. What does it have to remember about the text so far to do that well, and what is the cheapest story about how the text was produced that makes the memory unnecessary? Part A answers with the oldest tools there are: count the words, count the pairs, add a hidden variable. Each tool is a *story about how the data were generated*, and each is measured against one small language, ToyLang, where the best possible answer is known exactly. **What it assumes.** Probability (conditional probability, expectation), calculus (derivatives, a Lagrange multiplier), and nothing about language models. Nothing from the other parts is used. The notation is $w_t$ for the token at position $t$, $\mathcal V$ for the vocabulary, $P^\star$ for the process that really generated the text and $p$ for a model of it.
### A1. To explain a text is to say how it could have been generated
Here is a sentence from ToyLang, the language used as a ruler in every part of this note:
> Ann gave Ben a book and so well then Ben thanked Ann .
Cover the last name: *Ann gave Ben a book and so well then Ben thanked* ___. Any reader says "Ann" at once, and no reader needs to think about it. Nothing about the world is involved. The name was written eleven tokens earlier in this very sentence, and the rule of the language says that the one who gave is the one who is thanked. A model that has never seen the pair "Ann, Ben" before can still get it right if it carries the right thing from the start of the sentence to the end. The question that organises this whole note is what that *right thing* is, how big it is, and how a model comes to hold it.
**What a model is.** A model of text is a probability distribution over token sequences. Writing $p(w_1,\dots,w_n)$ for the probability it gives to a sequence is the same as telling a *story* about how the sequence came to be: draw the tokens like this, one after another. Different stories give different distributions, and the story decides what the model can notice. A story in which each token is drawn on its own cannot notice order. A story in which the next token depends on the previous two cannot notice a name written eleven tokens ago. This part walks through four stories of increasing richness and asks what each one can and cannot explain.
**The chain rule says that predicting the next token is the whole job.** Whatever story we tell, the probability of a whole sequence can be written as a product of next-token probabilities:
$$
p(w_1,\dots,w_n)\;=\;\prod_{t=1}^{n}\htmlClass{eq-der}{p}\big(\htmlClass{eq-obs}{w_t}\;\big|\;\htmlClass{eq-obs}{w_1,\dots,w_{t-1}}\big).
$$
This is an identity of probability, not an assumption (the derivation is below), so a model of whole sequences *is* a family of next-token predictors, one for each prefix. The models of this part differ only in how much of the prefix they look at. That is what "what does it remember" means.
Derivation: the chain rule of probability
The claim is that for any distribution over sequences, $p(w_1,\dots,w_n)=\prod_{t}p(w_t\mid w_1,\dots,w_{t-1})$, where the factor for $t=1$ is $p(w_1)$.
**The definition we start from.** The conditional probability of an event $B$ given an event $A$ with $p(A)>0$ is $p(B\mid A)=p(A\cap B)/p(A)$. Rearranged, it says that the probability of both is the probability of the first times the probability of the second given the first:
$$
p(A\cap B)=p(A)\,p(B\mid A).
$$
**The base step.** Take $A$ to be the event that the first $n-1$ tokens are $w_1,\dots,w_{n-1}$ and $B$ the event that the last token is $w_n$. Then $A\cap B$ is the event that the whole sequence is $w_1,\dots,w_n$, and
$$
p(w_1,\dots,w_n)=p(w_1,\dots,w_{n-1})\;p(w_n\mid w_1,\dots,w_{n-1}).
$$
**The repeat.** The first factor is again the probability of a sequence, now of length $n-1$. Apply the same step to it, then to the result, and so on, peeling off the last token each time:
$$
p(w_1,\dots,w_{n-1})=p(w_1,\dots,w_{n-2})\;p(w_{n-1}\mid w_1,\dots,w_{n-2}).
$$
After $n-1$ peelings only $p(w_1)$ is left, and collecting the factors gives the product.
**Check.** For $n=2$ the product is $p(w_1)\,p(w_2\mid w_1)$, which is the definition above with $A=\{w_1\}$ and $B=\{w_2\}$.
**Likelihood, nats and cross-entropy.** To compare stories we ask how much probability each gave to the text that really occurred. The *log-likelihood* of a sequence is $\ln p(w_1,\dots,w_n)$, and its negative is a cost in *nats* (natural-log units; one nat is $1/\ln 2\approx1.44$ bits). By the chain rule the cost is a sum of one term per token, $-\ln p(w_t\mid w_{Derivation: cross-entropy is never below entropy
The claim is that $H(P^\star,p)-H(P^\star)\ge0$, with equality only if $p=P^\star$ wherever $P^\star>0$. Write $P$ for $P^\star$ and sum over the sequences with $P>0$ (the others contribute nothing to either term).
**Step 1: the gap as one sum.** Subtract the two definitions:
$$
H(P,p)-H(P)=-\sum P\ln p+\sum P\ln P=-\sum P\,\ln\frac{p}{P}.
$$
**Step 2: a bound on the logarithm.** For every $u>0$, $\ln u\le u-1$, with equality only at $u=1$. (The curve $\ln u$ lies below its tangent at $u=1$, which is the line $u-1$, because $\ln$ is concave.) Take $u=p/P$ (if $p=0$ the left side below is infinite and the bound holds):
$$
-\ln\frac{p}{P}\;\ge\;-\Big(\frac pP-1\Big)=1-\frac pP .
$$
**Step 3: sum with the weights $P$.** Multiplying by $P>0$ keeps the direction of the inequality, and summing gives
$$
H(P,p)-H(P)\;\ge\;\sum P\Big(1-\frac pP\Big)=\sum P-\sum p=1-\sum p\;\ge\;0,
$$
because $p$ sums to at most 1 over the sequences where $P>0$.
**Step 4: when it is tight.** Equality needs $p/P=1$ for every sequence with $P>0$, which is $p=P$.
**Check.** If $P=(\tfrac12,\tfrac12)$ and $p=(\tfrac14,\tfrac34)$ then $H(P,p)=\tfrac12\ln4+\tfrac12\ln\tfrac43=0.837$ nats and $H(P)=\ln2=0.693$, so the gap is $0.144>0$.
**The ruler.** ToyLang has a sentence template
$$
X\ \text{gave}\ Y\ \text{a}\ T\ \;f_1\dots f_L\;\ \text{then}\ Y\ \text{thanked}\ X\ \text{.}
$$
with $N$ names (default 8; $X$ and $Y$ are two different ones), $T$ one of 3 things and $L$ fillers, each drawn from $F=6$ words. Three choices and the $L$ fillers are random: $X$ uniformly among $N$ names, $Y$ uniformly among the other $N-1$, $T$ among 3, and each filler among $F$. Everything after the fillers is determined: *then*, *thanked* and the full stop are fixed, and the second $Y$ and the last $X$ are copies. The length of a sentence is $10+L$ tokens, and $L$ is the length of the gap that a copy has to cross.
Because every sentence is equally likely, the true process can be written down completely. A sentence has probability $1/\big(N(N-1)\cdot3\cdot F^{L}\big)$, so it costs
$$
H(P^\star)\;=\;\ln N+\ln(N-1)+\ln 3+L\ln F\quad\text{nats per sentence.}
$$
At the default ruler with $L=3$ this is $2.079+1.946+1.099+5.375=10.499$ nats, which is $\ln 36{,}288$: there are 36,288 sentences. With no fillers ($L=0$) it is the first three terms alone, $5.124$ nats. At $L=10$ it is $23.04$ nats, or $1.152$ nats per token and a perplexity of $3.16$. The copies add nothing, and neither do *gave*, *a*, *then*, *thanked* and the full stop. The cost the true process pays at each position is exactly the logarithm of the number of choices left at that position, because the conditional distribution of every token given its prefix is uniform over those choices.
The true process pays only for the random choices and nothing for the copies. Each cell is a token with its cost in nats; shaded copper cells are random choices, dashed copper cells are copies (X and Y mark the same name), grey cells are fixed by the template. Drag the ring to lengthen the gap $L$: the total grows by $\ln 6=1.792$ nats per filler, and the copies stay free however far back their source is.Derivation: what ToyLang costs, position by position
The claim is that the true process pays $\ln N,\ 0,\ \ln(N-1),\ 0,\ \ln3,\ \ln F\ (L\text{ times}),\ 0,\ 0,\ 0,\ 0,\ 0$ at positions $0,1,2,\dots,9+L$, and that the sum is $\ln N+\ln(N-1)+\ln3+L\ln F$.
**Step 1: the probability of a sentence.** The generator makes $3+L$ choices in a fixed order: $X$ among $N$, then $Y$ among the $N-1$ names different from $X$, then $T$ among 3, then $L$ fillers among $F$. The number of distinct sentences is therefore $N(N-1)\cdot3\cdot F^{L}$, and each is produced with the same probability, the product of the choice probabilities $\frac1N\cdot\frac1{N-1}\cdot\frac13\cdot\frac1{F^L}$.
**Step 2: the cost of a whole sentence.** Taking minus the logarithm,
$$
-\ln P^\star(\text{sentence})=\ln N+\ln(N-1)+\ln3+L\ln F .
$$
This is the same for every sentence, so it is also the entropy.
**Step 3: the same number from the chain rule.** Apply the chain rule position by position. At position 0 nothing has been chosen, so $X$ is uniform over $N$ names and the cost is $\ln N$. At position 1 the word *gave* has probability 1. At position 2, given $X$, the name $Y$ is uniform over $N-1$ names. At position 3, *a* has probability 1. At position 4, $T$ is uniform over 3. At each filler position the filler is uniform over $F$ words whatever came before. After the last filler the next token is *then* with probability 1, then $Y$ with probability 1 (it is the name already at position 2), then *thanked*, then $X$ with probability 1, then the full stop. Every factor equal to 1 costs $-\ln1=0$.
**Step 4: the two views agree.** The costs of Step 3 add up to the expression in Step 2.
**Check.** For $N=8$, $F=6$, $L=3$: $\ln 8+\ln7+\ln3+3\ln6=10.499$, and $8\cdot7\cdot3\cdot6^3=36{,}288$ with $\ln36288=10.499$.
**What the figure does not show, and the question for the rest of the note.** The true process is a model with all the structure of the language built in, and it pays 10.499 nats for an $L=3$ sentence. A model that is *learned* from sentences must pay more, and the extra is the interesting number. Every model class below is placed on the same two axes: how many nats it pays on the two copied names and what it has to store to pay less. To keep that tractable, we look only at the two copies: the second $Y$ (position $L+6$, whose source is $L+4$ tokens back) and the last $X$ (position $L+8$, whose source is $L+8$ tokens back). The true process pays 0 for each.
**What this chapter's idea still cannot do.** It has not yet explained anything. It says what a good explanation would be (a probability distribution that comes close to $H(P^\star)$ and, on the copies, close to 0) and gives the scale to read it by. The simplest story we can tell is the one with no order at all.
### A2. A bag of words explains how often, but not where
A model that has only counted words is asked the question of A1 on a sentence with a gap of $L=10$ fillers: after *then Ben thanked*, which name? Having seen only how often each token occurs, it spreads its belief over all of them, and pays $3.689$ nats for the right answer where the true process pays 0. This chapter shows where the number comes from, why no amount of data improves it, and what the model is good for in spite of that.
**The story: a bag.** Imagine one fixed bag of tokens in which token $w$ occurs with proportion $\theta_w$, where $\theta_w\ge0$ and $\sum_w\theta_w=1$. To write a text, draw a token from the bag, write it down, put it back, and repeat $n$ times. Every draw is independent of the others. The probability of the text is
$$
p(w_1,\dots,w_n\mid\theta)\;=\;\prod_{t=1}^{n}\theta_{w_t}\;=\;\prod_{w\in\mathcal V}\theta_w^{\,c_w},
$$
where $c_w$ is the number of times $w$ occurs in the text. The second form groups equal tokens. It says the probability depends on the text *only through the counts $c_w$*.
**The estimate: count.** Given a training text with counts $c_w$ and $n=\sum_wc_w$ tokens, the $\theta$ that makes the text most probable is
$$
\hat\theta_w=\frac{c_w}{n}.
$$
The proportion in the bag is the proportion in the text. It is a short result, but it is the first of three derivations in this part that use the same move (a Lagrange multiplier on the constraint that probabilities sum to one), so the argument is worth having in full.
Derivation: the maximum-likelihood estimate of a multinomial is the relative frequency
The claim is that among all $\theta$ with $\theta_w\ge0$ and $\sum_w\theta_w=1$, the log-likelihood $\ell(\theta)=\sum_wc_w\ln\theta_w$ is largest at $\hat\theta_w=c_w/n$.
**Step 1: the objective.** The log of the likelihood above is
$$
\ell(\theta)=\sum_{w\in\mathcal V}c_w\ln\theta_w .
$$
We want its maximum under the constraint $g(\theta)=\sum_w\theta_w-1=0$.
**Step 2: attach the constraint with a multiplier.** Form the Lagrangian with a number $\mu$ to be found:
$$
\mathcal L(\theta,\mu)=\sum_wc_w\ln\theta_w-\mu\Big(\sum_w\theta_w-1\Big).
$$
At a constrained maximum the gradient of $\ell$ is parallel to the gradient of the constraint, which is what setting every partial derivative of $\mathcal L$ to zero says.
**Step 3: set the derivative in each coordinate to zero.**
$$
\frac{\partial\mathcal L}{\partial\theta_w}=\frac{c_w}{\theta_w}-\mu=0\quad\Longrightarrow\quad\theta_w=\frac{c_w}{\mu}.
$$
**Step 4: find the multiplier from the constraint.** Sum the result over $w$ and use $\sum_w\theta_w=1$:
$$
1=\sum_w\frac{c_w}{\mu}=\frac{n}{\mu}\quad\Longrightarrow\quad\mu=n,
$$
and so $\theta_w=c_w/n$.
**Step 5: it is a maximum, and it is the only one (for positive counts).** The function $\ell$ is a sum of concave functions $c_w\ln\theta_w$, so it is concave, and the constraint set is convex. A stationary point of a concave function on a convex set is the global maximum, and the stationary point found above is unique.
**Step 6: tokens never seen.** The stationary-point argument divides by $\theta_w$, so Steps 3 to 5 are stated for $c_w>0$. If some $c_w=0$ the optimum sits on the boundary of the simplex, at $\theta_w=0$, where setting a derivative to zero is no longer the right test, and the log inequality of A1 does the job instead. Sum only over the tokens with $c_w>0$ (the others add $0$ to $\ell$ whatever $\theta_w$ is, with $0\ln0=0$). For every $\theta$ in the simplex, with $\hat\theta_w=c_w/n$ and $c_w/\hat\theta_w=n$,
$$
\ell(\theta)-\ell(\hat\theta)=\sum_{c_w>0}c_w\ln\frac{\theta_w}{\hat\theta_w}\;\le\;\sum_{c_w>0}c_w\Big(\frac{\theta_w}{\hat\theta_w}-1\Big)=n\Big(\sum_{c_w>0}\theta_w-1\Big)\;\le\;0,
$$
by $\ln u\le u-1$. So $\hat\theta$ is the maximum with or without zero counts. Its price is that a token never seen gets $\hat\theta_w=0$: the model then says that $w$ is impossible, and any text containing $w$ gets log-likelihood $-\infty$. This is the zero-count problem, and it returns in A3.
**Check.** Two tokens with counts $(3,1)$ give $\hat\theta=(0.75,0.25)$, and $\ell=3\ln0.75+\ln0.25=-2.249$; moving to $(0.7,0.3)$ gives $3\ln0.7+\ln0.3=-2.274$, which is lower.
**A bag is a model of a multiset.** The count form of the probability has a consequence worth stating twice. Shuffle the tokens of a text in any way: the counts do not change, so *the probability does not change*. The bag-of-words model gives the same probability to *Ann gave Ben a book ... Ben thanked Ann* and to every reordering of it, including the ones that are no longer sentences of the language. The data it can explain are sets of tokens with multiplicities, not sequences. Throughout this part "bag of words" means this *static unigram*, one fixed set of frequencies; it does not mean a model that computes features, such as occurrence counts, from the prefix in front of the blank (the end of A5 comes back to that difference). This returns in Part C, where a transformer layer with no position signal turns out to share this property, with a difference that matters: it can look things up.
The figure makes it concrete. A ToyLang sentence with three fillers has 13 tokens. Its bag has $13!/(2!\,2!)=1{,}556{,}755{,}200$ distinct orderings when the three fillers are different words, and exactly 12 of them are ToyLang sentences: the six orderings of the fillers, each with the roles of $X$ and $Y$ either way round.
Shuffling a sentence changes nothing for a bag of words and everything for the language. Filled cells have moved after the chosen number of random swaps ($L=3$). The bag's bar never moves; the true process pays 10.50 nats while the sequence is still a sentence and finds it impossible (dashed ghost) once it is not. Drag the handle from the opening 4 swaps.
**What the bag pays on the copies.** In ToyLang a sentence of $n=10+L$ tokens contains 4 name tokens ($X$ twice, $Y$ twice), spread over $N$ names. In the long run a given name makes up a fraction $4/(nN)$ of all tokens, and that is the probability the unigram model, fitted on the whole language, gives to it. A copied name therefore costs
$$
-\ln\frac{4}{nN}=\ln\frac{(10+L)N}{4}\ \text{nats},
$$
which is $3.258$ for $L=3$ and $3.689$ for $L=10$, against 0 for the true process. That is worse than a uniform guess over the 22 token types of ToyLang, which would pay $\ln22=3.091$: each name is individually rare, and the bag knows nothing that says a name is due. It grows with $L$, because a longer sentence dilutes the names among more fillers. Both copies together cost $7.378$ nats at $L=10$. The whole sentence costs the model $58.68$ nats ($20$ tokens at $2.934$ nats each) against $23.04$, a gap of $35.64$ nats, because the bag also pays about $\ln n$ for each of *gave*, *a*, *then*, *thanked* and the full stop, which the template fixes. The bag is not just missing the copy: it is blind to order everywhere. Nor does more data help, since the estimate $c_w/n$ converges to exactly the frequencies used here.
Derivation: what the fitted bag charges a ToyLang sentence
The claim is that the unigram model fitted on ToyLang charges every sentence of length $n=10+L$ the same number of nats, namely $n$ times the entropy of the token frequencies, and that this is $\ln\big((10+L)N/4\big)$ per copied name.
**Step 1: the long-run frequencies.** Every sentence contains the same number of tokens of each kind: 4 name tokens, 1 thing token, $L$ filler tokens and one each of *gave*, *a*, *then*, *thanked* and the full stop. By symmetry the $N$ names are equally frequent, the 3 things are equally frequent, and the $F$ fillers are equally frequent. So the proportions are
$$
\theta_{\text{name}}=\frac{4}{nN},\quad
\theta_{\text{thing}}=\frac{1}{3n},\quad
\theta_{\text{filler}}=\frac{L}{Fn},\quad
\theta_{\text{gave}}=\theta_{\text{a}}=\theta_{\text{then}}=\theta_{\text{thanked}}=\theta_{.}=\frac1n .
$$
They add up: $4/n+1/n+L/n+5/n=(10+L)/n=1$.
**Step 2: the cost of any sentence.** All sentences have the same composition, so they all cost the same: each token kind contributes its count times $-\ln\theta$,
$$
4\ln\frac{nN}{4}\;+\;\ln(3n)\;+\;L\ln\frac{Fn}{L}\;+\;5\ln n .
$$
This is $n$ times the entropy of the frequencies, $n\,H(\theta)=-n\sum_w\theta_w\ln\theta_w$, because the count of a token kind is $n$ times its proportion.
**Step 3: the cost of one copied name.** A copied name is a name token, so it costs $-\ln\theta_{\text{name}}=\ln\big(nN/4\big)=\ln\big((10+L)N/4\big)$.
**Check.** For $N=8$, $L=3$, $n=13$: $4\ln26+\ln39+3\ln26+5\ln13=13.032+3.664+9.774+12.825=39.29$ nats, and the entropy of the frequencies is $3.023$ nats per token, so $13\times3.023=39.29$.
**What a bag is good for: naive Bayes.** A bag of words is a poor model of a sentence and a good model of a *topic*. Suppose each document belongs to a class $z$ (sports, cooking, a customer complaint), and each class has its own bag. The story is: choose the class with probability $\pi_z$, then draw every token of the document from that class's bag $\theta_z$. The joint probability of a class and a document is
$$
p(\htmlClass{eq-hid}{z},w_1,\dots,w_n)\;=\;\pi_{\htmlClass{eq-hid}{z}}\prod_{t=1}^{n}\theta_{\htmlClass{eq-hid}{z},\,w_t}.
$$
The class is hidden in the story and has to be inferred from the words. Bayes' rule does it, and it gives a classifier that is *generative*: it explains the words and gets the class as a by-product. The posterior probability of the class is
$$
p(z\mid w_{1:n})\;=\;\frac{\pi_z\prod_w\theta_{zw}^{\,c_w}}{\sum_{z'}\pi_{z'}\prod_w\theta_{z'w}^{\,c_w}},
\qquad
\ln p(z\mid w_{1:n})=\text{const}+\ln\pi_z+\sum_wc_w\ln\theta_{zw},
$$
so, up to a constant that does not depend on $z$, the log-posterior is a *linear* function of the count vector. The model is called naive because of the assumption that, once the class is known, the tokens are independent. The training is the three-line argument above applied once per class: the log-likelihood of the training set separates into one term for the class proportions and one multinomial term per class, so $\hat\pi_z$ is the fraction of training documents in class $z$ and $\hat\theta_z$ is the relative frequency of tokens within the documents of class $z$.
Derivation: Bayes' rule for classes, and the training estimates
The claim is that $p(z\mid w_{1:n})$ is the expression above, and that the maximum-likelihood estimates are the class fractions and the within-class relative frequencies.
**Step 1: the posterior from the definition of conditional probability.** The joint probability of the class and the document is $p(z,w_{1:n})$. The marginal probability of the document adds this up over the classes, $p(w_{1:n})=\sum_{z'}p(z',w_{1:n})$. By the definition of conditional probability,
$$
p(z\mid w_{1:n})=\frac{p(z,w_{1:n})}{p(w_{1:n})}=\frac{p(z,w_{1:n})}{\sum_{z'}p(z',w_{1:n})}.
$$
**Step 2: insert the story.** Replace the joint probability by $\pi_z\prod_t\theta_{z,w_t}=\pi_z\prod_w\theta_{zw}^{c_w}$. This gives the fraction above.
**Step 3: the log form.** Take logarithms of the numerator. The denominator is the same for every $z$, so it is the constant:
$$
\ln p(z\mid w_{1:n})=\ln\pi_z+\sum_wc_w\ln\theta_{zw}-\ln\sum_{z'}\pi_{z'}\prod_w\theta_{z'w}^{c_w}.
$$
**Step 4: training.** Given labelled documents $(z_i,d_i)$, the log-likelihood of the whole set is
$$
\sum_i\ln\pi_{z_i}+\sum_i\sum_wc_w(d_i)\ln\theta_{z_iw}
=\sum_z n_z\ln\pi_z+\sum_z\sum_wC_{zw}\ln\theta_{zw},
$$
where $n_z$ is the number of documents in class $z$ and $C_{zw}$ is the total count of $w$ in the documents of class $z$. The parameters $\pi$ and $\theta_z$ for different $z$ appear in separate terms and are tied only by their own sum-to-one constraints. Each term is a multinomial log-likelihood, so the Lagrange argument of the previous derivation applies to each: $\hat\pi_z=n_z/\sum_{z'}n_{z'}$ and $\hat\theta_{zw}=C_{zw}/\sum_{w'}C_{zw'}$.
**Check.** Three sports documents and three cooking documents give $\hat\pi=(\tfrac12,\tfrac12)$. In the worked example below, the 13 sports tokens contain *goal* 3 times, so $\hat\theta_{\text{sports},\text{goal}}=3/13$.
**The Bernoulli alternative.** A second story for the same job writes a document as a list of yes-or-no decisions, one per vocabulary word: does this word occur in the document, with probability $\theta_{zw}$? The estimate of $\theta_{zw}$ is then the fraction of class-$z$ documents that contain $w$, and every word that is absent counts as evidence through the factor $1-\theta_{zw}$. McCallum and Nigam (1998) compared these two event models for text classification, the multi-variate Bernoulli (presence) and the multinomial (counts); an index record summarises the result as the multinomial generally better and the Bernoulli better at small vocabularies, and the paper itself was not read. The contrast is the point for us: the Bernoulli model also loses all order.
Worked example: three sports documents and three cooking documents
The training set has six short documents with a vocabulary of $|\mathcal V|=10$ words. Sports: *goal match team win goal*; *team coach match goal*; *win match team team*. Cooking: *bake oven flour sugar bake*; *sugar flour butter bake*; *oven butter bake sugar*. Each class has 13 tokens, and the class prior is $\tfrac12$ each. The test document is *coach bake bake bake*.
Without smoothing, *coach* never occurs in cooking, so $\hat\theta_{\text{cooking},\text{coach}}=0$, and *bake* never occurs in sports, so $\hat\theta_{\text{sports},\text{bake}}=0$. Both class likelihoods are exactly zero and the posterior is $0/0$: the model cannot answer, although three of the four words point to cooking.
With one pseudo-count added to every word and class, $\hat\theta_{zw}=(C_{zw}+1)/(13+10)$:
$$
\text{sports: }\ \tfrac12\cdot\tfrac{2}{23}\Big(\tfrac{1}{23}\Big)^3,\qquad
\text{cooking: }\ \tfrac12\cdot\tfrac{1}{23}\Big(\tfrac{5}{23}\Big)^3 .
$$
The ratio cooking to sports is $5^3/2=62.5$, so $p(\text{cooking}\mid\text{document})=62.5/63.5=0.984$. The Bernoulli model with the same pseudo-counts gives $0.80$ for cooking: it gets there more slowly, because each of the eight vocabulary words absent from the document also counts as evidence.
**A task a bag cannot do, in ToyLang.** Let the class be the *direction* of the gift: class A when the giver's name comes earlier in the list of names than the receiver's, class B when it comes later. A person tells the classes apart by looking at which name comes first. A bag cannot. Every sentence has each of its two names twice, and the class changes nothing about which pairs of names are possible except which one is the giver. Counting over all ordered pairs, a given name occurs the same number of times in class A as in class B (for $N=8$ it is $2\times7=14$ tokens in each), so the two class bags are identical, the likelihood ratio is 1, and the posterior equals the prior, $\tfrac12$, for every sentence. Fitting the model on an enumerated ToyLang gives exactly $\tfrac12$ on every sentence.
Derivation: the two direction classes have identical bags
The claim is that in ToyLang with $N$ names, the unigram model fitted to class A (giver earlier in the list) and to class B (giver later) are the same, so naive Bayes cannot tell the classes apart.
**Step 1: what a class contains.** Number the names $1,\dots,N$. Class A consists of the sentences with $xy$. Within a class all pairs are equally likely, all things and fillers are chosen the same way, and the template is the same. So the proportion of every token kind other than names is the same in the two classes. It remains to compare the names.
**Step 2: count name tokens.** A sentence with the ordered pair $(x,y)$ contains $x$ twice and $y$ twice. Fix a name $i$. In class A, the pairs that contain $i$ are the $N-1$ pairs $\{i,j\}$, each of which appears in class A once (as $(i,j)$ if $i
**What the bag generates.** A story can be run forwards, and what it produces shows what it understood. Drawing 13 tokens from the bag fitted to ToyLang with $L=3$ gives, for one seed,
> now then Flo well Cy so yes a thanked Flo then ok now
Each type of token appears with its long-run frequency: the names make up 4 tokens in 13 together, the fillers 3 in 13, and each template word 1 in 13. But there is no *Ann gave Ben*, no pairing of *thanked* with a name that has been seen, and no sentence. It is word salad, and that is a literal reading of the sentence "the probability does not depend on the order".
**What this chapter's model still cannot do.** It explains how often a token occurs and, with a class variable, which kind of document it came from. It does not know that *thanked* follows a name, that *then* comes after the fillers, or that the name after *thanked* has been mentioned before. It pays 3.689 nats on each copied name at $L=10$, and because the estimate converges to the true long-run frequencies, collecting more text cannot reduce that number. The next story gives the model a short memory.
### A3. A window of $r$ tokens buys order, and a table that grows with the window
Take the same sentence, now with an $L=3$ gap, and give the model the two tokens before the blank: *Ann gave Ben a book and so well then Ben thanked* ___. The model sees *Ben thanked*. It can tell that the answer is a name and not *Ben*, so seven of the eight names are open: $\ln7=1.946$ nats. To say *Ann* with certainty it would have to see the start of the sentence, eleven tokens back. This chapter makes that trade exact: a window of $r$ tokens buys order, and the price is a table whose size grows by a factor of $F$ with every filler it must see across.
**The story: a Markov chain of order $r$.** Keep the chain rule of A1 and make one assumption: the next token depends on the previous $r$ tokens and on nothing earlier,
$$
p(w_t\mid w_1,\dots,w_{t-1})\;=\;p_r\big(w_t\mid w_{t-r},\dots,w_{t-1}\big).
$$
With $r=0$ this is the bag of A2. With $r=1$ it is a *bigram* model (a pair of consecutive tokens), with $r=2$ a trigram model, and in general an $(r{+}1)$-gram model. The conditional distributions $p_r(\cdot\mid\text{context})$ are the parameters: one multinomial for every context, that is one row of a table for every distinct window of $r$ tokens. (A window that reaches before the start of a sentence includes a start marker.) Such a model can notice order, because *Ben thanked* and *thanked Ben* are different contexts.
**The estimate: count, again.** The log-likelihood of a training text splits into one multinomial term per context, so the Lagrange argument of A2 applies row by row, and
$$
\hat p_r(w\mid u)\;=\;\frac{c(u,w)}{c(u)},
$$
where $c(u,w)$ counts how often the window $u$ was followed by $w$ and $c(u)=\sum_wc(u,w)$.
Derivation: the maximum-likelihood n-gram is a table of relative frequencies
The claim is that the order-$r$ Markov model that maximises the likelihood of a training text has $p_r(w\mid u)=c(u,w)/c(u)$.
**Step 1: write the log-likelihood with the context made explicit.** By the chain rule and the order-$r$ assumption, for a text $w_1,\dots,w_n$,
$$
\ell(p_r)=\sum_{t}\ln p_r\big(w_t\mid u_t\big),\qquad u_t=(w_{t-r},\dots,w_{t-1}).
$$
**Step 2: group the terms by context and token.** Each pair $(u,w)$ appears $c(u,w)$ times, so
$$
\ell(p_r)=\sum_{u}\sum_{w}c(u,w)\,\ln p_r(w\mid u).
$$
**Step 3: the contexts decouple.** The parameters $p_r(\cdot\mid u)$ for different $u$ are free of each other, each subject only to $\sum_wp_r(w\mid u)=1$. The sum over $u$ is therefore maximised by maximising each inner sum alone.
**Step 4: each inner sum is the multinomial problem of A2.** For a fixed $u$ the inner sum is $\sum_wc(u,w)\ln p_r(w\mid u)$ with the counts $c(u,w)$ in the role of $c_w$. The Lagrange argument gives $p_r(w\mid u)=c(u,w)/\sum_{w'}c(u,w')=c(u,w)/c(u)$.
**Check.** In the tiny corpus below, the context *the* occurs 6 times followed by *cat* twice, so $\hat p_r(\text{cat}\mid\text{the})=2/6$.
**What ToyLang says about the window.** Let $r$ be the window and let the model have unlimited data, so that every conditional distribution is the exact long-run one. What does it pay on the two copies? For the final name $X$ (source $L+8$ tokens back):
$$
\text{cost}_X(r)=\begin{cases}\ln N & r=1,\\ \ln(N-1) & 2\le rDerivation: the staircase of the n-gram cost on ToyLang
The claim is the staircase above. A model with unlimited data and window $r$ predicts the next token with the exact conditional distribution of the next token given the last $r$ tokens, so its cost is the conditional entropy of the token given that window. We compute it for the final name $X$, then for the second $Y$.
**Step 1: what the window can contain for the final name.** The final $X$ is at position $L+8$, and the window holds positions $L+8-r$ to $L+7$. Reading back from the end: position $L+7$ is *thanked* ($r\ge1$), position $L+6$ is the second $Y$ ($r\ge2$), and the earlier positions are *then*, the fillers, $T$, *a*, then the first $Y$, *gave*, and finally the first $X$ at position 0 ($r\ge L+8$).
**Step 2: $r=1$.** The window is the single token *thanked*. In ToyLang the token *thanked* is always followed by the final name, and that name is $X$, which is uniform over all $N$ names marginally. The cost is $\ln N$.
**Step 3: $2\le rThe cost of a copy falls in two steps, and the second step is $L+8$ tokens away. Cost in nats of a copied name for an $r$-token model with unlimited data, against the window $r$: solid violet is the final name $X$, dashed is $Y$, grey are other gap lengths. The small panel counts the contexts a table needs to reach 0 on $X$ (log scale). Drag the handle to change $L$: the steps move right one position per filler and the table grows by a factor of 6.
**A hidden state that moves: the hidden Markov model.** A different use of a hidden variable lets it change as the text goes. The story is that a hidden state $s_t$ evolves as a Markov chain, $\Pr(s_t=j\mid s_{t-1}=i)=T_{ij}$, starting from $\pi$, and at every step the current state emits a token, $\Pr(w_t=w\mid s_t=j)=o_j(w)$. A state is a kind of mental situation: "just read a name", "in the gap", "about to thank". The joint probability of a state path and a text is
$$
p(s_{1:n},w_{1:n})\;=\;\pi_{s_1}\,o_{s_1}(w_1)\prod_{t=2}^{n}T_{s_{t-1}s_t}\,o_{s_t}(w_t),
$$
and the probability of the text alone adds this up over every state path. With $J$ states that is $J^n$ paths, far too many to list. The *forward algorithm* does it in $nJ^2$ operations by carrying one number per state, the probability of the text so far *and* the state now:
$$
\varrho_1(j)=\pi_j\,o_j(w_1),\qquad
\varrho_t(j)=\Big[\sum_i\varrho_{t-1}(i)\,T_{ij}\Big]\,o_j(w_t),\qquad
p(w_{1:n})=\sum_j\varrho_n(j).
$$
The vector $\varrho_t$ is the whole memory the model needs about the text so far, and its length is $J$, whatever $t$ is. That is the idea we want to keep in mind.
Derivation: the forward algorithm
The claim is that $\varrho_t(j)=\Pr(w_{1:t},s_t=j)$ satisfies the recursion above, so that $p(w_{1:n})=\sum_j\varrho_n(j)$ is computed in $nJ^2$ steps.
**Step 1: the definition.** Let $\varrho_t(j)$ be the joint probability that the first $t$ tokens are $w_1,\dots,w_t$ and that the hidden state at time $t$ is $j$. Summing over $j$ at $t=n$ removes the condition on the last state and gives $p(w_{1:n})$.
**Step 2: the first step.** At $t=1$ the state is $j$ with probability $\pi_j$ and it then emits $w_1$ with probability $o_j(w_1)$. So $\varrho_1(j)=\pi_j\,o_j(w_1)$.
**Step 3: split a later step by the previous state.** For $t\ge2$, the event "tokens $w_{1:t}$ and state $j$ now" is the union over the previous state $i$ of the events "tokens $w_{1:t-1}$, state $i$ at $t-1$, then state $j$ at $t$, then token $w_t$". These events are disjoint for different $i$, so their probabilities add:
$$
\varrho_t(j)=\sum_i\Pr(w_{1:t-1},s_{t-1}=i,\,s_t=j,\,w_t).
$$
**Step 4: factor each term with the chain rule and the Markov structure.** By the chain rule, the term is $\Pr(w_{1:t-1},s_{t-1}=i)\cdot\Pr(s_t=j\mid w_{1:t-1},s_{t-1}=i)\cdot\Pr(w_t\mid s_t=j,\dots)$. In the story the next state depends only on the previous state, and the token depends only on the current state, so the last two factors are $T_{ij}$ and $o_j(w_t)$. The first factor is $\varrho_{t-1}(i)$. Hence
$$
\varrho_t(j)=\sum_i\varrho_{t-1}(i)\,T_{ij}\,o_j(w_t)=\Big[\sum_i\varrho_{t-1}(i)\,T_{ij}\Big]o_j(w_t).
$$
**Step 5: the cost.** Each of the $J$ entries of $\varrho_t$ takes a sum of $J$ terms, so a step costs $J^2$ and the whole text $nJ^2$.
**Step 6: prediction.** The probability that the next token is $w$, given the text so far, is $\sum_j\big[\sum_i\varrho_t(i)T_{ij}\big]o_j(w)\big/\sum_i\varrho_t(i)$.
**Check.** Take two states, $\pi=(0.6,0.4)$, rows of $T$ equal to $(0.7,0.3)$ and $(0.4,0.6)$, and token probabilities $o_0=(0.9,0.1)$, $o_1=(0.2,0.8)$ over the tokens $(a,b)$. For the text *a b a* the recursion gives $\varrho_1=(0.540,0.080)$, $\varrho_2=(0.041,0.168)$, $\varrho_3=(0.0863,0.0226)$ and so $p=0.10893$. Adding up all $2^3=8$ state paths one by one gives the same $0.10893$.
**A hidden state that remembers the binding.** An HMM can hold exactly what ToyLang needs. Let a state be a triple: the ordered pair of names $(x,y)$ the sentence is about, and the *phase*, the position inside the sentence. The phase advances by one at every step. A state emits $x$ or $y$ with probability 1 when the phase says so (positions $0$ and $L+8$ for $x$, $2$ and $L+6$ for $y$), the template word when the phase says so, and a uniform thing or filler otherwise. The start state is uniform over the pairs. The forward algorithm on this machine gives the sentence the probability $1/\big(N(N-1)\cdot3\cdot F^{L}\big)$, the same as the generator: there is exactly one state path for each sentence. After the prefix up to *thanked* only one state of positive probability is alive, and it emits the right $X$ with probability 1. It pays $0$ on both copies at every $L$.
The memory it needs is small. The number of states is $N(N-1)(10+L)$, which is $1{,}120$ at $N=8$, $L=10$, against $10{,}158{,}317{,}568$ table rows, a ratio of $9{,}069{,}926$ to one. The part that carries the binding across the gap is only the pair, and the next paragraph says how few situations that can be.
**How few states can carry the binding?** Count the situations a predictor must tell apart at one moment, just after *then*, when the next token is the second $Y$. Suppose the predictor keeps a state, updates it deterministically from each new token, and predicts from the state alone. Take two prefixes with different bindings $(x,y)\ne(x',y')$. If $y\ne y'$, the right next token differs; a cost of 0 means probability exactly 1 on the right name, and because the probabilities of one distribution sum to 1 the other name then gets probability 0, so one state cannot serve both. If $y=y'$ but $x\ne x'$, the next token agrees but the final name does not: a shared state would receive the same tokens *y thanked*, be updated to the same state, and predict the same final name, which cannot be both $x$ and $x'$. So all $N(N-1)$ bindings need different states: at least $56$ at $N=8$, that is $\ln56=4.025=\ln8+\ln7$ nats, the cost of choosing $X$ and $Y$ at the start, at any $L$.
Three scope notes. The bound counts *distinguishable states at one moment*; it is not the dimension of a real vector, which has uncountably many states and is limited by its precision instead (B1 and B3 come back to this). It assumes that the state update is deterministic. And it is needed for the two copies *jointly*: the second $Y$ alone can be predicted from $N$ states, one per value of $Y$. B1 derives the general version of the bound with the exact count.
This is the first model of the note that solves the problem, and it solves it by being the right size: a memory that holds $\ln56$ nats and not $F^L$ rows. But look at what was done. We, the designers, put the pair $(x,y)$ and the phase into the state by hand. Nothing in the HMM story says how a machine would find that structure from sentences, and the structure itself does not share anything between pairs: *Ann, Ben* and *Cy, Di* are different states with separate transition rules, and with 48 names there are $2{,}256$ pairs. The binding is held by *which slot is active*, not by *what was written in this sentence*.
**Sharing strength with embeddings.** A third answer replaces the table by a function. Bengio, Ducharme, Vincent and Jauvin (2003) learn a distributed representation for each word together with the probability function of the next word, and this allows the model to generalise to sequences it has not seen if they are made of words similar to words in sequences it has seen; the paper reports that it beats $n$-gram models on two corpora (abstract only). The recipe is: give every token $w$ a vector $e(w)\in\mathbb R^{d}$, concatenate the vectors of the last $r$ tokens ($rd$ numbers), pass them through a hidden layer of $d_h$ units, and read off a probability for each of the $|\mathcal V|$ tokens with a softmax. The number of parameters is
$$
|\mathcal V|\,d\;+\;(rd+1)\,d_h\;+\;(d_h+1)\,|\mathcal V|,
$$
which grows *linearly* with the window $r$, not exponentially. For ToyLang with $|\mathcal V|=22$ token types, $d=8$, $d_h=32$ and a window of $r=18$ (enough for $L=10$), it is $5{,}542$ parameters, against the $10^{10}$ rows of the table. The reason for the generalisation is that the same function is applied to every context: if the embeddings of *Ben* and *Cy* are close, what the model learned about *Ben thanked* carries over to *Cy thanked* with no extra data.
Two things are not claimed. The parameter count does not prove that a network of this size *can compute* the copy for every gap; it only says its size does not explode. And the window is still a window: the final name sits $L+8$ tokens back, so the reach is capped at $r$, and when $L$ varies from sentence to sentence the source sits at a different offset each time. One way to follow it is to dedicate hidden units to each possible offset, which makes the cost grow with the longest gap; it is not claimed to be the only way.
**What is still missing: a rule, not a table.** Look again at what the copy is. The rule *the last name is the first name* does not mention any name, and it does not depend on the fillers. A table that stores one row per window stores that one fact $|\mathcal U|$ times, once for each window, because every window looks like a new situation. A learner that has the *hypothesis* "the final token equals the token $d$ places back" for some $d\in\{1,\dots,L+8\}$ has only $L+8$ candidates to choose from. A single training sentence eliminates every wrong candidate, since the only position that holds $X$ is position 0: the second $Y$, the first $Y$, and every filler or template word differ from $X$. At $L=10$ there are $18$ candidates and one sentence settles it, against $10^{10}$ rows and about $2.4\times10^{11}$ sentences for the table.
A table counts every filler combination, a rule counts only where to look. Log-scale bars: the rows an $r$-token table needs to reach cost 0 on the final name, the states of the hand-built machine, the positions a copy rule could point at. Drag the handle to change $L$: the top bar climbs by a factor of 6 per step, the other two barely move.Derivation: one sentence identifies the copy rule
The claim is that if the hypotheses are "the final name equals the token $d$ places before it", for $d=1,\dots,L+8$, then a single sentence of ToyLang is consistent only with $d=L+8$.
**Step 1: where the final name is.** The final $X$ is at position $L+8$. The hypothesis with offset $d$ says that the token at position $L+8-d$ equals it.
**Step 2: which positions hold the name $X$.** Only position 0 holds $X$ before the final name: the second $Y$ at $L+6$ and the first $Y$ at position 2 are names different from $X$, and every other token (template words, a thing, a filler) is not a name at all.
**Step 3: so exactly one offset survives.** The position $L+8-d$ equals 0 only for $d=L+8$. Every other offset predicts a token that is not $X$, and is refuted by this one sentence.
**Check.** On all sentences of three small languages (up to $N=4$, $F=2$, $L=2$, and $N=3$, $L=4$) exactly one offset survives, namely $L+8$, for every sentence. In a drawn sentence for each $L$ from 0 to 30 and $N=2$, $8$, $48$ the surviving set is $\{L+8\}$.
**The rule class is a design choice too.** "The final name equals the token $d$ places back" was chosen for a fixed $L$: the final name sits $L+8$ back, so $d=11$ at $L=3$ and $d=13$ at $L=5$, and no single offset fits both. A learner with fixed offsets needs one hypothesis per length, and a hand-chosen rule class commits the sin of the hand-built HMM: the designer supplied the structure. The rule that survives any $L$ is stated by *content*, "the token that played the role of $X$", not by distance.
**Two kinds of information.** Some of what the model must know is the same in every sentence: that *thanked* is followed by a name, that the last name is the first, what the template looks like. It is learned once from many sentences and lives in the weights: **in-weights information**. Some is written in *this* sentence and gone after it: who $X$ and $Y$ are. It cannot be learned in advance, so it must be read from the text and used until the sentence ends: **in-context information**. In ToyLang the in-context information that must cross the gap is exactly the binding, $\ln\big(N(N-1)\big)=4.025$ nats at the default ruler. A table stores in-weights information for each possible in-context information, which is why it is so large; the right design stores the rule in the weights, keeps only the binding where it can be written during the sentence, and reads it when needed. Part B asks how big that place must be if its size is fixed, and Parts C and D what happens when it grows with the text.
**What this chapter's models still cannot do.**
Three ways to share strength, one thing missing. Each model shares something and stops short of reading the binding by content.
### A5. What a better model must do
Each story in this part was a reading of the same question: what must the model remember to predict the next token? Reading back over the four chapters, four requirements stand out. Each is a statement about ToyLang that a model can pass or fail, and for each there is a number.
- **R1. Order matters.** The probability of a text changes when the tokens are reordered. The bag fails it, by construction: all 1,556,755,200 orderings of the example in A2 have the same probability.
- **R2. Reach back far.** The cost on the copied names does not grow with the gap $L$, and the size of the model does not grow like $F^L$. The $r$-token table meets the first half only if $r\ge L+8$, and fails the second by a factor of $F$ per filler.
- **R3. Look up a binding written in this very text.** The model uses information that was written in the current sentence and was never in its training data, such as names it has never seen. The table, the mixture and the HMM do it only by having a slot for every binding in advance: a row, a component, a state.
- **R4. Train in parallel.** The losses of all positions of a text can be computed at once from the text, with no step waiting for the previous one. Counting is parallel. The forward algorithm of the HMM, as written, is a loop over $t$ in which $\varrho_t$ needs $\varrho_{t-1}$. (A product of matrices is associative, so this is a statement about the algorithm as written and not about every HMM; Part B returns to it for states that are updated through a nonlinearity, where no such regrouping is available.)
The table below puts every model of this part against the ruler at $L=10$ and $N=8$, with the two copies together ("copy nats"), and "stores" counts the table rows, bags or states needed to reach that cost (at most, for the short windows). "Yes" means the requirement is met in the sense above, "no" that it is not, and "part" that it is met under a condition stated in the chapter that introduced the model.
| Model | Story | Copy nats | Stores | R1 | R2 | R3 | R4 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Bag of words | draws, no order | 7.378 | 22 numbers | no | no | no | yes |
| Pair mixture | topic, then draws | 4.605 | 28 bags | no | no | part | yes |
| Window 1 | one token back | 4.159 | 22 rows | part | no | no | yes |
| Window 2 | two tokens back | 4.025 | 484 rows | part | no | no | yes |
| Window 18 | $L+8$ tokens back | 0 | 10,158,317,568 rows | yes | part | no | yes |
| HMM by hand | pair, phase | 0 | 1,120 states | yes | yes | part | part |
| Neural window | $r$ back, embedded | not computed | 5,542 weights | yes | part | part | yes |
**A caveat about counts.** At the position of the final name the prefix holds $Y$ twice and $X$ once, so a model that conditioned on the *counts* of the prefix could tell the two names apart: the name that occurs once is $X$. That is a feature of this template and not a way of binding (a rule that read occurrence counts would fail as soon as the roles occurred a different number of times), and Part C comes back to it. Neither the static unigram nor the pair mixture uses it. Read the "part" in the mixture's R3 cell this way: it uses *which two names* this sentence contains, and nothing about which is $X$ and which is $Y$.
The ranks are not all in one order. At $L=3$ the mixture pays $3.744$ nats on the two copies, better than a window of 1 ($4.159$). At $L=6$ the two tie exactly, since $(10+6)/2=8=N$ and both pay $2\ln8=4.159$. Beyond that the mixture is worse (at $L=10$ it pays $4.605$) because it is blind to order and a longer sentence dilutes the names. The figure shows the whole picture as the gap grows.
Only a table that spans the gap, or a state built to hold the binding, gets the copies for free. Each bar is what a model pays on the two copied names together, in nats, with unlimited data, grouped by family. The bag and the mixture get worse with the gap because they cannot tell where the names go; the short windows stay flat because they never see the source; the long window and the hand-built state pay nothing. Drag the handle to change the gap $L$ and watch the mixture cross the window of 1 at $L=6$.
**What the failures have in common.** Put the table next to the HMM. The table fails R3 because it holds an answer for each binding; the HMM fails it in the same way with a smaller table. In both, the binding is held by *which slot of a fixed list is active*. What is needed instead is a place to *write* the binding while the sentence is read and to *read* it back when it is needed, so that nothing about the binding has to be anticipated in the model's size. The cost of getting it wrong is visible in the numbers: for the table, the training data and the size grow with $F^L$; for the pair-mixture and the bag, the cost on the copy grows with $L$; for the HMM, with $N^2$.
**Where Part B starts.** The HMM idea is right: carry a summary of the past whose size does not depend on how long the past is. The next part keeps that and drops the hand-design. A recurrent network learns its own state $h_t$ from data, as a vector of $d$ numbers updated at every token, and so its parameters do not depend on $L$. That meets R1 and R2. Then come the questions that the HMM bound has already sharpened. How big must the state be? For ToyLang with $N$ names it must tell $N(N-1)$ bindings apart, and a vector of $d$ real numbers can do that in principle; but a long document has many facts, and a state of fixed size must hold all of them. What happens when it cannot? And can such a model be trained without computing the states one after the other? Those are the questions of Part B, and the answer to the last, no, is what makes the move to Part C worth the trouble.
## B. A state that remembers: recurrence and its bottleneck
**The question of this part.** Part A ended with four requirements: order matters (R1), the answer may lie far back (R2), a binding written in this very text must be looked up (R3), and training must run in parallel (R4). Can one fixed-size state, carried from token to token, meet them? This part assumes Part A: the chain rule, the cross-entropy in nats, ToyLang and the n-gram table whose size grows like $F^L$. It goes quickly.
### B1. What does a model have to carry across the gap?
Take ToyLang with $L=40$ fillers. The final name sits 48 tokens after its source, and an n-gram table that answers it needs $N(N-1)\cdot 3\cdot F^{L}\approx 2.2\times10^{33}$ contexts . Before building a better model, ask what the answer actually depends on. The fillers are random and never used again, and so is the thing $T$. The names $X$ and $Y$ are used, and nothing else is. A model does not have to remember the sentence; it has to remember what the *future* depends on.
**The state a prediction needs.** Call the *future* of a prefix the distribution of the rest of the sentence given that prefix. Two prefixes of the same length with the same future can share one summary, because no later question can tell them apart. The least a model can carry after $t$ tokens is the number of distinct futures among prefixes of length $t$. This count is exact in ToyLang, and the figure shows it as nats, $\ln(\text{number of distinct futures})$.
The state a prediction needs is small and flat across the gap. The step line is the logarithm of the number of distinct futures among prefixes that end at that token, in nats ($N=8$ names). After the first name it is $\ln 8$, after the second name it is $\ln 56$, and it stays there for every filler, however many. Drag the ring at the end of the shaded filler band (the gap, drawn compressed) to change the gap length $L$; the line gets longer and no higher.Derivation: the number of distinct futures at each position
The claim is that after $t$ tokens of a ToyLang sentence the number of distinct futures is $1$ at $t=0$, $N$ at $t=1,2$, $N(N-1)$ for $3\le t\le 6+L$, $N$ at $t=7+L,\,8+L$, and $1$ at the last two positions. The generator draws $X$, then $Y\neq X$, then $T$, then the $L$ fillers, uniformly and independently, and everything after the gap is a copy.
**Step 1: what a future is.** A sentence is fixed by $(X,Y,T,f_1,\dots,f_L)$. A prefix of length $t$ fixes some of these. Its future is the uniform distribution over all completions, so two prefixes have the same future exactly when the sets of completions are equal.
**Step 2: before the gap.** At $t=0$ nothing is fixed: one future. At $t=1$ the prefix is $X$ and the completions depend on $X$ (the final name must equal it): $N$ futures. At $t=2$ the prefix is "$X$ gave", still $N$. At $t=3$ the prefix contains $Y$, whose second copy appears later, so the pair $(X,Y)$ is needed: $N(N-1)$ futures.
**Step 3: inside the gap.** The tokens $a$, $T$ and the fillers are drawn independently of everything after the gap, and no completion refers to them. Two prefixes that differ only in $T$ or in a filler have the same set of completions. So the count stays at $N(N-1)$ up to and including the token "then", position $t=6+L$.
**Step 4: after the copies start.** Once the second $Y$ has been written ($t=7+L$) the completion is "thanked $X$ ." and does not depend on $Y$ any more: $N$ futures. The same holds after "thanked". After the final name the completion is "." alone: one future.
**Step 5: the minimum state.** A state $s_t=g(w_{1:t})$ that predicts as well as the full prefix must separate prefixes with different futures, so it takes at least as many values as there are futures. At $t=3,\dots,6+L$ that is $N(N-1)$ values, i.e. at least $\ln\big(N(N-1)\big)$ nats, whatever $L$ is.
**Check.** Listing every sentence and grouping prefixes by their set of completions gives exactly these counts for $(N,F,L)=(3,2,1),(3,2,2),(4,2,2),(3,3,1),(5,2,0)$. For $N=8$ the plateau is $\ln 56=4.03$ nats, $\log_2 56=5.81$ bits. A real model must also know where in the sentence it is: at most $\ln(11+L)$ more nats, one per position.
So the state needed is a few nats, and it does not grow with $L$. An n-gram state is the last $r$ tokens, which contains $X$ only if $r\ge L+8$: that is why its table explodes. A *learned* summary may drop the fillers and keep the names.
**The recurrent network as a story about the data.** Give up the table and let the model compute its own summary, one token at a time:
$$
h_t=\phi\big(W_h\,h_{t-1}+W_x\,e(w_t)+b\big),\qquad p(w_{t+1}\mid w_{1:t})=\sigma\big(W_{\text{out}}\,h_t+b_{\text{out}}\big),\qquad h_0=0 .
$$
Here $h_t\in\mathbb R^{d}$ is the *state*, $\phi$ is an elementwise nonlinearity (tanh in this part), $W_h$ is $d\times d$, $W_x$ maps the input vector $e(w_t)$ into the state, $W_{\text{out}}$ maps the state to scores over the vocabulary $\mathcal V$ and $\sigma$ is the softmax. The story is: *the past matters only through $h_t$*. Unlike the hidden-state models of Part A, $h_t$ is not a latent variable to be inferred; it is a deterministic function of the text so far, computed by the same $W_h$, $W_x$ at every step. This is the recurrent network of Elman (1990) .
**The win over n-grams: the cost does not depend on $L$.** With one-hot inputs, the parameters are $W_h$ ($d^2$), $W_x$ ($d\,\lvert\mathcal V\rvert$), $b$ ($d$), $W_{\text{out}}$ ($\lvert\mathcal V\rvert\,d$) and $b_{\text{out}}$ ($\lvert\mathcal V\rvert$):
$$
\lvert\theta\rvert=d^{2}+2d\,\lvert\mathcal V\rvert+d+\lvert\mathcal V\rvert .
$$
ToyLang with $N=8$ has $\lvert\mathcal V\rvert=8+3+6+5=22$ words (names, things, fillers, and *gave, a, then, thanked, .*), so $d=32$ gives $\lvert\theta\rvert=2486$ numbers , for a gap of 4 tokens or of 400. The same weights are applied at every position (weight sharing), so reach is no longer paid for in parameters.
**What it cannot do yet.** Having *room* for $X$ and $Y$ in $h_t$ does not mean training will put them there. The signal that teaches the net to keep $X$ must travel from the final-name loss back through 48 applications of the same transition. B2 asks how big that signal is.
### B2. Learning through time: why the signal from far back fades
The loss on the final name is 48 steps from the token that decides it. How large is the gradient that reaches that token?
**The chain of Jacobians.** Write $u_t=W_hh_{t-1}+W_xe(w_t)+b$, so $h_t=\phi(u_t)$. A change in $h_s$ changes $h_{s+1}$ through the Jacobian $J_{s+1}=\operatorname{diag}\big(\phi'(u_{s+1})\big)\,W_h$, and so on up the chain. For a loss $\ell$ computed from $h_n$,
$$
\frac{\partial \ell}{\partial h_s}=\Big(J_{s+1}^{\top}J_{s+2}^{\top}\cdots J_{n}^{\top}\Big)\frac{\partial \ell}{\partial h_n}.
$$
This is the chain rule applied $D=n-s$ times; the parameter gradient $\partial\ell/\partial W_h$ sums such terms over every $s$. Training through this sum is called backpropagation through time. Everything depends on the size of a product of $D$ matrices.
**The linear case is exact.** Take $\phi$ the identity, so $J_t=W_h$ and the product is $W_h^{D}$. If $W_h$ is diagonalizable, $W_h=Z\Lambda Z^{-1}$ with eigenvalue magnitudes at most $\rho$ (the *spectral radius*), then $W_h^{D}=Z\Lambda^{D}Z^{-1}$ and $\lVert W_h^{D}\rVert\le\operatorname{cond}(Z)\,\rho^{D}$; since $\rho(W_h^{D})=\rho^{D}\le\lVert W_h^{D}\rVert$, the size is pinned between $\rho^{D}$ and $\operatorname{cond}(Z)\rho^{D}$. For the simplest case, $W_h=g\,Q$ with $Q$ orthogonal, every direction is scaled by exactly $g^{D}$ .
The gradient that reaches back $D$ steps shrinks or grows like the gain to the power $D$. Gradient size on a log scale in a 16-unit network with $W_h$ equal to a gain $g$ times a random orthogonal matrix: violet is the linear net ($g^D$ exactly), ink is tanh units with the band over 20 seeds, dashed is an additive path with a forget gate of 0.99; the dotted line marks 48 steps, the distance of the final name at $L=40$. Drag the handle to change $g$.
At $g=0.9$ the linear gradient 48 steps back is $0.9^{48}=0.0064$; at $g=0.5$ it is $3.6\times10^{-15}$, which is lost when added to an $O(1)$ term in single precision (resolution $1.2\times10^{-7}$) and is only 16 times the resolution of double precision ($2.2\times10^{-16}$); at $g=1.05$ it is $10.4$ and at $g=1.1$ it is $97.0$. In both directions the net is hard to train: a vanishing gradient carries no information about the far past, and an exploding one swamps the near past. Only the knife-edge $g=1$ keeps the size at 1.
**A nonlinearity makes it worse, not better.** For tanh, $0<\phi'\le 1$, so $\lVert\operatorname{diag}(\phi')W_h\rVert_2\le\lVert W_h\rVert_2$ and the product of $D$ Jacobians is at most $\lVert W_h\rVert_2^{D}$. With $\lVert W_h\rVert_2=g<1$ the gradient must vanish; that is a theorem, not a tendency . For $g\ge1$ the bound allows growth, but every step multiplies by factors $1-h^2$ that are below 1: in the figure's net (input scale 0.1, one seed) the tanh gradient at $g=1$ is $0.037$ at distance 48 although the linear one is exactly 1. Across 20 seeds the tanh value at $g=1$ lies between $0.021$ and $0.127$ with median $0.037$, so the gradient is of order 0.02 to 0.13 . Pascanu, Mikolov and Bengio analyse both the vanishing and the exploding gradient, propose clipping the gradient norm to tame the exploding one and a soft constraint for the vanishing case (abstract). Earlier analyses are Hochreiter's 1991 thesis and Bengio, Simard and Frasconi (1994) .
Derivation: what is remembered and what can be learned share one factor
The claim is that in the linear net the influence of an old input on the present state, and the gradient that would teach the net to use it, are the same matrix $W_h^{D}$.
**Step 1: unroll the state.** With $\phi$ the identity, $b=0$ and $h_0=0$, $h_t=W_hh_{t-1}+W_xe(w_t)$. Substituting repeatedly,
$$
h_n=\sum_{s=1}^{n}W_h^{\,n-s}\,W_x\,e(w_s).
$$
**Step 2: forward influence.** The part of $h_n$ that came from the token at position $s$ is $W_h^{\,n-s}W_xe(w_s)$, whose size is at most $\lVert W_h^{\,n-s}\rVert\,\lVert W_xe(w_s)\rVert$.
**Step 3: backward influence.** Differentiating the same sum, $\partial h_n/\partial\big(W_xe(w_s)\big)=W_h^{\,n-s}$, so for a loss $\ell(h_n)$ the gradient with respect to that input is $\big(W_h^{\,n-s}\big)^{\top}\,\partial\ell/\partial h_n$.
**Step 4: conclude.** Both are controlled by $W_h^{n-s}$. If the matrix shrinks the token's trace, the net no longer holds the token; and the same shrinkage removes the gradient that would have said the token mattered.
**Check.** For $W_h=gQ$ with orthogonal $Q$, a unit gradient at step $n$ has size exactly $g^{D}$ at step $n-D$, and the figure's violet line reproduces it for $g=0.5,\dots,1.1$.
**Gates give the signal a path that does not multiply.** The long short-term memory of Hochreiter and Schmidhuber (1997) keeps a second vector, the cell $c_t$, and updates it additively; the record describes constant error carousels in memory cells . The 1997 cell had that carousel and, as recalled, no forget gate; the forget gate was added in later work, which is not in the reading list . The widely used form with a forget gate is written from memory :
$$
c_t=f_t\odot c_{t-1}+i_t\odot \tilde c_t ,\qquad f_t,i_t\in(0,1)^{d}\ \text{computed from } h_{t-1}\text{ and } w_t .
$$
One chain-rule step: the derivative of $c_t$ with respect to $c_{t-1}$ contains the term $\operatorname{diag}(f_t)$ (the other terms pass through the dependence of the gates and of $\tilde c_t$ on $h_{t-1}$). Along this direct path the product over $D$ steps is $\operatorname{diag}\big(\prod_t f_t\big)$: no matrix power, no $\phi'$ factor . If the net learns $f\approx1$ for the unit that holds $X$, the gradient along this path is $f^{D}$: $0.99^{48}=0.62$, $0.98^{48}=0.38$, and $1$ at $f=1$ (dotted line). The decay does not vanish by design; it becomes a *learned* quantity, one number per unit, rather than a property of the whole matrix $W_h$ (this is a property of the forget-gate form; the 1997 cell without a forget gate has the carousel at $f=1$ exactly).
**What it cannot do yet.** The state still has $d$ numbers, however well the gradient reaches them. B3 asks what a fixed state can hold.
### B3. The bottleneck made exact
ToyLang asks the state to hold two names. Now ask for a state that can answer *any* of $m$ questions of the form "which name was bound to key $j$?", with each answer one of $N=8$ names: a text with $m$ bindings instead of two. How large must the state be? There are three separate answers, and they say different things.
**(i) A counting bound.** Suppose the state is $d$ numbers stored with $b$ bits each. It takes at most $2^{bd}$ distinct values. The $m$ answers can form $N^{m}$ different patterns, and if the state must determine every pattern then it must take at least $N^{m}$ values:
$$
2^{\,bd}\ \ge\ N^{m}\quad\Longleftrightarrow\quad bd\ \ge\ m\log_2 N\quad\Longleftrightarrow\quad m\le \frac{bd}{\log_2 N}.
$$
This holds for any design, linear or not, gated or not, but only for finite precision $b$: with exact real numbers, one number could encode everything. It is also weak. ToyLang's two names need $\log_2 56=5.81$ bits, so by counting one number would do; counting does not explain why recurrent nets struggle on ToyLang (B2 does). For a text with bindings: $d=32$ half-precision numbers ($b=16$) allow $m\le170$, and $b=32$ allows $m\le341$.
**(ii) The best linear memory with random keys.** Give each binding a key $k_i\in\mathbb R^{d}$, a random unit vector (a name or position the model has not coordinated with the others), and a value $v_i\in\mathbb R^{N}$, the code of one of the $N$ names (one-hot codes; repeats are allowed, so for $m>N$ several pairs share a code). Store the pairs by adding, and read by multiplying with a key:
$$
S=\sum_{i=1}^{m}\htmlClass{eq-der}{v_i}\,\htmlClass{eq-obs}{k_i}^{\top},\qquad S\,\htmlClass{eq-hid}{k_j}=\htmlClass{eq-der}{v_j}+\underbrace{\sum_{i\ne j}\big(k_i^{\top}k_j\big)\,v_i}_{\varepsilon_j,\ \text{the interference}} .
$$
The state $S$ has $N\cdot d$ numbers whatever $m$ is. The first term is the answer; $\varepsilon_j$ is noise made of the other pairs. $\mathbb E\lVert\varepsilon_j\rVert^{2}=(m-1)/d$ exactly, so the interference-to-signal ratio is $\sqrt{(m-1)/d}$.
Derivation: the retrieval noise is $\sqrt{(m-1)/d}$
The claim is that if the keys are independent and uniform on the unit sphere of $\mathbb R^{d}$, and independent of the unit-norm one-hot values, then $\mathbb E\lVert\varepsilon_j\rVert^{2}=(m-1)/d$.
**Step 1: one overlap.** Fix $k_j$. By rotation invariance of the uniform distribution we may rotate so that $k_j=e_1$. Then $k_i^{\top}k_j$ is the first coordinate of $k_i$. The $d$ coordinates of $k_i$ are exchangeable and $\sum_a k_{ia}^2=1$, so each has mean square $1/d$:
$$
\mathbb E\big[(k_i^{\top}k_j)^{2}\big]=\frac1d .
$$
**Step 2: expand the squared norm.** Write the interference as a sum and square it:
$$
\lVert\varepsilon_j\rVert^{2}=\sum_{i\ne j}(k_i^{\top}k_j)^{2}\,\lVert v_i\rVert^{2}+\sum_{i\ne i'}(k_i^{\top}k_j)(k_{i'}^{\top}k_j)\,v_i^{\top}v_{i'} .
$$
**Step 3: the cross terms vanish.** For $i\ne i'$ the two keys are independent, so $\mathbb E[(k_i^{\top}k_j)(k_{i'}^{\top}k_j)]=\mathbb E[k_i^{\top}k_j]\,\mathbb E[k_{i'}^{\top}k_j]=0$, because a uniform vector on the sphere has mean zero. (This holds whether or not two values are equal, since $v_i^{\top}v_{i'}$ is then a constant multiplying a mean-zero quantity.)
**Step 4: add the diagonal terms.** With $\lVert v_i\rVert=1$ there are $m-1$ diagonal terms, each of mean $1/d$:
$$
\mathbb E\lVert\varepsilon_j\rVert^{2}=\frac{m-1}{d}.
$$
**Step 5: ratio to the signal.** The signal $v_j$ has squared norm 1, so the root-mean-square interference-to-signal ratio is $\sqrt{(m-1)/d}$. It equals 1 at $m=d+1$.
**Check.** In 200 trials per point with 8 queries each, the measured ratio is $0.98$ at $(d,m)=(32,32)$ against $\sqrt{31/32}=0.98$, and within 2.1% of the formula at every one of the 64 grid points of the figure. (The second moment $1/d$ also follows from the density of $k_i^\top k_j$, proportional to $(1-t^2)^{(d-3)/2}$, whose second moment is $1/d$.)
**From noise to a wrong answer.** The answer is read by taking the name whose score is largest. Treat the $N$ scores as Gaussian with the noise variance $(m-1)/(Nd)$ in each coordinate; then the probability of reading the right name is
$$
\Pr(\text{correct})\ \approx\ \int_{-\infty}^{\infty}\mathcal N(z)\,\Phi\!\Big(z+\tfrac{1}{s}\Big)^{N-1}dz,\qquad s=\sqrt{\tfrac{m-1}{N\,d}},
$$
with $\mathcal N$ the standard normal density and $\Phi$ its distribution function. It depends on $m$ and $d$ only through $(m-1)/d$. This is an approximation (the noise is a sum of many terms, but not exactly Gaussian), and the figure puts it next to a seeded simulation.
A linear memory answers correctly until the number of pairs reaches its width, then fades, and the width only sets the scale. Top: the chance of reading back the right name out of 8 against $m/d$; blue dots are simulations at $d=8,16,32,64$, the violet curve the Gaussian closed form at $d=32$. Bottom: the interference-to-signal ratio, measured (blue) against $\sqrt{(m-1)/d}$ (violet). Drag the ring to change $m/d$.
The curves for $d=8,16,32,64$ lie nearly on top of each other when plotted against $m/d$. That is the scaling. Exactly, the second moment of the noise is $(m-1)/d$. Approximately, the accuracy is a function of $(m-1)/d$ alone: at a given $m/d$ the four widths differ by at most 6 points (at $m/d=1.75$), and by about 1 point at $m=d$; each simulated point carries sampling noise of a few points, because the 8 queries of a trial share one memory . **The retrieval noise of a linear random-key memory grows like $\sqrt{m/d}$, and the number of pairs it holds scales like $d$.** At $m=d$ the chance of the right name is about 90% for every width (0.896, 0.891, 0.902, 0.891 for $d=8,16,32,64$), against 12.5% for a guess. At $d=32$ it falls to 72% at $m=64$ and 53% at $m=128$. "Capacity $\approx d$" is a convention: the interference equals the signal at $m=d+1$, where the right name still wins about 90% of the time at $N=8$.
**What is and is not claimed here.** The claim is about *this design*: a state that is a sum of outer products, with keys that are random and uncoordinated. It is not a bound on every state, and it is a different argument from the counting bound (i), which is design-free and says $m\lesssim bd/\log_2N$. If the keys are orthonormal ($m\le d$) the interference is exactly zero, since $k_i^\top k_j=0$ for $i\ne j$; the $\sqrt{m/d}$ is the price of not coordinating the keys. A recurrent network is not this memory: an LSTM is nonlinear and can overwrite, select and forget, and the numbers above predict neither its capacity nor its failures. Part D returns to this memory with better write rules.
**The encoder-decoder, as a reading of this model.** Sutskever, Vinyals and Le encode a source sentence with a multilayer LSTM into a fixed-length vector and decode from it ; Cho et al. introduced an encoder-decoder with a fixed-length vector for translation (abstract only). A source sentence of $m$ tokens is $m$ position-to-token bindings that the decoder can reach only through the final state, so in the model above the fraction that can be recovered falls once $m$ is of the order of the state size. That is an analogy for the direction, not a prediction of the size. Evidence from the abstracts is mixed. Cho, van Merriënboer, Bahdanau and Bengio, in a later analysis of such models, report that performance degrades as the sentence length grows (abstract only), and Bahdanau, Cho and Bengio conjecture that the fixed-length vector is a bottleneck (abstract only). Sutskever et al., however, state that their LSTM did not have difficulty on long sentences, and that reversing the source sentence helped by introducing many short-term dependencies (abstract only): that is an optimization explanation, not a capacity one, and it sits in tension with the capacity narrative of this section. Capacity and optimization are separate causes, and nothing here separates them.
**(iii) Sequential dependence.** The state $h_t$ depends on $h_{t-1}$, which depends on $h_{t-2}$. Evaluating $h_n$ needs $n$ transitions one after another, however many processors are available, and training sees the same chain even though every target word is known in advance. A sentence with $L=40$ is 50 tokens: 50 steps in series, against a single parallel layer in Part C. The exception is a *linear* recurrence, which can be folded in logarithmic depth (derivation below); the nonlinear and gated networks of this part cannot. Backpropagation through time also stores every $h_t$: $n\,d$ numbers per sequence .
Derivation: a linear recurrence can be evaluated in logarithmic depth, a nonlinear or gated one cannot be folded
The claim is that $h_t=\lambda_th_{t-1}+\nu_t$ (scalars or diagonal; the coefficients $\lambda_t,\nu_t$ do not depend on $h$) can be computed for all $t$ in about $\log_2 n$ rounds, which shows that seriality is not forced by recurrence as such.
**Step 1: a step is an affine map.** Write $\tau_t(h)=\lambda_th+\nu_t$. Then $h_t=\tau_t(h_{t-1})$.
**Step 2: composition stays affine.** Apply $\tau_1$ then $\tau_2$:
$$
\tau_2(\tau_1(h))=\lambda_2(\lambda_1h+\nu_1)+\nu_2=(\lambda_2\lambda_1)\,h+(\lambda_2\nu_1+\nu_2),
$$
so a pair $(\lambda,\nu)$ summarizes any run of steps and the rule $(\lambda_2,\nu_2)\circ(\lambda_1,\nu_1)=(\lambda_2\lambda_1,\ \lambda_2\nu_1+\nu_2)$ combines two runs.
**Step 3: the rule is associative, so runs can be combined in parallel.** Composition of functions is associative. In round $r$, let every position combine its run with the run that ends $2^{r-1}$ positions earlier, all positions at once. After round $r$ each position holds the composition of the last $2^{r}$ steps; after $\lceil\log_2 n\rceil$ rounds it holds the composition of all steps up to it, which is $h_t$ (with $h_0=0$ the affine map is just evaluated at 0). This is the parallel scan.
**Step 4: why gates and $\tanh$ do not fold.** The scan needs coefficients that are known in advance. In a gated network the gates are computed from $h_{t-1}$, so $\lambda_t$ is itself a function of the state being computed; in $h_t=\tanh(W_hh_{t-1}+W_xe(w_t))$ the composition of two steps is not of the same form. In both cases there is no fixed-size summary of a run to combine. This does not prove that no faster evaluation exists, only that this one does not carry over, which is why the recurrent nets of B1 and the LSTM are evaluated in series.
**Check.** For $n=8$ the tree needs 3 rounds against 8 sequential steps; for $n=50$, 6 rounds against 50. A linear recurrence meets R4; a nonlinear or gated one does not.
**What it cannot do yet.** Gradients fade (B2), the state is full once $m\approx d$ (B3 ii), and training is serial (B3 iii). Two of three are about the same decision: forcing the whole past through $d$ numbers.
### B4. Do not compress: keep every pair and look back
If compression is the problem, stop compressing. Keep all $m$ keys $k_i$ and values $v_i$, and answer a query by finding the stored key most like it. For the memory of B3 this is exact: for unit keys, $k_i^{\top}k_j\le1$ with equality only when $k_i=k_j$ (Cauchy–Schwarz), so the key most similar to $k_j$ is $k_j$ itself and its value is returned with no interference . In the simulation of the figure, the look-back found the right pair in every query.
Keeping every pair removes the errors and pays for them in storage. Top: the chance of reading back the right name out of 8 against the number of stored pairs $m$, for the fixed state with 32-dimensional keys (dashed, simulated) and for looking back over all stored keys (solid, always 1). Bottom: the numbers kept at the chosen $m$, on one scale: 256 for the fixed state (light bar), $m(d+N)$ for keeping every key and value (dark bar); the dashed frame is the size at $m=256$. Drag the handle on the dashed curve to change $m$.
At $m=256$ and $d=32$ the fixed state is right 39% of the time and holds 256 numbers; keeping everything is right every time and holds $m(d+N)=10{,}240$ numbers. **The price of not compressing is memory that grows with the text, and a read that touches all $m$ items.** This is the trade of Part C.
**Soft look-back.** Bahdanau, Cho and Bengio let the decoder look back over all the source states instead of one final vector: the model learns to focus on the relevant source positions for each target word (abstract only). In the generic form, each decoder step $i$ scores every source state $h_j$ against its own previous state $q_{i-1}$, normalizes the scores with a softmax into weights $a_{ij}$, and reads a weighted sum:
$$
g_i=\sum_{j=1}^{n}\htmlClass{eq-der}{a_{ij}}\,h_j,\qquad a_{ij}=\frac{\exp s_{ij}}{\sum_{j'}\exp s_{ij'}},\qquad s_{ij}=\text{score}(q_{i-1},h_j),
$$
Look back instead of compressing. The query scores every stored state; the weighted sum of the states is the context.
The score in the original paper is a small feed-forward network ; Luong, Pham and Manning compare global and local variants . For ToyLang this is the missing mechanism: the query at the final name is "the name that played role $X$", the source states hold names, and the weighted sum should put almost all its weight on position 0 whether $L$ is 4 or 400.
**Why the gradient now reaches.** The context $g_i$ is a sum with the term $a_{ij}h_j$, so $\partial g_i/\partial h_j$ contains $a_{ij}I$ with no product of Jacobians over the distance between $i$ and $j$. The path from the loss to the stored state has length one, not 48 . (The state $h_j$ itself is still produced by a recurrence, so what it holds is learned through the chain of B2; only the *read* has become direct.)
**Where the requirements stand.**
| | R1 order | R2 reach | R3 look up a binding | R4 parallel training |
| --- | --- | --- | --- | --- |
| Recurrent net, fixed state | built in | parameters do not depend on $L$; gradient fades as $\rho^{D}$ | only through $N\cdot d$ numbers (a linear memory holds about $d$ pairs; the general limit is the counting bound of B3 i) | no: $n$ steps in series (a linear recurrence would fold, B3 iii; nonlinear and gated ones do not) |
| Gated net (LSTM) | built in | learnable decay along the cell path | still $d$ numbers; no clean bound is claimed here | no (the gates depend on the state) |
| Recurrent net with look-back | built in | the read is direct; the encoder is still serial | yes, at a cost growing with $m$ | no: the encoder and decoder remain recurrent |
**What the argument establishes and what it does not.** *Derived here:* the state a prediction needs in ToyLang is $\ln 56$ nats and does not grow with $L$ (B1); in the linear and tanh nets the gradient through $D$ steps is bounded by powers of the gain (B2); a linear random-key memory has retrieval noise $\sqrt{(m-1)/d}$ and holds a number of pairs that scales like $d$ (B3 ii); any state of $bd$ bits holds at most $bd/\log_2N$ independent names (B3 i). *Not claimed:* that an LSTM or GRU is limited to $d$ pairs (gates can select, overwrite and compress nonlinearly; the counting bound is the only general limit derived here); that real encoder-decoders fail because of capacity rather than optimization; that a trained recurrent net cannot solve ToyLang (the state needed is small, so nothing derived here says it cannot). The honest summary is narrower than the slogan: a state that must hold *many independent bindings* in a fixed number of numbers, with keys not coordinated in advance, is capacity-limited; a state that holds *few things that matter* is not.
**What is still missing.** The look-back made the read direct, but it sits on top of two recurrent networks, so training is still serial (R4). It also leaves open what the keys, values and queries really are, and whether the recurrence is needed at all. Part C removes it: each position keeps its key and value, every position queries all earlier ones, and the whole sentence is computed at once.
## C. Attention: a set-to-set map
**The question of this part.** Part A ended with a model that counts and cannot copy a name written in this very sentence; Part B with a model that carries one fixed-size state and must squeeze every name into it. Here the third answer: keep every earlier token as it is, and let each new token *point* at the ones it needs. This part derives that operation (attention), shows what kind of function it is (a map on a set, which is what makes it a bag of words that can look things up), and builds one transformer block from it. It assumes the requirements R1 to R4 of Part A (order, reach, lookup of a binding written in this text, parallel training), the bag-of-words model as a model of a set, and the Bahdanau soft alignment of B4. One convention for ToyLang positions, used throughout: in a sentence with gap $L$ the second name sits at position $L+6$ and the final name at position $L+8$; the final name is predicted from position $L+7$, the token “thanked”, and the second name from position $L+5$, the token “then”. Notation: a sentence is an $n\times d$ matrix $X$ with one row $x_i$ per token; the colour code is the one of “How to read the notes”: blue for observed tokens and keys, copper for the query and hidden quantities, violet for values and derived quantities.
### C1. Attention is a latent pointer: the answer is copied from somewhere, and the model does not know where
**A dictionary that answers by similarity.** A programmer's dictionary stores pairs (key, value). A query is compared with the keys, the one that matches returns its value, and the others return nothing. Attention is the same device made soft. Every stored item $j$ keeps a key $k_j$ (what it answers to) and a value $v_j$ (what it hands over). The query $q$ is the question. Each key scores its similarity to the query, $s_j=q\cdot k_j$; the scores become weights $a_j$ that are positive and add to 1 (the softmax of the scores, with an inverse temperature $\beta$ that C1 derives below); and the answer is the blend $o=\sum_ja_jv_j$. In the figure the six keys are unit vectors on a circle and the query is a unit vector you drag around it; every value is a single number, so the blend is a number and the arithmetic fits on the screen.
Attention is a soft dictionary: the query asks, the keys answer by similarity, the values are blended. Drag the copper query around the circle; dot size and table bar are the weight, the violet marker is the answer $o$. The buttons set β: hard (β = ∞, one key answers), sharp (8) and soft (1).
With the query at $60^\circ$ the best key is the one at $85^\circ$ with value $-0.5$. A hard lookup returns exactly $-0.50$. At β = 8 the key at $20^\circ$ still receives $0.245$ of the weight against $0.754$ for the best, and the answer is $-0.18$; at β = 1 six keys share the weight and the answer, $0.06$, is near the plain average of the six values, $0.02$. Everything that follows is this one picture, with tokens in place of the points and learned vectors in place of the angles. The rest of the chapter asks where the softmax weights come from, in which sense the blend is an expectation, and why it can be trained.
**Many queries at once: the attention matrix.** A sequence does not ask one question; every token asks its own. Stack the $n$ queries as the rows of $Q$ and the $n$ keys as the rows of $K$: the scores of all pairs are the matrix $QK^\top$, entry $(i,j)=q_i\cdot k_j$, and the softmax applied to each row gives the attention matrix $A$, with $A_{ij}=a_j$ of the dictionary above for the query $q_i$. Row $i$ is the weight query $i$ puts on every key, so it adds to 1; column $j$ says how much the key $j$ is read, and has no such constraint. The figure draws an $8\times8$ case with two-dimensional arrows, so you can see why a cell is bright: the weight is large where the angle between the row's arrow and the column's arrow is small.
Attention is a matrix whose every row is its own softmax over the keys. Rows are queries (copper), columns keys (blue); the left grid holds the weights, the right the raw scores $QK^\top$, Σ the row sums. Click a row to spell out its lookup, press causal to mask the future, and drag β: at 0 all rows are equal, and as β grows each row concentrates on its best key.
**What to see.** At β = 3, row 6 puts $0.506$ on key 5, $0.326$ on key 6 and the other six share $0.17$; its lookup gives $o=+0.18$, and the other rows disagree because their arrows do (row 3 puts $0.52$ on key 4, row 2 puts $0.49$ on key 1 and $0.21$ on key 8). Press the causal mask: every score above the diagonal becomes $-\infty$, its weight exactly 0, and each row renormalises. Row 1 keeps only itself ($1.00$), row 3 loses its favourite key 4 and puts $0.90$ on key 3 instead of $0.22$, and row 6 barely moves ($+0.18$ to $+0.16$). Compare the grids: $QK^\top$ has rank 2 whatever $n$ is, but the weights do not (rank 8 for every $\beta>0$ in range, 1 at $\beta=0$). The reason is short: $\exp(\beta x)=\sum_k\beta^kx^k/k!$, and with $x=\cos(\theta_i-\varphi_j)$ each power adds angular frequencies up to $k$, each a rank-2 piece, so the rank grows by two per frequency until it reaches $n$. The softmax is what makes attention more than a low-rank similarity table; C2 asks what happens to this matrix when the tokens are shuffled.
**The prompt.** Take the ToyLang sentence with ten filler words,
> Ann gave Ben a pen so so so now ok yes and yes so and then Ben thanked **?**
The missing token is *Ann*, eighteen tokens back. Part A showed that a table needs $8\cdot7\cdot3\cdot6^{10}\approx1.0\times10^{10}$ distinct contexts to answer this by counting, and Part B that a state of $d$ numbers has to share its room among every binding it carries. Neither is how a person does it: the answer is not computed from the sentence, it is *found* in the sentence. This chapter makes “found” a probabilistic statement and derives the operation that implements it.
**A story about how the answer came to be.** Let the model hold $m$ earlier items, each a pair (key $k_j$, value $v_j$) as in the dictionary above, and let the current token pose a query $q$. The story has one hidden step. A position $J\in\{1,\dots,m\}$ is chosen, with probability that depends on how well its key matches the query, and then the answer is the value stored there:
$$
\Pr(J=j\mid \htmlClass{eq-hid}{q})\;=\;a_j\;=\;\frac{\exp(\beta\,\htmlClass{eq-hid}{q}\cdot \htmlClass{eq-obs}{k_j})}{\sum_{i=1}^{m}\exp(\beta\,\htmlClass{eq-hid}{q}\cdot \htmlClass{eq-obs}{k_i})},\qquad
\text{answer}=\htmlClass{eq-der}{v_J} .
$$
The number $s_j=q\cdot k_j$ is the *score* of position $j$; the weights $a_j$ are its softmax with inverse temperature $\beta$. The hidden $J$ is the **latent pointer**. We never see which position the answer was copied from, only the answer, so what the model can report is the average over the pointer:
$$
o\;=\;\sum_{j=1}^{m}a_j\,\htmlClass{eq-der}{v_j}\;=\;\mathbb E_{J\sim a}\,[\,\htmlClass{eq-der}{v_J}\,].
$$
That is attention. When the values are one-hot token vectors, $v_j=\mathbf e_{w_j}$, $o$ is exactly the mixture $\Pr(\text{token}=w)=\sum_j a_j\mathbf 1[w_j=w]$, a probability distribution over words. When the values are general vectors, $o$ is their mean under the pointer: a first moment, not a sample. A later layer reads that mean through nonlinear maps, so the whole network is not literally a mixture of copy processes; the pointer is an exact reading of *one layer*, and an analogy for the stack. The same device, a softmax over source positions whose weight mass on equal tokens is added up, is how copy mechanisms in pointer networks and pointer-generator summarisers are built.
The same dictionary, applied to a ToyLang sentence: the query at the answer slot finds the first name. Each token carries a key tag (S subject slot, O object slot); the copper card is the query at “?”, the violet bars are the weights $a_j$ on the $m=18$ earlier tokens (the ten fillers share two columns; the dotted line is $1/m$). Drag the handle to change β: at 0 every token gets $1/m$, and as β grows the weight moves to the source and the cost of the copied name falls.
**Why softmax?** Any rule that turns scores into probabilities makes choices, and softmax is what remains when we refuse to make any we cannot justify; read it as treating the scores as log-potentials of an unnormalised distribution over positions, which the normaliser $Z$ turns into probabilities. The derivation below gets it two ways. The first: among all distributions over the $m$ positions that have a given expected score, the softmax is the one that commits to nothing else (maximum entropy), and $\beta$ is the price of sharpness. That route is motivation only: nothing in attention fixes an expected score in advance, so the constraint has no operational meaning here. The second route is the stronger one: if the odds between two positions may depend only on their two scores and not on their absolute level, the odds must be exponential in the score difference.
Derivation: softmax is the maximum-entropy choice, and β is a price
The claim is that among all distributions $a$ on $m$ positions with $\sum_ja_js_j=\bar s$ (an expected score fixed in advance), the one with the largest entropy $H(a)=-\sum_ja_j\ln a_j$ is $a_j=e^{\beta s_j}/Z$, where $Z=\sum_je^{\beta s_j}$ and $\beta$ is the number that makes the expected score equal $\bar s$. Such a $\beta$ exists only if $\bar s$ lies strictly between the smallest and the largest score (at the extremes only a distribution supported on the best-scoring, or the worst-scoring, positions reaches it, which is the limit $\beta\to\pm\infty$).
**Step 1: set up the constrained problem.** We maximise $H(a)$ subject to $\sum_ja_j=1$ and $\sum_ja_js_j=\bar s$. Write the Lagrangian with multipliers $\mu$ and $\beta$:
$$
\mathcal L(a)=-\sum_ja_j\ln a_j\;+\;\beta\Big(\sum_ja_js_j-\bar s\Big)\;-\;\mu\Big(\sum_ja_j-1\Big).
$$
**Step 2: differentiate in one weight.** For each $j$,
$$
\frac{\partial\mathcal L}{\partial a_j}=-\ln a_j-1+\beta s_j-\mu=0
\quad\Longrightarrow\quad
a_j=e^{\beta s_j}\,e^{-1-\mu}.
$$
**Step 3: fix the constant.** The factor $e^{-1-\mu}$ does not depend on $j$. Normalisation $\sum_ja_j=1$ forces it to equal $1/Z$, so
$$
a_j=\frac{e^{\beta s_j}}{Z}.
$$
**Step 4: it is the maximum.** $H$ is strictly concave in $a$ and the constraints are linear, so the one stationary point is the global maximum.
**Step 5: what β means.** Define $\bar s(\beta)=\sum_ja_js_j$. Since $\ln Z$ has derivative $\bar s$ in $\beta$,
$$
H=\ln Z-\beta\bar s,\qquad
\frac{dH}{d\beta}=\bar s-\bar s-\beta\frac{d\bar s}{d\beta}=-\beta\,\frac{d\bar s}{d\beta}.
$$
Differentiating $\bar s=\sum_js_je^{\beta s_j}/Z$ once more gives $d\bar s/d\beta=\operatorname{Var}_a(s)=\sum_ja_j(s_j-\bar s)^2\ge0$. Hence
$$
\frac{dH}{d\bar s}=\frac{dH/d\beta}{d\bar s/d\beta}=-\beta .
$$
Each unit of expected score the pointer gains costs $\beta$ nats of entropy. A large $\beta$ is a sharp pointer, bought with entropy; a small $\beta$ is a blur. The sign matters: $\beta>0$ exactly when the required mean score $\bar s$ is above the mean of the uniform weights, $\tfrac1m\sum_js_j$ (then sharper means a higher mean score, and $\beta$ is a price per unit of score); at $\bar s$ equal to that mean, $\beta=0$ and the entropy is at its maximum $\ln m$, so the rate $dH/d\beta=-\beta\operatorname{Var}_a(s)$ is zero there.
**Step 6: a second route.** Suppose only the *odds* of choosing $j$ over $i$ are allowed to depend on the scores, and only through their difference, because adding the same constant to every score changes nothing a pointer should care about: $a_j/a_i=h(s_j-s_i)$. Multiplying the odds along a chain $j\to i\to k$ gives
$$
\frac{a_j}{a_i}\cdot\frac{a_i}{a_k}=\frac{a_j}{a_k}
\quad\Longrightarrow\quad
h(u)\,h(v)=h(u+v)\qquad\text{for all }u,v,
$$
with $u=s_j-s_i$ and $v=s_i-s_k$.
The only positive continuous solutions are $h(u)=e^{\beta u}$. Normalising the odds gives the same $a_j=e^{\beta s_j}/Z$.
**Check.** Three positions with scores $(1,0,0)$ and a required mean score $\bar s=\tfrac12$: the softmax with $\beta=\ln2$ gives $a=(\tfrac12,\tfrac14,\tfrac14)$ with entropy $1.0397$ nats; the alternative $(\tfrac12,0.3,0.2)$ has the same mean score and entropy $1.0297$, smaller, as the claim says.
**What the temperature does.** At $\beta=0$ all weights are $1/m$ and the output is the plain average of the values. As $\beta\to\infty$ the ratio $a_j/a_i=e^{\beta(s_j-s_i)}$ blows up for any positive score gap, the weight collapses onto the best-scoring position (tied positions share it equally), and attention becomes the hard dictionary lookup of a programmer. In between, the softmax is a *smooth, differentiable* lookup, and that is why it can be trained. For $\beta>0$, raising $\beta$ lowers the entropy of the pointer, at the rate $-\beta\operatorname{Var}_a(s)$ (the rate is 0 at $\beta=0$, where the entropy is at its maximum $\ln m$), and raises the expected score, at the rate $\operatorname{Var}_a(s)$; the second rate is zero only when the pointer is a point mass or all scores are equal.
**ToyLang by hand.** To see the pointer work, set the keys by hand and leave learning for C3. Give every position a key with two numbers, “I am the subject slot” and “I am the object slot”: position 0 has key $(1,0)$, position 2 has key $(0,1)$, every other position has $(0,0)$. The query at “thanked” is $(1,0)$ (give me the subject), the query at “then” is $(0,1)$ (give me the object); the values are the token themselves. The final name is predicted at position $L+7$, which sees $m=L+8$ earlier items; the source scores $1$ and the other $m-1$ items score $0$. So
$$
a_{\text{source}}=\frac{e^{\beta}}{e^{\beta}+(m-1)}=\frac{1}{1+(m-1)e^{-\beta}},
$$
$$
\text{cost}=-\ln a_{\text{source}}=\ln\!\big(1+(m-1)e^{-\beta}\big).
$$
At $L=10$ ($m=18$) this gives weight $0.960$ and cost $0.041$ nats at $\beta=6$, weight $0.998$ and cost $0.0021$ nats at $\beta=9$; at $\beta=0$ the weight is $1/18=0.056$ and the cost $\ln18=2.890$ nats. Compare the other models of Part A on the same token: a bag of words pays $\ln(20\cdot8/4)=3.689$ nats, and an $n$-gram with an 8-token window pays $\ln7=1.946$ nats because the source is eighteen tokens back, out of reach.
**Distance does not appear; the number of competitors does.** The formula has no distance in it. Moving the source to another place among the other items, or putting the other items in another order, changes nothing, because the softmax sees only the multiset of scores (C2 turns this into a theorem). What grows with $L$ is the count $m-1$ of items the pointer must outcompete, and it grows the cost only through the logarithm: the cost is below $\varepsilon$ as soon as
$$
\beta\;\ge\;\ln\frac{m-1}{e^{\varepsilon}-1}.
$$
For $\varepsilon=0.01$ nats this needs $\beta\ge6.55$ at $L=0$ ($m=8$), $7.43$ at $L=10$ and $9.27$ at $L=100$ ($m=108$). So the sharpness the pointer needs grows like $\ln m$, while the table of Part A grew like $6^{L}$ and the bottleneck of Part B had a capacity fixed at about $d$.
The pointer finds a name at any distance; a window and a bag of words do not. Bars show the attention weights at β = 8 on one ToyLang sentence, with the first name (X) and the answer slot (“?”) marked; the bracket is the eight cells an 8-token window sees. Drag the handle to lengthen the gap: the window loses the source after $L=0$, the bag pays more and more, and attention stays below 0.02 nats even at $L=40$.
**It is learnable because the gradient has a meaning.** A pointer should be trained by the answers it produces. Let the target token be $w^\star$ and let $\chi_j=\mathbf 1[w_j=w^\star]$ mark the positions that hold it. The probability of the target is $p=\sum_ja_j\chi_j$. Its log-gradient with respect to the scores is the difference between two distributions over positions: where the pointer *would have been* had we known the answer, minus where it *was*.
Derivation: the gradient of the log-likelihood is the posterior responsibility minus the weight
The claim is that $\dfrac{\partial\ln p}{\partial s_i}=\beta\,(r_i-a_i)$, where $r_i=\dfrac{a_i\chi_i}{p}=\Pr(J=i\mid\text{answer}=w^\star)$ is the posterior probability of position $i$ given the answer.
**Step 1: the derivative of one softmax weight.** Since $\ln a_j=\beta s_j-\ln Z$ and $\partial\ln Z/\partial s_i=\beta a_i$,
$$
\frac{\partial\ln a_j}{\partial s_i}=\beta(\delta_{ij}-a_i)
\quad\Longrightarrow\quad
\frac{\partial a_j}{\partial s_i}=\beta\,a_j(\delta_{ij}-a_i).
$$
**Step 2: the derivative of $p$.** Insert it in $p=\sum_j\chi_ja_j$:
$$
\frac{\partial p}{\partial s_i}=\beta\sum_j\chi_ja_j(\delta_{ij}-a_i)=\beta\,\chi_ia_i-\beta\,a_i\sum_j\chi_ja_j=\beta\,a_i(\chi_i-p).
$$
**Step 3: divide by $p$.**
$$
\frac{\partial\ln p}{\partial s_i}=\beta\Big(\frac{a_i\chi_i}{p}-a_i\Big)=\beta\,(r_i-a_i).
$$
**Step 4: why $r_i$ is a posterior.** By Bayes' rule, $\Pr(J=i\mid w^\star)=\Pr(J=i)\Pr(w^\star\mid J=i)/\Pr(w^\star)=a_i\,\chi_i/p=r_i$, because given $J=i$ the answer is $w_i$ with certainty.
**Step 5: the expectation-maximisation reading.** The joint log-probability of $(J,w^\star)$ is $\ln a_J+\ln \chi_J$, whose score-gradient is $\beta(\delta_{iJ}-a_i)$. Averaging it over $J\sim r$ gives $\beta\sum_jr_j(\delta_{ij}-a_i)=\beta(r_i-a_i)$, the same thing. This is Fisher's identity: the gradient of a marginal likelihood is the posterior average of the complete-data gradient, and its E-step is $r$.
**Step 6: vector values.** When the answer is $o=\sum_ja_jv_j$ and $g=\partial \ell/\partial o$ for a loss $\ell$, then $\partial o/\partial s_i=\beta a_i(v_i-o)$, so $\partial \ell/\partial s_i=\beta a_i\,(v_i-o)\cdot g$: the same shape, with “does this item explain the answer” replaced by “does this item point the output where the loss wants it”.
**Check.** Seven positions with scores $(1,0,0,0,\tfrac12,0,0)$, of which the second and the fifth hold the target, and $\beta=3$: here $p=0.185$ and the formula gives responsibilities $0.182$ and $0.818$ for the two holders (zero elsewhere); central finite differences of $\ln p$ agree with $\beta(r_i-a_i)$ to $10^{-6}$. At $\beta=0$ the weights are uniform and the responsibility splits the two holders evenly, $\tfrac12$ each.
Read the result in words. The score of every position that holds the target is raised, in proportion to its share of the posterior; the scores of positions that do not are lowered by the weight they were given; the gradient is zero exactly when the pointer already equals its posterior. In ToyLang the source is unique, so $r$ is a point mass on it, and every training signal says the same thing: *move weight to the source*.
**What this still does not do.** Three things were assumed. The keys were set by hand from absolute positions 0 and 2, which is a fact about ToyLang and not something a network knows; where keys come from is C3, and a pointer that finds the binding by content instead of by address is the construction of Part E. The output is one blend of values per query; whether the model can ask two questions at once is C4. And the picture treats the $m$ items as a collection: the formula above never used their order. That last remark is the next chapter, and it is both the strength and the weakness of attention.
### C2. A transformer layer is a function on sets: what that gives and what it takes away
**The prompt.** Shuffle the ten tokens of “Ann gave Ben a pen then Ben thanked Ann .” in any way you like, and feed them to the pointer of C1. What changes? The scores $q\cdot k_j$ do not mention order; the normaliser $\sum_je^{\beta s_j}$ is a sum, and sums do not care about order. So the weights follow the tokens wherever they go. Part A met this property in the bag-of-words model: its probability of a sentence depended only on the counts, and any shuffle had the same probability. This chapter shows that attention inherits it exactly, says precisely what survives (equivariance, not invariance), and finds out what it costs.
**The layer as a map on a matrix.** Stack the tokens as rows of $X\in\mathbb R^{n\times d}$. One attention layer maps $X$ to $n$ new rows:
$$
f(X)\;=\;A\,X W_V,\qquad A=\operatorname{softmax}_{\text{row}}\!\big(\beta\,XW_QW_K^\top X^\top\big),
$$
where row $i$ of $A$ is how token $i$ spreads its attention over all tokens (C3 explains the three matrices). A permutation of the tokens is a matrix $\Pi$ with one 1 in each row and column; $\Pi X$ is $X$ with its rows reordered. Two words are needed. A map $f$ is **equivariant** if reordering the input reorders the output in the same way, $f(\Pi X)=\Pi f(X)$ for every $\Pi$. A map $g$ is **invariant** if reordering changes nothing, $g(\Pi X)=g(X)$.
**Claim.** Self-attention is equivariant. It is not invariant: each output row travels with its token.
Derivation: attention is permutation-equivariant
The claim is that $f(\Pi X)=\Pi f(X)$ for the layer above, for every permutation matrix $\Pi$. The proof is three small facts.
**Step 1: the scores permute on both sides.** The score matrix of the reordered input is
$$
\mathcal M(\Pi X)=\beta\,(\Pi X)W_QW_K^\top(\Pi X)^\top=\Pi\,\big[\beta XW_QW_K^\top X^\top\big]\,\Pi^\top=\Pi\,\mathcal M(X)\,\Pi^\top,
$$
using $(\Pi X)^\top=X^\top\Pi^\top$. Here $\mathcal M(X)=\beta XW_QW_K^\top X^\top$ is the score matrix of $X$. In index form, if $\varpi$ is the permutation that $\Pi$ encodes, $\mathcal M(\Pi X)_{ij}=\mathcal M(X)_{\varpi(i)\varpi(j)}$.
**Step 2: the row-wise softmax commutes with the reordering.** Let $A'$ be the attention matrix of $\Pi X$. Then
$$
A'_{ij}=\frac{e^{\mathcal M_{\varpi(i)\varpi(j)}}}{\sum_{l}e^{\mathcal M_{\varpi(i)\varpi(l)}}}=\frac{e^{\mathcal M_{\varpi(i)\varpi(j)}}}{\sum_{l'}e^{\mathcal M_{\varpi(i)l'}}}=A_{\varpi(i)\varpi(j)},
$$
where the denominators agree because $l'=\varpi(l)$ runs over all positions once: a sum does not depend on the order of its terms. In matrix form $A'=\Pi A\Pi^\top$.
**Step 3: the output.** Use $\Pi^\top\Pi=I$:
$$
f(\Pi X)=A'(\Pi X)W_V=\Pi A\,\Pi^\top\Pi\,XW_V=\Pi\,AXW_V=\Pi f(X).
$$
**Step 4: the rest of a block.** The residual add, the per-token MLP and layer norm act on each row separately, so $g(\Pi X)=\Pi g(X)$ for them trivially. A composition of equivariant maps is equivariant, so a stack of any depth is.
**Step 5: where invariance does hold.** With the query held outside the set, $o(q)=\sum_ja_jv_j$ is invariant to any reordering of the pairs $(k_j,v_j)$, because the sum is. Pooling the rows, $\mathbf 1^\top f(X)$, is invariant, because $\mathbf 1^\top\Pi=\mathbf 1^\top$. Neither is what the layer itself computes.
**Check.** A two-layer random model (width 8, residual and a tanh MLP in each layer) on the ten tokens “Flo gave Ann a pen then Ann thanked Flo .” with five swaps applied: every output row, compared with the row of the same token before the move, changes by less than $10^{-14}$.
Without positions the outputs travel with their tokens; with positions they do not. Each strip is one token's output vector (violet positive, grey negative, depth is magnitude) for the two “Flo” tokens (solid bracket) and the two “Ann” tokens (dashed), moved by the handle's number of swaps. Left, no positions: strips arrive unchanged and copies of a name are identical (gap 0.00). Right, sinusoidal positions: copies now differ (gaps 2.02 and 2.33).
**What it means.** A transformer layer without positions is a function on a *set* of vectors. It is the bag-of-words model of Part A with one difference that matters: the bag-of-words model scored each word by its count, while the attention layer lets every item look at the content of every other and carry the result with it. It is a bag of words that can look things up. That is why the pointer of C1 works for any $L$: lookup by content never needed the order. Zaheer et al. (Deep Sets) characterise the permutation-invariant functions on sets and the conditions for equivariant layers, and Lee et al. (Set Transformer) build an attention-based permutation-invariant module; the equivariance of the plain transformer layer is not the subject of either abstract, so it is derived here.
**The price in ToyLang.** Equivariance forces a particular failure. If two tokens have identical input vectors, swapping them leaves $X$ unchanged, so $f(X)=f(\Pi X)=\Pi f(X)$, and the two outputs are equal: the first “Ann” (the giver) and the last “Ann” (the one thanked) cannot be told apart by any layer. Worse, the swap of the first two names exchanges two *different* valid sentences. Take the prefix up to “then”:
$$
\xi=\text{“X gave Y a T}\;f_1\dots f_L\;\text{then”},\qquad \xi'=\text{“Y gave X a T}\;f_1\dots f_L\;\text{then”}.
$$
The second is the first with positions 0 and 2 swapped, $\xi'=\Pi \xi$, and the next token is $Y$ after $\xi$ but $X$ after $\xi'$. The output at “then” is a fixed point of the swap (position $L+5$ does not move), so equivariance gives the same output for both, and the same distribution over names. Whatever that distribution is, it is wrong for one of the two.
Derivation: a position-free model pays at least ln 2 nats on the second name
The claim is that any model whose prediction at “then” is the same for $\xi$ and for $\xi'$ pays at least $\ln2=0.693$ nats on average over the two, for every $L$ and every $N$, and that a lookup-by-content attention layer achieves exactly $\ln2$.
**Step 1: one distribution for two sentences.** The model assigns the same distribution $p$ over names to $\xi$ and to $\xi'$. The correct next name is $Y$ after $\xi$ and $X$ after $\xi'$. The average cost over the two sentences, which are equally likely under ToyLang, is
$$
\bar c=-\tfrac12\big(\ln p(Y)+\ln p(X)\big).
$$
**Step 2: bound it.** The logarithm is concave, so by Jensen's inequality $\tfrac12(\ln p(Y)+\ln p(X))\le\ln\frac{p(Y)+p(X)}{2}$. The two probabilities are disjoint events, so $p(X)+p(Y)\le1$, and
$$
\bar c\;\ge\;-\ln\frac{p(X)+p(Y)}{2}\;\ge\;-\ln\tfrac12=\ln2 .
$$
Equality holds when $p(X)=p(Y)=\tfrac12$.
**Step 3: it is attainable by attention.** Give every name token the key 1 and every other token the key 0, use the query 1, and let the value of a token be its one-hot vector; read only the name coordinates. The weights are equal on all names present in the prefix, and the prefix up to “then” holds exactly two names, each once. The output is $(\tfrac12,\tfrac12)$ on $\{X,Y\}$, and the cost is $\ln2$ on either sentence, whatever $\beta$ and $L$ are.
**Check.** A grid search over $p(Y)=1-p(X)\in(0,1)$ gives a minimum average cost of $0.6931$. Attention over names (keys 1 on names and 0 elsewhere, $\beta=5$) at $L=10$, on a prefix and on the same prefix with its first two names swapped, gives cost $0.6931$ on both.
Compared with Part A, a position-free lookup model pays $0.693$ nats for the second name at every $L$, where the bag of words paid $\ln((10+L)N/4)=3.689$ nats at $L=10$. It is much better, and still not zero. (The *final* name is different: the prefix at “thanked” has $Y$ twice and $X$ once, so a count can tell them apart. That is a coincidence of the template, not a way out; with a different template the same symmetry bites.)
**How the symmetry is broken.** There are three ways, and a model can use any of them.
*Add a position code.* Give token $i$ the input $x_i+\mathrm{pe}_i$, where the position code $\mathrm{pe}_i$ depends on the slot. Collect the codes as the rows of a matrix $E_{\text{pos}}$. Now $f(\Pi X+E_{\text{pos}})\ne\Pi f(X+E_{\text{pos}})$ in general, because $\Pi E_{\text{pos}}\ne E_{\text{pos}}$: the swap moves the tokens and leaves the codes where they were, so the two sentences $\xi$ and $\xi'$ become different inputs. Vaswani et al. add sinusoidal codes, $\mathrm{pe}_{t,2j}=\sin(t/10000^{2j/d})$ and $\mathrm{pe}_{t,2j+1}=\cos(t/10000^{2j/d})$ for position $t$, and report that learned codes gave nearly identical results. Rotary embeddings (Su et al.) instead rotate the query and key by an angle that depends on position, so that the score depends on the relative offset directly (not checked against the original text here).
The sinusoidal code has one property worth deriving, because it is the reason to choose it: moving by a fixed offset is a fixed linear map on the code.
Derivation: a shift by b positions is a rotation of the sinusoidal code
The claim is that for each frequency $\omega$, the pair $(\sin\omega t,\cos\omega t)$ at position $t+b$ is obtained from the pair at $t$ by a rotation that depends on $b$ and $\omega$ but not on $t$.
**Step 1: the angle-addition formulas.**
$$
\sin\omega(t+b)=\sin\omega t\cos\omega b+\cos\omega t\sin\omega b,\qquad
\cos\omega(t+b)=\cos\omega t\cos\omega b-\sin\omega t\sin\omega b .
$$
**Step 2: write it as a matrix.**
$$
\begin{pmatrix}\sin\omega(t+b)\\\cos\omega(t+b)\end{pmatrix}
=\begin{pmatrix}\cos\omega b&\sin\omega b\\-\sin\omega b&\cos\omega b\end{pmatrix}
\begin{pmatrix}\sin\omega t\\\cos\omega t\end{pmatrix}.
$$
**Step 3: stack the pairs.** The full code is the stack of such pairs for the frequencies $\omega_i=10000^{-2i/d}$, so $\mathrm{pe}_{t+b}=R_b\,\mathrm{pe}_{t}$ for one block-diagonal rotation $R_b$ that does not depend on $t$. That is a statement about the code, not yet about the score. Take a score that is a bilinear form of the two codes, $\mathrm{pe}_t^\top B\,\mathrm{pe}_{t+b}=\mathrm{pe}_t^\top B R_b\,\mathrm{pe}_t$. It depends on the offset $b$ alone only if $B$ commutes with the block rotations: $B$ block-diagonal, each $2\times2$ block of the form $\alpha I+\eta J$ with $J=\begin{pmatrix}0&1\\-1&0\end{pmatrix}$ the quarter turn (so the block is a scaled rotation, and the code pairs are ordered as sine, cosine). Then block $i$ contributes $\alpha_i\cos\omega_ib-\eta_i\sin\omega_ib$ times the squared length of that block of the code, which is 1, whatever $t$ is. A general learned $B$ does not have this form, so the sinusoidal code makes relative offsets *easy to represent*; whether a trained head uses them is a separate question.
**Check.** With $d=8$, positions 3 and 7 and $b=4$, $R_b\,\mathrm{pe}_3$ equals $\mathrm{pe}_7$ to rounding error. Vaswani et al. give this as the reason they hypothesised the model could learn to attend by relative position.
*Mask the future.* A language model must not see the token it predicts, so row $i$ may attend only to positions $j\le i$. The mask is not symmetric under reordering, and it breaks equivariance, but less than one might think. In one layer, row $i$ is $\sum_{j\le i}a_{ij}v_j$, a sum over the earlier tokens, so it depends only on $x_i$ and on the *multiset* of the earlier tokens, not on their order: a single masked layer without positions is still a bag of the past. With two layers, where the mask sits matters.
- No mask: the whole model is equivariant, and the output at “then” is the same for $\xi$ and $\xi'$.
- A one-layer model with the mask: still a bag of the past, so the output at “then” is the same for $\xi$ and $\xi'$.
- The mask in layer 2 only: layer 1 is equivariant, so its rows are the same set of rows reordered, and the last row of layer 2 attends to every earlier row, that is, to the whole prefix. A sum over the whole prefix does not see the order, so the last row is the same for $\xi$ and $\xi'$.
- The mask in layer 1, read by a second layer: row $j
Two-layer random model, width 8, on the prefixes $\xi$ and $\xi'$ of the sentence with $L=2$, no positions: no mask, mask in layer 2 only and a one-layer masked model all give a difference below $10^{-15}$; the mask in layer 1 gives $0.60$ whether or not layer 2 is masked. With sinusoidal positions: full attention $0.10$, masks in both layers $0.52$. Shuffling the ten tokens of a sentence with five swaps changes no output without positions (largest change below $10^{-14}$) and changes them by up to $3.75$ with positions.
**What this still does not do.** The layer treats the tokens as a set and lets each ask a question of the others, but it has not yet been said what a key, a value and a query are, why the score is a dot product, or how big it should be. Those are learned maps from the vectors $x_i$, and they are C3. The question left over from this chapter is the order of the *words*: a position code gives every token an address, but a model that needs “the word before this one” or “the name that followed *gave*” must compute that by content, which takes two layers; that is the construction of Part E.
### C3. Queries, keys and values are three projections, and the scale of their product matters
**The prompt.** In C1 the key of position 2 was the vector $(0,1)$ and the query at “then” was also $(0,1)$. In a trained model nobody writes those down: each token arrives as a vector $x_i\in\mathbb R^d$ with $d=512$, say, and both the key and the query must be computed from it. What are the computations, which of their parts matter, and how large should the score be when the key has 64 numbers instead of two?
**Three projections.** With the rows-as-tokens convention, the transformer computes
$$
Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V,\qquad
O=\operatorname{softmax}_{\text{row}}\!\Big(\tfrac{1}{\sqrt{d_k}}\,QK^\top\Big)V\,W_O,
$$
with $W_Q,W_K\in\mathbb R^{d\times d_k}$, $W_V\in\mathbb R^{d\times d_v}$ and $W_O\in\mathbb R^{d_v\times d}$. Compared with C1, the inverse temperature is now $\beta=1/\sqrt{d_k}$ and the matrices carry the scale. This is the scaled dot-product attention of Vaswani et al. ($d_k=d_v=64$ in their base model).
One head, one read. Hover a box to follow one path: the query is scored against every key, and the softmax weights mix the values.
The same vector $x_i$ is projected to a key (how to find me), a value (what I hand over when found) and a query (what I look for). Nothing says which part of $x_i$ is address and which is content: training decides, where in C1 we chose by hand, keys from position and values from token.
**Only a product of the projections matters.** The score of token $i$ for token $j$ is $q_i\cdot k_j=x_i\,W_QW_K^\top\,x_j^\top$, a bilinear form in the two tokens with the single matrix $B=W_QW_K^\top$ of rank at most $d_k$. Replacing $W_Q$ by $W_QG$ and $W_K$ by $W_KG^{-\top}$ for any invertible $G\in\mathbb R^{d_k\times d_k}$ changes nothing, because $W_QG\,(W_KG^{-\top})^\top=W_QGG^{-1}W_K^\top=W_QW_K^\top$. There is therefore no canonical “query space” or “key space”, only a learned similarity between two tokens. The count agrees: rank-$d_k$ matrices in $\mathbb R^{d\times d}$ form a family of dimension $d_k(2d-d_k)$, while the two matrices carry $2dd_k$ numbers; the difference, $d_k^2$, is the dimension of the group of the $G$'s. The same holds on the output side: only $W_VW_O$, of rank at most $d_v$, matters. A head is two low-rank maps: one decides *where to look*, one decides *what to write*. (Checked on a random head: replacing the matrices by $W_QG$ and $W_KG^{-\top}$ changes the scores by less than $10^{-12}$.) The invariance is a property of the bare bilinear score. It is broken by anything that treats the query or key space as a coordinate system: biases added to $q$ and $k$, rotary embeddings (which rotate $q$ and $k$ by a position-dependent angle that a general $G$ does not commute with), and normalisation of $q$ and $k$ (QK-norm, ). In those models only the subgroup of $G$ that preserves the extra structure is a symmetry.
**Why the score is a dot product.** The promise made at the end of C2 was to say why. A score is a number that says how well a query matches a key. The simplest similarity that can be learned and differentiated is a bilinear form, $s(x_i,x_j)=x_i^\top B\,x_j$: it is linear in each argument, so a gradient is easy, and the matrix $B$ is free. Any smooth similarity $s(x_i,x_j)$ is locally of that form: an expansion around a pair of reference vectors, to first order in each argument separately, gives a constant, a term linear in $x_i$ (which cancels in the softmax over $j$), a term linear in $x_j$ (a bias for or against the key alone) and a bilinear cross term; the cross term is what compares this query with this key. Writing $B=W_QW_K^\top$ splits that matrix into the two projections, and the rank bound $d_k$ is a budget: the dot product is taken in $d_k$ dimensions, not $d$. This argues that a bilinear score is the natural first choice; it does not show that nothing better exists, and other score functions (additive attention, with a small network as the score) are in use.
In ToyLang terms: the hand-set keys of C1 are the choice $W_K$ that reads two coordinates of a position code, and the query at “thanked” is a $W_Q$ that maps the token “thanked” to the direction $(1,0)$. Any $G$ gives another equivalent pair.
**The scale.** Suppose $W_Q,W_K$ start random, so the entries of $q$ and $k$ are independent with mean 0 and variance 1. How large is a score? Larger scores sharpen the softmax, and a softmax that is already a point mass carries no gradient. The next derivation makes this exact.
Derivation: the score has variance equal to the key size, and a saturated softmax has no gradient
The claim is that $\operatorname{Var}(q\cdot k)=d_k$ under the assumptions above, so dividing by $\sqrt{d_k}$ restores variance 1; and that the softmax passes less and less gradient as its output approaches a point mass.
**Step 1: the mean.** $q\cdot k=\sum_{i=1}^{d_k}q_ik_i$. By independence $\mathbb E[q_ik_i]=\mathbb E[q_i]\,\mathbb E[k_i]=0$, so the score has mean 0.
**Step 2: the variance of one term.** For independent factors with mean 0,
$$
\operatorname{Var}(q_ik_i)=\mathbb E[q_i^2]\,\mathbb E[k_i^2]-\big(\mathbb E[q_i]\mathbb E[k_i]\big)^2=1\cdot1-0=1 .
$$
**Step 3: add the terms.** The $d_k$ terms are independent, so their variances add:
$$
\operatorname{Var}(q\cdot k)=\sum_{i=1}^{d_k}\operatorname{Var}(q_ik_i)=d_k .
$$
Dividing the score by $\sqrt{d_k}$ divides the variance by $d_k$, leaving 1, whatever the key size.
**Step 4: the Jacobian of the softmax.** With $a=\operatorname{softmax}(s)$ (here $\beta=1$), the derivative of the weight $a_i$ with respect to the score $s_j$ is $a_i(\delta_{ij}-a_j)$, from Step 1 of the derivation in C1. Its trace is
$$
\sum_ia_i(1-a_i)=1-\sum_ia_i^2 .
$$
**Step 5: the point mass is a dead end.** If $a$ is one-hot, every $a_i(1-a_i)$ is zero and so is the whole Jacobian: no change of any score changes the output to first order. A head that has locked onto the wrong key then receives almost no gradient to move off it. Large raw scores push $a$ toward a point mass: if the best of $m$ scores is a gap $\Delta$ above the rest, the largest weight is at least $1/(1+(m-1)e^{-\Delta})$, and the typical gap between the best score and the rest grows in proportion to the spread $\sqrt{d_k}$.
**Check.** For $m=16$ random keys and $d_k=64$, over 200 draws the mean of $1-\sum a_i^2$ is $0.221$ without the scale and $0.866$ with it (the maximum possible value is $1-1/16=0.9375$). The sample variance of $q\cdot k$ over 20,000 pairs is $63.76$ at $d_k=64$ and $3.99$ at $d_k=4$.
Without the scale the softmax freezes as the keys get longer. One query against sixteen random keys: top bars use the raw scores, bottom bars the scores divided by $\sqrt{d_k}$. Drag the handle to double the key size from 1 to 1024: raw weights collapse onto one key (this draw has entropy 0.00 at 1024, the mean over 200 draws is 0.08 nats), the scaled ones do not (about 2.4 nats at every size; uniform is $\ln16=2.77$).
At $d_k=64$ the raw softmax puts on average $0.84$ of its weight on a single key out of sixteen and has mean entropy $0.41$ nats; the scaled one puts $0.25$ on its largest and has entropy $2.35$ nats. Vaswani et al. give the same variance argument in a footnote, say they suspect that large dot products push the softmax into regions of very small gradient, and use additive attention as the comparison. What is derived here is the variance and the Jacobian; the assumption that $q$ and $k$ have independent unit-variance entries holds at initialisation and not after training, where norms differ, so the scale is a good starting point rather than an invariant.
**What this still does not do.** The projections explain *how* a pointer is computed from a vector, but one head makes one pointer, and a block needs more than a pointer: it needs several at once, somewhere to write what they fetch, and a place to compute with it. That is C4.
### C4. Heads, the residual stream and the block
**The prompt.** At “thanked” the model must hand over *Ann*; at “then”, *Ben*. In the block built below both lookups are made at every position by two pointers, and the current token decides which result to use: two pointers, somewhere to write what they fetched, and a gate.
**Several heads are several pointers.** A head is one copy of the machinery of C3 with its own matrices. With $n_h$ heads of key size $d_k=d/n_h$, head $u$ produces $H_u=A_uXW_V^{(u)}\in\mathbb R^{n\times d_k}$ with its own attention matrix $A_u$, the heads are concatenated, and one matrix $W_O\in\mathbb R^{d\times d}$ mixes them. Split $W_O$ into the $n_h$ row-blocks $W_O^{(u)}\in\mathbb R^{d_k\times d}$ that multiply each head's slice.
Heads are pointers; their writes add. Concatenating the heads and applying $W_O$ is the same as adding each head's own write.Derivation: concatenating heads is the same as adding independent head outputs
The claim is that $[H_1\ \cdots\ H_{n_h}]\,W_O=\sum_{u=1}^{n_h}H_uW_O^{(u)}$, so a layer is a sum of $n_h$ independent pointer-and-write terms.
**Step 1: block multiplication.** A row of the concatenation is $(z_1,\dots,z_{n_h})$, where $z_u\in\mathbb R^{d_k}$ is the slice that comes from head $u$, and $W_O$ is the vertical stack of the $W_O^{(u)}$. A row vector times a stacked matrix is the sum of each piece times its block:
$$
(z_1,\dots,z_{n_h})\begin{pmatrix}W_O^{(1)}\\\vdots\\W_O^{(n_h)}\end{pmatrix}=\sum_{u=1}^{n_h}z_u\,W_O^{(u)} .
$$
**Step 2: apply it to every row.** The identity holds row by row, so $[H_1\cdots H_{n_h}]W_O=\sum_uH_uW_O^{(u)}$.
**Step 3: unpack one term.** $H_uW_O^{(u)}=A_u\,\big(XW_V^{(u)}W_O^{(u)}\big)$: the attention matrix of head $u$ (where to look) times a rank-$d_k$ write map (what to put into the stream). Each head has its own $B_u=W_Q^{(u)}W_K^{(u)\top}$ and its own $W_V^{(u)}W_O^{(u)}$.
**Step 4: the parameter count does not depend on $n_h$.** The four matrices have $4d^2$ numbers for any $n_h$ when $d_k=d/n_h$. More heads buy more independent pointers with shorter keys; fewer buy one pointer with a richer key.
**Check.** On a random two-head layer ($n_h=2$) with $n=5$, $d=8$, the concatenated form and the sum of the per-head terms differ by less than $10^{-12}$.
**A shared workspace: the residual stream.** A block does not replace its input; it adds to it. With $\operatorname{LN}$ for layer normalisation (below),
$$
x\;\leftarrow\;x+\operatorname{Attn}\big(\operatorname{LN}x\big),\qquad
x\;\leftarrow\;x+\operatorname{MLP}\big(\operatorname{LN}x\big).
$$
The vector $x_i$ that each position carries is the **residual stream**. Every sublayer reads it through its projections, computes, and *adds* its result back; layers talk to one another only through it. Elhage et al. call it a communication channel that does no processing itself. The stream is a set of *subspaces*, some holding the token, some a position, some what a head fetched; a head reads some to build its query and keys and writes into others. We choose which holds what in the hand-built block below, and Part E returns to the subspaces as the “registers” of a small computer.
**Layer normalisation.** For one position, $\operatorname{LN}(x)=(x-\bar x\mathbf 1)/\operatorname{sd}(x)$ followed by a learned gain and shift, where $\bar x$ and $\operatorname{sd}(x)$ are the mean and standard deviation over the $d$ coordinates. It is invariant to adding a constant to every coordinate and to positive rescaling of the vector, and its output, before the gain, has mean 0 and norm exactly $\sqrt d$ (the squared norm of $x-\bar x\mathbf1$ is $d\operatorname{sd}(x)^2$). For $x=(2,0,-2,4)$: $\bar x=1$, $\operatorname{sd}(x)=\sqrt5=2.236$, and $\operatorname{LN}(x)=(0.447,-0.447,-1.342,1.342)$, whose norm is $2=\sqrt4$. The original transformer normalises after each residual add; many current models normalise before each sublayer, as written above.
**The MLP is a per-token memory.** The second sublayer is $\operatorname{MLP}(x)=\phi(xW_1)\,W_2$ for a nonlinearity $\phi$: it reads the stream at one position, scores it against the columns of $W_1$ (one per hidden unit), and returns a combination of the rows of $W_2$ weighted by those scores. It never looks at other positions; that is attention's job. Part D develops this as a slow associative memory. In the block below it does one job, a gate.
**The shapes.**
| Object | Shape | Meaning |
| --- | --- | --- |
| $X$ | $n\times d$ | the residual stream: one row per token |
| $W_Q,\;W_K$ | $d\times d_k$ | read the stream into queries and keys |
| $W_V$ | $d\times d_v$ | read the stream into values |
| $A=\operatorname{softmax}(QK^\top/\sqrt{d_k})$ | $n\times n$ | row $i$: how token $i$ spreads its attention |
| $O=AVW_O$ | $n\times d$ | what the head writes back into the stream |
| $W_1,\;W_2$ | $d\times d_{\text{ff}},\;d_{\text{ff}}\times d$ | the per-token memory |
The base model of Vaswani et al. has 6 layers, 8 heads, $d=512$, $d_{\text{ff}}=2048$ and $d_k=d_v=64$. A layer then has $4d^2=1{,}048{,}576$ numbers in attention and $2\,d\,d_{\text{ff}}=2{,}097{,}152$ in the MLP, $3{,}145{,}728$ in all (biases and norms ignored). The MLP holds two thirds of the layer only because $d_{\text{ff}}=4d$ was chosen; the memory view does not say it must.
**A block that solves ToyLang, by hand.** One block with explicit matrices and no training. It is deliberately cheap: it finds the names by their *absolute positions* 0 and 2, a fact about ToyLang and not something a network can rely on; the content-addressed construction is in E2. The stream of each token has four subspaces: the token (a one-hot over the $N+14$ words), a position code (a one-hot over positions 0 to 4, enough because both names sit in the first three slots), a *subject copy* and an *object copy* (vectors over the $N$ names). There are two heads, with $d_k=5$ and $d_v=N$.
Two pointers, two copies, one gate. Each head is triggered by one token and writes one copy; the MLP keeps only the copy that belongs to the current token.
- **Head 1:** “thanked” gives the query $\beta\,e_0$ (look at position 0) and every other token the zero query; keys read the position code, values the name coordinates; it writes into the subject copy.
- **Head 2:** the same with “then” mapped to $\beta\,e_2$ (position 2), writing into the object copy.
- **The MLP** has $2N$ hidden units, $\operatorname{ReLU}(\text{subject copy}_n+\text{token}_{\text{thanked}}-1)$ and $\operatorname{ReLU}(\text{object copy}_n+\text{token}_{\text{then}}-1)$, an AND of “this name is in the copy” and “the token asks for it”. The logit of name $n$ is $\kappa$ times the sum of its two hidden units ($\kappa$ plays the role of $\beta$ for the output softmax), and the prediction is the softmax over the names only. The idle head has the zero query and writes the plain average of the names it sees; the gate removes it, since for a copy value in $[0,1]$, $\operatorname{ReLU}(\text{copy}+0-1)=0$.
One block holds two pointers, a workspace and a gate; each step changes one thing. “Flo gave Gus a book then Gus thanked Flo .” with β = 6 and κ = 10, at the two positions that need a lookup. Drag the handle through embed, read, write, gate and predict, and read each panel against the faint one before it; the three bars are the first name, the second name and the sum of the six others.Worked example: the forward pass at “thanked” and “then”
The sentence is “Flo gave Gus a book then Gus thanked Flo .” ($L=0$, $N=8$, $n=10$), $\beta=6$ and $\kappa=10$. Flo is name 5 and Gus is name 6.
**At “thanked” (position 7, eight tokens visible).** Head 1's query is $(\beta,0,0,0,0)$; the scores are $6$ at position 0 and $0$ at the other seven, so the weights are $0.9829$ at position 0 (“Flo”) and $0.0024$ at each of the others. The values are the name one-hots, so the subject copy is $0.9829$ on Flo and $0.0049$ on Gus (Gus sits at positions 2 and 6). Head 2's query is zero, so its weights are $\tfrac18$ each and the object copy is $0.125$ on Flo and $0.250$ on Gus. The gate for the subject group is $\operatorname{ReLU}(\text{copy}+1-1)=\text{copy}$, and for the object group $\operatorname{ReLU}(\text{copy}+0-1)=0$. The logits are $\kappa\times0.9829=9.829$ for Flo and $\kappa\times0.0049=0.049$ for Gus, $0$ for the six others. The softmax over the eight names (the readout is restricted to names) gives $0.9996$ on Flo.
**At “then” (position 5, six tokens visible).** Head 2's query is $(0,0,\beta,0,0)$: the weights are $0.9878$ at position 2 (“Gus”) and $0.0024$ at each of the other five. The object copy is $0.9878$ on Gus and $0.0024$ on Flo. Head 1 averages: the subject copy is $0.1667$ on Flo and $0.1667$ on Gus. The object gate gives $0.9878$ for Gus and $0.0024$ for Flo, the subject gate gives zero. The logits are $9.878$ for Gus, and the probability of Gus is $0.9996$.
**What to take from it, and what not.** The construction shows that one block with two heads and a gate can solve ToyLang for every gap $L$, and it makes the roles concrete: the token and position subspaces are read, the copies are written, the MLP gates. It does *not* show that a trained network finds this, and its trick is cheap: the names sit at absolute positions 0 and 2, so a truncated position code suffices. A block that cannot rely on that, and finds the binding by *content* with two layers, is the point of Part E.
**The whole decoder.**
A block repeated, with a way in and a way out. Step through: each block adds a lookup and then a per-token computation to the stream.
The way in is the embedding $e(w_t)$ plus the position code, giving the first stream $X^{(0)}$. The body is $N_b$ blocks. The way out is the unembedding $W_U\in\mathbb R^{d\times|\mathcal V|}$, which turns the last stream vector at position $t$ into one logit per word, whose softmax is the model's distribution for $w_{t+1}$. Training minimises the cross-entropy of Part A, in nats per token,
$$
\ell(\theta)=-\frac1{n-1}\sum_{t=1}^{n-1}\ln p_\theta\big(w_{t+1}\mid w_{\le t}\big),\qquad
p_\theta(\cdot\mid w_{\le t})=\operatorname{softmax}\big(x^{(N_b)}_tW_U\big).
$$
Many models tie the unembedding to the embedding, $W_U=E^\top$, so that one table serves as both the dictionary of words and the readout. Blocks are stacked because one block is one round of lookups followed by one round of per-token computation: a later block can look up by what an earlier one wrote, which a single block cannot do (E returns to this).
**What this still does not do.** A single block is one round of lookups and one round of per-token computation. Whether many blocks together amount to more is Part D (what the lookups and the MLP are, as memories) and Part E (what stacks of them compute). What can be asked here is what a block costs.
### C5. What attention costs, and what it buys back
**The prompt.** The ToyLang sentence with $L=100$ has $n=110$ tokens. Predicting its final name needs one lookup into 108 earlier items. Predicting *every* token needs every token to be able to look at all the earlier ones, and that is where the cost lives.
**What is computed.** With $d_k=d/n_h$ for $n_h$ heads, one layer on $n$ tokens does the following multiply-adds, counting each one once.
| Part | Count per layer | Why |
| --- | --- | --- |
| projections $Q,K,V,O$ | $4\,n\,d^2$ | four products of an $n\times d$ matrix with a $d\times d$ matrix |
| scores $QK^\top$, all heads | $n^2d$ | per head $n\times d_k$ times $d_k\times n$, times $n_h$ heads |
| weighted sum $AV$, all heads | $n^2d$ | per head an $n\times n$ matrix times an $n\times d_k$ matrix, for each of the $n_h$ heads |
| MLP with $d_{\text{ff}}=4d$ | $8\,n\,d^2$ | two products with a $d\times4d$ matrix |
The terms linear in $n$ add to $12\,n\,d^2$ and the attention matrix adds $2\,n^2d$; they are equal at $n=6d$, which is $3072$ for $d=512$. For short sequences the projections and the MLP dominate; beyond $6d$ tokens the quadratic term does. At $n=110$, $d=512$ the linear part is $346{,}030{,}080$ multiply-adds and the attention matrix $12{,}390{,}400$; at $n=4096$ they are $12{,}884{,}901{,}888$ and $17{,}179{,}869{,}184$. (A count, not a timing; actual speed depends on the hardware and on how the computation is arranged. Memory-efficient implementations avoid writing the whole $n\times n$ matrix, but do not remove the quadratic number of scores.) For a causal model only the entries with $j\le i$ are needed, which is $n(n+1)/2$ per head: $6{,}105$ for $n=110$.
**Depth in time: sequential steps.** An RNN computes $h_t$ from $h_{t-1}$: a chain of $n$ dependent steps, and information between tokens at distance $\delta$ crosses $\delta$ recurrent maps (Part B: a product of $\delta$ Jacobians, where the gradient vanishes). In a transformer all rows of a layer are independent given the layer below, so the chain has as many steps as layers, 6 in the base model whatever $n$ is, and information crosses one softmax.
**Training in parallel.** Because of the mask, row $i$ of every layer depends only on tokens up to $i$, so one forward pass gives all $n$ next-token probabilities and one backward pass their gradients: the $n$ losses of a sentence are computed at once. This is requirement R4 of Part A, and the RNN cannot meet it for the reason above.
Training reuses nothing; generating reuses the keys and values. The cache trades the quadratic recomputation for memory that grows with $n$.Attention trades sequential steps for pairs. Log scale against $n$: the causal scores $n(n+1)/2$ of one head and layer (violet, slope 2), the key-value cache of a 6-layer, $d=512$ decoder (grey, slope 1), the sequential steps of an RNN (black) and of a 6-layer transformer (dashed, flat). Drag the ring along the score curve; $n=20$ and $n=110$ are ToyLang with $L=10$ and $L=100$, and at $n=12{,}287$ one head's scores outnumber the whole cache.
**Generating: the key-value cache.** When generating, the keys and values of the earlier tokens do not change: a causal layer's $k_j$ and $v_j$ depend only on tokens $\le j$. So they are stored once and reused, and the new token needs only its own query against $n$ cached keys: $O(nd+d^2)$ work per layer (the $nd$ for the scores and the weighted sum, the $d^2$ for the new token's projections) instead of recomputing $O(n^2d)$. The price is memory that grows with $n$: $2\,n\,d$ numbers per layer (the keys and values of all heads together, $n_hd_k=d$), $2\cdot n\cdot6\cdot512$ for the base configuration, which is $675{,}840$ at $n=110$ and $25{,}165{,}824$ at $n=4096$. The score matrix of one head in one layer, written out, has $n(n+1)/2$ entries; it equals the whole six-layer cache at $n=4\cdot6\cdot512-1=12{,}287$ (both are $75{,}491{,}328$) and is larger beyond. This is the cost named in Part B: a fixed-size state forgets, and a state that never forgets grows with the text. Attention chooses the second. Part D returns to the same memory as a store that is written once per token and read in full.
**The four requirements, again.** Part A listed what a model of text must do. For the transformer block of this part:
| Requirement | Bag of words | $n$-gram | RNN | Transformer block |
| --- | --- | --- | --- | --- |
| R1 order matters | no (a set) | within the window | yes | yes, if positions are added (C2) |
| R2 reach far back | whole text, unordered | $r$ tokens | any, if the state keeps it (B3) | all earlier tokens, one hop |
| R3 look up a binding written in this text | no | only inside the window | within the capacity of the state | yes, by content (C1) |
| R4 train in parallel | yes | yes | no | yes |
| Cost of the final name, $L=10$ | $3.689$ nats | $1.946$ nats ($r=8$) | set by the state size (Part B) | $0.0021$ nats at $\beta=9$ |
**What it costs, honestly.** Quadratic time in $n$ for the scores, quadratic memory if the whole score matrix is written out, and a cache that grows with every token. Attention also needs the order to be supplied (C2), and it holds nothing between sentences: what it knows about the world is in $W_Q,\dots,W_2$, written by training. The bag-of-words model of Part A counted words and could not look anything up; the RNN of Part B looked everything up through a single bottleneck; the block here keeps the whole past and pays for it in the square of its length.
**Back to the prompt.** To predict “Ann” after “thanked”, with a hundred filler words between, the model carries nothing across the gap: the name is still in the cache, and the query “thanked” finds it, at a cost of $\ln(1+107\,e^{-\beta})$ nats, $0.013$ at $\beta=9$. What it cannot yet explain is *why* this is a memory at all, what is written when, what it forgets and what it would cost to keep less. Those are the questions of Part D; whether a stack of such blocks can do more than look things up is Part E.
## D. Two memories: weights and fast weights
**The question of this part.** In Part C a token read the tokens before it by content. That pile of earlier tokens is a memory, and so are the weight matrices of the network itself. What is the difference? This part treats both as one object, a store of (key, value) pairs that is written and then read, and asks the same three things of each: who writes it, how is it read, and what does it forget. The answer to the first question is most of the story. The weights are written by training, over millions of sentences, and change slowly. The attention state is written by the forward pass, one token at a time, and lasts one sentence. When the read has no softmax, the second store is a weight matrix that the earlier tokens write into: a **fast weight**. **What it assumes.** Part A (the ToyLang ruler, and the split between information that lives in the weights and information that lives in the sentence), Part B (a state of $d$ numbers holds about $d$ random pairs; restated below, so this part stands alone) and Part C (queries, keys, values and the attention formula). The notation is $k_i,v_i,q$ for key, value and query, $S_t$ for a fast-weight state, $m$ for the number of stored pairs, $\beta$ only for inverse temperature, $\eta$ for a write strength, $\gamma$ for a decay factor and $\psi$ for a kernel feature map. **Three words.** A **state** is a fixed-size summary (the matrix $S_t$); a **cache** is the growing list of stored keys and values (softmax attention's KV cache); **memory** is either.
### D1. How much can one matrix remember?
**The prompt.** `Flo gave Ann a pen and yes well well and so then Ann thanked` and the next word is `Flo`. The name was written at position 0 and is wanted at position 14, with thirteen words in between. Nothing in the weights can know that Flo is the answer, because the next sentence will have other names. Whatever holds it was written during this sentence. How many names can such a memory hold before it starts returning the wrong one?
**The smallest memory.** Take pairs $(k_i,v_i)$, with keys $k_i\in\mathbb R^{d_k}$ and values $v_i\in\mathbb R^{d_v}$, and store them in one matrix $S$ by adding an outer product per pair. This is the same object as the fast-weight state $S_t$ of D3. A query reads the matrix by multiplying:
$$
S=\sum_{i=1}^{m}\htmlClass{eq-der}{v_i}\,\htmlClass{eq-obs}{k_i}^{\top},
\qquad
S\,\htmlClass{eq-hid}{q}=\sum_{i=1}^{m}\htmlClass{eq-der}{v_i}\,\big(\htmlClass{eq-obs}{k_i}\cdot\htmlClass{eq-hid}{q}\big).
$$
The read is a similarity-weighted sum of the stored values. Adding a pair is the rank-one update $S\leftarrow S+vk^\top$, and the matrix has $d_v\times d_k$ numbers however many pairs it holds. This is the **Hebbian** write, named after the idea that units active together get connected. Correlation-matrix memories of this kind are old: Willshaw, Buneman and Longuet-Higgins (1969), Kohonen (1972) and Anderson (1972) are the usual references . To keep the algebra light, take $d_k=d_v=d$ in this chapter.
**When the keys are orthonormal the memory is exact.** Read with a stored key $k_j$:
$$
Sk_j=\sum_i v_i\,(k_i\cdot k_j)=v_j
\quad\text{if } k_i\cdot k_j=\delta_{ij}.
$$
But $\mathbb R^{d}$ has only $d$ orthonormal vectors. Past $d$ pairs the keys must overlap, and a read picks up the values of the neighbours:
$$
Sk_j=\underbrace{v_j}_{\text{signal}}+\underbrace{\sum_{i\ne j}(k_i\cdot k_j)\,v_i}_{\text{interference}} .
$$
**How large the interference is.** Let the keys be random unit vectors and the values independent random unit vectors. Measure a read by its squared error $\lVert Sk_j-v_j\rVert^2$, with $\lVert v_j\rVert=1$. An answer of zero has error 1, and that is the yardstick for "worse than remembering nothing". The error of the read is exactly the interference term, $\sum_{i\ne j}\rho_i\,v_i$ with overlaps $\rho_i=k_i\cdot k_j$. Each of the $m-1$ other pairs adds $\mathbb E\rho_i^2=1/d$ (the overlap of two random unit vectors has mean square $1/d$, and the cross terms between independent zero-mean values vanish), so the expected squared error of the Hebbian read is $(m-1)/d$. This is derived in B3(ii), where it is the retrieval noise $\sqrt{(m-1)/d}$ of a linear memory. At $m=1$ the error is 0: a single pair is stored exactly. It reaches the yardstick at $m=d+1$, and that is where the capacity of this memory sits: about $d$ pairs.
A read returns a weighted sum of all the stored values, and only one weight should be large. Each bar is the weight a store puts on one stored value when the first pair (copper, should be 1) is asked for. Drag the number of stored pairs: Hebbian weights on the others are key overlaps that add up to noise, while least squares keeps them at zero until the pairs outnumber the 8 dimensions. Each panel prints this draw's squared error and the mean over draws.
**A better write: least squares.** The Hebbian write is blind: it adds each pair without asking what the matrix already returns. Ask instead for the matrix that recalls all stored pairs as well as possible, $\min_S\sum_i\lVert Sk_i-v_i\rVert^2$. Stack the keys as the rows of $K\in\mathbb R^{m\times d}$ and the values as the rows of $V$. For $m\le d$ with independent keys, $S=V^\top(KK^\top)^{-1}K$ recalls every pair exactly. For $m>d$, $S=V^\top K(K^\top K)^{-1}$, and the read of pair $j$ is $\sum_i (H_K)_{ji}v_i$ where $H_K=K(K^\top K)^{-1}K^\top$ is the $m\times m$ orthogonal projection onto the $d$-dimensional column space of $K$ in $\mathbb R^{m}$. The expected squared error is $1-(H_K)_{jj}$, whose mean over $j$ is $1-d/m$.
Derivation: the least-squares store and its error $1-d/m$
The claim is that the minimiser of $\sum_i\lVert Sk_i-v_i\rVert^2$ recalls $m\le d$ independent keys exactly, and that for $m>d$ its expected squared error on pair $j$ is $1-(H_K)_{jj}$, so $1-d/m$ on average.
**Step 1: the loss in matrix form.** With $K$ and $V$ as above, $\sum_i\lVert Sk_i-v_i\rVert^2=\lVert KS^\top-V\rVert_F^2$.
**Step 2: the normal equations.** Differentiate with respect to $S^\top$ and set the gradient to zero:
$$
2K^\top\big(KS^\top-V\big)=0
\quad\Longrightarrow\quad
K^\top K\,S^\top=K^\top V .
$$
**Step 3: the case $m>d$.** Now $K^\top K$ is $d\times d$ and invertible, so $S^\top=(K^\top K)^{-1}K^\top V$, that is $S=V^\top K(K^\top K)^{-1}$. The recalled values are $KS^\top=H_KV$ with $H_K=K(K^\top K)^{-1}K^\top$.
**Step 4: the case $m\le d$.** Now $KK^\top$ is $m\times m$ and invertible. Take $S^\top=K^\top(KK^\top)^{-1}V$. Then $KS^\top=V$: every pair is recalled exactly, and among all exact solutions this one has the smallest norm.
**Step 5: $H_K$ is a projection.** Write $H=H_K$. It is symmetric, and $H^2=K(K^\top K)^{-1}K^\top K(K^\top K)^{-1}K^\top=H$. Reading the diagonal of $H^2=H$ gives
$$
\sum_iH_{ji}^2=H_{jj},
\qquad\text{so}\qquad
\sum_{i\ne j}H_{ji}^2=H_{jj}-H_{jj}^2 .
$$
**Step 6: the error of pair $j$.** Row $j$ of $HV$ is $\sum_iH_{ji}v_i$. Average over independent unit values (the cross terms vanish, as in B3(ii)):
$$
\mathbb E_V\lVert (HV)_j-v_j\rVert^2=(1-H_{jj})^2+\sum_{i\ne j}H_{ji}^2=(1-H_{jj})^2+H_{jj}-H_{jj}^2=1-H_{jj} .
$$
**Step 7: average over $j$.** The trace of a projection is its rank, here $d$, so $\sum_jH_{jj}=d$ and
$$
\frac1m\sum_j\big(1-H_{jj}\big)=1-\frac dm .
$$
**Check.** At $m=2d$ the error is $1/2$. At $m\to\infty$ it tends to 1: with very many pairs the best linear memory approaches remembering nothing.
Worked example: three pairs in two dimensions
Take $d_k=2$ keys $k_1=(1,0)$, $k_2=(0,1)$, $k_3=(1,1)/\sqrt2$ with scalar values $v_1=3$, $v_2=5$, $v_3=2$. The Hebbian store is $S=3(1,0)+5(0,1)+2(1,1)/\sqrt2=(4.414,\,6.414)$.
Reading $k_1$ gives $4.414$ instead of 3: the extra $1.414$ is $v_3\,(k_3\cdot k_1)=2\times0.707$. Reading $k_2$ gives $6.414$ instead of 5, and reading $k_3$ gives $(4.414+6.414)/\sqrt2=7.657$ instead of 2. The summed squared error is $2+2+32=36$.
The least-squares store solves $K^\top K\,S^\top=K^\top V$ with $K^\top K=\begin{pmatrix}1.5&0.5\\0.5&1.5\end{pmatrix}$ and gives $S=(1.707,\,3.707)$. It reads $1.707$, $3.707$ and $3.828$ against the targets $3$, $5$, $2$, and the summed squared error is $6.69$. Three pairs cannot be stored exactly in two dimensions, but the best matrix is more than five times better than the Hebbian one.
**Whatever the rule, at most $d$ arbitrary pairs.** The Hebbian and least-squares stores differ in how they fail, not in where. To recall every one of $m$ pairs exactly we need $Sk_j=v_j$ for all $j$, which is $m\,d_v$ linear equations in the $d_v\,d_k$ entries of $S$. For arbitrary values these are solvable only if the keys are linearly independent, and $\mathbb R^{d_k}$ holds at most $d_k$ independent vectors. So a matrix read by multiplication holds at most $d_k$ arbitrary pairs exactly, whatever rule wrote it, and it does so by spending $d_vd_k$ numbers on $d_k$ values of $d_v$ numbers each: one stored number per parameter, with nothing saved. Least squares reaches that bound; the Hebbian write does not.
**Where this comes from in training.** The outer-product write is not an exotic operation: gradient descent does it. For a linear layer $y=Sx$ with loss $\tfrac12\lVert Sx-y^\star\rVert^2$, the gradient is $(Sx-y^\star)x^\top$, so a step of size $\eta$ is
$$
S\leftarrow S+\eta\,(y^\star-Sx)\,x^\top .
$$
That is a write of the pair (key $x$, value $y^\star$) by the *delta rule* of Widrow and Hoff (1960) , which D4 studies. The weights of a linear layer are a memory written by SGD, one outer product per training example. The Hebbian write is the same step on a different loss, $-v^\top Sk$, whose gradient is $-vk^\top$.
Derivation: a gradient step on a linear layer is an outer-product write
The claim is that the gradient of $\ell(S)=\tfrac12\lVert Sx-y^\star\rVert^2$ with respect to $S$ is $(Sx-y^\star)x^\top$.
**Step 1: one entry.** Let $e=Sx-y^\star$, so $\ell=\tfrac12\sum_ie_i^2$ and $e_i=\sum_jS_{ij}x_j-y^\star_i$.
**Step 2: chain rule.** Row $i$ of $S$ enters only $e_i$, and $\partial e_i/\partial S_{ij}=x_j$. So
$$
\frac{\partial\ell}{\partial S_{ij}}=e_i\,x_j .
$$
**Step 3: collect the entries.** The matrix with entries $e_ix_j$ is the outer product $e\,x^\top$.
**Check.** A finite-difference derivative of $\ell$ in each entry of a random $4\times 5$ matrix agrees with $e\,x^\top$ to $10^{-8}$.
**What this memory cannot do.** It holds about $d_k$ pairs, and it needs someone to supply the keys and values. In a network there are two suppliers. The weights supply pairs that were written by training, and the sentence supplies pairs that are written while it is read. D2 takes the first, D3 the second.
### D2. Is the MLP a slow memory, and can one fact in it be edited?
**The prompt.** Which word follows `The capital of France is`? That fact is the same in every sentence, so it can sit in the weights. Where? Write a two-layer MLP as $f(x)=W_2\,\phi(W_1x)$ with $d_h$ hidden units, and name the rows of $W_1\in\mathbb R^{d_h\times d}$ as keys $k_i^\top$ and the columns of $W_2\in\mathbb R^{d\times d_h}$ as values $v_i$. Multiplying out the matrix product column by column gives
$$
f(x)=\sum_{i=1}^{d_h}\htmlClass{eq-der}{v_i}\;\phi\big(\htmlClass{eq-obs}{k_i}\cdot x\big).
$$
This is the read of D1, $\sum_iv_i\,(k_i\cdot x)$, with a score function $\phi$ applied to each similarity before the weighted sum. The linear memory is the case $\phi(s)=s$. The difference from D1 is who wrote the pairs: here the rows and columns are weights, written by gradient descent over the whole training set, and they stay the same for every sentence the model ever reads. That is why the MLP is the model's **slow memory**.
**What the nonlinearity does.** A linear score returns, for a query near key $i$, a little of every other value as well. A nonlinearity such as a ReLU switches off the keys that match weakly, so that the pairs a query really resembles dominate the sum. It selects; D5 measures how sharply a softmax separates the pairs it keeps. Geva, Schuster, Berant and Levy (2021) argue that the feed-forward layers of a transformer are key-value memories of this form: their keys match input patterns, their values shift the next-token distribution, and lower layers match shallow patterns while upper layers match semantic ones. Only the abstract was read; no number from the paper is used.
**What the capacity of an MLP counts.** A memory with $d_h$ hidden units keeps $d_h$ explicit key rows and $d_h$ explicit value columns, which is $2d_hd$ numbers. Claims that a sharper score gives a capacity of order $d^2$ or $d^3$ count hidden units, not facts per parameter: the cost of each unit is still $2d$ numbers. D1 already gave the floor for a linear memory, one stored number per parameter, and nothing here beats it for arbitrary values. What a sharper score buys is a cleaner read, not free storage. (With the usual hidden width $4d$ the MLP has $8d^2$ parameters against $4d^2$ for attention, two thirds of a block; that is a design choice, not a consequence of the memory view.)
**ToyLang.** What could the MLP of a trained ToyLang model hold? The fact that the word after `gave` is a name, that `then` is followed by a name, that the fillers come from a list of six. These are the same in every sentence, so they live in the weights, in the sense of Part A. It cannot hold that this sentence's first name is Flo, because that pair did not exist when the weights were written.
**Editing one fact.** If the MLP is a memory, one fact in it should be changeable as one pair is changed in D1. Let $W$ be the second matrix, so a key $k$ (the hidden activation) produces $Wk$. Let $\Sigma_k=\mathbb E[kk^\top]$ be the second-moment matrix of the keys the layer sees on typical text. We want a new matrix $W'$ with $W'k^\star=v^\star$ for one key $k^\star$, and as little change as possible on typical keys. The change on a typical key $k$ is $\Delta k$ with $\Delta=W'-W$, and its expected square is $\mathbb E\lVert\Delta k\rVert^2=\operatorname{tr}(\Delta \Sigma_k\Delta^\top)$. So the edit is the solution of
$$
\min_{\Delta}\ \tfrac12\operatorname{tr}\big(\Delta \Sigma_k\Delta^\top\big)
\quad\text{subject to}\quad
\Delta k^\star=v^\star-Wk^\star .
$$
The answer is a rank-one update:
$$
W'=W+\frac{\big(v^\star-Wk^\star\big)\,\big(\Sigma_k^{-1}k^\star\big)^\top}{k^{\star\top}\Sigma_k^{-1}k^\star}.
$$
It is the delta rule of D1 with the key $k^\star$ replaced by $\Sigma_k^{-1}k^\star$, a key *decorrelated* with respect to the typical keys: the directions that typical keys already use heavily are down-weighted, so that the edit lands where little else lives. (This is a preconditioned key, not a whitened one; whitening would use $\Sigma_k^{-1/2}$.) The naive edit along $k^\star$ alone, $\Delta=(v^\star-Wk^\star)k^{\star\top}/(k^{\star\top}k^\star)$, also sets $W'k^\star=v^\star$, but it disturbs typical keys more by the factor
$$
\frac{\big(k^{\star\top}\Sigma_kk^\star\big)\big(k^{\star\top}\Sigma_k^{-1}k^\star\big)}{\big(k^{\star\top}k^\star\big)^2}\ \ge\ 1 .
$$
Derivation: the rank-one edit and the factor by which the naive edit is worse
The claim is that the minimiser of $\tfrac12\operatorname{tr}(\Delta \Sigma_k\Delta^\top)$ subject to $\Delta k^\star=r$ is $\Delta=r\,(\Sigma_k^{-1}k^\star)^\top/(k^{\star\top}\Sigma_k^{-1}k^\star)$, with $r=v^\star-Wk^\star$ and $\Sigma_k$ symmetric positive definite.
**Step 1: the Lagrangian.** Attach a multiplier vector $\lambda$ to the constraint:
$$
\mathcal L(\Delta,\lambda)=\tfrac12\operatorname{tr}\big(\Delta \Sigma_k\Delta^\top\big)-\lambda^\top\big(\Delta k^\star-r\big).
$$
**Step 2: the stationarity condition.** The gradient of $\tfrac12\operatorname{tr}(\Delta \Sigma_k\Delta^\top)$ with respect to $\Delta$ is $\Delta \Sigma_k$ (as $\Sigma_k=\Sigma_k^\top$), and the gradient of $\lambda^\top\Delta k^\star$ is $\lambda k^{\star\top}$. Setting the sum to zero,
$$
\Delta \Sigma_k=\lambda k^{\star\top}
\quad\Longrightarrow\quad
\Delta=\lambda\,k^{\star\top}\Sigma_k^{-1}.
$$
**Step 3: fix the multiplier from the constraint.** Substitute $\Delta$ into $\Delta k^\star=r$:
$$
\lambda\,\big(k^{\star\top}\Sigma_k^{-1}k^\star\big)=r
\quad\Longrightarrow\quad
\lambda=\frac{r}{k^{\star\top}\Sigma_k^{-1}k^\star}.
$$
**Step 4: the edit.** Putting $\lambda$ back, and using $\Sigma_k^{-1}=(\Sigma_k^{-1})^\top$,
$$
\Delta=\frac{r\,\big(\Sigma_k^{-1}k^\star\big)^\top}{k^{\star\top}\Sigma_k^{-1}k^\star}.
$$
**Step 5: the cost of each edit.** For a rank-one $\Delta=r\,u^\top$ the cost is $\operatorname{tr}(\Delta \Sigma_k\Delta^\top)=\lVert r\rVert^2\,u^\top \Sigma_ku$. With $u=\Sigma_k^{-1}k^\star/(k^{\star\top}\Sigma_k^{-1}k^\star)$,
$$
u^\top \Sigma_ku=\frac{k^{\star\top}\Sigma_k^{-1}\Sigma_k\Sigma_k^{-1}k^\star}{(k^{\star\top}\Sigma_k^{-1}k^\star)^2}=\frac1{k^{\star\top}\Sigma_k^{-1}k^\star}.
$$
With the naive $u=k^\star/(k^{\star\top}k^\star)$ it is $u^\top \Sigma_ku=k^{\star\top}\Sigma_kk^\star/(k^{\star\top}k^\star)^2$.
**Step 6: the ratio.** Divide the naive cost by the optimal one:
$$
\frac{\text{naive}}{\text{optimal}}=\frac{\big(k^{\star\top}\Sigma_kk^\star\big)\big(k^{\star\top}\Sigma_k^{-1}k^\star\big)}{\big(k^{\star\top}k^\star\big)^2}.
$$
By the Cauchy-Schwarz inequality $\big(a^\top b\big)^2\le(a^\top a)(b^\top b)$ with $a=\Sigma_k^{1/2}k^\star$ and $b=\Sigma_k^{-1/2}k^\star$, so $a^\top b=k^{\star\top}k^\star$, the numerator is at least $(k^{\star\top}k^\star)^2$. The ratio is at least 1, with equality exactly when $k^\star$ is an eigenvector of $\Sigma_k$.
**Check.** If $\Sigma_k=I$ the two edits coincide and the ratio is 1. For $d=8$ and $\Sigma_k=\operatorname{diag}(1,e^{-4/7},\dots,e^{-4})$ with $k^\star\propto(1,\dots,1)$ the ratio is $4.41$.
Edit along the preconditioned direction and the keys you did not mean to touch move less. Left: six of ten typical keys (blue) in the plane of the two edit directions; a grey ring is the change the naive edit (along $k^\star$) makes on a key, a violet disc the change of the preconditioned edit (along $\Sigma_k^{-1}k^\star$). Right: the two totals over all ten keys. Drag the ring: at 0 the keys form a round cloud and the edits coincide; as the cloud flattens its long axis turns toward the violet line (within $10^\circ$ from unevenness 3) and the discs shrink while the rings stay.
Meng, Bau, Andonian and Belinkov (2022) locate factual associations in mid-layer feed-forward modules of a GPT model and edit them by a rank-one change to a feed-forward weight (ROME). The closed form above is derived here from the objective stated; it is not checked against the objective of their paper, whose abstract describes a rank-one edit of this kind. Two limits are worth saying. The edit protects typical keys on average, and not any particular other fact. And its premise, that a fact sits in one weight matrix of one layer, is a hypothesis about trained models that the edit's success supports without proving.
**What the MLP cannot do.** Its pairs are fixed when training ends. It can know the capital of France, and it cannot hold the name that the sentence in front of it introduced a moment ago. For that the model needs a memory that is written by the forward pass.
### D3. Attention writes a memory as it reads: fast weights
**The prompt.** `Flo gave Ann a pen and yes well well and so then Ann thanked`. Position 14 must return `Flo`. The pair (first name, Flo) exists in no weight matrix. But the sentence has already passed through the network once, and each layer computed a key and a value for every word. That is a store of pairs, written by the forward pass.
**Attention as a store.** In Part C every token $x_i$ was turned into a key $k_i=W_Kx_i$ and a value $v_i=W_Vx_i$, and the token at the current position into a query $q=W_Qx_t$. Read this as D1 reads a matrix: the sentence has written the pairs $(k_i,v_i)$ for $i\le t$, and the attention output is a read with query $q$,
$$
o_t=\sum_{i\le t}\htmlClass{eq-der}{v_i}\;a_i,
\qquad
a_i=\frac{e^{\beta\,\htmlClass{eq-hid}{q}\cdot\htmlClass{eq-obs}{k_i}}}{\sum_{j\le t}e^{\beta\,\htmlClass{eq-hid}{q}\cdot\htmlClass{eq-obs}{k_j}}} .
$$
It is the read of D2, a weighted sum of values with a score function on the similarities, with two changes. The score is a softmax, which normalises the weights over all stored pairs, and the pairs come from the sentence. The two memories differ in everything except the form of the read:
Same read, different table. The MLP's pairs are written by training and last as long as the model; attention's are written by the forward pass and last one sequence.
**Take the softmax away.** If the weights are the raw similarities, $a_i=k_i\cdot q$, the output is $o_t=\sum_{i\le t}v_i\,(k_i\cdot q_t)$, which is exactly the Hebbian read of D1 over the pairs written so far. Reorder the sum, and the whole history collapses into one matrix:
$$
o_t=\Big(\sum_{i\le t}v_ik_i^\top\Big)q_t=S_t\,q_t,
\qquad
S_t=S_{t-1}+v_t\,k_t^\top,\quad S_0=0 .
$$
The sum is a matrix of fixed shape $d_v\times d_k$ however long the sentence is. Causal linear attention is therefore a recurrent network whose hidden state is a matrix: write the token's pair with one outer product, read with one matrix-vector product. That is the observation of Katharopoulos, Vyas, Pappas and Fleuret (2020): with a kernel feature map and associativity the cost of attention becomes linear in the length, and the causal version is a recurrent network. (Their $\phi$ is the $\psi$ of this note, and their update $\phi(k)v^\top$ is the transpose of the convention here.) With the normaliser $z_t=\sum_i\psi(k_i)$ the read is $o_t=S_t\psi(q_t)\big/\big(z_t^\top\psi(q_t)\big)$ with $S_t=\sum_{i\le t}v_i\psi(k_i)^\top$, equal to causal attention with the kernel $\psi(q)\cdot\psi(k)$; a positive map such as $\psi(x)=\mathrm{elu}(x)+1$ keeps the weights positive.
Derivation: causal attention with a feature kernel is a recurrent network with a matrix state
The claim is that causal attention with weights $\psi(q_t)\cdot\psi(k_i)$, normalised over $i\le t$, equals $S_t\psi(q_t)/(z_t^\top\psi(q_t))$ with $S_t=S_{t-1}+v_t\psi(k_t)^\top$ and $z_t=z_{t-1}+\psi(k_t)$.
**Step 1: the explicit form.** Causal attention with kernel $w(q,k)=\psi(q)\cdot\psi(k)$ returns
$$
o_t=\frac{\sum_{i\le t}w(q_t,k_i)\,v_i}{\sum_{i\le t}w(q_t,k_i)} .
$$
**Step 2: the numerator is a matrix product.** The scalar $\psi(k_i)\cdot\psi(q_t)=\psi(k_i)^\top\psi(q_t)$ is a $1\times1$ matrix, and matrix multiplication is associative, so
$$
v_i\big(\psi(k_i)^\top\psi(q_t)\big)=\big(v_i\psi(k_i)^\top\big)\psi(q_t).
$$
**Step 3: pull $\psi(q_t)$ out of the sum.** Summing over $i\le t$,
$$
\sum_{i\le t}w(q_t,k_i)\,v_i=\Big(\sum_{i\le t}v_i\psi(k_i)^\top\Big)\psi(q_t)=S_t\,\psi(q_t).
$$
**Step 4: the denominator the same way.**
$$
\sum_{i\le t}w(q_t,k_i)=\Big(\sum_{i\le t}\psi(k_i)\Big)^{\!\top}\psi(q_t)=z_t^\top\psi(q_t).
$$
**Step 5: the recursion.** Both sums grow by one term per token, so $S_t=S_{t-1}+v_t\psi(k_t)^\top$ and $z_t=z_{t-1}+\psi(k_t)$, with $S_0=0$ and $z_0=0$. With the identity feature map and no denominator this is the recurrence above.
**Step 6: what it costs.** The explicit form at step $t$ touches $t$ pairs. The recurrence touches one outer product ($d_vd_k$ multiply-adds to write) and one matrix-vector product ($d_vd_k$ to read), whatever $t$ is.
**Check.** For 12 random tokens with $d_k=5$, $d_v=4$ the recurrence equals the explicit sum, and the masked attention matrix $\operatorname{tril}(QK^\top)V$, to $10^{-10}$ at every step, for the identity map and for $\psi=\mathrm{elu}+1$ with its denominator.
Write the sentence in token by token and ask two readers who came first. The 8 × 8 state adds one outer product per word (violet positive, grey negative). Drag the cursor along the sentence: each bar is the cosine between a reader's answer and a name, the true first name (copper) against the strongest wrong one, and the taller wins. At the end the linear read of the state loses the first name while the softmax read of the cache keeps it.
**This is a weight matrix that the tokens write.** Look at what $S_t$ is. It is the weight matrix of a linear layer, $q\mapsto S_tq$. The update $S_t=S_{t-1}+v_tk_t^\top$ is one gradient step, of size 1, on the linear objective $-v_t^\top Sk_t$, whose gradient with respect to $S$ is $-v_tk_t^\top$: a Hebbian write. So after $t$ tokens, $S_t$ is a weight matrix that has taken $t$ such steps, one for each earlier token, inside the forward pass. D4 replaces this linear objective by the error $\tfrac12\lVert Sk-v\rVert^2$, and the resulting delta rule is the natural reading of "trained". A weight matrix that changes at every token is a **fast weight**. The parameters that are ordinary weights, $W_Q$, $W_K$ and $W_V$, change only by gradient descent over the whole training run, and they decide what each token writes into $S_t$ and under what key. They are the slow weights that program the fast ones. For linear attention this is a definition plus an exact equivalence (Schlag et al., below); for a softmax model it is a reading of the architecture, an analogy, and nothing here shows that the $W_K$ and $W_V$ of a trained softmax model write something like a Hebbian state.
The idea is old. Hinton and Plaut (1987) gave each connection both a slow weight and a fast weight, the fast one decaying quickly and cancelling interference. Schmidhuber (a 1991 report, then 1992) proposed a slow network that generates context-dependent weight changes for a fast one, the *fast weight programmer* . Ba, Hinton, Mnih, Leibo and Ionescu (2016) used fast weights to attend to the recent past, and Schlag, Irie and Schmidhuber (2021) showed that linear transformers are formally fast weight programmers and replaced the additive update by a delta-rule update (D4). Of these four, the abstracts were read.
**What is derived here, and what is an analogy.** The equality of causal linear attention and a recurrent matrix state is derived above and checked; so is the reading of the update as a gradient step on a linear objective. The statement that a *softmax* transformer works this way is not derived, because softmax does not compress (below). What the program is, in the figure and in the experiments below: every word is written under a random unit key that stands for its position (a new key for each sentence), with the word's embedding as the value, and the question asks for the key of position 0. It is a fixed rule with no learning, the simplest one that could work; Part E builds a program that works for every gap.
**Closing the loop with Part B.** The state has $d_vd_k$ numbers however long the sentence is, and by D1 no write rule can make it hold more than $d_k$ arbitrary pairs exactly. So a fixed-size fast state is a bottleneck of the kind Part B found in the recurrent network: a gap of $L$ filler words writes $L$ pairs that nobody asked for, each adding noise to the first word. In the experiment below the state is $8\times8$ and the sentence is ToyLang with 8 names, so that guessing the first name is right 1 time in 8.
A fixed state loses the first word as the gap grows, and a cache that grows with the sentence does not. Each curve is the share of 300 random ToyLang sentences in which the reader names the first word correctly. Drag the ring on the Hebbian curve along the gap; the bars compare the 64 numbers of the state with the numbers a cache holds there. The delta rule keeps the latest writes and the first word is the oldest, so it does worst here.
**Softmax keeps every pair, and that is the price.** A softmax read cannot be collapsed into a fixed state as the linear read was. The normaliser $\sum_je^{s_j}$ couples every stored pair to the query, so the read is not a sum of independent contributions $v_i\,f(k_i)$ that can be accumulated one pair at a time into a matrix. The exact statement is narrower: the weight $e^{\beta q\cdot k}$ cannot be written as an inner product $\psi(q)\cdot\psi(k)$ of vectors of any fixed dimension, so no fixed-size state of the linear-attention form (a sum of $v\,\psi(k)^\top$) reproduces the softmax read for all sentences. This does not exclude a state of another form. The bit-counting argument of B3(i) covers arbitrary states at finite precision: a state of $bd$ bits holds at most $bd/\log_2N$ independent names, whatever its form.
Derivation: the exponential weight has no finite feature map
The claim is that for no dimension $D$ is there a map $\psi:\mathbb R^{d_k}\to\mathbb R^D$ with $\psi(q)\cdot\psi(k)=e^{\beta q\cdot k}$ for all $q,k$. A state $\sum_iv_i\psi(k_i)^\top$ with feature dimension $D$ can only produce reads of this form.
**Step 1: what a finite feature map would imply.** If such a $\psi$ existed (one shared map for queries and keys), the Gram matrix $G_{ij}=e^{\beta k_i\cdot k_j}$ of any $n$ keys would equal $\Psi\Psi^\top$ with $\Psi\in\mathbb R^{n\times D}$, so it would have rank at most $D$, whatever $n$ is. With separate maps $\psi_Q$ for queries and $\psi_K$ for keys the matrix is $\Psi_Q\Psi_K^\top$, and the rank bound $D$ is the same.
**Step 2: rewrite the weight.** Since $-\tfrac12\lVert k_i-k_j\rVert^2=k_i\cdot k_j-\tfrac12\lVert k_i\rVert^2-\tfrac12\lVert k_j\rVert^2$,
$$
e^{\beta k_i\cdot k_j}=e^{\beta\lVert k_i\rVert^2/2}\;e^{-\beta\lVert k_i-k_j\rVert^2/2}\;e^{\beta\lVert k_j\rVert^2/2}.
$$
So $G=\Lambda\,G_{\mathrm{gauss}}\,\Lambda$ with $\Lambda$ a positive diagonal matrix and $G_{\mathrm{gauss}}$ the Gram matrix of the Gaussian kernel.
**Step 3: the Gaussian kernel matrix is nonsingular.** It is a standard fact that the Gaussian kernel matrix of $n$ distinct points is positive definite, hence of rank $n$; the conclusion of this derivation is exactly as safe as that fact, and a read rescaled by a query-dependent factor $c(q)$ leaves the normalised read unchanged and the rank argument intact, since row scaling keeps rank $n$. A positive diagonal factor does not change the rank, so $G$ has rank $n$.
**Step 4: contradiction.** For $n>D$ distinct keys, $\operatorname{rank}G=n>D$. So no finite $D$ works, and an exact state of this form would have to grow with the number of distinct keys.
**What can be done instead.** Expand the exponential, $e^{\beta q\cdot k}=\sum_{p\ge0}\frac{\beta^p}{p!}(q\cdot k)^p$. Each term is an inner product of tensor powers, $(q\cdot k)^p=\langle q^{\otimes p},k^{\otimes p}\rangle$, because multiplying out $(\sum_aq_ak_a)^p$ gives $\sum q_{a_1}\cdots q_{a_p}\,k_{a_1}\cdots k_{a_p}$. So $\psi(x)=\big(\sqrt{\beta^p/p!}\;x^{\otimes p}\big)_{p\ge0}$ is an exact feature map, with infinitely many entries. Cut at degree $P$ and the state has $d_v\sum_{p\le P}d_k^{\,p}$ numbers and approximates the read. Approximations with fewer features also exist, for example random-feature maps in the style of Performer.
**Check.** For $d_k=3$ and $P=3$ the feature vector has $1+3+9+27=40$ entries, and a recurrent state over these features reproduces the truncated read to $10^{-10}$. With $d_k=5$ and $\beta=0.5$ on one random draw, the distance of the truncated read from the exact softmax read is $0.61$ at $P=0$, $0.095$ at $P=3$ and $0.0016$ at $P=8$. The Gram matrix of $\exp(k_i\cdot k_j)$ for 10 random keys in $\mathbb R^3$ is nonsingular, so it has rank 10 in three dimensions.
**Memory against cost.** The two designs buy different things. Per head, with $d_k=d_v=d$:
| | Linear fast state | Softmax cache |
| --- | --- | --- |
| Numbers held after $t$ tokens | $d^2$ | $2dt$ |
| Multiply-adds for token $t$ | $2d^2$ | $2dt$ |
| Pairs held exactly | at most $d$ | all $t$ |
The cache holds more numbers than the state beyond $t=d_kd_v/(d_k+d_v)=d/2$ tokens, and softmax costs more per step beyond $t=2d_kd_v/(d_k+d_v)=d$. At $d=64$ and a sequence of $n=4096$ tokens one head costs $33{,}554{,}432$ multiply-adds in total with the recurrent state ($8192$ per token, whatever the position) and $1{,}074{,}003{,}968$ with the cache, which is $32$ times as much. At the end the cache holds $2dn=524{,}288$ numbers against the state's $d^2=4096$ (which equals $n$ here only by coincidence). These are per sequence per head; per token the linear cost is constant and the softmax cost grows with the position. Linear attention's saving is a saving of memory and time, bought with a capacity of $d_k$ pairs. Softmax attention buys unbounded recall with a state that grows with the sentence, at a cost that is quadratic over the sequence.
Derivation: the counts in the table
The claim is that one head with key and value dimensions $d_k$ and $d_v$ costs $2d_kd_v$ multiply-adds per token with a recurrent state and $t\,(d_k+d_v)$ for token $t$ with a cache.
**Step 1: the state.** A write is the outer product $v\,k^\top$: $d_vd_k$ multiplications. Adding it to $S$ is $d_vd_k$ additions. A read $Sq$ is $d_v$ dot products of length $d_k$: $d_vd_k$ multiply-adds. Counting a multiply-add as one operation, one token is a write and a read, $2d_kd_v$.
**Step 2: the cache at position $t$.** The $t$ scores $q\cdot k_i$ are $t$ dot products of length $d_k$: $t\,d_k$. The weighted sum $\sum_ia_iv_i$ is $t$ vectors of length $d_v$ scaled and added: $t\,d_v$. Together $t\,(d_k+d_v)$ (the exponentials are $t$ more operations, ignored).
**Step 3: totals over a sequence.** Summing over $t=1,\dots,n$,
$$
\sum_{t=1}^n2d_kd_v=2nd_kd_v,
\qquad
\sum_{t=1}^nt\,(d_k+d_v)=(d_k+d_v)\frac{n(n+1)}2 .
$$
**Step 4: the numbers.** For $d_k=d_v=64$, $n=4096$: $2\cdot4096\cdot4096=33{,}554{,}432$ and $128\cdot4096\cdot4097/2=1{,}074{,}003{,}968$. Their ratio is $(n+1)/(2d)=32.0$.
**Step 5: the crossings.** The state holds $d_kd_v$ numbers and the cache $t\,(d_k+d_v)$, which are equal at $t=d_kd_v/(d_k+d_v)$. The per-token costs $2d_kd_v$ and $t\,(d_k+d_v)$ are equal at $t=2d_kd_v/(d_k+d_v)$. For $d_k=d_v=64$ these are $32$ and $64$.
**Check.** A double loop adding $2d_kd_v$ per token and $t\,(d_k+d_v)$ for token $t$ gives the same totals.
| Statement | Status |
| --- | --- |
| Causal linear attention equals a recurrent matrix state read by $S_tq_t$ | , and (Katharopoulos et al.) |
| $S_t$ is the result of $t$ gradient steps on per-token linear losses | |
| A matrix read by multiplication holds at most $d_k$ arbitrary pairs exactly | |
| No fixed-size state of the linear-attention form reproduces the softmax read | |
| For linear attention, $W_K,W_V,W_Q$ are the program that writes the fast weights | a definition plus an exact equivalence, (Schlag et al., abstract) |
| For softmax models, the same reading | analogy |
| Trained softmax transformers use their cache like this Hebbian state | not established here |
**Recap before the write rules.** Exact: causal linear attention is a recurrent network with a matrix state, and its write $S_t=S_{t-1}+v_tk_t^\top$ is a gradient step on a linear objective. That recurrence is linear in the sense of B3(iii): the coefficients (here 1 on $S_{t-1}$, and the term $v_tk_t^\top$) do not depend on $S$, so the whole sequence of states can be computed by a parallel scan, and linear attention meets requirement R4 of Part A (training in parallel). The DeltaNet line of work (Yang et al. 2024, abstract) trains the delta-rule variant of D4 in parallel over the sequence. Analogy: that the slow weights of a softmax model program a fast memory. Not established: what trained models write.
**What the fast state cannot do.** It writes everything it sees, whether or not anyone will ask for it, and it can only add. A name that changes cannot be replaced, a filler word costs as much space as a name, and old pairs can only fade by adding noise. D4 asks what a better write looks like.
### D4. Add, decay or correct: what should a fixed-size state write?
**The prompt.** `Ann gave Ben a pen`, and then the sentence says it was Cy who gave it. The pair (giver, Ann) is in the state and a new pair (giver, Cy) arrives under the same key. What should the state return when asked for the giver? With the Hebbian write it returns both, added: Ann plus Cy. A state of $d_k$ numbers per value that only adds cannot change its mind. There are three write rules, and each repairs a different fault.
Forget, then write. Each rule chooses what to erase and what to put in its place.
**Add.** $S_t=S_{t-1}+v_tk_t^\top$. It forgets nothing and so interferes with everything.
**Decay.** $S_t=\gamma\,S_{t-1}+v_tk_t^\top$ with $0<\gamma\le1$. Old pairs fade: the pair written $\tau$ steps ago is scaled by $\gamma^\tau$. Fading removes noise, and it also removes signal. This is the idea of Hinton and Plaut's fast weights, which decay quickly and so cancel interference , and Ba et al.'s fast weights are a decayed sum of outer products of hidden states, the same recurrence with this $\gamma$ (form). The expected squared error is
$$
\mathbb E\lVert\hat v-v\rVert^2=\underbrace{(1-\gamma^{a})^2}_{\text{lost signal}}+\underbrace{\frac1d\Big(\sum_{b=0}^{n-1}\gamma^{2b}-\gamma^{2a}\Big)}_{\text{noise from the other pairs}}
$$
for a pair of age $\tau$ among $n$ writes in $d$ dimensions. At $\gamma=1$ this is the Hebbian error $(n-1)/d$. A smaller $\gamma$ lowers the noise term and raises the lost-signal term, so for old pairs it is no better than remembering nothing.
Derivation: the error of a decayed state
The claim is that after $n$ writes with decay $\gamma$, random unit keys and independent random unit values, the pair written $\tau$ steps before the end is read back with expected squared error $(1-\gamma^\tau)^2+\frac1d\big(\sum_{b=0}^{n-1}\gamma^{2b}-\gamma^{2\tau}\big)$.
**Step 1: the state.** Unrolling $S_t=\gamma S_{t-1}+v_tk_t^\top$, the pair of age $b$ appears with weight $\gamma^b$:
$$
S=\sum_{b=0}^{n-1}\gamma^b\,v_{(b)}k_{(b)}^\top ,
$$
where $(b)$ names the pair written $b$ steps before the end.
**Step 2: read the pair of age $\tau$.** Multiplying by its key $k_{(\tau)}$ and subtracting the target,
$$
Sk_{(\tau)}-v_{(\tau)}=-(1-\gamma^\tau)\,v_{(\tau)}+\sum_{b\ne\tau}\gamma^b\,\big(k_{(b)}\cdot k_{(\tau)}\big)\,v_{(b)} .
$$
**Step 3: average over the values.** The values are independent with mean zero and unit expected squared norm, so the cross terms vanish:
$$
\mathbb E_V\lVert Sk_{(\tau)}-v_{(\tau)}\rVert^2=(1-\gamma^\tau)^2+\sum_{b\ne\tau}\gamma^{2b}\,\big(k_{(b)}\cdot k_{(\tau)}\big)^2 .
$$
**Step 4: average over the keys.** As in D1, $\mathbb E(k_{(b)}\cdot k_{(\tau)})^2=1/d$. Summing over $b\ne \tau$ gives $\frac1d\big(\sum_b\gamma^{2b}-\gamma^{2\tau}\big)$.
**Check.** At $\gamma=1$ it is $(n-1)/d$, the Hebbian error. For $d=8$, $n=24$, $\gamma=0.9$ the formula gives $0.529$, $0.718$, $1.159$ and $1.483$ at ages $0,4,12,23$; the mean of 2000 random streams gives $0.521$, $0.723$, $1.159$ and $1.476$. The formula crosses the line at 1 between age 8 ($0.955$) and age 9 ($1.010$).
**Correct.** The delta rule is a gradient step on $\tfrac12\lVert Sk-v\rVert^2$, the error the state makes on the pair it is about to store. The gradient is $(Sk-v)k^\top$, so a step of size $\eta$ is
$$
S_t=S_{t-1}+\eta\,\big(v_t-S_{t-1}k_t\big)\,k_t^\top
=S_{t-1}\big(I-\eta k_tk_t^\top\big)+\eta\,v_tk_t^\top .
$$
Read the second form as erase, then write: remove from $S$ the part along $k_t$, then write the new value there. With a unit key and $\eta=1$ the new pair is stored exactly, $S_tk_t=v_t$, and writing the same pair a second time changes nothing, because the error is already zero. Here the write strength is $\eta$; $\beta$ is kept for the inverse temperature. Schlag, Irie and Schmidhuber (2021) put this update in place of the additive one; the DeltaNet line of work (Yang et al. 2024) trains it in parallel over the sequence, and Gated DeltaNet (Yang, Kautz and Hatamizadeh) adds a decay gate (abstracts read for those two).
Worked example: Ann is replaced by Cy under the same key
Take $d_k=2$, the unit key $k=(1,0)$, and scalar values with Ann $=3$ and Cy $=5$.
**Hebbian.** After writing Ann, $S=3k^\top=(3,0)$. Writing Cy adds $5k^\top$: $S=(8,0)$, and the read $Sk$ is $8$, the sum of the two.
**Delta rule, $\eta=1$.** After writing Ann, $S=(3,0)$. The error on Cy is $5-Sk=5-3=2$, so the write is $2\,k^\top$ and $S=(5,0)$. The read is $5$: Cy replaced Ann. Writing Cy again: the error is $5-5=0$ and nothing changes.
**A neighbour.** Suppose a second pair with the unit key $k'=(0.6,0.8)$ and value 7 was written first, so the state before Ann is $S=7k'^\top=(4.2,\,5.6)$. Writing Ann $=3$ under $k$ with the delta rule: the error is $3-Sk=3-4.2=-1.2$, so $S'=(4.2-1.2,\,5.6)=(3,\,5.6)$ and $S'k=3$, exactly. But $S'k'=3\times0.6+5.6\times0.8=6.28$ instead of $7$. The neighbour is disturbed by $-0.72=0.6\times(-1.2)$: the overlap of the keys times the error that was corrected.
**What it guarantees, and what it does not.** The guarantee is local: the latest pair is stored exactly, and if the keys are orthonormal every pair is exact, because writing along $k_t$ does not move the read of any orthogonal key. Nothing more is promised. With random unit keys and one pass, writing pair $t$ moves the read of an earlier pair $i$ by $(v_t-S_{t-1}k_t)\,(k_t\cdot k_i)$, so each later write disturbs the old ones in proportion to its error. The exact storage of $d$ pairs needs more than one pass, or an orthogonal set of keys, or a different algorithm.
**The online algorithm that does store $d$ pairs.** Least squares can be computed one pair at a time. Recursive least squares replaces the key by the preconditioned key $R_{t-1}k_t$, where $R_{t-1}=(\lambda I+\sum_{i0$ that solution is shrunk: it stores $d$ pairs exactly only in the limit $\lambda\to0$ (on one draw at $d=16$ and $m=16$ pairs the squared error of the read is $0.339$ for $\lambda=0.5$ and $1.3\times10^{-8}$ for $\lambda=10^{-6}$; over many draws the $\lambda=10^{-6}$ error has a median near $2\times10^{-8}$ but a mean dominated by ill-conditioned key matrices).
Derivation: recursive least squares is the delta rule with a preconditioned key
The claim is that the ridge solution $S_t=G_tR_t$, with $G_t=\sum_{i\le t}v_ik_i^\top$ and $R_t=(\lambda I+\sum_{i\le t}k_ik_i^\top)^{-1}$, satisfies the update above.
**Step 1: the batch solution.** Minimising $\sum_i\lVert Sk_i-v_i\rVert^2+\lambda\lVert S\rVert_F^2$ gives, as in D1 with the ridge term, $S=G\,(\lambda I+\sum_ik_ik_i^\top)^{-1}=GR$.
**Step 2: update the inverse.** Write $R=R_{t-1}$, $k=k_t$. The Sherman-Morrison formula for a rank-one change of an inverse gives
$$
R_t=R-\frac{R\,kk^\top R}{1+k^\top Rk}.
$$
**Step 3: update the product.** $G_t=G_{t-1}+vk^\top$, so $S_t=G_{t-1}R_t+vk^\top R_t$. Multiply out each term:
$$
G_{t-1}R_t=S_{t-1}-\frac{S_{t-1}k\,k^\top R}{1+k^\top Rk},
\qquad
vk^\top R_t=\frac{v\,k^\top R}{1+k^\top Rk}.
$$
The second uses $k^\top R_t=k^\top R\big(1-\tfrac{k^\top Rk}{1+k^\top Rk}\big)=\tfrac{k^\top R}{1+k^\top Rk}$.
**Step 4: add.**
$$
S_t=S_{t-1}+\frac{\big(v-S_{t-1}k\big)\,(Rk)^\top}{1+k^\top Rk}.
$$
**Check.** Running the recursion over 20 random pairs with $d_k=6$ and $\lambda=0.5$ reproduces the batch ridge solution $V^\top K(K^\top K+\lambda I)^{-1}$ to relative error below $10^{-12}$.
The same pairs, four ways to store and read them. Squared error of the recalled value (1 is the error of answering nothing) against the number of stored pairs in 16 dimensions, mean of 200 draws. Hebbian crosses 1 at 17 pairs and least squares is exact up to 16; the softmax cache stays near 0 because it keeps every pair, and the delta rule in a single pass is better than Hebbian but far from exact. Drag the ring on the Hebbian curve.A fixed-size state must forget something, and each rule chooses differently. Squared error of the recalled value against the age of the pair, 24 pairs in 8 dimensions (lines are closed forms where they exist, dots the mean of 200 streams). Adding forgets nothing and is wrong everywhere; the delta rule is exact for the latest pair and worsens with age; decay trades lost signal for less noise, so drag the ring to see where its line crosses 1; the softmax cache keeps everything.
**What to take from the two figures.** Past its capacity, a fixed-size state blurs the pairs it holds, and three of the four rules compared here (add, decay, one-pass delta) end up worse than answering nothing for old pairs. The rules differ in what they sacrifice: adding sacrifices the answer to noise, decay sacrifices old pairs on purpose, the delta rule sacrifices old pairs by overwriting along new keys. In practice the gain is found in combining them and in learning, per token, how much to write ($\eta_t$) and how much to keep ($\gamma_t$), which is what gated and delta-rule architectures do (; not checked against any particular model). Only the cache avoids forgetting by construction, and it pays for that in memory.
**The ToyLang payoff: write less.** The failure in D3 was that the 8 + L words of a ToyLang sentence were all written, the fillers included. The memory view suggests the repair of B1's minimal state: write only what will be asked for. Give each write a strength $\eta_t$ (the write strength of the delta rule, here applied to the Hebbian write) and set it to 0 on the fillers and the function words and to 1 on the content tokens, the names and the thing: $S_t=S_{t-1}+\eta_t\,v_tk_t^\top$. In every sentence this writes 4 pairs, whatever $L$ is, into the same $8\times8$ state of 64 numbers. The gate is hand-set: it is read off the template, and what training finds is not tested here. The question is unchanged (name the first word from the state, 300 sentences per gap, 8 names). With the gate the Hebbian state names the first word in $93\%$ of sentences at $L=0$ and $93\%$ at $L=40$ (between $93\%$ and $96\%$ at every gap tried), against $84\%$ falling to $40\%$ without it. A least-squares write of the same 4 pairs is exact at every gap, $100\%$, because 4 pairs do not outnumber the 8 dimensions. The one-pass delta rule is worse than the Hebbian write on this question, $74\%$ to $82\%$: the first name is the oldest of the four writes. Accuracy no longer depends on $L$ because the state no longer does: the four keys are fresh random unit vectors in every sentence, whatever the gap.
Writing only what will be asked for makes the state's accuracy independent of the gap. The write strength $\eta_t$ is 1 on the names and the thing and 0 on fillers and function words. Below, the share of 300 ToyLang sentences in which the Hebbian state names the first word, with every word written (grey) and with the gate (black); drag the ring along the gap. The gate is set by hand; the figure does not say that training finds it.
**What these states still cannot do.** The write rule changes how the state fails; it does not change where. Beyond $d_k$ arbitrary pairs the state blurs. To keep many pairs and still recall one sharply we need a read that separates them better than the linear one, and that is what the softmax is. A gate helps only when the template tells in advance what to keep.
### D5. How much does a softmax read separate the pairs it keeps?
**The prompt.** The cache of D3 holds all 14 pairs of the Flo sentence, and a linear read of them names the wrong first word in about three sentences of ten at this gap. A softmax read of the same pairs names it right in all 300 sentences tried. Why does swapping the score function change the capacity so much, and what is the capacity?
**The soft dictionary.** With stored pairs $(k_i,v_i)$ and a query $q$, the softmax read is
$$
o=\sum_{i=1}^{m}a_i\,v_i,
\qquad
a_i=\frac{e^{\beta\,k_i\cdot q}}{\sum_je^{\beta\,k_j\cdot q}},
$$
where the inverse temperature $\beta$ multiplies the scores. C1 derives this read and its two limits (the average of all values at $\beta=0$, a hard lookup as $\beta\to\infty$), and here it is taken as given. What matters below is that $\beta$ is the worth of one unit of similarity: the weight ratio of two keys is $e^{\beta\,(s_i-s_j)}$.
**Capacity, with the right yardstick.** Use unit keys, ask for stored pair $j$ (so $q=k_j$ and its score is $1$), and keep the values independent and random as in D1. The signal has weight $e^{\beta}$ and pair $i$ has weight $e^{\beta s_i}$ with $s_i=k_i\cdot k_j$; the common normaliser cancels in the ratio. The expected squared interference-to-signal ratio is
$$
\mathbb E\Big[\frac{\text{interference}^2}{\text{signal}^2}\Big]=(m-1)\,\mathbb E\big[e^{2\beta s}\big]\,e^{-2\beta},
$$
and we call the **capacity** the number of pairs where it equals 1:
$$
m^\star(\beta)=1+\frac{e^{2\beta}}{\mathbb E\big[e^{2\beta s}\big]} .
$$
The same yardstick applied to the Hebbian store gives $(m-1)/d=1$, that is $m^\star=d+1$. The yardstick has a geometric meaning: when the interference is as long as the signal and points elsewhere, the read makes an angle of 45 degrees with the wanted value. It is a convention, and the number of pairs at which a read is *usable* depends on how it is decoded. For the softmax it is a different threshold from the error-of-nothing of D1, because there the signal weight $a_j$ is itself below 1. For $s$ the dot product of two random unit vectors in $\mathbb R^d$, treating $s$ as Gaussian with variance $1/d$ gives $\mathbb E[e^{2\beta s}]\approx e^{2\beta^2/d}$ and
$$
m^\star(\beta)\approx1+e^{2\beta(1-\beta/d)} .
$$
This approximation is good for $\beta$ small against $d$, and it fails before $\beta$ reaches $d$: the exact value comes from integrating $e^{2\beta s}$ against the true density of $s$. At $d=16$ the approximation gives $34.1$ against the exact $34.5$ at $\beta=2$, and $404$ against $483$ at $\beta=4=d/4$ (16% low); at $\beta=8$ it gives $2{,}982$ against $19{,}591$, so the exact capacity is $6.6$ times larger. At $d=64$ and $\beta=8$ it gives $1.20$ million against $1.27$ million (6% low).
Derivation: the capacity of the softmax read, exactly and in the Gaussian approximation
The claim is that for unit keys the read of stored pair $j$ has expected squared interference-to-signal ratio $(m-1)\mathbb E[e^{2\beta s}]e^{-2\beta}$, and that it equals 1 at $m^\star=1+e^{2\beta}/\mathbb E[e^{2\beta s}]$.
**Step 1: the read.** With $q=k_j$, $a_i=e^{\beta s_i}/Z$, $s_i=k_i\cdot k_j$ and $s_j=1$. So $o=a_jv_j+\sum_{i\ne j}a_iv_i$: the signal has weight $a_j=e^{\beta}/Z$ and the interference is $\sum_{i\ne j}a_iv_i$.
**Step 2: average over the values.** As in D1, for independent unit values with mean zero the expected squared norm of the interference is $\sum_{i\ne j}a_i^2$. The ratio to $a_j^2$ is $\sum_{i\ne j}e^{2\beta s_i}\big/e^{2\beta}$, and $Z$ has cancelled.
**Step 3: average over the keys.** The $m-1$ overlaps $s_i$ are identically distributed, so the expectation of the sum is $(m-1)\mathbb E[e^{2\beta s}]$. Dividing by $e^{2\beta}$ gives the ratio above, and setting it to 1 gives $m^\star$.
**Step 4: the exact moment.** For a uniformly random unit vector in $\mathbb R^d$ the dot product with a fixed unit vector has density proportional to $(1-s^2)^{(d-3)/2}$ on $[-1,1]$. With $s=\cos\theta$ this is $\sin^{d-2}\theta$, so
$$
\mathbb E[e^{ts}]=\frac{\int_0^\pi e^{t\cos\theta}\sin^{d-2}\theta\,d\theta}{\int_0^\pi\sin^{d-2}\theta\,d\theta},
\qquad t=2\beta .
$$
**Step 5: the Gaussian approximation.** For large $d$, $s$ is close to a Gaussian with mean 0 and variance $1/d$, whose moment generating function gives $\mathbb E[e^{ts}]\approx e^{t^2/(2d)}$. At $t=2\beta$ this is $e^{2\beta^2/d}$. Then $m^\star\approx1+e^{2\beta-2\beta^2/d}=1+e^{2\beta(1-\beta/d)}$.
**Step 6: the capacity grows with $\beta$.** Dividing the defining equation, $m^\star-1=1/\mathbb E\big[e^{2\beta(s-1)}\big]$. Since $s\le1$, the factor $e^{2\beta(s-1)}$ falls pointwise as $\beta$ rises, so its mean falls and $m^\star$ rises. At $\beta=0$ the mean is 1 and $m^\star=2$. The Gaussian form peaks at $\beta=d/2$ and then falls, because it ignores that $s\le1$.
**Step 7: what the capacity at large $\beta$ measures.** It is the point where the mean squared ratio is 1, and at large $\beta$ that mean is dominated by rare close neighbours. At $d=16$ and $\beta=8$ the $2.0\%$ of keys with overlap above $0.5$ carry $84\%$ of $\mathbb E\big[e^{2\beta(s-1)}\big]$, and the $0.5\%$ above $0.6$ carry $62\%$. So the large-$\beta$ capacities are statements about the mean, a heavy-tailed quantity, and not about a typical draw (no median check was run at $\beta=8$). At $\beta=4$ the check below finds the mean not driven by a few rare draws.
**Check.** At $d=16$ the exact capacity is $34.5$, $483$, $19{,}591$ and $244{,}492$ for $\beta=2,4,8,12$, against the Gaussian form $34.1$, $404$, $2982$ and $404$. At $d=64$ and $\beta=4,8$ the exact values are $1816$ and $1.27$ million against $1809$ and $1.20$ million.
A sharper read moves the capacity by orders of magnitude. Interference over signal against the number of stored pairs in 16 dimensions (log axes); each curve grows like the square root of the pairs and its capacity is where it crosses 1. Hebbian crosses at 17. Drag the softmax dot along the line at 1 to change the inverse temperature; the Gaussian closed form is trustworthy only while it is small against the dimension.
**Why not make $\beta$ huge?** Capacity grows with $\beta$, but so does the cost of learning. The gradient of a weight with respect to a score is $\partial a_i/\partial s_l=\beta\,a_i(\delta_{il}-a_l)$, and as one weight approaches 1 this tends to 0 for every score. A saturated softmax passes almost no gradient to the keys and queries that produced it, so $\beta$ trades capacity against trainability. What value trained models use is not addressed here.
**The softmax is a kernel smoother, for unit keys.** Since $\lVert q-k\rVert^2=\lVert q\rVert^2+\lVert k\rVert^2-2q\cdot k$, the weight $e^{\beta q\cdot k}$ equals $e^{\beta(\lVert q\rVert^2+\lVert k\rVert^2)/2}\,e^{-\beta\lVert q-k\rVert^2/2}$. If every key has the same norm, the first factor is the same for all keys and cancels in the normaliser, and the softmax read is exactly the Nadaraya-Watson kernel smoother with a Gaussian kernel of bandwidth $1/\sqrt\beta$. With keys of different norms the equality fails. This is the sense in which attention averages the values of nearby keys.
**Hopfield networks.** A modern continuous Hopfield network has the energy $E(\xi)=-\tfrac1\beta\ln\sum_ie^{\beta x_i\cdot\xi}+\tfrac12\xi\cdot\xi$ for stored patterns $x_i$. Its gradient is $\xi-\sum_ia_ix_i$ with $a=\operatorname{softmax}(\beta X\xi)$, so a stationary point satisfies $\xi=X^\top\operatorname{softmax}(\beta X\xi)$: one softmax read of the stored patterns with query $\xi$, with keys and values both equal to the patterns. Ramsauer et al. (2020) define this network and show that its update rule is equivalent to transformer attention. The equivalence is exact in this algebra, up to the value projection that a transformer applies after the read. Whether the attention layers of a trained model do associative denoising of stored patterns is a separate question that the equivalence does not answer.
Derivation: the kernel smoother, and the Hopfield energy gradient
The claim is that for unit-norm keys softmax attention equals the Gaussian Nadaraya-Watson smoother, and that $\nabla E(\xi)=\xi-X^\top\operatorname{softmax}(\beta X\xi)$.
**Step 1: expand the squared distance.** $\lVert q-k_i\rVert^2=\lVert q\rVert^2+\lVert k_i\rVert^2-2\,q\cdot k_i$, so
$$
e^{\beta q\cdot k_i}=e^{\beta\lVert q\rVert^2/2}\;e^{\beta\lVert k_i\rVert^2/2}\;e^{-\beta\lVert q-k_i\rVert^2/2}.
$$
**Step 2: cancel the common factors.** If $\lVert k_i\rVert=1$ for all $i$, the first two factors do not depend on $i$. They multiply every numerator term and every denominator term of $o=\sum_ie^{\beta q\cdot k_i}v_i\big/\sum_je^{\beta q\cdot k_j}$, and cancel. What remains is the smoother with weights $e^{-\beta\lVert q-k_i\rVert^2/2}$.
**Step 3: the energy gradient.** Differentiate $E$ term by term. The first term is $-\tfrac1\beta\ln Z$ with $Z=\sum_ie^{\beta x_i\cdot\xi}$, and $\nabla_\xi Z=\beta\sum_ie^{\beta x_i\cdot\xi}x_i$, so
$$
\nabla\Big(-\tfrac1\beta\ln Z\Big)=-\frac{1}{\beta Z}\nabla Z=-\sum_i\frac{e^{\beta x_i\cdot\xi}}{Z}\,x_i=-X^\top a .
$$
The second term has gradient $\xi$. Together $\nabla E=\xi-X^\top a$.
**Step 4: the fixed point.** $\nabla E=0$ is $\xi=X^\top\operatorname{softmax}(\beta X\xi)$, which is the softmax read with query $\xi$ and keys and values both equal to the rows of $X$.
**Check.** For 20 random unit keys in $\mathbb R^6$ and $\beta=3$ the two sides of Step 2 agree to $10^{-16}$. With the keys scaled by factors between $0.7$ and $1.24$ the reads differ by $0.11$ on a read of norm $0.19$. A central finite difference of $E$ agrees with $\xi-X^\top a$ to $10^{-8}$.
**What the softmax read still cannot do.** It separates the pairs it keeps, and the price is that it keeps them all. Its capacity is large only because its state is the whole cache, which grows with the sentence. Everything in D1 to D5 is a different trade between the three things a memory can spend: numbers held, work per token, and how sharply it can read.
### D6. Two clocks: what is written when, and what each memory forgets
**The prompt.** `Flo gave Ann a pen ... then Ann thanked` `Flo`. The model that answers it holds the pair (first name, Flo) in the fast memory, and the rule "after `thanked`, ask for the first name" in the slow memory. Which one is written when?
The comparison of the two stores by keys, values, size and cost is the table of D3. What the clocks add:
| | Slow weights | Fast weights |
| --- | --- | --- |
| Time scale | the training run (millions of sentences) | one sentence |
| Write rule | delta rule on (input, target) pairs (D1) | Hebbian or delta rule on (key, value), or a cache (D3, D4) |
| Forgets | later training overwrites earlier along shared directions (same interference law as D1) | the end of the sentence; interference when full; decay by design |
| Holds in ToyLang | the template: names follow `gave`, fillers come from six words, the final name copies an earlier one | this sentence's names |
**In-weights and in-context.** Part A named two kinds of information: the kind learned once from many sentences, and the kind written in this sentence and gone after it. The slow weights hold the first kind and the fast memory holds the second. The division of labour in the Flo sentence is clean. The slow weights hold a program: write each name under a key that says which role it plays, and when you see `thanked`, read the key of the first name. The fast memory holds the data the program wrote: Flo, and Ann.
**What the picture does not explain.** It does not say what a trained model actually writes under which key; the ToyLang experiments use a fixed program (random position keys, and in D4 a hand-set gate). It does not say how the slow weights learn a program that works; for that see the induction heads of Part E. And it does not say how to choose between the designs: the trade between numbers held, work per token and sharpness of the read is the one thing every long-sequence architecture negotiates.
**What Part E does with this.** The slow weights are a program and the fast memory is its data. Part E makes that precise by writing a program by hand: two attention heads, with weights set by hand, that solve ToyLang for every gap length, and then asks what a trained model can and cannot do with the same instruction set.
## E. A small computer: reading, writing and running programs
**The question.** Can the forward pass of a transformer be read as a small computer that runs a program, and what is the smallest program we can write for it by hand? ToyLang is the test. With $L=30$ fillers the second Y sits 34 tokens after the first, and the final X sits 38 tokens after its source. Whatever predicts them exactly, with a fixed number of parameters, has to hold two names across the gap and fetch them on cue. This part assumes the earlier ones: ToyLang and the split between in-weights and in-context information (Part A), the fixed-size state and its bottleneck (Part B), attention as a latent pointer, positions, heads and the residual stream (Part C), and fast weights, linear attention and the delta rule (Part D). It builds one machine by hand (E2), shows a second program that learns from examples inside the forward pass (E3), names the pattern that trained models seem to share with the first machine (E4), and closes the whole note with one table, the places where the picture breaks, open questions, exercises, the answer to the question the note started with, and a reading list (E5).
### E1. What would it mean for a forward pass to run a program?
**A prompt with numbers in it.** Take a ToyLang sentence with $L=30$ and $N=48$ names. A bag of words pays $\ln\!\big((10+L)N/4\big)=\ln 480=6.17$ nats on each copied name. An n-gram model that sees fewer than 38 previous tokens pays $\ln(N-1)=3.85$ nats on the final X, and a table that answers it exactly needs $N(N-1)\cdot 3\cdot 6^{30}\approx1.5\times10^{27}$ contexts. A recurrent net must carry both names in a state of fixed size (at least $N(N-1)=2{,}256$ distinct situations here, Part B). The transformer keeps one key and one value per position and looks up what it needs. The question of this chapter is what, precisely, we mean when we say that this is a computer.
**The vocabulary.** The reading is a dictionary between three things: the parts of a computer, the parts of the forward pass, and the parts of ToyLang.
| In a computer | In the forward pass | In ToyLang |
| --- | --- | --- |
| memory | the context: one key and one value per earlier position (the KV cache of C5) | the names written at positions 0 and 2 |
| registers | named subspaces of the residual stream, one set per position | where a token, its left neighbour and a fetched name are kept |
| instruction | one head: a read addressed by content, then a write into the residual stream | "fetch the left neighbour", "fetch the name after gave" |
| steps (depth) | layers; one step per layer, and more steps only by looping | two steps |
| program | the weights ($W_Q,W_K,W_V,W_O$ and the MLPs) | the explicit matrices of E2 |
| data | the tokens, and the positions they sit at | the sentence |
The memory is not a fast-weight matrix here but its softmax counterpart, the set of all pairs (D3). Each head's write goes into the residual stream by addition, so a later layer can read what an earlier one wrote (C4). That is what makes depth a sequence of instructions: layer 1 can leave a note at every position, and layer 2 can read the notes.
**What is established.** Each item below is stated only as far as the abstract of its paper states it.
- Weiss, Goldberg and Yahav propose RASP, a programming language for transformer-encoder computation: a way to state, in program form, what a stack of attention layers computes.
- Giannou et al. build a 13-layer looped transformer that emulates a small instruction-set computer. It is a construction: the weights are written by hand, not learned.
- Pérez et al. show that the transformer is Turing complete without external memory. The result is for hard attention and unbounded precision, which is far from the softmax and the floating-point numbers of an actual model.
- Merrill and Sabharwal show the opposite side: transformers with logarithmic precision can be simulated by constant-depth threshold circuits, so under a standard complexity assumption (L is not P) they cannot solve linear equalities or decide membership for an arbitrary context-free grammar.
- Dehghani et al. add recurrence in depth with a per-position halting rule (Universal Transformers): the way to give the program more steps than the number of layers is to reuse a layer.
**What the reading buys, and what it does not.** The reading gives us a way to write down a program and to check, by running it, that it does what we claim. That is what E2 does. It is also a way to ask a question about a trained model: which instruction does this head perform? E4 shows one such question that has an answer. What it does not give is a reason to expect that training finds programs like the ones we write. A hand-set construction shows that the architecture can express a computation. It does not show that gradient descent will find it, or that a given model uses it. The two are kept apart wherever they could be confused. The vocabulary above is an analogy that becomes exact only for the explicit machines built in this part.
**Where this leaves us.** The reading says what the parts are. It does not say whether two layers are enough for ToyLang, how sharp attention has to be, or what breaks when we take a part away. Those are questions about one concrete machine, so we build one.
### E2. Build one: a two-layer program that solves ToyLang for every gap
**A prompt with a distractor.** Here is a sentence with a clause in front of the ToyLang sentence, and a gap of four fillers:
`Ann met Ben . Cy gave Di a pen and so and so then`: who comes next? The rule is: after then, the name that followed gave; after thanked, the name that came before gave. The answer after then is Di, and after the further tokens `Di thanked` it is Cy. The machine we build gets this right for every gap length and every one of up to 48 names, and it has 479 nonzero numbers in its weights.
**The program, in two instructions.** Write $w_t$ for the token at position $t$.
1. *Layer 1, at every position $t$:* copy the token to the left into a register called PREV. After this step position $t$ holds both $w_t$ and $w_{t-1}$.
2. *Layer 2, at then and at thanked:* find the position whose registers say "gave" and copy a name from it into a register called OUT. At then, the position is the one whose PREV is gave, and the name is its own token (the name right after gave). At thanked, the position is the one whose own token is gave, and the name is its PREV (the name right before gave).
The prediction is read from OUT. This is a pointer composed with a read in the sense of C1: the address of the name we want is defined by a property of its *neighbour*, and a single layer can only address a position by that position's own contents. Layer 1 turns the neighbour relation into a property of the position, and layer 2 addresses by it.
Two instructions, four registers. Step through: layer 1 leaves a note at every position; layer 2 addresses by that note.
**The residual stream.** Put $\lvert\mathcal V\rvert=N+15$ for the size of the vocabulary (the $N$ names, gave, a, three things, six fillers, then, thanked, the full stop and met) and $n_{\max}=64$ for the largest number of positions. The vector at position $t$ is made of four blocks, and for $N=48$ it has $d=3\lvert\mathcal V\rvert+n_{\max}=253$ numbers.
| Block | Size | Holds | Written by |
| --- | --- | --- | --- |
| TOK | $\lvert\mathcal V\rvert=63$ | one-hot of the token $w_t$ | the embedding |
| POS | $n_{\max}=64$ | one-hot of the position $t$ | the embedding |
| PREV | $\lvert\mathcal V\rvert=63$ | scratch: what layer 1 puts here, the token to the left | layer 1 |
| OUT | $\lvert\mathcal V\rvert=63$ | scratch: what layer 2 puts here, the fetched name | layer 2 |
Vectors are columns, so a projection acts as $q_t=W_Qx_t$, and attention is $a_{t,j}=\sigma_j(q_t\cdot k_j)$ over $j\le t$, with output $\sum_j a_{t,j}v_j$ added to the residual stream (C4). Let $\mathrm{tok}_w$, $\mathrm{pos}_t$, $\mathrm{prev}_w$, $\mathrm{out}_w$ be the unit vectors of single coordinates of the four blocks, and $\mathrm{val}_w$, $\mathrm{key}_j$ the unit vectors of the value space (one coordinate per token) and of the key space of the layer-1 head (one coordinate per position).
**Layer 1: one head, the previous-token head.** The key of position $j$ is its position, the query of position $t$ points at $t-1$, and the value is the token:
$$
\begin{gathered}
W_Q^{(1)}=\beta\sum_{t=1}^{n_{\max}-1}\mathrm{key}_{t-1}\,\mathrm{pos}_t^\top,\qquad
W_K^{(1)}=\sum_{j=0}^{n_{\max}-1}\mathrm{key}_j\,\mathrm{pos}_j^\top,\\[4pt]
W_V^{(1)}=\sum_{w}\mathrm{val}_w\,\mathrm{tok}_w^\top,\qquad
W_O^{(1)}=\sum_{w}\mathrm{prev}_w\,\mathrm{val}_w^\top .
\end{gathered}
$$
The score of key $j$ for the query at $t$ is $\beta\,[\,j=t-1\,]$. As $\beta$ grows, attention concentrates on position $t-1$ and the head writes $\mathrm{tok}_{w_{t-1}}$ into PREV. Position 0 has no left neighbour, so its query is zero and it attends to itself. Each matrix is a sum of outer products of unit vectors, so it is literally a table of key-to-value associations, like the memories of Part D, with the associations set by hand.
**Layer 2: one head, a key with two numbers.** The key of position $j$ has two coordinates, $\mathrm{at}$ (its own token is gave) and $\mathrm{aft}$ (its left neighbour is gave). The query at then asks the second, and the query at thanked asks the first. The value is the name part of TOK plus the name part of PREV, and it is written to OUT.
$$
\begin{gathered}
W_K^{(2)}=\mathrm{at}\,\mathrm{tok}_{\text{gave}}^\top+\mathrm{aft}\,\mathrm{prev}_{\text{gave}}^\top,\\[4pt]
W_Q^{(2)}=\beta\big(\mathrm{at}\,\mathrm{tok}_{\text{thanked}}^\top+\mathrm{aft}\,\mathrm{tok}_{\text{then}}^\top\big),\\[4pt]
W_V^{(2)}=\sum_{i\in\text{names}}\mathrm{val}_i\,(\mathrm{tok}_i+\mathrm{prev}_i)^\top,\qquad
W_O^{(2)}=\sum_{w}\mathrm{out}_w\,\mathrm{val}_w^\top,
\end{gathered}
$$
where $\mathrm{at},\mathrm{aft}$ are the two unit vectors of the two-dimensional key space of the layer-2 head. A query from any other token is zero and asks nothing; those rows are unused. The unembedding reads OUT: $W_U=\kappa\sum_w\mathrm{lg}_w\,\mathrm{out}_w^\top$, with $\mathrm{lg}_w$ the logit coordinate of token $w$ and the readout scale $\kappa=10$. The probability of the next token is the softmax of the logits over the whole vocabulary. Counting nonzero entries: layer 1 has $63+64+63+63=253$ (query, key, value, output), layer 2 has $2+2+96+63=163$, and the unembedding has 63, for 479 in all.
The program runs as a change of registers. Rows are positions, columns registers of the residual stream. Drag the ring through three steps: at 0 only TOK is filled, at 1 layer 1 has put each token's left neighbour in PREV, at 2 layer 2 has put the two fetched names in OUT (rows framed in copper, with the probability of the name). Position 0 has no left neighbour, so it copies itself.
**What the attention looks like.** The next figure shows the attention matrices for the actual sentence, one at a time, with the fillers merged into one row and column so that a gap of thirty is as readable as a gap of two. Layer 1 is a band one step below the diagonal. In layer 2 only the two framed rows matter. The row at then puts its weight on the name right after gave; the row at thanked puts its weight on gave itself, because the name it needs is stored there, in PREV.
Layer 1 writes who is on the left; layer 2 uses it to find the binding, whatever the gap. Rows are the position asking, columns the position looked at, shade is attention weight (numbers from 0.15; the causal mask blocks the upper triangle; fillers are merged). In layer 2 the framed rows are the two questions. Drag the ring to change $L$; the buttons set β (12, 4, 2): at 12 the answer is right and sure at every $L$ (Ben 0.997, Ann 0.997 at $L=8$), at 2 the weights spread (0.324 and 0.505). The mean over 100 plain sentences at $L=20$ is 0.17 at β = 2 (Exercise 2).
**Exact in the hard-attention limit.** As $\beta\to\infty$ the output of the construction is exactly the one-hot vector of the right name at then and at thanked, for every $L$, every $N$ and any clause in front, as long as the sentence has at most $n_{\max}$ positions and gave occurs once.
Derivation: the two layers fetch the right names when attention is hard
The claim is that, with attention concentrated on the largest score, $\mathrm{OUT}$ at then equals the one-hot vector of the name after gave, and at thanked equals the one-hot vector of the name before gave. Let gave sit at position $p_\star$ (the only occurrence).
**Step 1: layer 1 copies the left neighbour.** For $t\ge1$ the score of key $j$ at query $t$ is
$$
\beta\,\mathrm{key}_{t-1}\cdot\mathrm{key}_j=\beta\,[\,j=t-1\,].
$$
The largest score is at $j=t-1$, alone, so the softmax puts all weight there. The value is $\mathrm{tok}_{w_{t-1}}$, and $W_O^{(1)}$ writes it into PREV. Hence
$$
\mathrm{PREV}_t=\mathrm{tok}_{w_{t-1}}\quad(t\ge1),\qquad \mathrm{PREV}_0=\mathrm{tok}_{w_0}.
$$
The second line is position 0 attending to itself, which does no harm below.
**Step 2: the keys of layer 2.** The key of position $j$ is
$$
k_j=\Big(\,[\,w_j=\text{gave}\,],\ \ [\,w_{j-1}=\text{gave}\,]\,\Big),
$$
because the first coordinate reads $\mathrm{TOK}_j$ and the second reads $\mathrm{PREV}_j$. Since gave occurs only at $p_\star$, we get $k_{p_\star}=(1,0)$, $k_{p_\star+1}=(0,1)$ and $k_j=(0,0)$ for every other $j$. (Position 0 has $\mathrm{PREV}_0=\mathrm{tok}_{w_0}$, which is a name, so $k_0=(0,0)$ as well.)
**Step 3: the query at then.** The query is $\beta\,(0,1)$, so the score of key $j$ is $\beta\,[\,j=p_\star+1\,]$, and all the weight goes to $j=p_\star+1$. The value there is the name part of TOK plus the name part of PREV at that position. The token at $p_\star+1$ is the name $Y$, so $\mathrm{tok}$ contributes $\mathrm{val}_Y$; the left neighbour is gave, which is not a name, so $\mathrm{prev}$ contributes nothing. After $W_O^{(2)}$:
$$
\mathrm{OUT}_{\text{then}}=\mathrm{out}_Y .
$$
**Step 4: the query at thanked.** The query is $\beta\,(1,0)$, the score is $\beta\,[\,j=p_\star\,]$, and all the weight goes to $p_\star$. The token at $p_\star$ is gave, not a name, so $\mathrm{tok}$ contributes nothing; its left neighbour is the name $X$, so $\mathrm{prev}$ contributes $\mathrm{val}_X$:
$$
\mathrm{OUT}_{\text{thanked}}=\mathrm{out}_X .
$$
**Step 5: the unembedding.** The logit vector is $\kappa\,\mathrm{lg}_Y$ at then, which gives the next-token probability
$$
\Pr(Y)=\frac{e^{\kappa}}{e^{\kappa}+\lvert\mathcal V\rvert-1}.
$$
**Check.** For $N=48$, $\kappa=10$, $\lvert\mathcal V\rvert=63$ this is $22026/(22026+62)=0.9972$. Nothing in the steps used $L$, $N$ or what comes before $p_\star$, which is why a clause in front changes nothing. $\square$
**At finite sharpness.** A real attention is soft. Let $n$ be the number of positions a query can see and put $\varepsilon=n\,e^{-\beta}$ and $\delta=n\,e^{-\beta(1-2\varepsilon)}$. If $\varepsilon\le\tfrac14$, the OUT block at a copy position differs from the one-hot answer by at most $\Delta=2\varepsilon+3\delta$ in total absolute value, and the right name wins the argmax whenever $\Delta<1$. The error is exponentially small in $\beta$, and the sharpness needed grows only like $\ln n$. The derivation uses that no name is adjacent to another name, which holds in ToyLang with or without the clause.
Derivation: the error at finite beta
The claim: if $\varepsilon=n e^{-\beta}\le\frac14$, then at a copy position $\lVert \mathrm{OUT}-\mathrm{out}_{\text{ans}}\rVert_1\le 2\varepsilon+3\delta$ with $\delta=n e^{-\beta(1-2\varepsilon)}$.
**Step 1: layer 1 is almost hard.** At position $t\ge1$ the score is $\beta$ at $j=t-1$ and $0$ at the other $t$ visible positions, so
$$
a_{t,t-1}=\frac{e^{\beta}}{e^{\beta}+t},\qquad \varepsilon_t:=1-a_{t,t-1}=\frac{t}{e^{\beta}+t}\le t\,e^{-\beta}\le\varepsilon .
$$
The register is a convex combination, so $\mathrm{PREV}_t=\mathrm{tok}_{w_{t-1}}+r_t$ with $\lVert r_t\rVert_1\le 2\varepsilon_t\le 2\varepsilon$ (mass $\varepsilon_t$ is removed from the right token and put on others). Position 0 is exact.
**Step 2: the keys of layer 2 move a little.** The first key coordinate reads TOK and is exact. The second reads $\mathrm{PREV}_j[\text{gave}]$, which is the weight that position $j$ gave to the position holding gave. It is $1-\varepsilon_j\ge1-\varepsilon$ at $j=p_\star+1$ and at most $\varepsilon_j\le\varepsilon$ at every other $j$.
**Step 3: layer 2 is almost hard.** At then the target score is at least $\beta(1-\varepsilon)$ and every other score is at most $\beta\varepsilon$. The total weight $\omega$ off the target is at most the number of other positions times the largest ratio:
$$
\omega\le (n-1)\,\frac{e^{\beta\varepsilon}}{e^{\beta(1-\varepsilon)}}\le n\,e^{-\beta(1-2\varepsilon)}=\delta .
$$
At thanked the keys are exact (the first coordinate), the target score is $\beta$ and the others are $0$, so $\omega\le n\,e^{-\beta}\le\delta$.
**Step 4: the output.** Write the value of position $j$ as $v_j=v_j^\star+\rho_j$, where $v_j^\star$ is the exact value (a one-hot name or zero) and $\lVert\rho_j\rVert_1\le2\varepsilon$ is the name part of $r_j$. The value $v_j$ is the name part of TOK plus the name part of PREV. Because no name is adjacent to another name, at most one of the two parts holds a name, so $\lVert v_j\rVert_1\le 1+2\varepsilon\le1.5$ for $j\ge1$. Position 0 is the exception: it attends to itself, so its PREV is its own token, a name, and its value holds that name twice, $\lVert v_0\rVert_1=2$. Hence $\lVert v_j\rVert_1\le2$ for every $j$. With target $t^\star$ and answer $e$ (a one-hot vector), the output $o=\sum_ja_jv_j$ satisfies
$$
o-e=(a_{t^\star}-1)\,e+a_{t^\star}\rho_{t^\star}+\sum_{j\ne t^\star}a_jv_j ,
$$
and therefore
$$
\Delta:=\lVert o-e\rVert_1\le\omega+2\varepsilon+2\omega\le 2\varepsilon+3\delta .
$$
(Without position 0 the middle term would be $1.5\,\omega$ and the constant 2.5; the stated $3\delta$ is the one that holds for every position.)
**Step 5: the argmax.** Call the bound $\Delta$. The coordinate of the right name is at least $1-\Delta_{\text{right}}$ and every other coordinate is at most $\Delta_{\text{other}}$, with $\Delta_{\text{right}}+\Delta_{\text{other}}\le \Delta$. If $\Delta<1$ then $1-\Delta_{\text{right}}>\Delta_{\text{other}}$, so the right name has the largest logit (the non-name coordinates are exactly zero). $\square$
**Check.** For $n=44$ ($L=30$ with the clause in front): at $\beta=12$ the bound is $\Delta=0.0014$; at $\beta=8$, $\Delta=0.0856$; at $\beta=6$, $\Delta=1.43$, which certifies nothing. The smallest $\beta$ for which the bound certifies the argmax at $n=44$ is 6.19. The value vectors of the layer-2 head have $\ell_1$ norm 2 at position 0 and at most 1.013 elsewhere (at $\beta=6$).
**Why two layers.** The tempting reading is that ToyLang needs depth. It does not, quite, and the numbers say so. On plain ToyLang a single layer solves both copies, by using the *template*: absolute positions (at then, attend to position 2; at thanked, to position 0) or recency (at then, the nearest name; at thanked, the farthest name). What a single layer cannot do is follow the *rule*, "the name next to gave", once the template moves. The distractor clause moves it, and then the one-layer machines fail while the two-layer program does not care.
Derivation: a single attention layer cannot address a position by its neighbour
The claim is about one attention head followed by a linear read-out, with hard attention: the head copies the value of the position with the largest score. The head gives position $j$ a score that depends only on the query's token and position and on the candidate's own token and position, $s(w_p,p;\,w_j,j)$, and a value $v(w_j,j)$ that depends only on the candidate's own token and position. Nothing in $s$ or $v$ can refer to $w_{j-1}$. The claim is: no such head answers correctly for **every assignment of names to the four name slots** in the two sentences below.
**Step 1: two sentences with the same names in the same places.** For any four distinct names $n_0,n_2,n_4,n_6$, take
$$
\begin{gathered}
\text{(a)}\quad n_0\ \text{met}\ n_2\ .\ n_4\ \text{gave}\ n_6\ \text{a pen and so then}\\
\text{(b)}\quad n_0\ \text{gave}\ n_2\ .\ n_4\ \text{met}\ n_6\ \text{a pen and so then}
\end{gathered}
$$
The names sit at positions 0, 2, 4 and 6 in both. The rule gives the answer $n_6$ for (a) and $n_2$ for (b). The verbs swap places, and nothing else changes.
**Step 2: the head must copy a name position.** The value of a verb position depends on the verb and the position only, so it is the same vector for every assignment of names, while the correct answer changes with the assignment. If the head copied a verb position in (a), its output could not equal $n_6$ for every assignment. So in (a) it copies a name position $j$, and since the value of a name position depends on $(w_j,j)$ and the four names are distinct, the output is $n_6$ only if $j=6$. By the same argument it must copy position 2 in (b).
**Step 3: the contradiction.** In (a) the head needs the score at position 6 to beat the score at position 2. In (b) it needs the reverse. But the score of a name position depends on the query token (then, position 11) and on the candidate's own $(w_j,j)$, and these are identical in (a) and (b) for positions 2 and 6. The same two numbers cannot be ordered both ways. So the head fails on (a) or on (b). $\square$
**What the lemma leaves open.** The quantifier over assignments is needed: for one fixed pair of sentences a value could hard-code a name. With $H$ heads, a verb can switch heads on and off, so a one-layer model can handle a finite list of layouts, roughly one head per place gave can sit. That is the shape of the cost: composition costs two heads for all layouts, enumeration costs a head per layout. This is a sketch, not a theorem: the single-head case with hard attention and a linear read-out is what is proved, and no matching lower bound for $H$ heads is shown here. A wide MLP after the layer is a different trade: the value vectors can be made to carry (name, position) pairs and the MLP can sort them out, at a width of order $N\cdot n_{\max}$. That is not analysed here.
**Taking parts away.** The figure runs the ablations at the chosen sharpness. Reading the numbers at $\beta=12$:
- *No position signal.* The POS block is zeroed, so the layer-1 head cannot tell positions apart and averages the whole left context into PREV, with weight $1/(t+1)$ on each token. The causal mask alone leaks a weak sense of order (position 1 has only the first token to its left), which happens to rescue one of the two copies: on plain ToyLang the final X is right 100 of 100 and the second Y 0 of 100; with the clause it is the other way round, 99 of 100 for the second Y and 1 of 100 for the final X.
- *No layer 1.* PREV is empty. At then the keys are all zero in the coordinate the query asks about, so attention is uniform and the output holds X and Y with equal weight, a tie that the argmax breaks by token index (56 of 100 right on plain ToyLang, a coin flip). At thanked the value, the PREV of the gave position, is empty and the output is zero (0 of 100).
- *Lower sharpness.* On plain ToyLang at $L=20$ with 48 names the program is right on all 200 copies from $\beta=1.5$ upward. At $\beta=1$ or lower the second Y is lost (0 of 100) and the final X survives. With the clause in front the same cliff sits higher, between $\beta=1.5$ and $\beta=2$ (numbers in the box above). The clause pushes every position four to the right, where layer 1 spreads its attention over more candidates (Step 1: the weight left on the neighbour is $e^{\beta}/(e^{\beta}+t)$). Even where the argmax is right, the probability of the answer falls as attention blurs (the maps figure, third button).
Each part of the program is needed, and the one-layer machines only look sufficient on the plain template. Each row is one machine, scored on 60 sentences at L = 20 with 48 names: the grey bar is plain ToyLang, the hatched bar is the same machine when a clause of four tokens is put in front. Drag the ring to change beta. Three rows move with beta (the program, no position with a clause, one layer with fixed slots); the other three are flat over this range, and the linear-attention version has no softmax and no beta.
**The same program with a fast-weight memory.** Replace the softmax by a plain dot product (no normalisation, scores of size 1 instead of $\beta$). Every score in the construction is exactly 0 or 1, so the linear attention returns exactly the value that hard attention returned, with no sharpness to tune. The layer-1 head is then the fast-weight memory of D3: its state after $t$ writes is the $\lvert\mathcal V\rvert\times n_{\max}$ matrix $S_t=\sum_{j\le t}\mathrm{tok}_{w_j}\,\mathrm{pos}_j^\top$, written with the pair (value = token, key = position), and the read $S_t\,\mathrm{pos}_{t-1}$ returns $\mathrm{tok}_{w_{t-1}}$ exactly because the position keys are orthogonal. The layer-2 state has two key columns and, for 48 names, 48 name rows read from TOK and PREV plus the output block (the 163 nonzero entries counted above). In the ablation numbers above the linear version is right on 200 of 200 copies both ways. Its price is the capacity limit of D1 and D3: the state has $63\times64=4032$ numbers and the position code has $n_{\max}=64$ distinct keys, so a sentence longer than 64 tokens has no address. Orthogonal keys avoid interference only while there are as many keys as positions.
**What the construction shows, and what it does not.** It shows that two attention layers, a two-dimensional key and a one-hot position code are enough to realise the rule "the name next to gave" for every gap length and every name, and that the rule survives a clause in front. It also shows which parts matter: the position signal, the layer-1 note, the sharpness. It does not show that a trained network learns this program, that it uses these registers if it learns a program of this kind, or that sentences far longer than those in training would still work (the one-hot position code has 64 slots; a sinusoidal code has no such limit, but the certified sharpness above grows with $\ln n$). The construction is in the same family as the hand-set constructions in the literature (E1), not a statement about what training finds.
**Where this leaves us.** The program has fixed weights and does one thing. A trained model also *learns* things from the prompt: it can fit a function from a few examples placed in the context. If the forward pass is a computer, what algorithm is it running when it does that?
### E3. Learning inside the forward pass: one layer, one gradient step
**A prompt with numbers in it.** Eight labelled points $(x_i,y_i)$ are in the context, drawn from a line with an unknown slope, and then a query $x_q$. The best prediction from the examples is the least-squares line through them. Can a stack of attention layers with fixed weights produce a prediction that improves layer by layer towards it? Yes, and what the stack runs is a gradient descent written into the forward pass.
**The task and the target.** The examples are $(x_i,y_i)_{i=1}^m$ with $y_i=w^\top x_i$ plus a little noise, and the query is $x_q$. A model $\hat y=\Theta x_q$ with regression weights $\Theta$ starts at $\Theta_0=0$ and is trained by gradient descent on the in-context loss $\frac1{2m}\sum_i(\Theta x_i-y_i)^2$ with step $\eta$. One step gives
$$
\begin{gathered}
\Theta_1=\Theta_0-\frac{\eta}{m}\sum_{i=1}^m(\Theta_0x_i-y_i)\,x_i^\top=\frac{\eta}{m}\sum_{i=1}^m y_i\,x_i^\top,\\[4pt]
\hat y=\Theta_1x_q=\frac{\eta}{m}\sum_{i=1}^m y_i\,(x_i\!\cdot x_q).
\end{gathered}
$$
Read the last expression as attention: the keys are the inputs $x_i$, the values are the labels $y_i$, the query is $x_q$, the score is the plain dot product with no softmax, and the output is the sum of values weighted by $\eta/m$. That is the linear-attention read of Part D, and the Hebbian write $\sum_i y_ix_i^\top$ of D1. One step of gradient descent on the examples is a linear attention over them.
Two routes, one answer. Reading the context with linear attention gives exactly the prediction of one gradient step on the examples.
**Three gradient steps, one family.** The Hebbian state of D3 is this first step from the zero state, taken on the linear objective $-y\,\Theta x$. The delta rule of D4 is a gradient step on the in-context error, one token at a time. The layer below is the batch version: a gradient step on the error over all examples at once.
**A layer that does it.** Give each example the token $z_j=(x_j,y_j)$ and the query the token $z_q=(x_q,0)$. One linear self-attention layer with the hand-set matrices $W_K=W_Q=\mathrm{diag}(1,\dots,1,0)$ (they read $x$ and ignore the label) and $W_V=\mathrm{diag}(0,\dots,0,-1)$ (it reads minus the label slot), with a write of size $\eta/m$, updates every token as
$$
\begin{gathered}
z_j\;\leftarrow\;z_j+\frac{\eta}{m}\sum_{i=1}^{m}(W_Vz_i)\,\big(W_Kz_i\cdot W_Qz_j\big),
\end{gathered}
$$
where the sum runs over the example tokens only. For the label slot this reads $y_j\leftarrow y_j-\frac{\eta}{m}\sum_i y_i\,(x_i\cdot x_j)=y_j-\Theta_1x_j$: every example's label becomes its residual after one step, and the query's slot becomes $-\Theta_1x_q$. Stacking $\ell$ identical layers makes the label slots the residuals of $\Theta_\ell$.
The restriction of the sum to the example tokens is part of the construction, not ordinary causal attention: the examples attend to each other in both directions and the query token is not a key. Once the query is allowed to be a key, its label slot is no longer zero after the first layer ($-\Theta_1x_q$), it writes a nonzero value into every token in the second layer, and the stack stops following gradient descent.
Derivation: $\ell$ linear-attention layers equal $\ell$ gradient steps
The claim is that after $\ell$ layers the label slot of token $j$ holds $y_j-\Theta_\ell x_j$ (with $y_q=0$), where $\Theta_\ell=\Theta_{\ell-1}-\frac{\eta}{m}\sum_i(\Theta_{\ell-1}x_i-y_i)x_i^\top$ and $\Theta_0=0$.
**Step 1: the base case.** Before any layer the label slot of token $j$ is $y_j=y_j-\Theta_0x_j$, since $\Theta_0=0$.
**Step 2: what a layer does to a label slot.** Suppose the slots hold the residuals $r_i^{(\ell-1)}=y_i-\Theta_{\ell-1}x_i$. The value $W_Vz_i$ is $-r_i^{(\ell-1)}$ in the label slot, and the score is $x_i\cdot x_j$. The slot of token $j$ becomes
$$
r_j^{(\ell)}=r_j^{(\ell-1)}-\frac{\eta}{m}\sum_{i=1}^m r_i^{(\ell-1)}\,(x_i\cdot x_j).
$$
**Step 3: define the new weights.** Let
$$
\Theta_\ell=\Theta_{\ell-1}+\frac{\eta}{m}\sum_{i=1}^m r_i^{(\ell-1)}\,x_i^\top .
$$
This is the gradient step, because $r_i^{(\ell-1)}=-(\Theta_{\ell-1}x_i-y_i)$.
**Step 4: the slot is the new residual.** Using the definition of $\Theta_\ell$,
$$
y_j-\Theta_\ell x_j=r_j^{(\ell-1)}-\frac{\eta}{m}\sum_i r_i^{(\ell-1)}\,(x_i\cdot x_j)=r_j^{(\ell)} .
$$
The induction closes. For the query token, $y_q=0$, so its slot after $\ell$ layers is $-\Theta_\ell x_q$ and the prediction is minus that slot.
**Step 5: a closed form.** Let $\Sigma=\frac1m\sum_ix_ix_i^\top$ and $g=\frac1m\sum_iy_ix_i$. The recursion reads $\Theta_\ell=\Theta_{\ell-1}+\eta\,(g-\Sigma\Theta_{\ell-1})$, that is $\Theta_\ell-\Theta_\star=(I-\eta\Sigma)(\Theta_{\ell-1}-\Theta_\star)$ with $\Theta_\star=\Sigma^{-1}g$ the least-squares solution. Hence
$$
\Theta_\ell=\big(I-(I-\eta\Sigma)^\ell\big)\,\Sigma^{-1}g ,
$$
which tends to $\Theta_\star$ when all eigenvalues of $\eta\Sigma$ lie in $(0,2)$.
**Check.** One dimension, examples $(1,1),(2,2),(1,3)$, $\eta=0.3$: $\Sigma=2$, $g=8/3$, so $1-\eta\Sigma=0.4$ and $\Theta_1=(1-0.4)\cdot\frac43=0.8$, $\Theta_2=(1-0.16)\cdot\frac43=1.12$, tending to $\Theta_\star=4/3$. $\square$
Each layer is one gradient step, and the stack walks to the least-squares line. The blue dots are eight examples in the context and the copper dotted line marks the query. The violet line is the model after the chosen number of linear-attention layers and the dashed ink line is the least-squares fit. Drag the ring up the right end of the violet line: the slope climbs from 0 towards 0.97, and the loss on the examples falls towards its least-squares floor on the right.
For the eight examples of the figure, one layer gives slope 0.215, three layers 0.512, ten layers 0.890, against 0.970 for least squares; the loss on the examples falls from 0.210 to 0.128, 0.048 and 0.0033, against a floor of 0.0019.
**What is known.** von Oswald et al. (abstract only) construct exactly this for self-attention-only transformers on simple regression tasks: a linear self-attention layer with hand-set weights that performs a gradient-descent step on a regression loss, and they describe the trained models they study as mesa-optimizers. Garg et al. train transformers from scratch to learn simple function classes in context and find accuracy close to least squares on linear functions. Akyürek et al. show that a transformer can implement gradient descent and ridge regression and that trained models behave like those estimators on linear problems.
**What is not known.** The scope of the construction is narrow: self-attention-only models, simple regression, a hand-set linear attention without softmax, a sum over the example tokens only, and a token format with the label next to its input. A single layer with softmax matches gradient descent less well than a linear one. Nothing here shows that the in-context learning of a large language model is gradient descent. The defensible summary is: a stack of attention layers can run gradient descent on the examples in its context, trained small models on linear regression behave as if they did, and for language models the question is open.
**Where this leaves us.** E3 is an algorithm run by a stack of layers with identical weights, on tokens prepared in a convenient format. The natural question is which small programs trained models contain. One of them is the circuit of E2.
### E4. Induction heads: the same two-head composition, found in trained models
**A prompt with numbers in it.** Read `Ann gave Ben a pen` once. When the sentence starts again, `Ann gave Ben a pen Ann ...`, the next token is "gave": it is the token that came after the last "Ann". Do that for every token of the second copy and a model makes no mistakes on a sequence it has never seen. That rule, find the earlier occurrence of the current token and copy what came after it, is the *induction* rule.
**The circuit.** It is E2 with the anchor gave replaced by the current token. Layer 1 is the same previous-token head, and it is what puts each token's left neighbour beside it: after layer 1, position $t$ holds both $w_t$ and $w_{t-1}$, which is the kind of note that E3 needs to find a label next to its input. In layer 2 the query is the current token $w_p$, the key of position $j$ is its PREV (the token to its left), and the value is its own token:
$$
\begin{gathered}
W_K^{(2)}=\sum_{w}\mathrm{key}_w\,\mathrm{prev}_w^\top,\qquad
W_Q^{(2)}=\beta\sum_{w}\mathrm{key}_w\,\mathrm{tok}_w^\top,\\[4pt]
W_V^{(2)}=\sum_{w}\mathrm{val}_w\,\mathrm{tok}_w^\top .
\end{gathered}
$$
The score of position $j$ is $\beta\,[\,w_{j-1}=w_p\,]$, so the head attends to the positions that *followed* an earlier occurrence of the current token and copies their token. If the earlier tokens are distinct, there is exactly one such position and the output is exact. Position 0 attends to itself and would match its own token, so the sequences start with a full stop that never occurs again. The proof is Step 1 to Step 3 of the E2 derivation with "gave" replaced by "$w_p$".
The second copy is predicted by finding what followed the first. The grey row under each token is what layer 1 wrote into PREV. Drag the ring along the second copy: the copper arc runs from the query to the position whose PREV equals the query token, with weight 1.000, and the program gives the token there probability 0.997.
**How E2 is an induction circuit.** The second-Y copy in E2 is induction with a fixed anchor: the query is always the verb gave (set by the then token), the key is PREV, the value is the token. The final-X copy is its mirror: the key is the token itself and the value is PREV. In the vocabulary of the circuits literature the first uses the layer-1 output through the key of the second head and the second uses it through the value; the names for these are K-composition and V-composition (Elhage et al.). Elhage et al. reverse-engineer small attention-only models and introduce induction heads as heads formed by composing two heads across layers.
**What the trained models show.** Olsson et al. describe an induction head as one that completes $[A][B]\dots[A]\to[B]$ by combining a head that copies the previous token with a head that finds earlier occurrences of the current token and attends to what followed. They report a phase change in training in which induction heads form at the same point as a jump in in-context learning. The authors call the evidence preliminary: it is causal for small attention-only models and correlational for larger ones. The circuit above serves as a definition of what such a head computes, and as an existence proof that two layers are enough for it. Whether the heads found in trained models use the registers of E2 is not tested here. Whether induction is a special case of the gradient-descent mechanism of E3 or a different mechanism is not settled by the cited abstracts.
**Where this leaves us.** Two layers, one previous-token note and one content-addressed read give the copy rule for any sequence, and they are the part of the picture that appears both in a hand construction and in trained small models. What remains is to put everything on one page and to say where it fails.
### E5. The whole picture, where it breaks, and what is open
**One table.** Each row is a model class of this note, read along the same axes: what it remembers, how it writes, how it reads and what it forgets, then what it costs and what it achieves on ToyLang.
| Model | Remembers | Writes | Reads | Forgets |
| --- | --- | --- | --- | --- |
| bag of words | counts of tokens | add a count | the count | the order |
| n-gram | counts per $r$-token context | add a count | the count for the last $r$ tokens | anything older than $r$ |
| RNN | a state of fixed size | overwrite through the gates | the state | what does not fit in the state |
| linear attention | a $d_v\times d_k$ matrix of fast weights | add $v_tk_t^\top$ (or correct) | $S_tq_t$ | interference beyond about $d_k$ pairs |
| softmax attention | every key and value | append | softmax-weighted average | nothing, but it blurs when many keys compete |
| transformer block, stacked | the cache, plus the residual stream | each head writes into the stream, the MLP writes per token | heads read the cache and the earlier writes | nothing in the cache; the stream is overwritten by later writes |
| Model | Cost | ToyLang: nats on a copied name, or whether it is exact |
| --- | --- | --- |
| bag of words | a table of size $\lvert\mathcal V\rvert$ | $\ln\frac{(10+L)N}{4}$ for each copied name: 3.69 at $L=10$, $N=8$ |
| n-gram | contexts grow like $F^L$ | $\ln(N-1)$ on the final X for $r
At $L=10$ and $N=8$ the bag of words pays 3.69 nats on a copied name and an n-gram with $r<18$ pays 1.95 on the final X. At $L=30$ and $N=48$ the bag of words pays 6.17, the n-gram pays 3.85 on the final X below $r=38$, and the table needs $1.5\times10^{27}$ contexts.
**Where the picture breaks.**
- **Positions.** The construction of E2 stops at $n_{\max}=64$ tokens, because the position code has 64 slots. A code with unbounded range (sinusoidal, as in the original transformer, or rotary) removes that limit but not the next one.
- **Sharpness against length.** The guarantee in E2 needs $\beta$ to grow like $\ln n$: 6.19 at $n=44$, 9.54 at $n=1000$. A fixed temperature stops working past some length; attention that was sharp on short contexts is diffuse on long ones, which is the temperature and capacity trade-off of Part D.
- **Length generalisation.** Trained transformers are widely reported to do worse on sequences longer than those seen in training. This is not measured here and no cause is traced.
- **Cost.** Softmax attention does $O(n)$ work for each new token, $O(n^2)$ over the sequence, and its cache grows with $n$ (C5). Linear attention removes that and pays in capacity (D3).
- **Expressivity.** A fixed-depth transformer with logarithmic precision sits inside constant-depth threshold circuits (Merrill and Sabharwal), so the depth of E1 is a real limit: problems that need many sequential steps need many layers or loops (Dehghani, Giannou).
**Exercises.** Each answer is folded.
**1.** How many nonzero numbers does the layer-2 head of the E2 program hold for $N=48$, and how many does the whole program hold?
Worked example: counting the program
The query matrix has 2 nonzero entries (the rows for thanked and then), the key matrix has 2 (the TOK and PREV coordinates of gave), the value matrix has $2\times48=96$ (each of the 48 names is read from TOK and from PREV) and the output matrix has 63. That is $2+2+96+63=163$. Layer 1 has $63+64+63+63=253$ (query, key, value, output) and the unembedding 63. In all $253+163+63=479$.
**2.** What is the smallest sharpness $\beta$ for which the bound of E2 certifies the argmax at $n=100$, and why is it so much larger than the sharpness at which the program is already right on every plain ToyLang sentence?
Worked example: certified versus observed sharpness
Solving $\varepsilon\le\frac14$ and $2\varepsilon+3\delta<1$ at $n=100$ gives $\beta\ge7.08$. On plain ToyLang the program is right on every seeded sentence from $\beta=1.5$. The bound adds up the weight on every other position (up to $n-1$ of them) and assumes it all lands on the wrong name. In ToyLang only a handful of positions carry a name, and the others contribute zero to every name coordinate, so the argmax among names is far safer than the sum of the weights suggests. The probability of the answer, though, does depend on how much weight leaks: averaged over 100 plain sentences at $L=20$ it is 0.17 at $\beta=2$ and 0.56 at $\beta=3$, against 0.997 in the hard limit.
**3.** Three in-context examples in one dimension, $(1,1),(2,2),(1,3)$, and $\eta=0.3$. What does one linear-attention layer predict for $x_q=1$, what do two predict, and what is the least-squares answer?
Worked example: one and two gradient steps
With $m=3$ the first layer gives $\hat y_1=\frac{\eta}{m}\sum_iy_ix_ix_q=0.1\,(1+4+3)=0.8$. The residuals are then $1-0.8=0.2$, $2-1.6=0.4$ and $3-0.8=2.2$, and the second layer adds $0.1\,(0.2\cdot1+0.4\cdot2+2.2\cdot1)=0.32$, so $\hat y_2=1.12$. Least squares gives $8/6=1.333$. The closed form $\Theta_\ell=(1-0.4^\ell)\cdot\frac43$ agrees: $0.8$, $1.12$, $1.248$ for $\ell=1,2,3$.
**4.** In the sequence `. Ann gave Ann pen Ann`, why does the induction circuit give gave and pen a probability of 0.415 each, and not 0.5? What changes if the leading full stop is dropped?
Worked example: the split
The two earlier occurrences of Ann were followed by gave and by pen, so layer 2 splits its weight $\frac12,\frac12$ and OUT holds $\frac12$ on each. The logits are $\kappa/2=5$ for each of the two tokens and 0 for the other 61, so each gets $e^{5}/(2e^{5}+61)=0.415$. The remaining 0.17 is spread over the other tokens, which have the same logit 0: a probability is spread over the whole vocabulary, not over the two candidates. Without the full stop, position 0 holds Ann with Ann as its own PREV (it attended to itself), so it matches the query too; the weight splits three ways, the logits are $\kappa/3$ for Ann, gave and pen, and each of the three gets about 0.19 (the exact values are 0.195, 0.195 and 0.194).
**5.** The position code of E2 has 64 slots. What breaks for a sentence of 70 tokens, and what does the fix cost?
Worked example: past the position code
Positions 64 and above have no unit vector, so the shift matrix has no entry for them and layer 1 would copy nothing there. Widening the code to 128 slots fixes it for sequences up to 128 and costs $n_{\max}$ more key dimensions in layer 1 and a larger one-hot block. A sinusoidal code removes the limit, at the price that nearby positions are no longer orthogonal and the sharpness has to grow with $n$ (the certified $\beta$ is 7.08 at $n=100$ and 9.54 at $n=1000$).
**Back to the first sentence.** The note began with a sentence a person completes without effort.
> Flo gave Ann a pen and yes well then Ann thanked ___
Run it through the program of E2 and the machine writes Flo with probability 0.997 (and Ann after then with the same probability), using 479 nonzero numbers and no knowledge of these two names. The one question the person asked, *who gave the pen?*, is the query of the layer-2 head. Against the four requirements of A5, the stacked transformer reads as follows.
| Requirement | The stacked transformer | Where it is shown |
| --- | --- | --- |
| R1. Order matters | yes: the position code (the POS block) and the causal mask break the symmetry of a set; with the code zeroed one of the two copies is lost (100 of 200 right) | C2, E2 |
| R2. Reach back far | any distance inside the position code: both copies are right for every gap $L\le30$ with no change to the weights; the price is $O(n^2)$ time over the sequence, an $O(n)$ cache per token, and a sharpness $\beta$ that grows like $\ln n$ | E2, C5 |
| R3. Look up a binding written in this very text | yes: the pair is written into the cache (one key and one value per position) and into the PREV register of layer 1, and read by content; right on all 13,536 predictions over every ordered pair of 48 names at $L=0,17,30$ | E2 |
| R4. Train in parallel | yes: the loss at every position comes from one pass over the text, with no step waiting for the previous one | C5 |
**What had to be carried across the gap, and where it is written.** What had to cross the gap is the pair of names: at least $\ln 56=4.03$ nats for $N=8$ (the 56 situations of A4 and B1, and $N(N-1)=2{,}256$ for 48 names), which a fixed-size state has to keep, and keep for the whole gap. The transformer keeps nothing in a state: it writes the names into the cache when it sees them (Flo at position 0, Ann at position 2) and into the layer-1 register PREV (Flo, on the position that holds gave), and at thanked a content-addressed read finds the gave position and copies Flo through OUT, so the memory grows with the sentence instead of staying at 56 states.
**Reading list.** Entries with a link are those whose abstract or text states the claim cited above; the others are listed for further reading and were not read for this note.
- Vaswani et al. 2017, *Attention Is All You Need*, [arXiv:1706.03762](https://arxiv.org/abs/1706.03762).
- Weiss, Goldberg, Yahav 2021, *Thinking Like Transformers* (RASP), [arXiv:2106.06981](https://arxiv.org/abs/2106.06981).
- Pérez, Marinković, Barceló 2019, *On the Turing Completeness of Modern Neural Network Architectures*, [arXiv:1901.03429](https://arxiv.org/abs/1901.03429).
- Giannou et al. 2023, *Looped Transformers as Programmable Computers*, [arXiv:2301.13196](https://arxiv.org/abs/2301.13196).
- Dehghani et al. 2019, *Universal Transformers*, [arXiv:1807.03819](https://arxiv.org/abs/1807.03819).
- Merrill, Sabharwal 2023, *The Parallelism Tradeoff: Limitations of Log-Precision Transformers*, [arXiv:2207.00729](https://arxiv.org/abs/2207.00729).
- von Oswald et al. 2023, *Transformers learn in-context by gradient descent*, [arXiv:2212.07677](https://arxiv.org/abs/2212.07677) (abstract).
- Garg, Tsipras, Liang, Valiant 2022, *What Can Transformers Learn In-Context?*, [arXiv:2208.01066](https://arxiv.org/abs/2208.01066).
- Akyürek et al. 2023, *What learning algorithm is in-context learning?*, [arXiv:2211.15661](https://arxiv.org/abs/2211.15661).
- Elhage et al. 2021, *A Mathematical Framework for Transformer Circuits*, Transformer Circuits Thread.
- Olsson et al. 2022, *In-context Learning and Induction Heads*, [arXiv:2209.11895](https://arxiv.org/abs/2209.11895).
- Katharopoulos et al. 2020, *Transformers are RNNs*, [arXiv:2006.16236](https://arxiv.org/abs/2006.16236); Schlag, Irie, Schmidhuber 2021, *Linear Transformers Are Secretly Fast Weight Programmers*, [arXiv:2102.11174](https://arxiv.org/abs/2102.11174).
- Schmidhuber 1992, *Learning to control fast-weight memories*, Neural Computation 4(1), 131-139.
- Widrow, Hoff 1960, *Adaptive switching circuits*; Kohonen 1972 and Anderson 1972 on correlation matrix memories.