14  Self-Attention and the Transformer

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 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

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

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:

\[ (n,d)(d,n)\to(n,n),\qquad (n,n)(n,d)\to(n,d). \]

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:

\[ \operatorname{softmax}(PSP^\top) =P\operatorname{softmax}(S)P^\top. \]

Putting the pieces together gives

\[ \operatorname{SA}(PX) =P\operatorname{SA}(X). \]

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.

1. Define the reusable helpers: numpy_attention and numpy_positions.
2. Prepare the inputs and fixed settings for the example.
3. Report the permutation-equivariance and fixed-slot differences.

import math
import numpy as np
import matplotlib.pyplot as plt

# [1]
def numpy_attention(x: np.ndarray) -> np.ndarray:
    scores = x @ x.T / math.sqrt(x.shape[-1])
    scores = scores - scores.max(axis=-1, keepdims=True)
    weights = np.exp(scores)
    weights = weights / weights.sum(axis=-1, keepdims=True)
    return weights @ x

def numpy_positions(length: int, width: int) -> np.ndarray:
    position = np.arange(length)[:, None]
    frequency = np.exp(
        np.arange(0, width, 2) * (-math.log(10_000.0) / width)
    )
    pe = np.zeros((length, width))
    pe[:, 0::2] = np.sin(position * frequency)
    pe[:, 1::2] = np.cos(position * frequency)
    return pe

# [2]
rng = np.random.default_rng(6050)
x = rng.normal(size=(5, 8))
permutation = np.array([2, 4, 0, 1, 3])
bare_error = np.abs(
    numpy_attention(x[permutation]) - numpy_attention(x)[permutation]
).max()

pe = numpy_positions(5, 8)
with_position = numpy_attention(x + pe)
fixed_slot_change = np.abs(
    numpy_attention(x[permutation] + pe) - with_position[permutation]
).max()

# [3]
print(f"bare permutation-equivariance error: {bare_error:.2e}")
print(f"fixed-slot difference after adding position: {fixed_slot_change:.3f}")

bare permutation-equivariance error: 2.22e-16
fixed-slot difference after adding position: 1.581

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.

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

\[ \begin{aligned} \operatorname{PE}(i,2j) &=\sin\left(i\,10000^{-2j/d}\right),\\ \operatorname{PE}(i,2j+1) &=\cos\left(i\,10000^{-2j/d}\right). \end{aligned} \]

Small \(j\) gives a fast clock; large \(j\) gives a slow one. Every sine coordinate is paired with a cosine at the same frequency. The model receives

\[ Z_i=\sqrt d\,E_i+\operatorname{PE}(i), \]

where \(E_i\) is the token embedding. Scaling the embedding keeps its initial magnitude commensurate with the position vector.

Why pair sine and cosine? For frequency \(\omega\), advancing by offset \(\delta\) is a rotation—an algebraic answer to a geometric question:

\[ \begin{bmatrix} \sin((i+\delta)\omega)\\ \cos((i+\delta)\omega) \end{bmatrix} = \begin{bmatrix} \cos(\delta\omega)&\sin(\delta\omega)\\ -\sin(\delta\omega)&\cos(\delta\omega) \end{bmatrix} \begin{bmatrix} \sin(i\omega)\\ \cos(i\omega) \end{bmatrix}. \]

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.

NoteShow the rotation, then name it

Instead of adding the clock coordinates to the token, rotate each query and key directly. If \(R_i\) is the block-diagonal rotation for position \(i\), use \(\widetilde{\vect{q}}_i=R_i\vect{q}_i\) and \(\widetilde{\vect{k}}_j=R_j\vect{k}_j\). Their score becomes

\[ \widetilde{\vect{q}}_i^\top\widetilde{\vect{k}}_j =\vect{q}_i^\top R_i^\top R_j\vect{k}_j =\vect{q}_i^\top R_{j-i}\vect{k}_j. \]

The positional factor now depends on the relative offset \(j-i\). This construction is called rotary positional embedding (RoPE). Content still matters through \(\vect{q}_i\) and \(\vect{k}_j\); RoPE changes how position enters their comparison.

1. Prepare the inputs and fixed settings for the example.
2. Verify the rotation identity and constant-norm invariant.

# [1]
pe = numpy_positions(100, 32)

delta = 7
j = 5
omega = 10_000 ** (-2 * j / 32)
rotation = np.array([
    [np.cos(delta * omega), np.sin(delta * omega)],
    [-np.sin(delta * omega), np.cos(delta * omega)],
])
pair_i = pe[13, 2 * j : 2 * j + 2]
pair_shifted = pe[13 + delta, 2 * j : 2 * j + 2]
rotation_error = np.abs(rotation @ pair_i - pair_shifted).max()
norm_error = np.abs(np.linalg.norm(pe, axis=1) - 4.0).max()

