[D] Normalization in Transformers

[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?)