Batch Normalization

Batch normalization (BN, BatchNorm) is a layer that standardizes each channel of a network’s activations with the mean and variance computed over the current mini-batch, then applies a learned scale and shift. Introduced by Ioffe and Szegedy in 2015, it was soon taken up across convolutional networks, because a batch-normalized network can be trained with much higher learning rates and is less sensitive to its initialization. Its costs: an image’s output depends on the other images in its batch, training and inference compute differently, and it degrades with few images per batch, as in detection and segmentation.

Purpose

In a deep network, the inputs of each layer depend on all the weights below it, so their distribution changes as those weights are trained. Ioffe and Szegedy (2015) argued that this forces low learning rates and careful initialization, and that with saturating activation functions such as the sigmoid it pushes inputs into the flat regions where gradients vanish. They called this change internal covariate shift and proposed to reduce it by normalizing layer inputs inside the network, where the gradient accounts for the normalization. Whether that is why BN works is disputed (see Evidence and Impact); its effect on optimization is not.

When to Use It

BN suits convolutional networks trained with many images per device, as in image classification. It is a poor fit when batches are small or not independent, as for detectors and segmentation networks that fit only one or two high-resolution images per GPU; the variants below then replace it.

How It Works

The transform

Consider one channel and a mini-batch $\mathcal{B} = \{x_1, \dots, x_m\}$ of its values. The batch normalizing transform (Ioffe and Szegedy, Algorithm 1) computes

$$\mu_{\mathcal{B}} = \frac{1}{m}\sum_{i=1}^{m} x_i, \qquad \sigma_{\mathcal{B}}^2 = \frac{1}{m}\sum_{i=1}^{m} (x_i - \mu_{\mathcal{B}})^2, \qquad \hat{x}_i = \frac{x_i - \mu_{\mathcal{B}}}{\sqrt{\sigma_{\mathcal{B}}^2 + \epsilon}}, \qquad y_i = \gamma \hat{x}_i + \beta .$$

The constant $\epsilon$ (PyTorch’s default is $10^{-5}$) keeps the division stable. The learned scale $\gamma$ and shift $\beta$ let the layer represent the identity: with $\gamma = \sqrt{\sigma^2 + \epsilon}$ and $\beta = \mu$ the original activations come back, so normalization does not reduce what the network can express. Because the mean is subtracted, a bias in the preceding layer has no effect, and $\beta$ takes over its role; Ioffe and Szegedy drop the bias and place BN between the convolution or linear map and the nonlinearity.

The convolutional form

A convolution applies the same filter everywhere, so BN should treat every position of a feature map of shape $(N, C, H, W)$ alike. Ioffe and Szegedy (§3.2) therefore pool each channel’s statistics over the batch and over all spatial positions, an effective batch of $m' = N \cdot H \cdot W$ values, and learn one $\gamma$ and one $\beta$ per channel rather than per activation. PyTorch’s documentation calls this spatial batch normalization.

Training and inference

During training a prediction depends on which other images share the batch; at inference it should depend on the input alone. Ioffe and Szegedy therefore normalize at inference with population statistics: the average of the batch means and the unbiased estimate $\frac{m}{m-1}\,\mathbb{E}_{\mathcal{B}}[\sigma_{\mathcal{B}}^2]$ of the variance, and they note that moving averages can be used during training. With fixed statistics $\mu_c$ and $\sigma_c^2$, BN at inference is a per-channel affine map. If channel $c$ of the preceding convolution computes $W_c * u + b_c$, the two fold into one convolution (our algebra):

$$y_c = \frac{\gamma_c\,(W_c * u + b_c - \mu_c)}{\sqrt{\sigma_c^2 + \epsilon}} + \beta_c = W'_c * u + b'_c, \qquad W'_c = \frac{\gamma_c}{\sqrt{\sigma_c^2 + \epsilon}}\, W_c, \qquad b'_c = \beta_c + \frac{\gamma_c\,(b_c - \mu_c)}{\sqrt{\sigma_c^2 + \epsilon}} .$$

