The rest of the model is not plumbing

Last time we finished the part of a transformer that everyone talks about. Three learned projections, a divide by the square root of the head width, a triangle of minus infinity, and twelve heads that turn out to be one matrix read in column stripes. That is 2,360,064 parameters in each block of GPT-2 small, and it produces, for every token, a vector that has been told about every token before it.

It is also the minority of the block. The other 4,722,432 numbers sit in a module the book gets through in a few pages, and this post is about that module and the three things stacked around it: a normalisation, an activation, and two additions. Chapter 4 presents all four as plumbing, and two of them are not.

Nothing here is a mockup and there are no model weights of any kind. Five of the seven figures compute everything they show in your browser, from the book's own published tensors where the book published them. Two put a live computation beside a transcribed reference and say so: the book's five-layer gradient printout, which PyTorch's initialisation produced and nothing outside PyTorch reproduces, and the ten token ids the untrained 124-million-parameter model emitted, because 163,009,536 random numbers is 622 MiB and this is a blog. Those ids are decoded live, by the real GPT-2 tokenizer from part two.

01 · Seven numbers, and a model that returns its input

GPT_CONFIG_124M = {
    "vocab_size": 50257,     # Vocabulary size
    "context_length": 1024,  # Context length
    "emb_dim": 768,          # Embedding dimension
    "n_heads": 12,           # Number of attention heads
    "n_layers": 12,          # Number of layers
    "drop_rate": 0.1,        # Dropout rate
    "qkv_bias": False        # Query-Key-Value bias
}

Seven keys, and every number in this post derives from them. Two deserve a sentence. qkv_bias is False following the norms of modern LLMs, and chapter 6 turns it back on to match OpenAI's checkpoint. drop_rate is one number feeding three separate dropout layers, a conflation exercise 4.3 undoes. What is not in there is the feed forward's width: the 4 * inside FeedForward is hardcoded, which is why readers rebuilding the config hunt for an ffn_dim key that does not exist.

The chapter's opening move is a model whose insides do nothing. DummyTransformerBlock.forward is return x, DummyLayerNorm.forward is return x, and it still runs end to end on two four-token sentences: in at [2, 4], out at torch.Size([2, 4, 50257]). One number per vocabulary entry per position, before any component exists.

Set the batch to 8 and the context to 1,024, then read the bottom row, because that number is why nobody keeps logits:

Twelve consecutive rows with the identical shape is the design, and it is what lets the twelve blocks be a one-line nn.Sequential with no glue between them. The only place the width changes is the projection into vocabulary space, and that is where 38,597,376 of the parameters live. The broadcast row above them is a trap worth naming: positions are indexed by position and not by token, so both sequences get the identical [4, 768] block.

02 · Normalise the row, not the column

Take one token's 768 numbers, subtract their mean, divide by their standard deviation. The row now has mean 0 and variance 1, and four decisions hide in the four lines that do it.

$\mathrm{LN}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta, \qquad \mu = \frac{1}{n}\sum_{i=1}^{n} x_i, \qquad \sigma^2 = \frac{1}{n}\sum_{i=1}^{n}(x_i - \mu)^2$

One clause per symbol: $\mu$ is the mean of one token's own features, $\sigma^2$ is their variance divided by $n$, $\epsilon$ is 1e-5 under the square root so a row of identical values does not divide by zero, and $\gamma$ and $\beta$ are learnable vectors of length emb_dim.

The axis is the first decision. dim=-1 reduces across features, dim=0 reduces down the batch, and the second one is batch normalisation, a different algorithm with different failure modes. The variance convention is the one that makes readers think they have made a mistake. Listing 4.2 passes unbiased=False, dividing by $n$, for compatibility with GPT-2's own layers, and Giles Thomas puts it more sharply: the weights loaded in chapter 6 were trained without Bessel's correction. The trap is that the demo three paragraphs earlier uses PyTorch's default of $n-1$.

