Due to the scale invariance of RMS, an RMSNorm layer followed by a linear projection followed by another RMSNorm allows the first RMSNorm to be eliminated entirely — a mathematically lossless simplification.
For MLA-models that use both latent normalization and pre-attention norm, this means the pre-attention RMSNorm can be removed with no change to model outputs, see FlashNorm paper.
However, the pre-attention norm's learned weights are still needed. These can be eliminated cleanly by folding them into the Q and KV down-projection weights using the FlashNorm weight-folding trick — again with no loss in model accuracy.
For reference, we have applied this weight folding trick to a few LLMs (Llama, Qwen, SMolLM) here:
https://huggingface.co/models?other=weightless-rmsnorm
Due to the scale invariance of RMS, an RMSNorm layer followed by a linear projection followed by another RMSNorm allows the first RMSNorm to be eliminated entirely — a mathematically lossless simplification.
For MLA-models that use both latent normalization and pre-attention norm, this means the pre-attention RMSNorm can be removed with no change to model outputs, see FlashNorm paper.
However, the pre-attention norm's learned weights are still needed. These can be eliminated cleanly by folding them into the Q and KV down-projection weights using the FlashNorm weight-folding trick — again with no loss in model accuracy.
For reference, we have applied this weight folding trick to a few LLMs (Llama, Qwen, SMolLM) here:
https://huggingface.co/models?other=weightless-rmsnorm