Batch Normalization and the Train-Eval Gap

A layer with batch normalization computes two different functions depending on a flag. In training mode it normalises each channel using the mean and variance of the current mini-batch. In evaluation mode it normalises using two buffers that were accumulated during training and then leaves them fixed. Both produce a number of the same shape from the same weights, which makes the difference easy to overlook until a model that converged smoothly reports worse accuracy the moment it is evaluated.

The discrepancy is a real change of function, not a numerical artefact. Understanding it requires knowing exactly what is stored, when it is updated, and how the batch size at training time constrains what the stored values can be.

What the layer stores

For an input with $C$ channels and a batch of $N$ samples, training mode computes per-channel statistics over the $N$ samples and their spatial extent (for a 2D feature map, $N \cdot H \cdot W$ values per channel):

$$\mu_c = \frac{1}{N H W}\sum x_{c},\qquad \sigma^2_c = \frac{1}{N H W}\sum (x_c - \mu_c)^2 + \epsilon$$

and normalises to $\hat{x}_c = (x_c - \mu_c)/\sqrt{\sigma^2_c + \epsilon}$, then applies the learnable affine transform $\gamma_c \hat{x}_c + \beta_c$. Alongside these parameters the layer owns two non-trainable buffers initialized to zero mean and unit variance, updated after each forward pass by an exponential moving average:

$$\hat{m}_c \leftarrow (1-p)\hat{m}_c + p,\mu_c,\qquad \hat{v}_c \leftarrow (1-p)\hat{v}_c + p,\sigma^2_c$$

The momentum parameter $p$ is usually around $0.1$ (frameworks may define it as the weight on the new value; check which convention applies, because the same name is used for both). At evaluation time the layer performs $(\hat{m}_c, \hat{v}_c)$ normalisation with no dependence on the batch.

Three consequences follow directly from this byte-level description. First, the buffers are part of the checkpoint: a model saved without them cannot be evaluated correctly, and loading weights without buffers silently gives the moving averages their initialization values. Second, an evaluation run performed before any training steps uses exactly $\mu=0, \sigma^2=1$, so the layer is the identity up to $\gamma, \beta$ until the buffers have seen data. Third, a buffer update with $p$ small means the stored values reflect a long window of history, weighted toward recent batches but never completely forgetting early ones.

Why the two paths diverge

With running averages, the layer assumes the distribution of activations at evaluation is the distribution it saw during training. When that holds and the averages have converged, the two paths agree in expectation. They diverge in three ordinary situations.

The first is a batch size difference. The batch statistics are estimates with error of order $1/\sqrt{N_{\text{eff}}}$ where $N_{\text{eff}}$ is the number of values per channel. With a batch of $8$ images at $7\times 7$ resolution the per-channel estimate uses $392$ values, which is noisy; the running average over thousands of such noisy estimates is a much better estimate of the true mean. The training-time function is therefore noisier than the evaluation-time function, and a small batch makes training-time normalisation closer to sampling random scales. The affine parameters $\gamma, \beta$ are learned under that noise.

The second is a distribution shift between training and evaluation data. If the evaluation set’s activations have a different mean, the fixed buffers apply a normalisation that no longer centres them and the following layers see pre-activations offset by an amount that is constant rather than random.

The third is sequence length and padding in recurrent and attention models, where the layer normalises over a different set of positions in each mode depending on masking. Any implementation that computes masked statistics differently between the two modes produces a mismatch that is invisible in training.

import torch
from torch import nn

bn = nn.BatchNorm1d(4)
bn.train()
x = torch.randn(16, 4) * 3.0 + 1.0
bn(x)                      # buffers update, output uses batch stats
before = bn.running_mean.clone()
y_train = bn(x)            # second call: output uses this batch's stats
bn.eval()
y_eval = bn(x)             # output uses running_mean, running_var
print(float((y_train - y_eval).abs().max()) > 0.0,
      bn.running_mean.shape, before.shape)

The snippet’s point is not the numbers but the shapes: two buffers of size $C$, updated in training, and a different function in evaluation with identical weights.

