Google's Newly Discovered Optimizer Lion: The "Training Lion" That Has Both Efficiency and Effectiveness
Yesterday I came across a new paper from Google on Arxiv, Symbolic Discovery of Optimization Algorithms, which is mainly about automatically searching for optimizers. At first glance it doesn't seem all that interesting, since there have been quite a few similar works before, and most of the results are rather bland. However, on closer reading I found there was more to it than meets the eye: it turns out the authors, using thousands of TPU-hours of compute combined with manual intervention, discovered an optimizer called Lion (EvoLved Sign Momentum — I have to say the name feels a bit forced) that is faster and more memory-efficient. They ran extensive experiments on image classification, image-text matching, diffusion models, language model pretraining and fine-tuning, and many other tasks, and in most of these tasks Lion outperformed mainstream optimizers such as AdamW.
Better memory efficiency and better performance — truly having your cake and eating it too. What kind of optimizer could have such formidable properties? Let's take a look at the paper's results together.
Results First
This post is mainly concerned with the discovered optimizer itself, so I won't go into the details of the search process — interested readers can consult the original paper. The update rule of the Lion optimizer is
\begin{equation}\text{Lion}:=\left\{\begin{aligned} &\boldsymbol{u}_t = \text{sign}\big(\beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\big) \\ &\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}) \\ &\boldsymbol{m}_t = \beta_2 \boldsymbol{m}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t \end{aligned}\right.\end{equation}more
where $\boldsymbol{g}_t = \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{t-1})$ is the gradient of the loss function and $\text{sign}$ is the sign function, which maps positive numbers to 1 and negative numbers to -1. We can compare this to the update process of the currently mainstream optimizer AdamW:
\begin{equation}\text{Adam}\color{skyblue}{\text{W}}:=\left\{\begin{aligned} &\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\ &\boldsymbol{v}_t = \beta_2 \boldsymbol{v}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t^2\\ &\hat{\boldsymbol{m}}_t = \boldsymbol{m}_t\left/\left(1 - \beta_1^t\right)\right.\\ &\hat{\boldsymbol{v}}_t = \boldsymbol{v}_t\left/\left(1 - \beta_2^t\right)\right.\\ &\boldsymbol{u}_t =\hat{\boldsymbol{m}}_t\left/\left(\sqrt{\hat{\boldsymbol{v}}_t} + \epsilon\right)\right.\\ &\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}) \end{aligned}\right.\end{equation}
The comparison is quite clear: Lion has fewer hyperparameters than AdamW (missing one $\epsilon$), caches one fewer set of parameters $\boldsymbol{v}$ (hence saving memory), and removes the division and square-root operations that dominate the compute cost in AdamW's update (hence being faster).
Before this, the optimizer most similar to Lion was probably SIGNUM, whose update process is
\begin{equation}\text{SIGNUM}:=\left\{\begin{aligned} &\boldsymbol{m}_t = \beta \boldsymbol{m}_{t-1} + \left(1 - \beta\right) \boldsymbol{g}_t \\ &\boldsymbol{u}_t = \text{sign}\big(\boldsymbol{m}_t\big) \\ &\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t \boldsymbol{u}_t \end{aligned}\right.\end{equation}
Like Lion, SIGNUM also applies the sign function to the update, and is even simpler than Lion (equivalent to a special case of Lion with $\beta_1=\beta_2$ and $\lambda_t=0$). Unfortunately, SIGNUM did not achieve better performance — it was designed mainly to reduce transmission costs in distributed computing. Lion's update rule differs in an important way, in particular placing the momentum update after the variable update, and extensive experiments have demonstrated its advantage in effectiveness.
Experiments in the Paper
As mentioned at the beginning of this post, Lion has been tested on quite a number of tasks, with a wealth of experimental results. Below I list some of the results I consider most important.
Lion's results on NLU and NLG tasks, mostly better than AdamW and Adafactor
Comparison of Lion against numerous optimizers on Vision Transformers
On CV classification tasks, Lion converges faster
On NLP autoregressive generation, Lion converges faster
The upper-right figure shows training curves on ImageNet, showing that while Lion achieves better validation performance, its training-set performance is not necessarily better than AdamW's
Hyperparameter Settings
Seeing such striking results, I couldn't help but want to try it myself. Before running experiments, of course, one needs to understand how the various hyperparameters should be set. First, for $\beta_1,\beta_2$, the combination automatically found by the paper is $\beta_1=0.9,\beta=0.99$, which was reused in most experiments, but for NLP tasks the combination $\beta_1=0.95,\beta_2=0.98$ was used instead (the paper's detailed experimental configurations are given in Table 12 on the last page).
As for the more critical learning rate $\eta$ and weight decay rate $\lambda$: since every component of Lion's update $\boldsymbol{u}$ has absolute value exactly 1, this is typically much larger than in AdamW, so the learning rate needs to be reduced by more than a factor of 10 to achieve roughly the same update magnitude. And since the learning rate has been lowered, in order to keep the magnitude of weight decay unchanged, the weight decay rate should be scaled up correspondingly. The last page of the original paper gives reference hyperparameter values for each experiment: for small (Base-scale) models, $\eta = 3\times 10^{-4}$ and $\lambda=0.01$ are used, while for large models (with over a billion parameters), the learning rate is reduced further to $\eta = 2\times 10^{-4}$ or even $\eta = 10^{-4}$.
In fact, we previously derived a scheme for jointly setting the learning rate and weight decay rate in Some "Training Strategies" Derived from the Ideas Behind the Amos Optimizer, and it's most convenient to refer to that scheme here. In that scheme the update is written as (the notation differs slightly from what's used above, but shouldn't cause confusion, so I won't try to force them into full consistency)
\begin{equation}\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - (\alpha_t \boldsymbol{u}_t + \rho_t\boldsymbol{\theta}_t)\end{equation}
where
\begin{equation}\alpha_t \approx \frac{\alpha_0\Vert\boldsymbol{\varepsilon}_0\Vert}{\Vert\boldsymbol{u}_t\Vert} \frac{1}{\kappa t + 1},\quad \rho_t \approx \frac{\alpha_0^2}{2q} \frac{1}{\kappa t + 1}\end{equation}
Here $\boldsymbol{u}_t$ is the original update quantity; $\alpha_0$ is the (initial-stage) relative magnitude of the parameter change, typically on the order of $10^{-3}$, meaning that the relative change in the norm of the parameters after each update step is roughly one part in a thousand; $q$ is a hyperparameter that, absent special circumstances, can simply be set to 1; and $\kappa$ is a hyperparameter controlling the decay speed of the learning rate, which can be set according to factors such as the size of the training data.
Since $\boldsymbol{u}_t$ has gone through the $\text{sign}$ operation, we have $\Vert\boldsymbol{u}_t\Vert=\sqrt{k}$, where $k$ is the dimensionality of the parameter; and $\Vert\boldsymbol{\varepsilon}_0\Vert\approx\sqrt{k}\sigma$ — this we already derived in Some "Training Strategies" Derived from the Ideas Behind the Amos Optimizer, where $\sigma$ is the scale of the parameter's variation, and for multiplicative matrices, $\sigma^2$ is simply its initialization variance. So, after a series of simplifications, we get
\begin{equation}\alpha_t \approx \frac{\alpha_0\sigma}{\kappa t + 1},\quad \rho_t \approx \frac{\alpha_0^2}{2(\kappa t + 1)}\end{equation}
Here $\alpha_t$ is the same as the earlier $\eta_t$, and $\lambda_t = \rho_t / \alpha_t = \alpha_0 / 2\sigma$. Taking BERT base's $d=768$ as an example, the initialization variance is roughly on the order of $1/d$, so $\sigma = \sqrt{1/d}\approx 0.036$; suppose $\alpha_0$ is taken to be $1.11 \times 10^{-3}$ (just to round the result to a nice number), then according to the formula above, the learning rate is approximately $4\times 10^{-5}$ and the decay rate is approximately $0.015$. In my own MLM pretraining experiments, this combination worked quite well.
Personal implementation: https://github.com/bojone/bert4keras
Further Thoughts
Overall, Lion is quite impressive — whether in the original paper or in my own experiments, it holds its own against AdamW. Combined with its speed and memory advantages, it's conceivable that it could earn itself a place among the mainstream optimizers of the future.
Since Adam was proposed, its property of fast convergence has made it the default optimizer for many models. Some scholars have even suggested that this could create a kind of evolutionary effect in reverse: all model improvements tend to evolve in directions favorable to Adam. In other words, because we chose Adam as our optimizer, we may have discarded many changes that are genuinely effective in practice but happen not to work well with Adam, leaving only the improvements that are favorable to Adam. A detailed discussion of this can be found in NEURAL NETWORKS (MAYBE) EVOLVED TO MAKE ADAM THE BEST OPTIMIZER. Given this backdrop, discovering an optimizer that is simpler and more effective than Adam is quite a remarkable achievement, even if it was found through massive computational search.
Readers might wonder: why is Lion able to achieve better generalization performance? The original paper's explanation is that the $\text{sign}$ operation introduces extra noise (compared to using exact floating-point values), which pushes the model into a flatter (though not necessarily lower) region of the loss landscape, leading to better generalization. To verify this, the authors compared the robustness to perturbations of model weights trained by AdamW versus Lion, and found that Lion's weights were more robust to perturbation. However, in theory, this only demonstrates that Lion indeed reaches a flatter region — it doesn't prove that this result is actually caused by the $\text{sign}$ operation. That said, even after all these years, the mechanism behind Adam's success still isn't fully understood, and Lion has only just been proposed, so we shouldn't be too nitpicky about it.
My own guess is that the $\text{sign}$ operation in Lion treats every component equally, allowing the model to fully exploit the contribution of every component, which leads to better generalization. With SGD, the update magnitude is proportional to the gradient, but some components may have small gradients simply because they weren't initialized well, not because they're unimportant. Lion's $\text{sign}$ operation, in effect, gives every parameter a chance to "revive" or even "shine again." In fact, one can show that in the early stages of training, Adam's updates are also close to $\text{sign}$, and it only gradually deviates from this as training progresses.
Is Lion perfect, then? Clearly not. For instance, the original paper points out that it performs worse than AdamW with small batch sizes (below 64), which isn't hard to understand: the $\text{sign}$ operation already introduces noise, and a small batch size adds even more noise on top of that. Noise is something that needs to be kept at just the right level — too much of it, stacked together, can easily degrade performance. Also, precisely because $\text{sign}$ amplifies the noise in the optimization process, improperly set hyperparameters can more easily cause the loss to blow up or otherwise diverge; in such cases, it may help to introduce warmup, or increase the number of warmup steps. Furthermore, Lion still needs to cache a momentum parameter, so its memory usage is higher than AdaFactor — whether this part of the parameter count can be further optimized remains to be seen.
Summary
This post introduced Google's newly proposed optimizer, Lion, discovered through large-scale computational search combined with manual intervention. Compared to the mainstream AdamW, it is faster and more memory-efficient, and extensive experimental results show that on most tasks it performs on par with, or even better than, AdamW.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.