Batch Normalisation

Normalise across the batch instead of across the features.

In training mode, for input x of shape (batch, d):

mean = mean(x, over the batch axis)      shape (d,)
var  = mean((x - mean)^2, over the batch axis)
out  = gamma * (x - mean) / sqrt(var + eps) + beta

Also maintain running statistics for use at inference time:

running_mean = momentum * running_mean + (1 - momentum) * mean
running_var  = momentum * running_var  + (1 - momentum) * var

Return (out, new_running_mean, new_running_var). Do not modify the arrays you were given.

The axis is the whole difference from layer norm. Batch norm reduces over axis 0, so each feature is normalised using statistics gathered from the other examples in the batch. That couples the examples together, which is why it behaves differently at training and inference time, degrades with small batches, and needs the running averages at all.

Use the biased variance, dividing by batch. The eps goes inside the square root.

Function signature

def batch_norm(x, gamma, beta, running_mean, running_var, momentum=0.9, eps=1e-5):

Example

Input

x = np.array([[1.0], [3.0]])
batch_norm(x, np.ones(1), np.zeros(1), np.zeros(1), np.ones(1))

Output

(array([[-0.999995], [ 0.999995]]), array([0.2]), array([1.]))

Reasoning

The batch mean is 2 and the biased variance is 1, so the rows normalise to almost exactly -1 and 1 (the small difference is eps). The running mean moves 10% of the way from 0 to 2, and the running variance stays at 1.

Constraints and notes

• Reduce over axis 0, the batch axis.
• Use the biased variance, which divides by batch.
• Do not change the arrays you are given.