PyTorch’s BatchNorm2d updates running_mean and running_var after each training batch as $\hat{s} \leftarrow (1 - \text{momentum})\,\hat{s} + \text{momentum}\cdot s_t$ with momentum 0.1 by default; it normalizes with the biased batch variance but stores the unbiased one in running_var. model.train() selects batch statistics and updates the running averages; model.eval() uses the running averages.

The backward pass

Every $\hat{x}_i$ depends on every $x_j$ in the batch through $\mu_{\mathcal{B}}$ and $\sigma_{\mathcal{B}}^2$, so backpropagation must go through the statistics too; Ioffe and Szegedy point out that treating them as constants makes training blow up. Writing $g_i = \partial \ell / \partial y_i$, their chain-rule derivation gives

$$\frac{\partial \ell}{\partial \gamma} = \sum_{i} g_i \hat{x}_i, \qquad \frac{\partial \ell}{\partial \beta} = \sum_{i} g_i, \qquad \frac{\partial \ell}{\partial \hat{x}_i} = g_i\, \gamma,$$

plus terms through $\sigma_{\mathcal{B}}^2$ and $\mu_{\mathcal{B}}$. Collecting those terms (our algebra, checked numerically below) gives a compact form,

$$\frac{\partial \ell}{\partial x_i} = \frac{\gamma}{m\sqrt{\sigma_{\mathcal{B}}^2 + \epsilon}} \left( m\, g_i - \sum_{j} g_j - \hat{x}_i \sum_{j} g_j \hat{x}_j \right),$$

which removes from the incoming gradient its mean and its component along $\hat{x}$: shifting or (ignoring $\epsilon$) scaling a channel’s inputs leaves $\hat{x}$ unchanged, so those directions receive no gradient.

How to Apply It

The script implements BN in NumPy for a feature map of shape $(N, C, H, W) = (4, 3, 8, 8)$: four of scikit-learn’s 8×8 digit images filtered by three random 3×3 kernels. The gradients are checked against centered finite differences.

import numpy as np
from scipy.signal import correlate2d
from sklearn.datasets import load_digits

rng = np.random.default_rng(0)

# A conv feature map from real images: 4 digit images (8x8), 3 random 3x3 filters.
imgs = load_digits().images[:4] / 16.0                      # (N, H, W), pixels in [0, 1]
filters = rng.normal(size=(3, 3, 3))                         # (C, 3, 3)
x = np.stack([[correlate2d(im, f, mode="same") for f in filters] for im in imgs])
print("feature map", x.shape)                                # (N, C, H, W) = (4, 3, 8, 8)

def bn_forward(x, gamma, beta, eps=1e-5):
    """Training-mode BN: statistics per channel over (N, H, W)."""
    mu = x.mean(axis=(0, 2, 3), keepdims=True)
    var = x.var(axis=(0, 2, 3), keepdims=True)               # biased (divide by M)
    xhat = (x - mu) / np.sqrt(var + eps)
    y = gamma.reshape(1, -1, 1, 1) * xhat + beta.reshape(1, -1, 1, 1)
    return y, (xhat, var, gamma, eps)

def bn_backward(dy, cache):
    """Gradients of the loss w.r.t. x, gamma, beta, given dL/dy."""
    xhat, var, gamma, eps = cache
    M = dy.shape[0] * dy.shape[2] * dy.shape[3]              # values per channel
    dgamma = (dy * xhat).sum(axis=(0, 2, 3))
    dbeta = dy.sum(axis=(0, 2, 3))
    g = gamma.reshape(1, -1, 1, 1) / np.sqrt(var + eps)
    dx = g / M * (M * dy - dbeta.reshape(1, -1, 1, 1) - xhat * dgamma.reshape(1, -1, 1, 1))
    return dx, dgamma, dbeta

