Repository: google-deepmind/optax

Issue:
Implement support for Yogi optimiser
See paper:

https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf


Discussion:

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -23,6 +23,7 @@
 from optax._src.alias import noisy_sgd
 from optax._src.alias import rmsprop
 from optax._src.alias import sgd
+from optax._src.alias import yogi
 from optax._src.combine import chain
 from optax._src.control_variates import control_delta_method
 from optax._src.control_variates import control_variates_jacobians
@@ -63,6 +64,7 @@
 from optax._src.transform import scale_by_schedule
 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
 from optax._src.transform import ScaleByAdamState
 from optax._src.transform import ScaleByFromageState
 from optax._src.transform import ScaleByRmsState
@@ -134,6 +136,7 @@
     "scale_by_schedule",
     "scale_by_stddev",
     "scale_by_trust_ratio",
+    "scale_by_yogi",
     "ScaleByAdamState",
     "ScaleByFromageState",
     "ScaleByRmsState",
@@ -147,6 +150,7 @@
     "trace",
     "TraceState",
     "Updates",
+    "yogi",
 )
 
 #  _________________________________________
diff --git a/optax/_src/alias.py b/optax/_src/alias.py
--- a/optax/_src/alias.py
+++ b/optax/_src/alias.py
@@ -109,3 +109,13 @@ def sgd(learning_rate: float,
       transform.trace(decay=momentum, nesterov=nesterov),
       transform.scale(-learning_rate),
   )
+
+
+def yogi(learning_rate: float,
+         b1: float = 0.9,
+         b2: float = 0.999,
+         eps: float = 1e-3) -> GradientTransformation:
+  return combine.chain(
+      transform.scale_by_yogi(b1=b1, b2=b2, eps=eps),
+      transform.scale(-learning_rate),
+  )
diff --git a/optax/_src/transform.py b/optax/_src/transform.py
--- a/optax/_src/transform.py
+++ b/optax/_src/transform.py
@@ -408,6 +408,50 @@ def update_fn(updates, state, params=None):
   return GradientTransformation(init_fn, update_fn)
 
 
+def scale_by_yogi(
+    b1: float = 0.9, b2: float = 0.999,
+    eps: float = 1e-3, eps_root: float = 0.0, 
+    initial_accumulator_value: float = 1e-6) -> GradientTransformation:
+  """Rescale updates according to the Adam algorithm.
+
+  References:
+    [Zaheer et al, 2018](https://papers.nips.cc/paper/2018/hash/90365351ccc7437a1309dc64e4db32a3-Abstract.html) #pylint:disable=line-too-long
+
+  Args:
+    b1: decay rate for the exponentially weighted average of grads.
+    b2: decay rate for the exponentially weighted average of variance of 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.
+    initial_accumulator_value: The starting value for accumulators.
+      Only positive values are allowed.
+
+  Returns:
+    An (init_fn, update_fn) tuple.
+  """
+
+  def init_fn(params):
+    value_like = lambda p: jnp.full_like(p, initial_accumulator_value)
+    mu = jax.tree_map(value_like, params)  # First moment
+    nu = jax.tree_map(value_like, params)  # Second Central 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)
+    signed_sq = jax.tree_multimap(
+        lambda g, v: jnp.sign(v - g**2)*g**2, updates, state.nu)
+    nu = _update_moment(signed_sq, state.nu, b2, 2)
+    count_inc = _safe_int32_increment(state.count)
+    mu_hat = _bias_correction(mu, b1, count_inc)
+    nu_hat = _bias_correction(nu, b2, count_inc)
+    updates = jax.tree_multimap(
+        lambda m, v: m / (jnp.sqrt(v + eps_root) + eps), 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)),
+      ('yogi', alias.yogi(1.0)),
   )
   def test_parabel(self, opt):
     initial_params = jnp.array([-1.0, 10.0, 1.0])
@@ -113,6 +114,7 @@ def step(params, state):
       ('rmsprop', alias.rmsprop(5e-3)),
       ('fromage', transform.scale_by_fromage(-5e-3)),
       ('adabelief', alias.adabelief(1e-1)),
+      ('yogi', alias.yogi(1.0)),
   )
   def test_rosenbrock(self, opt):
     a = 1.0