Repository: google-deepmind/optax

Issue:
Implement scaling by AdaBelief
See paper:

https://arxiv.org/abs/2010.07468

This should be a trivial modification of the existing `scale_by_adam` transform.


Discussion:
I would like to implement this issue!
Awesome! Looking forward to your contribution!

Couple of suggestions:

Please move the bias correction code out of scale_by_adam into a standalone fn:

```
def _bias_correction(moment, decay, count):
  """Perform bias correction. This becomes a no-op as count goes to infinity."""
  bias_correction = 1 - decay**count
  return jax.tree_map(lambda t: t / bias_correction.astype(t.dtype), moment)
```

Then add both a `scale_by_belief` in `transform.py` and a `adabelief` optimiser in `alias.py`.

Thanks!

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -15,6 +15,7 @@
 # ==============================================================================
 """Optax: composable gradient processing and optimization, in JAX."""
 
+from optax._src.alias import adabelief
 from optax._src.alias import adagrad
 from optax._src.alias import adam
 from optax._src.alias import adamw
@@ -55,6 +56,7 @@
 from optax._src.transform import Params
 from optax._src.transform import scale
 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_rms
 from optax._src.transform import scale_by_rss
@@ -81,6 +83,7 @@
 __version__ = "0.0.1"
 
 __all__ = (
+    "adabelief",
     "adagrad",
     "adam",
     "adamw",
@@ -124,6 +127,7 @@
     "rmsprop",
     "scale",
     "scale_by_adam",
+    "scale_by_belief",
     "scale_by_fromage",
     "scale_by_rms",
     "scale_by_rss",
diff --git a/optax/_src/alias.py b/optax/_src/alias.py
--- a/optax/_src/alias.py
+++ b/optax/_src/alias.py
@@ -20,6 +20,16 @@
 GradientTransformation = transform.GradientTransformation
 
 
+def adabelief(learning_rate: float,
+         b1: float = 0.9,
+         b2: float = 0.999,
+         eps: float = 1e-8) -> GradientTransformation:
+  return combine.chain(
+      transform.scale_by_belief(b1=b1, b2=b2, eps=eps),
+      transform.scale(-learning_rate),
+  )
+
+
 def adagrad(
     learning_rate: float,
     initial_accumulator_value: float = 0.1,
diff --git a/optax/_src/transform.py b/optax/_src/transform.py
--- a/optax/_src/transform.py
+++ b/optax/_src/transform.py
@@ -298,6 +298,10 @@ def _safe_int32_increment(count):
   one = jnp.array(1, dtype=jnp.int32)
   return jnp.where(count < max_int32_value, count + one, max_int32_value)
 
+def _bias_correction(moment, decay, count):
+  """Perform bias correction. This becomes a no-op as count goes to infinity."""
+  bias_correction = 1 - decay**count
+  return jax.tree_map(lambda t: t / bias_correction.astype(t.dtype), moment)
 
 def scale_by_adam(
     b1: float = 0.9, b2: float = 0.999,
@@ -328,10 +332,8 @@ def update_fn(updates, state, params=None):
     mu = _update_moment(updates, state.mu, b1, 1)
     nu = _update_moment(updates, state.nu, b2, 2)
     count_inc = _safe_int32_increment(state.count)
-    mu_bias_correction = 1 - b1**count_inc
-    nu_bias_correction = 1 - b2**count_inc
-    mu_hat = jax.tree_map(lambda t: t / mu_bias_correction.astype(t.dtype), mu)
-    nu_hat = jax.tree_map(lambda t: t / nu_bias_correction.astype(t.dtype), nu)
+    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)
@@ -364,6 +366,45 @@ def update_fn(updates, state, params=None):
   return GradientTransformation(init_fn, update_fn)
 
 
+def scale_by_belief(
+    b1: float = 0.9, b2: float = 0.999,
+    eps: float = 1e-8, eps_root: float = 0.0) -> GradientTransformation:
+  """Rescale updates according to the Adam algorithm.
+
+  References:
+    [Zhuang et al, 2020](https://arxiv.org/abs/2010.07468)
+
+  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.
+
+  Returns:
+    An (init_fn, update_fn) tuple.
+  """
+
+  def init_fn(params):
+    mu = jax.tree_map(jnp.zeros_like, params)  # First moment
+    s = jax.tree_map(jnp.zeros_like, params)  # Second Central moment
+    return ScaleByAdamState(count=jnp.zeros([], jnp.int32), mu=mu, nu=s)
+
+  def update_fn(updates, state, params=None):
+    del params
+    mu = _update_moment(updates, state.mu, b1, 1)
+    prediction_error = jax.tree_multimap(lambda g, m: g-m, updates, state.mu)
+    nu = _update_moment(prediction_error, 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
@@ -82,6 +82,7 @@ def step(updates, state):
       ('lamb', alias.adamw(1e-1)),
       ('rmsprop', alias.rmsprop(1e-1)),
       ('fromage', transform.scale_by_fromage(-1e-2)),
+      ('adabelief', alias.adabelief(1e-1)),
   )
   def test_parabel(self, opt):
     initial_params = jnp.array([-1.0, 10.0, 1.0])
@@ -111,6 +112,7 @@ def step(params, state):
       ('lamb', alias.adamw(1e-1)),
       ('rmsprop', alias.rmsprop(5e-3)),
       ('fromage', transform.scale_by_fromage(-5e-3)),
+      ('adabelief', alias.adabelief(1e-1)),
   )
   def test_rosenbrock(self, opt):
     a = 1.0