Deep Learning · Batch normalisation

Batch normalisation (Ioffe and Szegedy, 2015) normalises each feature over the current mini-batch:

$\hat{x} = \frac{x-\mu_B}{\sqrt{\sigma_B^2+\epsilon}}, \qquad y = \gamma\hat{x} + \beta$

γ and β are learned, so the network can undo the normalisation if that helps.

• During training it uses the batch statistics and updates running averages.
• During evaluation it uses those running averages. Forgetting model.eval() changes the outputs.
• In practice it lets you train with higher learning rates and makes deep nets less sensitive to initialisation.

The comparison: a 10-layer MLP on a synthetic classification task, trained with and without BatchNorm, using the same seed and the same SGD settings.

import torch, torch.nn as nn
from sklearn.datasets import make_classification

X, y = make_classification(n_samples=4000, n_features=20, n_informative=10, random_state=0)
X, y = torch.tensor(X, dtype=torch.float32), torch.tensor(y)

def make(bn):
    layers, d = [], 20
    for _ in range(10):
        layers += [nn.Linear(d, 64)] + ([nn.BatchNorm1d(64)] if bn else []) + [nn.ReLU()]
        d = 64
    return nn.Sequential(*layers, nn.Linear(64, 2))

for bn in (False, True):
    torch.manual_seed(0)
    net = make(bn)
    opt = torch.optim.SGD(net.parameters(), lr=0.05)
    for epoch in range(5):
        for i in range(0, 4000, 64):
            opt.zero_grad()
            nn.functional.cross_entropy(net(X[i:i + 64]), y[i:i + 64]).backward()
            opt.step()
    net.eval()
    with torch.no_grad():
        acc = (net(X).argmax(1) == y).float().mean().item()
    print(f"BatchNorm={bn}: training accuracy after 5 epochs {acc:.3f}")

bnl = nn.BatchNorm1d(3); bnl.train(); _ = bnl(torch.randn(8, 3) * 5 + 10)
print("running mean after one batch:", [round(v, 2) for v in bnl.running_mean.tolist()])

Output:

BatchNorm=False: training accuracy after 5 epochs 0.500
BatchNorm=True: training accuracy after 5 epochs 0.972
running mean after one batch: [0.92, 1.17, 1.26]