The Transformer

Residuals, LayerNorm, and FFN

The three extras that make deep stacks trainable

~11 minLesson 29 of 57Written by Claude Opus 4.8

A 96-layer GPT is just one block stacked ninety-six times. Three humble pieces of glue - a skip connection, a normalization, and a small MLP - are the entire reason that tower can be trained at all.

Attention gets the headlines, but on its own it cannot build a deep network. Stack a dozen raw sublayers and the signal either explodes to infinity or fades to nothing on the way up, and - worse - the gradient dies long before it reaches the early layers on the way back down. The fix is not a cleverer attention; it is three unglamorous pieces of plumbing that appear around every sublayer. Get them and depth becomes free. Miss them and nothing trains.

Why depth is the enemy

Each sublayer is roughly a matrix multiply. Chain $L$ of them and you are multiplying the signal by $L$ matrices in a row - if their typical gain is a hair above $1$ the magnitude grows like $gain^{L}$ and blows up; a hair below $1$ and it collapses to zero. The backward pass is the same product in reverse (Lesson 9): the gradient at layer $0$ is a long chain of Jacobians multiplied together, and a long product of small numbers underflows to nothing. This is the vanishing / exploding problem, and it is what the glue on this page exists to defeat.

The residual shortcut

The first piece is almost embarrassingly simple. Instead of replacing a token's representation with what a sublayer computes, you add the sublayer's output back onto the input:

y = x + Sublayer (x) .

The sublayer only has to learn a small correction to $x$ , never the whole thing from scratch - if it has nothing useful to add it can output zero and the block becomes the identity. But the real magic is in the derivative. Because differentiation is linear, the local Jacobian of that block is

The identity term is a gradient highway

\frac{\partial y}{\partial x} = I + \frac{\partial Sublayer ( x )}{\partial x} .

That $I$ is the whole game. In the backward pass the gradient reaching layer $0$ is the product of every block's Jacobian. With the identity term, each factor is $I + J_{ℓ}$ - so even if a sublayer's own Jacobian $J_{ℓ}$ shrinks toward zero, the gradient still passes through that block with a clean factor of $1$ . The skip connections form a DAG of shortcut edges from the loss straight back to the input; the gradient always has an uninterrupted road home. Delete the residuals and that road becomes a chain of tiny numbers multiplied together - which is exactly the vanishing gradient that kills deep training.

LayerNorm: reset every token to unit scale

Residuals rescue the gradient, but they let the forward signal drift - every add pushes the magnitude around. Layer normalization pins it back. For a single token vector $x \in R^{d}$ it computes the mean and variance across that token's own features, standardizes, then rescales with two learned vectors $γ$ and $β$ :

Normalize across features, one token at a time

μ = \frac{1}{d} i = 1 \sum d x_{i}, σ^{2} = \frac{1}{d} i = 1 \sum d (x_{i} - μ)^{2},

LN (x) = γ ⊙ \frac{x - μ}{σ ^{2} + ϵ} + β .

After the middle term every token has mean $0$ and variance $1$ - no matter how wild it arrived. The learned $γ$ (scale) and $β$ (shift) hand back the freedom to be non-standard if that helps. Crucially the statistics are computed per token, over the feature axis - never across the batch. That is why LayerNorm works with a batch size of $1$ and with variable-length sequences, where BatchNorm would choke.

See it break - and heal

Here is a stack of up to 12 sublayers with a real token vector flowing through it. The left panel is a live histogram of that vector's values after the last layer; the right panel is a schematic training-loss curve for the same configuration. Flip the two switches and watch which panel each one controls.

Ablation playground - toggle the glue, watch it explode or stabilize

stable - unit varianceexploded / divergedtraining loss

depth12 layers

12 layers - Residual ON - LayerNorm ON - act. std = 1 - loss 0.09

Both switches on: unit-variance activations and a smooth descent. This is a real Transformer block.

Two switches, two different jobs

Start with both switches off at depth 12. The histogram detonates - the standard deviation rockets past 20 and every value slams into the edge of the axis - while the loss curve refuses to descend and instead climbs off the top and NaNs out. Now flip LayerNorm on: the histogram snaps back to a tidy bell with std $= 1.00$ (that is literally all LayerNorm promises), yet the loss still crawls and plateaus high. Stable activations alone don't fix the gradient. Flip Residual on too and the loss finally dives smoothly toward zero. LayerNorm tames the forward distribution; residuals tame the backward gradient. You need both.

