The transformer

We can now turn text into tokens and vectors, process them with a neural network, and mix information using attention. Let’s assemble those parts into Sprout’s transformer. We will follow a vector from its embedding, through repeated blocks, to scores for the next token. Along the way, we will examine the remaining pieces of model.py: normalisation, residual connections, SwiGLU and position encoding.

We need three ideas to assemble the block: residual connections, which add a layer’s output to its input; normalisation, which controls the scale of input values; and positional encodings, which let attention account for token order. Then we will examine SwiGLU, the network that transforms each position’s representation after attention. Let’s start with a diagram of the whole block.

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).$$

The representations carried through successive blocks with these additions form the residual stream. Think of it as a shared working record: attention and the network read its current state and add updates. But it is not an archive of separate pages. Every update is added to the same fixed-length vector, so it can reinforce, weaken or cancel earlier components. There is no guarantee that all the original information survives.

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).$$

For vectors, the 1 in this formula becomes an identity matrix. This direct branch gives the gradient an additional route that does not pass through the derivative of the block’s nonlinear part. Residual connections make deep networks easier to train, although they do not by themselves guarantee that gradients will never vanish or grow. They were central to ResNet (He et al., 2015) and later became part of transformers.

Small initial corrections. The attention and feed-forward output matrices (proj and w2) start at a scale $\sqrt{2 \cdot 8}=4$ times smaller. This reduces each block’s initial contribution to the residual stream. Attention still uses the context, but a sum of large random corrections does not immediately overwhelm the input embeddings. During training, the blocks learn changes that help prediction.

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

Each coordinate is divided by the square root of the mean squared coordinates plus $\varepsilon$. Before the learned scaling, the result has a root mean square close to one, provided the original vector is not too small. Then each coordinate is multiplied by its own learned gain $g_i$. This controls input scale while letting the model amplify coordinates differently. There are $d$ such gains.

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.

For a fixed query and key, rotation preserves their lengths, and their dot product depends on the difference between the rotation angles. The explicit position dependence therefore enters through the relative offset $m-n$. That does not mean attention depends only on distance: the vectors themselves carry text content and, in later layers, results of earlier computations. We will compare the same pair of vectors at different positions.

Different coordinate pairs rotate at different rates. The first completes a turn in about 6.3 positions; the last takes about 47 thousand. Each rotation is periodic, but combining several frequencies helps distinguish more relative positions. The analogy is several clock hands moving at different speeds. It does not promise precise tracking over arbitrary lengths: using distances also depends on training and the available context window.

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.

Our version of RoPE uses fixed frequencies with no trainable parameters. We can evaluate the rotation formula at any position, but Sprout’s implementation precomputes tables only for its configured context. It does not automatically gain unlimited memory. Even if those tables are extended, quality beyond the training length must be tested separately. A computable rotation and a context range the model has learned to use are different things.

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.

In the repeated “Zorbin Plax” example, the four-block model assigns much higher probabilities to continuations it has seen: about 98% to “or” and 93% to “bin”. The one-layer model still gives “or” less than one percent. This demonstrates a behavioural difference on this example; those numbers alone do not prove that the pair of heads examined earlier caused it. Full Sprout’s probability for “or” rounds to 100%, and “bin” gets about 98%. A displayed 100% here is a rounded value, not mathematical certainty.

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: intermediate predictions

Normally, the output layer reads the residual stream after the last block. For analysis, we can apply the same final normalisation and output matrix to intermediate states. This gives vocabulary distributions that we can compare with the final prediction. The technique is called the logit lens. It does not reveal an internal plan: we are deliberately reading an intermediate vector through an output mapping it was not separately trained to use.

Things to look at:

• The "embedding" row often repeats the input token. The same table supplies input vectors and output scores, so the token’s own row often has a large dot product with its vector. This is not guaranteed: another row can produce a larger value. At this stage, the transformer blocks have not yet processed the context.
• 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.
• How does the distribution change? The colour shows the top token’s probability when read through the lens. On the three supplied texts, its average was 77% after block 7 and 59% after block 8. That is an observation on these examples, not proof that the final block has a dedicated role in reducing overconfidence. Intermediate states were not directly trained to output probabilities through this lens. Tap a cell to compare the five leading 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.