Batch Norm and Layer Norm: What They Are Actually Doing to the Loss Landscape

Batch Normalisation is described as fixing "internal covariate shift." The explanation sounds plausible but is not the full story — and later research showed that the covariate shift explanation is mostly wrong. What Batch Norm actually does is smooth the loss landscape. Layer Norm does the same thing differently. Here is the geometry that both are exploiting.

Normalisation layers are universally used in deep networks, but the explanation for why they work has evolved significantly since Batch Norm was introduced. The covariate shift framing is familiar but empirically weak. The loss landscape explanation is more accurate and more useful for intuition.

The problem: deep networks are sensitive to scale

Without normalisation, a weight in an early layer that grows slightly too large causes the activations in subsequent layers to grow, which causes gradients to explode or saturate. The chain of transformations amplifies small perturbations. Training is brittle: it requires careful learning rate tuning, careful initialisation, and is prone to collapse for deep networks.

Batch Normalisation: normalise over the batch

For a layer producing activations x of shape (batch_size, features), Batch Norm computes: μ = mean over batch dimension, σ^2 = variance over batch dimension, x_norm = (x - μ) / sqrt(σ^2 + ε), output = γ * x_norm + β.

γ and β are learnable per-feature scale and shift parameters. They are necessary: without them, normalisation would constrain every layer to produce zero-mean unit-variance activations, removing the expressive capacity that the layer is supposed to have. γ and β allow the network to learn any mean and variance — they just make it explicit and trainable rather than implicit.

The effect on training: mean and variance are controlled at every layer boundary. Learning rate can be much higher (less risk of activation explosion). Gradients are better conditioned. Training is faster and more stable.

What Batch Norm actually does to the loss landscape

Santurkar et al. (2018) showed experimentally that Batch Norm does not primarily reduce internal covariate shift (the activations still shift; the paper showed that networks with Batch Norm and injected covariate shift train fine). What it does do is smooth the loss landscape — the loss function becomes more Lipschitz, meaning its gradient does not change rapidly from step to step. A smoother landscape allows larger learning rates and more predictable gradient updates.

The problems with Batch Norm

Batch statistics: the normalisation statistics (μ, σ^2) are computed over the batch. With large batches this is stable. With small batches (batch size 1 or 2) the statistics are noisy and training is unstable. At inference with batch size 1, the batch statistics are meaningless — Batch Norm maintains running estimates of μ and σ during training to use at inference, but these estimates drift and can cause train/inference discrepancies.

Recurrent networks: computing batch statistics across time steps is problematic because statistics vary with sequence position.

Layer Normalisation: normalise over the features

For the same activation tensor of shape (batch_size, features), Layer Norm computes: μ = mean over feature dimension (per example), σ^2 = variance over feature dimension (per example), x_norm = (x - μ) / sqrt(σ^2 + ε), output = γ * x_norm + β.

Each example is normalised independently of other examples in the batch. This means Layer Norm works for batch size 1, works across variable-length sequences, and produces the same result at training and inference (no running statistics needed). These properties make it ideal for Transformers.

Why Transformers use Layer Norm

Transformers process variable-length sequences with arbitrary batch sizes. Batch Norm would require tracking separate statistics for each sequence position, and inference behaviour would depend on batch composition. Layer Norm avoids both problems. Every token embedding is normalised over its feature dimension independently. The γ and β parameters are shared across sequence positions (same feature dimension), keeping the parameter count manageable.

The placement matters too: Pre-LN (layer norm applied before the sub-layer, as in the original Transformer) vs Post-LN (after the sub-layer, as in GPT-2). Pre-LN makes training more stable for very deep models because the residual pathway is unscaled — gradients flow back through the residual connection without passing through a layer norm.

Try on Colab: train a 10-layer MLP on MNIST without any normalisation, then with Batch Norm, then with Layer Norm. Plot the distribution of activations at each layer across training epochs. The activation explosion in the unnormalised case and the stability introduced by either normalisation variant will be visible in the histograms.