Symbolic Discovery of Optimization Algorithms — part 1/4

Cleaned-text code points [0, 15886). Bibliography omitted here; source.md includes it.

# 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 ( Evo L ved S i gn M o me 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% 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 1 1 https://github.com/google/automl/tree/master/lion . Lion is also successfully deployed in production systems such as Google’s search ads CTR model. † † Correspondence: xiangning@cs.ucla.edu , crazydonkey@google.com .

## 1 Introduction

Optimization algorithms, i.e., optimizers, play a fundamental role in training neural networks. There are a large number of handcrafted optimizers, mostly adaptive ones, introduced in recent years ( Zhuang et al., 2020 ; Balles and Hennig, 2018 ; Liu et al., 2020 ; Bernstein et al., 2018 ; Dozat, 2016 ; Anil et al., 2020 ) . However, Adam ( Kingma and Ba, 2014 ) with decoupled weight decay ( Loshchilov and Hutter, 2019 ) , also referred to as AdamW, and Adafactor with factorized second moments ( Shazeer and Stern, 2018 ) , are still the de facto standard optimizers for training most deep neural networks, especially the recent state-of-the-art language ( Brown et al., 2020 ; Vaswani et al., 2017 ; Devlin et al., 2019 ) , vision ( Dosovitskiy et al., 2021 ; Dai et al., 2021 ; Zhai et al., 2021 ) and multimodal ( Radford et al., 2021 ; Saharia et al., 2022 ; Yu et al., 2022 ) models.

Discovered optimizer Lion. β 1 = 0.9 \beta_{1}=0.9 and β 2 = 0.99 \beta_{2}=0.99 by default are derived from Program . It only tracks momentum and uses the sign operation to compute the update. The two gray lines compute the standard decoupled weight decay, where λ \lambda is the strength.

Another direction is to automatically discover such optimization algorithms. The learning to optimize (L2O) approach proposes to discover optimizers by training parameterized models, e.g., neural networks, to output the updates ( Andrychowicz et al., 2016 ; Metz et al., 2019 ; Li and Malik, 2017 ; Metz et al., 2022 ) . However, those black-box optimizers, typically trained on a limited number of small tasks, struggle to generalize to state-of-the-art settings where much larger models are trained with significantly more training steps. Another line of methods ( Bello et al., 2017 ; Wang et al., 2022 ) apply reinforcement learning or Monte Carlo Sampling to discover new optimizers, where the search space is defined by trees composed from predefined operands (e.g., gradient and momentum) and operators (e.g., unary and binary math operations). However, to make the search manageable, they often limit the search space by using fixed operands and restricting the size of the tree, thereby limiting the potential for discovery. For example, they are unable to modify the tracking of momentum or how it contributes to the update, which is an essential component of Lion. Consequently, the algorithms discovered have not yet reached the state-of-the-art. AutoML-Zero ( Real et al., 2020 ) is an ambitious effort that attempts to search every component of a machine learning pipeline while evaluating on toy tasks. This work follows the research direction of automatic discovering optimizers and is in particular inspired by AutoML-Zero, but aims at discovering effective optimization algorithms that can improve the state-of-the-art benchmarks.

In this paper, we present a method to formulate algorithm discovery as program search and apply it to discover optimization algorithms. There are two primary challenges. The first one is to find high-quality algorithms in the infinite and sparse program space. The second one is to further select out the algorithms that can generalize from small proxy tasks to much larger, state-of-the-art tasks. To tackle these challenges, we employ a range of techniques including evolutionary search with warm-start and restart, abstract execution, funnel selection, and program simplification.