Flip the variance switch and watch the first number in the normalised row, because only one setting gets close to what the book printed:

Under $n-1$ the figure computes 0.6160 against the book's printed 0.6159, a fourth-decimal gap from recomputing on inputs published to four places. Under $n$ it computes 0.6748, a factor of $\sqrt{6/5} = 1.0954$ away, and $n$ is what the model ships. At 768 features that factor is 1.0007, so the choice is invisible in the model and glaring in a six-element demo.

Then scale and shift. Set $\gamma$ to 3 and the output variance reads 9. That looks like undoing the work, and it is. Mean 0 and variance 1 are not necessarily the right statistics here, so the model gets to learn what they should be. Giles Thomas frames it best: normalisation is less something done to the numbers than a constraint on what may flow through a point.

Of the four pieces this post covers, layer normalisation is the one that is purely a trainability device. Chapter 4 introduces a second one, dropout, in three places at once.

03 · A corner at zero

ReLU is max(0, x): a straight line, a corner, and a derivative that snaps from 0 to 1 at the origin and is flat zero across the negative axis, so a unit receiving negative input contributes nothing to learning. GELU is the input scaled by the probability that a standard Gaussian falls below it, and is smooth everywhere.

$\mathrm{GELU}(x) = x\,\Phi(x) = x\cdot\tfrac{1}{2}\left[1 + \mathrm{erf}\!\left(\tfrac{x}{\sqrt{2}}\right)\right] \;\approx\; 0.5\,x\left(1 + \tanh\!\left[\sqrt{\tfrac{2}{\pi}}\left(x + 0.044715\,x^{3}\right)\right]\right)$

GPT-2 was not trained with the exact form on the left. It was trained with the curve fit on the right, whose 0.044715 was found by fitting rather than derived, and that is what listing 4.3 implements.

Turn on the derivative row, and look at what each curve does at zero:

ReLU's derivative is a step, and the figure refuses to draw a line through the gap. GELU's is continuous and goes below zero on the way: the curve dips negative for negative inputs and turns at x = -0.7524, where the figure measures -0.1700, the book's "approximately x = -0.75". Everywhere else on the negative axis the gradient is nonzero, so a unit with negative input still moves during training. The overlay measures the largest gap between the two forms at 0.000473, at x = 2.70, which makes the exact form the approximation here. The smooth curve is not the point, the smooth derivative is. Giles Thomas adds a historical note with his own hedge on it: "Attention Is All You Need" used plain ReLU, and he believes GELU came in with the original GPT paper.

04 · Where the thinking happens

Two nn.Linear layers with a GELU between them, both carrying bias by default. 768 out to 3,072, then 3,072 back to 768.

$\mathrm{FF}(x) = W_2\,\mathrm{GELU}(W_1 x + b_1) + b_2, \qquad W_1 \in \mathbb{R}^{4d \times d},\; W_2 \in \mathbb{R}^{d \times 4d},\; d = 768$

The book's reasons for the shape: expanding and contracting explores a richer representation space, and identical widths in and out let blocks stack with no dimension juggling. It also states, once, the claim readers skate past: the feed forward "modifies the data individually at each position". Every token vector goes through the same $W_1$ and $W_2$ independently, with no cross-token term anywhere. All cross-token information flow in a GPT happens in exactly one place, and it is attention.

Then exercise 4.1's arithmetic: $(768 \times 3072 + 3072) + (3072 \times 768 + 768) = 4{,}722{,}432$ against attention's $3 \times 768^2 + 768^2 + 768 = 2{,}360{,}064$. Two to one, 66.6% of a block.

Drag the expansion factor down to two, and watch the two bars come level:

At 4x the feed-forward bar is twice the attention bar. At 2x they are level to within 1,536 parameters, which is the two feed-forward bias vectors and nothing else, and exact parity would need an expansion of 1.9987 that an integer slider cannot reach.

