optax.contrib.dadapt_adamw

Contents

optax.contrib.dadapt_adamw#

optax.contrib.dadapt_adamw(
learning_rate: base.ScalarOrSchedule = 1.0,
betas: tuple[jax.typing.ArrayLike, jax.typing.ArrayLike] = (0.9, 0.999),
eps: jax.typing.ArrayLike = 1e-08,
estim_lr0: jax.typing.ArrayLike = 1e-06,
weight_decay: jax.typing.ArrayLike = 0.0,
weight_decay_mask: Any | Callable[[base.Params], Any] | None = None,
) → base.GradientTransformationExtraArgs[source]#

Learning rate free AdamW by D-Adaptation.

Adapts the baseline learning rate of AdamW automatically by estimating the initial distance to solution in the infinity norm. This method works best when combined with a learning rate schedule that treats 1.0 as the base (usually max) value.

Parameters:

  • learning_rate – Learning rate scheduling parameter. The recommended schedule is a linear_schedule with init_value=1.0 and end_value=0, combined with a 0-20% learning rate warmup.

  • betas – Betas for the underlying AdamW Optimizer.

  • eps – eps for the underlying AdamW Optimizer.

  • estim_lr0 – Initial (under-)estimate of the learning rate.

  • weight_decay – AdamW style weight-decay. To use Regular Adam decay, chain with add_decayed_weights.

  • weight_decay_mask – A tree with same structure as (or a prefix of) the params PyTree, or a Callable that returns such a pytree given the params/updates. The leaves should be booleans, True for leaves/subtrees you want to apply the weight decay to, and False for those you want to skip. The mask must be static for the gradient transformation to be jit-compilable. Passed through to optax.add_decayed_weights().

Returns:

The corresponding optax.GradientTransformation.

References

Defazio et al, Learning-Rate-Free Learning by D-Adaptation, 2023

Contents