Repository: keras-team/keras-core

Issue:
From KerasNLP to KerasCore : Add `RMSLayerNormalization`
Context: [kers-team/keras-nlp/#1207](https://github.com/keras-team/keras-nlp/issues/1207)

KerasNLP's T5 (now Llama too) uses a custom layer norm which is indeed RMSLayerNormalization. It's being used as it is from [here](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/t5/t5_layer_norm.py) to [here](https://github.com/shivance/keras-nlp/blob/llama-backbone/keras_nlp/models/llama/llama_layernorm.py) 
As it is getting popular, it would be good to officially support it as a Keras layer.

I can open a PR!


Discussion:
As @mattdangerw says [here](https://github.com/keras-team/keras-nlp/issues/1207#issuecomment-1676527644), we could either have a new layer `RMSLayerNormalization` or an option to the existing layer, something like LayerNormalization(rms_scaling=True).
> something like LayerNormalization(rms_scaling=True).

I think that would be preferable. Thanks for the suggestion!
Something like this?

```
Arguments

    Args:
        ...
        center: If True, add offset of `beta` to normalized tensor. If False,
            `beta` is ignored. Defaults to `True`.
        scale: If True, multiply by `gamma`. If False, `gamma` is not used.
            When the next layer is linear (also e.g. `nn.relu`), this can be
            disabled since the scaling will be done by the next layer.
            Defaults to `True`.
        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.
```
Sure, sounds good to me. Thanks!

Resolving patch:
diff --git a/keras_core/layers/normalization/layer_normalization.py b/keras_core/layers/normalization/layer_normalization.py
--- a/keras_core/layers/normalization/layer_normalization.py
+++ b/keras_core/layers/normalization/layer_normalization.py
@@ -82,6 +82,10 @@ class LayerNormalization(Layer):
             When the next layer is linear (also e.g. `nn.relu`), this can be
             disabled since the scaling will be done by the next layer.
             Defaults to `True`.
+        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.
         beta_initializer: Initializer for the beta weight. Defaults to zeros.
         gamma_initializer: Initializer for the gamma weight. Defaults to ones.
         beta_regularizer: Optional regularizer for the beta weight.
@@ -106,6 +110,7 @@ def __init__(
         epsilon=1e-3,
         center=True,
         scale=True,
+        rms_scaling=False,
         beta_initializer="zeros",
         gamma_initializer="ones",
         beta_regularizer=None,
@@ -128,6 +133,7 @@ def __init__(
         self.epsilon = epsilon
         self.center = center
         self.scale = scale
+        self.rms_scaling = rms_scaling
         self.beta_initializer = initializers.get(beta_initializer)
         self.gamma_initializer = initializers.get(gamma_initializer)
         self.beta_regularizer = regularizers.get(beta_regularizer)
@@ -143,7 +149,7 @@ def build(self, input_shape):
         else:
             shape = (input_shape[self.axis],)
             self.axis = [self.axis]
-        if self.scale:
+        if self.scale or self.rms_scaling:
             self.gamma = self.add_weight(
                 name="gamma",
                 shape=shape,
@@ -155,7 +161,7 @@ def build(self, input_shape):
         else:
             self.gamma = None
 
-        if self.center:
+        if self.center and not self.rms_scaling:
             self.beta = self.add_weight(
                 name="beta",
                 shape=shape,
@@ -196,27 +202,31 @@ def _broadcast(v):
             # this is at least as numerically stable as the fused version.
             inputs = ops.cast(inputs, "float32")
 
-        # Calculate the mean and variance last axis (layer activations).
-        mean = ops.mean(inputs, axis=self.axis, keepdims=True)
-        # Don't use ops.var, as that would re-compute the mean from scratch.
-        variance = ops.mean(
-            ops.power(inputs - mean, 2), axis=self.axis, keepdims=True
-        )
-
-        scale, offset = _broadcast(self.gamma), _broadcast(self.beta)
+        # Calculate the variance last axis (layer activations).
+        variance = ops.var(inputs, axis=self.axis, keepdims=True)
 
         # Compute the batch normalization.
         inv = 1 / ops.sqrt(variance + self.epsilon)
-        if scale is not None:
-            scale = ops.cast(scale, inputs.dtype)
-            inv = inv * scale
-        x = -mean * inv
-        if offset is not None:
-            offset = ops.cast(offset, inputs.dtype)
-            x = offset + x
-        outputs = inputs * ops.cast(inv, inputs.dtype) + ops.cast(
-            x, inputs.dtype
-        )
+
+        if self.rms_scaling:
+            # Calculate outputs with only variance and gamma if rms scaling
+            # is enabled
+            outputs = inputs * ops.cast(inv, inputs.dtype) * self.gamma
+        else:
+            # Calculate the mean last axis (layer activations).
+            mean = ops.mean(inputs, axis=self.axis, keepdims=True)
+            scale, offset = _broadcast(self.gamma), _broadcast(self.beta)
+            if scale is not None:
+                scale = ops.cast(scale, inputs.dtype)
+                inv = inv * scale
+            x = -mean * inv
+            if offset is not None:
+                offset = ops.cast(offset, inputs.dtype)
+                x = offset + x
+
+            outputs = inputs * ops.cast(inv, inputs.dtype) + ops.cast(
+                x, inputs.dtype
+            )
 
         outputs = ops.cast(outputs, input_dtype)
 


Test patch:
diff --git a/keras_core/layers/normalization/layer_normalization_test.py b/keras_core/layers/normalization/layer_normalization_test.py
--- a/keras_core/layers/normalization/layer_normalization_test.py
+++ b/keras_core/layers/normalization/layer_normalization_test.py
@@ -49,6 +49,17 @@ def test_ln_basics(self):
             expected_num_losses=0,
             supports_masking=True,
         )
+        self.run_layer_test(
+            layers.LayerNormalization,
+            init_kwargs={"rms_scaling": True},
+            input_shape=(3, 3),
+            expected_output_shape=(3, 3),
+            expected_num_trainable_weights=1,
+            expected_num_non_trainable_weights=0,
+            expected_num_seed_generators=0,
+            expected_num_losses=0,
+            supports_masking=True,
+        )
         self.run_layer_test(
             layers.LayerNormalization,
             init_kwargs={"axis": (-3, -2, -1)},
@@ -95,3 +106,13 @@ def test_output(self):
         inputs = np.arange(5).astype("float32")[None, :]
         out = layer(inputs)
         self.assertAllClose(out, [[-0.41386, 0.29307, 1.0, 1.70693, 2.41386]])
+
+    def test_output_with_rms_scaling(self):
+        layer = layers.LayerNormalization(
+            dtype="float32",
+            rms_scaling=True,
+            gamma_initializer="ones",
+        )
+        inputs = np.arange(5).astype("float32")[None, :]
+        out = layer(inputs)
+        self.assertAllClose(out, [[0., 0.70693, 1.41386, 2.12079, 2.82772]])