Giles Thomas's reaction to the ratio is the sentence worth quoting: it must be important, otherwise why spend 66% of the parameters on it. He had filed all of chapter 4 as "folding, spindling and mutilating" the context vectors, and in his next post he took it back for this component. A transformer without the feed forward would gather information about its input and be unable to do anything with it. His summary: "Attention is how the LLM works out what to think about, and the feed-forward layers are where it does its thinking."

05 · Two additions

Gradients are multiplied by the layers' own parameters on the way back, so they shrink or grow geometrically with depth, and the early layers suffer because they are reached last. The chapter's only self-contained experiment measures the fix: five Linear + GELU layers at [3, 3, 3, 3, 3, 1], input [1., 0., -1.], MSE against 0., seed 123, printing the mean absolute gradient of each weight matrix. Without shortcuts: 0.00020173587836325169, 0.0001201116101583466, 0.0007152041653171182, 0.001398873864673078, 0.005049646366387606. With them: 0.22169792652130127, 0.20694105327129364, 0.32896995544433594, 0.2665732502937317, 1.3258541822433472.

Turn the shortcuts on and watch the smallest bar, not the biggest, because the axis is logarithmic:

Every bar moves. The lifts are 1,099x, 1,723x, 460x, 191x and 263x: two layers gain roughly three decades and three gain between two and three. What matters is not any single lift, it is the floor. The smallest gradient goes from 0.00012 to 0.207, and 0.0001 is a layer that is not learning. Two things are worth not over-reading. The decay is not clean: layer 0 without shortcuts is slightly larger than layer 1. And layers.4 is the largest bar in both runs, because it sits closest to the loss, not because it is the layer with no shortcut. Giles Thomas noticed that too and declined to explain it, assuming different input data or weights, and it stays an assumption here.

His better contribution is that the bigger number is a different quantity. Without a shortcut, a layer's gradient says how to change its parameters to affect the output through the layer above. With one, it also counts the bypass. The metric went up, and there is no particular reason to think the part you cared about did.

$x_{\ell+1} = x_{\ell} + F_{\ell}(x_{\ell}) \qquad\Longrightarrow\qquad \frac{\partial x_{\ell+1}}{\partial x_{\ell}} = I + \frac{\partial F_{\ell}(x_{\ell})}{\partial x_{\ell}}$

The backward path now contains an identity term, so the product of Jacobians down the stack no longer decays geometrically by construction. It can still come out small if the branch Jacobian works against the identity: what is gone is the guarantee, not the smallness.

What follows is a change to what a block is for. Giles Thomas reads the residual stream through the Talmud, flagging that he is not Jewish and welcomes correction: a page is the core text ringed by commentary, and commentary on commentary, and the core text is never erased. A block does not hand back replacement vectors, it annotates a stream that keeps carrying the original.

    def forward(self, x):
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut

        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut
        return x

Ten lines, four of them the stash and the add. Same shape in and out, and the book's own sentence for the subtle half: the dimensions are unchanged and every output vector has been re-encoded to carry context from the whole sequence. One earlier sentence needs correcting against this listing. The book says layer normalisation is "typically applied before and after the multi-head attention module", but there is one norm before attention, one before the feed forward, and nothing on attention's output before the add.

06 · Where 163 million numbers live

Swap the placeholders for the real classes and count. 163,009,536 parameters, all accountable: token embedding 38,597,376, positional 786,432, twelve blocks at 7,085,568 each for 85,026,816, a final layer norm at 1,536, an output head of 38,597,376. The two embedding-shaped matrices are 77,194,752 of that, 47.4% of the model, and the twelve blocks everyone spends the time on are 52.2%.

Press the XL preset, then flip weight tying on and off:

XL is 1,637,792,000 untied and 1,557,380,800 tied, and the head that switch removes is 4.9% of the model there against 23.7% here, because a head grows with the width and a block with the width squared. Dragging n_layers moves the block segments and nothing else; emb_dim moves all six.

