Normalization

Normalization layers standardize intermediate activations and then apply trainable scale and shift. They reduce sensitivity to initialization and learning rate, but the axis being normalized matters: batch norm couples examples in a minibatch, while layer norm normalizes features within each example and is therefore natural in transformers.

Standardize, then scale and shift

Normalization layers all share one shape: standardize a set of activations, then apply a learned scale and shift. They differ only in which axis supplies the statistics:

MethodNormalizes overNatural setting
Batch normthe minibatch (and spatial positions)convolutional networks
Layer normone example’s featurestransformers

For a set of activations , normalization computes

Here the values are the activations in whichever axis is being normalized: a batch-feature column for batch norm, or the feature coordinates of one example for layer norm. The mean recenters those values and measures their spread.

Normalization then returns

The normalized value is scaled by trainable and shifted by trainable , so the layer can recover a useful activation scale instead of forcing every downstream feature to stay standardized. The small prevents division by zero.

Batch normalization usually estimates across the minibatch and spatial positions, common in convolutional networks. Layer normalization estimates them across a single example’s feature dimension, so train and inference use the same statistics.

Batch normalization

Batch normalization standardizes each feature or channel using statistics from the current training minibatch. For a feature/channel , let be the set of values used to estimate its moments: in an MLP this is usually the minibatch examples for feature ; in a CNN this is often minibatch examples plus spatial positions for channel .

The normalized activation is

Here is one activation value for item or position and feature/channel , is a small numerical constant, and are learned scale and shift parameters. The learned affine parameters are important: they let the network choose the activation scale it needs after the standardization step.

During training, BatchNorm uses minibatch statistics and updates running estimates of the mean and variance. During inference, it uses those running estimates so predictions do not depend on which other examples happen to be in the same batch.

BatchNorm helps with vanishing and exploding gradients indirectly. By keeping intermediate activations in a controlled range, it reduces sensitivity to weight scale and learning rate, which usually makes the local backward Jacobians less erratic. It is not a proof that gradients cannot vanish or explode, but it was one of the key techniques that made deeper CNNs easier to optimize.

Layer normalization

Layer normalization uses the same standardize-then-affine pattern, but computes the moments across the feature coordinates of one example rather than across the minibatch. That makes train and inference behavior the same and avoids coupling different examples together. This is why LayerNorm is the default normalization style inside transformers, where sequence lengths, batch sizes, and autoregressive inference patterns often make BatchNorm inconvenient.

Worked example

This snippet applies batch normalization and layer normalization to the same tensor and prints the moments each normalization controls.

import torch
 
x = torch.tensor([[1., 2., 7.], [3., 4., 9.]])
batch = (x - x.mean(0)) / torch.sqrt(x.var(0, unbiased=False) + 1e-5)
layer = (x - x.mean(1, keepdim=True)) / torch.sqrt(x.var(1, unbiased=False, keepdim=True) + 1e-5)
print("batch_norm_means", torch.round(batch.mean(0), decimals=4).tolist())
print("batch_norm_vars", torch.round(batch.var(0, unbiased=False), decimals=4).tolist())
print("layer_norm_row0", torch.round(layer[0], decimals=4).tolist())

Observed output:

batch_norm_means [0.0, 0.0, 0.0]
batch_norm_vars [1.0, 1.0, 1.0]
layer_norm_row0 [-0.8889999985694885, -0.5080000162124634, 1.3969999551773071]

Batch normalization makes each feature column zero-mean and unit-variance across the two examples. Layer normalization instead standardizes the first row across its three features.

Caveats

BatchNorm’s train/eval split is a real failure mode: stale running statistics can break inference after distribution shift or very small batches. It is also awkward for online inference and variable batch composition. LayerNorm avoids batch coupling but does not preserve feature scale information unless the learned affine parameters recover it. Normalization also changes the effective optimization geometry, so it is not merely preprocessing.

References