optax.scale_by_belief

optax.scale_by_belief#

optax.scale_by_belief(
b1: jax.typing.ArrayLike = 0.9,
b2: jax.typing.ArrayLike = 0.999,
eps: jax.typing.ArrayLike = 1e-16,
eps_root: jax.typing.ArrayLike = 1e-16,
*,
nesterov: bool = False,
) → optax.GradientTransformation[source]#

Rescale updates according to the AdaBelief algorithm.

See optax.adabelief() for more details.

Parameters:

  • 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 second moment of the prediction error to improve numerical stability. If backpropagating gradients through the gradient transformation (e.g. for meta-learning), this must be non-zero.

  • nesterov – Whether to use Nesterov momentum.

Returns:

A optax.GradientTransformation object.