6.7 Normalization layers

What this is and why it exists

Normalisation layers keep activations in a range where training works, and they are the reason very deep networks became routine rather than heroic. The choice between them is architectural rather than a matter of taste: convolutional networks use one, sequence architectures use another, and the reason is a dependence on batch size that decides the question. This topic also covers the failure that produces wrong answers with no error message at all — a layer that behaves differently in training and in evaluation, used in the wrong mode.

The vocabulary

• Normalisation — rescaling values to a standard centre and spread.
• Batch statistics — the mean and spread computed across the examples in the current batch.
• Running estimates — averages accumulated during training, kept for use at inference.
• Scale and shift — the two learned parameters that let a normalisation layer undo itself if that helps.
• Train mode and eval mode — the two states a module can be in.
• Internal covariate shift — the original explanation for why normalisation helps, now considered incomplete.

The mental model

The problem is that every layer's input distribution keeps moving. As the layers beneath it update, what arrives changes in centre and in spread, so each layer spends part of its effort chasing a moving target, and activations can drift into ranges where gradients vanish. Normalisation removes the drift: subtract a mean, divide by a spread, and hand the next layer something with a predictable shape. Then two learned parameters — a scale and a shift — let the network restore whatever distribution it actually wanted, so nothing is lost by normalising. The early explanation for the benefit was that it reduced shifting distributions between layers; the current understanding is that it also smooths the optimisation surface, allowing larger learning rates. Both stories agree on the practice.

Batch normalisation normalises each feature across the examples in the batch, which works very well for convolutional networks and carries a consequence that is the single most common source of confusing bugs in this module. At training time it uses the current batch's statistics. At inference you are frequently predicting for one example, and a batch of one has no meaningful spread — so the layer must use something else. PyTorch states the arrangement plainly: "during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation."

That dual behaviour is the trap. The layer only knows which set of statistics to use because the module has been put in the right mode. Forget to switch to evaluation mode and the layer normalises your test batch by its own statistics — so a prediction depends on which other examples happened to be alongside it, results change when the batch size changes, and single-example inference is nonsense. Nothing raises an error. The habit that prevents it is to switch modes explicitly at the top of each phase, every time, and to combine evaluation mode with switching off gradient recording:

model.train()
# ... one epoch of training ...

model.eval()
with torch.no_grad():
    # ... validation or inference ...
    ...

The second consequence of depending on the batch is small batches. With four examples, the computed mean and spread are noisy estimates, and that noise goes straight into every activation; below a certain size, batch normalisation actively hurts. It also complicates distributed training, where each device sees only its share of the batch unless the statistics are explicitly combined across devices. There is a setting that turns off the running estimates altogether — PyTorch notes that with it switched off "this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well" — which is occasionally useful and does not remove the underlying dependence.

Layer normalisation removes the batch from the picture entirely. It normalises across the features of each example independently, so one example is normalised the same way whether it arrives alone or among a thousand. There is no train-versus-eval distinction, no batch-size sensitivity, and no complication under distribution. That is the practical reason transformers use it, along with the fact that sequences have variable length, which makes batch statistics awkward in a second way. If you take one rule from this topic: normalise across the batch for convolutional vision work, and across the features for sequence architectures.

The variants are each a different choice of what to group over. Group normalisation splits the channels into groups and normalises within each group per example — designed for the case where you need convolutional behaviour but the batch is too small for batch normalisation, so detection and segmentation at high resolution use it. Instance normalisation normalises each channel of each example on its own, which removes per-example contrast and is standard in style transfer for that reason. RMSNorm is a stripped-down layer normalisation that rescales by the root mean square without subtracting a mean, which is slightly cheaper and works as well in practice, and is now common in large sequence models. Knowing which one a paper used saves you misreading its results, because these choices change how a result reproduces at a different scale.

Placement is the last detail worth knowing. In residual architectures, normalising before the block rather than after it makes deep stacks train more stably and reduces the need for a long warmup; normalising after was the original arrangement and needs more care. If you are reproducing a result, copy the placement exactly — it is not cosmetic.

What you should now be able to explain or do

Say what a normalisation layer computes and what the learned scale and shift are for. Describe batch normalisation's two modes and quote the arrangement in your own words. Recognise the symptoms of a model left in the wrong mode. Explain why small batches and distributed training complicate batch normalisation. Say what layer normalisation changes and why sequence architectures adopted it. Name the variants and the grouping each uses. State the placement rule for residual blocks and why copying it matters.

Check yourself