📝 Batch Normalization

Description

< What is it? >

[Batch normalization (batch norm)]
Batch norm normalizes a layer's inputs across the current mini-batch, then rescales the result with two learnable parameters. It is a layer, not a preprocessing step: it runs inside the forward pass, every step, on activations whose distribution keeps shifting as the weights below it change.

[Mini-batch]
A mini-batch is a small, randomly selected subset of the training dataset. Instead of passing the entire dataset through the model at once, we divide it into smaller batches. This allows for more efficient training and helps in stabilizing the learning process.

• Why not the whole dataset (full batch) or one sample (stochastic)?
  • Full batch (entire dataset): Slow, memory-heavy, and often gets stuck in sharp local minima.
  • Single sample (stochastic, batch size = 1): Fast updates but very noisy, making training unstable.

[Normalization]
Normalization is a technique used to adjust the values of input data to a common scale, without distorting differences in the ranges of values. In the context of batch normalization, it involves scaling and shifting the inputs to have a mean of zero and a standard deviation of one.

z ──▶ normalize with the batch's μ, σ² ──▶ scale by γ, shift by β ──▶ y
        (zero mean, unit variance)          (learnable, per feature)

Without the second half, forcing every layer to zero mean and unit variance would throw away representational power — a sigmoid unit could never operate anywhere but its linear middle. The learnable $\gamma$ and $\beta$ give that power back: the layer can recover any mean and scale it needs, including the identity.

< The computation >

For a mini-batch of $m$ examples, computed per feature (per channel, for images):

$$\mu_B = \frac{1}{m}\sum_{i=1}^{m} z_i, \qquad \sigma_B^2 = \frac{1}{m}\sum_{i=1}^{m}\left(z_i - \mu_B\right)^2$$ $$\hat{z}_i = \frac{z_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}, \qquad y_i = \gamma\,\hat{z}_i + \beta$$

Reading the symbols

— the subscript $B$ just means "over the batch".

Symbol	Name	What it is
$\mu_B$	miu	The mean of the feature across the mini-batch
$\sigma_B^2$	sigma	The variance across the mini-batch. $\sigma_B$ on its own is the standard deviation, which is what the formula's square root recovers
$\epsilon$	epsilon	A small constant, typically 1e-5, that guards the division when a feature is constant across the batch
$\gamma$	gamma	The learned scale, one per feature
$\beta$	beta	The learned shift, one per feature

$\gamma$ and $\beta$ are trained by backpropagation like any other parameter — one pair per feature, so a BatchNorm2d(64) layer adds 128 trainable values.

Key points

< Why it works >

• Higher learning rates become safe. Consistent activation scales between layers mean an update in one layer cannot blow up the input distribution of the next, so training tolerates a learning rate that would otherwise diverge — see Gradient Descent.
• Less sensitivity to initialization. The layer renormalizes whatever scale it is handed, which widens the range of workable starting weights — though it does not remove the need for sane initialization.
• Activations stay in the responsive range. For sigmoid and tanh it keeps inputs away from the saturated tails where derivatives are near zero — one of the ways it mitigates vanishing gradients.
• It regularizes mildly. Each example's normalization depends on whichever examples happened to share its batch, which injects noise. That slightly reduces overfitting and is why batch norm and dropout are often not both needed.

< Covariate shift, and the "internal" kind >

Covariate shift is a change in the distribution of the input features between the data a model was trained on and the data it later sees. The defining condition is that the rule survives the change:

$$P_{\text{train}}(x) \neq P_{\text{deploy}}(x), \qquad P(y \mid x) \text{ unchanged}$$

Train a self-driving model on sunny daytime photographs and deploy it on foggy night streets: a stop sign is still a stop sign, so $P(y \mid x)$ is the same rule as before — but the pixel distribution is nothing like the training set, and the model has no reason to generalize into that unfamiliar input space. Only $P(x)$ moved, and that was enough.

Internal covariate shift applies the same idea one layer down. Every layer's input is the previous layer's output, and the previous layer's weights change on every step — so each layer is chasing an input distribution that keeps moving underneath it. The argument is that this slows training, because later layers spend their capacity re-adapting instead of learning.

That was the original motivation for batch normalization: hold each layer's input distribution steady, and the layers above it can stop chasing. Later work (Santurkar et al., 2018) found the explanation does not hold up under measurement — batch norm helps even when noise is deliberately injected to restore the shift — and attributes the benefit instead to a smoother loss landscape with more predictable gradients. The technique works either way; the smoothing account is the better answer in an interview.

The two are worth keeping apart. Internal covariate shift happens during training and batch norm addresses it; ordinary covariate shift happens after deployment and normalization does nothing for it. That one is a data problem — retraining on representative data, or detecting the drift and flagging it. It also breaks calibration, whose guarantees hold only on the distribution the calibrator was fitted to.

< Training vs. inference — the distinction that bites >

Mode	Statistics used	Set by
Training	$\mu$, $\sigma^2$ of the current batch	model.train()
Inference	Running averages accumulated during training	model.eval()

During training the layer also updates a running mean and variance (an exponential moving average, controlled by momentum, default 0.1). At inference those fixed running values are used instead, so a prediction depends only on the input — not on whatever else happened to be batched alongside it.

Forgetting model.eval() is the classic bug: predictions change depending on batch composition, evaluation scores move between runs, and a batch of size 1 normalizes an example against itself, producing zeros. model.eval() also switches dropout off, which is why it is a single call for both.

< In PyTorch >

Layer	Input shape	Normalizes over	Typical use
nn.BatchNorm1d(C)	(N, C) or (N, C, L)	N (and L)	Fully connected layers, 1-D signals
nn.BatchNorm2d(C)	(N, C, H, W)	N, H, W	Images — the common case in CNNs
nn.BatchNorm3d(C)	(N, C, D, H, W)	N, D, H, W	Volumetric data, video

import torch.nn as nn

model = nn.Sequential(
    # bias=False: BatchNorm's β makes the preceding bias redundant
    nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
    nn.BatchNorm2d(64),
    nn.ReLU(),
)

model.train()   # batch statistics; running averages updated
model.eval()    # running averages; nothing updated

The argument is the feature/channel count, not the batch size. Two practical notes: the preceding layer's bias is redundant (β already provides the offset, so use bias=False), and the conventional order is Linear/Conv → BatchNorm → activation, following the original paper.

< Batch size sensitivity >

Batch norm estimates $\mu$ and $\sigma^2$ from the batch, so the estimates get noisier as the batch shrinks. Below roughly 8–16 examples the noise starts to hurt more than it regularizes; at a batch size of 1 the layer is undefined in training mode. This is the direct link between batch norm and your choice of mini-batch size, and it is why memory-constrained training (large images, detection, segmentation) often switches to a batch-size-independent alternative.

For those cases — and for transformers and RNNs, where sequence examples vary in length — layer norm or RMSNorm normalize within a single example instead. See Normalization inside the network for the comparison.

< Related ideas >

• Data Normalization scales the inputs before training; batch norm scales activations during it. Doing one does not remove the need for the other.
• Vanishing & Exploding Gradients covers what batch norm can and cannot fix.
• Regularization puts its noise-based side effect in context.
• Calibration is one of the first things ordinary covariate shift breaks after deployment.

Crash course

Batch Norm vs Layer Norm - Explained

Why Batch Normalization Works

Why Does Batch Norm Work? (C2W3L06)

Reference

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift — Ioffe & Szegedy (2015)
• How Does Batch Normalization Help Optimization? — Santurkar et al. (2018)
• PyTorch nn.BatchNorm2d documentation