8. Normalization
Deep networks train faster and more reliably when the activations flowing between layers stay well-scaled. This lesson introduces normalization layers, which standardize a layer's inputs on the fly, then learn to rescale them. We cover batch normalization and layer normalization, where each computes its statistics, how they behave at inference, and where to place them.
Objectives
- Explain why normalizing activations inside the network stabilizes and accelerates training.
- Derive the batch normalization transform: normalize, then scale and shift with learned \(\gamma, \beta\).
- Understand why running (moving-average) statistics replace batch statistics at inference.
- Define layer normalization and see why it suits recurrent networks and Transformers.
- Decide where to place a normalization layer relative to the activation \(g^{[l]}\).
- Compare batch and layer normalization along their normalization axis and use cases.
8.1 Why normalize inside the network
Recall a layer computes \(z^{[l]} = W^{[l]} a^{[l-1]} + b^{[l]}\) and \(a^{[l]} = g^{[l]}(z^{[l]})\). As training updates every \(W^{[l]}\), the distribution of each layer's input \(a^{[l-1]}\) keeps shifting. This moving target, sometimes called internal covariate shift, forces later layers to constantly re-adapt and slows the whole network down.
Normalizing the activations at each layer keeps their mean and variance stable across updates. The immediate benefits:
- The loss surface becomes smoother, so we can use a higher learning rate without diverging.
- Training converges in fewer epochs and is less sensitive to the weight initialization.
- The learned scale and shift give the network back the freedom to undo the normalization if that helps.
Remark: normalization is applied to the pre-activation \(z^{[l]}\) or the activation \(a^{[l]}\), not to the parameters. It is a layer inserted into the forward pass, with its own learnable parameters.
8.2 Batch normalization
Batch normalization (BatchNorm) standardizes each feature across the examples of a mini-batch, then applies a learned affine transform. It operates per feature, so every feature keeps its own statistics.
Normalization recenters and rescales a layer input to zero mean and unit variance before the learned scale and shift.
8.2.1 Batch statistics
For a feature \(x\) over a mini-batch \(\mathcal{B} = \{x^{(1)}, \dots, x^{(m)}\}\) of size \(m\), compute the batch mean and variance:
\[\boxed{ \mu_\mathcal{B} = \frac{1}{m}\sum_{i=1}^{m} x^{(i)}, \qquad \sigma_\mathcal{B}^2 = \frac{1}{m}\sum_{i=1}^{m}\left(x^{(i)} - \mu_\mathcal{B}\right)^2 }\]8.2.2 Normalize, scale, and shift
Standardize each value to zero mean and unit variance, using a small constant \(\epsilon > 0\) for numerical stability:
\[\boxed{ \hat{x}^{(i)} = \frac{x^{(i)} - \mu_\mathcal{B}}{\sqrt{\sigma_\mathcal{B}^2 + \epsilon}} }\]Then rescale with two learned parameters per feature, a scale \(\gamma\) and a shift \(\beta\):
\[\boxed{ y^{(i)} = \gamma\, \hat{x}^{(i)} + \beta }\]Remark: \(\gamma\) and \(\beta\) are learned by gradient descent like any weight. If the optimal behaviour is the raw input, the network can recover it by learning \(\gamma = \sqrt{\sigma_\mathcal{B}^2 + \epsilon}\) and \(\beta = \mu_\mathcal{B}\). Normalization never removes capacity, it only reparameterizes it.
8.2.3 Inference with running statistics
At inference we often score a single example, so a batch mean and variance are undefined or meaningless. Instead BatchNorm uses population estimates accumulated during training as exponential moving averages, with momentum \(\alpha \in [0, 1)\):
\[\boxed{ \mu \leftarrow \alpha\, \mu + (1 - \alpha)\, \mu_\mathcal{B}, \qquad \sigma^2 \leftarrow \alpha\, \sigma^2 + (1 - \alpha)\, \sigma_\mathcal{B}^2 }\]At test time the transform is fixed and deterministic, using these running statistics in place of the batch ones:
\[\boxed{ y = \gamma\, \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta }\]Remark: this train/inference split is the source of most BatchNorm bugs. Forgetting to switch the layer to evaluation mode leaves it computing batch statistics at test time, which corrupts predictions.
8.3 Layer normalization
Layer normalization (LayerNorm) keeps the same normalize-scale-shift recipe but changes the axis it averages over. Rather than pooling across the batch, it computes the statistics over the features of a single example. Each example is therefore normalized on its own, independent of the others in the batch.
8.3.1 Per-example statistics
For one example with feature vector \(a \in \mathbb{R}^{H}\) (its \(H\) activations in a layer), average over the features:
\[\boxed{ \mu = \frac{1}{H}\sum_{k=1}^{H} a_k, \qquad \sigma^2 = \frac{1}{H}\sum_{k=1}^{H}\left(a_k - \mu\right)^2 }\]The normalization, scale, and shift are identical in form to BatchNorm, applied per example:
\[\boxed{ \hat{a}_k = \frac{a_k - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad y_k = \gamma_k\, \hat{a}_k + \beta_k }\]8.3.2 Why LayerNorm for sequences
Because the statistics come from a single example, LayerNorm behaves the same in training and inference, and it does not depend on the batch size. This matters when the batch is tiny or when examples have variable length, as in text. LayerNorm is the normalization of choice for recurrent networks and Transformers, where the sequence length varies and a per-timestep batch mean would be ill-defined.
Remark: LayerNorm needs no running statistics, so there is no train/inference discrepancy to manage. This alone makes it simpler to deploy than BatchNorm.
8.4 Placement and practical effects
A normalization layer sits between the linear step \(W^{[l]} a^{[l-1]} + b^{[l]}\) and the nonlinearity \(g^{[l]}\). Two orderings are common.
Normalization is inserted between the linear map and the activation inside each layer.
- Before activation (normalize \(z^{[l]}\), then apply \(g^{[l]}\)): the original and most common placement. It keeps the input to the nonlinearity centred, which is where saturation hurts most.
- After activation (normalize \(a^{[l]}\)): sometimes used and occasionally better in practice, though it is less standard.
Two more practical points:
- Bias becomes redundant. The shift \(\beta\) replaces the layer bias, since normalization subtracts the mean and would cancel \(b^{[l]}\) anyway. Layers followed by normalization are often written without their own bias.
- BatchNorm depends on batch size. Its statistics are noisier with small batches, which acts as a mild regularizer but degrades badly when the batch is very small. LayerNorm is immune to this, which is another reason sequence models prefer it.
Remark: the batch-dependent noise in BatchNorm can partly substitute for other regularizers, so networks using it sometimes need less dropout.
8.5 BatchNorm versus LayerNorm
The two layers share the normalize-scale-shift transform and differ only in the axis of the statistics and the consequences that follow.
| Aspect | Batch normalization | Layer normalization |
|---|---|---|
| Normalization axis | across the batch, per feature | across the features, per example |
| Depends on batch size | yes | no |
| Train vs inference | batch statistics vs running statistics | identical in both |
| Running statistics needed | yes | no |
| Typical use | CNNs and feedforward vision models | RNNs and Transformers |
Batch normalization computes statistics down a feature column across the batch, layer normalization across the features of a single example.
Remark: the Transformer block in lesson 16 places a LayerNorm before or after each sublayer, precisely because it removes the batch dependence that would otherwise couple examples of different lengths.
With activations kept well-scaled, the network trains stably at higher learning rates. The next lesson turns to controlling overfitting through regularization and dropout.