Group Normalisation

BatchNorm fails at small batch size, and several things in vision force small batches — detection and segmentation, which need high-resolution inputs, and diffusion training, where the model is large and the images are not small. GroupNorm is the answer that vision arrived at, and it is a within-example scheme like everything that works on sequences.

with $\sigma$ over the same index set, and $\mathcal{G}_g$ one of $G$ contiguous blocks of channels. The batch index $b$ appears on the left, so each example is normalised alone: no running statistics, no train/eval divergence, no dependence on what else is in the batch.

One knob, three known schemes

The group count recovers the whole family.

Both endpoints are worse than the middle, and the reason is a fact about what convolutional channels are. LayerNorm at $G = 1$ assumes every channel should share a scale, which is wrong — an edge detector and a colour channel have no business being normalised together. InstanceNorm at $G = C$ removes each channel’s magnitude entirely, which discards contrast information the classifier wants. Grouping keeps related filters together and lets unrelated ones differ.

Where it is actually used now

Diffusion. Every UNet in the DDPM lineage uses GroupNorm at 32 groups, in each residual block, before the activation — a detail inherited essentially unchanged from Ho et al. and reproduced in Stable Diffusion and its descendants. It is there because diffusion training runs at batch sizes that would make BatchNorm’s statistics useless, and because the sampler runs at batch size 1.

Transformer-based diffusion (DiT and after) drops it for LayerNorm, which is the same thing at $G = 1$ — the channel-grouping argument does not apply once the tensor is a sequence of tokens rather than a stack of feature maps.

import torch
from torch import Tensor, nn


class GroupNorm(nn.Module):
    def __init__(self, groups: int, channels: int, eps: float = 1e-5):
        super().__init__()
        if channels % groups:
            raise ValueError(f"{channels} channels does not divide into {groups} groups")
        self.groups, self.eps = groups, eps
        self.weight = nn.Parameter(torch.ones(channels))
        self.bias = nn.Parameter(torch.zeros(channels))

    def forward(self, x: Tensor) -> Tensor:                   # [B, C, H, W]
        b, c, h, w = x.shape
        x = x.reshape(b, self.groups, -1)                     # fold C/G, H, W together

        mu = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True, unbiased=False)
        x = (x - mu) * torch.rsqrt(var + self.eps)

        x = x.reshape(b, c, h, w)
        return x * self.weight[:, None, None] + self.bias[:, None, None]

The reshape assumes groups are contiguous in channel order, which means the grouping is decided by however the preceding convolution happened to order its output filters — arbitrary, and it does not matter, because the network learns an ordering that suits the grouping rather than the other way round.