4 — The transformer block

Chapter 3 ended with causal multi-head attention working. Each position can now look back at the tokens before it, decide which ones matter, and pull in a weighted mix of their vectors. " it" can finally gather something from " Time Traveller".

That is still a long way from predicting the next token. Three things are missing. First, attention only mixes. Its output at each position is a weighted average of value vectors followed by a linear projection, so it moves information between positions but does very little computation on what it gathers. Second, one round of mixing isn’t enough. GPT-2 does it twelve times, and a stack that deep has to stay numerically well-behaved and trainable. Third, the output is still a 768-number vector per position, not a score for each of the 50,257 tokens that could come next.

The transformer block solves the first two. It wraps attention with three parts: a layer norm that keeps the numbers in range, a small feed-forward network that computes on each position separately, and residual connections that let each part add to its input instead of replacing it. Twelve identical blocks, a final layer norm and an output head make the whole model. This chapter builds each part on numbers small enough to check, then assembles a model with exactly GPT-2 124M’s shape and counts its parameters.

All runs were on an Apple M3 Pro with 18 GB of memory, with Python 3.12 and PyTorch 2.15.0, on the CPU unless a run says MPS (the Mac’s GPU). Seeds are given with each run.

Layer norm, on four numbers

A vector passing through many layers can drift. Its numbers can grow large, shrink toward zero, or develop a few huge entries that swamp the rest. Each layer then sees inputs on a different scale from the ones it was tuned for, and training gets slow and unstable. Layer normalization fixes the scale of each vector before a layer uses it.

The recipe works on one vector at a time. Subtract the vector’s mean so its numbers center on zero. Divide by its standard deviation so they have a spread of one. Then multiply by a learned scale and add a learned shift, one of each per position in the vector, so the network can undo the normalization wherever that helps.

$\text{LayerNorm}(x)_i = \gamma_i \cdot \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta_i$

Here $\mu$ is the mean of the vector’s entries, $\sigma^2$ their variance, $\gamma$ and $\beta$ the learned scale and shift, and $\epsilon$ a tiny constant (GPT-2 uses $10^{-5}$) that stops a division by zero when every entry is equal.

Take the vector [2, 4, 6, 8] by hand. The mean is 5. The differences from the mean are −3, −1, 1, 3, whose squares average to (9 + 1 + 1 + 9) / 4 = 5. The standard deviation is √5 ≈ 2.2361. Dividing gives −1.3416, −0.4472, 0.4472, 1.3416. With the scale at its starting value of 1 and the shift at 0, that is the output.

In PyTorch:

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))   # learned scale, gamma
        self.bias = nn.Parameter(torch.zeros(d))    # learned shift, beta

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        return self.weight * (x - mean) / torch.sqrt(var + self.eps) + self.bias

x = torch.tensor([[2.0, 4.0, 6.0, 8.0],
                  [1.0, 1.0, 1.0, 50.0]])
print(LayerNorm(4)(x))

tensor([[-1.3416, -0.4472,  0.4472,  1.3416],
        [-0.5774, -0.5774, -0.5774,  1.7321]])

nn.Module is PyTorch’s base class for anything with learnable parts. Wrapping a tensor in nn.Parameter registers it as one of those parts, so training will update it. x.mean(dim=-1, keepdim=True) averages along the last dimension, the 4 numbers of each vector, and keeps that dimension with size 1 so the result lines up with x for the subtraction. unbiased=False divides by 4 rather than 3, which is the definition layer norm uses. The output matches PyTorch’s built-in nn.functional.layer_norm to every printed digit.

The second row shows what normalization does to an outlier. [1, 1, 1, 50] has a mean of 13.25 and one entry far from the others. After normalizing, the outlier is 1.73 and the rest are −0.58. Both rows now have mean 0 and variance 1, whatever their scale was before. Every vector reaching a layer has the same overall size, and only the pattern of its numbers carries information.

