[Keras Ops and Layer] Add keras.ops.rms_norm() and keras.layers.RMSNormalization() (#20911)

* Add RMSNorm and rms_norm

* math.square -> numpy.square

* Update docstrings

* Add RMSNormalization Layer

* Update docstrings

* Lint with new ruff version

* Add tests for layer

* Address comments

* Convert to tensor if not - avoid openvino and torch typing issues if scale is scalar

* address comments

* Fix tests

* Add reference to paper

* Fix docstring to remove input_dim argument

* Update layer_normalization.py

---------

Co-authored-by: François Chollet <[email protected]>

Author: DavidLandup0

keras/api/_tf_keras/keras/layers/__init__.py

@@ -83,6 +83,7 @@
 from keras.src.layers.normalization.layer_normalization import (
     LayerNormalization,
 )
+from keras.src.layers.normalization.rms_normalization import RMSNormalization
 from keras.src.layers.normalization.spectral_normalization import (
     SpectralNormalization,
 )

keras/api/_tf_keras/keras/ops/__init__.py

@@ -91,6 +91,7 @@
 from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
+from keras.src.ops.nn import rms_normalization
 from keras.src.ops.nn import selu
 from keras.src.ops.nn import separable_conv
 from keras.src.ops.nn import sigmoid

keras/api/_tf_keras/keras/ops/nn/__init__.py

@@ -35,6 +35,7 @@
 from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
+from keras.src.ops.nn import rms_normalization
 from keras.src.ops.nn import selu
 from keras.src.ops.nn import separable_conv
 from keras.src.ops.nn import sigmoid

keras/api/layers/__init__.py

@@ -83,6 +83,7 @@
 from keras.src.layers.normalization.layer_normalization import (
     LayerNormalization,
 )
+from keras.src.layers.normalization.rms_normalization import RMSNormalization
 from keras.src.layers.normalization.spectral_normalization import (
     SpectralNormalization,
 )

keras/api/ops/__init__.py

@@ -91,6 +91,7 @@
 from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
+from keras.src.ops.nn import rms_normalization
 from keras.src.ops.nn import selu
 from keras.src.ops.nn import separable_conv
 from keras.src.ops.nn import sigmoid

keras/api/ops/nn/__init__.py

@@ -35,6 +35,7 @@
 from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
+from keras.src.ops.nn import rms_normalization
 from keras.src.ops.nn import selu
 from keras.src.ops.nn import separable_conv
 from keras.src.ops.nn import sigmoid

keras/src/layers/__init__.py

@@ -57,6 +57,7 @@
 from keras.src.layers.normalization.layer_normalization import (
     LayerNormalization,
 )
+from keras.src.layers.normalization.rms_normalization import RMSNormalization
 from keras.src.layers.normalization.spectral_normalization import (
     SpectralNormalization,
 )

keras/src/layers/normalization/layer_normalization.py

@@ -86,7 +86,9 @@ class LayerNormalization(Layer):
         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.
+            approach that avoids ever computing the mean of the input. Note that
+            this *isn't* equivalent to the computation that the
+            `keras.layers.RMSNormalization` layer performs.
         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.

keras/src/layers/normalization/rms_normalization.py

@@ -0,0 +1,98 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.layers.layer import Layer
+
+
+@keras_export("keras.layers.RMSNormalization")
+class RMSNormalization(Layer):
+    """Root Mean Square (RMS) Normalization layer.
+
+    This layer normalizes the input tensor based on its RMS value.
+
+    The Keras layer performs the operation as described in
+    [Root Mean Square Layer Normalization](https://arxiv.org/pdf/1910.07467)
+    by Biao Zhang et al.
+
+
+    If `scale` is enabled, the layer will scale the normalized outputs via
+    a learnable scaling factor.
+
+    So, with scaling enabled, the normalization equations
+    are as follows:
+
+    Let the intermediate activations for a mini-batch to be the `inputs`.
+
+    ```python
+    rms_normalization(x) = x * rsqrt(mean(square(x))) * scale
+    ```
+
+    For example:
+
+    >>> layer = keras.layers.RMSNormalization()
+    >>> layer.build([5, 20, 30, 10])
+    >>> print(layer.scale.shape)
+    (10,)
+    >>> layer(np.random.rand(1, 10)).numpy()
+    array([[0.35098287, 1.0495652 , 1.4645109 , 1.2944688 , 0.31124955,
+            1.2768592 , 1.184331  , 0.17474432, 0.49955517, 1.2428929 ]],
+        dtype=float32)
+
+    Args:
+        axis: int. The axis on which to perform the normalization.
+        epsilon: float. A small number to add to avoid division by zero.
+    """
+
+    def __init__(self, axis=-1, epsilon=1e-6, **kwargs):
+        super().__init__(**kwargs)
+        self.axis = axis
+        self.epsilon = epsilon
+
+    def build(self, input_shape):
+        if isinstance(self.axis, list):
+            shape = tuple([input_shape[dim] for dim in self.axis])
+        else:
+            shape = (input_shape[self.axis],)
+            self.axis = [self.axis]
+
+        self.scale = self.add_weight(
+            name="scale", shape=shape, initializer="ones"
+        )
+
+        self.built = True
+
+    def call(self, x):
+        """Applies RMS normalization to the input tensor.
+
+        Args:
+            x: Input tensor of shape (batch_size, input_dim).
+
+        Returns:
+            The RMS-normalized tensor of the same shape (batch_size, input_dim),
+            scaled by the learned `scale` parameter.
+        """
+        return ops.rms_normalization(
+            x, scale=self.scale, axis=self.axis, epsilon=self.epsilon
+        )
+
+    def compute_output_shape(self, input_shape):
+        if isinstance(self.axis, int):
+            axes = [self.axis]
+        else:
+            axes = self.axis
+
+        for axis in axes:
+            if axis >= len(input_shape) or axis < -len(input_shape):
+                raise ValueError(
+                    f"Axis {axis} is out of bounds for "
+                    f"input shape {input_shape}. "
+                    f"Received: axis={self.axis}"
+                )
+        return input_shape
+
+    def get_config(self):
+        config = {
+            "axis": self.axis,
+            "epsilon": self.epsilon,
+        }
+        base_config = super().get_config()
+        return {**base_config, **config}

keras/src/layers/normalization/rms_normalization_test.py

@@ -0,0 +1,69 @@
+import numpy as np
+import pytest
+
+from keras.src import layers
+from keras.src import ops
+from keras.src import testing
+
+
+class RMSNormalizationTest(testing.TestCase):
+    @pytest.mark.requires_trainable_backend
+    def test_ln_basics(self):
+        self.run_layer_test(
+            layers.RMSNormalization,
+            init_kwargs={},
+            input_shape=(4, 2),
+            expected_output_shape=(4, 2),
+            expected_num_trainable_weights=1,
+            expected_num_seed_generators=0,
+        )
+        self.run_layer_test(
+            layers.RMSNormalization,
+            init_kwargs={
+                "axis": -1,
+            },
+            input_shape=(4, 2),
+            expected_output_shape=(4, 2),
+            expected_num_trainable_weights=1,
+            expected_num_seed_generators=0,
+        )
+
+    def test_correctness(self):
+        layer = layers.RMSNormalization()
+        layer.build(input_shape=(2, 2, 2))
+        inputs = np.random.normal(
+            loc=5.0, scale=10.0, size=(1000, 2, 2, 2)
+        ).astype("float32")
+
+        inputs = ops.convert_to_tensor(inputs)
+
+        out = layer(inputs)
+        expected = (
+            inputs
+            * ops.rsqrt(ops.mean(ops.square(inputs), axis=-1, keepdims=True))
+            * layer.scale
+        )
+
+        self.assertAllClose(out, expected, atol=1e-1)
+
+    def test_output(self):
+        layer = layers.RMSNormalization()
+        inputs = np.arange(10).astype("float32")[None, :]
+        out = layer(inputs)
+        self.assertAllClose(
+            out,
+            [
+                [
+                    0.0,
+                    0.18731716,
+                    0.37463433,
+                    0.5619515,
+                    0.74926865,
+                    0.9365858,
+                    1.123903,
+                    1.3112202,
+                    1.4985373,
+                    1.6858544,
+                ]
+            ],
+        )