Count the groups before counting the parameters

STATEMENT

Audit the axes, statistics and parameter shapes of batch and layer normalisation. An output tensor with the expected shape is not sufficient evidence that the intended axes were used.

GIVEN

First a dense matrix of shape $(N,D)=(2,3)$. Then a channels-first image tensor of shape $(N,C,H,W)=(2,3,4,4)$. Use learned affine parameters. Compare spatial batch normalisation, layer normalisation over all $(C,H,W)$, and channel-only layer normalisation performed separately at each spatial position.

FIND

The number of groups, entries per group, learned parameters and stored running-statistic scalars. Identify an implementation that silently normalises the wrong axis.

STRATEGY

Statistics have one value per reduction group. Learned affine parameters have one value per specified feature position and need not have the same shape as the statistics. Count mean and variance separately from gamma and beta.

SOLUTION

For the dense matrix:

Both return shape $(2,3)$. Their identical parameter counts hide different dependencies.

For the image tensor:

These layer-normalisation counts use an independent affine parameter for each position of the normalised shape, shared over the non-normalised axes. A deliberately tied affine parameterisation would have different counts.

The broadcast shapes of the spatial batch statistics and affine parameters are $(1,C,1,1)$. For full-example layer normalisation, the computed statistics have shape $(N,1,1,1)$, but gamma and beta each have shape $(C,H,W)$. For channel-only layer normalisation, move channels to the last axis, normalise that axis with shape $C$, then restore the original axis order.

Applying a last-axis normalisation of width $W$ directly to channels-first data instead computes a mean across each row of pixels. It returns the same overall tensor shape. Nothing about that successful return makes it channel normalisation.

Every entry participates in a constant number of reductions and elementwise operations, so forward and efficient backward arithmetic are $O(NCHW)$. The full dense Jacobian is unnecessary.

Check — sanity

Groups multiplied by entries per group must equal the tensor’s $96$ entries for each rule. Running-statistic buffers are not learned parameters, and a framework’s bookkeeping counter is separate from the mean and variance counted here.