How does Batch Normalization work mathematically and why does it stabilize training?

Batch Normalisation (BN) normalises the pre-activation values within a mini-batch to have zero mean and unit variance, then rescales them with learnable parameters γ (scale) and β (shift): BN(x) = γ · (x - μ_B) / √(σ²_B + ε) + β, where μ_B and σ²_B are the batch mean and variance, and ε is a small constant for numerical stability.

BN addresses internal covariate shift — the distribution of each layer's inputs changes during training as the preceding layers' weights update, forcing each layer to continuously adapt to a moving target. By renormalising inputs at each layer, BN stabilises this distribution. In practice, BN also provides a mild regularisation effect (similar to adding noise via the mini-batch statistics), reduces sensitivity to learning rate, and substantially reduces the need for dropout in many architectures.

import torch
import torch.nn as nn

# BatchNorm1d: for fully-connected layers (normalises over batch dim)
# BatchNorm2d: for conv layers (normalises per channel over batch+spatial)

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.BatchNorm1d(256),    # BN BEFORE or AFTER activation — varies by paper
    nn.ReLU(),
    nn.Linear(256, 128),
    nn.BatchNorm1d(128),
    nn.ReLU(),
    nn.Linear(128, 10),
)

# BatchNorm behaves DIFFERENTLY in train vs eval mode!
model.train()   # uses batch mean/var during forward pass
model.eval()    # uses running mean/var (exponential moving avg)

# Always call model.eval() at inference time:
with torch.no_grad():
    model.eval()
    preds = model(torch.randn(1, 784))  # inference — correct behavior

# Manual: BN keeps running stats during training
bn = nn.BatchNorm1d(256)
print(bn.running_mean.shape)  # torch.Size([256]) — updated each forward call