google-deepmind / google-deepmind/gemma
[Feature request]: Eliminate pre-attention RMSNorm in Gemma 4 via scale invariance + weight folding
- Dominant language
- Python
- Stars
- 5.7k
- Forks
- 1k
- Avg merge
- 10h 33m
- Merged PRs (30d)
- 2
Description
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. In models that use QKV-normalization (such as Gemma 4), this means the pre-attention RMSNorm can be removed with no change to model outputs, see [FlashNorm paper](https://arxiv.org/pdf/2407.09577).
However, the pre-attention norm's learned weights are still needed. These can be eliminated cleanly by folding them into the QKV projection weights using the [FlashNorm](https://arxiv.org/pdf/2407.09577) 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
Contributor guide
Assessment
This issue has not been assessed yet.