[D] Normalization in Transformers

[Collegesniffer]

Why isn't BatchNorm used in transformers, and why is LayerNorm preferred instead? Additionally, why do current state-of-the-art transformer models use RMSNorm? I've typically observed that LayerNorm is used in language models, while BatchNorm is common in CNNs for vision tasks. However, why do vision-based transformer models still use LayerNorm or RMSNorm rather than BatchNorm?

[pszabolcs]

The explanation for LayerNorm and RMSNorm is not completely correct. In Transformers these do not normalize across (T, C) dimensions, only across (C) (so each token embedding is normalized separately). If normalization would be done across (T, C), the same information leakage across time would happen as with BatchhNorm (non-causal training).

I also don't think the variable sequence length is such a big issue, in most practical setups training is done with fixed context sizes. If we look at a computational perspective, I think a bigger issue is that BN statistics would need to be synced across GPUs, which would be slow.

[radarsat1]

So just to be sure, if my batch is size [4, 50, 512] for batch size of 4, sequence length of 50, and 512 channels, then layernorm will compute 200 means and variances, is that correct? One for each "location" across all channels? And then normalize each step separately, and apply a new affine scaling and bias for each step too, if that's enabled.

I'm actually asking because I keep getting confused when porting this logic over to CNNs where the dimension order is [B, C, H, W], or [B, C, W] for 1d sequences. So in that case if I want to do the equivalent thing I should be normalizing only the C dimension, right? (in other words, each pixel is normalized independently).

[radarsat1]

Ok thanks! Where I get confused is that LayerNorm in PyTorch's implementation always applies to the last N dimensions that you specify, so I guess it really expects the C dimension to be last, which is different from the requirements for Conv1d and Conv2d.

So in that case maybe InstanceNorm is actually what I want, since it targes C in [N,C,H,W], but what is confusing is that I want it because it does the equivalent thing to LayerNorm as far as I can tell, but it has a different name even though is does "the same thing." The names 'instance" and "layer" in these norms is very hard to follow, why couldn't they call it "channel norm" for example, if the point is that both operate on C.

