LM4-T2 — Layer Normalization and Residual Connections

Learning objectives

After this lesson, you should be able to:

• Compute layer normalisation by hand, including scale, shift and the eps constant.
• Explain how layer norm differs from batch norm and why Transformers use it.
• Show with a derivative why a residual connection gives gradients a direct path.
• Distinguish pre-LN from post-LN blocks and say which GPT-2 uses.

The idea: keep the scale steady, keep a shortcut open

Layer normalisation. Take the vector of d numbers for one token. Subtract its mean so that it is centred, divide by its standard deviation so that its typical size is 1. Then let the model undo some of this if it wants, with a learned scale and shift per feature. The next layer therefore sees inputs of a predictable size, however large the previous layers made them. Crucially each token is normalised on its own, using only its own d numbers.

Residual connections. Instead of computing y = f(x), compute y = x + f(x). The layer only has to learn a correction to its input. If a layer has nothing useful to add, it can learn f ≈ 0 and pass the input through unchanged. A block that starts out near the identity is easy to stack.

Why they help: a signal view

Lab LM4-L2 stacks twelve random layers. Without help, the typical size of the activations is multiplied by a constant at every layer: a factor of 0.7 per layer leaves 0.014 of the signal after 12 layers, a factor 1.3 leaves 23 times too much. Layer norm removes the exponential drift. Residual connections keep the input visible in every later layer. Together they give a stack that can be made deep without the initial scale being tuned.

Why they help: a gradient view

For y = x + f(x) the derivative of the output with respect to the input is the identity matrix plus the derivative of f. When many layers are stacked, the gradient reaching an early layer is a sum of paths, and one of them is the pure identity path that crosses no layer at all. That path cannot vanish, so early layers always receive a learning signal.

Layer norm versus batch norm

Batch normalisation, common in image networks, computes statistics across the examples of a batch. For text that is awkward: sequence lengths vary, padding pollutes statistics, and generation runs with a batch of one. Layer norm computes statistics across the features of one token, so its result does not depend on the other examples or on the batch size, and it behaves identically in training and in generation.

The mathematics: normalise, scale, shift, add back

For one token with feature vector x = (x₁, …, x_d):

μ, σ²

mean and biased variance of the d entries of this one vector (divide by d, not d − 1).

ε

a small constant, typically 10⁻⁵, that prevents division by zero when all entries are equal.

γ, β

learned scale and shift, one pair per feature (2d parameters in total); initialised to 1 and 0.

y

the output, same shape as x. Before γ and β it has mean 0 and variance about 1.

For a tensor of shape (B, n, d) the statistics are taken over the last axis only, giving B · n separate normalisations.

Residual connection. For a sublayer f:

where I is the identity matrix. For N stacked residual blocks the end-to-end Jacobian is (I + J_N) ⋯ (I + J₁), which expands into the identity plus products of the individual Jacobians; the identity term is always there.

Two placements. With f an attention or feed-forward sublayer:

GPT-2 and nearly all later GPT-style models use pre-LN, plus one more LN at the very end of the stack. In pre-LN the residual stream is never normalised inside the stack, so its scale grows with depth (roughly like the square root of the number of blocks at initialisation); the final LN restores a fixed scale before the output head. In post-LN the norm sits on the main path; very deep post-LN models are harder to start training and usually need a warm-up of the learning rate.

Parameter count. One layer norm has 2d parameters. A block has two (4d), and the model has one more at the end: for GPT-2 small, 12 · 3,072 + 1,536 = 38,400 parameters, 0.03% of the model.

A worked example: normalising four numbers

Let x = [2, 4, 6, 8], d = 4, ε = 10⁻⁵.

Mean: μ = (2 + 4 + 6 + 8)/4 = 5. Deviations: [−3, −1, 1, 3]. Biased variance: σ² = (9 + 1 + 1 + 9)/4 = 5. The standard deviation including ε is √5.00001 ≈ 2.23607. Normalised: [−3, −1, 1, 3]/2.23607 = [−1.3416, −0.4472, 0.4472, 1.3416]. Check: the mean is 0, and the average of the squares is (1.8 + 0.2 + 0.2 + 1.8)/4 = 1.

With learned γ = [1, 1, 2, 2] and β = [0, 0, 1, 1]: y = [−1.3416, −0.4472, 2·0.4472 + 1, 2·1.3416 + 1] = [−1.3416, −0.4472, 1.8944, 3.6833]. The model is free to give some features a larger scale or an offset.