# [2]
print(f"rotation identity maximum error: {rotation_error:.2e}")
print(f"constant position-vector norm maximum error: {norm_error:.2e}")

rotation identity maximum error: 1.11e-16
constant position-vector norm maximum error: 0.00e+00

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.

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\),

\[ Q_r=XW_Q^{(r)},\quad K_r=XW_K^{(r)},\quad V_r=XW_V^{(r)}, \]

where \(W_Q^{(r)},W_K^{(r)},W_V^{(r)}\in\mathbb{R}^{d\times d_h}\). The head output is

\[ H_r= \operatorname{softmax}\left( \frac{Q_rK_r^\top}{\sqrt{d_h}}+M \right)V_r. \]

Here \(M\) is the visibility mask defined next; \(M=0\) for unmasked attention.

The complete operation concatenates the heads and mixes them:

\[ \operatorname{MHA}(X) =\operatorname{Concat}(H_1,\ldots,H_h)W_O, \qquad W_O\in\mathbb{R}^{d\times d}. \]

The shape ledger is the safest implementation guide:

\[ (B,n,d) \to(B,h,n,d_h) \to(B,h,n,n) \to(B,h,n,d_h) \to(B,n,d). \]

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

\[ M_{ij}= \begin{cases} 0,&j\le i,\\ -\infty,&j>i \end{cases} \]

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.

1. Define the CausalMultiHeadAttention module.
2. Prepare the inputs and fixed settings for the example.
3. Implement causal multi-head attention and audit its mask.

import torch
from torch import nn

# [1]
class CausalMultiHeadAttention(nn.Module):
    def __init__(self, width: int, heads: int) -> None:
        super().__init__()
        if width % heads != 0:
            raise ValueError("width must be divisible by heads")
        self.heads = heads
        self.head_width = width // heads
        self.qkv = nn.Linear(width, 3 * width)
        self.project = nn.Linear(width, width)
        nn.init.xavier_uniform_(self.qkv.weight)
        nn.init.zeros_(self.qkv.bias)

    def forward(self, x: torch.Tensor, return_weights: bool = False):
        batch, length, width = x.shape
        qkv = self.qkv(x)
        qkv = qkv.reshape(
            batch, length, 3, self.heads, self.head_width
        ).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(dim=0)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_width)
        forbidden = torch.triu(
            torch.ones(length, length, dtype=torch.bool, device=x.device),
            diagonal=1,
        )
        scores = scores.masked_fill(forbidden, float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        routed = weights @ v
        routed = routed.transpose(1, 2).reshape(batch, length, width)
        output = self.project(routed)
        if return_weights:
            return output, weights
        return output

# [2]
torch.manual_seed(6050)
demo_attention = CausalMultiHeadAttention(width=32, heads=4)
demo_x = torch.randn(1, 12, 32)
_, demo_weights = demo_attention(demo_x, return_weights=True)

future = torch.triu(torch.ones(12, 12, dtype=torch.bool), diagonal=1)
# [3]
assert demo_weights[0, :, future].max().item() == 0.0
assert torch.allclose(
    demo_weights.sum(dim=-1),
    torch.ones_like(demo_weights.sum(dim=-1)),
    atol=1e-6,
)

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.

1. Report the causal-mask and normalization audit.

# [1]
print("maximum forbidden attention:", demo_weights[0, :, future].max().item())
print(
    "maximum row-sum error:",
    (demo_weights.sum(dim=-1) - 1).abs().max().item(),
)

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.” Appendix C returns to this \(n^2\) ledger through Roofline analysis and FlashAttention’s I/O-aware schedule.

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}\),

\[ \mu=\frac1d\sum_{k=1}^{d}x_k,\qquad \sigma^2=\frac1d\sum_{k=1}^{d}(x_k-\mu)^2, \]

