Batch Normalization

Consider a layer producing a batch of activations $x_1, \dots, x_m$ for a single feature (channel), where $m$ is the mini-batch size. Batch Normalization computes the batch mean $\mu_B$ and variance $\sigma_B^2$, normalizes each $x_i$ to have zero mean and unit variance, and then applies a learned affine transform:

Here $\epsilon$ is a small constant (e.g. $10^{-5}$) preventing division by zero, and $\gamma, \beta$ are learnable parameters per feature, trained jointly with the rest of the network by gradient descent. This normalization is applied independently to each feature/channel: for a fully connected layer, each of the $d$ output units gets its own $(\gamma, \beta)$ pair; for a convolutional layer, each output channel gets its own pair, with statistics pooled over the batch and the spatial dimensions (height, width) of that channel. The role of $\gamma, \beta$ is crucial and often misunderstood: without them, forcing every activation to zero mean and unit variance would restrict what the layer can represent (for instance, a sigmoid fed only near-zero inputs behaves almost linearly). The affine parameters let the network undo the normalization if useful — in particular, setting $\gamma = \sqrt{\sigma_B^2+\epsilon}$ and $\beta = \mu_B$ exactly recovers the original activations. So BatchNorm does not remove representational power; it reparameterizes the layer so that the scale and mean of the activations are controlled by dedicated, easy-to-optimize parameters, decoupled from the complicated interactions among all upstream weights. This decoupling is why gradients become better behaved: a change in an upstream weight that previously would have rescaled an entire layer's activations (and hence needed to be compensated by every downstream layer) is now largely absorbed by the normalization step itself.

Placement. BatchNorm is typically inserted right after a linear or convolutional layer and before the nonlinearity: `Linear → BatchNorm → ReLU`. Because BatchNorm re-centers its input, the bias term of the preceding linear layer becomes redundant (it would be cancelled by mean subtraction) and is usually dropped.

Two computational modes.

• Training mode: statistics $\mu_B, \sigma_B^2$ are computed from the current mini-batch. This makes the output for a given example depend on the other examples in its batch — an unusual but essential property.
• Inference mode: mini-batches may not exist (e.g., single-example inference), and even when they do, using batch statistics would make predictions non-deterministic and dependent on which other examples happen to share the batch. Instead, BatchNorm uses fixed population estimates $\hat\mu, \hat\sigma^2$, obtained by tracking an exponential moving average of the batch statistics seen throughout training:

with momentum $\alpha$ typically around $0.1$. At inference time the normalization becomes a fixed, deterministic per-feature affine map:

Forward pass pseudocode

``` function BatchNormForward(X, gamma, beta, eps, mode, running_mean, running_var, momentum):

# X: shape (m, d) -- m examples, d features (for conv, d channels, pool over spatial dims too) if mode == TRAIN: mu = mean(X, axis=0) # shape (d,) var = mean((X - mu)^2, axis=0) # biased variance, shape (d,) Xhat = (X - mu) / sqrt(var + eps) Y = gamma Xhat + beta running_mean = (1 - momentum) running_mean + momentum mu running_var = (1 - momentum) running_var + momentum var cache = (X, Xhat, mu, var, gamma, eps) return Y, cache, running_mean, running_var else: # INFERENCE Xhat = (X - running_mean) / sqrt(running_var + eps) Y = gamma Xhat + beta return Y ```

Backward pass pseudocode (exact gradients)

