Repository: google/flax

Issue:
Feature request: spectral normalization
### Discussed in https://github.com/google/flax/discussions/2902

<div type='discussions-op-text'>

<sup>Originally posted by **rdilip** February 23, 2023</sup>
Are spectral normalization layers something that anyone's working on at the moment? They're present in Haiku, so there's a template to work off of, and they'd be very useful (e.g., to enable things like https://www.tensorflow.org/tutorials/understanding/sngp). </div>


Discussion:

Resolving patch:
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -102,6 +102,7 @@
     GroupNorm as GroupNorm,
     LayerNorm as LayerNorm,
     RMSNorm as RMSNorm,
+    SpectralNorm as SpectralNorm,
 )
 from .pooling import (avg_pool as avg_pool, max_pool as max_pool, pool as pool)
 from .recurrent import (
diff --git a/flax/linen/normalization.py b/flax/linen/normalization.py
--- a/flax/linen/normalization.py
+++ b/flax/linen/normalization.py
@@ -14,9 +14,14 @@
 
 """Normalization modules for Flax."""
 
+import dataclasses
+import functools
 from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union
+
 from flax.linen.dtypes import canonicalize_dtype
 from flax.linen.module import Module, compact, merge_param  # pylint: disable=g-multiple-import
+from flax.linen.transforms import map_variables
+import jax
 from jax import lax
 from jax.nn import initializers
 import jax.numpy as jnp
@@ -187,6 +192,24 @@ def _normalize(
   return jnp.asarray(y, dtype)
 
 
+def _l2_normalize(x, axis=None, eps=1e-12):
+  """Normalizes along dimension `axis` using an L2 norm.
+
+  This specialized function exists for numerical stability reasons.
+
+  Args:
+    x: An input ndarray.
+    axis: Dimension along which to normalize, e.g. `1` to separately normalize
+      vectors in a batch. Passing `None` views `t` as a flattened vector when
+      calculating the norm (equivalent to Frobenius norm).
+    eps: Epsilon to avoid dividing by zero.
+
+  Returns:
+    An array of the same shape as 'x' L2-normalized along 'axis'.
+  """
+  return x * jax.lax.rsqrt((x * x).sum(axis=axis, keepdims=True) + eps)
+
+
 class BatchNorm(Module):
   """BatchNorm Module.
 
@@ -196,8 +219,8 @@ class BatchNorm(Module):
     BN = nn.BatchNorm(use_running_average=False, momentum=0.9, epsilon=1e-5,
                       dtype=jnp.float32)
 
-  The initialized variables dict will contain in addition to a 'params'
-  collection a separate 'batch_stats' collection that will contain all the
+  The initialized variables dict will contain, in addition to a 'params'
+  collection, a separate 'batch_stats' collection that will contain all the
   running statistics for all the BatchNorm layers in a model::
 
     vars_initialized = BN.init(key, x)  # {'params': ..., 'batch_stats': ...}
@@ -618,3 +641,211 @@ def __call__(self, x):
         self.bias_init,
         self.scale_init,
     )
+
+
+class SpectralNorm(Module):
+  """Spectral normalization. See:
+
+  - https://arxiv.org/abs/1802.05957
+  - https://arxiv.org/abs/1805.08318
+  - https://arxiv.org/abs/1809.11096
+
+  Spectral normalization normalizes the weight params so that the spectral
+  norm of the matrix is equal to 1. This is implemented as a layer wrapper
+  where each wrapped layer will have its params spectral normalized before
+  computing its ``__call__`` output.
+
+  Usage Note:
+  The initialized variables dict will contain, in addition to a 'params'
+  collection, a separate 'batch_stats' collection that will contain a
+  ``u`` vector and ``sigma`` value, which are intermediate values used
+  when performing spectral normalization. During training, we pass in
+  ``update_stats=True`` and ``mutable=['batch_stats']`` so that ``u``
+  and ``sigma`` are updated with the most recently computed values using
+  power iteration. This will help the power iteration method approximate
+  the true singular value more accurately over time. During eval, we pass
+  in ``update_stats=False`` to ensure we get deterministic behavior from
+  the model. For example::
+
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x, train):
+        x = nn.Dense(3)(x)
+        # only spectral normalize the params of the second Dense layer
+        x = nn.SpectralNorm(nn.Dense(4))(x, update_stats=train)
+        x = nn.Dense(5)(x)
+        return x
+
+    # init
+    x = jnp.ones((1, 2))
+    y = jnp.ones((1, 5))
+    model = Foo()
+    variables = model.init(jax.random.PRNGKey(0), x, train=False)
+
+    # train
+    def train_step(variables, x, y):
+      def loss_fn(params):
+        logits, updates = model.apply(
+            {'params': params, 'batch_stats': variables['batch_stats']},
+            x,
+            train=True,
+            mutable=['batch_stats'],
+        )
+        loss = jnp.mean(optax.l2_loss(predictions=logits, targets=y))
+        return loss, updates
+
+      (loss, updates), grads = jax.value_and_grad(loss_fn, has_aux=True)(
+          variables['params']
+      )
+      return {
+          'params': jax.tree_map(
+              lambda p, g: p - 0.1 * g, variables['params'], grads
+          ),
+          'batch_stats': updates['batch_stats'],
+      }, loss
+    for _ in range(10):
+      variables, loss = train_step(variables, x, y)
+
+    # inference / eval
+    out = model.apply(variables, x, train=False)
+
+  Attributes:
+    layer_instance: Module instance that is wrapped with SpectralNorm
+    n_steps: How many steps of power iteration to perform to approximate the
+      singular value of the weight params.
+    epsilon: A small float added to l2-normalization 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).
+    error_on_non_matrix: Spectral normalization is only defined on matrices.
+      By default, this module will return scalars unchanged and flatten
+      higher-order tensors in their leading dimensions. Setting this flag to
+      True will instead throw an error if a weight tensor with dimension
+      greater than 2 is used by the layer.
+    collection_name: Name of the collection to store intermediate values used
+      when performing spectral normalization.
+  """
+
+  layer_instance: Module
+  n_steps: int = 1
+  epsilon: float = 1e-12
+  dtype: Optional[Dtype] = None
+  param_dtype: Dtype = jnp.float32
+  error_on_non_matrix: bool = False
+  collection_name: str = 'batch_stats'
+
+  @compact
+  def __call__(self, *args, update_stats: bool, **kwargs):
+    """Compute the largest singular value of the weights in ``self.layer_instance``
+    using power iteration and normalize the weights using this value before
+    computing the ``__call__`` output.
+
+    Args:
+      *args: positional arguments to be passed into the call method of the
+        underlying layer instance in ``self.layer_instance``.
+      update_stats: if True, update the internal ``u`` vector and ``sigma`` value
+        after computing their updated values using power iteration. This will help
+        the power iteration method approximate the true singular value more
+        accurately over time.
+      **kwargs: keyword arguments to be passed into the call method of the
+        underlying layer instance in ``self.layer_instance``.
+
+    Returns:
+      Output of the layer using spectral normalized weights.
+    """
+
+    def layer_forward(layer_instance):
+      return layer_instance(*args, **kwargs)
+
+    return map_variables(
+        layer_forward,
+        trans_in_fn=lambda vs: jax.tree_util.tree_map_with_path(
+            functools.partial(
+                self._spectral_normalize,
+                layer_instance_name=self.layer_instance.name,
+                update_stats=update_stats,
+            ),
+            vs,
+        ),
+        init=self.is_initializing(),
+        mutable=True,
+    )(self.layer_instance)
+
+  def _spectral_normalize(self, path, vs, layer_instance_name, update_stats):
+    """Compute the largest singular value using power iteration and normalize
+    the variables ``vs`` using this value. This is intended to be a helper
+    function used in this Module's ``__call__`` method in conjunction with
+    ``nn.transforms.map_variables`` and ``jax.tree_util.tree_map_with_path``.
+
+    Args:
+      path: dict key path, used for naming the ``u`` and ``sigma`` variables
+      vs: variables to be spectral normalized
+      layer_instance_name: name of the underlying ``self.layer_instance``,
+        used for naming the ``u`` and ``sigma`` variables
+      update_stats: if True, update the ``u`` vector and ``sigma`` variables
+        after computing their updated values using power iteration. This will
+        help the power iteration method approximate the true singular value
+        more accurately over time.
+    """
+    value = jnp.asarray(vs)
+    value_shape = value.shape
+
+    # Skip and return value if input is scalar, vector or if number of power iterations is less than 1
+    if value.ndim <= 1 or self.n_steps < 1:
+      return value
+    # Handle higher-order tensors.
+    elif value.ndim > 2:
+      if self.error_on_non_matrix:
+        raise ValueError(
+            f'Input is {value.ndim}D but error_on_non_matrix is True'
+        )
+      else:
+        value = jnp.reshape(value, (-1, value.shape[-1]))
+
+    u_var_name = (
+        layer_instance_name
+        + '/'
+        + '/'.join((dict_key.key for dict_key in path[1:]))
+        + '/u'
+    )
+    u_var = self.variable(
+        self.collection_name,
+        u_var_name,
+        jax.random.normal,
+        self.make_rng('params')
+        if not self.has_variable(self.collection_name, u_var_name)
+        else None,
+        (1, value.shape[-1]),
+        self.param_dtype,
+    )
+    u0 = u_var.value
+    sigma_var_name = (
+        layer_instance_name
+        + '/'
+        + '/'.join((dict_key.key for dict_key in path[1:]))
+        + '/sigma'
+    )
+    sigma_var = self.variable(
+        self.collection_name, sigma_var_name, jnp.ones, (), self.param_dtype
+    )
+
+    # Power iteration for the weight's singular value.
+    for _ in range(self.n_steps):
+      v0 = _l2_normalize(
+          jnp.matmul(u0, value.transpose([1, 0])), eps=self.epsilon
+      )
+      u0 = _l2_normalize(jnp.matmul(v0, value), eps=self.epsilon)
+
+    u0 = jax.lax.stop_gradient(u0)
+    v0 = jax.lax.stop_gradient(v0)
+
+    sigma = jnp.matmul(jnp.matmul(v0, value), jnp.transpose(u0))[0, 0]
+
+    value /= jnp.where(sigma != 0, sigma, 1)
+    value_bar = value.reshape(value_shape)
+
+    if update_stats:
+      u_var.value = u0
+      sigma_var.value = sigma
+
+    dtype = canonicalize_dtype(vs, u0, v0, sigma, dtype=self.dtype)
+    return jnp.asarray(value_bar, dtype)


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
@@ -16,15 +16,17 @@
 
 import copy
 from absl.testing import absltest, parameterized
+from typing import Any
 
 from flax import ids
 from flax import linen as nn
+from flax.training import train_state
 
 import jax
 from jax import random
 from jax.nn import initializers
 import jax.numpy as jnp
-
+import optax
 import numpy as np
 
 # Parse absl flags test_srcdir and test_tmpdir.
@@ -294,6 +296,165 @@ def __call__(self, x):
     (y1, y2), variables = model.init_with_output(key, x)
     np.testing.assert_allclose(y1, y2, rtol=1e-4)
 
+  def test_spectral_norm_train(
+      self,
+  ):
+    class FooDense(nn.Module):
+
+      @nn.compact
+      def __call__(self, x, train):
+        x = nn.Dense(8)(x)
+        x = nn.SpectralNorm(nn.Dense(6))(x, update_stats=train)
+        x = nn.Dense(4)(x)
+        return x
+
+    class FooConv(nn.Module):
+
+      @nn.compact
+      def __call__(self, x, train):
+        x = nn.Dense(9)(x)
+        x = x.reshape((1, 3, 3))
+        x = nn.SpectralNorm(nn.Conv(2, kernel_size=(2, 2)))(
+            x, update_stats=train
+        )
+        x = x.reshape(1, -1)
+        x = nn.Dense(4)(x)
+        return x
+
+    class FooAttention(nn.Module):
+
+      @nn.compact
+      def __call__(self, x, train):
+        a = nn.Dense(4)(x)
+        b = nn.Dense(4)(x)
+        x = nn.SpectralNorm(nn.attention.MultiHeadDotProductAttention(4))(
+            a, b, update_stats=train
+        )
+        x = nn.Dense(4)(x)
+        return x
+
+    key1, key2, key3 = random.split(random.PRNGKey(0), 3)
+    x = random.normal(key1, (1, 4))
+    y = random.normal(key2, (1, 4))
+
+    for model_cls, var_paths in (
+        (FooDense, ('Dense_1/kernel/',)),
+        (FooConv, ('Conv_0/kernel/',)),
+        (
+            FooAttention,
+            (
+                'MultiHeadDotProductAttention_0/key/bias/',
+                'MultiHeadDotProductAttention_0/key/kernel/',
+                'MultiHeadDotProductAttention_0/out/kernel/',
+                'MultiHeadDotProductAttention_0/query/bias/',
+                'MultiHeadDotProductAttention_0/query/kernel/',
+                'MultiHeadDotProductAttention_0/value/bias/',
+                'MultiHeadDotProductAttention_0/value/kernel/',
+            ),
+        ),
+    ):
+      variables = model_cls().init(key3, x, train=False)
+      params, batch_stats = variables['params'], variables['batch_stats']
+      for var_path in var_paths:
+        self.assertTrue(var_path + 'u' in batch_stats['SpectralNorm_0'].keys())
+        self.assertTrue(
+            var_path + 'sigma' in batch_stats['SpectralNorm_0'].keys()
+        )
+
+      class TrainState(train_state.TrainState):
+        batch_stats: Any
+
+      state = TrainState.create(
+          apply_fn=model_cls().apply,
+          params=params,
+          batch_stats=batch_stats,
+          tx=optax.adam(1e-3),
+      )
+
+      @jax.jit
+      def train_step(state, batch):
+        def loss_fn(params):
+          logits, updates = state.apply_fn(
+              {'params': params, 'batch_stats': state.batch_stats},
+              x=batch['image'],
+              train=True,
+              mutable=['batch_stats'],
+          )
+          loss = jnp.mean(
+              optax.l2_loss(predictions=logits, targets=batch['label'])
+          )
+          return loss, updates
+
+        grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+        (loss, updates), grads = grad_fn(state.params)
+        state = state.apply_gradients(grads=grads)
+        state = state.replace(batch_stats=updates['batch_stats'])
+        return state, loss
+
+      prev_loss = float('inf')
+      for _ in range(10):
+        state, loss = train_step(state, {'image': x, 'label': y})
+        self.assertTrue(loss < prev_loss)
+        prev_loss = loss
+
+  def test_spectral_norm_sigma(self):
+    for n_steps, update_stats, result in (
+        (1, True, 4.0),
+        (3, True, 4.0),
+        (10, True, 4.0),
+        (1, False, 1.0),
+    ):
+
+      class Foo(nn.Module):
+
+        @nn.compact
+        def __call__(self, x, train):
+          x = nn.SpectralNorm(nn.Dense(8, use_bias=False), n_steps=n_steps)(
+              x, update_stats=train
+          )
+          return x
+
+      x = jnp.ones((1, 8))
+      model_cls = Foo()
+      variables = model_cls.init(random.PRNGKey(0), x, train=False)
+      params, batch_stats = variables['params'], variables['batch_stats']
+      params = jax.tree_map(lambda x: 4 * jnp.eye(*x.shape), params)
+      logits, updates = model_cls.apply(
+          {'params': params, 'batch_stats': batch_stats},
+          x=x,
+          train=update_stats,
+          mutable=True,
+      )
+      np.testing.assert_allclose(
+          updates['batch_stats']['SpectralNorm_0']['Dense_0/kernel/sigma'],
+          result,
+          atol=1e-3,
+      )
+
+  def test_spectral_norm_3d_tensor(self):
+    for error_on_non_matrix in (True, False):
+
+      class Foo(nn.Module):
+
+        @nn.compact
+        def __call__(self, x, train):
+          x = nn.SpectralNorm(
+              nn.DenseGeneral((3, 4), use_bias=False),
+              error_on_non_matrix=error_on_non_matrix,
+          )(x, update_stats=train)
+          return x
+
+      x = jnp.ones((1, 2))
+      model_cls = Foo()
+
+      if error_on_non_matrix:
+        with self.assertRaisesRegex(
+            ValueError, 'Input is 3D but error_on_non_matrix is True'
+        ):
+          variables = model_cls.init(random.PRNGKey(0), x, train=False)
+      else:
+        variables = model_cls.init(random.PRNGKey(0), x, train=False)
+
 
 class StochasticTest(absltest.TestCase):