Symbolic Discovery of Optimization Algorithms

Abstract

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, Lion (EvoLved Sign Momentum). 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% zero-shot and 91.1% 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. The implementation of Lion is publicly available. 1 * Work done as a student researcher at Google Brain. † Work done while at Google.

Robust fine-tuning of zero-shot models

“…• Another question is, why not use an optimizer such as Lion [7] which does not divide updates by any value, and is therefore immune to the stuck-in-the-past scenario. We believe this may be a promising path forward.…”

“…We study these two directions in the context of contrastive language-image pre-training (CLIP) [44]. We examine CLIP-style models because of their importance in computer vision: CLIP-style models reach state-ofthe-art performance on a wide range of image classification tasks [44,63,42,7] and underlie image generation methods such as DALLE•2 [47] and Stable Diffusion [49]. Our contributions towards fast training and stable training are as follows.…”

EVA-02: A Visual Representation for Neon Genesis

“…Precisions and optimizers. In Table 18, we show that sufficiently pre-trained EVA-02 representations are robust enough that can be fine-tuned using various numerical precisions (e.g., fp16 and bf16) and optimizers (e.g., Lion [25], AdamW [64,84], and SGD [87]). Remarkably, the fine-tuning can be done using the SGD optimizer with only little performance drop.…”

“…The final IN-1K fine-tuning for all-sized models (including EVA-02-Ti and -S) can be done without using strong regularization such as cutmix [141], mixup [143] and random erasing [146]. In the Appendix, we show that our pre-trained representations are robust enough that can be fine-tuned using various numerical precisions (e.g., fp16 and bf16) and optimizers (e.g., Lion [25], AdamW [64,84], and SGD [87]). Remarkably, the fine-tuning can be done even using the SGD optimizer with only 0.1-point performance drop.…”

A Survey on Efficient Training of Transformers

“…By default, AdamW [61], a variant of Adam which decouples the L 2 regularization and the weight decay, is the most widely used optimizer for Transformers. More recently, Google searches optimization algorithms and discovers a simple and effective optimizer called Lion [18]. Lion only keeps track of the momentum with the first-order gradient, and its update only considers the sign direction and has the same magnitude for each parameter, which is very different from the adaptive optimizers like AdamW.…”