Let $N = m$ (batch size) for a given feature column. Given upstream gradient $dY$ (same shape as $X$), we need $dX$, $d\gamma$, $d\beta$. Because $\mu_B$ and $\sigma_B^2$ are themselves functions of every $x_i$ in the batch, the backward pass must differentiate through the batch statistics — this is the one subtlety that distinguishes BatchNorm's backward pass from an ordinary elementwise layer. ``` function BatchNormBackward(dY, cache): (X, Xhat, mu, var, gamma, eps) = cache N = X.shape[0] std_inv = 1 / sqrt(var + eps) dgamma = sum(dY Xhat, axis=0) # shape (d,) dbeta = sum(dY, axis=0) # shape (d,) dXhat = dY gamma # shape (m, d) dvar = sum(dXhat (X - mu), axis=0) (-0.5) std_inv^3 dmu = sum(dXhat (-std_inv), axis=0) + dvar mean(-2 (X - mu), axis=0) dX = dXhat std_inv \ + dvar 2 (X - mu) / N \ + dmu / N return dX, dgamma, dbeta ``` This decomposition is exact: it comes from applying the multivariate chain rule to the computational graph $x \to \mu_B \to \sigma_B^2 \to \hat x \to y$, treating $\mu_B(x_1,\dots,x_m)$ and $\sigma_B^2(x_1,\dots,x_m)$ as functions of the entire* batch, so $\partial \hat x_i/\partial x_j \neq 0$ even for $i\neq j$.

Time complexity. For a mini-batch of size $m$ and $d$ features, computing $\mu_B$ and $\sigma_B^2$ each cost $O(md)$ (a single pass, or two passes, over the batch); normalizing and applying the affine transform costs another $O(md)$. So the forward pass is $O(md)$, exactly the same asymptotic order as the linear/convolutional layer it augments — BatchNorm adds a small constant-factor overhead, not a new complexity class. For a convolutional layer with batch size $m$, spatial size $H\times W$, and $C$ channels, statistics are pooled over $m\cdot H\cdot W$ values per channel, giving $O(mHWC)$ work, again the same order as the convolution's own activation volume. The backward pass performs a fixed number of elementwise and reduction operations over the same $O(md)$ (or $O(mHWC)$) data, so it is also $O(md)$ — no asymptotic blowup relative to a naive elementwise nonlinearity's backward pass.

Space complexity. The cache needed for backward — $X$, $\hat X$, $\mu$, $\sigma^2$, $\gamma$ — is $O(md)$ (dominated by storing the batch itself, which the rest of the network stores anyway), plus $O(d)$ for the running statistics and learnable parameters. This is the same order as storing the layer's activations for any other elementwise operation, so BatchNorm does not change the network's overall memory asymptotics, though in practice it roughly doubles the activation memory of the layer it wraps (raw and normalized copies must both be kept for the backward pass).

Correctness of the backward pass. The gradient formula above is not a heuristic; it is the exact chain rule applied to the computational graph

Because $\mu_B = \frac1N\sum_k x_k$ and $\sigma_B^2 = \frac1N\sum_k (x_k-\mu_B)^2$ both depend on every $x_k$ in the batch, the Jacobian $\partial \hat x_i / \partial x_j$ is a full (not diagonal) $N\times N$ matrix for each feature. Writing $L$ for the scalar loss and using $dY = \partial L/\partial Y$ as given upstream, the multivariate chain rule gives

and expanding $\partial \hat x_k/\partial x_i$ through both the direct path ($x_i$ appearing in $\hat x_i$ itself) and the indirect paths through $\mu_B$ and $\sigma_B^2$ (since $x_i$ contributes to every $\hat x_k$ via the shared statistics) produces exactly the three-term sum in the pseudocode: a term from the direct normalization, a term proportional to $dvar$, and a term proportional to $dmu$. One can verify correctness numerically via the standard technique of comparing analytic gradients to finite differences $\frac{L(x+\delta e_i) - L(x-\delta e_i)}{2\delta}$; the formula above matches to numerical precision, confirming there is no missing term.