The word “layer” is historical: the statistics are computed across the features of one vector, not across the batch. Each token is normalized on its own, so a batch of one behaves the same as a batch of a hundred.

GPT-2 puts a layer norm at the input of each sub-layer, before attention and before the feed-forward network, and adds one more after the last block. The GPT-2 paper describes this as moving layer normalization “to the input of each sub-block”. The original transformer normalized after each sub-layer instead. Putting it first is called pre-norm, and it matters for the residual connections below.

GELU, a softer ReLU

A network made only of matrix multiplies and additions is linear, and any stack of linear layers collapses into a single linear layer. A nonlinearity between them is what lets depth add power. The classic choice is ReLU, which keeps positive numbers and sets negative ones to zero.

GPT-2 uses GELU, the Gaussian error linear unit. It multiplies its input by the probability that a standard normal variable falls below it:

$\text{GELU}(x) = x \cdot \Phi(x)$

where $\Phi$ is the standard normal cumulative distribution. For large positive $x$, $\Phi(x)$ is close to 1 and GELU passes $x$ through like ReLU. For large negative $x$, $\Phi(x)$ is close to 0 and the output is close to zero. In between it bends smoothly instead of snapping at zero. GPT-2’s code computes it with a tanh approximation, which is what this model uses too:

$\text{GELU}(x) \approx \tfrac{1}{2}\, x \left(1 + \tanh\!\left(\sqrt{2/\pi}\,\left(x + 0.044715\, x^3\right)\right)\right)$

class GELU(nn.Module):
    def forward(self, x):
        c = (2.0 / torch.pi) ** 0.5
        return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * x ** 3)))

Side by side on a few inputs:

These match PyTorch’s nn.functional.gelu(x, approximate="tanh") exactly.

Two differences matter. GELU is smooth, so its slope changes gradually instead of jumping from 0 to 1, which gives the optimizer a gradient everywhere. And it lets small negative inputs through as small negative outputs, dipping to about −0.17 near x = −0.75, so a unit that is slightly negative still passes some signal and still receives some gradient. A ReLU unit stuck on the negative side gets exactly zero gradient and can stop learning entirely. The GELU paper reports it matching or beating ReLU on a range of tasks. It is a modest improvement, not a requirement; a GPT with ReLU would still work.

The feed-forward network

Attention moves information between positions. The feed-forward network works on each position alone. It takes one 768-number vector, expands it to 3,072 numbers, applies GELU, and projects back down to 768:

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.fc = nn.Linear(cfg.emb_dim, 4 * cfg.emb_dim)
        self.act = GELU()
        self.proj = nn.Linear(4 * cfg.emb_dim, cfg.emb_dim)

    def forward(self, x):
        return self.proj(self.act(self.fc(x)))

nn.Linear(a, b) is a weight matrix and a bias that turn a vector of a numbers into one of b numbers. Applied to a tensor of shape (batch, tokens, 768), it acts on the last dimension and treats every token independently, so the same small network runs at every position with the same weights.

The factor of four is GPT-2’s choice, inherited from the original transformer, which used 512 and 2,048. The wide middle layer gives the network room: 3,072 GELU units, each of which can switch on for some pattern in the vector it sees. A useful way to think of it is that attention decides what to gather and the feed-forward network decides what to make of it.

It is also where most of a block’s parameters live. The expanding layer has 768 × 3,072 weights plus 3,072 biases, and the contracting layer 3,072 × 768 weights plus 768 biases, for 4,722,432 in all. The attention layer from chapter 3, with its 768 × 2,304 query-key-value matrix and 768 × 768 output projection plus biases, has 2,362,368. The feed-forward network is two-thirds of the block.

Residual connections

The last part is the simplest to write and the most important. Instead of replacing its input with its output, each sub-layer adds its output to its input:

x = x + attn(norm1(x))
x = x + ffn(norm2(x))