gamma, beta = rng.normal(size=3), rng.normal(size=3)
y, cache = bn_forward(x, gamma, beta)
xhat = cache[0]
print("max |per-channel mean| of xhat %.1e" % np.abs(xhat.mean(axis=(0, 2, 3))).max())
print("per-channel var  of xhat", xhat.var(axis=(0, 2, 3)).round(6))

# Gradient check with a random linear loss L = sum(w * y).
w = rng.normal(size=y.shape)
loss = lambda x, gamma, beta: (w * bn_forward(x, gamma, beta)[0]).sum()
dx, dgamma, dbeta = bn_backward(w, cache)

def numgrad(f, a, h=1e-6):
    g = np.zeros_like(a)
    for i in np.ndindex(a.shape):
        old = a[i]
        a[i] = old + h; fp = f()
        a[i] = old - h; fm = f()
        a[i] = old
        g[i] = (fp - fm) / (2 * h)
    return g

rel = lambda a, b: np.abs(a - b).max() / np.abs(b).max()
print("max relative error  dx %.1e  dgamma %.1e  dbeta %.1e" % (
    rel(dx, numgrad(lambda: loss(x, gamma, beta), x)),
    rel(dgamma, numgrad(lambda: loss(x, gamma, beta), gamma)),
    rel(dbeta, numgrad(lambda: loss(x, gamma, beta), beta))))

Output:

feature map (4, 3, 8, 8)
max |per-channel mean| of xhat 9.7e-17
per-channel var  of xhat [0.999976 0.999997 0.999989]
max relative error  dx 1.9e-09  dgamma 1.2e-10  dbeta 1.2e-10

Each channel of $\hat{x}$ has zero mean; its variance is $\sigma_{\mathcal{B}}^2 / (\sigma_{\mathcal{B}}^2 + \epsilon)$, slightly below 1 because of $\epsilon$. The hand-written gradients agree with finite differences to about $10^{-9}$.

Continuing the same script, the next lines compare it with torch.nn.BatchNorm2d (PyTorch 2.14.1) in training and evaluation mode.

# ---- compare with PyTorch (continues the script above) ----
import torch
print("torch", torch.__version__)
bn = torch.nn.BatchNorm2d(3).double()                    # eps=1e-5, momentum=0.1
with torch.no_grad():
    bn.weight.copy_(torch.from_numpy(gamma)); bn.bias.copy_(torch.from_numpy(beta))
xt = torch.from_numpy(x).requires_grad_()

bn.train()
yt = bn(xt)
(yt * torch.from_numpy(w)).sum().backward()
print("train mode: max |y - y_torch| %.1e, max |dx - dx_torch| %.1e, max |dgamma - dgamma_torch| %.1e" % (
    np.abs(y - yt.detach().numpy()).max(), np.abs(dx - xt.grad.numpy()).max(),
    np.abs(dgamma - bn.weight.grad.numpy()).max()))

# Running statistics after one training step: start at mean 0, var 1, momentum 0.1.
mu = x.mean(axis=(0, 2, 3)); M = x.size // 3
var_b = x.var(axis=(0, 2, 3)); var_u = var_b * M / (M - 1)
print("running_var           ", bn.running_var.numpy().round(5))
print("0.9 + 0.1 * unbiased  ", (0.9 + 0.1 * var_u).round(5))
print("0.9 + 0.1 * biased    ", (0.9 + 0.1 * var_b).round(5))

# Eval mode normalizes with the running statistics instead of the batch's own.
bn.eval()
with torch.no_grad():
    ye = bn(xt).numpy()
rm, rv = bn.running_mean.numpy(), bn.running_var.numpy()
ye_np = gamma[None, :, None, None] * (x - rm[None, :, None, None]) / np.sqrt(rv[None, :, None, None] + 1e-5) \
        + beta[None, :, None, None]
print("eval mode: max |y_eval - numpy| %.1e; max |y_eval - y_train| %.2f" % (
    np.abs(ye - ye_np).max(), np.abs(ye - y).max()))

