Normalization

A deep stack has a moving-target problem. Every layer is trained assuming the distribution of its inputs is fixed, but those inputs are the outputs of the layers below, whose weights are changing on every step. The signal that arrives at layer $k$ drifts in mean and scale throughout training, so each layer must continually re-adapt to a drifting distribution. Normalization attacks this directly: rescale the activations to a fixed first and second moment inside the forward pass, then let the network learn back any scale it actually needs. The standardization is a differentiable layer, not a preprocessing step on the data; its statistics flow through backpropagation like everything else.

The one decision that separates the whole family is which axes you average over. Standardizing is always the same three operations — subtract a mean, divide by a standard deviation, apply a learnable affine map. What changes from batch norm to layer norm to group norm is only the set of tensor entries pooled into each mean and variance. This lesson derives batch normalization in full (forward pass, running statistics, the backward pass through the batch moments), then walks the family by re-choosing that axis set, and closes on why the whole construction speeds up and regularizes training.

Batch normalization

Fix one feature (one coordinate of the activation vector, or one channel of a convolutional map). Over a minibatch $\mathcal{B}=\{x_1,\dots ,x_m\}$ of that feature's values, batch normalization computes the batch mean and variance, standardizes, then applies a learnable affine map:

The small constant $ϵ$ (typically $10^{-5}$) keeps the denominator away from zero. The pair $(\gamma ,\beta )$ is learned per feature by gradient descent; they are the layer's only parameters.¹

Read the four equations in order. The first two summarize the batch: ${\mu }_{\mathcal{B}}$ is where this feature's values sit on average, ${\sigma }_{\mathcal{B}}^2$ is how spread out they are. The third centers and rescales every value so the batch has mean $0$ and variance $1$; after this step the feature is on a fixed, canonical scale no matter what the layers below produced. The fourth returns two learnable parameters to the network: $\gamma$ stretches the standardized feature and $\beta$ shifts it, so the network can place the final distribution anywhere.

A concrete pass. Take a minibatch of $m=4$ values of one feature, $x=(2,4,6,8)$. Then ${\mu }_{\mathcal{B}}=5$ and ${\sigma }_{\mathcal{B}}^2=\frac{1}{4}((-3{)}^2+(-1{)}^2+1^2+3^2)=5$, so (ignoring $ϵ$) the standardized values are $\hat{x}=(-3,-1,1,3)/\sqrt{5}\approx (-1.34,-0.45,0.45,1.34)$ — mean $0$, variance $1$, as promised. With learned $\gamma =2,\beta =1$ the outputs are $y\approx (-1.68,0.11,1.89,3.68)$: recentered on $\beta =1$ and spread by $\gamma =2$. Change $x$ to $(12,14,16,18)$ and $\hat{x}$ is unchanged: the standardization is invariant to any shift of the whole batch, which is the point.

Where the parameters live. Normalization runs per feature, so the statistics and parameters are vectors indexed by the channel axis. For a convolutional activation of shape $(N,C,H,W)$ — batch $N$, channels $C$, height $H$, width $W$ — batch norm computes one ${\mu }_{\mathcal{B}}$ and one ${\sigma }_{\mathcal{B}}^2$ per channel by pooling the $N\cdot H\cdot W$ entries of that channel, and it holds one $\gamma$ and one $\beta$ per channel. The parameter count is therefore $2C$, independent of batch size or spatial resolution.

The affine map is what keeps the layer from losing expressive power. Forcing unit variance is a constraint; the network can cancel it whenever the constraint hurts.

The standardization moves a feature's distribution to a canonical position and the affine map relocates it — now under the optimizer's explicit control through two parameters, rather than as a byproduct of the layers below.

The forward pass as an algorithm

Collecting the four equations gives the training-time forward pass for a single feature; in practice the same operations run vectorized over every feature at once.

The layer is differentiable end to end. Backprop must route gradients not only through the affine $y_i=\gamma {\hat{x}}_i+\beta$ but also through ${\mu }_{\mathcal{B}}$ and ${\sigma }_{\mathcal{B}}^2$, since both depend on every $x_j$ in the batch — so the gradient with respect to one input mixes in contributions from all the others.

