Skip to content

← Flax step 18 of 100

Medium Primitives

Implement GroupNorm

Jump to the task

Split a vector's channels into groups, normalize each group, then apply a learned scale and bias per channel. This exercise isolates the grouping calculation using a one-dimensional input.

Group consecutive channels

For C channels and G groups, require G >= 1 and C % G == 0. Each group contains K = C // G consecutive channels. Reshape the vector from (C,) to (G, K), then compute each row's statistics:

grouped = x.reshape(G, K)
mu = jnp.mean(grouped, axis=-1, keepdims=True)
var = jnp.mean((grouped - mu) ** 2, axis=-1, keepdims=True)
normalized = (grouped - mu) / jnp.sqrt(var + eps)
output = gamma * normalized.reshape(C) + beta

Variance divides by K, not K - 1. gamma and beta each have shape (C,), with one value per channel, rather than one per group. Initialize gamma to ones and beta to zeros. These parameters leave the normalized values unchanged at initialization; the entire operation is not an identity.

Worked example

For x = [1, 2, 3, 4, 5, 6] and G=2, the groups are [1, 2, 3] and [4, 5, 6]. Their means are 2 and 5. Both have variance 2/3. With eps=1e-5, each group becomes approximately [-1.2247, 0, 1.2247]. Restore the original channel order to get [-1.2247, 0, 1.2247, -1.2247, 0, 1.2247].

A constant group becomes zero before the affine scale and bias. Positive epsilon keeps the denominator nonzero. With G=C, every group has one element, so every centered value is zero. With G=1, the vector calculation matches LayerNorm across all channels, using the same epsilon and parameters.

How this relates to image GroupNorm

Standard GroupNorm on images pools statistics across spatial positions as well as the channels within a group, independently for each example. There are no spatial positions in this exercise's vector input.

The implementation below supports leading dimensions by treating each leading position separately. On an image tensor, that extension would not be standard image GroupNorm: it would omit the spatial reduction. Likewise, the usual relationship between one-channel groups and InstanceNorm depends on normalizing over spatial dimensions. Do not infer image behavior from the single-element groups in this vector exercise.

GroupNorm uses no running batch statistics. Its usefulness compared with BatchNorm depends on the model and task; batch independence alone does not guarantee better accuracy.

Apply the affine parameters after normalization

Restore the channel axis before multiplying by gamma and adding beta. Applying a nontrivial scale or shift to the input before computing statistics generally changes the result. It is not an interchangeable ordering.

Problem

Implement MyGroupNorm(nn.Module) with num_groups and eps fields and a compact __call__:

  1. Read c = x.shape[-1] and assert c % self.num_groups == 0.
  2. Declare gamma and beta with self.param, shape (c,), using nn.initializers.ones and nn.initializers.zeros respectively.
  3. Reshape to (*x.shape[:-1], self.num_groups, c // self.num_groups).
  4. Normalize across the final axis using the mean and centered variance formula above, with keepdims=True and epsilon inside the square root.
  5. Restore the original shape, then return gamma * normalized + beta.

Implement my_groupnorm_forward(seed, x, num_groups, eps). Construct the module with num_groups=int(num_groups) and eps=float(eps), initialize using jax.random.PRNGKey(int(seed)) and x, then apply with the full variables mapping and return .reshape(-1).

x is a finite float32 vector of length C >= 1, with values small enough that squared centered values remain finite. The positive integer-valued num_groups divides C. Epsilon is finite, positive and representable in float32. The nonnegative integer-valued seed and group count may arrive as floats. Return a JAX vector of length C, leaving x unchanged. Constant initializers make this initial forward result independent of the seed.

Loading visualization…