returnn.torch.optim.lion

Lion optimizer <https://arxiv.org/abs/2302.06675>

Code adapted from https://github.com/lucidrains/lion-pytorch/, which is adapted from https://github.com/google/automl/blob/master/lion/lion_pytorch.py.

class returnn.torch.optim.lion.Lion(params, lr: float = 0.0001, betas: Tuple[float, float] = (0.9, 0.99), weight_decay: float = 0.0, use_triton: bool | None = None, decoupled_weight_decay: bool = False)

Lion (Evolved Sign Momentum (Evo_l_ved S_i_gn M_o_me_n_tum)) optimizer <https://arxiv.org/abs/2302.06675>

step(closure: Callable | None = None)

update step

returnn.torch.optim.lion.update_fn(p, grad, exp_avg, lr, wd, beta1, beta2)

Lion update function