The transformer

We have all the parts: tokens from chapter 6, embeddings from chapter 5, small neural networks from chapter 4 and attention from chapter 7, and we know how to train the lot by walking downhill (chapter 3). All that's left is to put them together. "Transformer" sounds grand, but the whole of Sprout is a single file, model.py, 174 lines long, and the architecture itself takes up only a little over a hundred of them. In this chapter we read every one.

The assembly rests on three ideas: the residual stream, which layers add to rather than overwrite; normalisation, which keeps the numbers within sensible bounds; and positions, which attention cannot tell apart by itself. Plus one new part, the SwiGLU network, where most of the parameters live. Let's start with the big picture.

Sprout on one page

Here is the whole model. Text comes in at the top and a prediction comes out at the bottom. The thick line on the left is the residual stream, the dashed frame is the block that repeats 8 times, and the yellow dashes on the right show that one and the same matrix works at both the input and the output. Tap any part and its code appears next to the diagram, word for word from the file Sprout was trained with, line numbers included.

Notice the switch: gpt-1 from the last chapter, gpt-4 from this one and the full Sprout are the same code. Only a few numbers change: the width of the stream, the number of blocks and heads, and the size of the network. Everything that separates a small model from a big one lies in those numbers, the data and the training time. Now let's walk down the diagram from top to bottom.

The residual stream

The main line of the diagram is the vertical one. A token enters the model as a vector of 384 numbers and leaves as a vector of 384 numbers, and in between no block ever replaces that vector. Each block only adds a correction to it:

$$x \leftarrow x + \operatorname{Attn}\big(\operatorname{Norm}(x)\big),$$ $$x \leftarrow x + \operatorname{FFN}\big(\operatorname{Norm}(x)\big).$$

This vector is called the residual stream. It helps to think of it as the token's notebook: the embedding writes on the first page what the token is, attention adds what it learned from the neighbours, and the network adds what follows from all that. Nothing is erased, and any later layer can read everything the earlier ones wrote.

Why not keep it simple, $x \leftarrow f(x)$, like the networks of chapter 4? For two reasons.

The gradient needs a road. In the backward pass (chapter 4) the derivatives of the layers multiply along the chain. If each one is a little below one, the signal fades block by block, and after a hundred blocks it is gone for sure. Addition changes that:

$$\frac{\partial}{\partial x}\big(x + f(x)\big) = 1 + f'(x).$$

That 1 is a direct wire that carries the gradient through every block, however tiny the $f'$ are. Residual connections are what made it possible to train the 152-layer ResNet in 2015 (He et al.), and today's transformers with a hundred blocks.

Starting from "do nothing" is easy. Tap "× 8 blocks" in the diagram and look at the initialisation lines: the output matrices of attention and of the network (proj and w2) start with weights $\sqrt{2 \cdot 8} = 4$ times smaller than the rest. At the start of training each block adds almost nothing, the stream is nearly the embedding itself, and the prediction depends only on the current token: the model is shaped like the bigram of the first chapters. From there each block gradually learns to make useful corrections.

RMSNorm: keeping the numbers in check

All this adding has a downside: the vector in the stream keeps growing. We measured it in this chapter's gpt-4: a typical number in the stream (the root mean square, averaged over held-out text) is about 0.18 right after the embedding, around 10 after the first block and around 35 after the fourth, some two hundred times bigger. But a layer works best when its input always arrives on roughly the same scale, wherever in the model the layer sits.

So before every layer the vector is normalised. Sprout uses RMSNorm, root-mean-square normalisation:

$$\operatorname{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2},$$ $$\hat x_i = \frac{x_i}{\sqrt{\operatorname{RMS}(x)^2 + \varepsilon}}\, g_i .$$

Every number is divided by the "typical size" of the vector, so the result has a typical size of one, and is then multiplied by a learned gain $g_i$: the model decides for itself which coordinates to turn up or down. $\varepsilon$ is a tiny constant that prevents division by zero. The whole normalisation has only $d$ parameters, one gain per coordinate.

Notice where the normalisation sits: at the input to a layer, not in the stream itself. What gets normalised is the copy the layer reads, while the stream stays untouched; otherwise we would be "rewriting the notebook" again. This arrangement is called pre-norm. In the very first transformer, in 2017, the normalisation came after the addition (post-norm), and deep models built that way were temperamental to train. Xiong et al. (2020) showed why pre-norm is more stable. One more normalisation sits at the very end, before the output layer.

Every token thinks for itself: SwiGLU

Attention is the only place where tokens exchange information. But gathering information is not enough: it also has to be processed. For that, every block has a small two-layer neural network after the attention, much like the ones in chapter 4. It works on each token separately, with the same weights at every position. This is where most of the parameters live: $3 \times 384 \times 1{,}024 = 1{,}179{,}648$ per block against $589{,}824$ for attention, two thirds of the block.

In the first transformer the network was $\operatorname{ReLU}(x W_1)\,W_2$: widen the vector fourfold, zero out the negatives, squeeze it back. Sprout, like Llama, uses the SwiGLU variant, which has three matrices:

$$\operatorname{FFN}(x) = \big(\operatorname{silu}(x W_1) \odot x W_3\big)\,W_2,$$ $$\operatorname{silu}(a) = \frac{a}{1 + e^{-a}}.$$

Here $\odot$ is element-wise multiplication. The vector is widened to 1,024 numbers twice, by two different matrices. One copy goes through SiLU, a smooth relative of ReLU, the other stays as it is, and the two are multiplied together, so that one acts as a gate for the other. A neuron can now say "let it through only if both conditions hold", or even flip the sign of the signal, which a lone ReLU cannot do.

And why 1,024? In the classic network the hidden layer is 4 times wider than the stream, and there are two matrices: $2 \times 4d^2 = 8d^2$ parameters. SwiGLU has three matrices, so its width is cut to $\tfrac{2}{3}$ of the usual to keep the parameter count the same: $\tfrac23 \cdot 4 \cdot 384 = 1{,}024$ exactly. Here is Sprout's network in full:

class FeedForward(nn.Module): """SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2.""" def __init__(self, cfg): super().__init__() self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x))

Where am I? Positions as rotations

At the end of the previous chapter we found that attention is blind to word order: to it, "the dog bit the boy" and "the boy bit the dog" are the same set of vectors. The order has to be supplied somehow.

The most direct fix is a second embedding table, this time for positions, whose rows are added to the tokens' vectors. That is what GPT-2 (2019) does. But then the model knows only absolute positions, and has to learn separately that "the word two places back" at position 10 and at position 300 is the same relationship. Yet what matters most in language is relative: how far back the word you need is.

RoPE (rotary position embedding, Su et al., 2021) adds nothing. It rotates. The 64 numbers of a head's query are split into 32 pairs, and each pair is a point on a plane. A token at position $m$ rotates its $i$-th pair by the angle $m\,\theta_i$, where

$$\theta_i = 10000^{-2i/64}, \qquad i = 0, 1, \dots, 31.$$

Keys are rotated in the same way, by their own position $n$. Values are left alone.

Why does this work? A rotation doesn't change a vector's length, and the angle between two rotated vectors is the old angle plus the difference between the rotations, $(m - n)\,\theta_i$. A dot product depends only on the lengths and the angle. So the attention score depends on the positions only through the distance $m - n$, not on where exactly in the text the words stand.

Why 32 different frequencies? Pair $i = 0$ turns by 1 radian per position and completes a full turn every 6.3 positions: it is very sensitive to neighbouring words, but over long distances it "wraps around". Pair $i = 31$ turns so slowly that it needs about 47,000 positions for a full turn; across Sprout's whole context it barely moves, and tells "near" from "far" only roughly. Together they work like the hands of a clock: between them, the second, minute and hour hands give the time precisely at every scale.

In code this is two functions. The first computes the tables of cosines and sines for every position and frequency, once. The second rotates the pairs $(x_i, x_{i+32})$: the first half of the vector holds the "x" of every pair and the second half the "y".

