Normalization Techniques

Why Normalize?

Imagine a factory assembly line where each station expects parts within a certain size range. If Station 3 suddenly starts outputting parts twice as big, Station 4’s machinery jams. Someone has to stop the line, recalibrate, and restart. Now imagine this recalibration happening after every batch of parts. That is what deep networks face without normalization: each layer’s input distribution shifts as the layers before it update their weights, and every layer has to constantly re-adapt instead of learning useful features. Deep networks suffer from internal covariate shift — the distribution of inputs to each layer changes during training as preceding weights update. This causes:

Slower training (need smaller learning rates to avoid instability)
Difficulty with saturating activations (inputs drift into flat regions of sigmoid/tanh)
Careful initialization requirements (poor init = immediate failure)

Normalization stabilizes these distributions by re-centering and re-scaling activations at each layer, enabling:

Higher learning rates (often 10x larger)
Faster convergence (typically 2-3x fewer epochs)
Reduced sensitivity to initialization (networks “just work” with standard init)

The honest truth about “internal covariate shift”: The original BatchNorm paper (Ioffe and Szegedy, 2015) attributed its success to reducing internal covariate shift. Later research (Santurkar et al., 2018) showed this explanation is incomplete — BatchNorm primarily works by smoothing the loss landscape, making it easier for optimizers to navigate. The mechanism matters less than the result: normalization is one of the most impactful techniques in modern deep learning.

Batch Normalization

The most influential normalization technique, introduced in 2015. The core idea: for each feature channel, compute the mean and variance across all examples in the current mini-batch, then normalize so the channel has zero mean and unit variance. Finally, apply a learnable scale (

\gamma

) and shift (

\beta

) so the network can undo the normalization if that is optimal. Why the learnable parameters? Without

\gamma

and

\beta

, normalization would force every layer’s output to have zero mean and unit variance, which might be too restrictive. The learnable parameters let the network decide the “ideal” mean and variance for each channel. If the network learns

\gamma = \sigma_B

and

\beta = \mu_B

, it has effectively undone the normalization — so BatchNorm can never hurt (in theory). Normalize across the batch dimension:

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

Where

\mu_B

and

\sigma_B

are batch statistics, and

\gamma

\beta

are learnable.

import torch
import torch.nn as nn

class BatchNorm(nn.Module):
    """Batch Normalization from scratch."""
    
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.eps = eps
        self.momentum = momentum
        
        # Learnable parameters
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        
        # Running statistics (not learnable)
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
    
    def forward(self, x):
        if self.training:
            # Compute batch statistics
            mean = x.mean(dim=(0, 2, 3))  # Mean over batch, H, W
            var = x.var(dim=(0, 2, 3), unbiased=False)
            
            # Update running statistics
            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
        else:
            mean = self.running_mean
            var = self.running_var
        
        # Normalize
        x = (x - mean.view(1, -1, 1, 1)) / torch.sqrt(var.view(1, -1, 1, 1) + self.eps)
        
        # Scale and shift
        return self.gamma.view(1, -1, 1, 1) * x + self.beta.view(1, -1, 1, 1)

BatchNorm behaves differently at train vs eval! During training, it uses the current batch’s statistics. During inference, it uses the running averages accumulated during training. Forgetting model.eval() before inference is one of the most common bugs in deep learning — your model will produce inconsistent outputs that depend on what other examples happen to be in the inference batch. Always call model.eval() before inference and model.train() before training resumes.

Pitfall — small batch sizes: BatchNorm estimates mean and variance from the current mini-batch. With a batch size of 2 or 4, these estimates are extremely noisy, which destabilizes training. If your GPU memory limits you to small batches (common with high-resolution images or large models), switch to GroupNorm or LayerNorm instead. A good rule of thumb: BatchNorm works well with batch sizes of 32 or larger; below 16, consider alternatives.

Layer Normalization

The critical difference from BatchNorm: LayerNorm normalizes across the feature dimension within a single sample, completely independent of other examples in the batch. Think of it this way — BatchNorm asks “how does this feature compare across all images in the batch?” while LayerNorm asks “how does this feature compare to all other features within this one image/token?” This independence from batch size is why LayerNorm became the standard for Transformers and RNNs, where batch sizes vary and sequences have different lengths.

