Batch and layer normalization

The idea

Every normalization layer makes the same two-step move: standardize the activations, then re-scale with learned parameters.

$\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad y = \gamma\, \hat{x} + \beta$

The $\epsilon$ term guards against division by zero. The learned per-feature pair $\gamma, \beta$ restores a learnable scale and shift after standardization. It does not generally undo normalization: the sample- or batch-dependent $\mu$ and $\sigma$ have already been removed, and fixed learned parameters cannot reconstruct them for every input. Normalization therefore changes the representation and optimization geometry; surrounding layers can adapt to that representation, but the transform is not generally invertible.

The entire design space is one question: which axis do you average over?

Batch norm: statistics across the batch. For each feature separately, treat the $m$ examples of the minibatch as the sample:

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

$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}, \qquad y_i = \gamma\, \hat{x}_i + \beta$

In convnets the "feature" is a channel and the statistics pool over batch and spatial positions, so each channel gets plenty of samples even at modest batch sizes.

The train/inference asymmetry — the #1 source of BN bugs. At inference there may be no batch at all (one user request), and predictions must not depend on whoever else is in the batch. So during training the layer also maintains exponential running averages of $\mu$ and $\sigma^2$, and at inference it normalizes with those frozen values instead — BN becomes a fixed affine map, which deployment toolchains routinely fold into the preceding layer's weights. The mode switch is where things go wrong: a forgotten model.eval() means eval-time predictions still use live batch statistics — they change with batch composition and the metrics look broken. The telltale signature: train metrics great, eval metrics garbage, and the gap "disappears" if you evaluate in train mode.

Why it helps. The 2015 paper's story was internal covariate shift: as upstream weights move, each layer's input distribution drifts, and normalizing supposedly spares downstream layers from chasing it. Later work (Santurkar et al. 2018) undercut that account — BN still helps when distribution shift is injected right after the norm — and located the benefit in optimization geometry: a smoother loss landscape with better-conditioned, more predictable gradients. That is what buys the practical wins: larger stable learning rates, far less sensitivity to initialization, very deep stacks that train reliably. Side effect worth naming: each example's normalization depends on its randomly drawn batchmates, and that noise acts as a mild regularizer — real, but never your main one.

Where BN strains. Batch statistics are estimates from $m$ samples. Below roughly 8 they get noisy enough to hurt accuracy outright; in distributed training it is the per-GPU batch that counts unless you pay for synchronized BN. BN also couples examples inside a batch — a leakage channel in contrastive learning and RL. For a plain dense BN feature with one scalar observation, batch size 1 gives zero variance and output $\beta$. Convolutional BN is different: it also pools over spatial positions, so a $1 \times C \times H \times W$ input can retain nonzero per-channel variance when $H W > 1$. Either way, tiny effective sample counts make statistics fragile.

Layer norm: statistics across features. Flip the axis. For each sample $i$ independently, average over its $d$ features:

$\mu_i = \frac{1}{d}\sum_{j=1}^{d} x_{ij}, \qquad \sigma_i^2 = \frac{1}{d}\sum_{j=1}^{d} (x_{ij} - \mu_i)^2$

then the same normalize-and-affine ($\gamma, \beta$ still per feature). No batch statistics anywhere: batch size 1 works, every token in a variable-length sequence is normalized on its own, and train and eval are identical — no running averages, no mode bug. That batch independence is why transformers and RNNs use LN: batch norm over a batch of padded, unrelated sequences would blend statistics across sentences and padding tokens.

Pre-LN vs post-LN. In a transformer block the norm can sit after the residual add (post-LN, the 2017 original) or inside the branch, before attention and the FFN (pre-LN). Post-LN places norms on the main signal path, so every gradient must pass through every norm on the way down — deep post-LN stacks are fragile at the start of training and need learning-rate warmup to avoid divergence. Pre-LN keeps the skip path a clean identity from the loss to the embeddings, so gradients flow unimpeded and deep stacks train stably with little or no warmup — the modern default (with RMSNorm, layer norm minus the mean subtraction, as the cheaper variant in current LLMs). Well-tuned post-LN can still edge ahead in final quality, which is why both survive in the wild.

