Batch Normalization — Deep Dive

Gradient Flow Through BatchNorm

BatchNorm is fully differentiable, so backpropagation passes through it normally. The gradients require care because the normalization couples the gradients of samples within a batch.

For a mini-batch of size $m$, let $y = \text{BN}(x)$:

$$\frac{\partial \mathcal{L}}{\partial \hat{x}_i} = \frac{\partial \mathcal{L}}{\partial y_i} \cdot \gamma$$

$$\frac{\partial \mathcal{L}}{\partial \sigma^2_\mathcal{B}} = \sum_{i=1}^m \frac{\partial \mathcal{L}}{\partial \hat{x}i} \cdot (x_i - \mu\mathcal{B}) \cdot \frac{-1}{2}(\sigma^2_\mathcal{B} + \epsilon)^{-3/2}$$

$$\frac{\partial \mathcal{L}}{\partial \mu_\mathcal{B}} = \sum_{i=1}^m \frac{\partial \mathcal{L}}{\partial \hat{x}i} \cdot \frac{-1}{\sqrt{\sigma^2\mathcal{B} + \epsilon}} + \frac{\partial \mathcal{L}}{\partial \sigma^2_\mathcal{B}} \cdot \frac{\sum_{i=1}^m -2(x_i - \mu_\mathcal{B})}{m}$$

$$\frac{\partial \mathcal{L}}{\partial x_i} = \frac{\partial \mathcal{L}}{\partial \hat{x}i} \cdot \frac{1}{\sqrt{\sigma^2\mathcal{B} + \epsilon}} + \frac{\partial \mathcal{L}}{\partial \sigma^2_\mathcal{B}} \cdot \frac{2(x_i - \mu_\mathcal{B})}{m} + \frac{\partial \mathcal{L}}{\partial \mu_\mathcal{B}} \cdot \frac{1}{m}$$

The gradient for each sample $x_i$ depends on all other samples in the batch through the $\mu$ and $\sigma^2$ terms. This is the coupling that creates the implicit regularization effect — and also why very small batches cause problems (high variance in statistics = high variance in gradients).

The Smooth Loss Landscape Hypothesis

Santurkar et al. (2018) “How Does Batch Normalization Help Optimization?” challenged the internal covariate shift explanation with two experiments:

1. Added random noise to BatchNorm outputs (reintroducing internal covariate shift) — networks still trained well
2. Measured the Lipschitz constant of the loss landscape with and without BatchNorm

Their finding: BatchNorm reduces the gradient’s Lipschitz constant — i.e., the gradient changes more slowly as you move in parameter space. Formally:

$$|\nabla_\mathbf{W} \mathcal{L}t| \leq |\nabla\mathbf{W} \mathcal{L}_{t-1}| \cdot (1 + \text{smoothing factor})$$

This allows larger gradient steps (higher learning rates) without overshooting the loss landscape.

The implication: BatchNorm’s benefits are primarily about optimization geometry, not about keeping activations in a specific range per se. This explains why LayerNorm (normalizing over different dimensions) achieves similar training benefits.

RMSNorm: A Simpler Variant

Zhang & Sennrich (2019) proposed Root Mean Square Layer Normalization (RMSNorm), which drops the mean-centering step:

$$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \cdot \gamma, \quad \text{RMS}(x) = \sqrt{\frac{1}{n}\sum_{i=1}^n x_i^2}$$

No bias parameter $\beta$, and no subtraction of the mean. RMSNorm is ~15% faster than LayerNorm (fewer operations) with minimal quality difference.

LLaMA (Meta, 2023) used RMSNorm throughout instead of LayerNorm, and it’s now standard in many efficient transformer implementations. The intuition: the re-centering in LayerNorm may be less important than the re-scaling; RMSNorm only does re-scaling.

Ghost BatchNorm and Large-Scale Training

At very large batch sizes (common in distributed training), BatchNorm can paradoxically hurt performance. The reason: with a batch of 4096 images, the batch statistics become near-perfect estimates of the true statistics — the noise that acts as a regularizer disappears. Models overfit more.

Ghost BatchNorm (Hoffer et al., 2017): During training with large batches, split the batch into “ghost batches” of size $k$ and compute normalization statistics per ghost batch. This preserves the noise property of small-batch normalization even when using large actual batches.

In practice, Sync BatchNorm (synchronizing batch statistics across all GPUs in distributed training) is the more common solution. It computes correct batch statistics globally but at the cost of inter-GPU communication overhead.

BatchNorm in ResNets: Pre- vs. Post-Activation

The original ResNet paper placed BatchNorm after conv, before ReLU (post-activation), with the skip connection bypassing both:

x → Conv → BN → ReLU → Conv → BN → + → ReLU
                                    ↑
                           x (identity shortcut)

He et al. (2016b) “Identity Mappings in Deep Residual Networks” proposed pre-activation:

x → BN → ReLU → Conv → BN → ReLU → Conv → + 
                                           ↑
                                  x (identity shortcut)

Pre-activation ensures the shortcut path (x →) is a clean identity — no BN or activation in the way. For networks with 100+ layers, pre-activation shows consistent improvement because gradients flow through the shortcut without being modified by BN or saturated by activations.

The tradeoff: pre-activation changes the meaning of the output layer (it’s now BN+ReLU+Conv instead of Conv+BN+ReLU), requiring careful handling of the first and last layers.

Normalization Choice by Architecture

Architecture	Normalization	Reason
ResNet, EfficientNet, VGG	BatchNorm	Large batches, image-level stats work well
GPT, LLaMA, BERT	LayerNorm/RMSNorm	Variable sequence lengths, single-sample inference
StyleGAN	AdaIN	Style-conditioned normalization per image
Object Detection (DETR, small batch)	GroupNorm	Small batches (2-4 images), stable statistics
Diffusion models (U-Net backbone)	GroupNorm	Small batch, needs per-sample stats
MobileNet/EfficientNet-lite	BatchNorm with calibration	Quantization-aware training needs stable running stats

Interaction With Dropout

BatchNorm and Dropout are somewhat redundant — both regularize training. Using both can hurt performance because:

1. Dropout adds multiplicative noise to activations; BatchNorm then normalizes out this noise
2. The interaction creates variance shift: the variance of normalized activations seen during training (with dropout noise) differs from inference (no dropout noise), causing a distribution mismatch in BatchNorm running statistics

For this reason, most modern architectures use BatchNorm or Dropout, not both. ResNets use BatchNorm without Dropout. Earlier architectures (AlexNet, VGG) used Dropout without BatchNorm.

When both are used (e.g., in EfficientNet), Dropout is placed after BatchNorm and only in the final layers, minimizing the interaction.

Practical Implementation Notes

import torch.nn as nn

# Standard usage in a CNN block
class ConvBNReLU(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride, padding=kernel_size//2, bias=False)
        # bias=False because BN has its own learnable shift (beta)
        self.bn = nn.BatchNorm2d(out_ch, momentum=0.1, eps=1e-5)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

Key implementation detail: set bias=False in the Conv layer when followed by BatchNorm. The BatchNorm’s $\beta$ parameter serves the same role as bias — adding it would be redundant.

At inference, call model.eval() to switch BatchNorm to using running statistics instead of batch statistics. Forgetting this is a common bug that causes different results in train vs. eval mode.

One thing to remember: BatchNorm’s mechanism is well-understood mathematically, but why it works so well remains partly empirical — the smooth loss landscape explanation is compelling but not complete, and choosing between BatchNorm, LayerNorm, and RMSNorm often comes down to your architecture and batch size constraints rather than first principles.