Repository: google-deepmind/optax

Issue:
Implement rectified Adam
See paper:

https://arxiv.org/abs/1908.03265


Discussion:
Hi, I would like to fix this once #13 is done
Awesome!
#13 is now approved and will be merged Monday morning.

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -21,6 +21,7 @@
 from optax._src.alias import adamw
 from optax._src.alias import lamb
 from optax._src.alias import noisy_sgd
+from optax._src.alias import radam
 from optax._src.alias import rmsprop
 from optax._src.alias import sgd
 from optax._src.alias import yogi
@@ -60,6 +61,7 @@
 from optax._src.transform import scale_by_adam
 from optax._src.transform import scale_by_belief
 from optax._src.transform import scale_by_fromage
+from optax._src.transform import scale_by_radam
 from optax._src.transform import scale_by_rms
 from optax._src.transform import scale_by_rss
 from optax._src.transform import scale_by_schedule
@@ -130,11 +132,13 @@
     "pathwise_jacobians",
     "piecewise_constant_schedule",
     "polynomial_schedule",
+    "radam",
     "rmsprop",
     "scale",
     "scale_by_adam",
     "scale_by_belief",
     "scale_by_fromage",
+    "scale_by_radam",
     "scale_by_rms",
     "scale_by_rss",
     "scale_by_schedule",
diff --git a/optax/_src/alias.py b/optax/_src/alias.py
--- a/optax/_src/alias.py
+++ b/optax/_src/alias.py
@@ -87,6 +87,17 @@ def noisy_sgd(learning_rate: float,
   )
 
 
+def radam(learning_rate: float,
+         b1: float = 0.9,
+         b2: float = 0.999,
+         eps: float = 1e-8,
+         threshold: float = 5.0) -> GradientTransformation:
+  return combine.chain(
+      transform.scale_by_radam(b1=b1, b2=b2, eps=eps, threshold=threshold),
+      transform.scale(-learning_rate),
+  )
+
+
 def rmsprop(learning_rate: float,
             decay: float = 0.9,
             eps: float = 1e-8,
diff --git a/optax/_src/transform.py b/optax/_src/transform.py
--- a/optax/_src/transform.py
+++ b/optax/_src/transform.py
@@ -459,6 +459,56 @@ def update_fn(updates, state, params=None):
   return GradientTransformation(init_fn, update_fn)
 
 
+def scale_by_radam(b1: float = 0.9,
+                  b2: float = 0.999,
+                  eps: float = 1e-8,
+                  eps_root: float = 0.0,
+                  threshold: float = 5.0) -> GradientTransformation:
+  """Rescale updates according to the Rectified Adam algorithm.
+  References:
+    [Liu et al, 2020](https://arxiv.org/abs/1908.03265)
+  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.
+    eps_root: term added to the denominator inside the square-root to improve
+      numerical stability when backpropagating gradients through the rescaling.
+    threshold: Threshold for variance tractability
+  Returns:
+    An (init_fn, update_fn) tuple.
+  """
+
+  ro_inf = 2./(1 - b2) - 1
+  def _radam_update(params):
+    ro = params[0]
+    mu_hat = params[1]
+    nu_hat = params[2]
+    r = jnp.sqrt((ro - 4)*(ro - 2)*ro_inf/((ro_inf - 4)*(ro_inf - 2)*ro))
+    updates = jax.tree_multimap(
+      lambda m, v: r*m / (jnp.sqrt(v + eps_root) + eps), mu_hat, nu_hat)
+    return updates
+
+  def init_fn(params):
+    mu = jax.tree_map(jnp.zeros_like, params)  # First moment
+    nu = jax.tree_map(jnp.zeros_like, params)  # Second moment
+    return ScaleByAdamState(count=jnp.zeros([], jnp.int32), mu=mu, nu=nu)
+
+  def update_fn(updates, state, params=None):
+    del params
+    mu = _update_moment(updates, state.mu, b1, 1)
+    nu = _update_moment(updates, state.nu, b2, 2)
+    count_inc = _safe_int32_increment(state.count)
+    b2t = b2**count_inc
+    ro = ro_inf - 2*count_inc*b2t/(1 - b2t)
+    mu_hat = _bias_correction(mu, b1, count_inc)
+    nu_hat = _bias_correction(nu, b2, count_inc)
+    updates = jax.lax.cond(ro >= threshold, _radam_update, lambda _: mu_hat,
+                    (ro, mu_hat, nu_hat))
+    return updates, ScaleByAdamState(count=count_inc, mu=mu, nu=nu)
+
+  return GradientTransformation(init_fn, update_fn)
+
+
 class AdditiveWeightDecayState(NamedTuple):
   """The decay transformation is stateless."""
 


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
@@ -83,6 +83,7 @@ def step(updates, state):
       ('rmsprop', alias.rmsprop(1e-1)),
       ('fromage', transform.scale_by_fromage(-1e-2)),
       ('adabelief', alias.adabelief(1e-1)),
+      ('radam', alias.radam(1e-1)),
       ('yogi', alias.yogi(1.0)),
   )
   def test_parabel(self, opt):
@@ -114,6 +115,7 @@ def step(params, state):
       ('rmsprop', alias.rmsprop(5e-3)),
       ('fromage', transform.scale_by_fromage(-5e-3)),
       ('adabelief', alias.adabelief(1e-1)),
+      ('radam', alias.radam(1e-3)),
       ('yogi', alias.yogi(1.0)),
   )
   def test_rosenbrock(self, opt):