SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiangning Chen $^{12§*}$  Chen Liang $^{1§}$  Da Huang $^{1}$  Esteban Real $^{1}$
Kaiyuan Wang<sup>1</sup> Hieu Pham<sup>1</sup> Xuanyi Dong<sup>1</sup> Thang Luong<sup>1</sup>
Cho-Jui Hsieh² Yifeng Lu¹ Quoc V. Le¹
$^{\S}$ Equal & Core Contribution
$^{1}$ Google  $^{2}$ UCLA
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\%$  finetuning 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.3\mathrm{x}$ . 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.
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 [2, 3, 7, 27, 58, 105]. However, Adam [48] with decoupled weight decay [60], also referred to as AdamW, and Adafactor with factorized second moments [86], are still the de facto standard optimizers for training most deep neural networks, especially the recent state-of-the-art language [12, 23, 90], vision [21, 26, 102] and multimodal [75, 84, 100] models.
*Work done as a student researcher at Google Brain.
Table 1: Accuracy of BASIC-L [72] on ImageNet and several robustness benchmarks. We apply Lion to both vision tower pre-training and vision-language contrastive training stages. The previous SOTA results on zero-shot and fine-tuning ImageNet accuracy are  $86.3\%$  and  $91.0\%$  [100].
Optimizer	Zero-shot	Fine-tune
ImageNet	V2	A	R	Sketch	ObjectNet	ImageNet
Adafactor	85.7	80.6	85.6	95.7	76.1	82.3	90.9
Lion	88.3	81.2	86.4	96.8	77.2	82.9	91.1
Figure 1: Left: ImageNet fine-tuning accuracy vs. pretraining cost of ViT models on JFT-300M. Right: FID of the diffusion model on  $256^2$  image generation. We use DDPM for 1K steps w/o guidance to decode image. As a reference, the FID of ADM is 10.94 [24].
Program 1: Discovered optimizer Lion.  $\beta_{1} = 0.9$  and  $\beta_{2} = 0.99$  by default are derived from Program 4. 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.
def train(weight, gradient, momentum, lr): update = interp(grad, momentum,  $\beta_{1}$  ) update  $=$  signupdate) momentum  $=$  interpgradient,momentum,  $\beta_{2})$  weight Decay  $=$  weight \*  $\lambda$  update  $=$  update  $^+$  weight decay update  $=$  update  $*\mathrm{lr}$  return update, momentum
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 [1, 54, 63, 64]. 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 [5, 95] 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 [80] 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. This optimizer 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 [72], 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. Limitations of our work are discussed in Appendix N.
Program 2: 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 (1r). The main output is the update to the weight. v1 and v2 are two additional variables for collecting historical information.
w = weight_initiate()
v1 = zero_initiate()
v2 = zero_initiate()
for i in range(num_train_steps):
    lr = learning_rate_schedule(i)
    g = compute-gradient(w, get_batch(i))
