Function adafactor_optimizer

axlearn/common/optimizers.py:1058–1146

Adafactor optimizer. References: https://arxiv.org/abs/1804.04235 https://github.com/deepmind/optax/blob/c4a4790b85ad69cda00a425cc3dcf9c9f9465120/optax/_src/alias.py#L77 WARNING: unlike adamw_optimizer, decay bias correction will *not* be applied if b1 or b2 is set to a constan

(
    learning_rate: schedule.Schedule,
    *,
    b1: Optional[schedule.Schedule],
    b2: schedule.Schedule,
    multiply_by_parameter_scale: bool,
    clipping_threshold: Optional[float],
    dtype_momentum: Any = jnp.float32,
    weight_decay: Optional[float] = None,
    weight_decay_scale_by_learning_rate_exponent: Optional[float] = None,
    weight_decay_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
    eps: float = 1e-30,
    factored: bool = True,
    apply_scale_by_trust_ratio: bool = False,
)

def adafactor_optimizer(
    learning_rate: schedule.Schedule,
    *,
    b1: Optional[schedule.Schedule],
    b2: schedule.Schedule,
    multiply_by_parameter_scale: bool,
    clipping_threshold: Optional[float],
    dtype_momentum: Any = jnp.float32,
    weight_decay: Optional[float] = None,
    weight_decay_scale_by_learning_rate_exponent: Optional[float] = None,
    weight_decay_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
    eps: float = 1e-30,
    factored: bool = True,
    apply_scale_by_trust_ratio: bool = False,
) -> PartitionedGradientTransformation:
    """Adafactor optimizer.

    References:
    https://arxiv.org/abs/1804.04235
    https://github.com/deepmind/optax/blob/c4a4790b85ad69cda00a425cc3dcf9c9f9465120/optax/_src/alias.py#L77

    WARNING: unlike adamw_optimizer, decay bias correction will *not* be applied if b1 or b2
    is set to a constant. However, users can enable bias correction by setting b1/b2 to
    config_for_function(schedule.decay_bias_correction).set(decay=<constant decay>).

    Args:
        learning_rate: (Schedule) The learning rate schedule.
        b1: (Schedule) first-moment exponential decay (beta1) schedule. If not None, enables
            momentum and uses extra memory.
        b2: (Schedule) second-moment exponential decay (beta2) schedule.
        multiply_by_parameter_scale: (bool): if True, then scale learning_rate by
            parameter norm. if False, provided learning_rate is absolute step size.
            Usually this should be set to False.
        clipping_threshold: (float>=1) optional value; if None, clipping disabled.
        dtype_momentum: (dtype) dtype of momentum buffers.
        weight_decay: (float) optional rate at which to decay weights.
        weight_decay_scale_by_learning_rate_exponent: (float) optional scale weight decay rate by
            (learning_rate ** exponent). Must not be None if weight_decay is not None.
            If set to 1, replicates the behavior of Adafactor weight decay in Lingvo and Praxis,
            https://github.com/google/praxis/blob/8fa3eb2e9ade0fd9a89a2ca56187882b12871605/praxis/optimizers.py#L1991-L1995.
            Set to 0 to disable scaling.
        weight_decay_per_param_scale: (optional) a Callable that returns a tree with same structure
            as the params PyTree, where each leaf is a float representing the per-param decay scale.
            The scale will be applied on top of the global decay rate:
            effective_decay_rate = global_decay_rate * per_param_scale.
            If None, all leaves will have a scale of 1.
        eps: (float) regularization constant for root mean squared gradient.
        factored: (bool) whether to use factored second-moment estimates.
        apply_scale_by_trust_ratio: (bool) whether to use variable-wise adaptive moments (LAMB):
            https://arxiv.org/abs/1904.00962.

    Returns:
        A PartitionedGradientTransformation representing an Adafactor optimizer.

    Raises:
        ValueError: If weight_decay_scale_by_learning_rate_exponent is not specified with
            weight_decay.
    """