Three matrices, and a triangle of minus infinity

Last time attention had no parameters in it at all. Six tokens, three dimensions each, inputs @ inputs.T, a row softmax, a weighted sum: 36 dot products, and the context vector for journey came out at [0.4419, 0.6515, 0.5683]. It worked, and it had barely done anything. The six inputs sat 0.6576 apart on average and the six outputs 0.0482, which is 7.3% of the spread they came from, with all six landing nearest the same embedding, starts.

It also could not learn. The score matrix was $XX^{\top}$, symmetric by necessity, so the raw affinity between journey and starts was one number, 1.4754, read the same way from either end. There was no dial anywhere for gradient descent to turn.

This post adds three matrices. Every one of them is 3 by 2.

All six figures here compute everything they show, in your browser, and there are no model weights of any kind on this page. One warning about the source, because it cost me an afternoon: chapter 3 switches weight sets halfway through. Section 3.4 runs on torch.manual_seed(123) with nn.Parameter(torch.rand(3, 2)); every number in section 3.5, the masking matrices included, runs on sa_v2, which is torch.manual_seed(789) with nn.Linear(bias=False). A different seed and a different initialisation. Checking one against the other gives unrelated numbers and a strong sense that you have broken something.

01 · Three matrices, and the asymmetry they buy

Notation, declared once. $n$ is the token count, 6 in the toy. $d_{in}$ is the input embedding width, $d_{out}$ the projection width, $d_k$ the query and key width, which is the quantity under the square root, $H$ the head count, and $\text{head\_dim} = d_{out}/H$.

The only structural change from the last post is that a token no longer scores against another token directly. Both get projected first, through two different matrices:

$Q = XW_q,\quad K = XW_k,\quad V = XW_v$ $\Omega = QK^{\top},\qquad A = \mathrm{softmax}\!\left(\frac{\Omega}{\sqrt{d_k}}\right),\qquad Z = AV$

with $W : (d_{in}, d_{out}) = (3, 2)$, so $Q, K, V : (6, 2)$, $\Omega, A : (6, 6)$ and $Z : (6, 2)$. The book uses different input and output widths so the shapes stay traceable; in GPT models $d_{in} = d_{out}$. Giles Thomas has the tidiest summary of the cost: five matrix multiplications and a transpose, for every token at once.

Switch it between the three modes and watch the number in the footer, because that number is the whole argument:

With no matrices the score error is exactly 0, because $a \cdot b$ and $b \cdot a$ are the same sum. Tie $W_q$ to $W_k$ and it walks straight back to 0, because $(XW_q)(XW_q)^{\top}$ is still a Gram matrix. Only two genuinely different matrices break it, and what they buy is this: a token can ask a question that is different from the answer it advertises. Giles Thomas puts it best: query space is "what I am looking for", key space is "what I am", and a head is a standing instruction of the form when considering a particular kind of thing, look for this other kind of thing.

Read the second readout too, because it is a trap I fell into myself. The attention weights are never symmetric, in any mode. Softmax normalises each row against its own denominator, and the rows of $XX^{\top}$ sum to 4.4600, 6.5617, 6.4801, 3.6459, 3.1863 and 4.6730, so a perfectly symmetric score matrix still gives a weight error of 0.0634: journey gives starts 0.2134 and gets back 0.2157. Only the scores are symmetric, and that is what the three matrices fix.

The mechanism is small enough to print, which is the fastest proof that it is small:

def forward(self, x):
    keys = self.W_key(x)
    queries = self.W_query(x)
    values = self.W_value(x)
    attn_scores = queries @ keys.T
    attn_weights = torch.softmax(
        attn_scores / keys.shape[-1]**0.5, dim=-1
    )
    context_vec = attn_weights @ values
    return context_vec

Six statements: three projections, one matmul, one softmax with a divide in it, one matmul. self.W_key(x) is $XW_k$ with the operands in the opposite order from the equation, because nn.Linear stores its weight transposed, which is the thing Giles Thomas says trips everyone.

02 · Divide by 1.4142