Our method discovers a simple and effective optimization algorithm: Lion, short for Evo L ved S i gn M o me n tum . This algorithm differs from various adaptive algorithms by only tracking momentum and leveraging the sign operation to calculate updates, leading to lower memory overhead and uniform update magnitudes across all dimensions. Despite its simplicity, Lion demonstrates outstanding performance across a range of models (Transformer, MLP, ResNet, U-Net, and Hybrid) and tasks (image classification, vision-language contrastive learning, diffusion, language modeling, and fine-tuning). Notably, we achieve 88.3% zero-shot and 91.1% fine-tuning accuracy on ImageNet by replacing Adafactor with Lion in BASIC ( Pham et al., 2021 ) , surpassing the previous best results by 2% and 0.1%, respectively. Additionally, Lion reduces the pre-training compute on JFT by up to 5x, improves training efficiency on diffusion models by 2.3x and achieves a better FID score, and offers similar or better performance on language modeling with up to 2x compute savings.

We analyze the properties and limitations of Lion. Users should be aware that the uniform update calculated using the sign function usually yields a larger norm compared to those generated by SGD and adaptive methods. Therefore, Lion requires a smaller learning rate l ​ r lr , and a larger decoupled weight decay λ \lambda to maintain the effective weight decay strength. For detailed guidance, please refer to Section 5 . Additionally, our experiments show that the gain of Lion increases with the batch size and it is more robust to different hyperparameter choices compared to AdamW. For limitations, the difference between Lion and AdamW is not statistical significant on some large-scale language and image-text datasets. The advantage of Lion is smaller if using strong augmentations or a small batch size ( < < 64) during training. See Section 6 for details.

An example training loop, where the optimization algorithm that we are searching for is encoded within the train function. The main inputs are the weight ( w ), gradient ( g ) and learning rate schedule ( lr ). The main output is the update to the weight. v1 and v2 are two additional variables for collecting historical information.

Initial program (AdamW). The bias correction and ϵ \epsilon are omitted for simplicity.

Discovered program after search, selection and removing redundancies in the raw Program . Some variables are renamed for clarity.

## 2 Symbolic Discovery of Algorithms

We present an approach that formulates algorithm discovery as program search ( Koza, 1994 ; Brameier et al., 2007 ; Real et al., 2020 ) . We use a symbolic representation in the form of programs for the following advantages: (1) it aligns with the fact that algorithms must be implemented as programs for execution; (2) symbolic representations like programs are easier to analyze, comprehend and transfer to new tasks compared to parameterized models such as neural networks; (3) program length can be used to estimate the complexity of different programs, making it easier to select the simpler, often more generalizable ones. This work focuses on optimizers for deep neural network training, but the method is generally applicable to other tasks.

### 2.1 Program Search Space

We adhere to the following three criteria while designing the program search space: (1) the search space should be flexible enough to enable the discovery of novel algorithms; (2) the programs should be easy to analyze and incorporate into a machine learning workflow; (3) the programs should focus on the high-level algorithmic design rather than low-level implementation details. We define the programs to contain functions operating over n-dimensional arrays, including structures like lists and dictionaries containing such arrays, in an imperative language. They are similar to Python code using NumPy / JAX ( Harris et al., 2020 ; Bradbury et al., 2018 ) as well as pseudo code of optimization algorithms. The details of the design are outlined below, with an example representation of AdamW in Program .

Input / output signature The program defines a train function, which encodes the optimization algorithm being searched for, where the main inputs are the model weight ( w ), the gradient ( g ) and the learning rate schedule value ( lr ) at the current training step. The main output is the update to the weight. The program also incorporates extra variables initialized as zeros to collect historical information during training. For example, AdamW requires two extra variables to estimate first and second moments. Note that those variables can be used arbitrarily, we use the name m and v in Program just for better readability. This simplified code snippet in Program uses the same signature as AdamW to ensure that the discovered algorithms have smaller or equal memory footprints. As opposed to previous optimizer search attempts ( Bello et al., 2017 ; Wang et al., 2022 ) , our method allows discovering better ways of updating the extra variables.

