Batch Normalization

Table of Contents

1. Learning Objectives
2. Prerequisites
3. Notation
4. Core Intuition
5. The Internal Covariate Shift Problem
6. Batch Normalization: Forward Pass
7. Backward Pass: Full Gradient Derivation
8. Scale Invariance and Gradient Stabilization
9. Inference Mode and Running Statistics
10. Batch Normalization as a Regularizer
11. Layer Normalization and RMSNorm
12. Worked Examples
13. Connection to the Broader Curriculum
14. Common Pitfalls and Misconceptions
15. Research Perspective
16. Summary of Takeaways
17. Exercises

Learning Objectives

After reading this chapter, you should be able to:

1. State the internal covariate shift hypothesis and explain how batch normalization addresses it.
2. Derive the batch normalization forward pass for a mini-batch.
3. Derive the complete backward pass through batch normalization, including gradients through batch statistics.
4. Prove scale invariance: normalized activations are unchanged under $\mathbf{w} \to c\mathbf{w}$.
5. Explain inference-mode normalization using running mean and variance.
6. Compare BatchNorm, LayerNorm, and RMSNorm and identify when each is appropriate.

Prerequisites

• Backpropagation — chain rule through computation graphs
• Multivariate Gaussian — mean, variance (optional)

Notation

• $\mathbf{x}^{(k)}$ — Mini-batch of activations
• $\mu_B, \sigma_B^2$ — Batch mean and variance
• $\hat{x}_i$ — Normalized activation
• $\gamma, \beta$ — Learnable scale and shift
• $\mathbb{E}[\cdot], \mathrm{Var}(\cdot)$ — Population moments at inference

Core Intuition

Deep networks compose many layers: $\mathbf{h}_l = \sigma(\mathbf{W}_l \mathbf{h}_{l-1} + \mathbf{b}_l)$. During training, as earlier layers update, the distribution of activations $\mathbf{h}_{l-1}$ shifts — each layer must continuously adapt to a moving input distribution. This internal covariate shift (Ioffe & Szegedy, 2015) slows training and can saturate activation functions, killing gradients.

Batch Normalization (BatchNorm) normalizes each layer's pre-activations to zero mean and unit variance within each mini-batch, then applies learnable scale and shift. This stabilizes the distribution of layer inputs, enables higher learning rates, and acts as a mild regularizer.

Modern Transformers use Layer Normalization instead (normalize across features, not batch), but BatchNorm remains essential for CNNs and provides the conceptual foundation for all normalization methods.

Series context. Chapter 7 (Training Deep Networks), Part I in Volume II.

The Internal Covariate Shift Problem

Definition 1 (Covariate Shift). Covariate shift occurs when the input distribution $P(\mathbf{x})$ changes between training and test time. In deep networks, the input to layer $l$ is the output of layer $l-1$, whose distribution changes as $\mathbf{W}_{l-1}$ updates during training.

Definition 2 (Internal Covariate Shift). The change in the distribution of layer inputs during training, caused by updates to preceding layers, is internal covariate shift.

Proposition 1 (Saturation). If pre-activations $z = \mathbf{w}^T \mathbf{h}$ have large magnitude, sigmoid/tanh activations saturate: $|\sigma'(z)| \approx 0$, causing vanishing gradients.

Proposition 2 (Scale Sensitivity). The gradient $\partial \mathcal{L}/\partial \mathbf{w}$ depends on $\|\mathbf{w}\|$ and $\|\mathbf{h}\|$. Large weight norms amplify gradient magnitudes, requiring careful learning rate tuning.

BatchNorm addresses both by ensuring normalized, well-conditioned inputs to each layer.

Batch Normalization: Forward Pass

Definition 3 (Batch Normalization). Given a mini-batch of pre-activations $\{z_i\}_{i=1}^B$ for a single feature (applied independently per feature dimension):

Step 1 — Batch statistics:

Step 2 — Normalize:

where $\epsilon > 0$ ensures numerical stability (typically $10^{-5}$).

Step 3 — Affine transform (learnable):

