Part III — The GPT

The transformer block — norm, GELU, residuals

Lesson 06 finished attention: a causal, multi-head layer that lets each token gather information from every earlier token, in several views at once. But look closely at what it does — and does not — do. It mixes vectors across positions with weighted sums; it never pushes a single token's vector through a nonlinear transform of its own, and if you naively stack a dozen of these layers the whole thing refuses to train. This lesson wraps attention in the three pieces that fix both problems and turn it into a unit you can stack: a feed-forward MLP, LayerNorm, and residual shortcuts.

1 · The piece attention is missing: a per-position MLP

Re-read the attention output from lesson 06 with a skeptical eye. For each position it returns z_i = Σ_j α_ij v_j — a weighted average of value vectors. Averaging is a linear operation. Stack two attention layers and, between the softmaxes, you are still only ever taking linear combinations of linear projections; there is no place where a token's own vector is bent through a nonlinearity. Attention answers "whom should I listen to?" brilliantly, but it never answers "now that I have listened, what should I compute?" That second question is per-position work, and it needs its own module.

The module is a small two-layer neural network — a feed-forward MLP — applied to each token vector independently and identically. It expands the vector to a wider hidden space, applies a nonlinearity, and contracts it back:

class FeedForward:            # applied to every position, same weights
    W1 = Linear(d,   4*d)     # expand:  768 -> 3072
    act = GELU()              # the nonlinearity
    W2 = Linear(4*d, d)       # contract: 3072 -> 768
    forward(x) = W2(act(W1(x)))

Why widen to 4d?

The expansion is where the nonlinear computation happens. A nonlinearity applied directly at width d can only carve d features; projecting up to 4d first gives GELU a much larger space to bend, letting the layer represent far richer per-position functions before compressing the result back down to the residual width. The factor 4× is the convention GPT-2 and most transformers use — big enough to add real capacity, small enough to stay affordable. The contraction back to d is mandatory, not optional: the block must return a (B,T,d) tensor of the same shape it received, or it could not be stacked (§4).

Why GELU, not ReLU?

The classic choice is ReLU — max(0, x) — which hard-clips every negative input to exactly zero. That has a nasty side effect: any neuron pushed slightly negative gets a gradient of exactly zero and stops learning (a "dead" unit), and the function has a sharp corner at the origin that makes optimization bumpier. GELU (Gaussian Error Linear Unit) is a smooth gate instead: it multiplies the input by the probability that a standard-normal draw falls below it,

where Φ is the standard-normal CDF. (GPT-2 was trained with exactly this tanh approximation, so we match it.) The payoff is a curve that is smooth everywhere and, crucially, has a small nonzero gradient for almost all negative inputs — it only fully zeroes out near x ≈ −0.75. So a neuron nudged negative can still recover and contribute, and the smooth shape gives the optimizer cleaner slopes to follow. It is one of the small, quiet changes that make deep transformers train more reliably than they would with ReLU.

2 · Depth breaks two ways

Now try to make the network deep — which we must, since one block is far too shallow to model language. Two independent failures appear, and each has its own fix.

Failure A: activations drift → LayerNorm

Every sublayer transforms its input, and the statistics of the output almost never match the input. Feed a layer a batch of vectors with some mean and variance and its outputs land at a different mean and variance; the next layer shifts them again. Over a dozen layers the scale can snowball toward huge values or collapse toward zero, and softmax, GELU, and the gradients all misbehave outside their healthy range. The training dynamics become unstable and the loss struggles to fall.

The fix is LayerNorm: before a sublayer acts, standardize each token vector across its own feature dimension so it has mean 0 and variance 1, then re-scale and shift it with two learned vectors γ (scale) and β (shift):

class LayerNorm:
    gamma = ones(d)                          # learned scale
    beta  = zeros(d)                         # learned shift
    eps   = 1e-5                             # keeps the denom finite
    forward(x):
        mu  = mean(x, axis=features)         # per-token, over the d dims
        var = var(x,  axis=features)         # biased (÷ n), GPT-2 style
        return gamma * (x - mu)/sqrt(var+eps) + beta