Output:

torch 2.14.1
train mode: max |y - y_torch| 8.9e-16, max |dx - dx_torch| 8.9e-16, max |dgamma - dgamma_torch| 1.8e-14
running_var            [0.94237 1.1904  0.99175]
0.9 + 0.1 * unbiased   [0.94237 1.1904  0.99175]
0.9 + 0.1 * biased     [0.94221 1.18926 0.99139]
eval mode: max |y_eval - numpy| 8.9e-16; max |y_eval - y_train| 1.41

In training mode the outputs and gradients agree to rounding error. The running variance follows the unbiased update, as documented. In evaluation mode the same four images give outputs up to 1.41 away from the training-mode ones, because after one step the running statistics are still far from this batch’s.

Effect on the learning rate

The next script trains a small CNN on the 1,797 digit images (70/30 split) for five epochs with SGD, momentum 0.9 and batches of 32, with and without BatchNorm2d after each convolution, at two learning rates. It runs in about four seconds on a CPU.

import torch, torch.nn as nn, torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
torch.set_num_threads(1)

d = load_digits()                                         # 1797 grayscale 8x8 images
X = torch.tensor(d.images / 16.0, dtype=torch.float32).unsqueeze(1)
y = torch.tensor(d.target)
Xtr, Xva, ytr, yva = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y)

def make_cnn(bn):
    norm = (lambda c: nn.BatchNorm2d(c)) if bn else (lambda c: nn.Identity())
    return nn.Sequential(                                  # conv bias is redundant before BN
        nn.Conv2d(1, 16, 3, padding=1, bias=not bn), norm(16), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(16, 32, 3, padding=1, bias=not bn), norm(32), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(), nn.Linear(128, 10))

def train(bn, lr, seed, epochs=5, B=32):
    torch.manual_seed(seed)
    model = make_cnn(bn)
    opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    for _ in range(epochs):
        model.train()
        for idx in torch.randperm(len(ytr)).split(B):
            loss = F.cross_entropy(model(Xtr[idx]), ytr[idx])
            opt.zero_grad(); loss.backward(); opt.step()
    return model

@torch.no_grad()
def accuracy(model, B=None):                               # B=None: whole set in one batch
    out = model(Xva) if B is None else torch.cat([model(Xva[i:i + B]) for i in range(0, len(Xva), B)])
    return (out.argmax(1) == yva).float().mean().item()

def val_accuracy(model):
    model.eval()                                           # use the running statistics
    return accuracy(model)

print("torch", torch.__version__)
for lr in (0.03, 0.3):
    for bn in (False, True):
        accs = [val_accuracy(train(bn, lr, seed)) for seed in range(5)]
        print(f"lr {lr:<4}  {'with BN   ' if bn else 'without BN'}  val acc " + " ".join(f"{a:.3f}" for a in accs))

# The model.eval() pitfall: a BN model left in training mode at test time.
model = train(True, 0.03, seed=0)
model.eval()
print(f"\neval mode, running statistics: val acc {accuracy(model):.3f}")
model.train()
print(f"train mode, one image per batch: val acc {accuracy(model, B=1):.3f}")

Output:

torch 2.14.1
lr 0.03  without BN  val acc 0.952 0.965 0.952 0.961 0.948
lr 0.03  with BN     val acc 0.991 0.985 0.993 0.969 0.985
lr 0.3   without BN  val acc 0.100 0.394 0.324 0.100 0.102
lr 0.3   with BN     val acc 0.924 0.641 0.924 0.844 0.909

eval mode, running statistics: val acc 0.991
train mode, one image per batch: val acc 0.961

At learning rate 0.03 both networks train, and the batch-normalized one is ahead after five epochs (96.9–99.3% against 94.8–96.5%). At 0.3 the plain network fails: in three of the five seeds it collapses to predicting a single class, with training loss stuck near $\ln 10 \approx 2.30$, the cross-entropy of a uniform guess. The batch-normalized network still reaches 64–92%. The last two lines show the evaluation-mode pitfall: the same network scores 99.1% with its running statistics but 96.1% when left in training mode and fed one image at a time, each image then being normalized by its own statistics.

