Skip to content

← Flax NNX step 25 of 100

Medium Primitives

NNX Implement RMSNorm

Jump to the task

Rescale without centering

RMSNorm divides each feature vector by its root mean square, with an epsilon added inside the square root, then applies a learned gain. Unlike LayerNorm, it does not subtract the vector's mean. This exercise uses one gain vector gamma of shape (D,), with no offset parameter. That is D trainable scalar values, not a single scalar gain.

For features on the last axis, compute ms = mean(x ** 2, axis=-1, keepdims=True) and return gamma * x / sqrt(ms + eps). Each vector is normalized independently, including when there are leading batch or token dimensions.

Mean square is not variance

In exact arithmetic, mean(x ** 2) = variance(x) + mean(x) ** 2. RMSNorm therefore retains a signal that LayerNorm removes by centering. For example, a constant nonzero vector becomes zero under LayerNorm before its affine transform, but stays nonzero under RMSNorm.

If a vector has mean zero, the two normalizations agree when they use the same gain and epsilon and LayerNorm's offset is zero. Residual addition does not guarantee mean-zero vectors, so this is a special case, not a reason to treat the layers as interchangeable.

Before the gain, the normalized vector has mean square ms / (ms + eps) in exact arithmetic. Its RMS is not exactly one when epsilon is positive. The gain can change this value further.

A complete example

import jax.numpy as jnp
from flax import nnx

class MyRMSNorm(nnx.Module):
    def __init__(self, d, eps, rngs):
        self.gamma = nnx.Param(jnp.ones((d,), dtype=jnp.float32))
        self.eps = eps

    def __call__(self, x):
        ms = jnp.mean(x ** 2, axis=-1, keepdims=True)
        return self.gamma[...] * (x / jnp.sqrt(ms + self.eps))

model = MyRMSNorm(2, 1.0, nnx.Rngs(0))
x = jnp.array([[2., 2.], [0., 0.]], dtype=jnp.float32)
print(model(x))  # approximately [[0.8944272, 0.8944272], [0.0, 0.0]]
model.gamma[...] = jnp.array([2., 3.], dtype=jnp.float32)
print(model(x))  # approximately [[1.7888544, 2.6832816], [0.0, 0.0]]

The first row has mean square 4 and denominator sqrt(4 + 1). The second row stays zero; positive epsilon prevents division by zero. keepdims=True makes the denominator broadcast over each row's features. The constructor accepts rngs, but initializing a vector of ones consumes no random key.

What the simpler formula does and does not imply

RMSNorm removes the centering calculation and, in this version, the offset. The original RMSNorm paper reports comparable performance to LayerNorm in its experiments. This does not promise equal training quality or a particular speedup for every model. Runtime depends on the implementation, hardware and surrounding workload.

Loading visualization…

Problem

Implement rmsnorm_forward(seed, x, eps):

  1. Define the module above with one nnx.Param gain initialized to ones. Store epsilon separately; do not add a bias or subtract the input mean.
  2. Instantiate it with d=x.shape[-1], the supplied epsilon, and nnx.Rngs(int(seed)).
  3. Return the normalized and scaled array, preserving its shape.

The fixtures use a nonempty float32 vector x. The calculation also supports leading batch or token dimensions. Use finite inputs whose squares and reductions stay within float32 range, with positive finite eps. seed is nonnegative and integer-valued, possibly supplied as a float. It does not affect the deterministic gain initialization.

Each call creates a fresh module; no training occurs. If JIT-compiling the entrypoint, make seed static because it is converted to a Python integer. Leave epsilon in the array calculation rather than calling Python float on a traced scalar. Output values alone cannot prove the use of an NNX parameter; inspect that structure as well.