Three details earn their place. Normalization is over the feature axis (each token independently), not the batch axis — so it behaves identically whether the batch has 1 sequence or 1,000, which is exactly what a language model needs. The ε guards against dividing by zero when a vector is nearly constant. And γ, β restore expressive freedom: forcing every vector to unit variance would be too rigid, so the model is allowed to learn the scale and offset that actually help — including, if it wants, to partly undo the normalization. (We use the biased variance, dividing by n rather than n−1, to match GPT-2's implementation so pretrained weights drop in cleanly later.)

Failure B: gradients vanish → residual shortcuts

The second failure is about the backward pass. In a plain deep stack, the gradient reaching layer ℓ is a long product of the per-layer Jacobians above it. If those factors are even slightly below 1 on average, their product shrinks geometrically, and by the time the signal reaches the first few layers it is effectively zero — the vanishing gradient problem. The book makes this vivid with a five-layer toy net: without shortcuts the mean gradient at the input layer is about 0.0002, four orders of magnitude smaller than at the top; the early layers simply do not learn.

The fix is the residual (shortcut / skip) connection: instead of replacing the input with the sublayer's output, add the sublayer's output to the input.

Each sublayer now computes only a correction to the running vector — the "residual stream" that flows straight up the network. The backward consequence is the whole point: differentiating x + f(x) gives 1 + f'(x), so there is always a clean +1 term — an identity highway — through which gradient flows to the bottom undiminished, no matter how small f' becomes. Add shortcuts to that same toy net and the input-layer gradient jumps from 0.0002 to around 0.22: the early layers come alive. Residual connections are what make depth trainable at all.

3 · Where the norm goes: pre-norm vs post-norm

LayerNorm and the residual can be wired in two orders, and the choice matters more than it looks. The original 2017 transformer put the norm after the sublayer and the addition — post-norm: x = LN(x + f(x)). GPT-2 and essentially every modern LLM put the norm inside the branch, before the sublayer, leaving the residual addition un-normalized — pre-norm:

The difference is what sits on the residual highway. In pre-norm the raw x passes to the next layer completely untouched — the identity path is truly clean, so the +1 gradient term survives all the way down. In post-norm every addition is immediately squashed by a LayerNorm, which repeatedly rescales the residual stream and re-introduces the shrinking-product problem the shortcut was meant to solve. The practical result: post-norm transformers need a careful learning-rate warmup and are finicky to train deep, while pre-norm transformers train stably from step one without warmup. That is why we normalize inside the branch. Everything below uses pre-norm.

4 · The full block

Assemble the pieces. A transformer block is two pre-norm residual sublayers stacked in series: first causal multi-head attention (the cross-position mixer from 06), then the feed-forward MLP (the per-position processor from §1). A dropout sits on each branch for regularization. Each sublayer follows the identical template x = x + dropout(sublayer(LN(x))):

class TransformerBlock:
    attn  = CausalMultiHeadAttention(d, heads, ...)   # from lesson 06
    ff    = FeedForward(d)                             # from section 1
    norm1 = LayerNorm(d);  norm2 = LayerNorm(d)
    drop  = Dropout(p)

    forward(x):                       # x : (B, T, d)
        x = x + drop(attn(norm1(x)))  # sublayer 1: mix across positions
        x = x + drop(ff(  norm2(x)))  # sublayer 2: transform per position
        return x                      # -> (B, T, d), same shape

The signature is the headline: (B, T, d) → (B, T, d). The block re-encodes each token's vector to fold in context and per-position computation, but it does not change the sequence length or the width. That invariance is exactly what makes the block a Lego brick — because the output is shape-compatible with the input, you can feed a block's output straight into another block, and another, with nothing to adjust between them. GPT-2 small stacks this block 12 times; the 1.5-billion-parameter GPT-2 stacks the same brick 48 times. Building the block right is building the whole model minus its two end-caps, which is precisely lesson 08's job.

5 · Watch depth live or die

The widget runs a deterministic deep stack of L toy layers and tracks two numbers as the signal moves up: the activation magnitude at each layer (forward pass) and a proxy for the gradient reaching layer 0 (backward pass). Each toy layer multiplies by a fixed factor slightly off from 1 and adds a fixed nonlinearity — no randomness, just the geometric-product effect that real depth suffers. Toggle residual and LayerNorm independently. With neither, watch the activation curve explode or vanish and the layer-0 gradient collapse to ~0 (the stack is dead). Turn on LayerNorm and the forward curve flattens; turn on residual and the gradient reaches the bottom. With both — the real block — the stack stays healthy at any depth. The bar underneath shows the per-block parameter split: attention 4d² versus MLP 8d².

Failure modes & checklist

Checkpoint

Where this points next

You now hold the transformer block: a self-contained unit that mixes context (attention), transforms each position (the GELU MLP), and — thanks to LayerNorm and residual shortcuts wired pre-norm — stays stable and trainable no matter how deep the stack. But a single block is shallow, and its output is still a tensor of d-dimensional vectors, not a next-token prediction. Two things are missing: depth (stack L blocks) and a way back to words (a final norm and a linear head that maps each vector to V logits over the vocabulary). Add those and run the autoregressive loop from lesson 00 for real. That is lesson 08, Assembling GPT & generating text, where the brick becomes a model.

Interview prompts

• What can the feed-forward MLP do that attention cannot, and where do most of a block's parameters live? (§1 — the MLP does nonlinear, per-position work (attention only takes linear weighted averages across positions), and it holds the bulk of the weights — 8d² vs attention's 4d² — acting as the block's memory while attention is the cheaper router.)
• Why does the MLP expand to 4d and then contract? (§1 — the wide hidden space gives the nonlinearity room to represent richer per-position functions; the contraction restores width d so the residual add and stacking work.)
• Why GELU instead of ReLU? (§1 — GELU is smooth and keeps a small nonzero gradient for almost all negative inputs, so neurons don't "die"; GPT-2 used its tanh approximation.)
• Name the two ways depth breaks and the fix for each. (§2 — activations drift in scale → LayerNorm; gradients vanish backward → residual shortcuts. Each cures one failure.)
• How does a residual connection help the backward pass? (§2 — x + f(x) differentiates to 1 + f'(x); the constant +1 is an identity path that carries gradient to early layers regardless of f'.)
• What is pre-norm vs post-norm, and why prefer pre-norm? (§3 — pre-norm normalizes inside the branch (x + f(LN(x))) leaving a clean residual highway that trains stably without warmup; post-norm (LN(x+f(x))) squashes the highway and is finicky at depth.)
• What is the input/output shape of a block and why does it matter? (§4 — (B,T,d) → (B,T,d); shape invariance is what lets identical blocks stack with nothing to adjust between them.)

Companion reads: cs336 · 02 The Transformer derives pre-norm and counts the block's parameters and FLOPs in detail; gpt_mini · 01 Architecture shows this exact block as runnable code.