Batch Normalization

The Real Interview Moment

Your model achieves 94% accuracy in training but 61% in production. The gap is too large to be normal overfitting. You dig into the inference code and find model.train() is never called - but neither is model.eval(). In PyTorch, model.eval() is not the default. If you do not call it explicitly, the model stays in training mode.

In training mode, batch normalization uses the current batch's statistics - mean and variance computed from the examples in that batch. In production, you are processing one request at a time. With a batch size of 1, the "batch mean" is just the single example's value, and the "batch variance" is 0. BN computes $(x - x) / \sqrt{0 + \epsilon} \approx 0$ for every pre-activation. The layer is outputting zero regardless of its input. Your "model" has been lobotomized by a missing model.eval() call.

The 33-point accuracy gap between training and production is entirely a batch normalization mode error. This is not a corner case - it is one of the most common production bugs in deep learning, and it comes from not understanding what batch normalization actually does during inference.

This lesson gives you a complete understanding of every normalization technique used in modern deep learning: the math, the train/eval distinction, the correct mental model for why BN works (it is not what the original paper claimed), and when to use each variant.

Why This Exists: The Internal Covariate Shift Hypothesis

Covariate shift occurs when the distribution of inputs to a model changes between training and test time. Internal covariate shift (Ioffe and Szegedy, 2015) is the hypothesized phenomenon where this happens inside the network: as weights in layer $l$ are updated, the distribution of inputs to layer $l+1$ changes - even when the external data distribution is fixed.

The hypothesis: layer 3 must continuously re-adapt to upstream changes in layer 2's output distribution. This makes training slow because each layer is "chasing a moving target."

The solution proposed: normalize each layer's inputs before applying them to the next layer. If the distribution of activations is normalized to fixed statistics (mean 0, variance 1), upstream weight changes do not shift the distribution seen by downstream layers.

:::note The Real Mechanism Subsequent research (Santurkar et al., 2018, "How Does Batch Normalization Help Optimization?") showed that internal covariate shift is NOT the primary mechanism. The paper demonstrated that BN does not actually reduce internal covariate shift significantly compared to controls - and that networks with BN still exhibit significant internal covariate shift. The real benefit appears to be smoothing the loss landscape: BN makes the loss function more Lipschitz-smooth, allowing larger learning rates and more stable gradient descent. The original hypothesis was wrong. The technique works extremely well regardless. :::

BN Algorithm: Normalize Then Scale and Shift

Given a mini-batch $\mathcal{B} = \{x_1, \ldots, x_B\}$ for a single feature:

Step 1: Batch mean

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

Step 2: Batch variance (biased estimate, divides by $B$ not $B-1$)

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

Step 3: Normalize

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

Where $\epsilon \approx 10^{-5}$ prevents division by zero when $\sigma_\mathcal{B}^2 = 0$.

Step 4: Scale and shift with learned affine transform

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

Where $\gamma$ (scale) and $\beta$ (shift) are learned parameters updated by gradient descent. This is critical: without $\gamma$ and $\beta$, BN would force every layer's activations to have exactly mean 0 and variance 1, potentially reducing the layer's representational capacity. The learned affine transform allows the network to restore any distribution it finds optimal - including the original un-normalized one if that is what the training process discovers.

For a layer with $d$ features, $\gamma$ and $\beta$ each have $d$ parameters. These are separate from the weight matrix - BN adds $2d$ learned parameters per layer.

Train Mode vs Eval Mode: The Critical Distinction

During training: BN uses the current mini-batch's statistics $(\mu_\mathcal{B}, \sigma_\mathcal{B}^2)$. These are computed fresh from the batch. Simultaneously, BN maintains running statistics via exponential moving average:

$\mu_{\text{running}} \leftarrow (1 - \alpha) \cdot \mu_{\text{running}} + \alpha \cdot \mu_\mathcal{B}$ $\sigma^2_{\text{running}} \leftarrow (1 - \alpha) \cdot \sigma^2_{\text{running}} + \alpha \cdot \sigma^2_\mathcal{B}$

Where $\alpha$ is the momentum parameter in PyTorch's BN (default 0.1 - confusingly named; this is the weight given to the new batch, not the optimizer momentum concept).

During evaluation (after model.eval()): BN uses the fixed running statistics accumulated during training. The running statistics represent population-level statistics of the training data - stable, deterministic, and independent of the current batch size.

import torch
import torch.nn as nn