The vector that flows down through the blocks is called the residual stream. Each attention layer and each feed-forward network reads a normalized copy of it, computes something, and writes the result back by addition. Nothing is ever overwritten. A sub-layer that has nothing useful to add can output something close to zero and leave the stream as it was.

The idea comes from residual networks for images, where it made networks of a hundred layers trainable. The usual explanation is about gradients. When training adjusts a weight in block 1, the signal telling it how has to travel backward through every block above. Through a plain stack, that signal is multiplied by each layer’s local slope in turn and can shrink to nothing or blow up. Through x + f(x), the derivative is 1 plus the derivative of f, so there is always a direct path, the ”+ 1”, that carries the gradient back unchanged.

To see what actually happens in this model, a 12-block version can be built twice from seed 123, once with residual connections and once with each block computing x = attn(norm1(x)); x = ffn(norm2(x)) instead, and given one batch of random tokens. The gradient norm of each block’s first feed-forward weight, block 1 to block 12, after one backward pass:

(To keep this quick, the blocks here are 256 numbers wide with 4 heads, the batch is 4 random sequences of 64 tokens, and dropout is off.)

The gradients did not vanish without the residuals. The layer norms in every block rescale each vector back to a standard size, so the signal doesn’t shrink layer by layer the way it does in an old-style deep network. On gradient size alone, the plain stack looks fine.

Something else breaks. Comparing the 64 output vectors of one sequence to each other, with the cosine similarity used in chapter 2, a value of 1 meaning they point the same way:

Without residuals, every position comes out of the stack as the same vector, to three decimal places. Each attention layer averages over earlier positions, and each average makes the vectors more alike. With nothing carrying the original token through, twelve rounds of averaging erase the difference between positions, so every position in a sequence ends up making the same prediction. Gradients still flow, but there is nothing position-specific left for them to adjust.

Training shows the consequence. A smaller model of the same design (4 blocks, 256-number vectors, 4 heads, a context of 128 tokens, dropout 0.1) was trained on The Time Machine from chapter 2 for 10 passes over the book on MPS, once with residuals and once without, seed 123, with everything else equal. Chapter 5 explains the loss and the training loop; for now, lower is better.

The dotted line is 6.346, the entropy of the book’s token frequencies: the loss of a model that ignores context entirely and always predicts each token at its overall rate (" the" 4.2% of the time, ”,” 5.0%, and so on). The plain model settles right on that line. It has learned which tokens are common and almost nothing else: by the time its 4 blocks are done, too little about the context survives to the output to predict from. The residual model goes straight through that floor. Doubling the depth to 8 blocks gives the same picture: with residuals the training loss after 10 epochs is 4.823, without them 6.369 again.

So residual connections do two jobs here. They give gradients a direct path, which matters more in deeper and un-normalized networks, and they keep each token’s own information in the stream so that every layer can refine it instead of replacing it. Pre-norm keeps the path clean: the layer norm sits on the branch, not on the stream itself, so the ”+ x” path runs from the embeddings to the output with nothing in the way.

One block

With all four parts, a block is short:

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = LayerNorm(cfg.emb_dim)
        self.attn = CausalSelfAttention(cfg.emb_dim, cfg.n_heads, cfg.context_length,
                                        cfg.dropout, cfg.qkv_bias)
        self.norm2 = LayerNorm(cfg.emb_dim)
        self.ffn = FeedForward(cfg)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x):
        x = x + self.drop(self.attn(self.norm1(x)))
        x = x + self.drop(self.ffn(self.norm2(x)))
        return x

CausalSelfAttention is chapter 3’s module, unchanged: one fused linear layer for queries, keys and values, the causal mask, softmax, dropout on the attention weights, and an output projection. The block adds dropout on each branch’s output before it joins the stream, which is where GPT-2 applies its “residual dropout”. nn.Dropout(p) sets a random fraction p of its input to zero during training and scales the rest up by 1/(1 − p) to keep the average the same. In evaluation mode it does nothing.

