NNX Implement GroupNorm

Normalize groups of features

Instead of normalizing a whole vector at once, split its C features into G equal, contiguous groups. Compute a separate mean and variance for each group, normalize its values, then apply a learned gain and offset to each original feature.

This exercise uses vectors, with features on the last axis. If there are leading dimensions, each vector is handled independently. There are no running statistics and no reduction across examples in a batch.

Which dimensions are included?

For a vector, reshape (C,) to (G, C // G) and reduce over the last axis. For several vectors shaped (..., C), reshape to (..., G, C // G) and still reduce only over the last axis.

This is the vector case of grouped normalization. Standard image GroupNorm also includes spatial positions in each group's statistics. For channels-last images (N, H, W, C), the grouped shape is (N, H, W, G, C // G) and the reduction includes axes (1, 2, 4). Simply passing images through this exercise's last-axis-only function would normalize each pixel separately, not perform image GroupNorm.

Two useful limits for this exercise are:

• G = 1: normalize the whole feature vector. This matches last-axis LayerNorm with the same epsilon, gain and offset.
• G = C: each group contains a single number. Its centered value and variance are zero, so the output is just the learned offset. Image InstanceNorm usually includes multiple spatial positions per channel and need not produce zero. Do not equate it with this vector case.

The GroupNorm paper studies its use in vision models, including small-batch training. Independence from batch statistics is useful, but there is no universal batch-size threshold at which BatchNorm fails or GroupNorm becomes the best choice.

A complete example

import jax.numpy as jnp
from flax import nnx

class MyGroupNorm(nnx.Module):
    def __init__(self, num_features, num_groups, eps, rngs):
        assert num_groups > 0 and num_features % num_groups == 0
        self.gamma = nnx.Param(jnp.ones((num_features,), dtype=jnp.float32))
        self.beta = nnx.Param(jnp.zeros((num_features,), dtype=jnp.float32))
        self.num_groups = num_groups
        self.num_features = num_features
        self.eps = eps

    def __call__(self, x):
        grouped = x.reshape(*x.shape[:-1], self.num_groups,
                            self.num_features // self.num_groups)
        mean = jnp.mean(grouped, axis=-1, keepdims=True)
        variance = jnp.mean((grouped - mean) ** 2, axis=-1, keepdims=True)
        normalized = (grouped - mean) / jnp.sqrt(variance + self.eps)
        return self.gamma[...] * normalized.reshape(x.shape) + self.beta[...]

model = MyGroupNorm(4, 2, 3.0, nnx.Rngs(0))
x = jnp.array([1., 3., 10., 10.], dtype=jnp.float32)
print(model(x))  # [-0.5, 0.5, 0.0, 0.0]
model.gamma[...] = jnp.array([2., 4., 1., 1.], dtype=jnp.float32)
model.beta[...] = jnp.array([1., -1., 3., 4.], dtype=jnp.float32)
print(model(x))  # [0.0, 1.0, 3.0, 4.0]

The first group [1, 3] has mean 2 and population variance 1. The chosen epsilon makes its denominator sqrt(1 + 3) = 2. The second group is constant, so it becomes zero before the affine transform.

Gain and offset each have shape (C,), not (G,): features in the same group share statistics but can learn different affine transforms. Ones and zeros need no random key, so this constructor accepts rngs without using it. The forward pass changes no parameters.

Details that affect the result

Compute population variance by taking the mean of squared centered values. Divide by the number of features in the group, not one less. Add positive epsilon inside the square root. Before the affine transform, each group's variance is var / (var + eps) in exact arithmetic, not exactly one. Constant groups remain well-defined.

Keep the original feature order: group 0 contains the first C // G features. Splitting into (C // G, G) would form different groups.

The visualization shows the broader image convention: each cell stands for a channel block that can contain spatial positions. For this exercise, imagine each cell contains just one scalar. The channel grouping is the same, but there are no spatial positions to include in the reduction.