Batch Normalization

Batch normalization re-centres and re-scales a layer's outputs using the statistics of the current mini-batch, which makes deep networks train faster and tolerate larger learning rates. After this lesson your program normalizes a batch, keeps running statistics, and switches to those statistics for evaluation.

The problem it fixes

While a network trains, the weights of early layers change, so the distribution of numbers reaching a later layer keeps drifting. Each layer is chasing a moving target, which forces small learning rates. Normalizing every layer's inputs to a steady distribution removes most of that drift.

The training-time formula

For one feature, take its values across the B samples of the batch:

mean   = (x1 + ... + xB) / B
var    = ((x1 - mean)^2 + ... + (xB - mean)^2) / B       divide by B, not B - 1
x_hat  = (x - mean) / sqrt(var + eps)
y      = gamma * x_hat + beta

eps (typically 1e-5) sits inside the square root so a batch whose values are all equal, with variance 0, does not divide by zero. gamma and beta are learned: if plain zero-mean, unit-variance were the best input for the next layer the network would leave them at 1 and 0, and if not, it can scale and shift (even undo the normalization) as the loss demands.

A worked trace

Batch [2, 4, 6, 10] with gamma = 1, beta = 0:

mean = 5.5
var  = (12.25 + 2.25 + 0.25 + 20.25) / 4 = 8.75
sqrt(8.75 + 0.00001) = 2.9580
y = [-1.1832, -0.5071, 0.1690, 1.5213]

Evaluation needs different statistics

At serving time there may be no batch at all, or a batch of one, and an answer must not depend on which other requests arrived together. So training keeps running estimates of the mean and variance, updated after every training batch with a momentum m (the PyTorch default is 0.1):

running_mean <- (1 - m) * running_mean + m * batch_mean
running_var  <- (1 - m) * running_var  + m * batch_var * B / (B - 1)

They start at mean 0 and variance 1. Note the B / (B - 1) factor: the running variance stores the unbiased estimate even though normalization during training used the biased one. (For a batch of one, use the batch variance as it is.) Evaluation then normalizes with running_mean and running_var instead of the batch's own, and updates nothing.

After the trace batch: running_mean = 0.1 * 5.5 = 0.55 and running_var = 0.9 * 1 + 0.1 * (8.75 * 4/3) = 2.0667. Evaluating the value 5.5 now gives (5.5 - 0.55) / sqrt(2.0667 + 0.00001) = 3.4433, far from the 0 you might expect. One batch is too little for the running averages to catch up, which is why evaluating a freshly started model with batch norm gives odd answers.

Common mistakes

• Dividing the variance by B - 1 during training normalization.
• Normalizing across features instead of across the batch. That is layer normalization, the next lesson.
• Updating the running statistics during evaluation, or normalizing training batches with running statistics.
• Using the unbiased variance in the training formula, or the biased one in the running update.
• Trusting batch norm with tiny batches (1 or 2): the batch statistics are too noisy to be meaningful.

Your exercise

Commands set a BATCH and the parameters, then run TRAIN (normalize with the batch statistics and update the running ones), EVAL (normalize with the running ones) or STATS. Implement normalize, batch_stats and update_running.