Batch Normalization › The Batch Normalization Algorithm

Batch normalization operates on a mini-batch of activations. For a single feature dimension (one neuron's pre-activation), consider a mini-batch $\mathcal{B} = \{x_1, x_2, \ldots, x_m\}$ of $m$ samples. Batch normalization transforms these values through four steps.

This section covers all four steps: computing the batch mean, computing the batch variance, normalizing, and then applying a learned scale and shift.

First, compute the mean activation across the mini-batch:

$$\mu_{\mathcal{B}} = \frac{1}{m} \sum_{i=1}^{m} x_i$$

where:

• $\mu_{\mathcal{B}}$: the mean of activations over the current mini-batch
• $m$: the number of samples in the mini-batch
• $x_i$: the activation value for sample $i$

This mean captures the "center" of the activation distribution for this batch. If activations are consistently large and positive, the mean will be large and positive. The goal is to shift this center to zero. Note that this is computed independently for each feature dimension: if a layer has 256 neurons, there are 256 separate mean computations, one for each neuron's activation values across the batch.

Next, measure how spread out the activations are around the mean:

$$\sigma^2_{\mathcal{B}} = \frac{1}{m} \sum_{i=1}^{m} (x_i - \mu_{\mathcal{B}})^2$$

where:

• $\sigma^2_{\mathcal{B}}$: the variance of activations over the current mini-batch
• $(x_i - \mu_{\mathcal{B}})^2$: the squared deviation of each activation from the batch mean

This variance tells us how much the activations fluctuate around their mean. High variance means some activations are very large and others very small. Our goal is to compress this spread to approximately one. Note that this uses the population variance formula (dividing by $m$, not $m-1$), which is a deliberate design choice: the original paper uses population statistics, and PyTorch's implementation follows this convention by default.

With the mean and variance in hand, we normalize each activation:

$$\hat{x}_i = \frac{x_i - \mu_{\mathcal{B}}}{\sqrt{\sigma^2_{\mathcal{B}} + \epsilon}}$$

where:

• $\hat{x}_i$: the normalized activation for sample $i$
• $\epsilon$: a small positive constant (typically $10^{-5}$) added for numerical stability to prevent division by zero when variance is very small

This step produces activations with zero mean and unit variance within the batch. The $\epsilon$ term prevents numerical problems when the variance is close to zero, which can happen if all activations in the batch are nearly identical. Concretely, if a batch happens to contain nearly identical inputs to a neuron (for example, if the preceding layer has collapsed its outputs to a constant), the variance would be zero and division would produce either infinity or NaN. The $\epsilon$ ensures the denominator is always at least $\sqrt{\epsilon} \approx 0.003$, which is small enough not to distort the normalization but large enough to prevent numerical failure.

Pure normalization would be too aggressive. Consider a layer followed by a sigmoid activation: if all inputs to the sigmoid are normalized to zero mean and unit variance, they land in the linear regime of the sigmoid, essentially turning the sigmoid into a linear function. The network would lose its nonlinearity. More generally, normalization restricts what each layer can represent: a layer that learned to produce large-scale outputs for particular inputs would have that learned information erased by normalization.

To restore representational flexibility, batch normalization introduces two learnable parameters per feature:

$$y_i = \gamma \hat{x}_i + \beta$$

where:

• $\gamma$: the scale parameter (learned during training), initialized to 1
• $\beta$: the shift parameter (learned during training), initialized to 0
• $y_i$: the final batch-normalized output for sample $i$

These parameters allow the network to learn the optimal output scale and shift for each feature. If the network determines that a particular feature dimension is most useful when centered at 5 with a spread of 2, gradient descent will drive $\beta \to 5$ and $\gamma \to 2$. In effect, $\gamma$ and $\beta$ allow the network to undo the normalization if that is beneficial, making normalization a learnable transformation rather than a rigid constraint.

This might seem paradoxical: why normalize and then allow the network to un-normalize? even if $\gamma$ and $\beta$ eventually learn to reproduce the original un-normalized distribution, the gradients flowing through the network during training are better conditioned because of the normalization. The gradient updates to $\gamma$ and $\beta$ happen in a well-behaved, smooth optimization landscape, whereas the gradient updates to the weights before batch normalization would have faced the full chaos of internal covariate shift. The normalization changes the geometry of the optimization landscape even if the final learned transformation looks similar to the pre-normalization case.

Each feature dimension in each layer has its own independent $\gamma$ and $\beta$. For a layer with $d$ features, batch normalization adds $2d$ learnable parameters. For large networks with many layers and many features per layer, this adds up, but the count is small relative to the weight matrices.

Putting all four steps together, the full batch normalization transformation for feature dimension $j$ across a mini-batch is:

$$\begin{aligned} \mu_j &= \frac{1}{m} \sum_{i=1}^{m} x_{i,j} \\ \sigma^2_j &= \frac{1}{m} \sum_{i=1}^{m} (x_{i,j} - \mu_j)^2 \\ \hat{x}_{i,j} &= \frac{x_{i,j} - \mu_j}{\sqrt{\sigma^2_j + \epsilon}} \\ y_{i,j} &= \gamma_j \hat{x}_{i,j} + \beta_j \end{aligned}$$

where $x_{i,j}$ is the activation of feature $j$ for sample $i$. Notice that normalization happens independently per feature dimension. Each neuron's output is normalized using only that neuron's statistics across the batch.

This independence is important. It means that batch normalization does not impose any relationship between different features or neurons. Each feature dimension is treated as its own mini-distribution, normalized independently, and then rescaled with its own learned parameters. This preserves the network's ability to learn arbitrary correlations between features through the weight matrices, while still ensuring each individual feature maintains a stable distribution.