Optimizer "Lion" in Symbolic Discovery of Optimization Algorithms

[towzeur]

🚀 The feature, motivation and pitch

A new optimizer "Lion" was proposed in Symbolic Discovery of Optimization Algorithms
https://arxiv.org/pdf/2302.06675.pdf

It is simple, memory efficient and have a faster runtime.
I think it can be a great addition to the Pytorch optimizer algorithms' collection as
it directly compares to AdamW which is already available.

the optimizer is a single file
https://github.com/google/automl/blob/master/lion/lion_pytorch.py

Pseudo Code:
$$
\begin{algorithm}
\caption{\name{} Lion}
\begin{algorithmic}
\State \textbf{given} $\beta_1$, $\beta_2$, $\lambda$, $\eta$, $f$
\State \textbf{initialize} $\theta_0$, $m_0\leftarrow 0$
\While{$\theta_t$ not converged}
\State $g_t \leftarrow \nabla_\theta f(\theta_{t-1})$
\State \textbf{update model parameters}
\State $c_t \leftarrow \beta_1 m_{t-1} + (1-\beta_1)g_t$
\State $\theta_t \leftarrow \theta_{t-1} - \eta_t(\text{sign}(c_t) + \lambda\theta_{t-1})$
\State \textbf{update EMA of $g_t$}
\State $m_t \leftarrow \beta_2 m_{t-1} + (1 - \beta_2)g_t$
\EndWhile
\State \textbf{return} $\theta_t$
\end{algorithmic}
\label{alg:ours}
\end{algorithm}
$$

Alternatives

No response

Additional context

No response

cc @vincentqb @jbschlosser @albanD @janeyx99

[janeyx99]

Thanks for raising this up! PyTorch Core has some pretty strict rules when it comes to adding a new Optimizer or Module: https://github.com/pytorch/pytorch/wiki/Developer-FAQ#i-have-a-new-function-or-feature-id-like-to-add-to-pytorch-should-it-be-in-pytorch-core-or-a-library-like-torchvision

In particular, this work might be a bit too new to be considered to be added into core just now, and the way to go may be to put it in a separate library. Please don't hesitate to let us know if there is any low level feature that you're missing to implement this though!