Normalisation: BatchNorm, LayerNorm, RMSNorm

Normalisation layers rescale activations so that downstream layers receive inputs with stable statistics. The technique cuts training time by 5-10x on deep networks and is now standard in every architecture. The choice of which normaliser is not a detail - it determines whether your model can be trained with batch size 1, whether you can serve with batches of varying size, and how fast each forward pass is.

BatchNorm

Introduced by Ioffe and Szegedy (2015) for CNNs. For each feature channel, compute mean and variance across the batch dimension, then normalise:

mu_c = mean over batch of x[:, c]
var_c = var over batch of x[:, c]
x_hat = (x - mu_c) / sqrt(var_c + eps)
y = gamma_c * x_hat + beta_c

gamma and beta are learnable per-channel scale and shift. At inference, running averages of mean and variance are used instead of batch statistics.

The original justification was "reducing internal covariate shift" - keeping the distribution of layer inputs stable across training. Santurkar et al (2018) showed that this story is wrong: BatchNorm works because it smooths the loss landscape, making gradients more predictable and allowing larger learning rates. The mechanism was right, the explanation was not.

Why transformers do not use BatchNorm

Three problems make BatchNorm a poor fit for sequence models:

1. Variable batch shapes at inference. Generation often runs with batch size 1. BatchNorm's running statistics drift from the training distribution.
2. Padding tokens corrupt statistics. In a batch of variable-length sequences, padded positions skew the mean and variance.
3. Distributed training pain. Computing batch statistics across data-parallel workers requires an all-reduce per BN layer.

LayerNorm (Ba, Kiros, Hinton 2016) normalises across the feature dimension for each token independently:

mu = mean over features of x_t
var = var over features of x_t
x_hat = (x_t - mu) / sqrt(var + eps)
y = gamma * x_hat + beta

No batch dependency, no padding issue, no all-reduce. It just works for sequence models. Every transformer from BERT and GPT-2 through GPT-4 uses LayerNorm.

Pre-norm vs post-norm

The original Transformer placed LayerNorm after the residual addition:

y = LayerNorm(x + Sublayer(x))   # post-norm

Modern transformers put it before:

y = x + Sublayer(LayerNorm(x))   # pre-norm

Pre-norm is dramatically easier to train at depth - gradients flow cleanly through the residual highway without going through a normaliser. The trade-off is a small accuracy hit on the original 6-layer translation benchmarks, irrelevant once you scale past 24 layers.

RMSNorm

Zhang and Sennrich (2019) observed that LayerNorm's mean-centring barely affects anything; the rescaling does most of the work. RMSNorm drops the mean:

rms = sqrt(mean(x^2) + eps)
y = gamma * x / rms

No subtraction, no beta shift. The compute saving is about 7-64% depending on hardware and model size. With no measurable accuracy loss it became the default in Llama (1, 2, 3), Mistral, Gemma, and most post-2023 open-weight LLMs.

A quick comparison

Trade-offs

• Batch size sensitivity. BatchNorm with batch size 1 is broken. GroupNorm handles small batches; LayerNorm and RMSNorm do not depend on batch size at all.
• Quantisation. RMSNorm is friendlier to int8 inference because there is no mean to track across calibration.
• Speed at scale. At Llama 70B sizes, swapping LayerNorm for RMSNorm shaves a few percent off wall-clock training time. Small per-step, large in aggregate.

Further reading

• Batch Normalization: Accelerating Deep Network Training - Ioffe and Szegedy, the original BN paper.
• How Does Batch Normalization Help Optimization? - Santurkar et al, the loss-landscape explanation.
• Layer Normalization - Ba, Kiros, Hinton.
• Root Mean Square Layer Normalization - Zhang and Sennrich, the RMSNorm paper.