optimizer that implements the lion algorithm the lion optimizer is a stochasticgradientdescent method that uses the sign operator to control the magnitude of the update unlike other adaptive optimizers such as adam that rely on secondorder moments this make lion more memoryefficient as it only keeps track of the momentum according to the s see reference its performance gain over adam grows with the batch size because the update of lion is produced through the sign operation resulting in a larger norm a suitable learning rate for lion is typically 310x smaller than that for adamw the weight decay for lion should be in turn 310x larger than that for adamw to maintain a similar strength lr wd args learningrate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta1 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the rate to combine the current gradient and the 1st moment estimate defaults to 0 9 beta2 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 1st moment estimate defaults to 0 99 baseoptimizerkeywordargs references chen et al 2023http arxiv orgabs2302 06675 s implementation http github comgoogleautomltreemasterlion initialize optimizer variables lion optimizer has one variable momentums args varlist list of model variables to build lion variables on update step given gradient and the associated model variable lr ops castlearningrate variable dtype gradient ops castgradient variable dtype beta1 ops castself beta1 variable dtype beta2 ops castself beta2 variable dtype m self momentumsself getvariableindexvariable self assignsub variable ops multiply lr ops sign ops add ops multiplym beta1 ops multiplygradient 1 0 beta1 self assign m ops add ops multiplym beta2 ops multiplygradient 1 0 beta2 def getconfigself config super getconfig config update beta1 self beta1 beta2 self beta2 return config lion doc lion doc replace baseoptimizerkeywordargs optimizer baseoptimizerkeywordargs optimizer that implements the lion algorithm the lion optimizer is a stochastic gradient descent method that uses the sign operator to control the magnitude of the update unlike other adaptive optimizers such as adam that rely on second order moments this make lion more memory efficient as it only keeps track of the momentum according to the s see reference its performance gain over adam grows with the batch size because the update of lion is produced through the sign operation resulting in a larger norm a suitable learning rate for lion is typically 3 10x smaller than that for adamw the weight decay for lion should be in turn 3 10x larger than that for adamw to maintain a similar strength lr wd args learning_rate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta_1 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the rate to combine the current gradient and the 1st moment estimate defaults to 0 9 beta_2 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 1st moment estimate defaults to 0 99 base_optimizer_keyword_args references chen et al 2023 http arxiv org abs 2302 06675 s implementation http github com google automl tree master lion initialize optimizer variables lion optimizer has one variable momentums args var_list list of model variables to build lion variables on update step given gradient and the associated model variable
from keras import ops
from keras.api_export import keras_export
from keras.optimizers import optimizer


@keras_export(["keras.optimizers.Lion"])
class Lion(optimizer.Optimizer):
    

    def __init__(
        self,
        learning_rate=0.001,
        beta_1=0.9,
        beta_2=0.99,
        weight_decay=None,
        clipnorm=None,
        clipvalue=None,
        global_clipnorm=None,
        use_ema=False,
        ema_momentum=0.99,
        ema_overwrite_frequency=None,
        name="lion",
        **kwargs,
    ):
        super().__init__(
            learning_rate=learning_rate,
            name=name,
            weight_decay=weight_decay,
            clipnorm=clipnorm,
            clipvalue=clipvalue,
            global_clipnorm=global_clipnorm,
            use_ema=use_ema,
            ema_momentum=ema_momentum,
            ema_overwrite_frequency=ema_overwrite_frequency,
            **kwargs,
        )
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        if beta_1 <= 0 or beta_1 > 1:
            raise ValueError(
                "Argument `beta_1` must be in the [0, 1] range. Otherwise, the "
                f"optimizer degenerates to SignSGD. Received: beta_1={beta_1}."
            )

    def build(self, var_list):
        
        if self.built:
            return
        super().build(var_list)
        self._momentums = []
        for var in var_list:
            self._momentums.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="momentum"
                )
            )

    def update_step(self, gradient, variable, learning_rate):
        
        lr = ops.cast(learning_rate, variable.dtype)
        gradient = ops.cast(gradient, variable.dtype)
        beta_1 = ops.cast(self.beta_1, variable.dtype)
        beta_2 = ops.cast(self.beta_2, variable.dtype)
        m = self._momentums[self._get_variable_index(variable)]

        self.assign_sub(
            variable,
            ops.multiply(
                lr,
                ops.sign(
                    ops.add(
                        ops.multiply(m, beta_1),
                        ops.multiply(gradient, (1.0 - beta_1)),
                    )
                ),
            ),
        )
        self.assign(
            m,
            ops.add(
                ops.multiply(m, beta_2), ops.multiply(gradient, (1.0 - beta_2))
            ),
        )

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "beta_1": self.beta_1,
                "beta_2": self.beta_2,
            }
        )
        return config


Lion.__doc__ = Lion.__doc__.replace(
    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
)