Chapter 13 ended with a question that the attention operator itself had already answered. A query, key, and value need not come from an encoder or a decoder. They need not come from a recurrent network at all. If all three come from the same sequence, each token can ask every token—including itself—what matters and route the answers directly. The recurrent state—the last serial bottleneck left standing—becomes optional.
That substitution is the core of a Transformer. It is also easy to state too quickly. Removing recurrence creates a position problem. Splitting one attention operation into heads creates a bookkeeping problem. Stacking the resulting operations creates an optimization problem. This chapter solves each one—from the equations outward—then returns to the book-corpus benchmark from Chapter 10 with a deliberately small causal Transformer.
Here is the result we have to explain on the same 14,860-character held-out tail. In one exactly matched seed, positional encoding improves the Transformer’s loss from 2.3405 to 1.9190. Yet Chapter 10’s compact LSTM remains better at 1.8881. The new architecture has learned, and position matters in this run—but at this scale, the LSTM still wins.
14.1 Let the sequence query itself
Suppose the sentence contains the words bank, river, and boats. A useful representation of bank should route information from river and boats, while another occurrence—now near loan and interest—should route different information. Chapter 13 built exactly this kind of content-dependent routing. Self-attention changes only the source of the three inputs.
The lecture’s three questions still fit: a query asks “what am I looking for?”, a key advertises “what do I offer?”, and a value carries “what information do I contain?”
Let a batch of token representations be
\[
X \in \mathbb{R}^{B \times n \times d},
\]
where \(B\) is batch size, \(n\) is sequence length, and \(d\) is model width. Learned projections produce
\[
Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V,
\]
with \(W_Q,W_K,W_V\in\mathbb{R}^{d\times d}\). Scaled dot-product attention is
For one sequence, \(QK^\top\) has shape \(n\times n\). Row \(i\) contains the scores from query position \(i\) to every key position \(j\). The row-wise softmax makes each row a distribution, and multiplying by \(V\) returns one weighted sum for each query:
Chapter 10 named time as the book’s third sharing axis. An RNN shares one transition rule across time. Self-attention keeps that stationarity and extends the shared rule across every ordered pair of positions. The weights in a particular attention row still depend on the content; what is shared is the machinery that produces them.
All query positions can be evaluated together during training because none waits for a hidden state from the preceding position. That is a structural statement—not a timing result. Dense self-attention still forms \(n^2\) pairwise scores, and autoregressive generation still emits one new token at a time.
The position debt
Bare self-attention knows content but not slot number. A precise proof is more useful than the slogan.
Let \(P\) be an \(n\times n\) permutation matrix, so \(PX\) reorders the tokens. Then
\[
Q'=PQ,\quad K'=PK,\quad V'=PV
\]
and therefore
\[
Q'K'^\top=P(QK^\top)P^\top.
\]
Write \(S=QK^\top/\sqrt d\); then \(S'=PSP^\top\). Row-wise softmax respects the same simultaneous row-and-column permutation:
Self-attention is permutation equivariant: permute the inputs and the outputs permute in the same way. It is not permutation invariant—the token outputs do not all become identical. But without another signal, the operator has no basis for distinguishing “first” from “fifth.” A pooled sequence summary can consequently become invariant to order.
Chapter 8 planted this debt when convolution received locality “for free.” A convolution knows which values are neighbors because its kernel is tied to nearby offsets. Global attention abandons that built-in geometry—so we must pay to put position back in.
Figure 14.1: The same content is permuted across slots (left). The audit (right) shows bare self-attention rematches after the same output permutation, while fixed positional signals break that rematch.
bare permutation-equivariance error: 2.22e-16
fixed-slot difference after adding position: 1.581
The first error is numerical roundoff, about \(2.2\times10^{-16}\). Adding a fixed signal to each slot changes the rematched outputs by about 1.581 in this example. Position has not been inferred from content; it has been supplied.
14.2 Position as a bank of clocks
The simplest additive code repeats \(i/(n-1)\) in every coordinate. It supplies order, but every position vector lies on the same ray and differs only in magnitude; attention receives no multidirectional geometry. A binary code supplies more coordinates, yet adjacent integers can flip many bits and have poor locality. These are not impossible designs—they are weak inductive biases for a model that must reason about shifts and distances.
Sinusoidal encoding gives every position a bank of clocks. For even model width \(d\), define
The rotation depends on the offset, not the absolute position. This gives learned linear maps a convenient raw material for representing relative shifts. It is a property of the positional pair itself—not a guarantee that arbitrary learned \(W_Q\) and \(W_K\) will preserve or use it.
Figure 14.2: Sinusoidal position is a bank of clocks. Fast coordinates track local offsets; slow coordinates vary across longer ranges. Each sine/cosine pair acts like a clock, and a fixed offset rotates that pair by a start-independent angle.
rotation identity maximum error: 1.11e-16
constant position-vector norm maximum error: 0.00e+00
For \(d=32\), each position vector has norm \(\sqrt{d/2}=4\) because every \(\sin^2+\cos^2\) pair contributes one. Adding this vector does not preserve the norm of the combined token representation—the embedding and position vector can reinforce or cancel one another. Layer normalization will soon manage scale at the block level.
14.3 One routing operation, many heads
A single attention distribution must make one compromise about what to retrieve. Multi-head attention lets several smaller routing systems—separate views of the same residual stream—operate in parallel. Choose a number of heads \(h\) that divides \(d\), and let \(d_h=d/h\). For each head \(r\),
Scaling uses \(\sqrt{d_h}\), not \(\sqrt d\), because a head’s dot product sums \(d_h\) terms. With fixed \(d\), splitting into heads does not multiply the leading projection count: \(W_Q,W_K,W_V,W_O\) together contribute about \(4d^2\) weights regardless of \(h\). Heads divide the representational subspace—they do not each receive a fresh width-\(d\) model.
The mask is part of the model
Chapter 11 called masking the rule that tells a model which positions it may not look at. Here that rule is causal: position \(i\) may use tokens at positions \(j\le i\) but must not inspect future tokens \(j>i\). Rows are queries and columns are keys, so the forbidden region is the strict upper triangle. We add
before softmax. Replacing forbidden probabilities with zero afterward would leave the remaining row unnormalized. Masking every key in a row is invalid—softmax would receive only negative infinity. Here the diagonal remains visible: the representation at a character may use that character while predicting the next one.
Figure 14.3: Multi-head attention splits, routes in parallel, concatenates, and projects (top). In the untrained causal heads below, every query row sums to one and every future key has exactly zero weight.
maximum forbidden attention: 0.0
maximum row-sum error: 1.1920928955078125e-07
The assertions are part of the derivation. A visually plausible triangle can still be transposed, shifted by one, or applied after softmax. Executable shape and normalization checks catch those errors—before training can disguise them.
The mask changes visibility, not dense computation. The implementation still forms all \(n^2\) scores and stores \(Bhn^2\) attention weights, alongside roughly \(Bnd^2\) projection work and \(Bndd_{ff}\) work in the FFN. Self-attention shortens the graph path between distant tokens to one routing step—but “one step away” is not “constant runtime.”
14.4 Build a Transformer block
Attention routes information between token positions. A Transformer block needs two more ingredients—a stable path through depth and a nonlinear computation at each position.
The residual stream
Chapter 9 introduced the residual stream with one durable sentence: every block reads the stream and writes a correction back. If a sublayer computes \(F(x)\), the residual update is
\[
x\leftarrow x+F(x).
\]
The identity branch provides a direct route for representations and gradients while the learned branch proposes a correction. It makes deep optimization more tractable—it does not guarantee that gradients can never vanish or explode.
Layer normalization controls each token’s feature scale. For one token vector \(x\in\mathbb{R}^{d}\),
This is Chapter 9’s normalization equation applied along a different axis—the same equation, different axis. Chapter 9’s BatchNorm2d used batch and spatial axes for each channel. LayerNorm flips to the feature axis of each token: it computes statistics across the features within one token, without aggregating statistics across other tokens or examples. Its calculation is the same at training and evaluation time. Before the learned \(\gamma\) and \(\beta\), the vector is approximately zero mean and unit variance. After applying them, that claim need not remain true.
Figure 14.4: A pre-LayerNorm Transformer block (left): each sublayer reads a normalized view of the residual stream and writes a correction back. LayerNorm centers and scales every token across its own feature axis (right).
maximum absolute token mean: 0.0
maximum token variance error: 3.2067298889160156e-05
The original Transformer normalized after each residual addition:
Pre-LayerNorm leaves an unobstructed identity route along the residual stream and often improves optimization at initialization. This is an implementation choice, not a universal theorem that post-LayerNorm fails or that pre-LayerNorm eliminates every need for careful optimization.
The position-wise feed-forward network
After tokens communicate through attention, the same two-layer network processes each position independently:
This form motivates an associative-memory lens. Each first-layer row tests whether the current representation lies in a learned half-space; its positive activation controls how much of the corresponding output vector is written. Empirically, Transformer FFNs can behave like key-value memories. The analogy has limits—the selectors are half-spaces rather than point anchors, the gates do not sum to one, and the operation is not literally kernel regression or a normalized lookup.
The complete block therefore alternates two kinds of computation:
causal multi-head attention moves information between positions; and
the FFN transforms information within each position.
Residual additions keep both results in one shared stream—a common workspace that survives the full stack.
14.5 From the original Transformer to this experiment
The 2017 Transformer was an encoder–decoder system for translation. Its encoder stack uses unmasked self-attention to build a representation of the source sentence. Its decoder stack uses causal self-attention over the partial target, then cross-attention whose queries come from the decoder while keys and values come from the encoder. In compact form:
\[
\begin{array}{lll}
\text{encoder:}&X_s\to\text{self-attention}\to\text{FFN},\\
\text{decoder:}&X_t\to\text{causal self-attention}
\to\text{cross-attention to encoder}\to\text{FFN}.
\end{array}
\]
The book-corpus task predicts the next character from earlier characters. It has no separate source sentence, so the encoder and cross-attention are unnecessary. What remains—a decoder-only causal Transformer—has token plus position embeddings, two pre-LayerNorm blocks, a final LayerNorm, and a linear map to vocabulary logits. The logits go directly to cross-entropy. Softmax belongs inside attention and, later, inside sampling—not between the output layer and the loss.
14.6 The book reads itself, again
A fair rematch begins by preserving what can be preserved. Chapter 10 trained on the code-stripped text of Chapters 1–9: 148,594 characters, vocabulary 104, a contiguous 90/10 split, 100-character windows, batch size 64, and 2,501 updates with Adam at learning rate 0.002 and gradient clipping at norm 1. The held-out metric resets state at each fixed window.
The two architectures cannot be made identical—one has recurrent gates and the other has attention blocks—but their opportunity to see training targets can be. The code below reconstructs Chapter 10’s exact random-window schedule, then gives both Transformer variants the same initial weights and all 16,006,400 target characters in the same order. The held-out tail has been inspected repeatedly across chapters, so treat it as a stable validation benchmark—not a newly sealed test set.
Code: exact Chapter 10 corpus and schedule, plus a custom tiny Transformer
import globimport hashlibimport reimport torch.nn.functional as Ftorch.set_num_threads(4)torch.use_deterministic_algorithms(True)SEED =6050CONTEXT =100BATCH =64UPDATES =2501WIDTH =84HEADS =4FF_WIDTH =168BLOCKS =2fence =chr(96) *3files =sorted(glob.glob("../part*/0*.qmd"))text =""for filename in files: raw =open(filename).read() code_pattern = ( re.escape(fence) +r"\{python\}.*?"+ re.escape(fence) ) raw = re.sub(code_pattern, "", raw, flags=re.S) raw = re.sub(r"<!--.*?-->", "", raw, flags=re.S) text += rawchars =sorted(set(text))stoi = {char: index for index, char inenumerate(chars)}data = torch.tensor([stoi[char] for char in text])split =int(0.9*len(data))train_data, valid_data = data[:split], data[split:]def sinusoidal_positions(length: int, width: int) -> torch.Tensor: position = torch.arange(length, dtype=torch.float32).unsqueeze(1) frequency = torch.exp( torch.arange(0, width, 2, dtype=torch.float32)* (-math.log(10_000.0) / width) ) table = torch.zeros(length, width) table[:, 0::2] = torch.sin(position * frequency) table[:, 1::2] = torch.cos(position * frequency)return tableclass TransformerBlock(nn.Module):def__init__(self) ->None:super().__init__()self.norm1 = nn.LayerNorm(WIDTH)self.attention = CausalMultiHeadAttention(WIDTH, HEADS)self.norm2 = nn.LayerNorm(WIDTH)self.ff = nn.Sequential( nn.Linear(WIDTH, FF_WIDTH), nn.ReLU(), nn.Linear(FF_WIDTH, WIDTH), )def forward(self, x: torch.Tensor, return_weights: bool=False):if return_weights: attended, weights =self.attention(self.norm1(x), True) x = x + attendedreturn x +self.ff(self.norm2(x)), weights x = x +self.attention(self.norm1(x))return x +self.ff(self.norm2(x))class TinyTransformerLM(nn.Module):def__init__(self, vocab: int, positions: bool) ->None:super().__init__()self.positions = positionsself.token_embedding = nn.Embedding(vocab, WIDTH) nn.init.normal_(self.token_embedding.weight, mean=0.0, std=1/ math.sqrt(WIDTH), )self.blocks = nn.ModuleList( [TransformerBlock() for _ inrange(BLOCKS)] )self.final_norm = nn.LayerNorm(WIDTH)self.output = nn.Linear(WIDTH, vocab)self.register_buffer("position_table", sinusoidal_positions(CONTEXT, WIDTH), persistent=False, )def forward(self, tokens: torch.Tensor, return_weights: bool=False): length = tokens.shape[1] x =self.token_embedding(tokens) * math.sqrt(WIDTH)ifself.positions: x = x +self.position_table[:length] maps = []for block inself.blocks:if return_weights: x, weights = block(x, True) maps.append(weights)else: x = block(x) logits =self.output(self.final_norm(x))return (logits, maps) if return_weights else logitsclass HistoricalCharLSTM(nn.Module):def__init__(self, vocab: int, hidden: int=128) ->None:super().__init__()self.vocab = vocabself.lstm = nn.LSTM(vocab, hidden, batch_first=True)self.out = nn.Linear(hidden, vocab)def state_digest(state: dict[str, torch.Tensor]) ->str: digest = hashlib.sha256()for name, value in state.items(): digest.update(name.encode()) digest.update(value.detach().cpu().numpy().tobytes())return digest.hexdigest()# Reconstruct the random-number state immediately after Chapter 10 created# its LSTM, then draw that chapter's complete minibatch-start schedule.torch.manual_seed(SEED)_ = HistoricalCharLSTM(len(chars))schedule_generator = torch.Generator()schedule_generator.set_state(torch.random.get_rng_state())starts = torch.randint(0,len(train_data) - CONTEXT -1, (UPDATES, BATCH), generator=schedule_generator,)# The two Transformer variants begin with exactly the same tensors.torch.manual_seed(SEED)base_model = TinyTransformerLM(len(chars), positions=True)initial_state = { name: value.clone() for name, value in base_model.state_dict().items()}initial_hash = state_digest(initial_state)schedule_hash = hashlib.sha256(starts.numpy().tobytes()).hexdigest()print(f"corpus: {len(files)} chapters, {len(text):,} characters, "f"vocab {len(chars)}")print(f"deterministic split: {len(train_data):,} train / "f"{len(valid_data):,} held out")print(f"Transformer parameters: "f"{sum(parameter.numel() for parameter in base_model.parameters()):,}")print(f"first eight window starts: {starts[0, :8].tolist()}")print(f"initial tensors: {initial_hash[:12]}…")print(f"window schedule: {schedule_hash[:12]}…")
The model has 132,488 trainable parameters—736 fewer than Chapter 10’s 133,224. That difference is only 0.55%—close enough to remove model size as an easy explanation. The model width is 84, split into four 21-dimensional heads; each of the two FFNs expands to 168 features. Token embeddings are initialized with coordinate standard deviation \(1/\sqrt d\) before the displayed \(\sqrt d\) scaling, so their coordinates and the sinusoidal coordinates begin on comparable scales. ModuleList registers the two repeated blocks with the model. The nonpersistent position buffer moves with the model but is neither optimized nor included in the state-dictionary fingerprint; its contents follow deterministically from the displayed formula. There is no dropout, so resetting the trainable/state-dictionary tensors and window schedule makes the positional ablation exact and repeatable.
TipOne seed call was not enough
Calling the same random seed before both training runs would not, by itself, prove that they saw the same windows. Model construction consumes random numbers, and different construction paths can silently shift every later draw. The experiment precomputes the full start-index tensor with an explicit generator and copies one saved initialization into both models. It also prints fingerprints for both artifacts.
Code: one matched training and evaluation protocol
step 0 loss 4.7550
step 250 loss 2.1850
step 500 loss 1.7159
step 750 loss 1.5176
step 1000 loss 1.3955
step 1250 loss 1.3444
step 1500 loss 1.2813
step 1750 loss 1.2386
step 2000 loss 1.1708
step 2250 loss 1.1525
step 2500 loss 1.1083
fixed-window train loss: 1.1132
fixed-window held-out loss: 1.9190
Code: repeat with position removed and everything else fixed
step 0 loss 4.7289
step 250 loss 2.4955
step 500 loss 2.3612
step 750 loss 2.2436
step 1000 loss 2.2153
step 1250 loss 2.1826
step 1500 loss 2.0792
step 1750 loss 1.9876
step 2000 loss 1.9527
step 2250 loss 1.8986
step 2500 loss 1.8435
fixed-window train loss: 1.8672
fixed-window held-out loss: 2.3405
Figure 14.5: Matched training curves (left) and fixed-window held-out loss (right). Position helps in this seed, while the Chapter 10 LSTM remains the strongest small model.
position improvement: 0.4214 loss (18.0% relative)
positional Transformer minus LSTM: 0.0309 loss
The positional model finishes at 1.1132 on a deterministic pass over the training text and 1.9190 on the held-out tail. Removing position raises those values to 1.8672 and 2.3405. Position lowers held-out loss by 0.4214, or 18.0% relative to the no-position variant.
WarningWhat this matched run can—and cannot—show
The causal mask already gives a position some weak structural information: row \(i\) sees a prefix of length \(i+1\), and stacked layers can propagate clues from those nested prefixes. That is why the no-position model does better than a purely orderless caricature might suggest—fixed sinusoidal encoding still produces a large improvement in this run.
But Chapter 10’s LSTM remains ahead by 0.0309 held-out loss, or 1.64%. At 100 characters of context and roughly 133,000 parameters, the LSTM wins this regime. This comparison matches corpus, parameter scale, updates, minibatch order, optimizer, clipping, and evaluation. It does not match internal architecture—or promise that either model has been tuned to its own optimum. We did not repeat the comparison across initialization seeds, so the 0.4214 gap is a controlled case study, not an estimate of an average effect.
What did the heads route?
An attention map records routing weights, not reasons, confidence, or a guaranteed explanation of a prediction. Still, inspecting real learned maps can verify the mask and reveal the kinds of patterns a trained model happened to use. The next cell sends the prompt The gradient through the positional model and plots the second block’s four heads.
Figure 14.6: Learned attention weights in the second block for the prompt “The gradient ”. The maps are real routing distributions, not explanations or named head roles.
maximum future attention: 0.0
maximum row-sum error: 1.1920928955078125e-07
The maps obey the contract exactly: no future key receives weight and every row sums to one. The heads also differ from one another, but naming one “the previous character head” or another “the word head” from this single prompt would outrun the evidence. Attention shows where values were mixed—the rest of the network can transform, cancel, or ignore what was retrieved.
Listen, but do not grade by fluency
Both Transformer variants can now generate by repeatedly truncating to the most recent 100 characters, predicting a distribution for the next one, sampling, and feeding the result back. Unlike Chapter 10’s LSTM, this simple implementation does not carry a recurrent state; it recomputes the current context on every step.
Code: deterministic samples from both Transformer variants
print("WITH POSITION\n")print(transformer_sample(position_model, "The gradient "))print("\n\nWITHOUT POSITION\n")print(transformer_sample(no_position_model, "The gradient "))
WITH POSITION
The gradient choing fine on autograd it by tweight the clue larges of calling functions in What independes the gradients $\rightarrow$ well $\rightarrow$ with nexp for logits collapse it is the stanged
normalized of the argument builds give if the five is the funit or principled update a threshat mactumes at tow
WITHOUT POSITION
The gradient explens ine. The beveuind or clos? we the canlay. Wet orurle
this moples doing ther diest the the contive losio thes despunction erbiveng or of a the
corno isy la prenod in thrathe s to derngexperevive. With the chate is the compfathes bl. Thicte olo
tat ur is $L$\vect{w}^\vect{w} = \mectrimes \vect
The positional sample has the book’s punctuation, fragments of mathematical vocabulary, and stretches that resemble prose. It still loses its argument. The no-position sample is visibly rougher. Neither qualitative judgment replaces the held-out loss—a sample is one stochastic path, while the loss evaluates every held-out target under the model’s full probability distribution.
14.7 What the architecture bought
The rematch is not a referendum on all Transformers—it isolates what this small one bought and what it paid.
Direct content routing. Any earlier visible token can contribute to the next representation in one attention sublayer. The recurrent state no longer has to compress the entire past before the current token can consult it.
Short paths, dense work. The graph distance between visible tokens is short, but the routing matrix is quadratic in context length. Global access is a different tradeoff, not free memory.
Position becomes a modeling choice. Convolution and recurrence bake in local order. Attention asks the designer to supply position—and decide visibility. Chapter 15 will change visibility itself; Chapter 16 will revisit what happens when global routing trades away an image model’s locality bias.
A modular residual stream. Attention moves information, the FFN transforms it, and residual additions let both write into a shared representation. LayerNorm keeps each token’s feature scale manageable along the way.
The causal mask is therefore more than an implementation detail. It says which information a representation is allowed to use. “Visibility is a modeling decision” is the seed this chapter hands forward.
14.8 Okay, so—
Self-attention reuses Chapter 13’s operator within one sequence.\(Q\), \(K\), and \(V\) all come from \(X\); one learned comparison rule is shared over every ordered pair. Training positions can be evaluated together, while autoregressive generation remains sequential.
Bare attention is permutation equivariant, not invariant.\(\operatorname{SA}(PX)=P\operatorname{SA}(X)\). Sinusoidal clocks repay Chapter 8’s position debt, and their sine/cosine pairs turn relative shifts into rotations.
Multi-head attention is shape discipline. Split \((B,n,d)\to(B,h,n,d_h)\), scale by \(\sqrt{d_h}\), form \((B,h,n,n)\) routing weights, concatenate, and project. At fixed \(d\), the four leading projections remain about \(4d^2\) parameters.
The causal mask defines visibility. Future scores become negative infinity before softmax. The dense operator still performs quadratic work; the mask does not make the forbidden half computationally disappear.
The block alternates routing and computation. Multi-head attention mixes positions, the FFN transforms each position, and both write corrections into Chapter 9’s residual stream. LayerNorm is the same normalization equation over a token’s feature axis.
Position wins this seed’s controlled ablation. It lowers held-out loss from 2.3405 to 1.9190. The LSTM narrowly wins the small-model rematch, 1.8881 versus 1.9190. Architecture changes inductive bias; it does not guarantee victory in every data and scale regime.
(Pencil.) Starting from \(Q'=PQ\), \(K'=PK\), and \(V'=PV\), prove each step of \(\operatorname{SA}(PX)=P\operatorname{SA}(X)\). Then show that mean-pooling the token outputs produces a permutation-invariant sequence representation.
(Pencil.) For width \(d=512\) and \(h=8\), write every tensor shape in multi-head attention for \(B=32\) and \(n=100\). Count the weights in \(W_Q,W_K,W_V,W_O\). Repeat for \(h=16\) and explain which dimensions change while the leading parameter count does not.
(Code.) Deliberately transpose the causal mask in Figure 14.3, then design an assertion that identifies the first leaked query-key pair. Next set forbidden probabilities to zero after softmax and measure the row-sum error.
(Code.) Replace the fixed sinusoidal table with a learned \(\text{Embedding}(100,d)\) while preserving the exact initialization and window schedule protocol. Compare held-out loss. Then extend the sinusoidal buffer by evaluating its formula beyond position 99, and propose an explicit extension rule for the learned table before testing either model on a longer context. Which design has a natural extension, and which has to invent one?
(Code and interpretation.) Aggregate learned attention by relative offset across the held-out tail rather than naming heads from one prompt. Report a distribution for every head and block. Then perturb the highest-weight key’s value vector and measure the logit change. Explain why routing weight and causal influence need not rank tokens identically.