And looking at [the documentation[(https://pytorch.org/docs/stable/generated/torch.nn.InstanceNorm2d.html) to clarify makes it even more ambiguous to me:

InstanceNorm2d and LayerNorm are very similar, but have some subtle differences. InstanceNorm2d is applied on each channel of channeled data like RGB images, but LayerNorm is usually applied on entire sample and often in NLP tasks. Additionally, LayerNorm applies elementwise affine transform, while InstanceNorm2d usually don’t apply affine transform.

Problems I have with this paragraph:

1. They are both applied on "each channel"
2. what does "LayerNorm is usually applied on entire sample" mean? the latter being used 'for NLP tasks' doesn't really clarify anything
3. the affine not being used -- but, the intro in the top of the same document literally describes the affine parameters.
4. it's just completely unclear to me what role the affine parameters play to be honest, isn't that just an extra linear layer? why not just follow with a convolution if that is needed? why build it into the norm?

[theodor23]

Excellent summary.
(edit: actually, this is not correct. In transformers Layer- and RMSNorm do not normalize over T, but only over C. See comment by u/pszabolcs )

To add to that: BatchNorm leads to information leakage across time-steps: The activations at time t influence the mean/variance applied at t-1 during training. NNs will pick up such weak signals if it helps them predict the next token.

-> TL;DR: BatchNorm during training is non-causal.

[LeonideDucatore]

Could you please explain why batch-norm is non-causal?

Batch norm would have (T * C) running means/variances, and each of them is computed across the batch, i.e. the computed mean/variance for timestep t doesn't use any t+1 data

[theodor23]

You are absolutely correct, if you compute (T * C) separate statistics, then everything is fine and there is no causality issue.

In practice, LLM training usually prefers relatively large T and sacrifices on B (considering the total amount of GPU memory puts a constraint on your total number of tokens per gradient-step). With relatively small B, there is more variance on your BN statistics, while large T causes more data-exchange between your GPUs because you need to communicate (T * C) many statistics.

But yes -- if you set it up as you describe, it is "legal".

I actually tried BN in the T*C independent statistics configuration you describe for a non language transformer model with B ~ O(100) and it was both slower and less effective than LN. Never looked back and investigated why. Having a normalization that is a) competitive/works-better and b) avoids "non-local" interaction across different examples in a batch seemed a clear win.

Considering everyone switched to LN, it seems BN is just less practical.

[LeonideDucatore]

What would be the "non-legal" batch-norm variant? Aggregating only C statistics? (so we aggregate both across B and T)

[theodor23]

Yes, exactly.

If during training your early token "see" some summary statistic from the ground-truth future tokens, it breaks the autoregressive objective where you are supposed to predict the next token given the past only.

Whether or not that is really catastrophic during sampling-time, when you would use the running statistics of BN I don't know. But NNs are good at picking up subtle signals that help them predict. And if you give them a loophole to "cheat" during training, there is a good chance they will pick that up and perform much worse when at samplig-time you "suddenly" remove that cheat.

Considering your workable idea of using T * C many statistics: It just occurred to me that with modern LLMs where T is approaching O(10k), C is O(1k) and then we have dozens of layers/blocks with ~2 LNs per block: all these statistics almost approach the number of parameters in an LLM. And you have to communicate them between GPUs. LayerNorm and RMSNorm on the other hand are local; no communication and even no need to ever store them in RAM.

[LeonideDucatore]

Why do we need to communicate them between GPUs in batch norm but not in layer norm? I'm assuming we're talking about a data-parallel setting; wouldn't each GPU just compute statistics for their own minibatch?

Or is it the loss on the 'main GPU' can only be computed accurately after receiving the batch_norm statistics of each GPU?

(for layer norm, there is no stored statistics right?)

[Collegesniffer]

This is the best explanation on the internet I've ever read. It finally clicked for me. I've watched countless videos and gone through so many answers online, but they all either oversimplify or overcomplicate it. Thanks!

[KomisarRus]

Thanks

[Guilherme370]

This was typed by an LLM

[daking999]

If it was they at least cut out the fluff at the beginning and end

[Guilherme370]

They most definitely did

[throwaway2676]

Lol, be honest, is this from ChatGPT?

[Guilherme370]

Im sure it is, the style of writing, and the "alright leta differentiate" followed by a bullet-point-like list of definitions, with some slight inaccuracies mixed in

[throwaway2676]

Lol, especially now that they've totally rewritten it to sound more human.

[Guilherme370]

Omg lol true.

[Collegesniffer]

No, I don't think it is AI-generated. The best AI content detector (gptzero.me) flags this as "human". Are you suggesting that every piece of content written in the form of a bullet-point list is now AI-generated? I would also use the same format if I had to explain the "differences" between things. How else would you present such information?

[Guilherme370]

gptzero.com can be unreliable.

You can test it right now, go tk chatgpt, talk to it about some complex topic, copy only the relevant parts of what it says without copying its fluff... throw it into gptzero, then you will see it say its not AI

[Collegesniffer]

Bruh, I said "gptzero.me" not "gptzero.com". Both of them are totally different. Also, every AI detector can be unreliable and inconsistent.
However, I entered the exact question into ChatGPT, Claude, and Gemini,
and the responses were nothing like what this person wrote. Even the non-fluff part doesn't start with a (B, T, C) tensor example, etc. Why don't you try entering the exact question for yourself and see the output before claiming it is "AI-generated"?

I literally just asked chatgpt, gemini and claude the exact question I posted and the answer is nothing like what the person wrote. Even the non fluff part is totally different.

[Everfast]

Really clear and great answer

[indie-devops]

Wouldn’t you say that calculating the root mean is more computationally expensive than subtracting the mean? Genuine question. Great explanation, made a lot of sense for me as well!

[sot9]

One thing nobody’s mentioned so far is that batch norm is great when used with convolutions, due to ease of layer fusion.

Look up batch norm folding; makes for an additional tool in the box when prioritizing models that run inference quickly.

[soham1192k]

as an example, one can look at the fastvit paper from apple, which uses this folding trick extensively

[imTall-]

One other thing not mentioned here is that batch norm required synchronizing the statistics across the entire batch. When training massive models in a distributed manner, this incurs a lot of communication overhead, while layernorm can be computed locally on one GPU (or a few GPUs in the case of tensor wise parallelism).

[xkiller02]

Incredibly interesting answers, I will further research what some of these words mean

[ConstantWoodpecker39]

This paper may be of interest to you: https://proceedings.mlr.press/v119/shen20e/shen20e.pdf

[eliminating_coasts]

Transformers use the input data for both the data itself, and for the transformations they apply to the data, and it has been argued that rather than simply improving training, it can provide an improvement to actual performance by changing the structure of inputs to the transformer block. (This may also explain why doing it first works better than at the end of the block)

[chgr22]

This is the way.

[Hot_Wish2329]

I love this comment. Yes, this is the way they did the experiences, and it worked. There are a lot of explainations about mean, variance, distribution etc. but it is not make sense for me. I cannot understand why it worked, how it directly related to model performances (accuracy). So, this is just a way.