ML coding · Deep learning layers
Batch normalisation in training mode
Given a batch X (N × D, N ≥ 2), parameters gamma, beta,
and running estimates running_mean, running_var (each length D), return
[out, new_mean, new_var] where
- out = γ ⊙ (X − μ) / √(σ² + eps) + β, with μ and σ² the per-feature mean and biased variance over the batch;
- new_mean = (1 − momentum)·running_mean + momentum·μ;
- new_var = (1 − momentum)·running_var + momentum·σ²unbiased, where σ²unbiased = σ²·N/(N − 1).
These are PyTorch's conventions for BatchNorm1d in training mode.
Examples
Input X = [[1, 10], [3, 30]], gamma = [1, 1], beta = [0, 0], running_mean = [0, 0], running_var = [1, 1], momentum = 0.1, eps = 0
Output [[[-1, -1], [1, 1]], [0.2, 2], [1.1, 20.9]]
Two examples: each feature becomes (−1, 1). The running variance moves towards 2 and 200, not 1 and 100.
Input X = [[1], [2], [3], [4]], gamma = [2], beta = [1], running_mean = [0], running_var = [1], momentum = 1, eps = 1.000e-5
Output [[[-1.6833], [0.1056], [1.8944], [3.6833]], [2.5], [1.6667]]
Momentum 1 replaces the running statistics outright.
Input X = [[0.5, -1], [0.5, 3]], gamma = [1, 1], beta = [0, 0], running_mean = [10, 10], running_var = [5, 5], momentum = 0, eps = 1.000e-5
Output [[[0, -1], [0, 1]], [10, 10], [5, 5]]
Momentum 0 leaves them untouched. The constant first feature normalises to 0.
Constraints
- 2 ≤ N ≤ 512, 1 ≤ D ≤ 512
- 0 ≤ momentum ≤ 1
- Checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
Statistics are over axis 0, the batch. Each feature has its own mean and variance.
Hint 2
Normalise with X.var(axis=0), which divides by N.
Hint 3
For the running variance, rescale the biased variance by N/(N − 1) before mixing it in.
Where candidates go wrong
- Normalising with the unbiased variance, or tracking the biased one. Both are plausible and both disagree with PyTorch.
- Reading momentum the Keras way: Keras keeps momentum·old + (1 − momentum)·new, the opposite convention.
- Mutating running_mean in place and returning the same list, so the caller's copy changes too.
The idea
Per-feature statistics come from axis 0. The batch is normalised with its own biased variance, which is the variance of the numbers actually in front of you. The running variance is an estimate of the population variance for use at inference, so it is fed the unbiased estimator. With N = 2 the two differ by a factor of 2, which is why the first example's running variance heads for 2 and 200 rather than 1 and 100.
At inference the layer uses the running estimates in place of μ and σ², so its output for an example no longer depends on what else is in the batch.
Time O(ND) · Space O(ND)
Reference solution
import numpy as np
def batch_norm_train(X, gamma, beta, running_mean, running_var, momentum, eps):
X = np.asarray(X, dtype=float)
n = X.shape[0]
mu = X.mean(axis=0) # per feature, over the batch
var = X.var(axis=0) # biased: normalises this batch
out = (X - mu) / np.sqrt(var + eps) * np.asarray(gamma, dtype=float) + np.asarray(beta, dtype=float)
unbiased = var * n / (n - 1) # unbiased: what the running estimate tracks
new_mean = (1 - momentum) * np.asarray(running_mean, dtype=float) + momentum * mu
new_var = (1 - momentum) * np.asarray(running_var, dtype=float) + momentum * unbiased
return [out, new_mean, new_var]Follow-ups an interviewer asks
- Write the inference-mode forward pass. Fold it into the preceding linear layer's weights and bias.
- Derive the backward pass. Why does the gradient with respect to X sum to zero over the batch?
- Why does batch norm behave badly with a batch size of 1 or 2, and what replaces it (group norm, layer norm)?
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
Batch norm has a training mode and an inference mode, and the bugs live in the difference. The question checks the axis, the momentum convention and the two different variances, which between them account for most real batch-norm bugs.
The theory behind it is in the lesson Attention, normalisation and convolution from scratch.
Keep going
Layer normalisation
Normalise each row over its features, then scale by gamma and shift by beta.
Open the problem → Medium15 minBatch normalisation in training mode
Normalise each feature over the batch and update the running mean and variance the way PyTorch does: biased variance to normalise, unbiased to track..
Open the problem → Medium20 minScaled dot-product attention with a causal mask
Implement softmax(QKᵀ/√d)V with an optional causal mask in NumPy.
Open the problem → Medium20 min2-D convolution with stride and padding
Implement a 2-D convolution (cross-correlation) with zero padding and stride, first with loops, then with sliding_window_view and einsum..
Open the problem →