The parameter gradients are the easy half: $\gamma$ and $\beta$ appear only in the final affine, so they just sum the upstream gradient over the batch, weighted by ${\hat{x}}_i$ for the scale and unweighted for the shift:

The input gradient is the interesting half. Writing $g_i=\partial L/\partial {\hat{x}}_i=\gamma \partial L/\partial y_i$ for the gradient after the affine, the chain rule carries $g_i$ back through the three paths by which $x_i$ reaches the loss: directly through its own ${\hat{x}}_i$, through the shared mean ${\mu }_{\mathcal{B}}$, and through the shared variance ${\sigma }_{\mathcal{B}}^2$. Collecting the three paths gives

with $g_i=\gamma \partial L/\partial y_i$. The first term is the naive gradient you would get if ${\mu }_{\mathcal{B}}$ and ${\sigma }_{\mathcal{B}}^2$ were constants. The second subtracts the batch mean of the upstream gradient — the correction from $x_i$'s influence on ${\mu }_{\mathcal{B}}$. The third subtracts an $\hat{x}$-weighted mean — the correction from its influence on ${\sigma }_{\mathcal{B}}^2$. Together the two corrections project the incoming gradient onto the subspace orthogonal to constant shifts and to rescalings of $\hat{x}$; that projection is what leaves the layer's output invariant to the mean and scale of its input, and it is the reason an example's gradient depends on which other examples share its batch.²

The dataflow makes the coupling visible: every input feeds the two batch statistics, which in turn feed the normalization of every input, before the per-feature affine finishes the map.

Train versus inference

The batch statistics ${\mu }_{\mathcal{B}},{\sigma }_{\mathcal{B}}^2$ are a property of the batch, not the example — they are unavailable at test time, where inputs may arrive one at a time and predictions must not depend on which other examples happen to be present. The fix is to accumulate running estimates during training (lines 6–7 of the algorithm, an exponential moving average with momentum $\alpha$) and freeze them for inference:

The running statistics $r_{\mu },r_v$ are buffers, not parameters: updated by a formula, never by a gradient. Each training step folds the current batch moments into the running estimate with momentum $\alpha$ (typically $\alpha \approx 0.1$, or its complement $0.9$ depending on the convention),

so that after enough steps $r_{\mu },r_v$ track the population mean and variance of the activation while the batch statistics still do the normalizing during training. At test time the layer becomes a fixed affine map of its input,

a single scale-and-shift with no batch and no per-example coupling, which is why it can be folded into the preceding linear layer at inference for free.³

The train/test asymmetry is the single richest source of batch-norm bugs. Confuse the two modes and the network silently misbehaves:

The first row is the canonical one: with the layer in training mode at test time, each prediction normalizes by the statistics of whatever batch it arrives in, so the same input yields different outputs depending on its neighbors.

Why it helps

Batch norm reliably lets you train deeper networks at higher learning rates, but why has been contested. Two explanations compete.

The original account was internal covariate shift: because lower layers keep changing, the distribution feeding each layer drifts, forcing upper layers to perpetually re-adapt. Fixing the mean and variance of every layer's input is meant to halt this drift, so each layer sees a stationary distribution and can learn faster.⁴

The modern account is that the covariate-shift explanation is at best incomplete: networks can be trained with deliberately injected post-norm distribution shift and still retain batch norm's benefits. The better-supported explanation is that normalization conditions the optimization — it makes the loss and its gradient smoother (smaller, more predictable curvature along the update direction), so larger, more reliable steps are safe. The decoupling of each weight's effective scale from the loss is a second mechanism: scaling a weight by $c$ leaves a normalized layer's output unchanged, which reduces the loss surface's sensitivity to weight magnitude.⁵

The family: which axes are normalized

Batch norm is one choice of which axes to average over. For a 4-D activation tensor of shape $(N,C,H,W)$ — batch $N$, channels $C$, spatial $H\times W$ — the variants differ only in the set of axes contributing to each $(\mu ,{\sigma }^2)$. Layer norm averages over the features of a single example (so it is batch- independent); instance norm normalizes each channel of each example on its own; group norm splits channels into groups and normalizes within a group, a middle ground that, like layer/instance norm, never touches the batch axis.

