Stretch: RMSNorm (alternative to LayerNorm)

The problem, from first principles

[01-layer-normalization-forward]'s LayerNorm does two separate things: it CENTERS a position's feature vector (subtracts the mean) and it RESCALES it (divides by the standard deviation), before applying a learned gamma/beta. Zhang & Sennrich (2019) asked a pointed empirical question: how much of LayerNorm's benefit actually comes from the re-scaling, versus the re-centering? Their finding, since adopted by LLaMA, PaLM, Mistral, and most other modern large language models, was that the RESCALING is doing almost all of the useful work, and the mean-centering step can simply be dropped, with training stability and final model quality both essentially unaffected.

The result, RMSNorm, is a strictly SIMPLER, cheaper normalization: no mean to compute, no beta shift to learn, just a single division by the vector's own root-mean-square. Fewer operations per call, fewer parameters to learn, and (because it computes only ONE statistic instead of two) a meaningfully cheaper operation to run billions of times over during training and inference, a real, measurable cost saving at the scale modern LLMs run at.

From theory to code

Implement rmsnorm_forward(x, gamma, eps). Compute the root-mean-square of x along its LAST axis (sqrt(mean(x^2) + eps)), divide x by that value, then apply the learned per-feature multiplicative scale gamma (there is no beta, unlike [01-layer-normalization-forward]).

Constraints

• No mean-subtraction anywhere: RMSNorm normalizes using only the SECOND moment (mean(x^2)), never the first (mean(x)).
• eps is added INSIDE the square root, exactly like [01-layer-normalization-forward], to avoid dividing by zero.
• There is no beta parameter: the only learned parameter is the multiplicative gamma.
• Statistics are computed along the LAST axis, independently per position, exactly like [01-layer-normalization-forward].

Hints