Batch norm in training mode
Medium14 min10 hidden tests
A defect detector for a factory camera uses batch normalization after a dense layer. In training mode, each feature is normalized with the mean and variance of the current batch, and the layer also updates running averages of those statistics for use at inference. Write batch_norm_train(x, gamma, beta, running_mean, running_var, momentum, eps) that returns the normalized batch and the two updated running statistics, following PyTorch's conventions. In the first example, the batch means [2.0, 15.0] move the running mean from [0.0, 0.0] to [0.2, 1.5].
Feed batches through
y = (x − mean) / √(var + eps); running = (1 − m) · running + m · batchThis batch of 2 has mean [2, 15] and variance ÷ n of [1, 25], used to normalize it. The running variance takes the sample variance ÷ (n − 1), [2, 50]: running = 0.9 × running + 0.1 × batch. Feed more batches and the running mean drifts toward [2, 15].
Running mean 0.2, 1.5.
Examples
Two examples: each column becomes about -1 and 1
batch_norm_train([[1.0, 10.0], [3.0, 20.0]], [1.0, 1.0], [0.0, 0.0], [0.0, 0.0], [1.0, 1.0])
returns ([[-0.999995, -1], [0.999995, 1]], [0.2, 1.5], [1.1, 5.9])
One feature, momentum 0.5
batch_norm_train([[0.0], [1.0], [2.0], [5.0]], [2.0], [1.0], [1.0], [4.0], 0.5)
returns ([[-1.13809], [-0.0690434], [1], [4.20713]], [1.5], [4.33333])
Requirements
xhas shape (n, d) with n >= 2.gamma,beta,running_meanandrunning_varhave shape (d,).- Normalize each column with the batch mean and the population variance (divide by n):
y = (x - mean) / sqrt(var + eps) * gamma + beta. - Update
running_mean = (1 - momentum) * running_mean + momentum * mean, andrunning_varthe same way using the sample variance of the batch (divide by n - 1). - Return a tuple
(y, running_mean, running_var): an (n, d) float array and two (d,) float arrays.
The idea behind it
Batch normalization standardizes each feature using the statistics of the current mini-batch, then applies a learned gain and bias. In training, y = (x - mean_batch) / sqrt(var_batch + eps) * gamma + beta, with statistics taken down each column. At inference there may be a single example, so the layer cannot rely on the batch. It keeps running averages instead, updated every training step: running = (1 - momentum) * running + momentum * batch_stat. In PyTorch, momentum 0.1 is the default and gives the new batch 10% weight. PyTorch normalizes with the population variance (divide by n), and it stores the unbiased sample variance (divide by n - 1) in the running estimate. Mixing these up, or forgetting model.eval(), is a classic source of a model that scores well in training and poorly in production.