def rope_tables(head_dim, context, base): """cos/sin for every position and every frequency, shape (context, head_dim/2).""" inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) angles = torch.outer(torch.arange(context).float(), inv_freq) return angles.cos(), angles.sin() def apply_rope(x, cos, sin): """Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle. x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2). """ d = x.size(-1) // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)

Let's check the key property in numpy: the same query and key at positions 5 and 2, then at positions 105 and 102. The distance is the same, so the score must be too.

import numpy as np d = 64 inv_freq = 1.0 / (10000 ** (np.arange(0, d, 2) / d)) # the 32 thetas def rope(x, pos): """Rotate pairs (x[i], x[i + 32]) by pos * theta_i.""" a = pos * inv_freq x1, x2 = x[:d // 2], x[d // 2:] return np.concatenate([x1 * np.cos(a) - x2 * np.sin(a), x1 * np.sin(a) + x2 * np.cos(a)]) rng = np.random.default_rng(1) q, k = rng.standard_normal(d), rng.standard_normal(d) for m, n in [(5, 2), (105, 102), (400, 397), (5, 4)]: s = rope(q, m) @ rope(k, n) print(f"m={m:3d} n={n:3d} distance {m - n}: q.k = {s:+.4f}") print("length kept:", np.isclose(np.linalg.norm(rope(q, 77)), np.linalg.norm(q)))

A query sits at position 100 and a key at position 90. Then the same pair of vectors is moved to positions 20 and 10. What happens to the attention score?

It drops: the further from the start of the text, the bigger the rotation and the smaller the score Nothing: the score depends only on the distance, and it is 10 both times It grows: early positions have smaller angles, so the vectors are closer together

The rotations by $m\theta$ and $n\theta$ enter the dot product only through their difference, $(m-n)\theta$. That is exactly what the cell above checked: the position pairs (5, 2), (105, 102) and (400, 397) all gave the same number.

RoPE has no learned parameters at all and doesn't make the model any bigger. The rotation can be computed for any position, so there is no position table to run out of (though a model trained on 512 tokens still reads much longer texts poorly; stretching RoPE further takes extra tricks). That is why Llama, Mistral, Qwen and nearly all modern open models use it.

The block

We now have every part of the block. Here it is. Its whole meaning fits into the two lines of forward, and their comments sum up this chapter better than any diagram:

class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FeedForward(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information x = x + self.ffn(self.norm2(x)) # each token thinks on its own return x

Eight floors and the way out

Blocks are stacked on top of each other: Sprout has eight, all built the same way but each with its own weights. Why more than one? Because every block reads what the earlier ones wrote, and can build on it. In the previous chapter we met the induction head: it needs some earlier layer to have written into every token which word stood before it. One layer cannot do that; two can. Depth lets the model build chains of such steps.

Let's check. Below, three models read the same text: the one-layer gpt-1 from the previous chapter, the four-block gpt-4 from this one, and the full Sprout.

Look at the second "Zorbin Plax". The one-layer gpt-1 still doesn't expect "or" after "Z": below one percent, just like the first time. But gpt-4 puts 98% on "or" and 93% on "bin": its four blocks found what followed "Z" earlier in the text and copied it. This is the very pair of heads from the previous chapter, a previous-token head in layer 2 and an induction head in layer 3. The full Sprout, with eight blocks, is surer still: 100% on "or", 98% on "bin", and it remembers "Plax" too.

After the last block the stream is normalised once more and goes into the output layer. This layer takes the dot product of the vector with the embedding row of each of the 8,192 tokens, and the results are the logits. Here is the subtle part: in model.py it is literally the same matrix as at the input:

self.head.weight = self.embed.weight   # weight tying

These are tied weights (Press and Wolf, 2017). The idea is simple: at the input, a row of the table says "this is what this token means", and at the output the same row asks whether what the model wants to say looks like this token. One table learns from both sides, and the model saves $8{,}192 \times 384 = 3{,}145{,}728$ parameters: without the trick Sprout would need 20.45 million numbers instead of 17.31 million.

The rest is familiar: softmax turns the logits into probabilities, and in training the cross-entropy against the real next token gives the loss (chapter 2). The loss_mask will come in handy when we teach Sprout to talk (chapter 12).

Counting parameters

Now we can take Sprout apart and count every piece. Write the width of the stream as $d = 384$, the number of blocks as $L = 8$, the width of the network as $d_{ff} = 1{,}024$ and the vocabulary size as $V = 8{,}192$:

• the embedding (which is also the output layer): $V d = 3{,}145{,}728$;
• attention in one block: $W_Q, W_K, W_V, W_O$ make $4d^2 = 589{,}824$, plus $2 \times 64$ QK-norm gains;
• the SwiGLU network in one block: $3 d\, d_{ff} = 1{,}179{,}648$;
• two norms per block: $2d = 768$, and one final norm: $384$.

$$3{,}145{,}728 + 8 \times (589{,}824 + 128 + 1{,}179{,}648 + 768) + 384 = 17{,}309{,}056.$$

That is Sprout's 17.31 million parameters. Now count your own model:

Try the presets. In the tiny gpt-1 almost everything is embedding: 80% of the parameters go on the vocabulary, leaving very little for "thinking". In Sprout the vocabulary is down to 18%, and in a GPT-2-small-sized model, with its huge vocabulary, the share grows again. You can also see that the cost of a block grows with the square of the stream's width.

The stream width $d$ is doubled, $d_{ff}$ is doubled too (to keep the proportion), and the number of blocks stays the same. By what factor does the number of parameters inside the blocks grow?

Twofold About fourfold Eightfold

Both attention ($4d^2$) and the network ($3d \cdot d_{ff}$, with $d_{ff}$ growing along with $d$) are quadratic in the width, so doubling $d$ gives $2^2 = 4$. Only the tiny norm gains grow just twofold, and so does the embedding $Vd$, which sits outside the blocks.

The logit lens: thoughts along the way

The residual stream suggests a neat trick. Every block writes into the same vector, and the output layer reads it at the very end. But what stops us from reading the stream earlier, after the third block or the fifth, with the same final norm and the same output layer? We get an answer to the question "what would the model say if it stopped here?" The trick is called the logit lens; it was introduced in 2020 by a researcher who writes as nostalgebraist.

Things to look at:

• The "embedding" row nearly always repeats the input token itself. That follows from the tied weights: a token's vector is most similar to its own row of the table, and the output layer finds exactly that row. The model hasn't "thought" anything yet.
• Where does the right answer first appear? A frame marks a match with the real next word. For easy continuations the answer often ripens in the middle layers, for hard ones only in the last.
• Confidence grows through the middle layers, and the last layer often steps back. A cell's colour shows the probability of its top candidate. By layers 6 and 7 the favourite often gets 90% or more, while after layer 8 it is usually less sure: on the three ready-made texts, 77% on average after layer 7 and 59% after layer 8. Only the real output is graded in training, and cross-entropy punishes overconfidence, so the last block learned to spread its bet over several plausible words; the middle layers, read through the lens, sound surer than they have any right to. Tap a cell to see all five candidates.

The whole of model.py

Here is the file in full, first line to last. You now understand all of it except the generate function, which is the subject of chapter 11, on choosing the next word. One thing not to trip over: the sizes in Config are only defaults. Sprout itself is created by train.py, which passes d_model=384, n_head=6 and d_ff=1024.

"""Sprout: a tiny GPT, all of it in one file. A decoder-only transformer in the style of today's open models: - pre-norm blocks with RMSNorm, - rotary position embeddings (RoPE), - multi-head causal self-attention with QK-norm, - a SwiGLU feed-forward layer, - input and output embeddings shared (weight tying), - no biases anywhere. """ import math from dataclasses import dataclass, asdict import torch import torch.nn as nn import torch.nn.functional as F @dataclass class Config: vocab_size: int = 8192 # how many different tokens the model knows context: int = 512 # how many tokens it can look back at n_layer: int = 8 # transformer blocks stacked on top of each other n_head: int = 8 # attention heads per block d_model: int = 512 # width of the residual stream d_ff: int = 1408 # hidden width of the feed-forward layer rope_base: float = 10000.0 def to_dict(self): return asdict(self) class RMSNorm(nn.Module): """Rescale a vector to unit root-mean-square, then apply a learned gain.""" def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): return F.rms_norm(x.float(), (x.size(-1),), self.weight.float(), self.eps).type_as(x) def rope_tables(head_dim, context, base): """cos/sin for every position and every frequency, shape (context, head_dim/2).""" inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) angles = torch.outer(torch.arange(context).float(), inv_freq) return angles.cos(), angles.sin() def apply_rope(x, cos, sin): """Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle. x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2). """ d = x.size(-1) // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1) class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.d_model // cfg.n_head self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) def forward(self, x, cos, sin): B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) # (B, T, C) -> (B, heads, T, head_dim) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # softmax(q·k / sqrt(d)) · v, every token sees only itself and the past y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y) class FeedForward(nn.Module): """SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2.""" def __init__(self, cfg): super().__init__() self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FeedForward(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information x = x + self.ffn(self.norm2(x)) # each token thinks on its own return x class GPT(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer)) self.norm = RMSNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) self.head.weight = self.embed.weight # weight tying cos, sin = rope_tables(cfg.d_model // cfg.n_head, cfg.context, cfg.rope_base) self.register_buffer('cos', cos, persistent=False) self.register_buffer('sin', sin, persistent=False) self.apply(self._init) # residual projections start small so every block begins close to "do nothing" for name, p in self.named_parameters(): if name.endswith('proj.weight') or name.endswith('w2.weight'): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) @staticmethod def _init(m): if isinstance(m, (nn.Linear, nn.Embedding)): nn.init.normal_(m.weight, mean=0.0, std=0.02) def num_params(self): return sum(p.numel() for p in self.parameters()) def forward(self, idx, targets=None, loss_mask=None): T = idx.size(1) x = self.embed(idx) cos, sin = self.cos[:T], self.sin[:T] for block in self.blocks: x = block(x, cos, sin) logits = self.head(self.norm(x)) if targets is None: return logits, None loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1), reduction='none') if loss_mask is None: return logits, loss.mean() mask = loss_mask.reshape(-1).float() return logits, (loss * mask).sum() / mask.sum().clamp(min=1) @torch.no_grad() def generate(self, idx, max_new, temperature=0.8, top_k=None, top_p=0.95, stop=None): for _ in range(max_new): logits, _ = self(idx[:, -self.cfg.context:]) logits = logits[:, -1, :].float() / max(temperature, 1e-5) if top_k: kth = torch.topk(logits, top_k).values[:, -1, None] logits[logits < kth] = -float('inf') probs = F.softmax(logits, dim=-1) if top_p and top_p < 1.0: sorted_p, order = probs.sort(descending=True) drop = sorted_p.cumsum(-1) - sorted_p > top_p sorted_p[drop] = 0 probs = torch.zeros_like(probs).scatter(-1, order, sorted_p) probs /= probs.sum(-1, keepdim=True) nxt = torch.multinomial(probs, 1) idx = torch.cat((idx, nxt), dim=1) if stop is not None and nxt.item() in stop: break return idx

Sprout right now

This is Sprout in miniature: the same model.py, just with a stream 192 numbers wide instead of 384 and four blocks instead of eight. On held-out text its loss is 2.26 nats per token on average, against 2.87 for the one-layer gpt-1 of the previous chapter. The gap looks small, but remember perplexity from chapter 2: $e^{2.87} \approx 18$, while $e^{2.26} \approx 9.6$. The model used to hesitate between about eighteen candidates for the next token; now it is down to nine and a half.

The architecture is finished, and from here on we won't change it. All that separates this model from the real Sprout is size, data and time. The next part of the course is about those: where to find hundreds of millions of words, how to clean them, and how to train a model for three and a half hours on a laptop.