Group Normalization - PyTorch Basics

Group Normalization

What you'll build: group_norm(x, num_groups, gamma, beta, eps) for 2D input (N, C): split the channels into groups, normalise within each group of each sample, then scale and shift per channel.

Background

BatchNorm's statistics come from the batch, so it degrades when batches are small (detection, segmentation, large models on few GPUs). Group Normalization (Wu & He, 2018) computes its statistics per sample over groups of channels, so it behaves identically for any batch size and in train and eval mode. It is the default in diffusion U-Nets.

Algorithm

With G = num_groups dividing C, and C // G channels per group:

1. Reshape x from (N, C) to (N, G, C // G), so each group of consecutive channels sits on its own row.
2. Compute the mean and the population variance (divide by the group size) over the last axis.
3. Normalise: x_hat = (x - mean) / sqrt(var + eps).
4. Reshape back to (N, C).
5. Scale and shift per channel: out = gamma * x_hat + beta.

Special Cases

Setting	Equivalent to
G = 1	LayerNorm over the channels
G = C	InstanceNorm (with a single value per channel, every output is 0, then beta)

Inputs

• x: float array (N, C).
• num_groups: integer that divides C.
• gamma, beta: float arrays (C,).
• eps: small positive float (default 1e-5).

Expected Return Format

An array of shape (N, C).

Edge Cases and Conventions

• Groups are consecutive channels: with C = 4, G = 2 the groups are channels {0, 1} and {2, 3}.
• The real layer works on (N, C, *spatial) and also pools over the spatial positions; here the input is (N, C) to keep the reshape simple.

Constraints

• Vectorised with a reshape; no loops over groups.
• PyTorch track: don't call F.group_norm or nn.GroupNorm.

Follow-up

• Extend it to (N, C, H, W). Which axes do you reduce over?
• Why is GroupNorm usually paired with weight standardisation in small-batch detection models?

Input: group_norm(np.array([[1.0, 2.0, 3.0, 5.0]]), 2, np.ones(4), np.zeros(4))
Output: array([[-1.,  1., -1.,  1.]])

Input: group_norm(np.array([[1.0, 2.0, 3.0, 5.0]]), 2, np.full(4, 2.0), np.ones(4))
Output: array([[-1.,  3., -1.,  3.]])

import numpy as np

def group_norm(x: np.ndarray, num_groups: int, gamma: np.ndarray, beta: np.ndarray, eps: float = 1e-5) -> np.ndarray:
    """
    Apply Group Normalization to (N, C) input.

    Args:
        x: Input array of shape (N, C)
        num_groups: Number of groups of consecutive channels (divides C)
        gamma: Scale parameter of shape (C,)
        beta: Shift parameter of shape (C,)
        eps: Small constant for numerical stability

    Returns:
        Normalized array of shape (N, C)
    """
    # Your code here
    pass