Layer & Batch Normalization - Batch Normalization

Batch Normalization

What It Does

BatchNorm normalizes each feature across the batch dimension.

Given a 2D input matrix where:

• Rows = batch samples
• Columns = features

BatchNorm normalizes each column (feature) independently across all samples in the batch.

How It Works

For each feature $j$ in the batch:

1. Compute mean and variance across batch:

$${\mu }_j=\frac{1}{m}\sum _{i=1}^m{x}_{ij}$$ $${\sigma }_j^2=\frac{1}{m}\sum _{i=1}^m(x_{ij}-{\mu }_j{)}^2$$

2. Normalize:

$${\hat{x}}_{ij}=\frac{x_{ij}-{\mu }_j}{\sqrt{{\sigma }_j^2+ϵ}}$$

3. Scale and shift (learnable parameters):

$$y_{ij}={\gamma }_j{\hat{x}}_{ij}+{\beta }_j$$

Code Example

import torch.nn as nn
 
model = nn.Sequential(
    nn.Linear(64, 128),
    nn.BatchNorm1d(128), # Normalize across batch for each of 128 features
    nn.ReLU(),
    nn.Linear(128, 10)
)

Advantages

• Stabilizes training (reduces internal covariate shift)
• Allows higher learning rates
• Reduces sensitivity to initialization
• Acts as regularization (slight noise from batch statistics)
• Accelerates convergence

Disadvantages

• Requires sufficiently large batch size (typically >32)
• Different behavior in training vs inference
• Doesn’t work well with small batches or online learning
• Complications with RNNs (sequence length varies)

Training Vs Inference

Training:

• Uses batch statistics (mean, variance of current batch)

Inference:

• Uses running statistics (exponential moving average from training)

model.train() # Use batch statistics
model.eval() # Use running statistics