Building blocks The train function consists of a sequence of assignment statements, with no restrictions on the number of statements or local variables. Each statement calls a function using constants or existing variables as inputs, and the resulting value is stored in a new or existing variable. For the program, we select 45 common math functions, most of which corresponds to a function in NumPy or an operation in linear algebra. Some functions are introduced to make the program more compact, such as the linear interpolation function interp(x, y, a) , which is made equivalent to (1 - a) * x + a * y . Preliminary experiments have investigated the inclusion of more advanced features such as conditional and loop statements, and defining and calling new functions, but these do not yield improved results, so we leave them out. A detailed description of the functions are summarized in Appendix H . When necessary, the types and shapes of the function arguments are automatically cast, e.g., in the case of adding a dictionary of arrays to a scalar.

Mutations and redundant statements The design of mutations utilized in evolutionary search is tightly intertwined with the representation of the program. We include three types of mutations: (1) inserting a new statement at a random location with randomly chosen functions and arguments, (2) deleting a random chosen statement, and (3) modifying a random statement by randomly altering one of its function arguments, which may be either variables or constants. To mutate an argument, we replace it with an existing variable or a newly generated constant obtained by sampling from a normal distribution X ∼ 𝒩 ⁡ ( 0 1 ) X\sim\mathcal{N}(0\,\,1) . Additionally, we can mutate an existing constant by multiplying it by a random factor 2 a 2^{a} , where a ∼ 𝒩 ⁡ ( 0 1 ) a\sim\mathcal{N}(0\,\,1) . These constants serve as tunable hyperparameters in the optimization algorithm, such as the peak learning rate and weight decay in AdamW. Note that we allow a program to include redundant statements during search, i.e., statements that do not impact the final program outputs. This is necessary as mutations are limited to only affecting a single statement. Redundant statements therefore serve as intermediate steps towards future substantial modifications in the program.

Infinite and sparse search space Given the limitless number of statements and local variables, as well as the presence of mutable constants, the program search space is infinite. Even if we ignore the constants and bound the program length and number of variables, the number of potential programs is still intractably large. A rough estimate of the number of possible programs is n p = n f l ​ n v n a ∗ l n_{p}=n_{f}^{l}n_{v}^{n_{a}*l} , where n f n_{f} is the number of possible functions, n v n_{v} is the number of local variables, n a n_{a} is the average number of arguments per statement, and l l is the program length. More importantly, the challenge comes from the sparsity of high-performing programs in the search space. To illustrate this point, we conduct a random search that evaluates over 2M programs on a low-cost proxy task. The best program among them is still significantly inferior to AdamW.

### 2.2 Efficient Search Techniques

We employ the following techniques to address the challenges posed by the infinite and sparse space.

Evolution with warm-start and restart We apply regularized evolution as it is simple, scalable, and has shown success on many AutoML search tasks ( Real et al., 2020 ; Real et al., 2019 ; Ying et al., 2019 ; So et al., 2019 ; Holland, 1992 ) . It keeps a population of P P algorithms that are gradually improved through cycles. Each cycle picks T < P T\!\!<\!P algorithms at random and the best performer is chosen as the parent , i.e., tournament selection ( Goldberg and Deb, 1991 ) . This parent is then copied and mutated to produce a child algorithm, which is added to the population, while the oldest algorithm is removed. Normally, evolutionary search starts with random candidates, but we warm-start the initial population as AdamW to accelerate the search. By default, we use a tournament size of two and a population size of 1K. To further improve the search efficiency, we apply two types of restart: (1) restarting from the initial program, which can lead to different local optima due to the randomness in evolution and encourage exploration. This can be done by running multiple searches in parallel. (2) restarting from the best algorithm found thus far to further optimize it, encouraging exploitation. Figure 2 (Left) displays the mean and standard error of five evolutionary search experiments. We run hyperparameter tuning based on AdamW by only allowing mutations of constants in the evolution, and run random search by sampling random programs, both with 4x more compute. Our search significantly outperforms the best results achieved by both baselines, demonstrated as the two dashed lines in the figure. The high variance in the search fitness necessitates running multiple repeats through restarting from the initial program. When the search fitness plateaus after ∼ \sim 300K progress, restarting from the best program found thus far further improves the fitness shown by the orange curve.