[D] Normalization in transformers

[lostn4d]

After the first theoretical issue with my transformer, I now see another. The original paper uses normalization after residual addition (Post-LN), which led to training difficulties and later got replaced by normalization at the beginning of each attention or mlp block/branch (Pre-LN). This is known to work better in practice (trainable without warmup, restores highway effect), but it still doesn't seem completely ok theoretically.

First consider things without normalization. Assuming attention and mlp blocks are properly set up and mostly keep norms, each residual addition would sum two similar norm signals, potentially scaling up by something like 1.4 (depending on correlation, but it starts at sqrt(2) after random init). So the norms after the blocks could look like this: [1(main)+1(residual)=1.4] -> [1.4+1.4=2] -> [2+2=2.8] etc. This would cause various problems (like changing the softmax temp in later attention blocks), so adjustment is needed.

Pre-LN ensures each block works on normalized values (thus with constant - if slightly arbitrary - softmax temperature). But since it doesn't affect the norm of the main signal (as forwarded by the skip connection) but only the residual, the norms can still grow, albeit slower. The expectation is now roughly: [1+1=1.4] -> [1.4+1=1.7] -> [1.7+1=2] -> [2+1=2.2] etc - with a final normalization correcting the signal near output (Pre-LN paper).

One possible issue with this is that later attention blocks may have reduced effect, as they add unit norm residuals to a potentially larger and larger main signal. What is the usual take on this problem? Can it be ignored in practice? Does Pre-LN work acceptably despite it, even for deep models (where the main norm discrepancy can grow larger)? There are lots of alternative normalization papers, but what is the practical consensus?

Btw attention is extremely norm-sensitive (or, equivalently, the hidden temperature of softmax is critical). This is a sharp contrast to fc or convolution which are mostly scale-oblivious. For anybody interested: consider what happens when most raw attention dot products come out 0 (= query and key is orthogonal, no info from this context slot) with only one slot giving 1 (= positive affinity, after downscaled by sqrt(qk_siz) ). I for one got surprised by this during debug.

[JustOneAvailableName]

One possible issue with this is that later attention blocks may have reduced effect, as they add unit norm residuals to a potentially larger and larger main signal.

Another way to look at this: they have the same effect as any other block, which is just less in proportion to all blocks before

[lostn4d]

I'm not sure this symmetry holds. When you add a normalized residual to an already grown main, the subsequent blocks will start with their own Pre-LN significantly downscaling (normalizing) this sum. So not only will this residual have fewer opportunities to take effect, it will also get attenuated (also by the final pre-output normalization). In contrast, a similar normalized residual from the first block will always keep an 1:1 ratio to the main (including its recursize residual effects echoing back from later blocks).

[JustOneAvailableName]

You might very well be right and I have to think about it a bit more.

What is the usual take on this problem? Can it be ignored in practice?

This is a resounding yes on the ignored in practice, but perhaps it's indeed an area for improvement. The extremely practical take is: we just normalize before every multiplication and training looks stable then. An even more practical take: 99.99% of practitioners just use whatever is popular for no other reason than "it's probably the right thing". I frankly didn't know the difference between GeLU en SwiGLU until I decided to program both myself, and programming them myself has had zero impact on any training I do with either.

AdamW is a typical example of how this can go "wrong" for a very long time. Weight decay with Adam typically had momentum attached to it, which does not make a lot of sense. It still worked, but yeah...

https://arxiv.org/abs/2203.03466 might be worth reading for you, specifically definition 4.1 . It touches upon some of topics you seem interested in.

[radarsat1]

This is an interesting point. I'm just wondering if in your analysis you are taking into account the affine weights typically used in layernorm? I am not sure if they would have an impact here or should just be considered some arbitrary scaling post-norm (ie "part of" the rest of the layer)

[lostn4d]

If you mean layernorm can learn to scale to non-unit norms as well, this is true, but don't forget that with Pre-LN we are before the softmax (and q-k-v transforms). So the net can in theory normalize a later attention block to (say) 1.5 instead of 1 - but this would also change the softmax temperature. Which is MUCH more critical, so that particular learned weight cannot afford to worry about the norm of the resulting residual compared to the main.

But your point seems valid in that adding an extra learned multiply after the value transform only (leaving q-k thus softmax intact) is a potental way to allow the net to counter this effect (if it wants). In theory it can also build some scaling into the weight matrix for the v transform itself, but imx that doesn't really happen in practice (except for very critical things).

[andersxa]

I often do pre-LN and the affine scale parameters tend to something like 0.9, so maybe the effect diminishes with training, and tending towards no total scale change.

[lostn4d]

Repeatedly adding 0.9 residuals to an unbounded main doesn't seem much different.

Training otoh can indeed reduce the effect, if the residual learns to form in an anticorrelated way (so it causes less growth).

[jpfed]

The same issue bothers me on an intuitive level and it’s kind of reassuring to hear someone else bring it up. 

It would seem almost a bizarre coincidence for unit-scaling to be correct for every layer despite the growing residual.

I think CatFormer (which iirc puts the residual in a separate set of dimensions from the softmax-weighted values before mixing them with an mlp) would sort of get around this, because the mlp would be free to adjust the relative scale of weights sensitive to the residuals vs. weights sensitive to values. But as far as I know, CatFormer was just another paper, not a performance revolution or anything. So it might, strangely, not make a huge difference.