I implemented BatchNorm, LayerNorm, and GroupNorm from scratch to see what they actually do to neuron activations [D]

I kept seeing BatchNorm/LayerNorm in model definitions without really understanding the assumptions behind each one, so I implemented all three (batch, layer, group) from scratch and ran them side by side on a simple 3-layer MLP on MNIST — comparing training speed, test accuracy, and neuron activations.

A few things that stuck with me:

Dead neurons are visible. If you plot post-activation values (neuron × batch sample) for the vanilla MLP, whole neuron rows stay white — dead across the entire batch. With any normalization, the same plot turns into a mixed texture: each neuron fires for some samples and not others, i.e. its output is actually input-dependent. That contrast was more stark than I expected.

All three help, and roughly equally here. Test accuracy: vanilla 84.1% → BatchNorm 96.6% / LayerNorm 95.4% / GroupNorm 96.3%. No meaningful gap between the three normalizations on this task — which makes sense: there's no conv feature extraction, so GroupNorm has no structure to exploit over LayerNorm, and the batch is large enough that BatchNorm's assumption holds fine.

The geometric framing helped me most. LayerNorm projects each sample onto the subspace where features sum to zero, then fixes the norm — so a d-dim vector loses 2 degrees of freedom (in 3-D, the normalized points literally live on a circle). GroupNorm generalizes this to d−2g. Framing them as "which degrees of freedom do we declare redundant" made the inductive biases click for me.

Full code and plots here.

Question for people who use these in anger: in practice, where does GroupNorm's per-group assumption actually earn its keep over LayerNorm outside of small-batch vision? Curious what real workloads make the difference show up.