Batch Normalization

What is Batch Normalization?

When training deep networks, weight updates in early layers constantly change the distribution of inputs feeding into later layers.

This moving input distribution forces deep layers to constantly adapt to shifting inputs, slowing down training.

Batch Normalization solves this by normalizing layer inputs across samples in each mini batch so features maintain a stable mean and variance.

  Mini Batch Activations ──► Compute Batch Mean & Variance ──► Standardize to (Mean 0, Var 1) ──► Apply Scale & Shift (gamma, beta)

Step by Step Equations

For a mini batch $B = {x_1, \dots, x_m}$ of size $m$:

1. Calculate Mini Batch Mean:

$$\mu_B = \frac{1}{m} \sum_{i=1}^m x_i$$

2. Calculate Mini Batch Variance:

$$\sigma_B^2 = \frac{1}{m} \sum_{i=1}^m (x_i - \mu_B)^2$$

3. Normalize Activations:

$$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$$

Where $\epsilon$ is a tiny number to prevent division by zero.

4. Scale and Shift:

$$y_i = \gamma \hat{x}_i + \beta$$

The parameters $\gamma$ (scale) and $\beta$ (shift) are learned automatically during backpropagation. This allows the network to undo normalization if doing so improves accuracy.

Key Benefits of Batch Normalization

1. Faster Training: Allows higher learning rates without risk of exploding or vanishing gradients.
2. Smooth Loss Landscape: Makes the loss surface smoother, helping gradient descent take direct paths to lower loss.
3. Mild Regularization: Because mini batch statistics vary slightly with each random batch, small noise is added to activations, reducing overfitting.

Training vs Inference Behavior

1. Training: Calculates mean and variance dynamically using the active mini batch.
2. Inference (Evaluation): Mini batch size might be 1. Instead of batch statistics, it uses running average mean and variance accumulated during training.

Say this out loud

Batch Normalization standardizes layer inputs across mini batch samples to zero mean and unit variance, followed by learned scale gamma and shift beta parameters. It speeds up training, smooths the loss landscape, and allows higher learning rates. During inference, mini batch statistics are replaced with running averages accumulated during training.

Followups to expect

1. Where should Batch Normalization be placed relative to activation functions? The original paper placed Batch Normalization right before non linear activations (Convolution -> BatchNorm -> ReLU). In modern practice, placing it either before or after ReLU works well.
2. Why does Batch Normalization struggle with small batch sizes? With very small batch sizes like 2 or 4, batch mean and variance estimates become extremely noisy, degrading model accuracy. Use Layer Normalization or Group Normalization instead.