Normalization (BatchNorm, LayerNorm, GroupNorm) and Pooling

Normalization is what made deep networks trainable

Normalization layers stabilize training by re-centering and re-scaling activations. Three common variants — and they're not interchangeable; pick by what shape you're working with:

BatchNorm2d — normalizes across the batch for each channel. Standard in CNNs. Needs reasonably large batch sizes (typically 16+) to estimate stats well. Has a behavior switch on train() vs eval(): training uses batch stats, eval uses running mean / running var (the buffers).
LayerNorm — normalizes per sample, across the feature dim. Standard in Transformers because batch size doesn't matter and inference behavior is stable.
GroupNorm — splits channels into groups and normalizes within each group. Works at any batch size including 1. Used in object detection (COCO, where batch=2 per GPU is common) and some vision Transformers.

Why not just BatchNorm everywhere?

BatchNorm leaks information across samples in a batch (each sample's normalization depends on the others), which is bad for some setups (RNN steps, multi-task learning, contrastive learning where you don't want batch coupling). It's also fragile at small batch size and at inference time on a single sample. LayerNorm sidesteps all of this — that's why Transformers settled on it.

Pooling

MaxPool2d picks the max in each window, AvgPool2d averages. AdaptiveAvgPool2d(out_size) is the modern hero — it figures out the right window size to produce the requested output. Set AdaptiveAvgPool2d(1) at the end of a CNN backbone to get a single feature vector regardless of input image size.

Code

BatchNorm2d — for CNNs·python

import torch
import torch.nn as nn

bn = nn.BatchNorm2d(64)             # normalizes per-channel across batch
x = torch.randn(32, 64, 16, 16)     # batch=32

bn.train()
out_train = bn(x)                    # uses batch stats
print(bn.running_mean.shape)         # torch.Size([64]) — buffer

bn.eval()
out_eval = bn(x)                     # uses running stats

LayerNorm — for Transformers·python

import torch
import torch.nn as nn

# Normalize the LAST dim (features), per-sample
ln = nn.LayerNorm(512)
x = torch.randn(32, 16, 512)         # batch=32, seq=16, features=512
out = ln(x)
print(out.mean(dim=-1).abs().max())  # ~0  (zero mean per-sample-per-position)
print(out.std(dim=-1).mean())        # ~1

GroupNorm — small batch, no batch coupling·python

import torch
import torch.nn as nn

# 64 channels split into 8 groups of 8 channels each
gn = nn.GroupNorm(num_groups=8, num_channels=64)
x = torch.randn(2, 64, 16, 16)       # batch=2 — too small for BatchNorm
out = gn(x)
print(out.shape)                      # torch.Size([2, 64, 16, 16])

Pooling — and the AdaptiveAvgPool trick·python

import torch
import torch.nn as nn

x = torch.randn(8, 64, 32, 32)

mp = nn.MaxPool2d(2)                  # halves spatial
print(mp(x).shape)                     # torch.Size([8, 64, 16, 16])

ap = nn.AvgPool2d(4)
print(ap(x).shape)                     # torch.Size([8, 64, 8, 8])

# AdaptiveAvgPool — outputs a fixed spatial size regardless of input
gap = nn.AdaptiveAvgPool2d(1)
print(gap(x).shape)                    # torch.Size([8, 64, 1, 1])

# Same module on a different input size still produces 1x1
y = torch.randn(8, 64, 224, 224)
print(gap(y).shape)                    # torch.Size([8, 64, 1, 1])

Normalization (BatchNorm, LayerNorm, GroupNorm) and Pooling

Normalization is what made deep networks trainable

Why not just BatchNorm everywhere?

Pooling

Code

External links

Exercise

Progress

Comments 0