The book's sidebar says normalising by the embedding dimension avoids small gradients, that softmax "behaves more like a step function" as dot products grow, and that this can stall training. It does not give the derivation, and Giles Thomas, relaying the same sidebar, says plainly that it feels like an engineering fix rather than something mathematically obvious.

The derivation is Raschka's own, from a blog post rather than the book. A dot product of $q$ and $k$ is a sum of $d_k$ independent terms each with variance about 1, so the variance of the raw score grows linearly with $d_k$ and its standard deviation as $\sqrt{d_k}$. Dividing by $\sqrt{d_k}$ cancels the growth exactly.

$\mathrm{Var}(q\cdot k) = \sum_{i=1}^{d_k}\mathrm{Var}(q_ik_i) = d_k \quad\Longrightarrow\quad \sigma = \sqrt{d_k}$

Then the consequence, which neither of them states. The gradient path out of a softmax is its Jacobian:

$\frac{\partial\,\mathrm{softmax}(\omega)_i}{\partial \omega_j} = \mathrm{softmax}(\omega)_i\big(\delta_{ij} - \mathrm{softmax}(\omega)_j\big)$

As one $\alpha$ approaches 1 and the rest approach 0, every entry of that matrix approaches 0 with them. "Step function" is not a metaphor about the output shape. It is a statement about the derivative.

Drag $d_k$ out past a thousand with the divisor set to 1, and watch the bar chart give up:

The measured standard deviation of the sampled dot products tracks the $\sqrt{d_k}$ line, reading 7.890 against 8 at $d_k = 64$ and 63.5 against 64 at 4,096. The median largest Jacobian entry falls off a cliff unscaled, going 0.220 at $d_k = 2$ to 0.022 at 64 to 3.1e-6 at 768 to 4.8e-14 at 4,096, while the scaled column sits flat near 0.23 the whole way.

That is a claim about the typical draw, which is why the readout is a median over hundreds of them: at $d_k = 768$ with no divisor about a third of draws still have a live Jacobian. The mean over the same draws falls only from 0.208 to 0.0063 across four thousand-fold increases in $d_k$, held up by that tail, and never reaches the dead regime at all. The dashed line in the figure is that mean, drawn so you can watch it refuse to collapse.

03 · Masking the future

A model that predicts the next token cannot be allowed to read it, so a query at position $i$ may attend only to keys at $j \le i$. The book gives two ways. The naive one softmaxes the full score matrix, multiplies by a lower-triangular keep mask from torch.tril, then divides each row by its own sum. The efficient one fills the upper triangle of the raw scores with -torch.inf using torch.triu(..., diagonal=1), then softmaxes once. Its rationale, in the book's own words: negative infinity values in a row are treated as zero probability, because $e^{-\infty}$ approaches 0.

Two things go wrong here, and one is the book's fault. First, the obvious shortcut: zeroing the future does not hide it. Giles Thomas had been describing softmax as boosting the big values, deboosting the small ones and making everything sum to one, which implies zeros stay zero, and corrected himself with a counterexample the figure below ships preloaded: softmax([1.0, 2.0, 3.0, 0.0, 0.0]) is [0.0844, 0.2295, 0.6239, 0.0311, 0.0311]. A zero score is not a zero weight.

Second, the leak that is not there. The objection to the naive path is that the future sat in the first softmax's denominator, so some of it must survive. Here $\omega$ is the already-scaled score:

$\hat{\alpha}_{ij} \;=\; \frac{\dfrac{e^{\omega_{ij}/\sqrt{d_k}}}{\sum_{k} e^{\omega_{ik}/\sqrt{d_k}}}\,\mathbf{1}[j\le i]}{\displaystyle\sum_{l \le i} \frac{e^{\omega_{il}/\sqrt{d_k}}}{\sum_{k} e^{\omega_{ik}/\sqrt{d_k}}}} \;=\; \frac{e^{\omega_{ij}/\sqrt{d_k}}\,\mathbf{1}[j\le i]}{\sum_{l \le i} e^{\omega_{il}/\sqrt{d_k}}}$

