Skip to content

← Normalization step 3 of 6

Medium Primitives

Implement Group Norm

Group normalization divides an example's channels into groups and normalizes each group separately. Its statistics use channels and spatial positions within one example, never values from other batch examples.

For a tensor shaped (N, C, H, W), N is the number of examples, C is the number of channels, and H and W are spatial dimensions. With C=4 and two groups, channels 0 and 1 belong together, and channels 2 and 3 belong together. Each group's statistics include all H-by-W positions in those channels.

For one example with H=W=1 and channel values [1, 3, 10, 14], the two groups have means 2 and 12, and population variances 1 and 4. With epsilon 0.01, the normalized values are approximately [-0.995037, 0.995037, -0.998752, 0.998752].

Your task

Implement group_norm(x, num_groups, eps=1e-5) for a floating-point tensor of shape (N, C, H, W). All dimensions are positive. num_groups is a positive integer dividing C, and eps is positive.

Split the channels into consecutive, equal-sized groups. For each example and group, compute the mean and population variance over its (C / num_groups) * H * W values. Normalize those values as (x - mean) / sqrt(variance + eps) and return the original shape.

Do not divide squared deviations by one fewer than the group size. Do not apply gamma, beta, or any other learned scale or shift. Constant groups, including a group containing one value, return zeros.

A useful layout is (N, G, C // G, H, W), with G = num_groups. Reduce its last three axes while retaining those axes for broadcasting, then reshape the normalized result back to (N, C, H, W).

With one group, the normalization statistics match LayerNorm over all (C, H, W) axes. With C groups, they match per-example, per-channel InstanceNorm over (H, W). These comparisons concern the normalization formula with matching epsilon and no affine parameters, not every behavior of the corresponding library modules.

The diagram shows which channel blocks share statistics. Each block includes all spatial positions; matching group numbers share a mean and variance. Switching the method lets you compare the groupings.

Loading visualization…