Batch normalization and layer normalization: Stabilising training

In short: Normalisation layers scale intermediate results in the network to a stable range; batch norm is typical for image networks, layer norm for transformers.

In more detail: During training the values a layer receives constantly change because the layers before it change. This makes learning slow and unstable. Normalisation keeps the mean and spread of the activations in a favourable range.

In Depth

Batch normalization

For each feature the mean μ and standard deviation σ are computed over the mini-batch and subtracted or divided:

x̂ = (x − μ) / √(σ² + ε)
y = γ · x̂ + β

γ and β are learnable parameters that can restore the scaling. Advantages: faster training, higher learning rates possible, slight regularisation. At use (inference) moving averages are used instead of batch values. Disadvantage: dependence on batch size (small batches are noisy).

Layer normalization

Normalises within a single example across all features of a layer, independent of batch size. Therefore standard in transformers and language models (and in RNNs). Variants: RMSNorm (scaling only, without mean, faster and used in many modern language models), pre-norm versus post-norm (normalisation before or after the sublayer; pre-norm usually trains more stably).

Other methods

Group norm (groups of channels, good with small batches, for example in image generation), instance norm (style transfer), weight norm, spectral norm (GANs). Which normalisation fits depends on architecture and batch size.

Practice

In PyTorch: nn.BatchNorm2d(channels), nn.LayerNorm(dim). Put the model into evaluation mode (model.eval()) before evaluating, otherwise batch norm uses batch statistics and results fluctuate. Combine with sensible initialisation, a suitable learning rate and possibly gradient clipping.

See also: Initialisation, normalisation and dropout, The transformer block, Gradient descent, Convolutional networks (CNN)