A block takes a tensor of shape (batch, tokens, 768) and returns the same shape. That is what makes stacking trivial: the output of one block is a valid input to the next.

The whole model

The model is chapter 2’s two embedding tables, a stack of blocks, a final layer norm and the output head, a linear layer that turns each 768-number vector into 50,257 scores, one per token in the vocabulary. These scores are called logits. They are not probabilities yet; softmax turns them into probabilities, and the next chapter does exactly that.

The sizes go in one small config object:

from dataclasses import dataclass

@dataclass
class GPTConfig:
    vocab_size: int = 50257      # GPT-2 BPE vocabulary
    context_length: int = 1024   # longest sequence the position table covers
    emb_dim: int = 768           # width of every token vector
    n_layers: int = 12           # transformer blocks
    n_heads: int = 12            # attention heads per block
    dropout: float = 0.1         # used only while training
    qkv_bias: bool = True        # GPT-2 has biases on the q, k, v projections
    tie_weights: bool = False    # share token embedding and output head?

A @dataclass is a Python class that is mostly fields with defaults; GPTConfig() gives the GPT-2 124M sizes, and GPTConfig(emb_dim=256, n_layers=4) changes just those two.

class GPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.emb_dim)
        self.pos_emb = nn.Embedding(cfg.context_length, cfg.emb_dim)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.Sequential(*[Block(cfg) for _ in range(cfg.n_layers)])
        self.final_norm = LayerNorm(cfg.emb_dim)
        self.head = nn.Linear(cfg.emb_dim, cfg.vocab_size, bias=False)
        if cfg.tie_weights:
            self.head.weight = self.tok_emb.weight
        self.apply(self._init)

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.zeros_(m.bias)

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        x = self.blocks(x)
        return self.head(self.final_norm(x))

nn.Sequential chains modules so that calling it runs them in order; the *[...] unpacks a list of 12 freshly made blocks, each with its own weights. torch.arange(T) is the tensor [0, 1, …, T − 1], the positions to look up. self.apply(fn) calls fn on every submodule, which here sets every weight to small random numbers with a standard deviation of 0.02 and every bias to zero, the scale GPT-2 starts from. (GPT-2 also shrinks the weights of the layers that write into the residual stream by a factor depending on depth; this model skips that, which makes no difference once real weights are loaded in chapter 6.)

Two short sentences, four tokens each, go through as a batch:

import tiktoken
enc = tiktoken.get_encoding("gpt2")

torch.manual_seed(123)
model = GPT(GPTConfig())
batch = torch.tensor([enc.encode("The cat sat on"),
                      enc.encode("A small model learns")])
print(batch)
model.eval()
with torch.no_grad():
    logits = model(batch)
print(logits.shape)

tensor([[  464,  3797,  3332,   319],
        [   32,  1402,  2746, 22974]])
torch.Size([2, 4, 50257])

model.eval() switches dropout off. torch.no_grad() tells PyTorch not to record the computation for training, which saves memory when only the output is wanted. In goes a (2, 4) tensor of token IDs; out comes (2, 4, 50,257): for each of the 2 sequences and each of the 4 positions, a score for every token in the vocabulary. The scores at position 3 of the first row are the model’s guess for what follows “The cat sat on”. Because of the causal mask, the scores at position 1 are its guess for what follows “The cat” and could not have looked at " sat" or " on". One forward pass makes a prediction at every position at once, which is what makes training efficient.

Counting the parameters

GPT-2’s smallest model is known as “124M”. Counting this model’s parameters one part at a time shows where every one of them is.

PyTorch agrees:

def count_params(model):
    return sum(p.numel() for p in model.parameters())

print(count_params(GPT(GPTConfig())))
print(count_params(GPT(GPTConfig(tie_weights=True))))

163037184
124439808

p.numel() is the number of elements in a tensor, and model.parameters() yields every parameter once.

