Source-linked AI summary
Symbolic Discovery of Optimization Algorithms
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, Quoc V. Le
TL;DR
Optimizer discovery is difficult because existing approaches either generalize poorly from small tasks or constrain the space of possible algorithms. The paper formulates discovery as program search and combines evolutionary search, validation-based selection, and simplification. It discovers Lion, a momentum-only sign-update optimizer with strong results across tasks, though improvements are smaller or nonsignificant in several settings.
Problem
Existing optimizer-discovery methods either generalize poorly from small tasks or restrict the search space, while high-performing programs are sparse in the broader program space.
Method
The paper formulates optimizer discovery as program search and uses evolutionary search, abstract execution, program selection, and simplification to address search and generalization challenges.
Results
Lion achieves strong or similar performance across vision, vision-language, diffusion, and language tasks, including up to 5x JFT pre-training compute savings and 2.3x diffusion training efficiency.
Takeaways & Limitations
Lion offers a simple, memory-efficient alternative to Adam and adaptive optimizers by tracking momentum and applying sign-based updates.
Takeaways & Limitations
The search is biased toward first-order algorithms, lacks constructs for advanced second-order methods, remains costly, and requires manual simplification.
Abstract
from arXiv · showhide
We present a method to formulate algorithm discovery as program search, and apply it to discover optimization algorithms for deep neural network training. We leverage efficient search techniques to explore an infinite and sparse program space. To bridge the large generalization gap between proxy and target tasks, we also introduce program selection and simplification strategies. Our method discovers a simple and effective optimization algorithm, $\textbf{Lion}$ ($\textit{Evo$\textbf{L}$ved S$\textbf{i}$gn M$\textbf{o}$me$\textbf{n}$tum}$). It is more memory-efficient than Adam as it only keeps track of the momentum. Different from adaptive optimizers, its update has the same magnitude for each parameter calculated through the sign operation. We compare Lion with widely used optimizers, such as Adam and Adafactor, for training a variety of models on different tasks. On image classification, Lion boosts the accuracy of ViT by up to 2% on ImageNet and saves up to 5x the pre-training compute on JFT. On vision-language contrastive learning, we achieve 88.3% $\textit{zero-shot}$ and 91.1% $\textit{fine-tuning}$ accuracy on ImageNet, surpassing the previous best results by 2% and 0.1%, respectively. On diffusion models, Lion outperforms Adam by achieving a better FID score and reducing the training compute by up to 2.3x. For autoregressive, masked language modeling, and fine-tuning, Lion exhibits a similar or better performance compared to Adam. Our analysis of Lion reveals that its performance gain grows with the training batch size. It also requires a smaller learning rate than Adam due to the larger norm of the update produced by the sign function. Additionally, we examine the limitations of Lion and identify scenarios where its improvements are small or not statistically significant. Lion is also successfully deployed in production systems such as Google search ads CTR model.
1 Introduction
The paper frames optimizer discovery as program search and introduces Lion, a momentum-only optimizer using sign-based updates. Lion shows strong results across vision, vision-language, diffusion, and language tasks, while requiring smaller learning rates and having task-dependent limitations.
- Lion: Lion tracks only momentum and uses sign-based updates, producing uniform update magnitudes with lower memory overhead than adaptive optimizers.Its default momentum coefficients are β1 = 0.9 and β2 = 0.99, with decoupled weight decay.
- Motivation: Existing optimizer-discovery methods struggle to generalize beyond small tasks or restrict the search space, limiting potential discoveries such as modified momentum tracking.The paper contrasts black-box learning-to-optimize methods with tree-based searches using fixed operands and bounded structures.
- Approach: The proposed method combines evolutionary program search with warm-starting, restarts, abstract execution, funnel selection, and program simplification.These techniques target both the infinite, sparse search space and the generalization gap between proxy and target tasks.
- Results: 88.3% zero-shot and 91.1% fine-tuning accuracy were achieved on ImageNet in vision-language contrastive learning, surpassing previous best results by 2% and 0.1%, respectively.Lion replaced Adafactor in BASIC and was applied during both image-encoder pre-training and contrastive learning.
- Results: Lion reduces JFT pre-training compute by up to 5x, improves diffusion training efficiency by 2.3x, and matches or exceeds Adam on several language-modeling tasks.The reported gains span image classification, diffusion, autoregressive modeling, masked language modeling, and fine-tuning.
- Properties and limitations: Lion usually requires a smaller learning rate and larger decoupled weight decay than Adam because sign-based updates have a larger norm.Its gains increase with batch size, but differences from AdamW are not statistically significant on some large-scale language and image-text datasets.
2 Symbolic Discovery of Algorithms
The paper represents optimizers as executable symbolic programs and searches this infinite, sparse space with evolutionary methods and pruning. It uses proxy-to-meta-validation selection to identify programs that generalize beyond small search tasks.
- Program Search Space: Symbolic programs are designed to be flexible, analyzable, transferable, and focused on high-level algorithmic design rather than implementation details.The representation uses an imperative language over n-dimensional arrays, similar to NumPy, JAX, and optimizer pseudocode.
- Program Search Space: The train function takes weights, gradients, and the current learning-rate schedule, returns a weight update, and may maintain historical variables.The signature preserves memory footprints no larger than AdamW while allowing discovered algorithms to update auxiliary variables differently.
- Program Search Space: Programs are sequences of assignment statements built from common mathematical functions, and mutations insert, delete, or modify statements and arguments.The search space has no fixed restriction on statement count or local variables.
- Search Challenges: The program space is infinite and sparse: after evaluating over 2M programs on a low-cost proxy, the best random-search program remained significantly inferior to AdamW.Sparsity, rather than only the number of possible programs, makes effective search difficult.
- Efficient Search Techniques: Warm-started regularized evolution uses AdamW initialization, tournament selection, mutations, and restarts to improve exploration and search efficiency.Restarting from the initial program explores different local optima, while later restarting from the best discovered program can continue progress.
- Efficient Search Techniques: Abstract execution prunes invalid, functionally duplicate, and redundant programs before full evaluation by checking types and shapes and caching semantic computations.Functional hashing identifies equivalent input-to-output computations for reuse.
- Generalization: Proxy evaluation uses smaller models, datasets, and training horizons, while meta-validation uses larger models and more steps to select programs that generalize.Search experiments generate 200–300K programs, evaluate roughly 20–30K after caching, and cost about 3K TPU V2 days with restarts.
- Generalization: Target tasks can be > 10^4 larger than proxy tasks, and later meta-overfitting tends to produce programs with better generalization.The gap includes days of training on 512 TPU V4 chips for targets versus roughly 20 minutes on one TPU V2 chip for proxies.
3 Derivation and Analysis of Lion
The analysis derives Lion through search, simplification, and comparison with related optimizers, then examines its update, momentum, hyperparameters, regularization, and resource requirements. Lion combines sign-based uniform updates with momentum tracking and achieves strong performance while reducing memory overhead.
- Update design: Lion tracks momentum and applies the sign operation, producing uniform update magnitudes across dimensions with lower memory overhead than adaptive optimizers.Its design retains momentum while avoiding the second-moment tracking used by popular adaptive methods.
- Derivation: Lion emerges from funnel selection and simplification of a searched program after redundant statements and functions are removed.The final algorithm is obtained by simplifying Program 4, itself derived from the raw searched program.
- Sign update and regularization: The sign update acts as a form of regularization: Lion has higher training error than AdamW but 2% higher validation accuracy for ViT-B/16 on ImageNet.The analysis also reports convergence in smoother regions, which usually results in better generalization.
- Momentum tracking: Lion uses β2 = 0.99 for longer momentum history and β1 = 0.9 to balance current gradients against that history before the sign operation.The β2 choice remembers approximately 10x longer gradient history than the commonly used 0.9 factor.
- Hyperparameter and batch size choices: Lion has fewer hyperparameters than AdamW and Adafactor, but its larger update norm requires a smaller learning rate and larger decoupled weight decay.The effective weight decay strength is maintained through the product lr * λ.
- Memory and runtime benefits: AdamW requires at least 16 TPU V4 chips while Lion requires 8 for ViT-B/16 at resolution 224 and batch size 4,096.The comparison uses bfloat16 momentum and illustrates Lion’s lower memory footprint for large-model training.
4 Evaluation of Lion
Lion is evaluated against AdamW and other optimizers across image, vision-language, diffusion, language-modeling, and fine-tuning tasks. It often improves accuracy or efficiency, with gains associated with larger models or batch sizes, while its ablations and operating requirements clarify the mechanism.
- Image Classification: Lion significantly outperforms AdamW on ImageNet classification, with gains of 1.96% for ViT-B/16 and 0.58% for ViT-S/16.The improvement shrinks under strong augmentations but remains 0.42% on CoAtNet-3.
- Image Classification: Lion preserves its advantage after ImageNet-21K pre-training, with larger-model gains of +0.52% on ImageNet, +0.57% on ReaL, and +0.74% on V2.The corresponding ViT-B/16 gains are +0.33%, +0.23%, and +0.25%.
- Image Classification: Lion matches larger AdamW-trained ViT models with less compute or fewer parameters, including 5x lower JFT pre-training cost on ImageNet ReaL.ViT-L/16 matches ViT-H/14 on ImageNet and ImageNet V2 with 3x less pre-training cost, while JFT-3B results improve with 1.8x fewer parameters.
- Vision-Language Contrastive Learning: Lion improves vision-language performance, reaching 88.3% zero-shot and 91.1% fine-tuning accuracy on ImageNet in BASIC-L.These results exceed the previous best by 2.0% and 0.1%, respectively.
- Diffusion Model: On diffusion image synthesis, Lion achieves better FID and faster convergence than AdamW, reaching AdamW’s final performance at 440K steps and reducing iterations 2.3x.For 256 × 256 generation, final FID is 4.1 for Lion versus 4.7 for AdamW.
- Language Modeling and Ablations: Language-modeling results are similar or slightly better than AdamW, while Lion’s 2.5% gain over AdamW appears at batch size 32K.Masked language modeling improves validation perplexity, whereas large-scale autoregressive training shows no perplexity difference throughout training.
- Language Modeling and Ablations: Lion’s sign-based update and momentum design matter: simple momentum sign ablations perform worse, and Lion prefers larger batches while requiring smaller learning rates and larger decoupled weight decay.The sign operation produces uniform-magnitude updates with a larger norm, motivating these hyperparameter requirements.
5 Hyperparameter Tuning
Lion requires different hyperparameter scaling from AdamW because its sign-based updates have a larger norm, while experiments indicate greater robustness to hyperparameter choices.
- Lion uses β1 = 0.9 and β2 = 0.99 by default, while AdamW uses β1 = 0.9 and β2 = 0.999.
- Lion typically uses a learning rate 3-10x smaller than AdamW because sign-based updates have a larger norm.Its weight decay λ is correspondingly 3-10x larger to maintain similar effective weight-decay strength.
- Example settings use lr = 1e-4 and λ = 10.0 for Lion versus lr = 1e-3 and λ = 1.0 for AdamW on ViT-B/16.
- Lion is more robust than AdamW to different learning-rate and weight-decay choices when training ViT-B/16 on ImageNet.
6 Limitations
The paper identifies limitations in both the search process and Lion’s empirical scope, including search-space bias, task-dependent gains, small-batch limitations, and momentum memory costs.
- Limitations of search: The search space is biased toward popular first-order optimizers and lacks functions needed to construct advanced second-order algorithms.
- Limitations of search: The search cost remains large, and simplifying discovered programs requires manual intervention.
- Limitations of Lion: Lion’s evaluation is limited to the chosen tasks, and its improvements are small or similar to AdamW on several large, high-quality datasets and models.
- Limitations of Lion: Lion may perform no better than AdamW when the batch size is small, below 64.
- Limitations of Lion: Tracking momentum in bfloat16 can be expensive for giant models, although factorizing the momentum is suggested as a possible solution.
7 Related Work
The paper situates its work within AutoML and meta-learning, alongside learning to learn, neural architecture search, and hyperparameter optimization.
- The work lies within AutoML and meta-learning research spanning learning to learn, neural architecture search, and hyperparameter optimization.
8 Conclusion
The paper proposes discovering optimization algorithms through program search and develops techniques for searching sparse spaces and bridging proxy-to-target generalization. It discovers Lion, a memory-efficient optimizer with strong generalization across architectures, datasets, and tasks.
- The paper discovers optimization algorithms through program search while addressing infinite sparse search spaces and large proxy-to-target generalization gaps.
- Lion is presented as a simple, effective, memory-efficient optimizer with strong generalization across architectures, datasets, and tasks.
B Image Classification Tasks
The image-classification evaluation spans standard, robustness, transfer, and fine-grained benchmarks. It includes ImageNet and several related datasets for assessing model performance across varied settings.
- Benchmark coverage: The evaluation includes ImageNet, ImageNet ReaL, and ImageNet V2.
- Benchmark coverage: Robustness-oriented benchmarks include ImageNet A, ImageNet R, ImageNet Sketch, and ObjectNet.
- Benchmark coverage: Additional evaluations use CIFAR-100 and Oxford-IIIT Pet.
C NLP Tasks
The NLP section evaluates large-scale language models across NLG and NLU tasks while also presenting optimizer programs and learning-curve analyses. The listed tasks span question answering, reasoning, comprehension, and language understanding.
- NLP evaluation: NLU evaluations cover commonsense reasoning, reading comprehension, inference, and multiple classification tasks.
- NLP evaluation: NLG evaluations include TriviaQA, Natural Questions, and Web Questions.
- Discovered algorithms: Program 6 tracks the second moment without EMA decay, matching AdaGrad.
- Discovered algorithms: Program 7 tracks the second moment using the difference between gradient and momentum, resembling AdaBelief.
- Learning curves: Figure 10 compares learning curves for ViT-S/16 and ViT-B/16, with five adaptive optimizers showing similar curves.
- Search behavior: Reducing proxy-task data favors algorithms with better regularization, while reducing search progress favors simple AdamW variants.
E Architecture Details for Language Modeling
The language-modeling architecture uses a feed-forward layer sized at four times the model dimension and different vocabulary sizes for small- and large-scale models.
- Architecture: The Transformer feed-forward layer has dimension 4 × dmodel.
- Vocabulary: Small-scale language models use a 32K vocabulary.
- Vocabulary: Large-scale language models use a 256K vocabulary.
F Details of Proxy Tasks
Proxy tasks use small vision and language models trained on reduced datasets for rapid optimizer evaluation. The section also reports validation perplexity for masked language modeling and training loss for ViT-B/16 on ImageNet.
- Vision proxy task: Vision proxy training uses a three-layer ViT with 96 hidden units and three heads on 10% ImageNet for 30k steps.The batch size is 64, with 64 × 64 images and patch size 16.
- Language proxy task: Language proxy training uses a two-layer Transformer with 128 hidden units and two heads on LM1B for 20K steps.The setup uses batch size 64, sequence length 32, and vocabulary size 3K.
- Evaluation curves: Figure 11 reports validation perplexity for masked language modeling on C4 and training loss for ViT-B/16 on ImageNet.
G Analysis of Loss Landscape
The analysis connects Lion’s generalization to loss-landscape geometry and the regularizing effect of sign-based updates. Lion reaches flatter regions and can improve validation accuracy despite higher training error than AdamW.
- Loss-landscape analysis: Lion helps ViT converge to flatter regions, retaining comparably lower error under Gaussian perturbations than AdamW.The comparison uses perturbed training loss averaged over 1K random noises.
- Generalization: AdamW achieves smaller training error, while Lion achieves better validation accuracy despite higher training error.This pattern is presented as evidence that flatter solutions may benefit generalization.
- Language-modeling analysis: On PG-19, Lion remains the best-performing optimizer across small, medium, and large Transformer models.The language-modeling comparison matches Lion’s β settings with corresponding ablations.