Using the unbiased variance (divide by 3) would give σ² = 6.667 and a slightly different answer, [−1.162, −0.387, 0.387, 1.162]. This is the most common reimplementation error: PyTorch’s LayerNorm and GPT-2 use the biased version.

A residual block with numbers. If the input is x = [1, −1] and the sublayer outputs f(x) = [0.1, 0.2], then y = [1.1, −0.8]: the block nudges the input by a small amount instead of replacing it.

import torch
import torch.nn as nn
import torch.nn.functional as F

x = torch.tensor([[2.0, 4.0, 6.0, 8.0]])
mean = x.mean(-1, keepdim=True)
var = x.var(-1, keepdim=True, unbiased=False)                 # biased variance, as in GPT-2
y = (x - mean) / torch.sqrt(var + 1e-5)
print([round(v, 4) for v in y[0].tolist()])                   # [-1.3416, -0.4472, 0.4472, 1.3416]
print(torch.allclose(y, F.layer_norm(x, (4,), eps=1e-5)))     # True: same as PyTorch

# a residual block whose sublayer outputs zero returns its input exactly
lin = nn.Linear(4, 4)
nn.init.zeros_(lin.weight); nn.init.zeros_(lin.bias)
block_out = x + lin(F.layer_norm(x, (4,)))                    # pre-LN: x + f(LN(x))
print(torch.equal(block_out, x))                              # True

# gradient through a skip connection: d(x + f(x))/dx contains the identity
w = torch.tensor([[0.5, 0.0], [0.0, 0.5]])
v = torch.tensor([1.0, -1.0], requires_grad=True)
out = v + w @ v                                               # f(v) = W v
out.sum().backward()
print(v.grad.tolist())                                        # [1.5, 1.5] = 1 (skip) + 0.5 (through W)

The last lines show the Jacobian argument in miniature: the gradient is 1.5 = 1 (the skip path) + 0.5 (through W). If W were tiny, the gradient would still be about 1.

Engineering implications and limitations

Epsilon placement. The ε goes inside the square root, next to the variance. Putting it outside, or omitting it, changes results slightly and can divide by zero for constant vectors.

Normalise over the last axis. A frequent bug is dim=0 or dim=1. The check y.mean(-1) ≈ 0 and y.var(-1, unbiased=False) ≈ 1 catches it.

Use the built-in in production. nn.LayerNorm(d) is faster and numerically careful (it accumulates in higher precision when training in half precision). Writing your own is for learning, as in the notebook.

The final norm matters. In pre-LN models the residual stream has grown by the end of the stack; leaving out the final norm makes logits badly scaled.

Dropout on the branch. Dropout is applied to the sublayer output before it is added to the stream, not to the stream itself, so the skip path remains intact. In large pretraining runs dropout is often set to zero.

Zero-initialised branches. Setting the last projection of each sublayer to zero (or small) makes every block start as an identity, which makes the start of training stable; some recipes do this deliberately.

Variants. Newer models use RMSNorm, which drops the mean subtraction and the shift, saving a little compute (Module 10). The pre-LN structure stays the same.

Common misconceptions

Engineering takeaways

1. Layer norm: y = γ (x − μ)/√(σ² + ε) + β, with μ and σ² over the d features of one token and σ² biased.
2. Residual connections compute y = x + f(x), so ∂y/∂x = I + ∂f/∂x and gradients always have an identity path.
3. Pre-LN, y = x + f(LN(x)), is used by GPT-2; post-LN, y = LN(x + f(x)), was the original design.
4. The residual stream grows with depth in pre-LN, so a final layer norm precedes the output head.
5. A layer norm has 2d parameters; the whole GPT-2 small has only 38,400 of them.

References and further reading

See Chapter 4 of the book (normalizing activations with layer normalization and adding shortcut connections) for the implementation used in this track.

• Ba, J. L., Kiros, J. R. and Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450.
• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
• Ba et al. (2016), Layer Normalization (research paper)
• Build a Large Language Model (From Scratch), Chapter 4: Implementing a GPT Model from Scratch (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: the feed-forward layer and its GELU nonlinearity, which hold most of a block’s parameters.

Related content

Nothing below is required to finish this lesson.

• Before
  LM4-T1: The GPT Architecture

  Describe the decoder-only stack from token IDs to next-token logits.

• Lab
  LM4-L2: Layer Normalization and Residual Numerics

  Implement layer norm and a residual update and see why activations stay stable.