Repository: google-deepmind/optax

Issue:
Implement the SM3 algorithm
See paper:
https://arxiv.org/abs/1901.11150


Discussion:
I'll try tackling this
Awesome!

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -25,6 +25,7 @@
 from optax._src.alias import radam
 from optax._src.alias import rmsprop
 from optax._src.alias import sgd
+from optax._src.alias import sm3
 from optax._src.alias import yogi
 from optax._src.base import EmptyState
 from optax._src.base import GradientTransformation
@@ -94,6 +95,7 @@
 from optax._src.transform import scale_by_rms
 from optax._src.transform import scale_by_rss
 from optax._src.transform import scale_by_schedule
+from optax._src.transform import scale_by_sm3
 from optax._src.transform import scale_by_stddev
 from optax._src.transform import scale_by_trust_ratio
 from optax._src.transform import scale_by_yogi
@@ -103,6 +105,7 @@
 from optax._src.transform import ScaleByRssState
 from optax._src.transform import ScaleByRStdDevState
 from optax._src.transform import ScaleByScheduleState
+from optax._src.transform import ScaleBySM3State
 from optax._src.transform import ScaleByTrustRatioState
 from optax._src.transform import ScaleState
 from optax._src.transform import trace
@@ -200,6 +203,7 @@
     "scale_by_rms",
     "scale_by_rss",
     "scale_by_schedule",
+    "scale_by_sm3",
     "scale_by_stddev",
     "scale_by_trust_ratio",
     "scale_by_yogi",
@@ -209,11 +213,14 @@
     "ScaleByRssState",
     "ScaleByRStdDevState",
     "ScaleByScheduleState",
+    "ScaleBySM3State",
     "ScaleByTrustRatioState",
     "ScaleState",
     "Schedule",
     "score_function_jacobians",
     "sgd",
+
+    "sm3",
     "sigmoid_binary_cross_entropy",
     "smooth_labels",
     "softmax_cross_entropy",
diff --git a/optax/_src/alias.py b/optax/_src/alias.py
--- a/optax/_src/alias.py
+++ b/optax/_src/alias.py
@@ -14,15 +14,10 @@
 # ==============================================================================
 """Aliases for popular optimisers."""
 
-from typing import Union, Optional
+from typing import Optional, Union
 
 import jax.numpy as jnp
-
-from optax._src import base
-from optax._src import combine
-from optax._src import privacy
-from optax._src import transform
-
+from optax._src import base, combine, privacy, transform
 
 ScalarOrSchedule = Union[float, base.Schedule]
 
@@ -359,6 +354,14 @@ def sgd(
   )
 
 
