Batch and layer normalization: keeping activations stable
Batch normalization normalizes activations across a mini-batch, reducing internal covariate shift and enabling faster training. Layer normalization normalizes across features within each sample, making it suitable for recurrent networks and independent of batch size. Both methods apply learnable scale and shift parameters to restore representational power.
The problem of internal covariate shift
During deep neural network training, the distribution of each layer's inputs changes as the parameters of previous layers are updated. This phenomenon, called internal covariate shift, slows down training by requiring lower learning rates and careful parameter initialization, as described in the batch normalization paper by Ioffe and Szegedy.
Normalization techniques address this by stabilizing the distributions of activations throughout the network. Batch normalization, introduced by Ioffe and Szegedy, normalizes layer inputs using statistics computed over each mini-batch, making normalization part of the model architecture and allowing much higher learning rates.
Layer normalization was later proposed by Ba, Kiros, and Hinton as an alternative that computes normalization statistics from all summed inputs to neurons in a layer on a single training case, making it straightforward to apply to recurrent neural networks and independent of batch size, as explained in the layer normalization paper.
How batch normalization works
Batch normalization normalizes layer activations by subtracting the mini-batch mean and dividing by the mini-batch standard deviation for each feature, then applies two learnable parameters: a scale parameter gamma and a shift parameter beta, as documented in the PyTorch BatchNorm2d documentation.
During training, the layer maintains running estimates of the mean and variance, updated with a momentum parameter, which are used for normalization during evaluation. The Dive into Deep Learning chapter on batch normalization notes that batch normalization also acts as a regularizer and can eliminate the need for Dropout in some cases.
A limitation of batch normalization is its dependence on mini-batch size: with very small batches the statistics become noisy, and with a batch size of one the method fails because each hidden unit would take value zero after subtracting the mean. It is also not straightforward to apply to recurrent neural networks.
How layer normalization works
Layer normalization computes the mean and variance used for normalization over the feature dimensions of each individual training sample, rather than across the batch. The PyTorch LayerNorm documentation specifies that the mean and standard deviation are calculated over the last D dimensions of the input, where D is the dimension of the normalized shape.
Unlike batch normalization, layer normalization performs the same computation at training and test times, and it does not require running statistics. This makes it independent of batch size and suitable for recurrent neural networks, where it stabilizes hidden state dynamics.
Layer normalization also uses learnable per-element affine parameters gamma and beta, which are applied after normalization. It is widely used in transformer architectures and other models where batch size may vary or be small.
Comparing batch and layer normalization
The fundamental difference between batch and layer normalization is the axis along which statistics are computed. Batch normalization normalizes across the batch dimension for each feature, while layer normalization normalizes across the feature dimension for each sample.
Batch normalization requires a sufficiently large batch size and behaves differently during training and evaluation due to the use of running statistics. Layer normalization uses the same computation in both phases and is independent of batch size.
The choice between them depends on the architecture. Batch normalization is common in convolutional neural networks, while layer normalization is preferred for recurrent networks and transformers. Both methods provide learnable scale and shift parameters to restore the network's representational capacity.
| Feature | Batch Normalization | Layer Normalization |
|---|---|---|
| Normalization axis | Across batch (per feature) | Across features (per sample) |
| Dependency on batch size | Yes | No |
| Train/test behavior | Uses running statistics at test | Same computation |
| Suitable for RNNs | Difficult | Well-suited |
| Learnable parameters | Per-channel scale and shift | Per-element scale and shift |
| Typical use | CNNs, MLPs | Transformers, RNNs |
Frequently asked questions
What is internal covariate shift?
Internal covariate shift is the change in the distribution of layer inputs during training as previous layer parameters are updated. It slows convergence and requires careful initialization. Batch normalization reduces this effect.
Does batch normalization work with small batch sizes?
Batch normalization can be less effective with very small batches because the computed statistics become noisy. With a batch size of one, the method fails entirely. Layer normalization is an alternative that does not depend on batch size.
Can layer normalization be used in convolutional networks?
Yes, layer normalization can be applied to convolutional networks by normalizing over the channel and spatial dimensions, as shown in the PyTorch documentation.
Why does batch normalization behave differently at training and test time?
During training, batch normalization uses mini-batch statistics for normalization. During evaluation, it uses running averages of mean and variance that were accumulated during training, ensuring stable inference.
Which normalization method is better for recurrent neural networks?
Layer normalization is preferred for recurrent neural networks because it normalizes per time step and is independent of batch size, making it straightforward to apply.