Weight tying is why two numbers name the same model. model.tok_emb.weight.shape and model.out_head.weight.shape are both torch.Size([50257, 768]), and subtracting the head turns 163,009,536 into 124,412,160, which is where the name comes from. Raschka builds it untied on experience rather than theory: separate layers train better for him. Giles Thomas supplies the missing reasoning: the vectors leaving the last block are not the ones that entered the first, so the projection out has no reason to be the transpose of the projection in. Chapter 6 ties them anyway, because OpenAI's checkpoint is tied, and that checkpoint also keeps the query, key and value biases this config turns off. Twelve blocks times three projections times 768, so the real gpt2 weights carry 27,648 more: 124,439,808 tied.

At four bytes a parameter the book prints 621.83 MB, which is 652,038,144 bytes divided by 1024 twice, so it is MiB by its own arithmetic. Two things fall out of the family table. head_dim is 64 in all four sizes, at 768/12, 1024/16, 1280/20 and 1600/25, so OpenAI scaled by adding heads and layers rather than widening them. And the book disagrees with itself about their names: chapter 4 says 345, 762 and 1,542 million, appendix B says 355 million, 774 million and 1.5 billion, and appendix B is the set the tied counts reproduce.

07 · Featureiman Byeswickattribute argue

With the block written, GPTModel is about twenty-five lines: embed, add positions, drop, twelve blocks, one norm, one projection. The bridge from logits to text is four operations.

        idx_cond = idx[:, -context_size:]
        logits = model(idx_cond)
        logits = logits[:, -1, :]
        probas = torch.softmax(logits, dim=-1)
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)

Crop to the last context_size tokens, take the last row of the logits and throw the other three away, softmax, argmax. The softmax is redundant: it is monotonic, so argmax(softmax(z)) and argmax(z) are the same index, and the book keeps it to show the full path and says so. That does real work, because it means the distribution over the vocabulary is not what picks the token here. It is what the token would be picked from if you were sampling. logits[:, -1, :] throws away three of the four rows, and the fix for that waste is to slice the hidden state before the output head rather than the logits after it, which is a different problem from the one KV caching solves.

Step it six times, and watch the two argmax markers stay on the same bar:

The two argmax markers land on the same index at every step, on a real forward pass through a real two-block GPT, which is the monotonicity claim made visible rather than asserted. Greedy decoding takes the top symbol every time, so the untrained model on the left usually finds two or three characters and cycles them.

The book's own run is beside it. "Hello, I am" encodes to [15496, 11, 314, 716], six greedy steps give [27018, 24086, 47843, 30961, 42348, 7267], and the ten ids decode to Hello, I am Featureiman Byeswickattribute argue. Figure 4.18 draws those six iterations arriving at "Hello, I am a model ready to help" instead. What came out is fragments: Byeswickattribute is three tokens landing next to each other rather than a word anything chose. Every structural thing in this model is right. The weights are 163,009,536 random numbers. A correct transformer and a language model are two different achievements, and only the first one is architecture.

08 · What's next

What this model is missing, in rough order of how much it hurts:

• A loss. Cross-entropy against the input shifted by one, the number every one of those 163 million parameters is going to be moved by. That plus an optimiser and a corpus is chapter 5, and the next post is the first half of it.
• Sampling that is not argmax. Greedy decoding takes the top token every time, which is why untrained gibberish is repetitive gibberish. Temperature and top-k turn a distribution into a choice.
• Everything here that is now historical. The normalisation, the activation, the positional table, the attention layout and the single dense feed forward have all been replaced since, and dropout has been dropped outright. That is the last post in the series. Raschka's position is that GPT-2 is worth implementing first, because the later changes are easier to understand once you have met the problems they solve.

Every figure here runs on a few hundred lines of dependency-free TypeScript, over the same Mat, softmaxStable and gelu the neural posts used. What cost the most time was not the model. It was the layer-norm figure: the book's demo and the LayerNorm class two paragraphs later use different variance conventions, and it had to reproduce both before its numbers agreed.