Why it helps optimization (the honest picture). The original motivation was reducing "internal covariate shift" — the claim that normalizing layer inputs keeps their distribution stable across training steps, letting downstream layers avoid constantly re-adapting. Subsequent analysis showed the primary mechanistic benefit is better captured through the optimization landscape: BatchNorm empirically makes the loss surface smoother and its gradients more Lipschitz (predictable), which permits larger stable learning rates and reduces sensitivity to the scale of initial weights. Concretely, for a layer output $y = \gamma \hat x + \beta$, rescaling an upstream weight by a constant $c$ leaves $\hat x$ unchanged (since normalization divides out the resulting change in $\sigma_B$), so gradients with respect to that weight scale as $1/c$ — the optimization automatically becomes invariant to weight rescaling, a self-stabilizing property that plain networks lack.

Tradeoffs. BatchNorm couples examples within a batch (a single example's output depends on its batch-mates), which is a correctness wrinkle for architectures needing per-example independence (contrastive setups, some generative models), and it introduces a train/inference mismatch that must be handled correctly (running statistics) or accuracy silently degrades. Its statistics become noisy and unreliable for very small batch sizes, and it requires synchronization overhead in distributed/multi-device training if you want statistics computed across all devices ("SyncBN").

Take a single feature (channel) with a mini-batch of $m=4$ raw pre-activation values, $\gamma = 2$, $\beta = 1$, $\epsilon = 10^{-5}$ (rounded to $0$ below for readability):

Step 1 — batch mean:

Step 2 — batch variance (biased, divide by $m$):

Step 3 — normalize ($\sqrt{2.0 + \epsilon} \approx 1.41421$):

Step 4 — affine transform $y_i = 2\hat x_i + 1$:

Step 5 — running statistics update (momentum $\alpha = 0.1$, assuming $\hat\mu = 0, \hat\sigma^2 = 1$ before this step):

Backward pass, suppose upstream gradient $dY = [1, -1, 1, -1]$ (arbitrary, for illustration):

$dXhat_i = dY_i \cdot \gamma = [2, -2, 2, -2]$. Then with $std\_inv = 1/\sqrt{2.00001} \approx 0.70711$:

The bracketed sum is $-4 + 0 + 4 + 0 = 0$, so $dvar = 0$ here — a consequence of the particular symmetric $dY$ chosen. Continuing,

Finally,

This trace deliberately used a non-symmetric raw input ($x_3 = 6$ breaking naive symmetry) and confirms every quantity — $\mu_B, \sigma_B^2, \hat x, y, d\gamma, d\beta, dX$ — flows consistently through the formulas, including the case where $dvar$ and $dmu$ happen to vanish (checkable independently by finite differences).

Batch Normalization was introduced by Sergey Ioffe and Christian Szegedy in 2015 while at Google, in the context of training very deep convolutional networks (like Inception-style architectures) faster and more reliably. It was one of the key innovations that, alongside residual connections, made it practical to train networks with dozens to over a hundred layers using ordinary stochastic gradient descent, directly enabling the rapid progress in image classification benchmarks in the following years.

BatchNorm (or a normalization descended from it) sits inside essentially every modern convolutional network: ResNet, Inception, DenseNet, and their descendants used in production image classifiers, object detectors, and segmentation systems all rely on it to train stably at depth. It made possible the standard recipe of training very deep vision models with plain SGD or momentum and comparatively simple learning-rate schedules, tying directly into the material on Learning-Rate Schedules — networks with BatchNorm tolerate much larger peak learning rates than those without it. In production ML pipelines, correctly switching between train and inference modes (and syncing running statistics across distributed workers) is a routine but critical engineering concern, closely tied to the broader discipline of Debugging the Training Process, since a mishandled BatchNorm mode is one of the most common silent bugs separating strong training-time metrics from poor deployed performance. Its direct descendants — Layer Normalization, Group Normalization, Instance Normalization, and RMSNorm — now dominate transformer-based language and vision models precisely because they were designed to fix BatchNorm's batch-size and sequence-length dependencies while keeping its core insight: explicitly controlling the scale and mean of activations, with learnable parameters, is enormously beneficial for training deep networks.

Ask the tutor: How does the derivation of BatchNorm's backward pass change if we normalize over the feature dimension instead of the batch dimension, as in Layer Normalization?