+def sm3(learning_rate: float,
+        momentum: float = 0.9) -> base.GradientTransformation:
+  return combine.chain(
+      transform.scale_by_sm3(momentum),
+      transform.scale(-learning_rate),
+  )
+
+
 def yogi(
     learning_rate: ScalarOrSchedule,
     b1: float = 0.9,
diff --git a/optax/_src/transform.py b/optax/_src/transform.py
--- a/optax/_src/transform.py
+++ b/optax/_src/transform.py
@@ -14,14 +14,12 @@
 # ==============================================================================
 """Gradient transformations."""
 
+from functools import reduce
 from typing import NamedTuple
 
 import jax
 import jax.numpy as jnp
-
-from optax._src import base
-from optax._src import clipping
-from optax._src import utils
+from optax._src import base, clipping, utils
 
 # pylint:disable=no-value-for-parameter
 
@@ -661,6 +659,77 @@ def update_fn(updates, state, params=None):
   return base.GradientTransformation(init_fn, update_fn)
 
 
+class ScaleBySM3State(base.OptState):
+  """State for the SM3 algorithm."""
+  mu: base.Updates
+  nu: base.Updates
+
+def scale_by_sm3(b1: float = 0.9,
+                 b2: float = 1.0,
+                 eps: float = 1e-8) -> base.GradientTransformation:
+  """Scale updates by sm3`.
+
+  References:
+    [Anil et. al 2019](https://arxiv.org/abs/1901.11150)
+
+  Args:
+    b1: decay rate for the exponentially weighted average of grads.
+    b2: decay rate for the exponentially weighted average of squared grads.
+    eps: term added to the denominator to improve numerical stability.
+
+  Returns:
+    An (init_fn, update_fn) tuple.
+  """
+
+  def zeros_for_dim(p):
+    return [jnp.zeros([s]) for s in p.shape]
+
+  def init_fn(params):
+    mu = jax.tree_map(zeros_for_dim, params)
+    nu = jax.tree_map(jnp.zeros_like, params)
+    return ScaleBySM3State(mu, nu)
+
+  def _expanded_shape(shape, axis):
+    # Replaces a `shape` of [M, N, K] with 1 in all dimensions except for i.
+    # For eg: i = 1 returns [1, N, 1].
+    rank = len(shape)
+    return [1] * axis + [shape[axis]] + [1] * (rank - axis - 1)
+
+  def _new_accum(g, v):
+    coeffs = ((1.0 - b2) if b2 != 1.0 else 1.0, b2)
+    if g.ndim < 2:
+      return coeffs[0]*g**2 + coeffs[1]*v[0]
+    else:
+      return coeffs[0]*g**2 + coeffs[1]*reduce(jnp.minimum, v)
+
+  def _new_mu(g, i):
+    if g.ndim < 2:
+      return g
+    else:
+      return jnp.max(g, axis=other_axes(i, g.ndim))
+
+  def other_axes(idx, ndim):
+    return list(range(idx)) + list(range(idx+1, ndim))
+
+  def update_fn(updates, state, params=None):
+    del params
+    mu = jax.tree_multimap(lambda g, v:
+                          [jnp.reshape(v[i], _expanded_shape(g.shape, i))
+                          for i in range(g.ndim)],
+                          updates, state.mu)
+    accum = jax.tree_multimap(_new_accum, updates, mu)
+    accum_inv_sqrt = jax.tree_map(
+        lambda t: jnp.where(t > 0, jax.lax.rsqrt(t + eps), 0.0), accum)
+    up = jax.tree_multimap(lambda g, a: g*a, updates, accum_inv_sqrt)
+    nu = _update_moment(up, state.nu, b1, 1)
+    mu = jax.tree_map(lambda g: [_new_mu(g, i)
+                      for i in range(g.ndim)], accum)
+
+    return nu, ScaleBySM3State(mu=mu, nu=nu)
+
+  return base.GradientTransformation(init_fn, update_fn)
+
+
 # TODO(b/183800387): remove legacy aliases.
 # These legacy aliases are here for checkpoint compatibility
 # To be removed once checkpoints have updated.


Test patch:
diff --git a/optax/_src/alias_test.py b/optax/_src/alias_test.py
--- a/optax/_src/alias_test.py
+++ b/optax/_src/alias_test.py
@@ -14,15 +14,11 @@
 # ==============================================================================
 """Tests for `alias.py`."""
 
-from absl.testing import absltest
-from absl.testing import parameterized
-
 import chex
 import jax
 import jax.numpy as jnp
-
-from optax._src import alias
-from optax._src import update
+from absl.testing import absltest, parameterized
+from optax._src import alias, update
 
 
 class AliasTest(chex.TestCase):
@@ -37,6 +33,7 @@ class AliasTest(chex.TestCase):
       ('fromage', lambda: alias.fromage(1e-2)),
       ('adabelief', lambda: alias.adabelief(1e-1)),
       ('radam', lambda: alias.radam(1e-1)),
+      ('sm3', lambda: alias.sm3(1.0)),
       ('yogi', lambda: alias.yogi(1.0)),
       ('dpsgd', lambda: alias.dpsgd(1e-2, 10.0, 0.001, 0))
   )
@@ -75,6 +72,7 @@ def step(params, state):
       ('fromage', lambda: alias.fromage(5e-3)),
       ('adabelief', lambda: alias.adabelief(1e-1)),
       ('radam', lambda: alias.radam(1e-3)),
+      ('sm3', lambda: alias.sm3(1.0)),
       ('yogi', lambda: alias.yogi(1.0)),
       ('dpsgd', lambda: alias.dpsgd(2e-3, 10., 0.001, 0, 0.2))
   )
diff --git a/optax/_src/equivalence_test.py b/optax/_src/equivalence_test.py
--- a/optax/_src/equivalence_test.py
+++ b/optax/_src/equivalence_test.py
@@ -14,17 +14,12 @@
 # ==============================================================================
 """Tests of equivalence between optax and other optimiser libraries."""
 
-from absl.testing import absltest
-from absl.testing import parameterized
-
 import chex
+import jax.numpy as jnp
+from absl.testing import absltest, parameterized
 from flax import optim
 from jax.experimental import optimizers
-import jax.numpy as jnp
-
-from optax._src import alias
-from optax._src import update
-
+from optax._src import alias, update
 
 STEPS = 50
 LR = 1e-2
@@ -62,6 +57,7 @@ def setUp(self):
        optimizers.rmsprop_momentum(LR, .9, 0.1, 0.9), 1e-5),
       ('adagrad', alias.adagrad(LR_SCHED, 0., 0.,),
        optimizers.adagrad(LR, 0.), 1e-5),
+      ('sm3', alias.sm3(LR), optimizers.sm3(LR), 1e-2),
   )
   def test_jax_optimizer_equivalent(self, optax_optimizer, jax_optimizer, rtol):