optax.contrib.cocob

Contents

optax.contrib.cocob#

optax.contrib.cocob(learning_rate: base.ScalarOrSchedule = 1.0, alpha: jax.typing.ArrayLike = 100.0, eps: jax.typing.ArrayLike = 1e-08, weight_decay: jax.typing.ArrayLike = 0.0, mask: Any | Callable[[base.Params], Any] | None = None) → base.GradientTransformation[source]#

Rescale updates according to the COntinuous COin Betting algorithm.

Algorithm for stochastic subgradient descent. Uses a gambling algorithm to find the minimizer of a non-smooth objective function by accessing its subgradients. All we need is a good gambling strategy. See Algorithm 2 of:

Parameters:

  • learning_rate – optional learning rate to e.g. inject some scheduler

  • alpha – fraction to bet parameter of the COCOB optimizer

  • eps – jitter term to avoid dividing by 0

  • weight_decay – L2 penalty

  • mask – mask for weight decay

Returns:

A GradientTransformation object.

References

Orabana et al, Training Deep Networks without Learning Rates Through Coin Betting, 2017

Contents