Batch Normalisation

By the end of this lesson you will be able to explain what batch normalisation does, write it in NumPy with its separate training and evaluation behaviour, and use nn.BatchNorm1d, nn.BatchNorm2d and nn.LayerNorm in PyTorch. You will also see it rescue a deep network that would otherwise learn badly.

In Weight Initialisation you saw signals shrink or explode as they pass through many layers. Good initialisation fixes the start of training, but weights change as the network learns, and the scale of each layer's output drifts. Batch normalisation fixes the scale at every layer, on every step.

The idea: rescale each feature using the batch

Batch normalisation (often "BatchNorm" or "BN", Ioffe and Szegedy, 2015) is a layer placed inside the network. For each feature (each neuron's output), it looks at the current mini-batch and:

1. computes the mean and variance of that feature across the batch;
2. subtracts the mean and divides by the standard deviation, so the feature has mean 0 and variance 1;
3. multiplies by a learned scale gamma and adds a learned shift beta.

Think of college admissions in India. Students arrive with marks from different boards, and one board marks more generously than another. Before making one merit list, an admissions office might convert each board's marks to a common scale, such as how far each student is above that board's average. BatchNorm does this for each neuron, on every batch.

Step 3 might look odd: why normalise and then rescale? Because forcing every feature to mean 0 and variance 1 may be too strict. With gamma and beta the network can choose any scale it likes, and it can even undo the normalisation if that helps. What it gains is control: the scale is now set by two simple parameters, not by a long chain of weights.

Step 1: normalise a batch by hand.

import numpy as np

rng = np.random.default_rng(0)
# 6 examples, 3 features with very different scales
batch = np.column_stack([rng.normal(50, 10, 6), rng.normal(0.2, 0.05, 6), rng.normal(-3, 2, 6)])
print("before: mean", batch.mean(axis=0).round(2), " std", batch.std(axis=0).round(3))

eps = 1e-5                                   # avoids dividing by zero
mean = batch.mean(axis=0)
var = batch.var(axis=0)
normalised = (batch - mean) / np.sqrt(var + eps)
print("after:  mean", normalised.mean(axis=0).round(2) + 0.0, " std", normalised.std(axis=0).round(3))

gamma, beta = np.array([1.0, 2.0, 0.5]), np.array([0.0, 1.0, -1.0])   # learned in a real network
out = gamma * normalised + beta
print("scaled: mean", out.mean(axis=0).round(2) + 0.0, " std", out.std(axis=0).round(3))

before: mean [50.94  0.2  -4.79]  std [3.694 0.046 1.44 ]
after:  mean [0. 0. 0.]  std [1.    0.998 1.   ]
scaled: mean [ 0.  1. -1.]  std [1.    1.995 0.5  ]

Before normalising, the three features live on wildly different scales. After, each has mean 0 and standard deviation 1. After the learned gamma and beta, each has exactly the mean and spread the network chose.

Training mode and evaluation mode

There is a catch. At test time you may predict for a single example, so there is no batch to take a mean over. And a prediction should not depend on which other examples happen to share its batch.

The solution: during training, BatchNorm also keeps a running mean and running variance, slowly updated averages of the batch statistics. In evaluation mode it uses those fixed values instead of the batch's own. PyTorch updates them as running = (1 - momentum) * running + momentum * batch_value, with momentum = 0.1 by default.

Step 2: a BatchNorm class with both modes.

import numpy as np

class BatchNorm:
    def __init__(self, n_features, momentum=0.1, eps=1e-5):
        self.gamma, self.beta = np.ones(n_features), np.zeros(n_features)
        self.running_mean, self.running_var = np.zeros(n_features), np.ones(n_features)
        self.momentum, self.eps, self.training = momentum, eps, True

    def forward(self, x):
        if self.training:
            mean, var = x.mean(axis=0), x.var(axis=0)
            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
        else:
            mean, var = self.running_mean, self.running_var
        return self.gamma * (x - mean) / np.sqrt(var + self.eps) + self.beta

rng = np.random.default_rng(0)
bn = BatchNorm(2)
for step in range(200):                                  # 200 training batches
    bn.forward(rng.normal([10.0, -4.0], [3.0, 0.5], size=(32, 2)))
print("running mean:", bn.running_mean.round(2), " running var:", bn.running_var.round(2))

one_example = np.array([[13.0, -4.5]])
bn.training = False
print("evaluation mode output:", bn.forward(one_example).round(3))
bn.training = True
print("training mode output:  ", bn.forward(one_example).round(3))

running mean: [ 9.85 -4.01]  running var: [8.2  0.22]
evaluation mode output: [[ 1.099 -1.04 ]]
training mode output:   [[0. 0.]]

After 200 batches, the running mean and variance are close to the true values (mean 10 and -4, variance 9 and 0.25). In evaluation mode, the single example is normalised with these: 13 is one standard deviation above 10, so it becomes about 1. In training mode, the "batch" is just that one example, so it is compared with itself and becomes all zeros. That is why forgetting model.eval() gives nonsense predictions.

Why it helps

Batch normalisation usually lets you:

• train faster, often with a larger learning rate, because each layer's input keeps a stable scale;
• train deeper networks that would otherwise vanish or explode;
• depend less on careful initialisation;
• get a little regularisation for free, because each example's output depends slightly on the random batch it lands in.

Researchers still debate exactly why it works so well. The original paper's explanation, "reducing internal covariate shift", has been questioned; later work points to a smoother loss surface. The practical benefits are not in doubt.

Step 3: see it rescue a deep network

Here is the effect on a real network, trained in PyTorch. We build a 10-layer ReLU network for digits with a deliberately poor setup: every weight starts with a small fixed standard deviation of 0.05, and we use plain SGD. We train it once without and once with BatchNorm1d after every hidden linear layer. This is plain code with its real output.

import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
X_train, X_val, y_train, y_val = train_test_split(X / 16.0, y, test_size=0.25, random_state=42, stratify=y)
X_train, X_val = torch.tensor(X_train, dtype=torch.float32), torch.tensor(X_val, dtype=torch.float32)
y_train, y_val = torch.tensor(y_train), torch.tensor(y_val)

def make_model(use_bn):
    layers, width = [], 64
    for _ in range(9):
        layers.append(nn.Linear(width, 64))
        if use_bn:
            layers.append(nn.BatchNorm1d(64))
        layers.append(nn.ReLU())
    layers.append(nn.Linear(64, 10))
    model = nn.Sequential(*layers)
    for m in model:
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, 0, 0.05)    # poor, too-small start
    return model

