Skip to content

← Neural Network Layers step 2 of 9

Medium Primitives

Implement Batch Normalization

For a batch of feature vectors, batch normalization computes separate statistics for each feature across the examples. With examples in rows, this means normalizing each column. The result for one example therefore depends on the other examples in its batch.

For x = [[1, 2], [3, 4]], the column means are [2, 3] and the column variances are [1, 1]. With eps = 1e-5, subtracting the means and dividing by sqrt(variance + eps) gives approximately [[-0.999995, -0.999995], [0.999995, 0.999995]].

Each feature then gets its own scale and shift: output[:, j] = gamma[j] * normalized[:, j] + beta[j].

Your task

Implement batch_norm(x, gamma, beta, eps=1e-5) for the current batch:

  • x is a floating-point tensor of shape (N, D), with positive dimensions.
  • gamma and beta each have shape (D,).
  • eps is a positive float. Add it inside the square root.
  • Return a tensor of shape (N, D).

For each column, compute its mean and population variance (sum of squared deviations divided by N, not N-1). Normalize as (x - mean) / sqrt(variance + eps), then multiply by gamma and add beta. Use basic reductions and arithmetic rather than a built-in BatchNorm layer.

This exercise has no running averages, momentum, or evaluation mode. Always use the provided batch's statistics. It also defines N=1: every variance is zero, so every output feature equals its beta value. You are implementing this formula, not all the behavior of a framework's BatchNorm module.

The diagram compares normalization groups. Here columns are features with no spatial axes. With BatchNorm selected, entries in one column share their statistics.

Loading visualization…