We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
← Normalization step 3 of 6
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…
Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.
Decide which values share a mean and variance
Many normalization methods subtract a mean and divide by a standard deviation:
$$\hat x=\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}.$$
The small positive epsilon keeps the denominator away from zero. The main
difference between methods is which values are used to compute the mean and
variance.
- Batch normalization usually computes one set of statistics per channel, combining examples and spatial positions during training. Its usual inference behavior uses running statistics collected during training.
- Layer normalization computes statistics over selected feature dimensions within each example. Different examples do not share statistics.
- Group normalization divides channels into groups. Within each example, it computes statistics over each group's channels and spatial positions.
For the feature vector [1, 3], the mean is 2 and population variance is 1.
Subtracting the mean gives [-1, 1]; dividing by sqrt(1 + epsilon) leaves values
close to [-1, 1]. Computing statistics over the batch axis instead would combine
values from different examples and produce a different operation.
RMSNorm rescales without centering
RMSNorm divides by the root mean square, with epsilon added for stability:
$$\hat x=\frac{x}{\sqrt{\operatorname{mean}(x^2)+\epsilon}}.$$
It does not subtract the mean, so its output need not have mean zero. For example,
[1, 1] stays close to [1, 1] rather than becoming [0, 0].
Normalization layers may apply learnable scale and bias after normalizing. Their shapes must broadcast over the intended axes. These extra parameters mean that the final output need not have mean zero or variance one.
Match the formula and the source of the statistics
Check whether epsilon goes inside or outside the square root. Also check whether the variance divides by the number of values (population variance) or one fewer (sample variance). These choices change the answer.
For feature preprocessing, estimate statistics on training data and reuse them on validation and test data. Using test data to estimate them leaks information into training. An individual exercise may instead ask you to use only the input it supplies. Follow that contract and handle constant features as specified.
Your code, saved here
Sign in to access your accepted solutions and saved draft history. You can keep reading the lesson without an account.
Sign in to your accountStuck?
JAX reference solution
Sign in to track your attempts toward unlocking the reference solution.