for use_bn in (False, True):
    torch.manual_seed(0)
    model = make_model(use_bn)
    opt = torch.optim.SGD(model.parameters(), lr=0.05)
    for epoch in range(10):
        model.train()
        perm = torch.randperm(len(X_train))
        for i in range(0, len(X_train), 32):
            idx = perm[i:i + 32]
            loss = nn.functional.cross_entropy(model(X_train[idx]), y_train[idx])
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    with torch.no_grad():
        acc = (model(X_val).argmax(dim=1) == y_val).float().mean().item()
    print(f"BatchNorm {'on ' if use_bn else 'off'}: validation accuracy after 10 epochs {acc:.3f}")

BatchNorm off: validation accuracy after 10 epochs 0.102
BatchNorm on : validation accuracy after 10 epochs 0.902

Without BatchNorm, the signal shrinks through ten small-weight layers, the gradients vanish, and the network stays at chance level (10% for ten digits). With BatchNorm, the same network, same learning rate and same number of epochs trains properly.

BatchNorm, LayerNorm and where to put them

Layer normalisation (LayerNorm) normalises across the features of one example instead of across the batch. It does not depend on the batch at all, so it behaves the same in training and evaluation and works with a batch size of 1. That is why transformers use it, as you will see in Attention & Transformers.

Where to place BatchNorm: the usual pattern is Linear (or Conv2d) → BatchNorm → activation. Because BatchNorm subtracts the mean, the bias of the layer before it is wasted, so people often write nn.Linear(64, 64, bias=False) or nn.Conv2d(..., bias=False).

import torch
from torch import nn

torch.manual_seed(0)
x = torch.randn(4, 3) * 5 + 2
bn, ln = nn.BatchNorm1d(3), nn.LayerNorm(3)
print("BatchNorm column means:", bn(x).mean(dim=0).detach().round(decimals=3))
print("LayerNorm row means:   ", ln(x).mean(dim=1).detach().round(decimals=3))
print("BN learnable:", [name for name, _ in bn.named_parameters()],
      " BN buffers:", [name for name, _ in bn.named_buffers()])

BatchNorm column means: tensor([-0., -0., -0.])
LayerNorm row means:    tensor([0., 0., 0., 0.])
BN learnable: ['weight', 'bias']  BN buffers: ['running_mean', 'running_var', 'num_batches_tracked']

BatchNorm makes each column (feature) average zero across the batch; LayerNorm makes each row (example) average zero. The running statistics are stored as buffers: saved with the model, but not trained by the optimiser.

Common mistakes

• Forgetting model.eval() at test time. BatchNorm then uses the test batch's own statistics, and single-example predictions break.
• Very small batches. With a batch of 2 or 4, the batch mean and variance are too noisy. Use bigger batches, or LayerNorm or GroupNorm.
• Calling model.train() while evaluating "just to check". Even a forward pass in training mode updates the running statistics.
• Wrong BatchNorm type. Use BatchNorm1d for (batch, features) data and BatchNorm2d for (batch, channels, height, width) images.

Key points

• Batch normalisation rescales each feature to mean 0 and variance 1 using the batch, then applies learned gamma and beta.
• In training it uses batch statistics and updates running averages; in evaluation it uses the running averages.
• It allows faster, deeper and less fragile training, with a small regularising effect.
• LayerNorm normalises each example across its features and does not depend on the batch; transformers use it.
• Always switch to model.eval() for testing and prediction.

import numpy as np

class BatchNorm:
    def __init__(self, n_features, momentum=0.1, eps=1e-5):
        self.gamma, self.beta = np.ones(n_features), np.zeros(n_features)
        self.running_mean, self.running_var = np.zeros(n_features), np.ones(n_features)
        self.momentum, self.eps, self.training = momentum, eps, True

    def forward(self, x):
        if self.training:
            mean, var = x.mean(axis=0), x.var(axis=0)
            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
        else:
            mean, var = self.running_mean, self.running_var
        return self.gamma * (x - mean) / np.sqrt(var + self.eps) + self.beta

rng = np.random.default_rng(5)
bn = BatchNorm(1)
for _ in range(200):
    bn.forward(rng.normal(50, 10, size=(16, 1)))
bn.training = False
print(bn.forward(np.array([[30.0], [50.0], [70.0]])).ravel().round(2))

[-2.18 -0.03  2.13]