Normalisation and residuals
Two building blocks without which deep networks would not train: the skip connection and layer normalisation, and why their order matters.
The idea
Two problems appear as soon as a network gets deep. First, the gradient vanishes on the long way back. Second, the values in the middle layers drift apart until nothing usable is left.
The residual connection solves the first: it lays a shortcut around every block along which the gradient can flow back undamped. Normalisation solves the second: it pulls values back to a middle range after every layer.
The variants
| Method | Normalises over | Used in |
|---|---|---|
| BatchNorm | The batch, per channel | CNNs, fixed batch sizes |
| LayerNorm | All features of one example | Transformers |
| RMSNorm | Like LayerNorm, no mean subtraction | Current language models |
| GroupNorm | Channel groups per example | Vision models with small batches |
BatchNorm has a practical drawback often overlooked: it behaves differently in training and in inference, because inference uses running averages. A model that looks good in training and not in production frequently has its problem exactly here.
import numpy as np
def layernorm(x, eps=1e-5):
# Over the feature axis, separately per example.
mu = x.mean(axis=-1, keepdims=True)
var = x.var(axis=-1, keepdims=True)
return (x - mu) / np.sqrt(var + eps)
def rmsnorm(x, eps=1e-6):
# No mean subtraction: saves one reduction and works about as well.
return x / np.sqrt((x ** 2).mean(axis=-1, keepdims=True) + eps)The formulas
RMSNorm saves the mean and therefore one reduction over the feature axis. On large models that is a measurable percentage of runtime at no quality cost: which is why current models have switched almost uniformly.
Pre-norm versus post-norm
The difference is decisive. With pre-norm the residual path x → x stays free of
any normalisation, so the gradient runs undamped from the last layer to the
first. With post-norm a normalisation sits on that path at every step, which
makes training deep models unusably unstable without a carefully tuned warmup.
The price of pre-norm is that activation norms grow across layers, because something is added in every block. The usual remedy is a final normalisation before the output layer.
Why residuals rescue the gradient
Across L blocks that is the product Π(I + ∂F/∂x), and the all-ones term
survives throughout. Without the residual it would be Π(∂F/∂x), and that product
goes to zero for small factors. That is the entire mathematical content of the
skip connection, and it explains why networks with over a hundred layers have
been trainable since it was introduced.
Related courses and sources
Batch Normalization
Why normalising intermediate values is what makes deep training stable in the first place.
For anyone asking why normalisation is what makes deep training possible.
Deep Residual Learning
The shortcut across layers that made hundred-layer networks trainable. Present in every architecture today.
For understanding how networks were able to get deep at all.