Instability is a disease of depth

Turn both switches off, then drag depth down to 1 or 2. The shallow network is fine - the bell barely spreads and the loss still descends. Crank depth back to 12 and the explosion returns. A 2-layer MLP genuinely doesn't need any of this; a 96-layer GPT cannot survive a single forward pass without it. That is why residuals and normalization went from "nice trick" to "non-negotiable" the moment networks got deep.

The feed-forward block

Attention lets tokens look at each other; the second sublayer in every Transformer block lets each token think on its own. It is a plain two-layer MLP applied position-wise - the same weights run independently on every token vector - that projects up to a wider hidden dimension (classically $4 d$ ), applies a nonlinearity, and projects back down:

FFN (x) = W_{2} ϕ (W_{1} x + b_{1}) + b_{2}, W_{1} \in R^{4 d \times d}, W_{2} \in R^{d \times 4 d} .

The nonlinearity $ϕ$ is almost always GELU - a smooth cousin of ReLU that gates each input by how likely it is under a standard normal, $GELU (x) = x Φ (x)$ , where $Φ$ is the Gaussian CDF. That wide-then-narrow MLP holds most of a Transformer's parameters and is where a great deal of its learned knowledge lives. Wrap attention and the FFN each in their own residual + LayerNorm and you have the whole block.

Pre-norm vs post-norm - where the LN sits

Post-norm (the original 2017 paper) normalizes after the residual add:

x_{ℓ + 1} = LN (x_{ℓ} + Sublayer (x_{ℓ})) .

Pre-norm (GPT-2 onward) normalizes the sublayer's input, leaving the residual stream itself untouched:

x_{ℓ + 1} = x_{ℓ} + Sublayer (LN (x_{ℓ})) .

Pre-norm keeps a completely clean identity path from the loss to layer $0$ - nothing is ever squashed on the highway - so it trains stably at great depth without the learning-rate warmup that post-norm needs. Nearly every modern LLM is pre-norm.

class TransformerBlock:
    def __init__(self, d, hidden=4*d):     # FFN hidden is 4x wider
        self.attn = MultiHeadAttention(d)
        self.ln1  = LayerNorm(d)          # one LN per sublayer
        self.ln2  = LayerNorm(d)
        self.W1, self.b1 = Linear(d, hidden)  # project up:   d  -> 4d
        self.W2, self.b2 = Linear(hidden, d)  # project down: 4d ->  d

    def ffn(self, x):
        return self.W2 @ gelu(self.W1 @ x + self.b1) + self.b2

    def forward(self, x):                   # x: one token's vector
        x = x + self.attn(self.ln1(x))      # residual around attention
        x = x + self.ffn(self.ln2(x))       # residual around the MLP
        return x                            # the residual stream is never overwritten

One pre-norm Transformer block, per token

Check yourself

Why does writing a block as x + Sublayer(x) stop gradients from vanishing in a deep stack?

What does LayerNorm compute its mean and variance over?

Recall: in a pre-norm block, write the two sub-updates that a token vector x passes through.

Lock in the full block structure. Try to state it, then check.

Lock it in

Writing each sublayer as $y = x + f (x)$ adds an identity term to the Jacobian, giving the gradient a clean highway from the loss all the way back to layer 0.
LayerNorm standardizes each token vector to mean 0 and variance 1 across its features - taming the forward signal so it never drifts with depth.
The feed-forward block is a position-wise MLP that projects up to $4 d$ , applies GELU, and projects back down - where most of the Transformer's parameters and learned knowledge live.
Pre-norm places LayerNorm inside the residual branch, keeping the highway completely clean - which is why nearly every modern LLM uses it.

Primary source

Read every equation on this page as runnable code in The Annotated Transformer (Harvard NLP) - its SublayerConnection is the residual + LayerNorm wrapper spelled out line by line, and its PositionwiseFeedForward is the FFN above. For the picture-first intuition of how these pieces sit inside the block, work through Jay Alammar's The Illustrated Transformer.

Ask your teacher

Want to see the backward pass through a residual worked out on a real graph, or why pre-norm removes the learning-rate warmup post-norm demands? Ask and I'll build the follow-up. The connection worth planting now: those residual connections form a DAG of skip edges, and the backward pass from Lesson 9 flows gradients along exactly those edges (the same graph traversal idea from BFS/DFS). The $+ I$ in each Jacobian is a shortcut edge straight home - which is the reason a network can be 96 layers deep and still learn.