BatchNorm and LayerNorm

Why normalize at all

As gradients flow through deep networks, the distribution of activations at each layer drifts (Internal Covariate Shift, in the original 2015 framing). Normalization layers force each layer's input to have a stable mean and variance, which keeps gradients well-conditioned and lets you train deeper, with higher learning rates, more reliably.

BatchNorm

Normalizes across the batch dimension: for each feature, subtract the batch mean and divide by the batch standard deviation. Add learned scale and shift per feature. Used in CNNs (especially ResNet-style). Failure mode: small batch sizes (batch=1, batch=8) give unreliable statistics; behaves differently in train vs eval (running statistics for eval), which is a famous source of bugs.

LayerNorm

Normalizes across the feature dimension within each example: for each example, subtract the per-example mean and divide by the per-example standard deviation. Used in transformers (because batch dimension is awkward for variable-length sequences). No train/eval distinction — same behavior either way.

Tip: If you can't decide which to use: BatchNorm for vision CNNs (large batches, fixed shapes), LayerNorm for transformers (small or variable batch shapes). RMSNorm (a stripped-down LayerNorm) is the modern variant in LLaMA and other current LLMs.

Where to put it

Two common patterns: Pre-norm (norm before the layer): x = x + sublayer(LayerNorm(x)). Trains more stably for very deep transformers. Post-norm (norm after the layer): x = LayerNorm(x + sublayer(x)). Original transformer convention. Modern transformers default to pre-norm.

Principle: Normalization is what makes deep networks trainable in 2026. Choosing the wrong kind (BN in a transformer, LN in a small CNN) usually still works but loses you 10-20% of accuracy or stability. Match the type to the architecture.

Code

BatchNorm vs LayerNorm in code·python

import torch, torch.nn as nn

# CNN with BatchNorm
cnn = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.BatchNorm2d(64),       # normalize across batch+spatial, per channel
    nn.ReLU(),
    nn.Conv2d(64, 128, 3, padding=1),
    nn.BatchNorm2d(128),
    nn.ReLU(),
)

# Transformer block with LayerNorm
class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn  = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn   = nn.Sequential(
            nn.Linear(d_model, 4 * d_model), nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )
    def forward(self, x):
        # Pre-norm pattern
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x), need_weights=False)[0]
        x = x + self.ffn(self.norm2(x))
        return x

Why normalize at all

BatchNorm

LayerNorm

Where to put it

Code

External links

Exercise

Progress

Comments 0