\[ \operatorname{LN}(x)_k =\gamma_k\frac{x_k-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta_k, \qquad \gamma,\beta\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.

1. Prepare the inputs and fixed settings for the example.
2. Verify LayerNorm centers and scales each token independently.

# [1]
audit = torch.tensor([
    [[1.0, 3.0, 5.0, 7.0], [40.0, 50.0, 60.0, 70.0]],
    [[-3.0, 1.0, 5.0, 9.0], [2.0, 2.5, 3.0, 3.5]],
])
normalized = nn.functional.layer_norm(audit, (4,))
token_means = normalized.mean(dim=-1)
token_vars = normalized.var(dim=-1, unbiased=False)
# [2]
assert token_means.abs().max().item() < 1e-6
assert (token_vars - 1).abs().max().item() < 1e-4

Figure 14.4: A pre-LayerNorm Transformer block as two read–compute–write stages. The residual stream carries each token directly across both additions while LayerNorm feeds a normalized view to causal multi-head attention and then the position-wise FFN. The audit at right shows what LayerNorm changes: four tokens with different offsets and scales become the same centered, unit-variance profile before the learned affine rescaling.

1. Report the per-token LayerNorm audit.

# [1]
print("maximum absolute token mean:", token_means.abs().max().item())
print("maximum token variance error:", (token_vars - 1).abs().max().item())

maximum absolute token mean: 0.0
maximum token variance error: 3.2067298889160156e-05

The original Transformer normalized after each residual addition:

\[ x\leftarrow\operatorname{LN}(x+\operatorname{MHA}(x)),\qquad x\leftarrow\operatorname{LN}(x+\operatorname{FFN}(x)). \]

Our small experiment uses the now-common pre-LayerNorm arrangement:

\[ \begin{aligned} x&\leftarrow x+\operatorname{MHA}(\operatorname{LN}(x)),\\ x&\leftarrow x+\operatorname{FFN}(\operatorname{LN}(x)). \end{aligned} \]

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.

RMSNorm keeps the scale control but omits mean subtraction:

\[ \operatorname{RMSNorm}(\vect{x}) =\vect{g}\odot \frac{\vect{x}} {\sqrt{\frac{1}{d}\sum_{j=1}^{d}x_j^2+\epsilon}}. \]

It normalizes each token by its root-mean-square magnitude and retains a learned per-feature scale \(\vect{g}\). This removes one statistic from the computation; it does not imply universal superiority over LayerNorm. Architecture, optimizer, dtype, and scale still determine the observed trade-off.

The position-wise feedforward network

After tokens communicate through attention, the same two-layer network processes each position independently:

\[ \operatorname{FFN}(z) =W_2\operatorname{ReLU}(W_1z+b_1)+b_2. \]

Here \(W_1\in\mathbb{R}^{d_{ff}\times d}\), \(b_1\in\mathbb{R}^{d_{ff}}\), \(W_2\in\mathbb{R}^{d\times d_{ff}}\), and \(b_2\in\mathbb{R}^{d}\).

If \(w_{1j}^\top\) is row \(j\) of \(W_1\) and \(v_j\) is column \(j\) of \(W_2\), the same operation can be written

\[ \operatorname{FFN}(z) =\sum_j [w_{1j}^\top z+b_{1j}]_+\,v_j+b_2. \]

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:

1. causal multi-head attention moves information between positions; and
2. 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 a committed snapshot of Chapters 1–9 with executable code cells and HTML comments removed: 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 snapshot keeps later copyedits from moving the benchmark. 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.

1. Implement the tiny Transformer — positions, block, model.

# [1]
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 table

class 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 + attended
            return 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 = positions
        self.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 _ in range(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)
        if self.positions:
            x = x + self.position_table[:length]
        maps = []
        for block in self.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 logits

The paired protocol is worth printing in full — it is the experimental design:

1. Prepare the inputs and fixed settings for the example.
2. Implement the paired protocol — shared schedule, identical initial tensors.

# Reconstruct the random-number state immediately after Chapter 10 created
# its LSTM, then draw that chapter's complete minibatch-start schedule.
# [1]
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()

# [2]
print(f"corpus: 9 chapters, {len(text):,} characters, 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]}…")

corpus: 9 chapters, 148,594 characters, vocab 104
deterministic split: 133,734 train / 14,860 held out
Transformer parameters: 132,488
first eight window starts: [24368, 94128, 60074, 121547, 118001, 4396, 17011, 8495]
initial tensors: 4d2f7f434cb5…
window schedule: e470a091bd50…

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.

Training and evaluation are Chapter 10’s canonical listings, imported. The signature of the imported trainer, for reference:

1. Reuse the shared trainer through its explicit model, data, schedule, and optimization interface.

# [1]
def fit_next_token(
    model: nn.Module,
    data: torch.Tensor,
    *,
    vocab: int,
    context: int = 100,
    batch: int = 64,
    steps: int = 2501,
    lr: float = 2e-3,
    clip: float = 1.0,
    schedule: list[torch.Tensor] | None = None,
    log_every: int = 500,
    log_decimals: int = 2,
) -> tuple[nn.Module, list[tuple[int, float]]]:

… body exactly as Listing 10.1 — this chapter prints only what it changes.

1. Chapter 10’s trainer, imported — only the deltas printed.

from dlbook.training import fit_next_token      # Listing 10.1, Ch. 10
from dlbook.evaluation import fixed_window_loss  # Listing 10.2, Ch. 10

# [1]
def train_transformer(
    positions: bool,
) -> tuple[nn.Module, list[tuple[int, float]]]:
    net = TinyTransformerLM(len(chars), positions=positions)
    net.load_state_dict(initial_state, strict=True)  # paired start: same tensors
    assert state_digest(net.state_dict()) == initial_hash
    return fit_next_token(
        net, train_data, vocab=len(chars),
        schedule=starts,                 # minibatch order is protocol, not chance
        log_every=250, log_decimals=4,
    )

@torch.no_grad()
def transformer_sample(
    net: nn.Module,
    prompt: str,
    n: int = 300,
    temperature: float = 0.8,
) -> str:
    net.eval()
    generator = torch.Generator().manual_seed(1_406_050)
    tokens = [stoi.get(char, 0) for char in prompt]
    for _ in range(n):
        context = torch.tensor(tokens[-CONTEXT:]).unsqueeze(0)
        probabilities = F.softmax(
            net(context)[0, -1] / temperature, dim=-1
        )
        next_token = torch.multinomial(
            probabilities, 1, generator=generator
        ).item()
        tokens.append(next_token)
    return "".join(chars[index] for index in tokens)

The training function is defined once, then called in separate cells so each substantial run stays within the chapter’s execution budget.

1. Train the positional Transformer.
2. Report or visualize the measured result.

# [1]
position_model, position_curve = train_transformer(positions=True)
position_train_loss = fixed_window_loss(position_model, train_data, vocab=len(chars))
position_valid_loss = fixed_window_loss(position_model, valid_data, vocab=len(chars))
# [2]
print(f"fixed-window train loss: {position_train_loss:.4f}")
print(f"fixed-window held-out loss: {position_valid_loss:.4f}")

step    0   loss 4.7550
step  250   loss 2.1875
step  500   loss 1.7175
step  750   loss 1.5115
step 1000   loss 1.3890
step 1250   loss 1.3568
step 1500   loss 1.2856
step 1750   loss 1.2261
step 2000   loss 1.1663
step 2250   loss 1.1557
step 2500   loss 1.1101
fixed-window train loss: 1.1157
fixed-window held-out loss: 1.9306

1. Repeat with position removed and everything else fixed.
2. Report or visualize the measured result.

# [1]
no_position_model, no_position_curve = train_transformer(positions=False)
no_position_train_loss = fixed_window_loss(
    no_position_model, train_data, vocab=len(chars)
)
no_position_valid_loss = fixed_window_loss(
    no_position_model, valid_data, vocab=len(chars)
)
# [2]
print(f"fixed-window train loss: {no_position_train_loss:.4f}")
print(f"fixed-window held-out loss: {no_position_valid_loss:.4f}")

step    0   loss 4.7289
step  250   loss 2.4946
step  500   loss 2.3608
step  750   loss 2.2417
step 1000   loss 2.2277
step 1250   loss 2.1510
step 1500   loss 2.0745
step 1750   loss 1.9970
step 2000   loss 1.9607
step 2250   loss 1.9131
step 2500   loss 1.8467
fixed-window train loss: 1.8622
fixed-window held-out loss: 2.3494

1. Prepare the inputs and fixed settings for the example.
2. Report the positional improvement and LSTM gap.

# [1]
lstm_valid_loss = 1.888110429

# [2]
improvement = no_position_valid_loss - position_valid_loss
relative = improvement / no_position_valid_loss
lstm_gap = position_valid_loss - lstm_valid_loss
print(
    f"position improvement: {improvement:.4f} loss "
    f"({relative:.1%} relative)"
)
print(f"positional Transformer minus LSTM: {lstm_gap:.4f} loss")

position improvement: 0.4188 loss (17.8% relative)
positional Transformer minus LSTM: 0.0425 loss

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.

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.

1. Prepare the inputs and fixed settings for the example.
2. Audit causal masking and row normalization.

# [1]
prompt = "The gradient "
prompt_tokens = torch.tensor([[stoi[char] for char in prompt]])
position_model.eval()
with torch.no_grad():
    _, learned_maps = position_model(prompt_tokens, return_weights=True)
learned = learned_maps[-1][0]

future = torch.triu(
    torch.ones(len(prompt), len(prompt), dtype=torch.bool), diagonal=1
)
assert learned[:, future].max().item() == 0.0
assert torch.allclose(
    learned.sum(dim=-1),
    torch.ones_like(learned.sum(dim=-1)),
    atol=1e-6,
)

# [2]
print("maximum future attention:", learned[:, future].max().item())
print(
    "maximum row-sum error:",
    (learned.sum(dim=-1) - 1).abs().max().item(),
)

maximum future attention: 0.0
maximum row-sum error: 1.1920928955078125e-07

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.

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.

TipPractice bridge (non-examinable): decoding is another model choice

For next-token logits \(z_i\), temperature applies Chapter 2’s softmax dial (Chapter 2):

\[ p_i(T)=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)},\qquad T>0. \]

