File size: 1,036 Bytes
4e316d6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Normalization Kernels

RMSNorm (Root Mean Square Layer Normalization) — a simplified alternative to LayerNorm that skips mean-centering and only rescales by the root-mean-square of activations. Cheaper to compute and performs comparably for transformer pre-norm architectures.

Two weight conventions exist across model families:

| Convention | Formula | Models |
|---|---|---|
| Direct | `norm(x) * weight` | Qwen3, LLaMA, Mistral |
| Unit offset | `norm(x) * (1 + weight)` | Gemma, Gemma3 |

The unit-offset convention initialises weights to zero so the initial scale is 1.0 (identity). Our implementation supports both via the `add_unit_offset` parameter.

Computation is always done in float32 regardless of input dtype to avoid numerical instability in half-precision, then cast back.

## References

| Paper | Link |
|---|---|
| Root Mean Square Layer Normalization (Zhang & Sennrich, 2019) | https://arxiv.org/abs/1910.07467 |
| Layer Normalization (Ba et al., 2016) — the predecessor | https://arxiv.org/abs/1607.06450 |