def demonstrate_bn_train_vs_eval():
    """
    Demonstrate the critical difference between train and eval mode.
    This reproduces the exact production bug from the opening scenario.
    """
    torch.manual_seed(42)
    bn = nn.BatchNorm1d(4)
    bn.train()

    # Build up running statistics over 100 training batches
    # Data with mean=5 (not zero) - simulates non-centered features
    for _ in range(100):
        x = torch.randn(32, 4) + 5.0
        _ = bn(x)

    print("After training - collected running statistics:")
    print(f"  running_mean: {bn.running_mean.tolist()}")
    print(f"  running_var:  {bn.running_var.tolist()}")

    # Production inference: single example
    x_single = torch.randn(1, 4) + 5.0

    # WRONG: training mode with batch size 1
    # batch_mean = x_single (the single value IS the batch mean)
    # batch_var  = 0 (variance of a single value is 0)
    # output = gamma * (x - x) / sqrt(0 + eps) + beta = gamma * 0 + beta
    bn.train()
    out_wrong = bn(x_single)
    print(f"\nTRAIN mode output (batch_size=1): {out_wrong}")
    print(f"  → This is just beta: {bn.bias}")
    print(f"  → Model ignores its input!")

    # CORRECT: eval mode uses running statistics
    bn.eval()
    out_correct = bn(x_single)
    print(f"\nEVAL mode output (batch_size=1): {out_correct}")
    print(f"  → Uses running stats - correct behavior")

    # Measure the difference
    diff = (out_wrong - out_correct).abs().mean().item()
    print(f"\nAbsolute difference between modes: {diff:.4f}")
    print(f"  → In a deep network with many BN layers, this error compounds catastrophically")


demonstrate_bn_train_vs_eval()

Manual BN Implementation: Understanding Every Step

import torch
import torch.nn as nn
from torch import Tensor


class ManualBatchNorm1d(nn.Module):
    """
    Manual implementation of BatchNorm1d for pedagogical understanding.
    Matches PyTorch's implementation exactly.

    Key distinction between training and eval modes:
    - Training: normalize using BATCH statistics (current mini-batch)
    - Eval: normalize using RUNNING statistics (accumulated from training)
    """

    def __init__(
        self,
        num_features: int,
        eps: float = 1e-5,
        momentum: float = 0.1,
        affine: bool = True,
    ):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum   # weight given to NEW batch in running average
        self.affine = affine

        if affine:
            # gamma (scale) initialized to 1: no initial scaling
            # beta (shift) initialized to 0: no initial shift
            self.weight = nn.Parameter(torch.ones(num_features))   # gamma
            self.bias = nn.Parameter(torch.zeros(num_features))    # beta

        # Running statistics: not parameters - not updated by gradient descent
        # Registered as buffers so they are saved/loaded with model state dict
        self.register_buffer("running_mean", torch.zeros(num_features))
        self.register_buffer("running_var", torch.ones(num_features))
        self.register_buffer("num_batches_tracked", torch.tensor(0, dtype=torch.long))

    def forward(self, x: Tensor) -> Tensor:
        assert x.dim() == 2, f"Expected 2D input (batch, features), got {x.dim()}D"

        if self.training:
            # TRAINING MODE: compute statistics from current batch
            batch_mean = x.mean(dim=0)                         # (num_features,)
            batch_var = x.var(dim=0, unbiased=False)           # biased variance (÷ B, not B-1)

            # Update running statistics via exponential moving average
            # detach() prevents gradients flowing into the running stats update
            self.running_mean = (
                (1 - self.momentum) * self.running_mean + self.momentum * batch_mean.detach()
            )
            self.running_var = (
                (1 - self.momentum) * self.running_var + self.momentum * batch_var.detach()
            )
            self.num_batches_tracked += 1

            # Normalize using BATCH statistics
            x_hat = (x - batch_mean) / (batch_var + self.eps).sqrt()

        else:
            # EVAL MODE: normalize using accumulated RUNNING statistics
            # These are stable population-level estimates, not dependent on current batch
            x_hat = (x - self.running_mean) / (self.running_var + self.eps).sqrt()

        # Apply learned affine transform
        if self.affine:
            x_hat = self.weight * x_hat + self.bias

        return x_hat


def verify_manual_bn():
    """Verify manual BN matches PyTorch's implementation."""
    torch.manual_seed(42)
    x = torch.randn(32, 16)

    official_bn = nn.BatchNorm1d(16)
    manual_bn = ManualBatchNorm1d(16)

    # Ensure same parameters
    with torch.no_grad():
        manual_bn.weight.copy_(official_bn.weight)
        manual_bn.bias.copy_(official_bn.bias)

    official_bn.train()
    manual_bn.train()

    out_official = official_bn(x)
    out_manual = manual_bn(x)

    max_diff = (out_official - out_manual).abs().max().item()
    print(f"Max difference from PyTorch BN (train mode): {max_diff:.2e}")  # Should be < 1e-6


