Implement BatchNorm Forward & Backward

Math Foundations

For mini-batch activations $X \in \mathbb{R}^{B \times D}$, Batch Normalization (Ioffe & Szegedy, 2015):

1. Batch Mean and Variance:

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

2. Standardization:

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

3. Scale and Shift (Learnable Parameters $\gamma, \beta \in \mathbb{R}^D$):

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

Forward and Backward Pass Computational Graph

Input X ──► [ Compute Mean & Var ] ──► [ Normalize x_hat ] ──► [ Scale gamma & Shift beta ] ──► Output Y
                                             │                              │
dInput dX ◄── [ Backprop Chain Rule ] ◄──────┴──────────────────────────────┴── dGamma, dBeta

During training, maintain Running Exponential Moving Averages of mean and variance for inference:

$$\text{running_mean} = \alpha \cdot \text{running_mean} + (1 - \alpha) \cdot \mu_B$$

$$\text{running_var} = \alpha \cdot \text{running_var} + (1 - \alpha) \cdot \sigma_B^2$$

NumPy Implementation from Scratch

import numpy as np


class BatchNormFromScratch:

  def __init__(self, dim, momentum=0.9, eps=1e-5):
    self.eps = eps
    self.momentum = momentum
    self.gamma = np.ones((1, dim))
    self.beta = np.zeros((1, dim))

    # Running statistics for inference
    self.running_mean = np.zeros((1, dim))
    self.running_var = np.ones((1, dim))

    # Saved variables for backprop
    self.cache = None

  def forward(self, x, is_training=True):
    if is_training:
      mu = np.mean(x, axis=0, keepdims=True)
      var = np.var(x, axis=0, keepdims=True)

      x_hat = (x - mu) / np.sqrt(var + self.eps)
      out = self.gamma * x_hat + self.beta

      # Update running statistics
      self.running_mean = (
          self.momentum * self.running_mean + (1 - self.momentum) * mu
      )
      self.running_var = (
          self.momentum * self.running_var + (1 - self.momentum) * var
      )

      # Save cache for backpropagation
      self.cache = (x, x_hat, mu, var)
    else:
      # Inference mode: use running statistics
      x_hat = (x - self.running_mean) / np.sqrt(self.running_var + self.eps)
      out = self.gamma * x_hat + self.beta

    return out

  def backward(self, dout):
    # Retrieve cache from forward pass
    x, x_hat, mu, var = self.cache
    B = x.shape[0]

    # Gradients for learnable parameters
    dgamma = np.sum(dout * x_hat, axis=0, keepdims=True)
    dbeta = np.sum(dout, axis=0, keepdims=True)

    # Gradient with respect to normalized x_hat
    dx_hat = dout * self.gamma

    # Analytical backpropagation for batch normalization input dx
    ivar = 1.0 / np.sqrt(var + self.eps)
    dvar = np.sum(dx_hat * (x - mu) * -0.5 * (ivar**3), axis=0, keepdims=True)
    dmu = np.sum(dx_hat * -ivar, axis=0, keepdims=True) + dvar * np.mean(
        -2.0 * (x - mu), axis=0, keepdims=True
    )

    dx = dx_hat * ivar + dvar * 2.0 * (x - mu) / B + dmu / B
    return dx, dgamma, dbeta

Say this out loud

Batch Normalization standardizes intermediate activations to zero mean and unit variance per mini-batch, applying learnable gamma scale and beta shift parameters. During training, running exponential moving averages are updated. During evaluation, fixed running statistics ensure deterministic predictions independent of mini-batch size.

Followups to expect

1. Why does BatchNorm fail on small batch sizes (e.g. B = 2)? Small batch sizes produce noisy, inaccurate estimates of mean and variance, degrading model training. Layer Normalization or Group Normalization is preferred for small batch sizes.
2. What is Layer Normalization vs Batch Normalization? Batch Normalization normalizes across samples within a mini-batch per channel. Layer Normalization normalizes across features within each individual sample instance, making it ideal for recurrent and Transformer models.