google-deepmind / google-deepmind/gemma

[Feature request]: Eliminate pre-attention RMSNorm in Gemma 4 via scale invariance + weight folding

Open
#638 3 comments 0 reactions 0 assignees View on GitHub
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).

Image

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.

Image

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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.