verify_manual_bn()

Why BN Actually Works: Loss Landscape Smoothing

Santurkar et al. (2018) showed empirically and theoretically that BN's primary benefit is smoothing the optimization landscape, not reducing internal covariate shift:

The Lipschitz smoothness argument: the loss is a function of all parameters. BN makes the loss function more Lipschitz-smooth - the gradient changes more predictably as parameters change. Formally, the $\beta$-smoothness constant of the loss (the Lipschitz constant of the gradient) is reduced by BN. Smaller smoothness constant means larger learning rates are safe - the gradient does not change too much between steps, so a large step remains valid.

Practical consequences:

1. You can use 10x–100x larger learning rates with BN than without - dramatically faster training
2. Loss curves are smoother with less variance across training runs
3. Less sensitivity to learning rate choice - BN provides a natural "guard rail" against divergence
4. Acts as an implicit regularizer: each example is normalized based on its relationship to the other examples in the batch, adding noise to individual activations

The regularization effect: because BN uses batch statistics (which depend on which other examples are in the batch), each example's normalized activations are slightly different across epochs. This is similar to noise injection - it prevents the network from memorizing exact activation values for specific training examples. This is why models with BN often need less dropout.

Layer Normalization: For Transformers and RNNs

Layer Normalization (Ba et al., 2016) normalizes across the feature dimension for each example independently, instead of across the batch dimension for each feature.

$\text{LN}(x) = \frac{x - \mathbb{E}[x]}{\sqrt{\text{Var}[x] + \epsilon}} \cdot \gamma + \beta$

Where the expectation and variance are computed over the feature dimension $d_{\text{model}}$ for each individual example, not across the batch:

$\mu = \frac{1}{d} \sum_{i=1}^{d} x_i \quad \text{(per example)}$ $\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2 \quad \text{(per example)}$

Why LayerNorm for transformers:

1. Variable batch size: auto-regressive generation processes one token at a time. With batch size 1, BN's batch statistics degenerate. LayerNorm computes statistics per-example and is completely independent of batch size.
2. Variable sequence length: BN would need to aggregate statistics across tokens with different semantic roles - positionally early tokens vs late tokens. LayerNorm normalizes each token's features independently.
3. No train/eval distinction: LayerNorm uses the same computation during training and inference - no running statistics, no train/eval mode difference. This eliminates an entire class of production bugs.
4. Empirically superior on language tasks: BERT, GPT-2, T5, LLaMA all use LayerNorm and achieve strong results.

import torch
import torch.nn as nn


def compare_bn_vs_ln():
    """
    Demonstrate the key behavioral differences between BN and LN.
    """
    batch_size, seq_len, d_model = 8, 64, 512

    # Simulate transformer hidden states: (batch, seq_len, d_model)
    x = torch.randn(batch_size, seq_len, d_model) + 2.0  # non-zero mean

    # LayerNorm: normalizes over d_model for each (batch, position) independently
    ln = nn.LayerNorm(d_model)
    out_ln = ln(x)

    print("LayerNorm output statistics:")
    print(f"  Mean over d_model (per token):  {out_ln.mean(dim=-1).abs().mean():.6f}  (should be ~0)")
    print(f"  Std over d_model (per token):   {out_ln.std(dim=-1).mean():.6f}  (should be ~1)")
    print(f"  No train/eval mode difference - same behavior always")

    # For 1D features (MLP hidden states): (batch, features)
    x_1d = torch.randn(batch_size, d_model) + 2.0
    bn = nn.BatchNorm1d(d_model)
    bn.train()
    out_bn_train = bn(x_1d)

    print("\nBatchNorm output statistics (training mode):")
    print(f"  Mean over batch (per feature):  {out_bn_train.mean(dim=0).abs().mean():.6f}  (should be ~0)")
    print(f"  Std over batch (per feature):   {out_bn_train.std(dim=0).mean():.6f}  (should be ~1)")

    # With batch_size=1: BN degenerates
    x_single = torch.randn(1, d_model) + 2.0
    out_bn_single = bn(x_single)  # still in train mode
    print("\nBatchNorm with batch_size=1 (train mode):")
    print(f"  Output mean: {out_bn_single.mean():.6f}  (should vary but is gamma*0+beta)")
    print(f"  → BN computes (x-x)/sqrt(0+eps) = 0 for every feature")

RMS Norm: Simplified Layer Norm (Used in LLaMA)

