916 lines
26 KiB
Markdown
916 lines
26 KiB
Markdown
|
||
# Normalization Techniques
|
||
|
||
## Context
|
||
|
||
You're designing a neural network or debugging training instability. Someone suggests "add BatchNorm" without considering:
|
||
- **Batch size dependency**: BatchNorm fails with small batches (< 8)
|
||
- **Architecture mismatch**: BatchNorm breaks RNNs/Transformers (use LayerNorm)
|
||
- **Task-specific needs**: Style transfer needs InstanceNorm, not BatchNorm
|
||
- **Modern alternatives**: RMSNorm simpler and faster than LayerNorm for LLMs
|
||
|
||
**This skill prevents normalization cargo-culting and provides architecture-specific selection.**
|
||
|
||
## Why Normalization Matters
|
||
|
||
**Problem: Internal Covariate Shift**
|
||
During training, layer input distributions shift as previous layers update. This causes:
|
||
- Vanishing/exploding gradients (deep networks)
|
||
- Slow convergence (small learning rates required)
|
||
- Training instability (loss spikes)
|
||
|
||
**Solution: Normalization**
|
||
Normalize activations to have stable statistics (mean=0, std=1). Benefits:
|
||
- **10x faster convergence**: Can use larger learning rates
|
||
- **Better generalization**: Regularization effect
|
||
- **Enables deep networks**: 50+ layers without gradient issues
|
||
- **Less sensitive to initialization**: Weights can start further from optimal
|
||
|
||
**Key insight**: Normalization is NOT optional for modern deep learning. The question is WHICH normalization, not WHETHER to normalize.
|
||
|
||
|
||
## Normalization Families
|
||
|
||
### 1. Batch Normalization (BatchNorm)
|
||
|
||
**What it does:**
|
||
Normalizes across the batch dimension for each channel/feature.
|
||
|
||
**Formula:**
|
||
```
|
||
Given input x with shape (B, C, H, W): # Batch, Channel, Height, Width
|
||
|
||
For each channel c:
|
||
μ_c = mean(x[:, c, :, :]) # Mean over batch + spatial dims
|
||
σ_c = std(x[:, c, :, :]) # Std over batch + spatial dims
|
||
|
||
x_norm[:, c, :, :] = (x[:, c, :, :] - μ_c) / √(σ_c² + ε)
|
||
|
||
# Learnable scale and shift
|
||
y[:, c, :, :] = γ_c * x_norm[:, c, :, :] + β_c
|
||
```
|
||
|
||
**When to use:**
|
||
- ✅ CNNs for classification (ResNet, EfficientNet)
|
||
- ✅ Large batch sizes (≥ 16)
|
||
- ✅ IID data (image classification, object detection)
|
||
|
||
**When NOT to use:**
|
||
- ❌ Small batch sizes (< 8): Noisy statistics cause training failure
|
||
- ❌ RNNs/LSTMs: Breaks temporal dependencies
|
||
- ❌ Transformers: Batch dependency problematic for variable-length sequences
|
||
- ❌ Style transfer: Batch statistics erase style information
|
||
|
||
**Batch size dependency:**
|
||
```python
|
||
batch_size = 32: # ✓ Works well (stable statistics)
|
||
batch_size = 16: # ✓ Acceptable
|
||
batch_size = 8: # ✓ Marginal (consider GroupNorm)
|
||
batch_size = 4: # ✗ Unstable (use GroupNorm)
|
||
batch_size = 2: # ✗ FAILS! (noisy statistics)
|
||
batch_size = 1: # ✗ Undefined (no batch to normalize over!)
|
||
```
|
||
|
||
**PyTorch example:**
|
||
```python
|
||
import torch.nn as nn
|
||
|
||
# For Conv2d
|
||
bn = nn.BatchNorm2d(num_features=64) # 64 channels
|
||
x = torch.randn(32, 64, 28, 28) # Batch=32, Channels=64
|
||
y = bn(x)
|
||
|
||
# For Linear
|
||
bn = nn.BatchNorm1d(num_features=128) # 128 features
|
||
x = torch.randn(32, 128) # Batch=32, Features=128
|
||
y = bn(x)
|
||
```
|
||
|
||
**Inference mode:**
|
||
```python
|
||
# Training: Uses batch statistics
|
||
model.train()
|
||
y = bn(x) # Normalizes using current batch mean/std
|
||
|
||
# Inference: Uses running statistics (accumulated during training)
|
||
model.eval()
|
||
y = bn(x) # Normalizes using running_mean/running_std
|
||
```
|
||
|
||
|
||
### 2. Layer Normalization (LayerNorm)
|
||
|
||
**What it does:**
|
||
Normalizes across the feature dimension for each sample independently.
|
||
|
||
**Formula:**
|
||
```
|
||
Given input x with shape (B, C): # Batch, Features
|
||
|
||
For each sample b:
|
||
μ_b = mean(x[b, :]) # Mean over features
|
||
σ_b = std(x[b, :]) # Std over features
|
||
|
||
x_norm[b, :] = (x[b, :] - μ_b) / √(σ_b² + ε)
|
||
|
||
# Learnable scale and shift
|
||
y[b, :] = γ * x_norm[b, :] + β
|
||
```
|
||
|
||
**When to use:**
|
||
- ✅ Transformers (BERT, GPT, T5)
|
||
- ✅ RNNs/LSTMs (maintains temporal independence)
|
||
- ✅ Small batch sizes (batch-independent!)
|
||
- ✅ Variable-length sequences
|
||
- ✅ Reinforcement learning (batch_size=1 common)
|
||
|
||
**Advantages over BatchNorm:**
|
||
- ✅ **Batch-independent**: Works with batch_size=1
|
||
- ✅ **No running statistics**: Inference = training (no mode switching)
|
||
- ✅ **Sequence-friendly**: Doesn't mix information across timesteps
|
||
|
||
**PyTorch example:**
|
||
```python
|
||
import torch.nn as nn
|
||
|
||
# For Transformer
|
||
ln = nn.LayerNorm(normalized_shape=512) # d_model=512
|
||
x = torch.randn(32, 128, 512) # Batch=32, SeqLen=128, d_model=512
|
||
y = ln(x) # Normalizes last dimension independently per (batch, position)
|
||
|
||
# For RNN hidden states
|
||
ln = nn.LayerNorm(normalized_shape=256) # hidden_size=256
|
||
h = torch.randn(32, 256) # Batch=32, Hidden=256
|
||
h_norm = ln(h)
|
||
```
|
||
|
||
**Key difference from BatchNorm:**
|
||
```python
|
||
# BatchNorm: Normalizes across batch dimension
|
||
# Given (B=32, C=64, H=28, W=28)
|
||
# Computes 64 means/stds (one per channel, across batch + spatial)
|
||
|
||
# LayerNorm: Normalizes across feature dimension
|
||
# Given (B=32, L=128, D=512)
|
||
# Computes 32×128 means/stds (one per (batch, position), across features)
|
||
```
|
||
|
||
|
||
### 3. Group Normalization (GroupNorm)
|
||
|
||
**What it does:**
|
||
Normalizes channels in groups, batch-independent.
|
||
|
||
**Formula:**
|
||
```
|
||
Given input x with shape (B, C, H, W):
|
||
Divide C channels into G groups (C must be divisible by G)
|
||
|
||
For each sample b and group g:
|
||
channels = x[b, g*(C/G):(g+1)*(C/G), :, :] # Channels in group g
|
||
μ_{b,g} = mean(channels) # Mean over channels in group + spatial
|
||
σ_{b,g} = std(channels) # Std over channels in group + spatial
|
||
|
||
x_norm[b, g*(C/G):(g+1)*(C/G), :, :] = (channels - μ_{b,g}) / √(σ_{b,g}² + ε)
|
||
```
|
||
|
||
**When to use:**
|
||
- ✅ Small batch sizes (< 8)
|
||
- ✅ CNNs with batch_size=1 (style transfer, RL)
|
||
- ✅ Object detection/segmentation (often use small batches)
|
||
- ✅ When BatchNorm unstable but want spatial normalization
|
||
|
||
**Group size selection:**
|
||
```python
|
||
# num_groups trade-off:
|
||
num_groups = 1: # = LayerNorm (all channels together)
|
||
num_groups = C: # = InstanceNorm (each channel separate)
|
||
num_groups = 32: # Standard choice (good balance)
|
||
|
||
# Rule: C must be divisible by num_groups
|
||
channels = 64, num_groups = 32: # ✓ 64/32 = 2 channels per group
|
||
channels = 64, num_groups = 16: # ✓ 64/16 = 4 channels per group
|
||
channels = 64, num_groups = 30: # ✗ 64/30 not integer!
|
||
```
|
||
|
||
**PyTorch example:**
|
||
```python
|
||
import torch.nn as nn
|
||
|
||
# For small batch CNN
|
||
gn = nn.GroupNorm(num_groups=32, num_channels=64)
|
||
x = torch.randn(2, 64, 28, 28) # Batch=2 (small!)
|
||
y = gn(x) # Works well even with batch=2
|
||
|
||
# Compare performance:
|
||
batch_sizes = [1, 2, 4, 8, 16, 32]
|
||
|
||
bn = nn.BatchNorm2d(64)
|
||
gn = nn.GroupNorm(32, 64)
|
||
|
||
for bs in batch_sizes:
|
||
x = torch.randn(bs, 64, 28, 28)
|
||
|
||
# BatchNorm gets more stable with larger batch
|
||
# GroupNorm consistent across all batch sizes
|
||
```
|
||
|
||
**Empirical results (He et al. 2018):**
|
||
```
|
||
ImageNet classification with ResNet-50:
|
||
batch_size = 32: BatchNorm = 76.5%, GroupNorm = 76.3% (tie)
|
||
batch_size = 8: BatchNorm = 75.8%, GroupNorm = 76.1% (GroupNorm wins!)
|
||
batch_size = 2: BatchNorm = 72.1%, GroupNorm = 75.3% (GroupNorm wins!)
|
||
```
|
||
|
||
|
||
### 4. Instance Normalization (InstanceNorm)
|
||
|
||
**What it does:**
|
||
Normalizes each sample and channel independently (no batch mixing).
|
||
|
||
**Formula:**
|
||
```
|
||
Given input x with shape (B, C, H, W):
|
||
|
||
For each sample b and channel c:
|
||
μ_{b,c} = mean(x[b, c, :, :]) # Mean over spatial dimensions only
|
||
σ_{b,c} = std(x[b, c, :, :]) # Std over spatial dimensions only
|
||
|
||
x_norm[b, c, :, :] = (x[b, c, :, :] - μ_{b,c}) / √(σ_{b,c}² + ε)
|
||
```
|
||
|
||
**When to use:**
|
||
- ✅ Style transfer (neural style, CycleGAN, pix2pix)
|
||
- ✅ Image-to-image translation
|
||
- ✅ When batch/channel mixing destroys information
|
||
|
||
**Why for style transfer:**
|
||
```python
|
||
# Style transfer goal: Transfer style while preserving content
|
||
# BatchNorm: Mixes statistics across batch (erases individual style!)
|
||
# InstanceNorm: Per-image statistics (preserves each image's style)
|
||
|
||
# Example: Neural style transfer
|
||
content_image = load_image("photo.jpg")
|
||
style_image = load_image("starry_night.jpg")
|
||
|
||
# With BatchNorm: Output loses content image's unique characteristics
|
||
# With InstanceNorm: Content characteristics preserved, style applied
|
||
```
|
||
|
||
**PyTorch example:**
|
||
```python
|
||
import torch.nn as nn
|
||
|
||
# For style transfer generator
|
||
in_norm = nn.InstanceNorm2d(num_features=64)
|
||
x = torch.randn(1, 64, 256, 256) # Single image
|
||
y = in_norm(x) # Normalizes each channel independently
|
||
|
||
# CycleGAN generator architecture
|
||
class Generator(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.conv1 = nn.Conv2d(3, 64, 7, padding=3)
|
||
self.in1 = nn.InstanceNorm2d(64) # NOT BatchNorm!
|
||
self.relu = nn.ReLU(inplace=True)
|
||
|
||
def forward(self, x):
|
||
x = self.conv1(x)
|
||
x = self.in1(x) # Per-image normalization
|
||
x = self.relu(x)
|
||
return x
|
||
```
|
||
|
||
**Relation to GroupNorm:**
|
||
```python
|
||
# InstanceNorm is GroupNorm with num_groups = num_channels
|
||
InstanceNorm2d(C) == GroupNorm(num_groups=C, num_channels=C)
|
||
```
|
||
|
||
|
||
### 5. RMS Normalization (RMSNorm)
|
||
|
||
**What it does:**
|
||
Simplified LayerNorm that only rescales (no recentering), faster and simpler.
|
||
|
||
**Formula:**
|
||
```
|
||
Given input x:
|
||
|
||
# LayerNorm (2 steps):
|
||
x_centered = x - mean(x) # 1. Center
|
||
x_norm = x_centered / std(x) # 2. Scale
|
||
|
||
# RMSNorm (1 step):
|
||
rms = sqrt(mean(x²)) # Root Mean Square
|
||
x_norm = x / rms # Only scale, no centering
|
||
```
|
||
|
||
**When to use:**
|
||
- ✅ Modern LLMs (LLaMA, Mistral, Gemma)
|
||
- ✅ When speed matters (15-20% faster than LayerNorm)
|
||
- ✅ Large Transformer models (billions of parameters)
|
||
|
||
**Advantages:**
|
||
- ✅ **Simpler**: One operation instead of two
|
||
- ✅ **Faster**: ~15-20% speedup over LayerNorm
|
||
- ✅ **Numerically stable**: No subtraction (avoids catastrophic cancellation)
|
||
- ✅ **Same performance**: Empirically matches LayerNorm quality
|
||
|
||
**PyTorch implementation:**
|
||
```python
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
class RMSNorm(nn.Module):
|
||
def __init__(self, dim, eps=1e-6):
|
||
super().__init__()
|
||
self.eps = eps
|
||
self.weight = nn.Parameter(torch.ones(dim))
|
||
|
||
def forward(self, x):
|
||
# Compute RMS
|
||
rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
|
||
|
||
# Normalize
|
||
x_norm = x / rms
|
||
|
||
# Scale (learnable)
|
||
return self.weight * x_norm
|
||
|
||
# Usage in Transformer
|
||
rms = RMSNorm(dim=512) # d_model=512
|
||
x = torch.randn(32, 128, 512) # Batch, SeqLen, d_model
|
||
y = rms(x)
|
||
```
|
||
|
||
**Speed comparison (LLaMA-7B, A100 GPU):**
|
||
```
|
||
LayerNorm: 1000 tokens/sec
|
||
RMSNorm: 1180 tokens/sec # 18% faster!
|
||
|
||
# For large models, this adds up:
|
||
# 1 million tokens: 180 seconds saved
|
||
```
|
||
|
||
**Modern LLM adoption:**
|
||
```python
|
||
# LLaMA (Meta, 2023): RMSNorm
|
||
# Mistral (Mistral AI, 2023): RMSNorm
|
||
# Gemma (Google, 2024): RMSNorm
|
||
# PaLM (Google, 2022): RMSNorm
|
||
|
||
# Older models:
|
||
# GPT-2/3 (OpenAI): LayerNorm
|
||
# BERT (Google, 2018): LayerNorm
|
||
```
|
||
|
||
|
||
## Architecture-Specific Selection
|
||
|
||
### CNN (Convolutional Neural Networks)
|
||
|
||
**Default: BatchNorm**
|
||
```python
|
||
import torch.nn as nn
|
||
|
||
class ConvBlock(nn.Module):
|
||
def __init__(self, in_channels, out_channels):
|
||
super().__init__()
|
||
self.conv = nn.Conv2d(in_channels, out_channels, 3, padding=1)
|
||
self.bn = nn.BatchNorm2d(out_channels) # After conv
|
||
self.relu = nn.ReLU(inplace=True)
|
||
|
||
def forward(self, x):
|
||
x = self.conv(x)
|
||
x = self.bn(x) # Normalize
|
||
x = self.relu(x)
|
||
return x
|
||
```
|
||
|
||
**Exception: Small batch sizes**
|
||
```python
|
||
# If batch_size < 8, use GroupNorm instead
|
||
class ConvBlock(nn.Module):
|
||
def __init__(self, in_channels, out_channels):
|
||
super().__init__()
|
||
self.conv = nn.Conv2d(in_channels, out_channels, 3, padding=1)
|
||
self.norm = nn.GroupNorm(32, out_channels) # GroupNorm for small batches
|
||
self.relu = nn.ReLU(inplace=True)
|
||
```
|
||
|
||
**Exception: Style transfer**
|
||
```python
|
||
# Use InstanceNorm for style transfer
|
||
class StyleConvBlock(nn.Module):
|
||
def __init__(self, in_channels, out_channels):
|
||
super().__init__()
|
||
self.conv = nn.Conv2d(in_channels, out_channels, 3, padding=1)
|
||
self.norm = nn.InstanceNorm2d(out_channels) # Per-image normalization
|
||
self.relu = nn.ReLU(inplace=True)
|
||
```
|
||
|
||
|
||
### RNN / LSTM (Recurrent Neural Networks)
|
||
|
||
**Default: LayerNorm**
|
||
```python
|
||
import torch.nn as nn
|
||
|
||
class NormalizedLSTM(nn.Module):
|
||
def __init__(self, input_size, hidden_size):
|
||
super().__init__()
|
||
self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
|
||
self.ln = nn.LayerNorm(hidden_size) # Normalize hidden states
|
||
|
||
def forward(self, x):
|
||
# x: (batch, seq_len, input_size)
|
||
output, (h_n, c_n) = self.lstm(x)
|
||
# output: (batch, seq_len, hidden_size)
|
||
|
||
# Normalize each timestep's output
|
||
output_norm = self.ln(output) # Applies independently per timestep
|
||
return output_norm, (h_n, c_n)
|
||
```
|
||
|
||
**Why NOT BatchNorm:**
|
||
```python
|
||
# BatchNorm in RNN mixes information across timesteps!
|
||
# Given (batch=32, seq_len=100, hidden=256)
|
||
|
||
# BatchNorm would compute:
|
||
# mean/std over (batch × seq_len) = 3200 values
|
||
# This mixes t=0 with t=99 (destroys temporal structure!)
|
||
|
||
# LayerNorm computes:
|
||
# mean/std over hidden_size = 256 values per (batch, timestep)
|
||
# Each timestep normalized independently (preserves temporal structure)
|
||
```
|
||
|
||
**Layer-wise normalization in stacked RNN:**
|
||
```python
|
||
class StackedNormalizedLSTM(nn.Module):
|
||
def __init__(self, input_size, hidden_size, num_layers):
|
||
super().__init__()
|
||
self.layers = nn.ModuleList()
|
||
|
||
for i in range(num_layers):
|
||
in_size = input_size if i == 0 else hidden_size
|
||
self.layers.append(nn.LSTM(in_size, hidden_size, batch_first=True))
|
||
self.layers.append(nn.LayerNorm(hidden_size)) # After each LSTM layer
|
||
|
||
def forward(self, x):
|
||
for lstm, ln in zip(self.layers[::2], self.layers[1::2]):
|
||
x, _ = lstm(x)
|
||
x = ln(x) # Normalize between layers
|
||
return x
|
||
```
|
||
|
||
|
||
### Transformer
|
||
|
||
**Default: LayerNorm (or RMSNorm for modern/large models)**
|
||
|
||
**Two placement options: Pre-norm vs Post-norm**
|
||
|
||
**Post-norm (original Transformer, "Attention is All You Need"):**
|
||
```python
|
||
class TransformerLayerPostNorm(nn.Module):
|
||
def __init__(self, d_model, num_heads):
|
||
super().__init__()
|
||
self.attn = nn.MultiheadAttention(d_model, num_heads)
|
||
self.ffn = nn.Sequential(
|
||
nn.Linear(d_model, 4 * d_model),
|
||
nn.ReLU(),
|
||
nn.Linear(4 * d_model, d_model)
|
||
)
|
||
self.ln1 = nn.LayerNorm(d_model)
|
||
self.ln2 = nn.LayerNorm(d_model)
|
||
|
||
def forward(self, x):
|
||
# Post-norm: Apply normalization AFTER residual
|
||
x = self.ln1(x + self.attn(x, x, x)[0]) # Normalize after adding
|
||
x = self.ln2(x + self.ffn(x)) # Normalize after adding
|
||
return x
|
||
```
|
||
|
||
**Pre-norm (modern, more stable):**
|
||
```python
|
||
class TransformerLayerPreNorm(nn.Module):
|
||
def __init__(self, d_model, num_heads):
|
||
super().__init__()
|
||
self.attn = nn.MultiheadAttention(d_model, num_heads)
|
||
self.ffn = nn.Sequential(
|
||
nn.Linear(d_model, 4 * d_model),
|
||
nn.ReLU(),
|
||
nn.Linear(4 * d_model, d_model)
|
||
)
|
||
self.ln1 = nn.LayerNorm(d_model)
|
||
self.ln2 = nn.LayerNorm(d_model)
|
||
|
||
def forward(self, x):
|
||
# Pre-norm: Apply normalization BEFORE sublayer
|
||
x = x + self.attn(self.ln1(x), self.ln1(x), self.ln1(x))[0] # Normalize before attention
|
||
x = x + self.ffn(self.ln2(x)) # Normalize before FFN
|
||
return x
|
||
```
|
||
|
||
**Pre-norm vs Post-norm comparison:**
|
||
```python
|
||
# Post-norm (original):
|
||
# - Less stable (requires careful initialization + warmup)
|
||
# - Slightly better performance IF training succeeds
|
||
# - Hard to train deep models (> 12 layers)
|
||
|
||
# Pre-norm (modern):
|
||
# - More stable (easier to train deep models)
|
||
# - Standard for large models (GPT-3: 96 layers!)
|
||
# - Recommended default
|
||
|
||
# Empirical: GPT-2, BERT (post-norm, ≤12 layers)
|
||
# GPT-3, T5, LLaMA (pre-norm, ≥24 layers)
|
||
```
|
||
|
||
**Using RMSNorm instead:**
|
||
```python
|
||
class TransformerLayerRMSNorm(nn.Module):
|
||
def __init__(self, d_model, num_heads):
|
||
super().__init__()
|
||
self.attn = nn.MultiheadAttention(d_model, num_heads)
|
||
self.ffn = nn.Sequential(
|
||
nn.Linear(d_model, 4 * d_model),
|
||
nn.ReLU(),
|
||
nn.Linear(4 * d_model, d_model)
|
||
)
|
||
self.rms1 = RMSNorm(d_model) # 15-20% faster than LayerNorm
|
||
self.rms2 = RMSNorm(d_model)
|
||
|
||
def forward(self, x):
|
||
# Pre-norm with RMSNorm (LLaMA style)
|
||
x = x + self.attn(self.rms1(x), self.rms1(x), self.rms1(x))[0]
|
||
x = x + self.ffn(self.rms2(x))
|
||
return x
|
||
```
|
||
|
||
|
||
### GAN (Generative Adversarial Network)
|
||
|
||
**Generator: InstanceNorm or no normalization**
|
||
```python
|
||
class Generator(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
# Use InstanceNorm for image-to-image translation
|
||
self.conv1 = nn.Conv2d(3, 64, 7, padding=3)
|
||
self.in1 = nn.InstanceNorm2d(64) # Per-image normalization
|
||
|
||
def forward(self, x):
|
||
x = self.conv1(x)
|
||
x = self.in1(x) # Preserves per-image characteristics
|
||
return x
|
||
```
|
||
|
||
**Discriminator: No normalization or LayerNorm**
|
||
```python
|
||
class Discriminator(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
# Often no normalization (BatchNorm can hurt GAN training)
|
||
self.conv1 = nn.Conv2d(3, 64, 4, stride=2, padding=1)
|
||
# No normalization here
|
||
|
||
def forward(self, x):
|
||
x = self.conv1(x)
|
||
# Directly to activation (no norm)
|
||
return x
|
||
```
|
||
|
||
**Why avoid BatchNorm in GANs:**
|
||
```python
|
||
# BatchNorm in discriminator:
|
||
# - Mixes real and fake samples in batch
|
||
# - Leaks information (discriminator can detect batch composition)
|
||
# - Hurts training stability
|
||
|
||
# Recommendation:
|
||
# Generator: InstanceNorm (for image translation) or no norm
|
||
# Discriminator: No normalization or LayerNorm
|
||
```
|
||
|
||
|
||
## Decision Framework
|
||
|
||
### Step 1: Check batch size
|
||
|
||
```python
|
||
if batch_size >= 8:
|
||
consider_batchnorm = True
|
||
else:
|
||
use_groupnorm_or_layernorm = True # BatchNorm will be unstable
|
||
```
|
||
|
||
### Step 2: Check architecture
|
||
|
||
```python
|
||
if architecture == "CNN":
|
||
if batch_size >= 8:
|
||
use_batchnorm()
|
||
else:
|
||
use_groupnorm(num_groups=32)
|
||
|
||
# Exception: Style transfer
|
||
if task == "style_transfer":
|
||
use_instancenorm()
|
||
|
||
elif architecture in ["RNN", "LSTM", "GRU"]:
|
||
use_layernorm() # NEVER BatchNorm!
|
||
|
||
elif architecture == "Transformer":
|
||
if model_size == "large": # > 1B parameters
|
||
use_rmsnorm() # 15-20% faster
|
||
else:
|
||
use_layernorm()
|
||
|
||
# Placement: Pre-norm (more stable)
|
||
use_prenorm_placement()
|
||
|
||
elif architecture == "GAN":
|
||
if component == "generator":
|
||
if task == "image_translation":
|
||
use_instancenorm()
|
||
else:
|
||
use_no_norm() # Or InstanceNorm
|
||
elif component == "discriminator":
|
||
use_no_norm() # Or LayerNorm
|
||
```
|
||
|
||
### Step 3: Verify placement
|
||
|
||
```python
|
||
# CNNs: After convolution, before activation
|
||
x = conv(x)
|
||
x = norm(x) # Here!
|
||
x = relu(x)
|
||
|
||
# RNNs: After LSTM, normalize hidden states
|
||
output, (h, c) = lstm(x)
|
||
output = norm(output) # Here!
|
||
|
||
# Transformers: Pre-norm (modern) or post-norm (original)
|
||
# Pre-norm (recommended):
|
||
x = x + sublayer(norm(x)) # Normalize before sublayer
|
||
|
||
# Post-norm (original):
|
||
x = norm(x + sublayer(x)) # Normalize after residual
|
||
```
|
||
|
||
|
||
## Implementation Checklist
|
||
|
||
### Before adding normalization:
|
||
|
||
1. ☐ **Check batch size**: If < 8, avoid BatchNorm
|
||
2. ☐ **Check architecture**: CNN→BatchNorm, RNN→LayerNorm, Transformer→LayerNorm/RMSNorm
|
||
3. ☐ **Check task**: Style transfer→InstanceNorm
|
||
4. ☐ **Verify placement**: After conv/linear, before activation (CNNs)
|
||
5. ☐ **Test training stability**: Loss should decrease smoothly
|
||
|
||
### During training:
|
||
|
||
6. ☐ **Monitor running statistics** (BatchNorm): Check running_mean/running_var are updating
|
||
7. ☐ **Test inference mode**: Verify model.eval() uses running stats correctly
|
||
8. ☐ **Check gradient flow**: Normalization should help, not hurt gradients
|
||
|
||
### If training is unstable:
|
||
|
||
9. ☐ **Try different normalization**: BatchNorm→GroupNorm, LayerNorm→RMSNorm
|
||
10. ☐ **Try pre-norm** (Transformers): More stable than post-norm
|
||
11. ☐ **Reduce learning rate**: Normalization allows larger LR, but start conservatively
|
||
|
||
|
||
## Common Mistakes
|
||
|
||
### Mistake 1: BatchNorm with small batches
|
||
|
||
```python
|
||
# WRONG: BatchNorm with batch_size=2
|
||
model = ResNet50(norm_layer=nn.BatchNorm2d)
|
||
dataloader = DataLoader(dataset, batch_size=2) # Too small!
|
||
|
||
# RIGHT: GroupNorm for small batches
|
||
model = ResNet50(norm_layer=lambda channels: nn.GroupNorm(32, channels))
|
||
dataloader = DataLoader(dataset, batch_size=2) # Works!
|
||
```
|
||
|
||
### Mistake 2: BatchNorm in RNN
|
||
|
||
```python
|
||
# WRONG: BatchNorm in LSTM
|
||
class BadLSTM(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.lstm = nn.LSTM(100, 256)
|
||
self.bn = nn.BatchNorm1d(256) # WRONG! Mixes timesteps
|
||
|
||
def forward(self, x):
|
||
output, _ = self.lstm(x)
|
||
output = output.permute(0, 2, 1) # (B, H, T)
|
||
output = self.bn(output) # Mixes timesteps!
|
||
return output
|
||
|
||
# RIGHT: LayerNorm in LSTM
|
||
class GoodLSTM(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.lstm = nn.LSTM(100, 256)
|
||
self.ln = nn.LayerNorm(256) # Per-timestep normalization
|
||
|
||
def forward(self, x):
|
||
output, _ = self.lstm(x)
|
||
output = self.ln(output) # Independent per timestep
|
||
return output
|
||
```
|
||
|
||
### Mistake 3: Forgetting model.eval()
|
||
|
||
```python
|
||
# WRONG: Using training mode during inference
|
||
model.train() # BatchNorm uses batch statistics
|
||
predictions = model(test_data) # Batch statistics from test data (leakage!)
|
||
|
||
# RIGHT: Use eval mode during inference
|
||
model.eval() # BatchNorm uses running statistics
|
||
with torch.no_grad():
|
||
predictions = model(test_data) # Uses accumulated running stats
|
||
```
|
||
|
||
### Mistake 4: Post-norm for deep Transformers
|
||
|
||
```python
|
||
# WRONG: Post-norm for 24-layer Transformer (unstable!)
|
||
class DeepTransformerPostNorm(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.layers = nn.ModuleList([
|
||
TransformerLayerPostNorm(512, 8) for _ in range(24)
|
||
]) # Hard to train!
|
||
|
||
# RIGHT: Pre-norm for deep Transformers
|
||
class DeepTransformerPreNorm(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.layers = nn.ModuleList([
|
||
TransformerLayerPreNorm(512, 8) for _ in range(24)
|
||
]) # Stable training!
|
||
```
|
||
|
||
### Mistake 5: Wrong normalization for style transfer
|
||
|
||
```python
|
||
# WRONG: BatchNorm for style transfer
|
||
class StyleGenerator(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.conv = nn.Conv2d(3, 64, 7, padding=3)
|
||
self.norm = nn.BatchNorm2d(64) # WRONG! Mixes styles across batch
|
||
|
||
# RIGHT: InstanceNorm for style transfer
|
||
class StyleGenerator(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.conv = nn.Conv2d(3, 64, 7, padding=3)
|
||
self.norm = nn.InstanceNorm2d(64) # Per-image normalization
|
||
```
|
||
|
||
|
||
## Performance Impact
|
||
|
||
### Training speed:
|
||
|
||
```python
|
||
# Without normalization: 100 epochs to converge
|
||
# With normalization: 10 epochs to converge (10x faster!)
|
||
|
||
# Reason: Larger learning rates possible
|
||
lr_no_norm = 0.001 # Must be small (unstable otherwise)
|
||
lr_with_norm = 0.01 # Can be 10x larger (normalization stabilizes)
|
||
```
|
||
|
||
### Inference speed:
|
||
|
||
```python
|
||
# Normalization overhead (relative to no normalization):
|
||
BatchNorm: +2% (minimal, cached running stats)
|
||
LayerNorm: +3-5% (compute mean/std per forward pass)
|
||
RMSNorm: +2-3% (faster than LayerNorm)
|
||
GroupNorm: +5-8% (more computation than BatchNorm)
|
||
InstanceNorm: +3-5% (similar to LayerNorm)
|
||
|
||
# For most models: Overhead is negligible compared to conv/linear layers
|
||
```
|
||
|
||
### Memory usage:
|
||
|
||
```python
|
||
# Normalization memory (per layer):
|
||
BatchNorm: 2 × num_channels (running_mean, running_std) + 2 × num_channels (γ, β)
|
||
LayerNorm: 2 × normalized_shape (γ, β)
|
||
RMSNorm: 1 × normalized_shape (γ only, no β)
|
||
|
||
# Example: 512 channels
|
||
BatchNorm: 4 × 512 = 2048 parameters
|
||
LayerNorm: 2 × 512 = 1024 parameters
|
||
RMSNorm: 1 × 512 = 512 parameters # Most efficient!
|
||
```
|
||
|
||
|
||
## When NOT to Normalize
|
||
|
||
**Case 1: Final output layer**
|
||
```python
|
||
# Don't normalize final predictions
|
||
class Classifier(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.backbone = ResNet50() # Normalization inside
|
||
self.fc = nn.Linear(2048, 1000)
|
||
# NO normalization here! (final logits should be unnormalized)
|
||
|
||
def forward(self, x):
|
||
x = self.backbone(x)
|
||
x = self.fc(x) # Raw logits
|
||
return x # Don't normalize!
|
||
```
|
||
|
||
**Case 2: Very small networks**
|
||
```python
|
||
# Single-layer network: Normalization overkill
|
||
class TinyNet(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.fc = nn.Linear(784, 10) # MNIST classifier
|
||
# No normalization needed (network too simple)
|
||
|
||
def forward(self, x):
|
||
return self.fc(x)
|
||
```
|
||
|
||
**Case 3: When debugging**
|
||
```python
|
||
# Remove normalization to isolate issues
|
||
# If training fails with normalization, try without to check if:
|
||
# - Initialization is correct
|
||
# - Loss function is correct
|
||
# - Data is correctly preprocessed
|
||
```
|
||
|
||
|
||
## Modern Recommendations (2025)
|
||
|
||
### CNNs:
|
||
- **Default**: BatchNorm (if batch_size ≥ 8)
|
||
- **Small batches**: GroupNorm (num_groups=32)
|
||
- **Style transfer**: InstanceNorm
|
||
|
||
### RNNs/LSTMs:
|
||
- **Default**: LayerNorm
|
||
- **Never**: BatchNorm (breaks temporal structure)
|
||
|
||
### Transformers:
|
||
- **Small models** (< 1B): LayerNorm + pre-norm
|
||
- **Large models** (≥ 1B): RMSNorm + pre-norm (15-20% faster)
|
||
- **Avoid**: Post-norm for deep models (> 12 layers)
|
||
|
||
### GANs:
|
||
- **Generator**: InstanceNorm (image translation) or no norm
|
||
- **Discriminator**: No normalization or LayerNorm
|
||
- **Avoid**: BatchNorm (leaks information)
|
||
|
||
### Emerging:
|
||
- **RMSNorm adoption increasing**: LLaMA, Mistral, Gemma all use RMSNorm
|
||
- **Pre-norm becoming standard**: More stable for deep networks
|
||
- **GroupNorm gaining traction**: Object detection, small-batch training
|
||
|
||
|
||
## Summary
|
||
|
||
**Normalization is mandatory for modern deep learning.** The question is which normalization, not whether to normalize.
|
||
|
||
**Quick decision tree:**
|
||
1. **Batch size ≥ 8?** → Consider BatchNorm (CNNs)
|
||
2. **Batch size < 8?** → Use GroupNorm (CNNs) or LayerNorm (all)
|
||
3. **RNN/LSTM?** → LayerNorm (never BatchNorm!)
|
||
4. **Transformer?** → LayerNorm or RMSNorm with pre-norm
|
||
5. **Style transfer?** → InstanceNorm
|
||
6. **GAN?** → InstanceNorm (generator) or no norm (discriminator)
|
||
|
||
**Modern defaults:**
|
||
- CNNs: BatchNorm (batch ≥ 8) or GroupNorm (batch < 8)
|
||
- RNNs: LayerNorm
|
||
- Transformers: RMSNorm + pre-norm (large models) or LayerNorm + pre-norm (small models)
|
||
- GANs: InstanceNorm (generator), no norm (discriminator)
|
||
|
||
**Key insight**: Match normalization to architecture and batch size. Don't cargo-cult "add BatchNorm everywhere"—it fails for small batches, RNNs, Transformers, and style transfer.
|