Train-mode evaluation as a diagnostic

Because the two paths use different statistics, comparing them on held-out data separates two properties that are otherwise entangled. Accuracy with training-mode statistics measures whether the network’s representation is linearly separable given the batch’s own centring. Accuracy with the running statistics measures whether the network is usable as a fixed function on single examples.

A model whose train-mode accuracy is much higher than its eval-mode accuracy has weights that encode useful features but buffers that do not describe the data. The usual causes are a momentum too large relative to the number of steps (the average never converged), a batch size so small that the statistics are dominated by noise, or a training distribution that does not match evaluation. The diagnostic is cheap: it requires no retraining, only two forward passes.

The reverse case, evaluation better than training, is also informative. It means the batch statistics are adding noise that hurts the classifier on average — the network is being regularised harder at training time than it needs to be, and the affine parameters have adapted to a noisier input than the one it will see.

Ordering of updates and the first batches

The moving average starts at $\hat{v}=1$, which is rarely the true variance of any real activation. With momentum around $0.1$, the weight on the initial value decays as $(1-p)^{t}$, so after $t$ updates the residual influence of the initialization is about $e^{-pt}$. Reaching a few percent of the initial value takes on the order of $30$ updates. A training run of a few dozen steps leaves the buffers dominated by a value that was never measured, and evaluation then normalises with an almost-arbitrary scale. This is not a subtle effect for short fine-tuning runs, which is why frameworks offer a way to freeze the buffers and let the affine parameters absorb the statistics instead.

Freezing is a legitimate alternative: with $\mu=0, \sigma^2=1$ fixed, the layer becomes a per-channel affine transform, and $\beta$ learns the offset that centring would have supplied. The cost is that the layer no longer adapts to shifts in the input distribution, so any layer before it that produces a drifting scale must be corrected elsewhere. This is the same trade that separates batch normalisation from mechanisms that normalise per example, and it reappears in the comparison with layer normalisation and residual scaling where the statistics are computed without a batch dimension.

The gradient path through batch statistics

A detail that matters for behaviour under frozen layers and for weight decay: the normalisation is differentiable with respect to the input, and the gradient is not the identity. For a normalised output $\hat{x} = (x - \mu)/\sigma$ with batch-computed $\mu$ and $\sigma$, the Jacobian of the vector-valued map has the form

$$J = \frac{1}{\sigma}\left(I - \frac{1}{N}\mathbf{1}\mathbf{1}^\top - \hat{x}\hat{x}^\top \frac{1}{N}\right)$$

up to the exact coefficients. The two subtracted projectors remove the components of the incoming gradient that would change the batch mean and the batch norm, so a gradient that would shift all samples equally is projected out. That projection is why a batch-normalised layer has a scale-invariant gradient with respect to the preceding weights: multiplying those weights by a constant leaves the normalised output unchanged and therefore the loss unchanged, and gradient descent on the preceding layer’s weights has no component along that radial direction. Weight decay, which does not see the normalisation, does have such a component, and the interaction between the two is why the effective update to those weights differs from the plain case.

In evaluation mode the buffers are constants, so $J$ is diagonal with entries $1/\sqrt{\hat{v}_c + \epsilon}$ per channel. The function is then a fixed affine map and its sensitivity to its input is determined entirely by the stored variance. A channel whose stored variance is small amplifies its input, and a channel whose stored variance is large attenuates it. If a channel’s activations collapse to a constant during training, the stored variance approaches the $\epsilon$ floor, and evaluation multiplies the (near-zero) signal by $1/\sqrt{\epsilon}$, a large constant. The classifier can then be extremely sensitive to a channel that carried no information, which is one mechanism behind evaluation-time instability on inputs that differ slightly from the training distribution.

The same reasoning explains the well-known sensitivity to batch composition at training time: the loss for one sample depends on the other samples in the batch, through the statistics. That coupling is the mechanism and the cost simultaneously, and it is what makes batch size a hyperparameter of the function being trained rather than only of the optimiser.