\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta

Where

\mu

and

\sigma

are computed per sample across all features.

class LayerNorm(nn.Module):
    """Layer Normalization -- the default choice for Transformers.
    
    Unlike BatchNorm, statistics are computed per-sample across
    all features, so behavior is identical at train and eval time.
    No running statistics to maintain, no batch-size sensitivity.
    """
    
    def __init__(self, normalized_shape, eps=1e-5):
        super().__init__()
        self.eps = eps
        # Learnable scale and shift, same role as in BatchNorm
        self.gamma = nn.Parameter(torch.ones(normalized_shape))
        self.beta = nn.Parameter(torch.zeros(normalized_shape))
    
    def forward(self, x):
        # Compute mean and variance across the last dimension (features)
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        return self.gamma * (x - mean) / torch.sqrt(var + self.eps) + self.beta

Used in: Transformers, RNNs, and any architecture where batch-size independence matters

Pre-norm vs post-norm placement matters a lot. In Transformers, placing LayerNorm before the attention/FFN block (pre-norm) generally trains more stably than the original post-norm placement. Most modern architectures (GPT-2 onward, LLaMA, etc.) use pre-norm. The reason: pre-norm ensures the residual stream stays well-conditioned, preventing gradient explosion in deep transformer stacks.

Comparison of Normalization Types

Type	Normalize Over	Best For
Batch Norm	Batch, H, W	CNNs with large batches
Layer Norm	C, H, W (per sample)	Transformers, RNNs
Instance Norm	H, W (per channel)	Style transfer
Group Norm	Groups of channels	Small batches
RMSNorm	Features (no mean)	LLMs (faster)

RMSNorm (Modern LLMs)

RMSNorm is a simplified variant of LayerNorm that skips the mean-centering step entirely and only divides by the root-mean-square of the activations. This saves both compute and memory (no mean calculation, no beta parameter). The empirical finding is surprising: the re-centering step in LayerNorm contributes very little to training stability; most of the benefit comes from the re-scaling alone. This is why LLaMA, Gemma, Mistral, and most modern LLMs use RMSNorm — it is approximately 10-15% faster than LayerNorm at scale, with no measurable loss in quality.

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

class RMSNorm(nn.Module):
    """Root Mean Square Normalization (used in LLaMA, Gemma, Mistral, etc.).
    
    Compared to LayerNorm:
    - No mean subtraction (no re-centering)
    - No learnable bias (beta parameter)
    - ~10-15% faster at scale with equivalent quality
    """
    
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        # Only a scale parameter, no bias -- this is the key simplification
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        # Root mean square: sqrt(mean(x^2))
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

When to Use What

Scenario	Recommendation
CNN with batch ≥ 32	Batch Norm
CNN with small batch	Group Norm
Transformer/Attention	Layer Norm or RMSNorm
RNN/LSTM	Layer Norm
Style Transfer	Instance Norm
Modern LLM	RMSNorm

Exercises

Exercise 1: BatchNorm Analysis

Train the same CNN with and without BatchNorm. Compare learning curves, final accuracy, and sensitivity to learning rate.

Exercise 2: Small Batch Study

Compare BatchNorm vs GroupNorm with batch sizes of 2, 4, 8, 16, 32.

Exercise 3: Pre-Norm vs Post-Norm

In transformers, compare placing LayerNorm before vs after attention/FFN blocks.

What’s Next

Module 17: Regularization for Deep Networks

Dropout, weight decay, and other techniques for preventing overfitting.

​Normalization Techniques

​Why Normalize?

​Batch Normalization

​Layer Normalization

​Comparison of Normalization Types

​RMSNorm (Modern LLMs)

​When to Use What

​Exercises

​What’s Next

Module 17: Regularization for Deep Networks

Normalization Techniques

Why Normalize?

Batch Normalization

Layer Normalization

Comparison of Normalization Types

RMSNorm (Modern LLMs)

When to Use What

Exercises

What’s Next