replace BatchNorm with LayerNorm

[bkmi]

Is your feature request related to a problem? Please describe.
BatchNorm violates the iid (independent and identically distributed) assumption (see here, which is fundamental to many of the objective functions in this library. This can potentially cause issues when training models that rely on this assumption.

Describe the solution you'd like
To address this, we propose replacing BatchNorm with LayerNorm or GroupNorm wherever feasible. These alternatives do not interfere with the iid assumption and are better suited to the requirements of our library.

This needs to adapt all the "factories" (e.g. here:

• Update all net builders where this applies here
• Think about other nets in here e.g. embedding nets which would like this feature too.

Describe alternatives you've considered
One alternative could be to allow users to choose between BatchNorm and LayerNorm through a configurable flag. However, this seems like a niche use case, and BatchNorm’s reliance on iid assumptions doesn’t offer a compelling reason to keep it, aside from legacy compatibility.

Additional context
Many classifiers provide an option like use_batch_norm: bool = False to control BatchNorm usage. We suggest updating this option to use_layer_norm to align with the new approach.

[bkmi]

@janfb wanted to see a side-by-side comparison of the two with LayerNorm and BatchNorm before making the change.

[francois-rozet]

To comment on that issue, with BatchNorm one should be very careful to always feed batches that represent "modes" in the same proportion at the original distribution. For example, for a binary classifier, it is invalid to evaluate the positive and negative batches separately during training as it therefore becomes enough to identify the class of a single element in the batch to decide the class of the entire batch.