The untied model has 163 million parameters, not 124 million. The difference, 38,597,376, is exactly one embedding-sized matrix. GPT-2 doesn’t have a separate output head. Its last step, in OpenAI’s model code, multiplies the final vectors by the transpose of the token embedding table wte. The same 50,257 × 768 matrix that turns token IDs into vectors at the bottom turns vectors back into token scores at the top. This is weight tying, and it makes sense: a token’s input vector and the direction that makes the model predict that token are naturally related, so one matrix can serve both. With tie_weights=True, the line self.head.weight = self.tok_emb.weight makes the head use the very same tensor, and model.parameters() counts it once. The total becomes 124,439,808, which is where the “124M” name comes from. (The GPT-2 paper’s table lists the smallest model as 117M; the 124M figure is the actual count of the released weights, which chapter 6 loads and checks.)

At 4 bytes per 32-bit float, the tied model is 497,759,232 bytes, just under 500 MB. The untied version would be 652 MB.

Notice how much of the model sits in the two ends. The token embedding alone is 31% of the tied model. That is a consequence of a 50,257-token vocabulary at a modest width; larger models spend a much smaller fraction on it because their blocks grow faster than their vocabulary. Inside the blocks, two-thirds of every block is the feed-forward network.

The model in this track keeps tie_weights as an option. Chapter 5 trains a small model from scratch with separate head weights, and chapter 6 turns tying on to load GPT-2.

Asking an untrained model to write

The model can already produce text. To generate, give it some tokens, take the logits at the last position, pick a token, append it, and repeat:

@torch.no_grad()
def generate(model, idx, max_new_tokens):
    model.eval()
    for _ in range(max_new_tokens):
        logits = model(idx[:, -model.cfg.context_length:])
        next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        idx = torch.cat([idx, next_id], dim=1)
    return idx

idx[:, -context_length:] keeps at most the last 1,024 tokens, because the position table has no row past that. logits[:, -1, :] takes only the last position’s scores, since the earlier positions predict tokens that are already known. argmax picks the highest-scoring token, a choice called greedy decoding. torch.cat joins the new token onto the end of the sequence along the token dimension. Every step re-runs the whole sequence through the model, which is wasteful; chapter 9 fixes that with a cache.

With the randomly initialized model from seed 123, twelve new tokens after two prompts:

The Time Traveller preventionの� Const preventionの� needy Chris IV DOJanton Richard doubt
Attention is retrospective programmer programmer programmer programmerLinkedIn programmer programmer regress programmer programmer programmer

Gibberish, as expected. Every weight is random, so the scores are noise, and greedy decoding follows that noise wherever it leads, sometimes into a loop of one token. (The � is half of a multi-byte character: the model picked a byte-level token that only makes sense next to its partner, as chapter 1 described.) The probabilities behind the choices are nearly flat. For “The cat sat on”, the most likely next token gets a probability of 0.000166, against 0.0000199 for a perfectly uniform guess, and the entropy of the whole distribution is 10.67 nats. A model that knew nothing at all, giving every token equal probability, would have an entropy of ln 50,257 ≈ 10.82; this one is barely more decided than that.

Where this leaves things

The architecture is complete. Token IDs go in; each block lets every position gather from earlier positions with attention, process what it gathered with a feed-forward network, and add both results into a residual stream kept in range by layer norm; twelve blocks later, an output head scores every token in the vocabulary. With weight tying, it is GPT-2 124M’s shape down to the last parameter: 124,439,808 of them.

And all 124 million are random. The model has the machinery to predict the next token but no knowledge of which token is likely. The scores it produces are noise, and the text it writes is noise.

What’s missing is a way to measure how wrong those scores are and to change the weights so they become less wrong. That is a loss function and a training loop, and chapter 5 builds both, trains a small version of this model on The Time Machine, and watches it go from noise to something that reads like the book, and then too much like it.