where $\gamma$ (scale) and $\beta$ (shift) are learnable parameters per feature.

Important equation. The affine parameters $\gamma, \beta$ allow the network to undo normalization if the optimal representation requires non-zero mean or non-unit variance. Without them, BatchNorm would constrain representational capacity.

Vector form. For a layer with $d$ features and batch $\mathbf{Z} \in \mathbb{R}^{B \times d}$, normalization is applied column-wise (across the batch dimension for each feature).

Backward Pass: Full Gradient Derivation

The backward pass is non-trivial because $\mu_B$ and $\sigma_B^2$ depend on all batch elements.

Setup. Let $L$ be the loss, $g_i = \partial L / \partial y_i$ the upstream gradient, and we seek $\partial L / \partial z_i$.

Step 1 — Gradients of affine parameters:

Step 2 — Gradient w.r.t. normalized activations:

Denote $\bar{g}_i = g_i \gamma$.

Step 3 — Gradients through normalization. Since $\hat{z}_i = (z_i - \mu_B)/\sigma$ where $\sigma = \sqrt{\sigma_B^2 + \epsilon}$:

The second term vanishes because $\sum_i (z_i - \mu_B) = 0$:

Step 4 — Gradient w.r.t. pre-activations:

Theorem 1 (Compact Backward Formula). Defining $\bar{g}_i = (\partial L / \partial y_i)\gamma$:

Proof. Substitute (6), (8) into (9) and simplify using $\sum_j \hat{z}_j = 0$ (normalized activations have zero mean). $\blacksquare$

Interpretation of (10). The gradient is centered (subtract mean $\bar{g}$), decorrelated (subtract component along $\hat{z}_i$), and rescaled by $1/\sigma$ — mirroring the forward normalization.

Scale Invariance and Gradient Stabilization

Theorem 2 (Scale Invariance of Normalized Activations). Let $z = \mathbf{w}^T \mathbf{h}$ and $\hat{z} = (z - \mu_B)/\sigma_B$. Scaling $\mathbf{w} \to c\mathbf{w}$ leaves $\hat{z}$ unchanged for all batch elements.

Proof. Under $\mathbf{w} \to c\mathbf{w}$: $z_i \to cz_i$, $\mu_B \to c\mu_B$, $\sigma_B \to c\sigma_B$. Thus $\hat{z}_i = (cz_i - c\mu_B)/(c\sigma_B) = \hat{z}_i$. $\blacksquare$

Corollary 1 (Weight-Independent Gradients). The gradient $\partial L / \partial \mathbf{w}$ scales as $1/\|\mathbf{w}\|$, making effective learning rates approximately independent of weight magnitude.

This provides automatic learning rate adaptation — a key reason BatchNorm enables training with larger learning rates.

Inference Mode and Running Statistics

During training, batch statistics $\mu_B, \sigma_B^2$ depend on the mini-batch. At inference, batches may be size 1 or have different composition.

Definition 4 (Running Statistics). During training, maintain exponential moving averages:

with momentum $m$ (typically 0.1).

Definition 5 (Inference Forward Pass).

Proposition 3 (Inference Fusion). Equation (12) can be rewritten as a single affine transform $y = \gamma' z + \beta'$ with:

No runtime overhead beyond a single affine map per feature.

Batch Normalization as a Regularizer

Proposition 4 (Stochastic Regularization). Mini-batch statistics $\mu_B, \sigma_B^2$ are random (depend on batch composition). This injects noise into the forward pass, acting as a regularizer similar to dropout.

Proposition 5 (Batch Size Dependence). Small batch sizes yield noisier statistics, increasing regularization strength. Large batches yield more stable estimates, reducing the regularization effect.

This explains why BatchNorm performance can degrade with very large batch sizes unless learning rate is adjusted.

Layer Normalization and RMSNorm

Definition 6 (Layer Normalization). Normalize across features for each sample independently:

Definition 7 (RMSNorm). Root Mean Square Normalization (used in LLaMA, modern LLMs) omits mean centering:

• Normalizes over — Batch — Features — Features
• Batch size dependence — Yes — No — No
• Mean subtraction — Yes — Yes — No
• Used in — CNNs — Transformers — Modern LLMs

Transformers prefer LayerNorm/RMSNorm because: (1) variable-length sequences make batch statistics unreliable; (2) inference with batch size 1; (3) no cross-sample dependency.

See Self-Attention and RoPE for normalization in Transformer blocks.

Worked Examples

Example 1: Forward Pass

Batch $B = 4$, pre-activations $z = (2, 4, 6, 8)$. Then $\mu_B = 5$, $\sigma_B^2 = 5$, $\sigma = \sqrt{5}$. Normalized: $\hat{z} = (-3/\sqrt{5},\,-1/\sqrt{5},\,1/\sqrt{5},\,3/\sqrt{5})$.

Example 2: Scale Invariance

If $\mathbf{w} \to 2\mathbf{w}$, all $z_i$ double, but $\hat{z}_i$ unchanged. The network's normalized representation is invariant to weight scaling.

Example 3: Inference vs. Training

At training with batch $\{1, 3, 5, 7\}$: $\mu_B = 4$. At inference with single sample $z = 5$: use $\mu_{\text{run}}$ accumulated over training, not $\mu_B = 5$.

Connection to the Broader Curriculum

• Backpropagation — gradients through BatchNorm nodes
• Self-Attention — Pre-LN and Post-LN Transformer architectures
• Flash Attention — orthogonal optimization axis

Common Pitfalls and Misconceptions

Pitfall 1: Using batch statistics at inference. Always switch to running statistics (11) for evaluation.

Pitfall 2: Very small batch sizes. Statistics are unreliable for $B = 1$ or $B = 2$; use GroupNorm or LayerNorm instead.

Pitfall 3: BatchNorm in RNNs. Hidden state statistics vary across time steps; LayerNorm is preferred.

Pitfall 4: Assuming BatchNorm always helps. With sufficient data and careful initialization, benefits diminish. Residual connections and better optimizers reduce the need.

Pitfall 5: Confusing $\gamma, \beta$ with mean and variance. They are learnable rescaling parameters, not the batch statistics.

Research Perspective

BatchNorm (Ioffe & Szegedy, 2015) enabled training of networks 10× deeper. Subsequent work questioned the internal covariate shift explanation (Santurkar et al., 2018), showing BatchNorm primarily improves the Lipschitz constant of the loss landscape.

LayerNorm (Ba et al., 2016) adapted normalization for sequence models. RMSNorm (Zhang & Sennrich, 2019) simplified LayerNorm for LLMs. Post-Norm vs. Pre-Norm Transformer architectures remain an active research area.

Summary of Takeaways

• Forward — $\hat{z} = (z - \mu_B)/\sigma_B$ — Normalize activations
• Affine — $y = \gamma\hat{z} + \beta$ — Restore capacity
• Backward — (10) — Gradients through statistics
• Scale invariance — $\hat{z}$ invariant to $\ — \mathbf{w}\ —$ — Stable training
• Inference — Running $\mu_{\text{run}}, \sigma^2_{\text{run}}$ — Batch-independent eval
• LayerNorm — Normalize over features — Transformers

Next article: Self-Attention →

Exercises

Exercise 1. Derive (10) from (9) without skipping algebraic steps.

Exercise 2. Prove Theorem 2 (scale invariance) for vector inputs $\mathbf{z} = \mathbf{W}\mathbf{h}$.

Exercise 3. Show that $\sum_i \hat{z}_i = 0$ and $\sum_i \hat{z}_i^2 = B$ (before affine transform).

Exercise 4. Compare gradient variance with and without BatchNorm for a deep linear network.

Exercise 5. Derive the inference fusion formulas (13).

Exercise 6. Explain why LayerNorm is preferred over BatchNorm for autoregressive language modeling.

Exercise 7. For RMSNorm (15), derive the backward pass (simpler than BatchNorm — no mean term).

Exercise 8 (Conceptual). BatchNorm creates dependency between samples in a batch. Why is this problematic for certain training regimes (e.g., small batches, reinforcement learning)?