Internal covariate shift, and what BatchNorm was actually designed to fix

The problem, from first principles

As a deep network trains, EVERY layer's weights are changing simultaneously, on every single step. This creates a subtle but real problem for any given layer, say, layer 5: the distribution of activations it RECEIVES from layer 4 keeps shifting around, not because layer 5's own input data changed, but because layers 1 through 4 upstream of it are all being updated too. Layer 5 essentially has to keep re-adapting to a constantly moving target, on top of the actual learning problem it's trying to solve, this shifting-distribution effect was named "internal covariate shift" in the original 2015 Batch Normalization paper (Ioffe & Szegedy), as the specific problem BatchNorm was designed to address (though later research has debated how much of BatchNorm's actual benefit really comes from directly fixing this effect, versus other side benefits, like smoothing the loss landscape, that come along with it).

[02-layers/04-weight-initialization] addresses a RELATED but distinct problem: getting activation variance right at the very START of training. BatchNorm addresses the SAME kind of variance-control problem, but continuously, DURING every single training step, by explicitly re-centering and re-scaling each layer's activations back to a controlled distribution (mean 0, variance 1, then rescaled by learnable gamma/beta) every time they pass through, regardless of how much the upstream layers have shifted since the last step.

From theory to code

Implement batchnorm_forward. In TRAINING mode, compute the CURRENT batch's own mean and (biased) variance, per feature, normalize x using them, then rescale by the learnable gamma (scale) and beta (shift) parameters. Simultaneously, update running_mean and running_var (exponential moving averages, using momentum) so they track training statistics over time. In EVAL mode, skip computing batch statistics entirely and normalize using the STORED running_mean/running_var instead, this is what makes a single test-time sample (where "this batch's mean" would be meaningless, being just that one sample) still normalize sensibly.

Constraints

• Training mode must use the CURRENT batch's mean and BIASED variance (ddof=0, NumPy's default for .var()) for the actual normalization of x.
• The running_var UPDATE specifically must use the batch's UNBIASED variance (ddof=1, i.e. batch_var * n / (n - 1)), matching torch.nn.BatchNorm1d's exact convention, even though the normalization itself uses the biased variance.
• Eval mode must use running_mean/running_var directly, computing NO batch statistics at all, and must leave running_mean/running_var unchanged (no update in eval mode).
• gamma and beta must be applied AFTER normalization, as gamma * x_norm + beta, per feature.

Hints