Batch Normalization Forward Pass

Implement the training-time forward pass of batch normalization, the layer that stabilizes deep-network training by keeping each feature’s activations centered and scaled. Given a batch X of shape (n, d), you normalize every feature using statistics computed across the n samples in the batch, then apply a learnable scale gamma and shift beta. Interviewers use this to check that you know which axis the statistics come from, why the eps term exists, and how the affine gamma/beta step restores representational power after normalization.

Definition

For each feature column j, compute the mean and (biased) variance across the batch, standardize, then scale and shift:

mu[j]    = (1/n) * sum_i X[i, j]
var[j]   = (1/n) * sum_i (X[i, j] - mu[j])^2
xhat[i, j] = (X[i, j] - mu[j]) / sqrt(var[j] + eps)
out[i, j]  = gamma[j] * xhat[i, j] + beta[j]

The statistics mu and var are per-feature, so both have shape (d,). At training time these come from the current batch (this is what you implement here). At inference time a real layer instead uses a running mean and variance accumulated during training, so predictions do not depend on how the eval batch happens to be grouped.

Task

Complete batchnorm_forward(X, gamma, beta, eps=1e-5) so it returns the (n, d) output. Take the mean and variance over axis=0 (across samples) to get per-feature statistics of shape (d,), normalize with (X - mu) / sqrt(var + eps), then apply gamma * xhat + beta, letting NumPy broadcast the (d,) parameters across all n rows. Use the biased variance (divide by n, i.e. np.var’s default ddof=0). Do not call any deep-learning framework; use plain NumPy only.

Example

X = np.array([[1.0, 2.0],
              [3.0, 6.0]])
gamma = np.array([1.0, 2.0])
beta = np.array([0.0, 1.0])
batchnorm_forward(X, gamma, beta, eps=0.0)
# -> array([[-1.,  -1.],
#           [ 1.,   3.]])

Feature 0 has mean 2 and std 1, so its column normalizes to [-1, 1]; feature 1 has mean 4 and std 2, so it normalizes to [-1, 1], then gamma=2, beta=1 maps it to [-1, 3].

Constraints

• 1 <= n <= 10^4, 1 <= d <= 10^3; use vectorized NumPy with reductions over axis=0, no Python loop over samples or features.
• gamma.shape == beta.shape == (d,) and X.shape[1] == d.
• eps is a small positive float (default 1e-5); it must be added inside the square root so a constant feature (var == 0) does not divide by zero.
• Use the biased variance (ddof=0), matching training-time batch norm; values fit in float64.

Approach

Batch norm standardizes each feature using statistics taken across the batch, then reapplies a learnable affine map. Reduce X over axis=0 to get a per-feature mean and variance of shape (d,), subtract the mean and divide by sqrt(var + eps) to standardize, then scale by gamma and shift by beta. Every step is a broadcast of a (d,) vector against the (n, d) batch, so no loops are needed. The eps lives inside the square root purely for numerical safety when a feature is constant.

Solution

import numpy as np

def batchnorm_forward(
    X: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    eps: float = 1e-5,
) -> np.ndarray:
    X = np.asarray(X, dtype=np.float64)
    gamma = np.asarray(gamma, dtype=np.float64)
    beta = np.asarray(beta, dtype=np.float64)

    mu = X.mean(axis=0)                    # (d,) per-feature mean
    var = X.var(axis=0)                    # (d,) biased variance, ddof=0
    xhat = (X - mu) / np.sqrt(var + eps)   # (n, d) standardized
    return gamma * xhat + beta             # (n, d) scale and shift

Walkthrough

On the example with X = [[1, 2], [3, 6]], gamma = [1, 2], beta = [0, 1], eps = 0:

• mu = X.mean(axis=0) reduces down each column: feature 0 mean (1 + 3)/2 = 2, feature 1 mean (2 + 6)/2 = 4, so mu = [2, 4].
• var = X.var(axis=0) uses the biased formula: feature 0 ((1-2)^2 + (3-2)^2)/2 = 1, feature 1 ((2-4)^2 + (6-4)^2)/2 = 4, so var = [1, 4] and sqrt(var) = [1, 2].
• xhat = (X - mu) / sqrt(var) broadcasts mu and the std across both rows. Column 0: [(1-2)/1, (3-2)/1] = [-1, 1]. Column 1: [(2-4)/2, (6-4)/2] = [-1, 1]. So xhat = [[-1, -1], [1, 1]].
• gamma * xhat + beta broadcasts (d,) params down every row. Column 0: 1*[-1, 1] + 0 = [-1, 1]. Column 1: 2*[-1, 1] + 1 = [-1, 3]. Result [[-1, -1], [1, 3]], matching the expected output.

Complexity & notes

• Time O(n * d): each reduction, the subtract/divide, and the affine step touch every element a constant number of times. Space O(n * d) for xhat and the output; the statistics mu, var add only O(d).
• Why eps is inside the sqrt: if a feature is constant across the batch, var is exactly 0 and 1 / sqrt(0) is infinity (or NaN once it hits the scale step). Adding a small eps before the root bounds the divisor away from zero. It goes inside the root, not outside, because it is regularizing the variance, not the standardized value; the default 1e-5 is small enough to leave normal features essentially unchanged.
• Biased vs sample variance: batch norm uses the biased estimate (divide by n, ddof=0), which is np.var’s default. Reaching for np.var(X, axis=0, ddof=1) would silently rescale the outputs and mismatch every framework’s implementation.
• Train time vs inference time: this forward pass computes statistics from the current batch, which is correct for training but makes each sample’s output depend on its batch-mates. At inference a real layer freezes a running (exponentially averaged) mean and variance collected during training and uses those instead, so a single example gets a deterministic, batch-independent output. Using batch statistics at eval time is a common bug that makes predictions jitter with batch composition.
• Axis discipline: the statistics must reduce over axis=0 (samples) to be per-feature, shape (d,). Reducing over axis=1, or forgetting the axis so NumPy reduces the whole array to a scalar, is the classic mistake and silently produces wrong shapes or wrong normalization.
• gamma and beta are learnable, letting the network undo the normalization if the identity mapping is what it needs (gamma = sqrt(var + eps), beta = mu); without them, forcing zero mean and unit variance would strip representational power from the layer.