Repository: google/flax

Issue:
Incorporate RMSNorm
Is there interest to incorporate RMSNorm (https://arxiv.org/pdf/1910.07467.pdf) either as a seperate normalization module or to make LayerNorm configurable to switch off the centering? It is becoming quite popular as a faster alternative to LayerNorm that maintains performance and is used for example in Palm.


Discussion:
Yes, I think it makes sense to support this in Flax. I think having a separate layer called RMSNorm makes sense, but it most likely can re-use use some of the shared normalization code we currently have for BatchNorm/LayerNorm.

Resolving patch:
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -108,7 +108,8 @@
 from .normalization import (
   BatchNorm as BatchNorm,
   GroupNorm as GroupNorm,
-  LayerNorm as LayerNorm
+  LayerNorm as LayerNorm,
+  RMSNorm as RMSNorm
 )
 from .pooling import (
   avg_pool as avg_pool,
diff --git a/flax/linen/normalization.py b/flax/linen/normalization.py
--- a/flax/linen/normalization.py
+++ b/flax/linen/normalization.py
@@ -49,7 +49,8 @@ def _abs_sq(x):
 def _compute_stats(x: Array, axes: Optional[Axes],
                    dtype: Optional[Dtype],
                    axis_name: Optional[str] = None,
-                   axis_index_groups: Any = None):
+                   axis_index_groups: Any = None,
+                   use_mean: bool = True):
   """Computes mean and variance statistics.
 
   This implementation takes care of a few important details:
@@ -68,6 +69,9 @@ def _compute_stats(x: Array, axes: Optional[Axes],
       are always at least float32 for stability (default: dtype of x).
     axis_name: Optional name for the pmapped axis to compute mean over.
     axis_index_groups: Optional axis indices.
+    use_mean: If true, calculate the mean from the input and use it when
+      computing the variance. If false, set the mean to zero and compute
+      the variance without subtracting the mean.
 
   Returns:
     A pair ``(mean, var)``.
@@ -79,8 +83,12 @@ def _compute_stats(x: Array, axes: Optional[Axes],
   dtype = jnp.promote_types(dtype, jnp.float32)
   x = jnp.asarray(x, dtype)
 
-  mean = jnp.mean(x, axes)
   mean2 = jnp.mean(_abs_sq(x), axes)
+  if use_mean:
+    mean = jnp.mean(x, axes)
+  else:
+    mean = jnp.zeros(mean2.shape, dtype=dtype)
+
   if axis_name is not None:
     concatenated_mean = jnp.concatenate([mean, mean2])
     mean, mean2 = jnp.split(
@@ -335,6 +343,76 @@ def __call__(self, x):
         self.bias_init, self.scale_init)
 
 
+class RMSNorm(Module):
+  """RMS Layer normalization (https://arxiv.org/abs/1910.07467).
+
+  RMSNorm normalizes the activations of the layer for each given example in a
+  batch independently, rather than across a batch like Batch Normalization.
+  Unlike LayerNorm which re-centers the mean to be 0 and normalizes by the
+  standard deviation of the activations, RMSNorm does not re-center at all
+  and instead normalizes by the root mean square of the activations.
+
+  Example::
+    >>> import jax.numpy as jnp
+    >>> import jax
+    >>> import flax.linen as nn
+    ...
+    >>> x = jax.random.uniform(jax.random.PRNGKey(0), (2, 3))
+    >>> layer = nn.RMSNorm()
+    >>> variables = layer.init(jax.random.PRNGKey(1), x)
+    >>> y = layer.apply(variables, x)
+
+  Attributes:
+    epsilon: A small float added to variance to avoid dividing by zero.
+    dtype: the dtype of the result (default: infer from input and params).
+    param_dtype: the dtype passed to parameter initializers (default: float32).
+    use_scale: If True, multiply by scale (gamma). 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.
+    scale_init: Initializer for scale, by default, one.
+    reduction_axes: Axes for computing normalization statistics.
+    feature_axes: Feature axes for learned bias and scaling.
+    axis_name: the axis name used to combine batch statistics from multiple
+      devices. See `jax.pmap` for a description of axis names (default: None).
+      This is only needed if the model is subdivided across devices, i.e. the
+      array being normalized is sharded across devices within a pmap.
+    axis_index_groups: groups of axis indices within that named axis
+      representing subsets of devices to reduce over (default: None). For
+      example, `[[0, 1], [2, 3]]` would independently batch-normalize over
+      the examples on the first two and last two devices. See `jax.lax.psum`
+      for more details.
+  """
+  epsilon: float = 1e-6
+  dtype: Optional[Dtype] = None
+  param_dtype: Dtype = jnp.float32
+  use_scale: bool = True
+  scale_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.ones
+  reduction_axes: Axes = -1
+  feature_axes: Axes = -1
+  axis_name: Optional[str] = None
+  axis_index_groups: Any = None
+
+  @compact
+  def __call__(self, x):
+    """Applies layer normalization on the input.
+
+    Args:
+      x: the inputs
+
+    Returns:
+      Normalized inputs (the same shape as inputs).
+    """
+    mean, var = _compute_stats(x, self.reduction_axes, self.dtype,
+                               self.axis_name, self.axis_index_groups,
+                               use_mean=False)
+
+    return _normalize(
+        self, x, mean, var, self.reduction_axes, self.feature_axes,
+        self.dtype, self.param_dtype, self.epsilon,
+        False, self.use_scale,
+        initializers.zeros, self.scale_init)
+
+
 class GroupNorm(Module):
   """Group normalization (arxiv.org/abs/1803.08494).
 


Test patch:
diff --git a/tests/linen/linen_test.py b/tests/linen/linen_test.py
--- a/tests/linen/linen_test.py
+++ b/tests/linen/linen_test.py
@@ -192,6 +192,23 @@ def test_layer_norm(self, reduction_axes):
                    jax.lax.rsqrt(x.var(axis=reduction_axes, keepdims=True) + e))
     np.testing.assert_allclose(y_one_liner, y, atol=1e-4)
 
+  @parameterized.parameters(
+      {'reduction_axes': -1},
+      {'reduction_axes': 1},
+      {'reduction_axes': (1, 2)})
+  def test_rms_norm(self, reduction_axes):
+    rng = random.PRNGKey(0)
+    key1, key2 = random.split(rng)
+    e = 1e-5
+    x = random.normal(key1, (2, 3, 4))
+    model_cls = nn.RMSNorm(use_scale=False, epsilon=e,
+                           reduction_axes=reduction_axes)
+    y, _ = model_cls.init_with_output(key2, x)
+    self.assertEqual(x.dtype, y.dtype)
+    self.assertEqual(x.shape, y.shape)
+    y_one_liner = (x * jax.lax.rsqrt(jnp.mean(jax.lax.square(x), axis=reduction_axes, keepdims=True) + e))
+    np.testing.assert_allclose(y_one_liner, y, atol=1e-4)
+
   def test_group_norm(self):
     rng = random.PRNGKey(0)
     key1, key2 = random.split(rng)