Function lion_optimizer

axlearn/common/optimizers.py:1709–1757

Lion optimizer with parameter scaling. https://arxiv.org/abs/2302.06675 Adapted from https://github.com/google/automl/blob/master/lion/lion_optax.py Args: learning_rate: The learning rate schedule. b1: The exponential decay rate for the 1st moment estimates. b2:

(
    learning_rate: schedule.Schedule,
    b1: float,
    b2: float,
    mu_dtype: Optional[jnp.dtype] = None,
    weight_decay: float = 0.0,
    multiply_by_parameter_scale: bool = False,
    weight_decay_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
)

def lion_optimizer(
    learning_rate: schedule.Schedule,
    b1: float,
    b2: float,
    mu_dtype: Optional[jnp.dtype] = None,
    weight_decay: float = 0.0,
    multiply_by_parameter_scale: bool = False,
    weight_decay_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
) -> PartitionedGradientTransformation:
    """Lion optimizer with parameter scaling.

    https://arxiv.org/abs/2302.06675
    Adapted from https://github.com/google/automl/blob/master/lion/lion_optax.py

    Args:
        learning_rate: The learning rate schedule.
        b1: The exponential decay rate for the 1st moment estimates.
        b2: The exponential decay rate for the 2nd moment estimates.
        weight_decay: Optional rate at which to decay weights.
        weight_decay_per_param_scale: 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.
        mu_dtype: Optional `dtype` to be used for the first order accumulator;
            if `None` then the dtype is inferred from params and updates.
        multiply_by_parameter_scale: If `True`, then scale learning_rate by
            parameter RMS. if `False`, provided learning_rate is absolute step size.
            Usually this should be left as False.

    Returns:
        A PartitionedGradientTransformation representing an Lion optimizer with parameter scalin.
    """
    tx = [scale_by_lion(b1=b1, b2=b2, mu_dtype=mu_dtype)]
    if multiply_by_parameter_scale:
        tx.append(scale_by_param_block_rms())
    tx.extend(
        [
            add_decayed_weights(
                weight_decay=weight_decay,
                # Weight decay updates will already be scaled by the learning rate below.
                learning_rate_exponent=None,
                per_param_scale=weight_decay_per_param_scale,
            ),
            scale_by_schedule(scale_from_learning_rate(learning_rate)),
        ]
    )

    return chain(*tx)


def adastar_optimizer(