Batch Norm Behaves Differently at Inference

The classic production incident: the model scores well in the notebook and badly behind an API. The notebook fed it batches of 256. The API sends one row at a time, and a batch of one has zero variance.

Write batch_norm(x, gamma, beta, running_mean, running_var, training, momentum=0.1, eps=1e-5) returning (out, new_running_mean, new_running_var). x is (N, D); normalisation is per feature, over the batch.

Training. Normalise with this batch's mean and variance, then update the running statistics:

running = (1 - momentum) * running + momentum * batch_stat

Inference. Normalise with the running statistics and leave them untouched.

One detail decides whether your numbers match PyTorch: the running variance is updated with the unbiased batch variance (dividing by N−1), while the normalisation itself uses the biased one (dividing by N).