RMS Norm (Zhang and Sennrich, 2019) is a simplified variant of LayerNorm that drops the mean-centering step:

$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \cdot \gamma$

$\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} x_i^2}$

The key simplification: no subtraction of mean, no bias parameter $\beta$. The theory is that the re-centering in LayerNorm is redundant - the scaling (RMS normalization) alone is sufficient to stabilize training. This reduces computation by approximately 7–15% on GPU.

Where used: LLaMA, LLaMA-2, LLaMA-3, GPT-NeoX, Mistral. The empirical results are comparable to LayerNorm with slightly lower compute cost.

import torch
import torch.nn as nn


class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization.
    Used in LLaMA, GPT-NeoX, and Mistral.

    Simpler than LayerNorm: no mean subtraction, no bias.
    Only normalizes by RMS and applies a learned scale.
    """

    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # gamma (no beta)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Compute RMS over the last dimension
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
        return x / rms * self.weight


# Verify against manual computation
dim = 512
rms_norm = RMSNorm(dim)
x = torch.randn(32, 64, dim)  # (batch, seq_len, d_model)

out = rms_norm(x)
# Check: |output| / |gamma| should be approximately 1
rms_out = out.pow(2).mean(dim=-1).sqrt() / rms_norm.weight.pow(2).mean().sqrt()
print(f"Normalized RMS (should be ~1): {rms_out.mean():.4f}")

Group Normalization: For Small-Batch Vision Tasks

Group Normalization (Wu and He, 2018) normalizes over groups of channels for each example, bridging BN (batch groups) and LN (single group of all channels).

$\text{GN}(x) = \frac{x - \mu_{g}}{\sqrt{\sigma_g^2 + \epsilon}} \cdot \gamma + \beta$

Where $\mu_g$ and $\sigma_g^2$ are computed over channels $C/G$ and spatial dimensions $H, W$ within group $g$, for each example independently.

When to use GN instead of BN: object detection and segmentation models (Mask R-CNN, Detectron) use image pairs or small batches for fine-grained tasks. With BN and batch size 2, the batch statistics are too noisy to be meaningful. GN provides stable normalization regardless of batch size because it normalizes within each example.

import torch
import torch.nn as nn

# Example: object detection backbone with 8 channels
# BN needs large batches; GN works with any batch size
x = torch.randn(2, 64, 32, 32)  # (N=2, C=64, H=32, W=32)

gn = nn.GroupNorm(num_groups=8, num_channels=64)  # 8 groups of 8 channels each
out_gn = gn(x)

# Statistics per group per example (not per batch)
print(f"GN output shape: {out_gn.shape}")  # (2, 64, 32, 32)
print(f"GN output mean (should be ~0 per group): {out_gn[:, :8, :, :].mean():.4f}")
print(f"GN output std (should be ~1 per group):  {out_gn[:, :8, :, :].std():.4f}")

Instance Normalization: For Style Transfer

Instance Normalization (Ulyanov et al., 2016) normalizes over the spatial dimensions $H \times W$ for each channel and each example independently.

$\text{IN}(x_{n,c}) = \frac{x_{n,c} - \mu_{n,c}}{\sqrt{\sigma_{n,c}^2 + \epsilon}}$

Where $\mu_{n,c}$ and $\sigma_{n,c}^2$ are computed over the spatial dimension for example $n$, channel $c$.

Why it works for style transfer: the "style" of an image is captured by the mean and variance of feature maps across spatial positions. By normalizing these per-instance per-channel, IN removes style information from content features, making it easy to re-apply a different style. This is the foundation of AdaIN (Adaptive Instance Normalization) used in neural style transfer.

Complete Normalization Comparison

Method	Normalize over	Batch size dep.	Seq length dep.	Use case
Batch Norm	Batch + spatial	Yes (strongly)	No	CNNs, large-batch training
Layer Norm	Feature dim	No	No	Transformers, NLP, any seq length
Group Norm	Channel groups + spatial	No	No	Small-batch vision (detection)
Instance Norm	Spatial (per sample, per channel)	No	No	Style transfer, generative models
RMS Norm	Feature dim (no mean center)	No	No	LLaMA, GPT-NeoX, efficient LLMs

Pre-Norm vs Post-Norm in Transformers

The original Transformer (Vaswani et al., 2017) used Post-LN: LayerNorm applied after each sub-layer's output plus residual connection:

$x_{l+1} = \text{LN}(x_l + F(x_l))$

GPT-2 and subsequent large language models switched to Pre-LN: LayerNorm applied to the input before each sub-layer:

$x_{l+1} = x_l + F(\text{LN}(x_l))$

Why Pre-LN trains more stably:

In Post-LN, the residual path $x_l + F(x_l)$ has un-normalized variance. As depth increases, the accumulated variance can drift significantly before LayerNorm corrects it. This causes training instability and requires very careful learning rate warmup for deep transformers.

In Pre-LN, the input to each sub-layer is always normalized. The residual $x_l$ is added after $F(\cdot)$, so the un-normalized signal flows directly through the residual connection. This makes gradient flow more stable - gradients can flow directly through residual connections without being distorted by LayerNorm. Pre-LN transformers generally train stably without warmup or with minimal warmup, even at large depth.

The tradeoff: Pre-LN models sometimes have slightly lower final performance than Post-LN models when both train to completion. The stability advantage of Pre-LN makes it preferred for large-scale training where instability has catastrophic cost (wasted GPU-months).

import torch
import torch.nn as nn


class PostLNBlock(nn.Module):
    """Original Transformer (Vaswani et al., 2017)."""
    def __init__(self, d_model: int, nhead: int, dim_ff: int):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.ff = nn.Sequential(
            nn.Linear(d_model, dim_ff),
            nn.GELU(),
            nn.Linear(dim_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Post-LN: normalize AFTER residual addition
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)       # LN after residual
        x = self.norm2(x + self.ff(x))     # LN after residual
        return x


class PreLNBlock(nn.Module):
    """GPT-2 / LLaMA style - more stable for deep models."""
    def __init__(self, d_model: int, nhead: int, dim_ff: int):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.ff = nn.Sequential(
            nn.Linear(d_model, dim_ff),
            nn.GELU(),
            nn.Linear(dim_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-LN: normalize BEFORE sub-layer, then add residual
        attn_out, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
        x = x + attn_out                   # residual without LN distortion
        x = x + self.ff(self.norm2(x))     # residual without LN distortion
        return x

BN for Fine-Tuning: Freezing Running Statistics

When fine-tuning a pretrained model on a small dataset, the mini-batch statistics from the small fine-tuning data can corrupt the running statistics that were carefully accumulated during large-scale pretraining. This is a subtle bug: the pretrained BN statistics were calibrated on the full pretraining distribution; the fine-tuning batches have different statistics.

The fix: freeze BN layers during fine-tuning by keeping them in eval mode even while the rest of the model trains.

import torch
import torch.nn as nn


def freeze_bn_running_stats(model: nn.Module) -> None:
    """
    Freeze BatchNorm running statistics for fine-tuning.
    BN layers stay in eval mode (using pretrained running stats) even when
    the model as a whole is in training mode.

    The learned gamma and beta parameters are still updated by gradient descent.
    Only the running mean and variance are frozen.

    When to use: fine-tuning on a dataset much smaller than the pretraining data,
    or when the fine-tuning domain differs significantly from pretraining.
    """
    for module in model.modules():
        if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
            module.eval()   # eval mode freezes running stats for this module
            # Note: module.eval() works even when parent model is in train mode


def train_with_frozen_bn(model: nn.Module, loader, optimizer, criterion):
    """Training loop that keeps BN in eval mode throughout."""
    model.train()               # model in train mode (dropout active, etc.)
    freeze_bn_running_stats(model)  # but BN layers stay in eval mode

    for x, y in loader:
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        # Important: call freeze_bn_running_stats after each optimizer step
        # because optimizer.step() doesn't change model mode, but if
        # model.train() is called elsewhere in the loop, BN mode would reset
        freeze_bn_running_stats(model)


# Alternative: convert BN to LN for fine-tuning (more drastic)
def replace_bn_with_ln(model: nn.Module) -> nn.Module:
    """
    Replace all BatchNorm1d layers with LayerNorm.
    More aggressive than freezing - removes BN's batch-size dependence entirely.
    Useful when fine-tuning with very small batches or online learning.
    """
    for name, module in model.named_children():
        if isinstance(module, nn.BatchNorm1d):
            num_features = module.num_features
            ln = nn.LayerNorm(num_features)
            # Copy learned parameters if they exist
            if module.weight is not None:
                with torch.no_grad():
                    ln.weight.copy_(module.weight)
                    ln.bias.copy_(module.bias)
            setattr(model, name, ln)
        else:
            replace_bn_with_ln(module)  # recurse into submodules
    return model

:::danger model.eval() is Not the Default PyTorch models start in training mode by default. model.eval() must be explicitly called before any validation loop or inference pass. The consequences of forgetting: (1) Dropout randomly zeros activations - your "validation" is actually measuring a noisy ensemble of different model states; (2) BatchNorm uses batch statistics - with small validation batches, these are noisy and produce different outputs than the running statistics; (3) With batch size 1 in production, BN completely ignores the input and outputs only its learned bias. None of these failures produce error messages. The model silently outputs wrong values. :::

:::warning BN Momentum is Confusingly Named PyTorch's nn.BatchNorm1d(momentum=0.1) uses momentum to mean the weight given to the new batch in the exponential moving average: running_mean = (1-momentum)*running_mean + momentum*batch_mean. This is the opposite of optimizer momentum convention, where momentum = 0.9 means 90% of the previous value is kept. BN momentum = 0.1 means 10% of the new batch is incorporated. The default 0.1 works well for most cases - higher values make running stats track the data more quickly but with more noise; lower values smooth more aggressively. :::

YouTube Resources

Video	Channel	Why Watch It
Batch Norm Explained	Andrej Karpathy	Practical implementation with debugging the train/eval distinction
How Does Batch Norm Help Optimization?	Yannic Kilcher	Paper walkthrough of Santurkar et al. 2018 debunking ICS
CS231n - Batch Normalization	Stanford CS231n	Derivation of forward and backward pass, practical guidance
Layer Norm vs Batch Norm	Machine Learning Mastery	Visual comparison across different architectures
RMSNorm - LLaMA Architecture	Andrej Karpathy	Pre-norm, RMSNorm, and transformer architecture choices

Interview Q&A

Q1: Explain batch normalization step by step and why each step is necessary.

BN has four steps, each with a specific purpose. Step 1 - compute batch mean and variance: creates a reference distribution for the current batch. This reference is used to normalize all examples identically, regardless of their absolute scale. Step 2 - normalize: subtracts mean and divides by standard deviation, making activations approximately zero-mean and unit-variance. This prevents extreme activation values from saturating downstream neurons. Step 3 - apply learned scale and shift: multiplies by $\gamma$ and adds $\beta$. This restores expressivity - without it, every layer would be forced to use unit-variance activations, potentially limiting what functions the network can express. The network can learn to undo the normalization if that is optimal. Step 4 (inference only) - use running statistics: ensures deterministic, batch-size-independent inference. Running statistics represent population-level statistics accumulated during training, not the potentially noisy statistics of a small inference batch.

Q2: What is the difference between training and eval mode in batch normalization?

In training mode, BN normalizes each example using the current mini-batch's mean and variance, and simultaneously updates the exponential moving average of running statistics. The batch statistics depend on which other examples are in the batch - this introduces stochasticity that acts as regularization. In eval mode, BN uses the fixed running statistics accumulated during training. This ensures deterministic, batch-size-independent behavior. The failure mode for forgetting model.eval(): with a single-example inference batch, batch mean equals that example's value and batch variance equals 0. Every activation normalizes to $(x - x)/\sqrt{0 + \epsilon} = 0$. The layer outputs only its learned bias parameter $\beta$ regardless of input - the model effectively ignores all its input features. This exact bug has caused major production incidents at multiple organizations.

Q3: Why is LayerNorm used in transformers instead of BatchNorm?

Three reasons: (1) Auto-regressive generation produces one token at a time, meaning batch size 1 during inference. BN with batch size 1 degenerates as described above - LayerNorm computes statistics per-example over the feature dimension and is independent of batch size. (2) Sequence length varies across examples and tasks. BN statistics computed across positions in a sequence mix semantically different roles (first token vs last token) - LayerNorm normalizes each position independently. (3) LayerNorm has no train/eval distinction - it uses the same computation always. This eliminates the entire class of train/eval mode bugs. Empirically, BERT, GPT-2, T5, and LLaMA all use LayerNorm (or RMSNorm) and achieve strong results - BN has not shown advantages in any published large-scale language model.

Q4: Explain the debate about why BN works. What does the 2018 paper show?

The original BN paper (Ioffe and Szegedy, 2015) claimed BN works by reducing internal covariate shift - the distribution shift of each layer's inputs caused by upstream weight updates. This made intuitive sense but was not rigorously verified. Santurkar et al. (2018) challenged this with controlled experiments: they added adversarial noise to BN's output to deliberately increase covariate shift, and found that models with BN still trained well despite the increased internal covariate shift. They also showed that models without BN exhibit similar levels of covariate shift as models with BN. Their positive finding: BN significantly smooths the optimization landscape - the loss is more Lipschitz-smooth, gradients are more predictable, and larger learning rates are stable. The mechanism is different from what was claimed, but the effect is real. The practical implication: use BN for its training stabilization and learning rate flexibility, not because you believe it is controlling distribution shift.

Q5: What happens when BN is applied with batch size 1 during training?

With batch size 1, the batch mean equals the single example's value for every feature. The batch variance is 0 (variance of a single point is undefined - PyTorch computes 0 with unbiased=False). The normalization computes $(x - x) / \sqrt{0 + \epsilon} = 0$ for every feature. The output is $\gamma \cdot 0 + \beta = \beta$ - a learned constant that is completely independent of the input. The model ignores its input entirely. In training mode, the running statistics also get corrupted by these degenerate batch statistics. BN is designed for mini-batch training with batch size ≥ 16 (ideally ≥ 32). For applications requiring batch size 1 - online learning, real-time inference, recurrent models - use LayerNorm, InstanceNorm, or GroupNorm. All of these normalize within each example independently.

Q6: Compare pre-norm and post-norm in transformers. Which is better and why?

Post-norm (original Transformer): $x_{l+1} = \text{LN}(x_l + F(x_l))$. The un-normalized residual accumulates variance across layers, requiring careful learning rate warmup for deep models. Post-norm can achieve slightly lower loss when trained to full convergence. Pre-norm (GPT-2, LLaMA): $x_{l+1} = x_l + F(\text{LN}(x_l))$. LayerNorm is applied before each sub-layer, so the input to each computation is always normalized. Gradients flow directly through the residual connections without LayerNorm distortion. Pre-norm trains stably without warmup even at large depth, making it strongly preferred for large-scale training where instability has catastrophic cost. The practical consensus: use Pre-LN for any transformer with more than 12 layers, any model where training instability would be costly, and any architecture where you want to avoid debugging convergence issues. Post-LN can be fine for small models where you can afford to tune warmup carefully, but the stability advantage of Pre-LN almost always outweighs its minor performance disadvantage.

Batch Normalization: Backward Pass Derivation

Understanding the backward pass is essential for implementing custom normalization or debugging gradient issues. Given the forward pass:

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

With loss $L$, the gradients are:

Gradient with respect to $\gamma$ and $\beta$ (easy - standard chain rule):

$\frac{\partial L}{\partial \gamma} = \sum_{i=1}^B \frac{\partial L}{\partial y_i} \hat{x}_i \qquad \frac{\partial L}{\partial \beta} = \sum_{i=1}^B \frac{\partial L}{\partial y_i}$

Gradient with respect to $x_i$ (harder - $\mu_B$ and $\sigma_B^2$ both depend on all $x_i$):

$\frac{\partial L}{\partial x_i} = \frac{\gamma}{\sqrt{\sigma_B^2 + \epsilon}} \left(\frac{\partial L}{\partial \hat{x}_i} - \frac{1}{B}\sum_j \frac{\partial L}{\partial \hat{x}_j} - \frac{\hat{x}_i}{B}\sum_j \frac{\partial L}{\partial \hat{x}_j} \hat{x}_j\right)$

The three terms correspond to: (1) the direct gradient through the normalized value, (2) the indirect gradient through the batch mean, and (3) the indirect gradient through the batch variance. All three must be accounted for - omitting the mean and variance gradient terms gives incorrect gradients.

import numpy as np


def batchnorm_backward(dout: np.ndarray, cache: tuple) -> tuple:
    """
    Backward pass through batch normalization.

    Args:
        dout: Gradient from upstream, shape (B, D)
        cache: Tuple from forward pass: (x_hat, gamma, beta, mu, var, eps, x)

    Returns:
        dx: Gradient wrt input x
        dgamma: Gradient wrt gamma (scale)
        dbeta: Gradient wrt beta (shift)
    """
    x_hat, gamma, beta, mu, var, eps, x = cache
    B, D = x.shape

    # Gradients for affine parameters - simple chain rule
    dbeta = dout.sum(axis=0)                    # (D,)
    dgamma = (dout * x_hat).sum(axis=0)         # (D,)

    # Gradient wrt x_hat (before affine transform)
    dx_hat = dout * gamma                        # (B, D)

    # Standard deviation (denominator in normalization)
    std_inv = 1.0 / np.sqrt(var + eps)

    # Gradient wrt variance
    dvar = (dx_hat * (x - mu) * -0.5 * std_inv**3).sum(axis=0)  # (D,)

    # Gradient wrt mean (two paths: direct and through variance)
    dmu = (-dx_hat * std_inv).sum(axis=0) + dvar * (-2.0 / B) * (x - mu).sum(axis=0)

    # Gradient wrt x (three paths: direct, through mean, through variance)
    dx = (dx_hat * std_inv
          + dvar * 2.0 / B * (x - mu)
          + dmu / B)

    return dx, dgamma, dbeta

This three-path gradient is why BN is expensive to differentiate and why frameworks implement it in optimized CUDA kernels rather than composing simpler operations.

Sync BatchNorm: For Distributed Training

In distributed training across multiple GPUs, each GPU processes a different mini-batch. Standard nn.BatchNorm2d computes statistics from its local batch (e.g., 8 examples per GPU). With 8 GPUs this means batch statistics from only 8 examples - too noisy for reliable normalization.

torch.nn.SyncBatchNorm synchronizes statistics across all GPUs before normalization:

import torch
import torch.nn as nn


def convert_to_sync_bn(model: nn.Module) -> nn.Module:
    """
    Convert all BatchNorm layers to SyncBatchNorm for distributed training.
    Must be called BEFORE wrapping model with DistributedDataParallel.
    """
    return nn.SyncBatchNorm.convert_sync_batchnorm(model)


# Distributed training setup
def setup_distributed_training(model: nn.Module, local_rank: int) -> nn.Module:
    """
    Full distributed training setup with SyncBatchNorm.

    local_rank: GPU index (0, 1, ..., n_gpus-1)
    """
    # Convert BN to SyncBN BEFORE DDP wrapping
    model = convert_to_sync_bn(model)
    model = model.cuda(local_rank)

    # Wrap with DistributedDataParallel
    model = nn.parallel.DistributedDataParallel(
        model,
        device_ids=[local_rank],
        output_device=local_rank,
    )

    return model


# Cost of SyncBatchNorm:
# - Requires all-reduce communication across GPUs at each BN layer
# - Adds ~10-30% training overhead on fast interconnects (NVLink)
# - Much higher overhead on slower interconnects (TCP/IP)
# - For small models or fast networks, SyncBN overhead may exceed the benefit
# - Consider GroupNorm (no communication needed) for detection/segmentation models

When SyncBN matters most: object detection models (Mask R-CNN, Faster R-CNN) trained with small local batch sizes (2–4 images per GPU). With 8 GPUs × 2 images = 16 images total, SyncBN uses all 16 for statistics instead of just 2. For image classification with large batches (64+ per GPU), standard BN is adequate.

Normalization in Vision Transformers

Vision Transformers (ViT) use LayerNorm throughout - the same Pre-LN pattern as language transformers. However, the normalization position relative to patch embeddings matters:

import torch
import torch.nn as nn


class ViTPatchEmbedding(nn.Module):
    """
    Vision Transformer patch embedding with normalization.
    Divides image into fixed-size patches, embeds each patch as a vector.
    """

    def __init__(self, image_size: int = 224, patch_size: int = 16,
                 in_channels: int = 3, d_model: int = 768):
        super().__init__()
        assert image_size % patch_size == 0
        self.n_patches = (image_size // patch_size) ** 2

        # Patch embedding: convolve with patch_size stride - no overlap
        self.proj = nn.Conv2d(
            in_channels, d_model,
            kernel_size=patch_size, stride=patch_size
        )

        # LayerNorm applied AFTER patch embedding
        # This normalizes the initial patch representations
        # WITHOUT this, patch embeddings can have very different magnitudes
        # depending on image brightness - causing training instability
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, C, H, W)
        x = self.proj(x)          # (B, d_model, H/P, W/P)
        x = x.flatten(2)          # (B, d_model, n_patches)
        x = x.transpose(1, 2)     # (B, n_patches, d_model)
        x = self.norm(x)          # normalize each patch embedding
        return x

The LayerNorm after patch projection is important for ViT training stability - raw pixel values from different image regions have very different statistical properties, and normalizing the patch embeddings before the transformer blocks ensures the attention layers receive consistently-scaled inputs.

:::tip Normalization Layer Parameter Count LayerNorm and BatchNorm both add $2d$ learnable parameters ($\gamma$ and $\beta$) per normalized dimension $d$. For a transformer with $d_{\text{model}} = 768$ and 12 layers (24 norms in a Pre-LN architecture), that is $24 \times 2 \times 768 = 36,864$ parameters - negligible compared to the total model size (typically 85M+ for BERT-base). RMSNorm removes $\beta$, reducing to $d$ parameters per layer. This saving is also negligible but removes one source of potential training instability (the shift parameter can occasionally cause numerical issues in very deep networks). :::

:::tip 🎮 Interactive Playground

Visualize this concept: Try the Batch Normalization Effect demo on the EngineersOfAI Playground - no code required.

:::