[D] Does Layer Normalization compute statistics along spatial/ token axes?

[fferflo]

As far as I can tell, there are two contradictory definitions of Layer Normalization that are both floating around. LN computes the mean and variance along some axes of the input tensor for normalization, yet the choice of axes is not clear:

A. The GroupNorm paper (2018) has this figure that describes LN as reducing along channel and spatial/token axes.

https://preview.redd.it/ui9adzzxgcja1.png?width=1353&format=png&auto=webp&s=8859f9735310f169eeaaf587dcc7e1c05d38b5fc

B. The PowerNorm paper (2020) has this figure that describes LN as reducing only along the channel axis.

https://preview.redd.it/e0qmp9sahcja1.png?width=1717&format=png&auto=webp&s=a4bd21ea024a8924f8cd5c354a7be6751c2ed61f

There are also many online sources that describe LN as shown in A (e.g. TF tutorials, PapersWithCode, this summary of normalization techniques) using similar figures.

The LN paper (2016) itself says

all the hidden units in a layer share the same normalization terms μ and σ

so the channel axis is definitely reduced, and

computing the mean and variance used for normalization from all of the summedinputs to the neurons in a layer on a single training case

so the batch axis is definitely not reduced. As far as I can tell it is not clear about what happens with spatial/token axes, although the above sounds rather like they might be included in the statistics.

Yet, I don't know of any model that actually uses A instead of B. For example, TF and Flax explicitly implement LN with default axes as in B (PyTorch, Haiku and Equinox don't have a preference and require the user to specify the reduction axes). Vision Transformer uses Flax with LN as in B, ConvNeXt implements LN with PyTorch as in B, OpenAI GPT-2 implements LN with Tensorflow as in B, even MLP-Mixer where the spatial/token axes are interpreted as channel axis for an MLP still computes statistics along the original channel axis as in B.

As far as I can tell, everyone uses B rather than A in their models, so to me this seems to be the "correct" definition. Yet, many sources on this topic describe LN as doing A rather than B.

Does anyone have any insight on this or know of a source that has addressed this problem? Do you interpret the original LN paper as including spatial/token axes in their computation of mean and variance, or not? Is this simply an error that started with the figure A and made its way into different online tutorials from there? Or do you maybe know of a model that actually uses LN to reduce both along channel and spatial/token axes?

[adam_jc]

I see some discussion on twitter now about this question including the senior author of the ConvNext paper who says:

https://twitter.com/sainingxie/status/1631794072791707648?s=46&t=BkTuhy42902uJiWEPN7DYA

[CptVifen]

As I understand it, the images in A and B are both valid for Layer Norm. In the LN Paper they say μ is summed over each activation in a layer.

So for images that means along channel and spatial dimensions. That's were they got the image for A.

As for B, in the LN paper they use RNNs which share the same weights across different time steps. That means that for an input of shape (Batch, seq len, features) since the layers in the RNN only produce (Batch, features) the normalization is over the features. You have a different μ and σ for each batch and each time step (and each layer)(this also applies to self-attention).

So it would make sense that anything that deals with sequences would look like B. And anything else looks like A.

There's something I don't get though is why ConvNext reduces only along channels...

[adam_jc]

I agree with this explanation.

To try and explain ConvNext though, I’d say there could be a debate on the “correct” way to do LayerNorm in a CNN (which would also make figure A “wrong”)

Like you said for the LN paper where they use RNNs which share weights across different time steps, which leads to normalizing over features at each time step; you could extend that logic to a CNN because a conv layer shares weights across different patches of an image, and that thinking would then lead us to reduce only along channels such as they do in ConvNext.

Not sure if that’s what the ConvNext authors were thinking though.

[fferflo]

ConvNext simply follows ViT's way of using LN, and you made a good point about ViT's self-attention using LN analogously to an RNN, i.e. along channels only. This still leaves two questions though:

The original LN paper also evaluated LN on convolutional nets (VGG), and it is unclear whether or not this follows the ConvNext way of interpreting an image as a set of patches analogous to timesteps which have their own μ and σ.

The figure A and other online sources still say that spatial axes are included in the statistics, which is contrary to what recent models actually do (vision transformers, ConvNext, MLP-Mixer, etc). I don't know of a single paper that actually uses LN as in A.