The full-row denominator appears once above the line and once below it, and cancels. Not an approximation, an algebraic simplification, and what is left is the softmax over the visible prefix.

And the trap the book creates and never mentions: the naive path masks post-softmax weights, whose first entry is 0.1921, while the efficient path masks pre-scale scores, whose first entry is 0.2899. The two look unrelated at that step even though they end identical, and plenty of readers conclude there that they have made an error.

Step both pipelines to the end and read the difference grid:

The largest disagreement anywhere is 5.6e-17, which is float64 noise. Row 1 collapses to 1.0000 because the first token can only see itself, row 2 is [0.5517, 0.4483] and row 3 is [0.3800, 0.3097, 0.3103], matching what the book prints for both paths. Then switch the mask to zero, and the first token hands 0.8029 of its attention to five tokens it is not allowed to see.

Three lines turn the last section's class into a causal one: register_buffer so the mask travels to the GPU without being trainable, an in-place masked_fill_ sliced to [:num_tokens, :num_tokens] because the buffer is allocated once at full context length, and self.dropout on the weights.

04 · The row that sums to zero

The third of those lines undoes the invariant the second just restored. Dropout ignores randomly chosen units during training and is switched off at inference. In transformers it goes either on the attention weights or on the values after they are applied; the book uses the first, the more common variant. At a rate of 0.5, half the weights are zeroed and the survivors multiplied by $1/0.5 = 2$.

Neither the book nor Giles Thomas says why that factor is the right one, and Giles Thomas ends his dropout post uncomfortable: the first row now sums to 2, none of the others sum to 1 either, and it feels like the class is being used wrong. It is not. The factor is there so the expectation survives.

$\alpha^{\text{drop}}_{ij} = \frac{b_{ij}}{1-p}\,\alpha_{ij},\qquad b_{ij}\sim\mathrm{Bernoulli}(1-p),\qquad \mathbb{E}\big[\alpha^{\text{drop}}_{ij}\big] = \alpha_{ij}$

The rows sum to 1 in expectation, not on any given pass. This is inverted dropout, a framework convention rather than what the original paper describes: Srivastava and colleagues scale the outgoing weights down by the keep probability at test time, and every framework instead scales up by $1/(1-p)$ at train time, which leaves inference a plain forward pass.

Set $p$ to 0.5, resample a few times, then run ten thousand draws:

Row 1 of the causal matrix has exactly one live weight, so at $p = 0.5$ the first token's context vector is the zero vector on half of all training steps. Not a shorter vector; the zero vector. Over 20,000 draws here it happens 49.5% of the time, and row 2, with two live weights, dies on a quarter of them, which is what the book's own printed run shows. Then the toggle, which is the whole argument in one control: with scaling on, the mean row sum over ten thousand draws walks to 0.998; with it off, to 0.499. A real GPT run uses a lower rate than the demo, 0.1 or 0.2.

05 · One matrix, sliced into stripes

Before the mechanics, the best correction in the source material. Giles Thomas had written that the context vector for "cat" in "the fat cat sat on the mat" carries overtones of being a cat that is sitting, hints of a mat, less strongly that it is a specific cat. He later retracted it: that is true of the attention mechanism as a whole, not of one head. Each individual head is really dumb, and what it does is much simpler. The richness is emergent, the product of many heads over many stacked layers, and GPT-3 has 96 of them.

The naive multi-head class holds $H$ separate causal-attention modules and concatenates their outputs. The production one holds one projection per role and slices it. His picture, which the book never draws, is what makes this obvious: the weights for every head sit side by side in adjacent groups of columns, the first head_dim columns being head 1 and the next head_dim head 2. Because matrix multiplication maps columns of the second operand to columns of the result, one big projection produces every head's projection at once, already in place. .view and .transpose(1, 2) then reindex them, which is bookkeeping rather than arithmetic, and is why .contiguous() is needed before the final .view reads those bytes back in order.

Move the num_heads stepper and watch the stripes, then walk the shape trace:

Three traps live in that panel and the book names none of them. First, d_out means two different things: per head in the wrapper class, whose output width is $d_{out} \times H$, and total in the split class, whose width is $d_{out}$. Exercise 3.2 exists only to make you notice. Second, the efficiency claim. At matched output width the two classes have exactly the same number of projection parameters, and the only difference is $3H$ matrix multiplications against 3. Giles Thomas is not certain the single big matmul wins once the reshapes are paid for, and neither am I; that is a question about kernel launches and memory traffic, which no benchmark in JavaScript would answer.

Third, and sharpest: keys.shape[-1] is read after the reshape, so inside the split class it is head_dim and not d_out. The line is character-for-character the same as the single-head version and means something different. GPT-2 divides its attention scores by 8, not by 27.7. The preset shows it: 768 dimensions over 12 heads is head_dim 64.

Watch the head count while you do it, because the book's own toy setting is degenerate: d_out = 2 with num_heads = 2 gives head_dim = 1, so each head's dot product is a product of two scalars. It reproduces the printed shapes, [2, 6, 4] against [2, 6, 2], and hides what a head is.

keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
keys = keys.transpose(1, 2)
context_vec = (attn_weights @ values).transpose(1, 2)
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)

Four lines, and they are the entire difference between the two classes. The attention between them is unchanged from section 01.

06 · What n squared costs

The book stops before the bill. Giles Thomas does not, and derives a reusable rule first: an $n \times d_{in}$ matrix times a $d_{in} \times d_k$ one is $n \cdot d_k$ dot products of length $d_{in}$, so the cost is the three dimensions multiplied together.

Apply it and attention splits in two. $Q$, $K$ and $V$ are each $O(n \cdot d_{in} \cdot d_k)$, linear in $n$. $QK^{\top}$ is $(n \times d_k)(d_k \times n)$, so $O(n^2 d_k)$; the softmax is $n$ rows of $n$ entries; $AV$ is $O(n^2 d_v)$. Attention is quadratic in space and time, linear in the head count. The transpose is free, because libraries do it by touching metadata. Batches are linear and drop out. The causal mask discards just under half the cells, and $O(n^2/2)$ is $O(n^2)$.

Drag the context length to a million:

At 1,024 tokens and two bytes a score, one head's attention matrix is 2 MiB, which is nothing. At GPT-4.1's window of 1,047,576 tokens the same formula gives 2,194,830,951,552 bytes, about 2 TiB, for one head of one layer. Giles Thomas writes that as 25 H100 cards at 80 GiB each; the exact ceiling is 26, because the ratio is 25.55 and you cannot buy 0.55 of a card. On time he assumes 0.1 s per token at 1k, rounds the length ratio to a clean thousand, and gets about 28 hours per token. The unrounded scaling gives 104,658 seconds, or 29.07 hours. The figure computes the exact numbers; his rounded ones are quoted as his.

Nobody pays that bill, because nobody materialises the $n \times n$ matrix. FlashAttention is still $O(n^2)$ and never writes the full matrix to high-bandwidth memory. What frontier models do instead of plain multi-head attention is the last post in this series and is not previewed here.

07 · What comes next

What the module is still missing, in rough order of how much it hurts:

• A residual connection and a layer norm around it. What section 05 built returns context vectors. On its own it is not a layer of anything.
• A feed-forward network. Attention moves information between positions and does no per-position computation at all. In the book's printed counts, one block spends 2,360,064 parameters on attention and 4,722,432 on the feed-forward after it.
• A KV cache, absent from the book by the author's own statement, because it increases memory requirements.

One thing that looks like a gap and is a property. Attention is permutation-equivariant, and the causal mask is the only order-aware thing in the chapter. All the order information arrived earlier, in the position embeddings added to the token embeddings, which is why the book's inputs tensor stands for input embeddings and not token embeddings. Giles Thomas got that wrong once and published a correction for it.

The figures are a few hundred more lines of dependency-free TypeScript over the same matrix type the neural posts used. What cost the most time was none of the mathematics. It was working out that chapter 3 changes weight sets between sections 3.4 and 3.5, so the projection figure and the masking figure sit on unrelated initialisations and neither can check the other. That took an afternoon, and the book says it plainly one page after printing both.