LayerNormalization with rms_scaling documentation is different from implementation

[mzhukova]

The documentation mentions that

rms_scaling: If True, center and scale are ignored, and the inputs are scaled by gamma and the inverse square root of the square of all inputs. This is an approximate and faster approach that avoids ever computing the mean of the input.

However, in the implementation, it actually does the following:

        if self.rms_scaling:
            # Calculate outputs with only variance and gamma if rms scaling
            # is enabled
            # Calculate the variance along self.axis (layer activations).
            variance = ops.var(inputs, axis=self.axis, keepdims=True)
            inv = ops.rsqrt(variance + self.epsilon)

            outputs = (
                inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)
            )

So the mean is indeed used, as variance is computed here rather than RMS norm.

There was also a discussion during the addition of RMS Normalization (#20911 (comment)) that confirms this behavior.

I think the docs could use an update to clarify this behavior. Right now, it sounds like the mean isn't used when rms_scaling is on, but the code suggests otherwise.

[mzhukova]

Alternatively, the implementation could be adjusted to match the docs, though that might make RMSNormalization behave just like LayerNormalization with rms_scaling=True, correct?

[sonali-kumari1]

Hi @mzhukova -

Thanks for pointing this out. The documenation suggests that when rms_scaling=true, the layer performs normalization without involving mean of the inputs, similar to RMS normalization in this paper which removes mean calculation and relies on RMS to reduce the computational overhead introduced by LayerNorm. However, as confirmed by #20911, it uses

variance = var(inputs, axis=self.axis, keepdims=True)
inv = rsqrt(variance + epsilon)
outputs = inputs * inv * gamma

which means it still involves mean(since variance depends on the mean) and it is not true RMS normalization. We will look into this and update you.

[abheesht17]

Thanks for spotting this! Using ops.var() is definitely not the intended behaviour for rms_scaling and seems like a bug; when rms_scaling = True, the layer should function like RMSNorm (keras-team/keras-core#715).

Also, the actual RMSNormalization layer will be a duplicate of LayerNormalization once we do the fix.

Will open a PR shortly to fix this.

Edit: Just raising a warning if rms_scaling = True. The correct way of using RMS norm is to use the RMSNormalization layer.