Small batches

After a fully connected layer, a batch of one image gives $\hat{x} = 0$ for every feature, and a batch of two gives $\hat{x} = \pm 1$ unless the two images nearly agree: each feature only records which image is larger. The script shows this on flattened digit images, then trains a small network with BN or layer normalization (LN, below) at batch sizes 32 and 2, scaling the learning rate with the batch size.

import torch, torch.nn as nn, torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
torch.set_num_threads(1)

d = load_digits()
X = torch.tensor(d.images / 16.0, dtype=torch.float32).reshape(-1, 64)   # flattened 8x8 images
y = torch.tensor(d.target)
Xtr, Xva, ytr, yva = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y)

# Batch statistics of two images: a normalized feature is +1 or -1 unless the images nearly agree.
torch.manual_seed(0)
bn = nn.BatchNorm1d(128)
with torch.no_grad():
    xhat = bn(nn.Linear(64, 128)(Xtr[:2]))                 # gamma = 1, beta = 0 at initialization
print("batch of 2: %.1f%% of normalized values have |xhat| > 0.99" % (100 * (xhat.abs() > 0.99).float().mean()))
print("           first image, first 6 features:", xhat[0, :6].numpy().round(3))
try:
    bn(Xtr[:1] @ torch.ones(64, 128))
except ValueError as err:
    print("batch of 1:", err)

def make_mlp(norm):
    N = nn.BatchNorm1d if norm == "BN" else nn.LayerNorm
    return nn.Sequential(nn.Linear(64, 128), N(128), nn.ReLU(),
                         nn.Linear(128, 128), N(128), nn.ReLU(), nn.Linear(128, 10))

def run(norm, B, seed, epochs=5):
    torch.manual_seed(seed)
    model = make_mlp(norm)
    opt = torch.optim.SGD(model.parameters(), lr=0.03 * B / 32, momentum=0.9)   # linear LR scaling
    for _ in range(epochs):
        model.train()
        for idx in torch.randperm(len(ytr)).split(B):
            if len(idx) < B: continue                       # drop the incomplete last batch
            loss = F.cross_entropy(model(Xtr[idx]), ytr[idx])
            opt.zero_grad(); loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        return (model(Xva).argmax(1) == yva).float().mean().item()

for norm in ("BN", "LN"):
    for B in (32, 2):
        accs = [run(norm, B, seed) for seed in range(3)]
        print(f"{norm}  batch {B:2d}  val acc " + " ".join(f"{a:.3f}" for a in accs))

Output:

batch of 2: 85.2% of normalized values have |xhat| > 0.99
           first image, first 6 features: [ 1.    -0.996 -0.941 -0.999 -0.977 -0.999]
batch of 1: Expected more than 1 value per channel when training, got input size torch.Size([1, 128])
BN  batch 32  val acc 0.980 0.987 0.980
BN  batch  2  val acc 0.719 0.787 0.746
LN  batch 32  val acc 0.965 0.972 0.972
LN  batch  2  val acc 0.965 0.965 0.976

With batches of two, the BN network drops from about 98% to 72–79% validation accuracy (71–81% over ten seeds): it trains on nearly binary features and is evaluated on continuous ones. LN, whose statistics come from each image alone, is unaffected. Convolutional layers suffered less here, presumably because their statistics also pool over spatial positions: the CNN of the previous example, trained with batches of two and the same learning-rate scaling, still reached 98.1–99.1% over five seeds.

Variants

