Normalisation

In twenty fifteen, two researchers added one small layer to their networks. Training became many times faster, and networks that wouldn't train at all suddenly worked. It was called batch normalisation. Today, some form of it is found in almost every serious network.

• One small layer, 2015
• Training got much faster
• Now it's almost everywhere

Two kinds of normalisation

As data passes through a network, its numbers can drift, which makes later layers chase a moving target. Normalisation resets each layer's inputs, to a mean of about zero, and a spread of about one. Here we have numbers around two hundred. Batch norm normalises each feature across the batch. Layer norm does it across the features in each example. Run it. Both bring the numbers back into a healthy range. Batch norm is standard in image networks. Layer norm is what transformers use, because it doesn't depend on the batch.

import torch
from torch import nn

x = torch.rand(256, 10) * 50 + 200
print(f"Before: mean {x.mean():.1f}  spread {x.std():.1f}")

batch = nn.BatchNorm1d(10)
out = batch(x)
print(f"After BatchNorm: mean {out.mean():.3f}  spread {out.std():.3f}")

layer = nn.LayerNorm(10)
out = layer(x)
print(f"After LayerNorm: mean {out.mean():.3f}  spread {out.std():.3f}")

Before: mean 225.1  spread 14.4
After BatchNorm: mean 0.000  spread 1.000
After LayerNorm: mean -0.000  spread 1.000

Using them correctly

Batch norm has one famous trap. It behaves differently in training and in testing, because it uses running averages at test time. So if you forget model dot eval, your test results will be wrong. It's also weak with very small batches. Layer norm has neither problem, so it's usually the choice for language models. Group norm is a good middle ground for vision tasks with tiny batches.

Summary