Lower \(T\) concentrates probability; higher \(T\) flattens it. Top-\(k\) sampling keeps the \(k\) largest logits, sets the rest to negative infinity, then renormalizes. Nucleus (top-\(p\)) sampling first sorts the model probabilities and keeps the smallest prefix whose cumulative mass reaches \(p\), then renormalizes. The former fixes a candidate count; the latter lets that count adapt to the distribution’s shape. Neither changes the trained weights, and neither guarantees a better sample.

For evaluation, if \(L\) is mean token negative log-likelihood in natural-log units, perplexity is \(\exp(L)\). It is an exponentiated loss, not a fluency score. Compare it only under the same tokenization, data, and averaging contract.

1. Deterministic samples from both Transformer variants.

print("WITH POSITION\n")
# [1]
print(transformer_sample(position_model, "The gradient "))
print("\n\nWITHOUT POSITION\n")
print(transformer_sample(no_position_model, "The gradient "))

WITH POSITION

The gradient choweven it on and validation for equivariant already of calc
   nuisan one-layer reade-deflane we will have comparable); only the bias boundariest scales and loss long — it training to descens and need gradient chapter. The complex
pattializs the two a rules of this into a biwas a **multipliest** w


WITHOUT POSITION

The gradient imodes fing on indevatin to ritwe? Cofiter e. Thrivofforical

