emmi.optim.lion

Lion Optimizer Paper: Symbolic Discovery of Optimization Algorithms - https://arxiv.org/abs/2302.06675 Original Impl: https://github.com/google/automl/tree/master/lion This implementation is adapted from timm: https://github.com/huggingface/pytorch-image-models/blob/e44f14d7d2f557b9f3add82ee4f1ed2beefbb30d/timm/optim/lion.py

Classes

Lion

Implements Lion algorithm.

Functions

lion(params, grads, exp_avgs[, maximize, foreach])

Functional API that performs Lion algorithm computation.

Module Contents

class emmi.optim.lion.Lion(params, lr, betas=(0.9, 0.99), weight_decay=0.0, caution=False, maximize=False, foreach=None)

Bases: torch.optim.optimizer.Optimizer

Implements Lion algorithm.

Initialize the hyperparameters.

Parameters:

  • params (torch.optim.optimizer.ParamsT) – iterable of parameters to optimize or dicts defining parameter groups

  • lr (float) – learning rate

  • betas (tuple[float, float]) – coefficients used for computing running averages of gradient and its square

  • weight_decay (float) – weight decay coefficient

  • caution (bool) – apply caution

  • maximize (bool)

  • foreach (bool | None)

step(closure=None)

Performs a single optimization step.

Parameters:

closure – A closure that reevaluates the model and returns the loss.

Returns:

the loss.

emmi.optim.lion.lion(params, grads, exp_avgs, maximize=False, foreach=None, *, beta1, beta2, lr, weight_decay, caution)

Functional API that performs Lion algorithm computation.

Parameters: