Why is variance the wrong denominator for RMSNorm?
My outputs were consistently close to the reference, but not close enough. Increasing the tolerance would have hidden the actual mistake.
The intuition
RMSNorm uses sqrt(mean(x squared) + epsilon). Variance first subtracts the mean and measures a different quantity. For [2, 2], the RMS is 2 while the variance is zero.
y = x * rsqrt(mean(x*x) + eps) * weightNormalization is not one interchangeable operation. The exact statistic, epsilon placement, and accumulation precision are part of the layer definition. Here, I want a scale estimate without centering the vector.
Related reading · Root Mean Square Layer Normalization