PyTorch Batch Normalization

Batch Normalization (BatchNorm) normalizes the activations of each layer across the current training batch. It stabilizes training, allows higher learning rates, and acts as a mild regularizer. Almost every modern CNN architecture uses BatchNorm after convolutional layers.

The Problem BatchNorm Solves

During training, each layer's output depends on all the weights in previous layers. As those weights update, the distribution of values flowing into the next layer shifts constantly. Each layer must constantly adapt to a moving target. This is called internal covariate shift. It slows down training and forces you to use very small learning rates.

Without BatchNorm:
  Layer 1 outputs → mean=2.3, std=0.8   (batch 1)
  Layer 1 outputs → mean=5.1, std=2.1   (batch 2)  ← distribution drifts
  Layer 2 must cope with wildly changing input distributions

With BatchNorm:
  Layer 1 outputs → normalized to mean≈0, std≈1    (every batch)
  Layer 2 always receives a stable distribution → trains faster and more stably

How BatchNorm Works

For each mini-batch:

1. Compute mean μ and variance σ² across the batch
2. Normalize: x_hat = (x - μ) / sqrt(σ² + ε)
3. Scale and shift: y = γ × x_hat + β
   γ and β are learnable parameters (initialized to 1 and 0)

γ lets the network learn the best output scale.
β lets the network learn the best output shift.

BatchNorm2d for CNNs

Use nn.BatchNorm2d after convolutional layers. It normalizes across the batch and spatial dimensions for each channel separately.

import torch
import torch.nn as nn

# Normalize after a conv layer with 32 output channels
block = nn.Sequential(
    nn.Conv2d(3, 32, kernel_size=3, padding=1, bias=False),
    nn.BatchNorm2d(32),   # 32 must match out_channels of previous conv
    nn.ReLU()
)

x = torch.rand(16, 3, 64, 64)   # batch of 16 color images
out = block(x)
print(out.shape)   # torch.Size([16, 32, 64, 64])

Note: When using BatchNorm after a Linear or Conv layer, set bias=False in that layer. BatchNorm already learns a shift parameter (β) that makes the bias redundant.

BatchNorm1d for Fully Connected Layers

fc_block = nn.Sequential(
    nn.Linear(256, 128, bias=False),
    nn.BatchNorm1d(128),
    nn.ReLU()
)

x = torch.rand(32, 256)   # batch of 32 samples, 256 features
out = fc_block(x)
print(out.shape)   # torch.Size([32, 128])

Training Mode vs Evaluation Mode

BatchNorm behaves differently in training vs evaluation mode. This is one of the most important differences between model.train() and model.eval().

model.train():
  BatchNorm uses mean and variance of the CURRENT BATCH
  These change from batch to batch → regularization effect

model.eval():
  BatchNorm uses RUNNING MEAN and RUNNING VARIANCE
  These are computed during training and are fixed at eval time
  → predictions are stable and deterministic

# During training
model.train()
out_train = model(x)   # uses current batch stats

# During evaluation
model.eval()
with torch.no_grad():
    out_eval = model(x)   # uses running stats accumulated during training

BatchNorm Learnable Parameters

bn = nn.BatchNorm2d(32)

print(bn.weight.shape)   # torch.Size([32])   ← γ (scale), one per channel
print(bn.bias.shape)     # torch.Size([32])   ← β (shift), one per channel
print(bn.running_mean.shape)   # torch.Size([32])   ← tracked during training
print(bn.running_var.shape)    # torch.Size([32])   ← tracked during training

Where to Place BatchNorm

Standard order (most common):
  Conv → BatchNorm → Activation (ReLU)

Alternative (some architectures):
  Conv → Activation → BatchNorm

In residual networks (ResNet):
  Input → Conv → BN → ReLU → Conv → BN → + (residual) → ReLU

Benefits of BatchNorm

• Allows much higher learning rates — training is 5–10× faster
• Reduces sensitivity to weight initialization
• Acts as regularization — often reduces the need for Dropout
• Makes gradient flow more stable through deep networks

GroupNorm and LayerNorm — Alternatives

BatchNorm requires a minimum batch size to compute meaningful statistics. It fails with batch size 1. Two alternatives avoid this limitation:

• LayerNorm — normalizes across the feature dimension for each sample individually. Standard in transformers and NLP models.
• GroupNorm — normalizes across groups of channels. Works with any batch size. Common in object detection and segmentation models.

ln = nn.LayerNorm(128)    # normalizes over last 128 dimensions
gn = nn.GroupNorm(8, 32)  # 8 groups, 32 channels total (4 channels per group)

Summary

Batch Normalization normalizes layer activations to a stable distribution during training. It reduces internal covariate shift, accelerates convergence, enables higher learning rates, and provides mild regularization. Use nn.BatchNorm2d after Conv2d layers and nn.BatchNorm1d after Linear layers. Set bias=False on the preceding layer since BatchNorm learns its own shift. Always call model.train() during training and model.eval() during validation to ensure BatchNorm uses the correct statistics. For small batches or transformers, use LayerNorm instead.