Worked example

Batch norm by hand: one feature, batch of $m = 4$ values $x = (1, 3, 5, 7)$, learned $\gamma = 2$, $\beta = 1$, $\epsilon = 10^{-5}$.

Batch statistics.

$\mu_B = \frac{1 + 3 + 5 + 7}{4} = 4, \qquad \sigma_B^2 = \frac{(-3)^2 + (-1)^2 + 1^2 + 3^2}{4} = \frac{20}{4} = 5$

The denominator $\sqrt{5 + 10^{-5}} \approx 2.2361$ — $\epsilon$ is invisible at this scale, exactly as designed.

Normalize.

$\hat{x} = \left(\frac{-3}{2.2361},\; \frac{-1}{2.2361},\; \frac{1}{2.2361},\; \frac{3}{2.2361}\right) = (-1.342,\; -0.447,\; 0.447,\; 1.342)$

Scale and shift with $y = 2\hat{x} + 1$:

$y = (-1.683,\; 0.106,\; 1.894,\; 3.683)$

Sanity check: the outputs have mean $1 = \beta$ and standard deviation $2 = \gamma$ (up to the vanishingly small $\epsilon$) — normalize-then-affine guarantees this on the training batch. (One framework wrinkle: PyTorch updates the running variance buffer with the unbiased estimate $\frac{m}{m-1}\sigma_B^2 = \frac{20}{3} \approx 6.67$, while the forward pass above divides by $m$ — a classic "why doesn't my hand calculation match the buffer" moment.)

Inference. Suppose the running averages have settled at $\mu = 3.8$, $\sigma^2 = 5.4$. A new input $x = 5$ is normalized with those frozen values, regardless of what else is in the batch:

$\hat{x} = \frac{5 - 3.8}{\sqrt{5.4 + \epsilon}} = \frac{1.2}{2.3238} \approx 0.516, \qquad y = 2(0.516) + 1 \approx 2.033$

Layer norm twist: reinterpret $(1, 3, 5, 7)$ as one sample's four features and the arithmetic above is exactly a layer norm computation — same mean 4, same variance 5, same outputs. The formulas are identical; the only thing that ever differs is the axis you average over.

Cheat sheet

Formulas

• BN (per feature, across batch): $\hat{x}_i = (x_i - \mu_B)/\sqrt{\sigma_B^2 + \epsilon}$, $y_i = \gamma \hat{x}_i + \beta$; $\mu_B, \sigma_B^2$ from the current minibatch (biased, divide by $m$).
• BN at inference: same form with frozen running $\mu, \sigma^2$ — a fixed affine map, foldable into the previous layer.
• LN (per sample, across $d$ features): identical form with $\mu_i, \sigma_i^2$ computed within each example.
• Degenerate scalar case: dense BN with one observation → $\sigma_B^2 = 0$, output $= \beta$. Convolutional BN also pools spatial positions, so batch size 1 alone does not force zero variance.

BN vs LN

Symptom → fix

Named variants worth recognizing

• GroupNorm: LN-style but over channel groups — the standard BN escape hatch for small-batch vision (detection, segmentation).
• InstanceNorm: per sample, per channel, over spatial positions — style transfer.
• RMSNorm: LN without mean subtraction or $\beta$ — cheaper, standard in modern LLMs.
• SyncBN: batch statistics all-reduced across GPUs — restores effective batch size in distributed training.
• BN folding: at export, merge frozen BN into the preceding conv/linear weights — free inference speedup.

Rules of thumb: convnets with healthy batches → BN; transformers, RNNs, batch size small or variable → LN; BN's regularization is a side effect — never budget for it.