• Synchronized BN (SyncBN) computes the statistics over the images on all GPUs rather than each GPU separately. Peng et al. (2018) used it, as Cross-GPU Batch Normalization, in the MegDet detector: with an all-reduce across the GPUs of one machine, two images per GPU gave statistics over 16 images, and together with learning-rate warmup it let them train with batches of up to 256 images on 128 GPUs. PyTorch provides torch.nn.SyncBatchNorm.
• Group normalization (Wu and He, 2018) divides the channels of each image into groups (32 by default) and normalizes over each group and all spatial positions, so it never uses the batch dimension and computes the same in training and inference.
• Batch renormalization (Ioffe, 2017) corrects the batch statistics toward the moving averages with an affine term that is treated as a constant in the backward pass, so that training and inference produce the same outputs; Ioffe reports that it trains substantially better than BN with small or non-i.i.d. mini-batches.
• Layer normalization (Ba, Kiros and Hinton, 2016) normalizes over all the units of a layer for each example, identically in training and testing, and was proposed with recurrent networks in mind. The transformer applies it after every residual sub-layer (Vaswani et al., 2017).
• Instance normalization (Ulyanov, Vedaldi and Lempitsky, 2016) normalizes each channel of each image over its spatial positions alone. Replacing BN with it in a feed-forward style-transfer network clearly improved the stylized images, the authors reported.

Best Practices

• Set the mode. Call model.eval() before validation or inference and model.train() before training; in PyTorch, torch.no_grad() does not change the mode, and a forward pass in training mode still updates the running statistics.
• Freeze BN when fine-tuning with small batches. Detectors such as Faster and Mask R-CNN train with one or two images per GPU and keep the pretrained BN layers fixed, turning them into a per-channel affine map; Wu and He describe this as BN being “frozen”. In PyTorch, keep those modules in evaluation mode.
• Reconsider dropout and weight decay. Ioffe and Szegedy removed dropout from their batch-normalized Inception network and reduced its L2 weight penalty fivefold, both of which improved validation accuracy.

Pitfalls

• Small or correlated batches give noisy statistics. Frames from one video, or crops from one image, are not independent samples, the non-i.i.d. case batch renormalization addresses.
• Mismatched statistics after a change of data: the running averages describe the training images, and a model applied to images with different statistics normalizes them with the wrong constants. Ioffe and Szegedy suggest, as future work, recomputing the population statistics for a new domain.

Evidence and Impact

On ImageNet, Ioffe and Szegedy (2015) reported that a batch-normalized Inception network matched the original’s 72.2% single-crop validation accuracy in 14 times fewer training steps, reached 74.8% with a 30 times higher learning rate, and, in an ensemble of six, 4.82% top-5 test error. They also argue that BN makes back-propagation through a layer independent of the scale of its weights, $\mathrm{BN}(W\mathbf{u}) = \mathrm{BN}(aW\mathbf{u})$ for a scalar $a$, with larger weights receiving smaller gradients, which stabilizes the growth of the weights at high learning rates.

The internal-covariate-shift explanation has been challenged. Santurkar et al. (2018) trained VGG networks on CIFAR-10 with noise of changing mean and variance injected after each BN layer; training was nearly as good as with plain BN and far better than without. By a gradient-based measure, networks with BN showed similar or even larger internal covariate shift. They argue instead that BN makes the loss landscape smoother: the loss and its gradient change less along the step direction, which makes gradients more predictive and allows larger learning rates. They prove that, under their assumptions, BN improves the Lipschitz constants of the loss and its gradient, and found that normalizing by $\ell_p$ norms, which does not guarantee control of the moments, performs comparably.

The batch dependence is BN’s main drawback for tasks that train with few images per GPU. In Wu and He’s ImageNet experiments with ResNet-50 on 8 GPUs, BN’s validation error rose from 23.6% at 32 images per GPU to 34.7% at 2 images per GPU, while group normalization stayed between 24.0% and 24.2%. At 32 images per GPU, GN was 0.5 points worse than BN, which the authors attribute to the regularizing noise of batch statistics that GN lacks. Ioffe and Szegedy gave the same reason for removing dropout: an image’s activations depend on which images share its batch.