update, v1, v2 = train(w, g, v1, v2, lr)
w = w - update
Program 3: Initial program (AdamW). The bias correction and  $\epsilon$  are omitted for simplicity.
def train(w, g, m, v, lr):
g2 = square(g)
$\mathfrak{m} =$  interp(g, m, 0.9)
$\mathbf{v} = \mathrm{interp}(g2, v, 0.999)$
sqrt_v = sqrt(v)
update  $=$  m/sqrt_v
wd = w * 0.01
update  $=$  update  $^+$  wd
$\mathrm{lr} = \mathrm{lr} * 0.001$
update  $=$  update  $\ast$  lr
return update, m, v
Program 4: Discovered program after search, selection and removing redundancies in the raw Program 8 (Appendix). Some variables are renamed for clarity.
g = clip(g, lr)
$\mathbf{g} = \arcsin (\mathbf{g})$
m = interp(g, v, 0.899)
m2 = m * m
$\mathbf{v} =$  interp(g, m, 1.109)
abs_m = sqrt(m2)
update  $=$  m / abs_m
wd = w * 0.4602
$\mathrm{lr} = \mathrm{lr} * 0.0002$
m = coshupdate)
2 Symbolic Discovery of Algorithms
We present an approach that formulates algorithm discovery as program search [11, 49, 80]. 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.
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 [10, 36] as well as pseudo code of optimization algorithms. The details of the design are outlined below, with an example representation of AdamW in Program 3.
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  $(\mathbf{w})$ , the gradient  $(\mathbf{g})$  and the learning rate schedule value  $(1r)$  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 3 just for better readability. This simplified code snippet in Program 2 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 [5, 95], 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  $\text{interp}(\mathbf{x}, \mathbf{y}, \mathbf{a})$ , which is made equivalent to  $(1 - \mathbf{a}) * \mathbf{x} + \mathbf{a} * \mathbf{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
Figure 2: Left: We run hyperparameter tuning on AdamW and random search, both with 4x more compute, to get the best results as two baselines (green and red lines). The evolutionary search, with mean and standard error calculated from five runs, significantly outperforms both of them. The use of multiple restarts from the initial program is crucial due to the high variance in the search fitness (blue curves), and restarting from the best program after 300K progress further improves the fitness (orange curves) when the original search plateaus. Right: Example curves of search fitness, the cache hit rate, and the percentage of redundant statements. The cache hit rate and the redundant statements percentage increase along with the search progress to  $\sim 90\%$  and  $\sim 70\%$ .
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 \sim \mathcal{N}(0, 1)$ . Additionally, we can mutate an existing constant by multiplying it by a random factor  $2^a$ , where  $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. The modification mutation makes it easier for the search to tune those constants while keeping most of the program unchanged. 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 bigger changes.
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. More importantly, the challenge comes from the sparsity of high-performing programs in the search space. To illustrate this point, we evaluate 2M randomly sampled 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 [42, 79, 80, 87, 99]. It keeps a population of  $P$  algorithms that are gradually improved through cycles. Each cycle picks  $T < P$  algorithms at random and the best performer is chosen as the parent, i.e., tournament selection [32]. 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 from both baselines, shown as the two dashed lines. The high variance in the search fitness necessitates running multiple repeats by restarting from the initial program. When the search fitness plateaus after  $\sim 300\mathrm{K}$  progress, restarting from the best program found thus far further improves the fitness shown by the orange curve.
Pruning through abstract execution We propose to prune the redundancies in the program space from three sources: programs with syntax or type / shape errors, functionally equivalent programs, and redundant statements in the programs. Before a program is actually executed, we perform an abstract execution step that (1) infers variable types and shapes to detect programs with errors; (2) produces a hash that uniquely identifies how the outputs are computed from the inputs, allowing us to cache and look up semantically duplicate programs [31]; (3) identifies redundant statements that can be ignored during actual execution and analysis. For instance, Program 4 is obtained after removing all redundant statements in Program 8 (Appendix). Abstract execution has negligible cost compared to the actual execution, with each input and function replaced by customized values, e.g., hash. See Appendix I for details of abstract execution. Preliminary experiments have shown that the search process can become overwhelmed with invalid programs and cannot make progress without filtering out invalid programs. As seen in Figure 2 (Right), the percentage of redundant statements and cache hit rate both increase as the search proceeds. Based on five search runs, each covering 300K programs, there are  $69.8 \pm 1.9\%$  redundant statements towards the end, implying that redundant statements removal makes the program  $\sim 3x$  shorter on average. The cache hit rate is  $89.1 \pm 0.6\%$ , indicating that using the hash table as cache brings  $\sim 10x$  reduction on the search cost.
Proxy tasks and search cost To reduce search cost, we create low-cost proxies by decreasing the model size, number of training examples, and steps from the target tasks. Evaluation on the proxies can be completed on one TPU V2 chip within  $20\mathrm{min}$ . We use the accuracy or perplexity on the validation set as the fitness. Each search experiment utilizes 100 TPU V2 chips and runs for  $\sim 72\mathrm{h}$ . There are a total of 200-300K programs generated during each search experiment. However, the number of programs that are actually evaluated is around 20-30K, thanks to the use of the cache through abstract execution. To incorporate restart, we start five repeats of search experiments, followed by another round of search initializing from the best algorithm found thus far. This results in a total cost of  $\sim 3\mathrm{K}$  TPU V2 days. See Appendix F for the details of proxy tasks.
2.3 Generalization: Program Selection and Simplification
The search experiments can discover promising programs on proxy tasks. We use performance on meta-validation tasks that are larger than the proxy tasks by increasing the model size and training steps, to select the programs that generalize beyond proxy tasks then further simplify them. The phenomenon of meta-overfitting occurs when the search fitness keeps growing, but the meta-validation metric declines, indicating that the discovered algorithms have overfit the proxy tasks. Two examples meta-validation curves are shown in Figure 11 (Left) in the Appendix.
Large generalization gap The discovered algorithms face a significant challenge due to the substantial gap between the proxy tasks during search and the target tasks. While proxy tasks can typically be completed within 20min on one TPU V2 chip, target tasks can be  $>10^{4}\mathrm{x}$  larger and require days of training on 512 TPU V4 chips. Furthermore, we expect the optimizer to perform well on different architectures, datasets and even different domains, so the discovered algorithms need to show strong out-of-distribution generalization. The sparse search space and inherent noise in the evolution process further compound this challenge, leading to inconsistent generalization properties between different runs. Our observation suggests that evolutionary search runs that meta-overfit later tend to uncover optimization algorithms that generalize better. See more details in Figure 11 (Right) in the Appendix.
Funnel selection To mitigate the generalization gap, we collect promising programs based on search fitness and add an extra selection step using a series of meta-validation tasks to select those generalize better. To save compute, we apply a funnel selection process that gradually increases the scale of the meta-validation tasks. For example, starting with proxy task A, we create a 10x larger task B by increasing the model size and the training steps. Only algorithms that surpass the baseline on task B will be evaluated on task C, which is 100x larger. This approach allows us to gradually filter out algorithms that show poor generalization performance, ultimately leading to the selection of algorithms that generalize well to larger tasks.
Simplification Simpler programs are easier to understand and our intuition is that they are more likely to generalize, so we simplify the programs with the following steps. Firstly, we remove redundant statements that do not contribute to the final output as identified through abstract execution. Secondly, we remove statements that are non-redundant but produce minimal differences when removed. This step can also be achieved through evolution by disabling the insertion of new statements in the mutation process. Finally, we rearrange the statements manually, assign clear and descriptive names to variables, and convert the program into its simpler, mathematically equivalent form.
3 Derivation and Analysis of Lion
We arrive at the Lion optimizer due to its simplicity, memory efficiency, and strong performance in search and meta-validation. The search also discovers other algorithms shown in Appendix D, e.g., some with better regularization and some resembling AdaBelief [105] and AdaGrad [29].
3.1 Derivation
The search and funnel selection process lead to Program 4, which is obtained by automatically removing redundant statements from the raw Program 8 (in the Appendix). We further simplify it to get the final algorithm (Lion) in Program 1. Several unnecessary elements are removed from Program 4 during the simplification process. The cosh function is removed since  $\mathfrak{m}$  would be reassigned in the next iteration (line 3). The statements using arcsin and clip are also removed as we observe no quality drop without them. The three red statements translate to a single sign function. Although both  $\mathfrak{m}$  and  $\mathbf{v}$  are utilized in Program 4,  $\mathbf{v}$  only changes how the momentum is updated (two interp functions with constants  $\sim 0.9$  and  $\sim 1.1$  is equivalent to one with  $\sim 0.99$ ) and does not need to be separately tracked. Note that the bias correction is no longer needed, as it does not change the direction. Algorithm 2 in the Appendix shows the pseudocode.
3.2 Analysis
Sign update and regularization The Lion algorithm produces update with uniform magnitude across all dimensions by taking the sign operation, which is in principle different from various adaptive optimizers. Intuitively, the sign operation adds noise to the updates, which acts as a form of regularization and helps with generalization [16, 30, 67]. An evidence is shown in Figure 9 (Right) in the Appendix, where the ViT-B/16 trained by Lion on ImageNet has a higher training error compared to AdamW but a  $2\%$  higher accuracy on the validation set (as shown in Table 8 from the Appendix). Additionally, the results in Appendix G demonstrate that Lion leads to the convergence in smoother regions, which usually results in better generalization.