The cleanest way to see the whole family is as different shadings of the same tensor. Lay the activation out as an $N\times C$ grid of spatial maps (each cell is one $H\times W$ slice). Every scheme computes its statistics over the shaded cells, and only the shading changes.

The table below makes the trade-offs explicit. The decisive column is batch- dependent? — the one property that determines whether a layer behaves identically at train and test, and whether it survives a batch of size one.

Layer norm, written out. For one example with feature vector $x\in {\mathbb{R}}^d$, layer norm pools all $d$ features into a single mean and variance:

The formulas are batch norm's with the sum running over features instead of over the batch. Because no other example enters, the computation is identical at train and test time, and a batch of size one is handled exactly like a batch of a thousand — the property that makes it the default in Transformers, where the batch axis carries variable-length sequences, and in recurrent networks, where the same statistics must apply at every timestep.

RMS norm drops the mean-centering: it divides by the root-mean-square of the features and keeps only a scale, no shift.

Skipping the subtraction of $\mu$ costs almost nothing in accuracy on large Transformers while removing one reduction and one buffer per layer, which is why recent large language models tend to use it in place of layer norm.

Batch norm's accuracy degrades when the batch is small because ${\mu }_{\mathcal{B}},{\sigma }_{\mathcal{B}}^2$ become noisy estimates of the population moments; the batch-independent variants sidestep this entirely, which is why detection and segmentation models (forced into tiny batches by memory) use group norm, and why sequence models, whose batch axis carries variable-length examples, use layer norm.⁶

Placement

Where the layer sits relative to the linear map and the activation is a genuine design choice. The original prescription was linear $\to$ norm $\to$ activation, so the normalized pre-activation feeds the nonlinearity; a common alternative normalizes the activation output instead. In residual architectures the analogous question is pre-norm versus post-norm: whether normalization sits inside the residual branch (before the sublayer) or after the residual addition.

Pre-norm leaves an unnormalized identity path from input to output of each block, so the gradient reaches early layers without passing through a normalization at every step — the reason very deep Transformers train more stably with pre-norm, often without learning-rate warmup. Post-norm can reach slightly better final quality when it trains at all, but is more fragile at depth.

Normalization, revisited

Goodfellow §8.7.1 introduces batch normalization (Ioffe & Szegedy, 2015) and gives the internal-covariate-shift motivation; the public literature since has both challenged that motivation and multiplied the variants. The covariate-shift explanation was directly tested by Santurkar et al. (2018), who trained batch-norm networks with deliberately injected post-normalization distribution shift and found the training speedup intact — evidence that the operative mechanism is the loss-surface smoothing described above, not the drift the name suggests. The technique works; the name is a historical accident.

The batch-independent variants each arose from a specific need. Layer norm (Ba et al., 2016) freed normalization from the batch axis for recurrent and sequence models; instance norm (Ulyanov et al., 2016) suited style transfer; group norm (Wu & He, 2018) filled the small-batch gap for detection and segmentation; and RMS norm (Zhang & Sennrich, 2019) dropped the mean-centering for a cheaper layer. The decisive outcome is architectural: modern large language models are built almost entirely on pre-norm layer or RMS normalization, because the clean identity path through each residual block (shown above) is what lets very deep transformers train stably, often without the learning-rate warmup post-norm requires. Batch norm remains the default for convolutional image classifiers with large batches, but the sequence world it never fit well moved to the batch-independent family and stayed there.

The regularizing side effect

Batch norm is filed under regularization for a concrete reason: it injects noise. Each example's normalized value ${\hat{x}}_i$ depends on the random minibatch it was sampled with, through ${\mu }_{\mathcal{B}}$ and ${\sigma }_{\mathcal{B}}^2$. Two identical inputs in different batches get slightly different activations — a stochastic perturbation closely analogous to dropout, which is why models with batch norm often need less of it.⁷

The effect cuts both ways: the noise that regularizes also destabilizes small batches and forces the train/test statistic split. Normalization trades a constraint on the activations for better conditioning and a mild dose of noise, at the bookkeeping cost of maintaining running statistics. It composes with the other tools in this module: parameter penalties, dropout and augmentation, and early stopping all attack overfitting from different angles, and normalization is the one that also reshapes the loss surface the optimizer has to descend.