Skip to content

← Flax step 19 of 100

Medium Primitives

Implement RMSNorm (Modern LLM Norm)

Jump to the task

RMSNorm rescales a feature vector using its root mean square (RMS), then applies a learned scale. Unlike LayerNorm, it does not subtract the mean. This exercise implements the version with a scale parameter and no bias.

Mean square is not variance

For D features, compute:

mean_square = sum(x^2) / D
normalized = x / sqrt(mean_square + eps)
output = gamma * normalized

gamma has shape (D,) and starts at ones. Each feature gets its own scale. The mean square includes the contribution from a nonzero mean: mean(x^2) = variance(x) + mean(x)^2 in exact arithmetic. Replacing it with variance would implement a different normalization.

For x = [1, 2, 3, 4], the mean square is 7.5. With eps=1e-6, the initial output is approximately [0.3651, 0.7303, 1.0954, 1.4606]. All entries remain positive; RMSNorm does not center this vector.

Epsilon and special inputs

A positive epsilon makes the denominator nonzero for an all-zero vector, whose output remains zero. A nonzero constant vector is not mapped to zero: each element becomes c / sqrt(c*c + eps) before the learned scale.

The normalized vector's mean square is mean_square / (mean_square + eps). Its RMS is therefore not exactly one with positive epsilon, though it approaches one when the mean square is much larger than epsilon. Learned gamma values can change the RMS further. Epsilon also means rescaling the input need not leave the result exactly unchanged, especially for inputs close to zero.

How it differs from LayerNorm

LayerNorm centers the vector and uses centered squared values. This RMSNorm formula skips centering and uses one mean-square reduction. That is fewer operations than the explicit two-reduction LayerNorm formula, but actual speed depends on the implementation and hardware.

A fixed downstream bias cannot generally reproduce input-dependent mean subtraction. The two normalizations are distinct design choices, not interchangeable formulas with guaranteed identical training behavior.

Use axis=-1, keepdims=True so each feature vector has its own denominator. For shape (batch, sequence, D), the leading positions remain independent. The exercise inputs are one-dimensional vectors.

Problem

Implement MyRMSNorm(nn.Module) with an eps field and a compact __call__(self, x):

  1. Declare gamma with self.param("gamma", nn.initializers.ones, (x.shape[-1],)).
  2. Compute ms = jnp.mean(x ** 2, axis=-1, keepdims=True).
  3. Return gamma * (x / jnp.sqrt(ms + self.eps)).

Do not subtract a mean or create a bias parameter. Implement my_rmsnorm_forward(seed, x, eps) around the module.

The wrapper constructs the module with eps=float(eps), initializes it with jax.random.PRNGKey(int(seed)) and x, then applies it with the full variables mapping and x. These constant parameter initializers make the result independent of the seed. There is no training update.

x is a finite float32 JAX vector of shape (D,), with D >= 1, and values small enough that the intermediate squared values remain finite. eps is a finite positive scalar representable in float32. seed is a nonnegative integer-valued scalar that may arrive as a float. Return a JAX vector of shape (D,), leaving x unchanged.

Loading visualization…