Batch Normalization - Neural Networks

Batch Normalization (training-mode forward)

What you'll build

The forward pass of batch normalization for a 2-D batch of activations, in training mode: normalise every feature with the mean and variance of the current batch, then apply a learned scale and shift. Return the output and a cache for the backward pass.

Background

BatchNorm (Ioffe & Szegedy, 2015) keeps each layer's inputs at a stable scale, which allows higher learning rates and makes deep CNNs much easier to train. Because the statistics come from the batch, every example's output depends slightly on the others, which also acts as a mild regulariser.

Formula

For X of shape (batch, features), per feature j:

1. mu_j    = mean(X[:, j])                       # batch mean,            shape (features,)
2. var_j   = mean((X[:, j] - mu_j)^2)            # biased batch variance, shape (features,)
3. X_norm  = (X - mu) / sqrt(var + eps)          # shape (batch, features)
4. out     = gamma * X_norm + beta               # learned scale and shift, broadcast over rows

Inputs

• X: (batch, features) float array
• gamma, beta: (features,)
• eps: small constant, default 1e-5

Expected return format

Return the tuple (out, cache):

• out: (batch, features), the normalised, scaled and shifted output
• cache: the tuple (X, X_norm, mu, var, gamma, eps), in this order, for the backward pass

Edge cases and conventions

• Use the biased variance (divide by batch, NumPy's default np.var), as BatchNorm does in training.
• eps goes inside the square root.
• A constant feature has var = 0; eps keeps the division finite and X_norm becomes 0.

Constraints

NumPy only; vectorised over the batch (reduce with axis=0).

Follow-up

• At inference time BatchNorm uses running averages of mu and var collected during training. What goes wrong if you forget model.eval()?
• Derive the backward pass from the cache; the compact form is dX = (gamma / (batch * std)) * (batch * dout - dout.sum(0) - X_norm * (dout * X_norm).sum(0)).
• LayerNorm normalises each example across its features instead (Transformers section). Why does it suit variable-length sequences better?

Input: batch_norm_forward(np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), np.ones(2), np.zeros(2))[0]
Output: [[-1.2247, -1.2247], [0.0, 0.0], [1.2247, 1.2247]] (rounded)

Input: batch_norm_forward(np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), np.array([2.0, 2.0]), np.array([0.5, -1.0]))[0]
Output: [[-1.9495, -3.4495], [0.5, -1.0], [2.9495, 1.4495]] (rounded)

import numpy as np

def batch_norm_forward(X, gamma, beta, eps=1e-5):
    """
    Batch normalization forward pass (training mode, batch statistics).

    Args:
        X: Input (batch, features)
        gamma: Scale (features,)
        beta: Shift (features,)
        eps: Added to the variance inside the square root

    Returns:
        out: Normalized, scaled and shifted output (batch, features)
        cache: Tuple (X, X_norm, mu, var, gamma, eps) for the backward pass
    """
    # Your code here
    pass