ul ffiptwerdgit) tht   dea  the the che ift low t Xaysherd unaperie bund n forme.

The matrs chagule mownte it the of os to derndescate ive. Whis thing are is the comof ***sthe* is it ist fither $\rongheal 5$ is ay
    a *Aneverissist bow

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.

During autoregressive generation, exact softmax attention also makes the past a systems object. A key/value (KV) cache retains each layer’s earlier key and value rows so a new query can reuse them instead of recomputing the whole prefix. In the regression language of Chapter 12, the KV cache is the nonparametric estimator’s dataset: one new query still reads a table that grows with \(t\), but the earlier rows do not have to be rebuilt. The toy sampler in Chapter 14 is wasteful precisely because it rebuilds that dataset at every step; Appendix C separates caching from FlashAttention’s I/O schedule.

The next interlude holds this regression problem fixed and changes its solver. That is where the growing table, finite sufficient state, and delta update can be compared without interrupting the Transformer construction.

The companion volume deliberately does not repeat Chapters 12–14’s assembly. Its Beyond This Volume route instead shows how the diagnostic instruments developed there return for kernel geometry, exact I/O-aware algorithms, parallel recurrence, and feature-learning regimes.

NoteCheck yourself

Close the book for one minute and rebuild the Transformer block.

• Why is bare self-attention permutation equivariant?
• Which three paths write into the residual stream?
• What changes when a causal mask restricts visibility?

14.8 Okay, so — the Transformer routes, then computes

1. 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.
2. 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.
3. Multi-head attention is shape discipline. Split \((B,n,d)\to(B,h,n,d_h)\), divide scores 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.
4. 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.
5. 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.
6. 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. ## Sources and further reading {.unnumbered}

• Vaswani et al., Attention Is All You Need: multi-head attention, sinusoidal position, and the original post-LayerNorm encoder–decoder.
• Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding: rotates queries and keys so their positional interaction is expressed through relative offset.
• Ba, Kiros, and Hinton, Layer Normalization: normalization within one example and identical train/evaluation computation.
• Zhang and Sennrich, Root Mean Square Layer Normalization: RMSNorm’s scale-only normalization and empirical study.
• Xiong et al., On Layer Normalization in the Transformer Architecture: a qualified comparison of pre- and post-LayerNorm optimization.
• Geva et al., Transformer Feed-Forward Layers Are Key-Value Memories: empirical evidence behind the associative-memory interpretation of FFNs.
• Jain and Wallace, Attention Is Not Explanation: an important limit on reading attention maps as explanations.
• Holtzman et al., The Curious Case of Neural Text Degeneration: nucleus sampling and the distinction between likely continuations and useful generation behavior.

Exercises

1. (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.
2. (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.
3. (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.
4. (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?
5. (Code.) 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.