Symbolic Discovery of Optimization Algorithms
Authors: Xiangning Chen
Chen Liang
Da Huang
Esteban Real
Kaiyuan Wang
Hieu Pham
Xuanyi Dong
Thang Luong
Cho-Jui Hsieh
Yifeng Lu
Quoc V Le
Venue: NeurIPS 2023 poster (2023)
Keywords: AutoML

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, $\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.

Symbolic Discovery of Optimization Algorithms
Xiangning Chen 1 23
Chen Liang &
Da Huang'
Esteban Real'
Kaiyuan Wang'
Hieu Pham
Xuanyi Dong'
Thang Luong'
Cho-Jui Hsieh
Yifeng Lu'
Quoe V. Le
Equaa & Core Contribution
GGoogl
UCLA
Abstract
We present 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 simple and effective optimization algorithm, Lion (EvoLved Sign Momentum). It is more memory -efficreef 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. %, respectively. On diffusion models, Lion outperforms Adam by achieving better FID score and reducing the training compute by up to 2.3x. For autoregressive, masked language modeling, and fine-tuning, Lion exhibits 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 smaller learning rate than Adam due to the larger norm of the update produced by the sign function. Additionally, examine the limitations of Lion and identify scenarios where its improvements are small or not statistically significant.
Introduction
Optimization algorithms, i.c., optimizers, play 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 588 105). However, Adam [48] with decoupled weight decay (60), also referred to as
Adam W. 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-ther language [22 23 90
vision
21 26 102 and multimodal 775 84 100 models.
Work done as a student researcher at Google Brain.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Table Accuracy of BASIC-L 72 on ImageNet and several robustness benchmarks. We apply Lion to both vision tower pre- pre-training and vision- laigunl contrastive training stages. The previous SOTA results on zero-shot and tieeiunnng ImageNet accuracy are 86.3% and 91.0% [00]
Zero-shot
Fine-tunc ImageNet
Optimizer
ImageNet
V2
Sketch
OhjectNet
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
Program 1: Discovered optimizer Lion. 81 = 0.9 and 83 = 0.99 by default are It only tracks momentum and uses the sign operation to compute the update. The two gray lines compute the standard decoupled weight decay, where À is the strength.
6.0
Lion (ours) Ademm
ImageNet Acc
888
5.5
derived from Program
87
3x speedup
E 5.0
86
Lion (ours)
85
AdamW
4.5
2.3x speedup
1000
2000
stipp
4000
0.25
0.50
0.75
1:00
Training Compute [exaFLOPs]
Iteration
le6
def train( gradient, nomentum, 1r): update interp(gradient, nonentun, (i)) update sign sign(update)
Figure I: Left: ImageNet fine-tuning accuracy vs. pre- training cost of ViT models on JFT-300M. Right: FID of the diffusion model on 256² image generation. We use DDPM for IK steps w/o guidance to decode image. As reference, the FID of ADM is 10.94 (24).
nomentun interp(gradient, monentur, 92)
weight_docay weight
update update veight_decay
update - update 1r
return update, monentun
Another direction is to automatically discover such optimization algorithms. The learning to optimize (L20) 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 limited number of small tasks, struggle to generalize state-of-the-art settings where much larger models are trained with sigmificantly more training steps. Another line of methods [5] apply reinforcement learning or Monte Carlo Sampling discover new optimizers, where the search space is defined by trees composed from predefined operands (c.g.. gradient and momentum) and operators (c.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 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 particular inspired by AutoML-Zeroo but aims discovering effective optimization algorithms that can improve the statc-of-the- benchmarks.
In this paper, we present a method to formulate algorithm discovery as program search and apply 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, tfte-offtheeart tasks. To tackle these challenges, we employ 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, Icading 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% 1 fine-tuning accuracy on ImageNet by replacing Adafactor with Lion 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

---

Program 2: An example training loop, where the optimization algorithm that we are searching for is encoded within the train function. The main inputs the weight (w), gradient (g) and learning rate schedule (1r). The main output is the update the weight. v1 and v2 are two additional variables for collecting historical information.
Program 3: Initial pro- gram (AdamW). The bias correction and are omit- ted for simplicity.
Program 4: Discovered program after search, se- lection and removing re- dundancies in the raw Pro- gram 8 (Appendix). Some variables are renamed for clarity.
daf train(s, 8, m, v, 1x): square (g)
m. interpís, n, 0.90
v interp(g2, v. 0.999)
aqrt_v eqrt(v)
def train(w, n, v, 1r):
update sqrt_y 0.01
g clip(g, 1r)
g arcsin(g)
a weightt_iiitielieed initialige()
vi zero sero_initializes
update update + wd Ir Ir 0.001
interp(g, v, 0.899)
m2
zero_snittallzaa)
range (oun, stepp)
update update 1r
v interp(g, m, 1.109) abs_n aqrt(n2)
1r seaplig__rtt__chhuueeeee coupute.gradient(w, (1)ch(())
return update, a,
updata aba_n
0.4602
update, vi, x2 train(v, 8, v1, v2, 1x) update
update update ed
Ir = Ir x 0.0002
coahattttte) update update 1r return update, n, V
Symbolic Discovery of Algorithms
We present an approach that formulates algorithm discovery program search [1] 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 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 110, 36 as well as pseudo code of optimization algorithms. The details of the design outlined below, with example representation of AdamW in Program 3
Input/ output signature The program defines 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 (1r) at the current training step. The main output 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 and in Program just for better readability. This simplified code snippet in Program Jses the same signature as AdamW W 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 sequence of assignment statements, with no restrictions on the number of statements or local variables. Each statement calls function using constants or existing variables inputs, and the resulting value is stored in new or existing variable. For the program, we select 45 common math functions, most of which corresponds to function in NumPy or an operation 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
3

---

-3.45
-2425
100%6
Hoaron baeeg
Fitness
80%
Percentage
-3.5
Fitness
2.430
3.435
3.55
3.70T
Redundant statemment
100K
200k
300k
400k
500k
-1.40
50k 100k
150k
200k
250k 300k
20%
Search progress
Search progress
We run hyperparameter tuning Adam W 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 enror 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 eurves), 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 ~00 and 700%
Figure 2: Left:
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, c.g., in the case of adding dictionary of arrays 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 new statement 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 normal distribution X ~ NO 1). Additionally, we can mutate an existing constant by multiplying it by random factor 29, where ~ NO 1). These constants serve 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 program to include redundant statements during search, i.e., statements that do not impact the final program outputs. This is necessary 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 the search space. To illustrate this point, we evaluates 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, 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 chosen as the parent, i.e., tournament selection (32). This parent is then copied and mutated 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 Adam W to accelerate the search. By default, we use tournament size of two and 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 (Left) displays the mean and standard error of five evolutionary search experiments. We run hyperparameter tuning based on Adam by only allowing mutations of constants in the evolution, and run random search by sampling random programs, both with 4x more
and has shown success on many AutoML search tasks [42 79 80

---

compute. Our search significantly outperforms the best results from both baselines, shown as the two dashed lines. The high variance in the search fitmess necessitates running multiple repeats by restarting from the initial program. When the search fitness plateaus after ~300K 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 [I]: (3) identifies redundant statements that can be ignored during actual execution and analysis. For instance, Program is obtained after removing all redundant statements in Program [ (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 [ 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 (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 + 1.9% redundant statements towards the end, implying that redundant statements removal makes the program ~3x shorter on average. The cache hit rate is 89.1 + 0.6%, indicating that using the hash table as cache brings 100 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 20min. We use the accuracy or perplexity on the validation set as the fitness. Each search experiment utilizes 100 TPU V2 chips and runs for 72h. There are 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 total cost of ~3K TPU V2 days. See Appendix [ 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-overfitinng 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 IT (Left) in the Appendix.
Large generalization gap The discovered algorithms face significant challenge due to the substan- tial 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 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 II (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 series of meta- validatiiiii 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.
5

---

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 [T05] and AdaGrad [29]
3.1 Derivation
The search and funnel selection process lead to Program which is obtained by automatically removing redundant statements from the raw Program g (in the Appendix). We further simplify it to get the final algorithm (Lion) in Program [ Several unnecessary elements are removed from Program 4 during the simplification process. The cosh function is removed since 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 m and v are utilized in Program 4 only changes how the momentum is updated (two interp functions with constants ~0.9 and ~1.1 is equivalent to one with ~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 form of regularization and helps with generalization [6] 30 67]. An evidence is shown in Figure 9 (Right) in the Appendix, where the VIT-B/16 trained by Lion ImageNet has a higher training error compared to AdamW W 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 the convergence in smoother regions, which usually results in better generalization.
Momentum tracking The default EMA factor used to track the momentum in Lion is 0.99 (32), compared to the commonly used 0.9 in AdamW and momentum SGD. The current gradient and momentum are interpolated with factor of 0.9 (3₂) before the sign operation is applied. This choice of EMA factor and interpolation allows Lion to balance between remembering a ~~00 longer history of the gradient in momentum and putting more weight on the current gradient in the update. The necessity of both 31 and B2 is further discussed in Appendix L
Hyperparameter and batch size choices Lion is simpler and has fewer hyperparameters compared to Adam W and Adafactor as it does not require e and factorization-related ones. The update is an glement-wise binary +1 if we omit the weight decay term, with larger norm than those produced by other optimizers like SGD and adaptive algorithms. As a result, Lion needs a smaller learning rate and in turn larger decoupled weight decay to achieve similar effective weight decay strength (Ir A). Detailed information on tuning Lion can be found in Appendix M Additionally, the advantage of Lion over AdamW enlarges the batch size increases, which fits the common practice of scaling up model training through data parallelism (Appendix D.
Memory and runtime benefits Lion only saves the momentum thus has smaller memory footprint than popular adaptive optimizers like AdamW, which is beneficial when training large models and or using a large batch size. As an example, AdamW needs at least 16 TPU V4 chips to train a VIT-B/16 with image resolution 224 and batch size 4,096, while Lion only needs 8 (both with bfloat16 momentum). Another practical benefit is that Lion has faster runtime (steps/ sec) in our experiments

---

due to its simplicity, usually 2-15% speedup compared to AdamW and Adafactor depending on the task, codebase, and hardware.
Relation to existing optimizers The sign operation has been explored in previous optimizers 7 g3 The closest ours is the handcrafted optimizer signSGD (7] (and its momentum variant) that also utilizes the sign operation to calculate the update but has a different momentum update rule from Lion. Their focus to mitigate communication costs between agents in distributed training, and they observe inferior performance when training ConvNets on image classification tasks. On the other hand, NAdam [27] combines the updated first moment and the gradient to compute the update, but Lion decouples the momentum tracking and how it is applied to the update through 32. A comparison of Lion with related optimizers can be found in Appendix K
4 Evaluation of Lion
In this section, we present evaluations of Lion, various benchmarks. We mainly compare it to AdamW (or Adafactor when memory is a bottleneck) as it is exceedingly popular and the de facto standard optimizer
majority of learning tasks. The result of momentum SGD is only included for ResNet since it performs worse than Adam elsewhere. We also benchmark other popular optimizers in Appendix IK including handcrafted and automatically discovered ones. We make sure that every optimizer is well-tuned for each task (see Appendix [M tuning details). By default, the learning rate schedule is cosine decay with 10K steps warmuu, and the momentum is saved bfloat16 to reduce the memory footprint. Ablations studies are performed in Appendix L
4.1 Image Classification
We perform experiments including various datasets and architectures the image classification task (see Appendix [ for dataset details). Apart from training from scratch on ImageNet, we also pre-train on two larger well-established datasets, ImageNet-21K, and JFT 89 The image size is 224 by default otherwise specified by the subscript.
Train from scratch ImageNet Following previous works 266 37.. we train ResNet-50 for 90 epochs with a batch size 1,024, and other models for 300 epochs with a batch size of 4,096. As shown in Table & (in the Appendix), Lion significantly outperforms AdamW various architectures. Empirically, the improvement is more substantial on models with larger capacity, with accuracy increases of 1.96% and 0.58% for ViT-B/16 and ViT-S/16, respectively. The performance gaps also tend to enlarger with fewer inductive biases. When strong augmentations are applied, the gain of Lion over AdamW shrinks, but still outperforms A by 0.42% on CoAtNet-3, despite the strong regularization during training (21).
Pre-train Ima We pre-train ViT-B/16 and ViT-L/16 ImageNet-2IK for 90 epochs with a batch size of 4.096. Table g shows that Lion surpasses Adam W when the training set is enlarged for 10x. The gaps larger models are bigger, with +0.52% vs. +0.33% (ImageNet), +0.57% vs. +0.23% (ReaL), and +0.74% vs. +0.25% (V2) for ViT-L/16 and ViT-B/16, respectively.
Pre-train on JFT To push the limit, we conduct extensive experiments on JFT. We follow the settings of Dosovitskiy al. (26) and Zhai et al. [102] for both pre-training and fine-tuning. Figure ( (Left) preseet the accuracy of three ViT models (ViT-B/16, ViT-L/16, and ViT-H/14) under different pre-training budgets on JFT-300M. Lion enables the ViT-L/16 to match the performance of ViT-H/14 trained by AdamW ImageNet and ImageNet V2 but with 3x less pre- pre-training cost. On ImageNet ReaL, the compute saving further becomes 5x. Another evidence is that even when a ViT-L/16 is trained by AdamW for 4M steps by Zhai et al. [022,, its performance still lags behind the same model trained by Lion for 1M steps.
Table in the Appendix shows the fine-tuning results, with higher resolution and EMA. Our VIT-L/16 matches the previous ViT-H/14 results trained by Adam W. while being 2x smaller. The advantage is larger on more challenging benchmarks, such +1.33% (V2), +6.08% (A), +5.54% (R) for ViT- I/16. When pretrained on JFT-3B, the ViT-g/14 trained by I Lion outperforms the previous ViT-G/14 results [102]. with 1.8x fewer parameters. Our ViT-G/I4 achieves 90.7 1%% accuracy ImageNet.
7

---

Table 2: Zero-shot aceuracy of LiTs on Im-
Acc
90.5
-5x speedup
MM ACC
80
speedup
ageNet, CIFAR-100, and Oxford-IIIT Pet. As reference, the zero-shot accuracy of
70
CLIP [75] on ImageNet is 76.2%.
(ours)
(ours)
AdammW
Adamiy
Model
Optimzzrr
ImageNet
C100
Pet
Training Compute [exaFLOPs]
1000 2030
3003
4006
Training Compute [exaFLOPs]
1000 2000
3030 4030
LiT-R/12-B
AdamW
68.78
71.41
86.62
Lion
69.88
71.78
87.36
Figure 3: ImageNet ReaL (Left) and ImageNet V2 (Right) accuracy after we pre-train VIT models on JFT-300M then fine-tune on ImageNet. See Table (in the Appendix) for the detailed numbers.
Adam
74.26
72.25
89.83
LiT-B/I6-B
Lion
75.39
72.49
91.20
LiT-g/14a--L
AdamW
83.43
80.93
94.88
Lion
84.09
81.43
95.86
Table One-shot evaluation averaged over three NLG and 21 NLU tasks. The results of GPT-3 [2]
and PaI.M [17] are included for reference. The LLMs trained by Lion have better in-context learning ability. See Table (in the Appendix) for detailed results on all tasks.
1.IB
2.1B
7.5B
6.7B
8B
Task
Adafactor
Lion
Adafactor
Lion
Adafactor
Lion
GPT-3
PaL.
#Tokens
300B
300B
780B
Avg NLG
11.1
12.1
15.6
16.5 57.4
24.1
24.7
23.1
23.9
Avg NLU
53.2
53,9
56.8
61.3
61.7
58.5
59.4
4.2
Vision-Language Contrastive Learning
This section focuses on the vision-language contrastive training 75 We compare Lion with Adam
(Adafactor) on zero-shot image classification and image-text retrieval benchmarks. We initialize the image encoder with strong pre- ve-trained model as it is suggested to be more efficient [103].
Locked-image text Tuning (LiT) We perform comparison between Lion and AdamW LiT 103 by training the text encoder [103] in a contrastive manner using the same frozen pre-trained VIT. All models trained for 1B image-text pairs with batch size of 16,384. Table[[]hoove the zero-shot image classification results on three model scales, with the name specifies the size, e.g., LiT-B/I6-B denotes ViT-B/16 and base size Transformer the encoder. Our method, Lion, demonstrates consistent improvement over Adam W with gains of +1.10%, +1. 1113%, and +0.66% on zero-shot ImageNet accuracy for LiT-B/32-B, LiT-B/16-B, and 422ss//142s---1 respectively. Figure 7 (Left) in the Appendix depicts an example zero-shot learning curve of LiT-B/16-B B. Similar results are obtained on the other two datasets. The zero-shot image-text retrieval results on MSCOCO (56) and Flickr30K [74] can be found in Figure 6 (in the Appendix). The evaluation metric is Recall @ K. calculated based on if the ground truth label of the query appears in the top-K retrieved examples. Lion outperforms Adam W on both datasets, with a larger gain in Recall than Recall 10 on Flicker30K, implying more accurate retrieval results: +1.70% vs. +0.60% for image -> text and +2.14% vs. +0.20% for text + image.
BASIC Pham et al. 722 propose scale up batch size, dataset, and model size simultaneously, achieving drastic improvements over CLIP. It uses a sophisticated CoAtNet (21] pre-trained on JFT-5B as the image encoder. Furthermore, the contrastive training is performed on 6.6B image-text pairs with larger 65,536 batch size. To push the limit, we only experiment on the largest BASIC-L, and use Lion on both image encoder pre-training and contrastive learning stages. As illustrated in Table T we achieve a significant 2.6% gain over the baseline, striking a 88.3% accuracy on zero-shot ImageNet classification. Note that this result 2.0% higher than the previous best result [100] The performance gain is consistent on five other robustness benchmarks. After fine-tuning the image encoder (CoAtNet-7) in BASIC-L obtained by Lion, we further achieve 91.1 % top-1 accuracy on ImageNet, which is 0. 1.1% better than the previous SOTA.
4.3 Diffusion Model
Recently, diffusion models achieve huge success on image generation 24 the performance of Lion on unconditional image synthesis and multimodal text-to-image generation.
88 so we test

---

2.90
0225
n.
perplexity
2.05
Lion (ours)
4.2
Lion lours)
AdamW
AdamW
(ours)
Lion (oursi
2.80
w1 1,55x speedup
AdamW
0.29
Asamn
2:75
-1 6x speedup
50.0
2.70
3.8
.22x speedup
FID. 100-30
18 1 38
Lion (oursi AdarnW
50000
100000
50000
(ours)
42.5
65.0
Iteration
Iteration
AdamW
Figure 5: Log perplexity on Wiki-40B (Left) and PG-19 (Right). The speedup brought by Lion tends to increase with the model scale. The largest model on Wiki-40B is omitted as we ob- serve severe overfitting.
Itoration
Iteration
Ice
Figure 4: Imagen text-to-image 64 (Left) and the 64 + 256 diffusion models (Right).
Image synthesis on ImageNet We utilize the improved U-Net architecture 24 and perform 64 x 64, 128 x 128, and 256 x 256 image generation on ImageNet. The batch size 2, 048 and the learning rate remains constant throughout training. For decoding, we apply DDPM 41 for IK sampling steps without classifier-free guidance. The evaluation metric is the standard FID score. Illustrated by and Figure 7 (Middle and Right) in the Appendix, Lion enables both better quality and faster convergence on the FID score. The gap between Lion and AdamW increases with the image resolution, where the generation becomes more challenging. When generating 256 256 images, Lion achieves the final performance of AdamW with 2.3x fewer steps. The final FID scores are 4.1 (Lion) vs. 4.7 (AdamW). For reference, the FID of ADM 24 is 10.94.
Text-to-image generation We follow the Imagen model and 64 x 64 + 256 x 256 super-resolution model. All models are trained on a high- quality internal image-text dataset with a batch size of 2,048 and constant learning rate. Due to computational constraints, our base U-Net has a width of 192 compared to 512 in the original 2B model, while the 600M super-resolution model identical to the original Imagen setup. Along with the training, 2K images are sampled from the MSCOCO (56) validation set for real-time evaluation. We use the CLIP score to measure image-text alignment and the zero-shot FID-30K to measure image fidelity. Classifier-free guidance [40] with weight of 5.0 is applied as it has been shown to improve image-text alignment. Figure depicts the learning curve. While there is no clear improvement on the base 64 64 model, Lion outperforms AdamW on the text-conditional super-resolutio model. It achieves a higher CLIP score and has a less noisy FID metric compared to AdamW.
84 setup to train base 64 x 64 text-to-image
4.4 Language Modeling and Fine-tuning
This section focuses on language modeling and fine-tuning. On language-only tasks, we find that tuning 31 and B2 can improve the quality for both AdamW and Lion. See Appendix M for tuning details. The masked language modeling and fine-tuning results are shown in Appendix
Autoregressive language modeling We first experiment on two smaller-scale academic datasets Wiki-40B [4 and PG-19 [76). The employed Transformer spans three scales: small (110M) medium (336M), and large (731M). The architecture details can be found Appendix [.] All models are trained with 2 tokens per batch for 125K steps, with learning rate schedule of 10K steps warmup followed by linear decay. The context length set to 512 for Wiki-40B and 1,024 for 19- Figure Sjillustrates the token-level perplexity for Wiki-40B and word-level perplexity for PG-19. Lion consistently achieves lower validation perplexity than Adam It achieves 1.6x and 1. 1.5x speedup when training the medium size model on Wiki-40B and PG-19, respectively. When the model is increased to the large size, the speedup on PG-19 further increases to 2x.
We then conduct larger experiments with the pre- pre-training dataset similar to GLaM [28]. Following GPT-3 [12), we train three models, from I. to 7.5B parameters, on 300B tokens with a batch size of 3M tokens and context length of 1K. We evaluate them on three natural language generative (NLG) and 21 natural language understanding (NLU) tasks (see Appendix Eopendix[Cfoo, for details). We observe no difference in perplexity throughout training. Nevertheless, Lion outperforms Adafactor on the average in-context learning ability, as shown in Table Our 7.5B baseline model, trained for 300B tokens, outperforms the 8B PaL.M, trained for 780B tokens, demonstrating the strength of our setup. Lion outperforms Adafactor on both NI. and NLU tasks, particularly on the NI. tasks, with an exact match improvement of +1.0, +0.9, and +0.6 for the 1.IB, BB, and 7.5B models, respectively.

---

Related Work
Our work lies in the area of AutoML and meta-learning that includes learning to learn 1 63 64
78 96 977. neural architecture search [1 15
94 |8 [06] and hyperparameter 55,, etc. There also a long history of using evolutionary methods to search for programs, I.c., genetic programming search space similar to AutoML-Zero [70] [0]. However, instead of discovering programs with fixed dimensional matrices, vector, and scalars for toy tasks, our goal is to develop programs that operate on n-dimensional arrays and can generalize to state-of-the-art tasks. Other related works include numerous handcrafted optimizers
57
79
87
93
optimization 225
42 49
Our approach builds upon symbolic
86 105
Conclusion
This paper aims at discovering optimization algorithms via program search. Despite the challenges from an infinite and sparse space, and large generalization gap between the proxy and target tasks, our method discovers simple and effective optimizer, Lion, that is memory-efficiett and achieves strong generalization across architectures, datasets and tasks.
Acknowledgements
We would like to thank (in alphabetical order) Angel Yu, Boging Gong, Chen Cheng, Chitwan Saharia, Daiyi Peng, David So, Hanxiao Liu, Hanzhao Lin, Jeff Lund, Jiahui Yu, Jingru Xu, Julian Grady, Junyang Shen, Kevin Regan, Li Sheng. Liu Yang. Martin Wicke, Mingxing Tan, Mohammad Norouzi, Qiqi Yan, Rakesh Shivanna, Rohan Anil, Ruiqi Gao, Steve Li, Vlad Feinberg, Wenbo Zhang, William Chan, Xiao Wang, Xiaohua Zhai, Yaguang Li, Yang Li, Zhuoshu Li, Zihang Dai, Zirui Wang for helpful discussions, and the Google Brain team at large for providing a supportive research environment.
References
Marcin Andrychowicz, Misha Denil, Sergio Gómez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. In D. Lee, M. Sugiyama, U. Luxburg. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
[2] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, 2020.
[3] Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance stochastic gradients. Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 404-133 413. PMLR, 10-15 Jul 2018
[4] Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenenbaum, and Boris Katz. Objectnet: A large-scale bias-controlled dataset for pushing the limits of object recognition models. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[5] Irwan Bello, Barret Zoph, Vijay Vasudevan, and Quoc V. Le. Neural optimizer search with reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML' 17, page 459 468. JMLR.org, 2017.
[6] Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing on Freebase from question- pairs. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 1533-1544, Seattle, Washington, USA, October 2013. Association for Computational Linguistics.
10

---

[7] Jeremy Bernstein, Yu-Xiang Wang. Kamyar Azizzadenesheli, and Animashree Anandkumar. signSGD: Compressed optimisation for non-convex problems. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 560-569. PMLR, 10-15 Jul 2018.
[8] Lucas Beyer, Olivier J. Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron den Oord. Are we done with imagenet?, 2020.
[9] Yonatan Bisk, Rowan Zellers, Ronan Le bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):77432-777 2241994492-77,,,... Apr. 2020. doi: 10.
v34i05.6239.
[10] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Landerman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:/ / thub. com com/google/jar
[11] Markus Brameier, Wolfgang Banzhaf, and Wolfgang Banzhaf. Linear genetic programminng volume 1. Springer, 2007.
[12] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
[13] Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, and Phillipp Koehn. One billion word benchmark for measuring progress in statistical language modeling. CoRR, abs/1312.3005, 2013.
[14] Xiangning Chen and Cho-Jui Hsich. Stabilizing differentiable architecture search via perturbation-baseed regularization. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning. volume 1 119 Proceedings of Machine Learning Research, pages 1554-1565. PMLR, 13-18 Jul 2020.
[15] Xiangning Chen, Ruochen Wang, Minhao Cheng, Xiaocheng Tang, and Cho-Jui Hsich. DrNAS: Dirichlet neural architecture search. In International Conference on Learning Representations, 2021.
[16] Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training strong data augmentations. In International Conference on Learning Representations, 2022.
[17] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope. James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson. Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayam Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang. Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier- HeieerHellsteem, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022.
[18] Christopher Clark, Kenton Lee, Ming- Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions.
11

---

In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Pupers), pages 2924-2936 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 111188//1//1111111 1300.
[19] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.
[20] Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Joaquin Quiñonero-Candee Ido Dagan, Bernardo Magnini, and Florence d' Alché Buc, editors, Machine Learning Challenges. Evaluating Predictive Uncertainty, Visual Object Classification, and Recognising Tectual Entailment, pages 177-190, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg.
[21] Zihang Dai. Hanxiao Liu, Quoc V Le, and Mingxing Tan. Coatnet:
Marrying convolution and attention for all data sizes. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 3965-3977. Curran Associates, Inc., 2021.
[22] Marie-Catherinn de Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. Proceedings of Sinn und Bedeutung, 1372):107-124, Jul. 2019. doi: 10.18148/sub/2019. v23.2..61....111..
[23] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguis- tics: Human Language Technologies, Volume (Long and Short Papers), pages 4171- 4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 116633/1/N119 1423.
[24] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 8780-8794, Curran Associates, Inc., 2021.
[25] Xuanyi Dong, Mingxing Tan, Adams Wei Yu, Daiyi Peng, Bogdan Gabrys, and Quoe Le. AutoHAS: Efficient hyperparameter and architecture search. In 2nd Workshop Neural Architecture Search at ICLR, 2021.
[26] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[27] Timothy Dozat. Incorporating Nesterov Momentum into Adam. Proceedings of the 4th International Conference on Learning Representatiinss pages 1 1-4, 2016.
[28] Nan Du Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu. Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Hellstee, Toju Duke, Lucas Dixon, Kun Zhang, Quoc Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with mixture-of- experts. In Kamalika Chaudhuri. Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference Machine Leurning, volume 162 of Proceedings of Machine Learning Research, pages 5547-5569. PMLR, 17-23 Jul 2022.
[29] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(11)212,,, 2011.
[30] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.
12

---

[31] Ryan Gillard, Stephen Jonany, Yingjie Miao, Michael Munn, Connal de Souza, Jonathan Dungay, Chen Liang. David R. So, Quoe V. Le, and Esteban Real. Unified functional hashing in automatic machine learning, 2023.
[32] David E Goldberg and Kalyanmoy Deb. A comparative analysis of selection schemes used in genetic algorithms. FOGA, 1991.
[33] Andrew Gordon, Zornitsa Kozareva, and Melissa Roemmele. SemEval-2012 task 7: Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In *SEM 2012: The First Joint Conference on Lexical and Computational Semantics Volume 1: Proceedings of the main conference and the shared lask, and Volume 2: Proceedingy of the Sixth International Workshop on Semantic Evaluation (SemEval 2012), pages 394-398, Montréal, Canada, 7-8 June 2012. Association for Computational Linguisties.
[34] Mandy Guo, Zihang Dai, Denny Vrandečić, and Rami AI-Rfou. Wiki-40B: Multilingual language model dataset. In Proceedings of the Twelfth Language Resources and Evaluation Conference, pages 1440-2452 Marseille, France, May 2020. European Language Resources Association. ISBN 979- 109-44--554--3444
[35] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1842-1850. PMLR, 10-15 Jul 2018.
[36] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Coumaaeau,, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-362 September 2020. doi: 10. 10.103///41886-11 11586644588-00022669922
[37] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pages 770-778, 2016. doi: 10. 10.1109/CVPR.2016.990
[38] Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang. Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of In Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV), pages 8340-8349, October 2021.
generalization.
[39] Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 15262-15271, June 2021.
[40] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022.
[41] Jonathan Ho, Ajay Jain, and Picter Abbecl. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 6840-6851. Curran Associates, Inc., 2020.
[42] John H Holland. Adaptation in natural and artificial systems: an introductory analysis with applications to biology, control, and artificial intelligence. MIT press, 1992.
[43] Frank Hutter, Holger Hoos, and Kevin Leyton--rowe Sequential model- based optimization for general algorithm configuration. In Carlos A. Coello Coello, editor, Learning and Intelligent Optimization, pages 507-523, Berlin, Heidelberg, 2011. Springer Berlin Heidelberg.
[44] Kevin Jamieson and Ameet Talwalkar. Non-stochastic best arm identification and hyperparam- eter optimization. In Arthur Gretton and Christian C. Robert, editors, Proceedings of the 19th International Conference on Artificial Intelligence and Statistics, volume Proceedings of Machine Learning Research, pages 240-248. Cadiz, Spain, 09-11 May 2016. PMLR
13

---

[45] Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 1601-1611, 1, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10. 101186/111/11711144
[46] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[47] Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Look- ing beyond the surface: challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long Papers), pages 252-262,, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10. 10.166531//N18/ 1023.
[48] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014.
[49] John Koza. Genetic programming as means for programming computers by natural selection. Statistics and computing, 4:87-112, 1994.
[50] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report,
2009.
[51] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai. Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452-466, 2019.
[52] Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. RACE: Large-scale ReAding comprehension dataset from examinations. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages 785-794, Copenhagen, Denmark, September 2017. Association for Computational Linguistics. doi: 10. 1/.18/////DDD1--0822
[53] Hector J. Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Proceedings of the Thirteenth International Conference on Principles of Knowledge Repre- sentation and Reasoning, KR 12, page 552-561. AAAI Press, 2012. ISBN 9781577355601.
[54] Ke Li and Jitendra Malik. Learning optimize. In International Conference Learning Representations, 2017.
[55] Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: A novel bandit-based approach to hyperparameter optimization. J. Mach. Learn. Res., 18(1):6765-688 6816, jan 2017. ISSN 1532-4435.
[56] Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitiick. Microsoft coco: Common objects in context. In David Fleet, Tomas Pajdla, Bernt Schiele, and Tinne Tuytelaars, editors, Computer Vision ECCV 2014, pages 755,755, Cham, 2014. Springer International Publishing.
[57] Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In International Conference Learning Representations, 2019.
[58] Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations, 2020.
[59] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach, 2019.
14

---

[60] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[61] Jerry Ma and Denis Yarats. Quasi-hyperbolic momentum and adam for deep learning. In International Conference on Learning Representations, 2019.
[62] James Martens and Roger Grosse. Optimizing neural networks with coonecker-factoree approximate curvature. In International conference on machine learning, pages 2408-2417. PMLR, 2015.
[63] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl- Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4556 4565. PMLR, 09-15 Jun 2019.
[64] Luke Metz, James Harrison, C. Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, and Jascha Training versatile learned optimizers by scaling up. 2022.
Velo:
[65] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of the 2018 Conference Empirical Methods in Natural Language Processing, pages 2381-2391, Brussels, Belgium, October-November 2018. Association for Computational Linguistics. doi: 10..88///v 1/D18- 1260.
[66] Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Vanderwende, Pushmeet Kohli, and James Allen. A corpus and cloze evaluation for deeper understanding commonsense stories. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, pages 339,,49,,60 San Diego, California, June 2016. Association for Computational Linguistics. doi: 10. 10.11835/vI/NT I/N16- 1098.
[67] Arvind Neelakantan, Luke Vilnis, Quoc V. Le, Lukasz Kaiser, Karol Kurach, Ilya Sutskever, and James Martens. Adding gradient noise improves learning for very deep networks, 2017.
[68] Yixin Nie, Adina Williams, Emily Dinan, Mohit Bansal, Jason Weston, and Douwe Kiela. Adversarial NLI: A new benchmark for natural language understanding. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4885- 4901, Online, July 2020. Association for Computational Linguistics. doi: 00.18853/v1/2022 acl-main.441 441.
[69] O. M. Parkhi, A. Vedaldi, A. Zisserman, and C. V. Jawahar. Cats and dogs. In IEEE Conference on Computer Vision and Patlern Recognition, 2012.
[70] Daiyi Peng, Xuanyi Dong, Esteban Real, Mingxing Tan, Yifeng Lu, Gabriel Bender, Hanxiao Liu, Adam Kraft, Chen Liang, and Quoe Le. Pyglove: Symbolic programming for automated machine learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 96-108. Curran Associates, Inc., 2020.
[71] Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 4095-4104. PMLR, 10-15 Jul 2018
[72] Hieu Pham, Zihang Dai, Golnaz Ghiasi, Kenji Kawaguchi, Hanxiao Liu, Adams Wei Yu, Jiahui Yu, Yi-Ting Chen, Minh- Minh-Thang Luong, Yonghui Wu, Mingxing Tan, and Quoc Le. Combined scaling for open- voccbuoaab image classification, 2021.
[73] Mohammad Taher Pilehvar and Jose Camacho-Collados, WiC: the word-in-context dataset for evaluating context-sensitive meaning representations. In Proceedings of the 2019 Conference
15

---

of the North American Chapter of the Association for Computational Linguistics; Human Lan- guage Technologies, Volume (Long and Short Papers), pages 1267-1273. Minneapolis. Min- nesota, June 2019, Association for Computational Linguistics. doi: 10. 10.18633/1//119-112
[74] Bryan A. Plummer, Liwei Wang. Chris M. Cervantes, Juan C. Caicedo, Julia Hockenmaier, and Svetlana Lazebnik. Flickr30k entities: Collecting region-to- -tr-sssse correspondences for richer image-to-sente senteence models. In 2015 IEEE International Conference on Computer Vision (ICCV), pages 2641 1249,, 2015. doi: 10.1 109/1CCV.2201.3333
[75] Alec Radford, Jong Wook Kim, Chris Hallacy, Adity Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pages 8748- -763. PMLR, 18-24 Jul 2021.
[76] Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long- rong-range sequence modelling. In International Conference on Learning Representations, 2020.
[77] Colin Raffel, Noam Shazeer. Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21 21(140):1667 2020.
[78] Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. In International Conference on Learning Representationss 2017.
[79] Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V. Le. Regularized evolution for image classifier architecture search. Proceedings of the AAAI Conference on Artificial Infelligence, 33(01))4780-4777 Jul. 2019. doi: 10. 10.1669/aaai..oo vt0.aaaaaa..3110333114177
[80] Esteban Real, Chen Liang, David So, and Quoe Le. Automl-zero: Evolving machine learning algorithms from scratch. In International Conference on Machine Learning, pages 8007-8019. PMLR, 2020.
[81] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do ImageNet classifiers generalize to ImageNet? In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5389-5400. PMLR, 09-15 Jun 2019.
[82] Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference Learning Representations, 2018.
[83] M. Riedmiller and H. Braun. A direct adaptive method for faster backpropagation learning: the rprop algorithm. In IEEE International Conference on Neural Nenworks, pages 586-591
vol.1, 1993. doi: 10.1
[84] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Raphael Gontijo-Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[85] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):8732 8740, Apr. 2020. doi: 10.1609/aaii. htt.///aaa..400..33333333333 334005.639.
[86] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 4596-4604. PMLR, 10-15 Jul 2018.
[87] David So, Quoe Le, and Chen Liang. The evolved transformer. International Conference Machine Learning, pages 5877-5886. PMLR, 2019.
16

---

[88] Jiaming Song. Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[89] Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effectiveness of data deep learning era. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
[90] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[91] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations, 2019.
[92] Haohan Wang, songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global repre- sentations by penalizing local predictive power. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buch E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[93] Ruochen Wang, Xiangning Chen, Minhao Cheng, Xiaocheng Tang. and Cho-Jui Hsieh, Rank- nosh: Efficient predictor-base based architecture search via non-uniform successive halving. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 10377-10386, October 2021.
[94] Ruochen Wang, Minhao Cheng, Xiangning Chen, Xiaocheng Tang, and Cho-Jui Hsieh. Re- thinking architecture selection in differentiable NAS. In International Conference on Learning Representations, 2021.
[95] Ruochen Wang,
Yuanhao Xiong, Minhao Cheng, and Cho-Jui Hsich. Efficient non-parametric optimizer search for diverse tasks. arXiv preprint arXiv:2209. 13575, 2022.
[96] Olga Wichrowska, Niru Maheswaranathan, Matthew W. Hoffman, Sergio Gómez Colmenarejo, Misha Denil, Nando de Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML' 17, page 3751 1751-37660 JMLR.org, 2017.
[97] Yuanhao Xiong, Li-Cheng Lan, Xiangning Chen, Ruochen Wang, and Cho-Jui Hsich. Learning to schedule learning rate with graph neural networks. In International Conference Learning Representations, 2022.
[98] Chengrun Yang, Gabriel Bender, Hanxiao Liu, Pieter-Jan Kindermans, Madeleine Udell, Yifeng Lu, Quoc V Le, and Da Huang. TabNAS: Rejection sampling for neural architecture search on tabular datasets. In Alice Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems. 2022.
[99] Chris Ying, Aaron Klein, Esteban Real, Eric Christiansen, Kevin Murphy, and Frank Hutter. Nas-bench-101: Towards reproducible neural architecture search. ICML, 2019.
[100] Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models. Transactions on Machine Learning Research, 2022.
[101] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791-4800 Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10111865/111999 1472.
[102] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision trans- formers, 2021.
[103] Xiaohua Zhai, Xiao Wang, Basil Mustafa, Andreas Steiner, Daniel Keysers, Alexander Kolesnikov, and Lucas Beyer. Lit: Zero-shot transfer with locked-image text tuning. In Proceedings of the IEFF/CE Conference on Computer Vision and Pattern Recognition (CVPR), pages 18123-18133, June 2022.
17

---

[104] Sheng Zhang, Xiaodong Liu, Jingjing Liu, Jianfeng Gao, Kevin Duh, and Benjamin Van Durme. Record: Bridging the gap between human and machine commonsense reading comprehension. 2018.
[105] Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18795-18806. Curran Associates, Inc., 2020.
[106] Barret Zoph and Quoc V. Le. Neural architecture search with reinforcement learning. In International Conference on Learning Representations (ICLR), 2017.
Acknowledgments and Disclosure of Funding
Use unnumbered first level headings for the acknowledgments. All acknowledgments go at the end of the paper before the list of references. Moreover, you are required to declare funding (financial activities supporting the submitted work) and competing interests (related financial activities outside the submitted work). More information about this disclosure can found at: https //neurips. //weurips..c//oofffreeeeeeeeeeee www.cc//oooooffeeee htc////nneeeennee///////pprrre http:///..o..aaaooo////////////a
Do not include this section in the anonymized submission, only in the final paper. You can use the ack environment provided in the style file to autmoatically hide this section in the anonymized submission.
18

---

Table Model performance when pre-trained on JFT-300M then fine-tuned on ImageNet. Those
numbers correspond to Figure (Left) and Figure 3 The fine-tuning resolution is 384 for ViT-B/16
and ViT-I/16, and 392² for ViT-H/14. Following Dosovitskiy et al. (26), Polyak averaging is not applied here.
Model
#Params
Epochs / Steps
Optimizer
ImageNet
ReaL
V2
A
R
AdamW
84.24
89.04
74.89
27.39
53.71
ViT-B/16384
86.86M
5/517,791
Lion
84.72
89.14
75.83
29.65
55.86
AdamW
86.69
89.95
78.03
40.55
64.47
7/517,791
Lion
87.32
90.43
79.29
47.13
68.49
ViT-L/16344
304.72M
AdamW
87.29
90.11
78.91
42.56
64.34
14/ 1,03,,55
Lion
88.09
90.62
80.48
51.55
70.72
AdamW
88.02
90.27
80.
53.14
69.48
ViT-H/14392
632.72M
14/ 1,035,583
Lion
88.78
90.68
81.41
58.21
73.09
Table 5: Model performance when pre-trained on JFT then fine-tuned on ImageNet. Two giant ViT models pre-trained JFT-3B while smaller ones pre-trained JFT-300M. The ViT-G/14 results are directly from Zhai et al. [122.
Model
ViT-L/16512
ViT-H/14518
ViT-g/145is
ViT-G/145is
#Params
305.18M
633.47M
1.04B
1.88B
Optimizer
AdamW
Lion
AdamW
Lion
Adafactor
Lion
Adafactor
Lion
ImageNet
87.72
88.50
88.55
89.09
90.25
90.52
90.45
90.71 90.71
Real.
90.46
90.91
90.62
91.02
90.84
91.11
90.81
91.06 / 91.25*
V2
79.80
81.13
81.12
82.24
83.10
83.39
83.33
83.54 83.83*
A
52.72
58.80
60.64
63.78
R
66.95
72.49
72.30
75.07
We observe overfitting in fine-tuning, therefore report both the last and oracle results.
A Pseudocode for Adam W and Lion
Algorithm 1 AdamW Optimizer
given 31, 11,32,,, 1). initialize Oo. mi f- 0, Up < L while 0, not converged do
Algorithm Lion Optimizer (ours)
given B1, A2, X, 77, ʃ
t ++11
initialize 9g. mo 4- 0
9t 1))(0,-1)
while O2 not converged do
update EMA g: and 92
St + (ef(((-1)
m4 +- 81m--1 + (1 31 gt
update model parameters
Vt - 1 + (1 B2)g
Ca < -1 + (1- 31)9t
bias correction
0, +- 01-1 77 (sign(())) 1
mn < m,/(1 39)
update EMA of 9t
- (1 35)
m4 332m-1
32)
update model parameters
end while return 0g
0. 0,-1 (//e/ A0.-1)
end while
return
B Image Classification Tasks
Our evaluation covers various benchmarks: ImageNet, ImageNet Real. 8 ImageNet V2 81
ImageNet (39), ImageNet R (38), ImageNet Sketch [92), ObjectNet 4. CIFAR-100 [50] and
Oxford-IIIT Pet [69].
19

---

imizzbxt R@1
(curs)
Lien (ours)
1x22mmg R@
22.4
(ours)
Adamm
AdamiW
AdamiV
AdamrW
68.3
66.
22.3
20000
20000
20000
40002
20000
40003
20000
40000
Iteration
Iteration
Lion (ours)
Lien (ours)
Lion (ours)
Lien (ours)
imgztet
AdamW
Adaii
Flick 1xt22mm
AdamW
AdamiW
2
22
Flickr
2:0c0 Iteration
Iteration
Iteration
Iteration
Figure 6: Zero-shot image-text retrieval results on MSCOCO (Top) and Flickr30K (Bottom) for LiT-B/16-B. Recall @ K is calculated based on if the ground truth label of the query appears in the top-K retrieved examples.
C NLP Tasks
This section shows all the natural language generation (NLG) and natural language understanding (NLU) tasks where we evaluate the large-scale language models Section 4.4 Those tasks include Open-Domain Question Answering. Cloze and Completion Tasks, Winogrud-Style Tasks, Common Sense Reasoning, In-Context Reading Comprehension, SuperGLUE, and Natural Language Inference.
NLG: TriviaQA 45 Natural Questions 511. Web Questions 6
NLU: HellaSwag [01], StoryCloze [66], Winograd [53]. Winogrande (85 RACE 52
PIQA 99. ARC [9]. OpenbookQA 655 BoolQ [18]. Copa [3), RTE 20 WIC
Multire 47 WSC 53, ReCoRD 104 CB
Adversarial NLI 68
Program 5: Algorithm with a better regularization. It dy- namically calculates the dot product between the weight and gradient, before comput- ing the weight decay.
Program 6: Algorithm that tracks the second mo- ment without EMA de- cay, which is the same as AdaGrad.
Program 7: Algorithm uses the difference be- tween gradient and mo- mentum to track the sec- ond moment, resembling AdaBelief.
def train(u, 8, v, 1x): m interp(u, 8, 0.1) g2 square(g)
der train(u, 8, 1r):
def train(u, 8, v, 1r): interp 0.1)
interp(n, 0.16)
g2 square(g)
intarpolate(v, g2, 0.001)
interp(v, g2, 0.0015)
g2 square (g)
v753 dot(g, egrt_v sqrt(v)
sqrt_y eqrt(v)
interp(v, g2, 0.001) eqrt_v egrt(v)
update m sqrt_y get_pi() min(v, v70)
update agrt_v
update agrt_v ed 0.0238
vd v753 update sin(update) update update 0.0216
update sinh (update)
update update vd Ir Ir 0.03721 update update return update,
Ir lr 0.0606 update updata 1r return update, a, v
update update
sin(v)
return update,
D Other Discovered Programs
By varying the task setting, different types of algorithms can be discovered. For example, if we reduce the amount of data the proxy task, we are more likely to discover algorithms with better regularization (Program 5), and if we reduce the search progress, we are likely to find simple variants of Adam (Program and 7). Future work can explore this potential discover optimizers specialized for different tasks
20

---

5.0
Zero-shot Net AcC
Lion lours)
Lion (ours)
Lion (ours)
AdamW
AdamW
AdamW
33
2.8
4.0
30000
40:03
250050 500000
250000
500000
750500
Iteration
Iteration
Iteration
Figure The zero-shot ImageNet aceuracy curve of LiT-B/16-B (Left)). FID comparison on 64 64 (Middle) and 128 x 128 (Right) image generation when training diffusion models.
Table 6: Architecture details for language modeling.
Table 7: Training error Ltrain and landscape flatness Lirain of ViT-B/16 trained from scratch on ImageNet.
Model
#Params
Bagerr Small-scale
amodel
Pheads
dread
Small Medium Large
110M
12
768
64
Optimizer
AdamW
Lion
336M
24
1024
64
ImageNet
75.48
77.44 44
731M
24
1536
16
96
ReaL.
80.64
82.57
Large-scale
V2
61.87
64.81
1.1B
1.07B
24
1536
16
96
Lirain
0.61
2.1B
2.14B
32
2048
16
128
Livain
3.74
1.37
7.5B
7.49B
32
4096
32
128
E
Architecture Details for Language Modeling
Table 6 shows the Transformer architecture details for language modeling (Section 4.4 The dimension of the cerd-forward layer is dmodet- We use vocabulary size 32K for small-scale and 256K for large-scale models.
F Details of Proxy Tasks
For vision tasks, we train a ViT with three layers, 96 hidden units and three heads, on 10% ImageNet
for 30k steps with batch size 64. The image size is 64 64 and the patch size is 16. For language tasks, we train a Transformer with two layers, 128 hidden units and two heads on LM1B [3] for 20K steps with batch size 64, sequence length 32 and vocabulary size 3K. The evaluation time may vary for different programs, but typically evaluation can be done on one TPU V2 chip within 20min. The validation accuracy perplexity is used the fitness.
Lion (ours)
80
Lon (ours)
ImageNet Acc
ImageNet Acc
AdamW
75
RAdam
78
RAdam
NAdam
NAdam
AdaBellet
76
AdaBelier
AMSGrad
AMSGrad
70
74
72
65
25000 50000
75000
50000
60000 70000 80000 90000
Iteration
Iteration
Figure& 8: Learning curve of ViT-S/16 (Left) and ViT-B/16 (Right) associated with Table 10 The curves of the five adaptive optimizers are similar to each other.
21

---

1.8
Lion (ours)
Log Perplexity
Training Loss
AdamW
1.6
1.4
Lion (ours)
AdamW
1.2
0
50000
100000
40000
60000
80000
Iteration
Iteration
Figure 9: Left: Validation perplexity when we perform masked language modeling on the C4 dataset.
Right: Training loss of ViT-B/16 on ImageNet.
G Analysis of Loss Landscape
In this section, we try to understand why our Lion optimizer achieves better generalization than AdamW from the lens of loss geometry. The convergence to a smooth landscape has been shown to benefit the generalization of deep neural networks [14 16 30, 46). Following Chen et al. [16), we measure the landscape flatness at convergence by L IK random noises) in Table 7 We observe that the ViT-B/16 trained by AdamW enjoys smaller training error Ltrain- However, Lion can enable ViT to converge to flatter regions, as it helps the model retain comparably lower error against Gaussian perturbations.
[e
Lerain + e)] (average over
H Available Functions
We include 43 available functions that can be used the program during search. Note that the input of the functions can be one n-dimensional array, dictionaries or lists of arrays, similar to the pytrees in JAX.
Program 8: Raw program of Lion be- fore removing redundent statements.
def train(w, 5. v, 1r):
elip(g, ir)
clip(m, v84s part(0..62706333388811557
Basic math functions from NumPy / JAX This includes unary functions like abs, cos, sin, tan, arcsin, arccos, arctan, exp, log, sinh, cosh, tanh, arcsinh, arccosh, arctanh, sign, exp2, exp10, expm1, log10, log2. loglp, square, sqrt, cube, cbrt, sign, reciprocal and binary functions like -. *, / power, maximum, minimum with the same semantic as the corresponding function in NumPy JAX.
v968 sign(v)
v968
& arcsin(g)
I interp(g, v, 0.8992999611811221
interp(g, 1.103330051122212)
v845 tanh (v845) 1r 1r 0.000217176178444488
Linear algebra functions commonly used in first-order optimization algorithms This includes: (1) unary function norm that computes the norm of each arrays in the input; (2) unary function global_nont norm that computes the global norm by treating all the numbers in the input as one vector; (3) binary function dot that treats the two inputs as two vectors and computes their dot product; (4) binary function cosine_ __iim that treats the two inputs as two vectors and computes their cosine similarity; (5) binary that clips the global norm of the first input to the value of the second input that is required to be scalar; (6) ternary function interpolate (interp) that uses the
update m.
aqrt(v1)
update update
vd aquara(vi)
0.460197838160855
ad ud * a cosh cosh (update) tan(1. 4572199883055588)
update update Ir cos(v845)
htsp_by_gllbbe _orm (clip)
retura update,
third argument a, required to be a scalar, to compute a linear interpolation of the first two arguments and y with (1 a)
y.
Functions producing commonly used constants This includes get_pi, get_e, get_eps that generates and 10-8 respectively.
22

---

0.15
0.5
77.5
ImageNet Acc
75.0
78
We ght
0.3
1.0
72.5
Lion (ours)
70.0
1
AdemW
9.6
2.0
67.5
1.5e-
Batch Size
Peak Learning Rate
Peak Learning Rate
Figure 10: Left: Ablation for the effect of batch size. Lion prefers larger batch than AdamW. ImageNet accuracy of ViT-B/16 trained from scratch when we vary Ir and for AdamW (Middle) and Lion (Right). Lion is robust different hyperparameter choices.
Early meto evenfitting excmptt: Lote meto overliting example
2.9
-2.00
Percer Percentag eentaaee
Meta-va taavvall
Meta-val val
30%
2.91
2.955
20%
-2.92
0%
20%
10%
60%
80%
100%
0%
20%
40%
Relative search progress
Relative search progress
Figure 11: Left: The valddvaaio (defined in Section 2.3 curves of two search runs measured on ~500 larger meta-valt validatio task compared to the proxy. The blue meta-overfits 155% of the search progress, while the orange one meta-overfits at ~90% and achieves a better metric. Right: Histogram of the search progress when meta-overfitting happens based on 50 runs. Half of the runs meta-overfit early but long tail of runs meta-overfit much later. Blue cross depicts the best meta-valin validation metric averaged within each bin, indicating that meta-overfitting happening later leads to programs that generalize better.
I Abstract Execution
We propose to prune the large search space with abstract execution. Our approach is motivated by the fact that large number of programs are invalid, functionally equivalent, contain redundant statements that waste compute during evaluation. To address this, we introduce an abstract execution step that checks the type and shape of each variable, and computes hash each unique computation from inputs outputs to detect redundant statements. The abstract execution can be seen static analysis the program, achieved by replacing functions and inputs with customized values. We outline the specifics of the customized values and abstract execution procedure for three use cases below. The cost of the abstract execution is usually negligible compared to the actual execution.
Detecting errors with type / shape inference To detect programs containing errors, we infer the type and shape of each variable the program through the following steps: (1) replace each input with abstract object that only contains type and shape information, and replace each statement with type and shape inference function; (2) iterate through all statements, Instead of executing the
4.25
Log Perplexity
3.95
3.85
4.20
3.90
3.80
Lion (ours)
Lion tours;
Lian (ourc)
4.15
AdamW
AdamW
AdamW Ablation0.95
Ablation0.95
3.85
AblationO.95
3.75
AblationO.90
Ablationa.
Ablatiora.se
4.10
0
50000
100000
0
50000
100000
50000
100000
Iteration
Iteration
Iteration
Figure 12: Log perplexity the small (Lefft) medium (Middle), and large (Right) size Transformer on PG-19. Since B1 0.95, Br = 0.98 in Lion when performing language modeling, we compare Ablationa.s and Ablationa with 0.95 and 0.98, respectively (see Section for the definition). Lion is still the best-performing one.
23

---

program is Table Accuracy on ImageNet. ImageNet ReaL, and ImageNet V2. Numbers in (+) are from Dai et al. 21 Dosovitskiy et al. [26]. Results are averaged from three runs.
RandAug
Model
#aaams
Optimizer
ImageNet
Real.
Mixup
Train from scratch on ImageNet
SGD
76.22
82.39
63.93
ResNet-50
25.56M
AdamW
76.34
82.72
64.24
Lion
76.45
82.72
64.02
AdamW
69.26
75.71
55.01
Mixer-S/16
18.53M
Lion
69.92
76.19
55.75
Adam W
68.12
73.92
53.37
Mixer-B/16
59.88M
Lion
70.11
76.60
55.94
AdamW
76.12
81.94
63.09
Lion
76.70
82.64
64.14
VIT-S/16
22.05M
Adam W
78.89
84.61
66.73
Lion
79.46
85.25
67.68
AdamW
75.48
80.64
61.87
Lion
77.44
82.57
64.81
VIT-B/16
86.57M
Adam W
80.12
85.46
68.14 14
Lion
80.77
86.15
69.19
Adam W
83.36 (83.3)
1oAtNet-1
42.23M
Lion
84.07
AdamW
84.45 (84.5)
CoAINet-3
166.97M
Lion
84.87
Pre-train ImageNet-21 then fine-tune ImageNet
Adam
84. 12 (83.97)
88.61 61 (88.35)
73.81
ViT-B/16w4
86.86M
Lion
84.45
88.84
74.06
AdamW
85.07 (85.15)
88.78 (88.40)
75.10
VIT-L/16w
304.72M
x
Lion
85.59
89.35
75.84
original statement, we validate a function call by checking the function signature and type and shape information of its arguments. If valid, we compute the type and shape information of the output and assign it to the new variable; (3) verify the validity of the derived type and shape of the output. This process essentially performs static analysis of the program, exposing errors caused by type and shape mismatch. Note that there are still run-time errors, such division by zero, that cannot be detected in this manner. Without such filtering of invalid programs, the search would be overwhelmed with invalid programs, making difficult to achieve meaningful progress.
Deduplicating with functional hash Among the valid programs that execute without errors, there are still lots of duplicates due to functionally equivalent programs that have different surface forms but the same underlying functionality. To address this issue, we calculate functional hash value for every unique computation from the inputs to the outputs as follows: (1) unique hash value is assigned to each input and function; (2) iterate through all statements, calculating the hash value of the outputs by combining the hash values of the functions and arguments; (3) compute the hash value of program by combining the hash values all outputs. We then build a hash table that maps each unique functional hash value to the fitness of the corresponding program. When generated, we first look up its hash value and only perform evaluation if it not found or if we want to evaluate it multiple times to reduce measurement noise. In experiments, this technique reduces the search cost by ~00, as depicted Figure 2 (Right).
Identifying redundant statements by tracking dependencies In program evolution, redundant statements are included enable combining multiple mutations to make larger program changes. However, these redundant statements increase the evaluation cost and make program analysis more challenging. To identify redundant statements, we need to determine the set of statements that the outputs depend on, which can be computed in recursive manner using the following steps: (1)

---

Table 9: Fine-tuning performance of the T5 Base, Large, and 11B on the GLUE dev set. Results reported are the peak validation scores per task.
MNLI
MNLI
Model
Optimizer
CoLA
SST-2
MRPC
STS-B
QQP
QNLI
RTE
Avg
AdamW
60.87
95.18
92.39
00.39/89.72 89,22
90.70/ /9.51
89.23 92.00
86.77
86.91
93.70
81,59
87
Base
Lion
61.07
95.18
92.52 89,46
90.61 / /90.40
89.52 / 92.20
87.27
87.25
93.85
85,56
87.91
AdamW
63.89
96.10
93.50/ /9.93
91.69 /91.6
90.08 92.57
89.69
89.92
94.45
89.17
89.46
Large
Lion
65.12
96.22
94,06/ 91.67
91.79/ / 91,60
90.23 / 92.67
89.85
89,94
94.89
90.25
89.86
1IB
Adam
69. 50
97.02
93.75
/91.18
92.57 /92.61
90.45 92.85
92.17
91.99
96.41
92.42
91.08
Lion
71.31
97.13
94.58 /92.6
93.04/ /9.04
90.57 92.95
91.88
91.65
96.56
93.86
91.60
replace the value of each input with an empty set, as they do not depend on any statement; (2) iterate through each statement. Note that each statement is an assignment that calls a function and assigns the result to variable, which in turn depends on the current statement and all the depending statements of the function arguments. Therefore we replace the value of the variable with its dependency, i.e., set of all depending statements; (3) compute the union of all statements that each output depends on, which contains all non non-redundant statements. By filtering out redundant statements, we obtain simplified version of the program that is cheaper to execute and easier to analyze. In our experiments, this reduces the program length by <3x on average, as shown in Figure (Right).
J Masked Language Modeling and Fine-tuning
Masked language modeling We also perform BERT training on the C4 dataset 77 It requires the language models to reconstruct randomly masked out tokens in the input sequence. We use the same architectures and training setups as the smaller-scale autoregressive experiments. Lion performs slightly better than AdamW regarding the validation perplexity: 4.18 vs. 4.25 (small), 3.42 vs. 3.54 (medium), and 3.18 vs. 3.25 (large). See Figure (Left) in the Appendix for the learning curves.
Fine-tuning We fine-tune Base (220M), Large (770M), and the largest 1BB T5 [77] the GLUE benchmark [9T]. Every model is fine-tuned for 500K steps with a batch size of 128 and a constant learning rate. Table shows the results the GLUE dev set. For MRPC and QQP, we report the F1 / Accuracy scores, for STS-B, we report the Pearson Spearman correlation, and for the other datasets, we report their default metric. On average, Lion beats AdamW across all three model scales. achieves 10, 12, and 10 wins out of 12 scores for T5 Base, Large, and 11B models, respectively.
K Comparison with Other Popular Optimizers
We also employ four popular handerafted optimizers:
RAdam (58] NAdam 27]. AdaBelief [05].
train ViT-S/16 and ViT-B/16 on ImageNet (with RandAug and Mixup). We thoroughly tune the peak learning rate Ir and decoupled weight decay X 600 of every optimizer, while other hyperparameters are set as the default values in Optax As shown in Table 10 Lion is still the best performing one. We notice that there is no clear winner amongst the baselines. AMSGrad performs the best on VIT-S/16 but the worst on ViT-B/16. The inferior performance of PowerSign and AddSign compared to other optimizers is consistent with previous observations that automatically discovered optimizers have difficulty generalizing real-world learning tasks. Figure g further shows that the learning curves of the five adaptive optimizers are pretty similar, whereas Lion has a unique one that learns faster.
AMSGrad [82] and two optimizers discovered by AutoML: PowerSign 5 and AddSign
L Ablations
Momentum tracking To ablate the effects of both 3, and 82, we compare to a simple update rule: m interp(g, m, (8); update sign(m). (m). Two optimizers, Ablationg and Ablationo.99. are created with 3 values of 0.9 and 0.99 respectively. Illustrated by Table T0, the two ablated optimization algorithms perform worse than all five compared basclines, let alone our Lion, Further
25

---

Table 10: The performance of various optimizers to train VIT-S/16 and ViT-B/16 on ImageNet (with AundAug and Mixup). Lion is still the best performing one.
Power- Sign
Add- Sign
Model
Tisk
Adam
RAdam
NAdam
Belief
AMSGrad
Ablationa
Ablationg
Lion
ImageNet
78.89
78.59
78,91
78.71
79.01
77.36
77.37
78.23
78.19
79.46
VIT-S/16
ReaL
84.61
84.47
84.62
84,56
85.01
83.39
83.36
84.28
84.17
85,25
V2
66.73
66.39
66.02
66.33
66.82
65.17
64.52
66.13
65.96
67.68
ImageNet
80.12
80.26
80.32
80.29
79.85
78.95
78.50
79.54
79.90
80.77
VIT-B/16
ReaL
85.46
85.45
85,44
85,48
85.16
84.76
84.49
85.10
85,36
86.15
68.14
67.76
68.46
68.19
68.48
67.46
65.95
68.07
68.20
69.19
ablation studies on the language modeling task (as depicted in Figure 12 the Appendix) yield similar conclusions. Those results validate the effectiveness and necessity of using two linear interpolation functions, letting Lion to remember longer gradient history meanwhile assign a higher weight to the current gradient.
Effect of batch size Some may question whether Lion requires large batch size accurately determine the direction due to the added noise from the sign operation. To address this concern we train a ViT-B/16 model on ImageNet using various batch sizes while maintaining the total training epoch as 300, and incorporating RandAug and Mixup techniques. As shown in Figure [[](ft))) the optimal batch size for Adam is 256, while for Lion is 4,096. This indicates that Lion indeed prefers larger batch size, but its performance remains robust even with small 64 batch size. Furthermore, when the batch size enlarges to 32K, leading to only 11K training steps, Lion achieves significant 2.5% accuracy gain over AdamW W (77.9% vs. 75.4%), demonstrating its effectiveness in the large batch training setting.
M Hyperparameter Tuning
To ensure a fair comparison, we tune the peak learning rate Ir and decoupled weight decay À for both AdamW (Adafactor) and our Lion using logarithmic scale. The default values for 31 and 83 in AdamW are set 0.9 and 0.999, respectively, with an 1e 8, while in Lion, the default values for 31 and B2 are discovered through the program search process and set as 0.9 and 0.99, respectively. We only tune those hyperparameters in Section 4.4] where 33 0.9, B.2 0.99 in Adam W. and 31 = 0.95, B2 0.98 in Lion. In our experience, reducing B2 results in shorter memorization of historical information and enhanced training stability. Additionally, the e Adam is set as - instead of the default le 8 as it improves stability in our experiments, similar to the observations in RoBERTa (59]
The update generated by Lion is an element-wise binary +1, as result of the sign operation, therefore it has larger norm than those generated by other optimizers. Based on our experience, suitable learning rate for Lion is typically 3-10x smaller than that for Adam W. Note that the initial value, peak value, and end value of the learning rate should be changed simultaneously with the same ratio compared to AdamW. We do not modify other training settings such as the learning rate schedule, gradient and update clipping. Since the effective weight decay is ir
update * X; update *= 1r, the value of X used for Lion is 3-10x larger than that for AdamW in order to maintain similar strength. For instance,
lr = 1e-ee 4, A = 10.0 in Lion and Ir = 1e - 3, X = 1.0 in Adam when training ViT-B/16 on ImageNet with strong augmentations,
(r = 3e 5, x = 0.1 in Lion and (r = 3e - 4, À = 0.01 in AdamW for diffusion models, Ir le 4, = 0.01 in Lion and (r 1e 3, X 0.001 in Adafactor for the 7.5B
language modeling.
Please see Table 12 (in the Appendix) for all hyperparameters.
Apart from the peak performance, the sensitivity to hyperparameters and the difficulty in tuning them are also critical for the adoption of an optimizer in practice. In Figure 10 (Middle and Right), we alter both Ir and when training ViT-B/16 from scratch on ImageNet. Suggested by the heatmaps, Lion is more robust to different hyperparameter choices compared to AdamW.
26

---

N
Limitations
Limitations of search Despite the efforts make the search space less restrictive, it remains inspired by the popular first-order optimization algorithms, leading to a bias towards similar algorithms. It also lacks the functions required to construct advanced second-order algorithms 2 35| 62]. The search cost is still quite large and the algorithm simplification requires manual intervention. Further reducing the bias in the search space to discover more novel algorithms and improving the search efficiency are important future directions. The current program structure is quite simplistic, we do not find good usage of more advanced program constructs such as conditional, loop statements, and defining new functions. Exploring how to incorporate these elements has the potential to unlock new possibilities.
Limitations of Lion While we endeavour to evaluate Lion on as many tasks as possible, the assessment limited the chosen tasks. On vision tasks, the discrepancies between Lion, AdamW, and momentum SGD are pretty small on ResNets, likely due to the fact that ConvNets are easier to optimize compared to Transformers. The performance gain brought by Lion decreases when strong augmentations are utilized. There are also several tasks where Lion performs similarly to AdamW, including: (1) the Imagen text-to-image base model, (2) the perplexity autoregressive language model trained on the large-scale internal dataset, which is arguably more reliable metric the in-context learning benchmarks, and (3) masked language modeling on C4. These tasks have common characteristic in that the datasets are massive and of high quality, which results in reduced difference between optimizers. Another potential limitation is the batch size. Though people often scale up the batch size to enable more parallelism, it is likely that Lion performs better than AdamW W the batch size is small (<64). Additional, Lion still requires momentum tracking in bfloat16, which can be expensive for training giant models. One potential solution is to factorize the momentum to save memory.
27

---

Table 11: One-shot evaluation on English NLP tasks, TriviaQA, NQs, and WebQs are NLG tasks and the are NLU tasks. This corresponds Table in the main
1.IB
2.1B
7.5B
6.7B
8B
Task
Adafactor
Lion
Adafactor
Lion
Adafactor
Lion
GPT-3
PaLM
#Tokens
300B
300B
780B
TriviaQA (EM)
21.5
25.1
32.0
33.4
47.9
48.8
44.4
48.5
NQs (EM)
4.3
4.8
6.3
7.3
12.3
12.1
9.8
10.6
WebQs (EM)
7.5
6.3
8.4
8.7
12.1
13.3
15.1
12.6
HellaSwag StoryCloze
50.7
50.3
59.4
59.3
68.2
68.3
66.5
68.2
74.8
74.4
78.2
78.3
81.2
81.5
78.7
78.7
Winograd Winogrande
75.1
80.2
81.3
82.1
85.3
84.2
84.6
85.3
59.7
60.5
64.8
65.7
71.4
71.0
65.8
68.3
RACE-m RACE-h
52.0
50.8
55.1
53.8
61.3
54.7
36.8
35.4
40.3
40.7
43.9
44.3
PIQA
71.3
72.1
75.5
74.5
76.3
76.1
ARC-e
62.0
69.5
68.9
72.7
62.6
ARC-c
31.2
32.9
37.3
38.0
43.3
42.6
41.5
42.3
OpenbookQA
44.8
48.0
48.4
49.0
51.4
52.4
53.0
47.4
BoolQ
54.3
56.7
64.1
62.9
73.5
73.9
68.7
64.7
Copa
75.0
78.0
83.0
84.0
85.0
87.0
82.0
82.0
RTE
55.6
52.4
49.8
59,2
63.9
62.5
54.9
$7.8
WIC
47.6
47.3
46.1
48.1
50.9
48.1
50.3
47.3
Multire (Fla)
35.9
44.3
45.0
48.8
44.7
59.2
64.5
50.6
WSC
76.5
75.4
79.6
79.3
$6.7
85.6
60.6
81.4
ReCoRD
73.4
73.7
77.8
77.7
81.0
81.1
88.0
87.8
CB
46.4
44.6
48.2
44.6
51.8
46.4
33.9
41.1
ANLI RI
33.3
30.1
32.4
31.2
31.5
34,0
31.6
32.4
ANLI R2
29.8
31,8
29.8
30.6
32.4
31.9
33.9
31.4
ANLI R3
29.8
31.8
31.4
31.9
33.6
34.2
33.1
34.5
Avg NLG
11.1
12.1
15.6
16.5
24.1
24.7
23.1
23.9
Avg NLU
53.2
53,9
56.8
57.4
61.3
61.7
58.5
59.4
28

---

Table 12: Hyperparameters for all the experiments.
Mode
Dropout
Depth
Augmentations
Optimizer
82
Ir
Train from scratch ImageNet
ResNet-50
Lion
0.9
0.909
0.99
3e
3e
1.0
Mixer-S/16
AdamW
0.999
le
0.3
Lion
0.99
3e
Mixci-B/16
0.1
AdamW
Lion
0.9
0.999
0.99
3e
le 2
0.3
3.0
NIT-S/16
0.1
0.1
Adami
Lion
0.9
0.9
0.999
0.99
le -
le 3
0.1
1.0
RandAug: 15
Adamm
0.999
3e
Mixup:
Lion
3e
VIT-B/16
0.1
0.1
AdamW
Lion
0.9
0.999
3e 3
0.3
RandAug: 15
Adam'
0.999
1e
1.0
Mixup: 0.5
Lion
0.9
0.99
le 4
10.0
RundAng:
AdamW
0.999
0.05
Mixup:
Lion
0.99
2x
1.0
3
0.7
RandAug: 15
Mixup:
AdamW
Lion
0.999
1e
0.05
Pre-irain on ImageNet-21K
VIT-B/16
AdamW
0.999
1c
0.1
VIT-L/16
0.1
0.1
AdamW
Lion
0.9
0.999
0.99
le 3
le
0.3
1.0
Pre-train on JFT
NIT-B/16
AdamW
Lion
0.9
0.999
6c 4
0.3
0.1
VIT-L/16
AdamW
Lion
0.9
0.9
0.999
0.99
3e 4
te
0.1
0.3
AdamW
De
Lion
0.99
de
0.3
V VIT-G/14
Adafactor
0.999
Se
0.03
Visiot-langunge contrastive leamnngg
LIT-B/--B
AdamW
0.999
1e
Lion
0.99
3e
LIT-g/14-L
AdamW
0.9
0.999
le 3
0.1
0.5
BASIC-L
Adafactor
Lion
0.9
0.9
0.999
0.99
Se 4
2e 4
0.01
0.1
Diffusion model
Imagen base & super-resolution
AdamnW
0.9
0.999
1 1e-3
Lion
0.9
0.99
1e
Image generation on ImageNet
0.1 128 129 & 256 x 256: 0.2
Lion
0.9
0.909
0.99
3e
3e
0.01
0.1
Autoregressive masked language modellng
Small Medium (PG-19, C4) & Large
Adami
0.99
3e
Lion
0.95
0.98
Se 4
Medium (Wiki-40B)
AdamW
3e
0.001
Lion
0.95
0.98
3e
0.01
1.1B 2.1B
Adafactor
Lion
0.9
0.99
2e 3
0.00015
7.58
Adafactor
Lion
0.95
0.9
0.99
0.98
1e 4
le
0.001
0.01
Language fine me-ttuning
TS-Base Large & 11B
0.1
AdamW Lion
0.95
0.9
0.99
0.98
Se 5
Se
29


Reviews:
[{"review_id": "TpyCxUvwr7", "reviewer": "Reviewer_hEv6", "summary": "\nSummary: The authors' work undoubtedly stands as an exceptional and noteworthy contribution to the field, showcasing a plethora of remarkable achievements that have the potential to significantly influence and shape the landscape of optimization and extend their impact well beyond its boundaries. As I delve into the intricacies of their research, it becomes evident that their dedication and innovative approach have yielded results that are both impressive and promising. \n\nHowever, amidst the numerous commendable aspects of their study, there also exist certain primary concerns that warrant attention and resolution. While I acknowledge the impressive strides made by the authors, it is essential to critically address certain aspects that could further enhance the credibility and applicability of their findings. In the following detailed comments, I will highlight specific areas that require closer scrutiny, offering constructive feedback and suggestions to fortify the integrity and comprehensiveness of their research.\n\n", "questions": "\nSee *Weaknesses.", "limitations": "\nNo. See *Weaknesses.\n\n", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I am elated to express my joy and satisfaction upon the efficiency of the authors' method has transcended that of the renowned Adam optimizer, exhibiting its superiority across various tasks. This notable achievement undoubtedly highlights the ingenuity and prowess of the authors in devising such a remarkable approach.\n\nThe profound impact of their work is vividly demonstrated through the astonishing results presented in Table 1 and Figure 1, which undoubtedly leave one in awe. In particular, the Discovered optimizer, aptly named \"Lion,\" showcased in Program 1, stands as a testament to the authors' ability to craft a solution that is both simplistic and elegant, evoking admiration from peers and fellow researchers alike.\n\nMoreover, the decision made by the authors to simplify Program 4 and refine it into the more efficient Program 1 is not only logical but also reinforces the practicality of their approach. This thoughtful choice not only streamlines the methodology but also adds to the elegance and cohesiveness of their research.\n\nIn their diligent pursuit of showcasing the sparsity of the high-performance optimizer, the authors undertook a remarkably extensive and rigorous experimental endeavor, involving the execution of a staggering 2 million randomly selected programs. With meticulous precision and analytical acumen, they thoroughly assessed the performance of each program, leaving no stone unturned in their quest for a superior alternative to AdamW. Upon meticulously scrutinizing the results obtained from this monumental experiment, the authors arrived at a rather intriguing and significant revelation: none of the 2 million randomly selected programs exhibited a performance that surpassed the prowess of AdamW. This particular finding stands as a testament to the exceptional capabilities and efficiency of AdamW as a high-performance optimizer, solidifying its position as a prominent and leading choice in the optimization landscape. The sheer scale and comprehensiveness of this experimental undertaking are truly commendable, underscoring the authors' unwavering commitment to pursuing scientific rigor and thoroughness. Their methodical approach and the impressive volume of data analyzed serve as an invaluable resource for the scientific community, contributing to a deeper understanding of optimizer performance and fostering future advancements in the realm of optimization algorithms.\n\nThe authors' insightful use of meta-validation to select programs with superior generalization ability is indeed a reasonable and effective approach. Especially noteworthy is their ingenious technique of progressively increasing the complexity of the meta-validation task, aptly termed \"Funnel selection.\" This strategic refinement undoubtedly enhances the efficiency of the selection process, allowing for the identification of programs with enhanced generalization capabilities.\n\nAdditionally, the thoughtful incorporation of sign updates and regularization is highly justified, especially when considering the context of adversarial attacks, where techniques like L_inf PGD attacks often rely on such sign update mechanisms. In fact, sign updates have also been previously employed as optimizers in prior works.\n\nUndoubtedly, the analysis presented in Section 3.2 is both captivating and intellectually stimulating, offering valuable insights that resonate with the broader research community. The robustness and trustworthiness of these findings are indeed commendable and a testament to the authors' thoroughness in their investigations.\n\nOne fascinating observation that delights me is the correlation between model size and Lion's validation accuracy. The authors' remark on the positive relationship between increased model size and enhanced validation accuracy raises intriguing possibilities. The plausible explanation that the sign update technique increases uncertainty and ultimately improves the model's generalization ability opens exciting avenues for future research.\n", "weaknesses": "\nWhile their research presents impressive insights and advancements, it is essential to address one concern that arises from the potential handcrafted nature of the simplification and derivation process for the programs. A more elegant solution would involve automating this process as an inherent aspect of defining the search space, further enhancing the overall methodology and augmenting its potential impact.\n\n\nSign updates have also been previously employed as optimizers in prior works.\n\n\nIt is worth kindly noting that, when the proposed method is combined with more potent data augmentations, the relative performance gains might appear less pronounced. This insightful observation highlights the need for further exploration and optimization to fully unlock the method's potential under varying experimental conditions. While considering this aspect, I must admit that I personally feel a bit disappointed.\n\nI am genuinely interested in exploring the application of the Lion algorithm to the domain of Stable Diffusion. It is pertinent to acknowledge that training stable diffusion models can be an extremely resource-intensive endeavor, often incurring substantial costs (for instance, approximately 1 million dollars for stable diffusion v1). Given this context, the potential of the proposed Lion method to be utilized in conjunction with stable diffusion holds tremendous appeal and significance. If the Lion algorithm proves to be compatible with stable diffusion models, it would indeed be a highly advantageous and desirable outcome. The ability to leverage the proposed method's capabilities in optimizing stable diffusion models could potentially mitigate the significant expenses associated with their training, thereby opening up new avenues for research and applications in this area. The prospect of such synergy between Lion and Stable Diffusion sparks a sense of optimism and excitement, as it has the potential to drive transformative advancements in the field. As such, I eagerly look forward to any insights or findings that may shed light on the feasibility and practicality of applying the Lion algorithm to the challenging task of stable diffusion, as this would undoubtedly represent a valuable and impactful contribution to the scientific community.\n\n\nOne of the  concerns in the field of AutoML revolves around the efficacy of the search process, causing a great deal of apprehension among researchers and practitioners. More precisely, a central aspect of this concern is rooted in the correlation between the performance exhibited by programs during the search phase, specifically on the small proxy, and their subsequent performance on the final task. This particular correlation is viewed with heightened attention, as it holds significant implications for the overall effectiveness of the AutoML approach. For instance, in the context of Neural Architecture Search (NAS), where the exploration of different architectural configurations plays a crucial role in achieving optimal results, there is a palpable sense of unease regarding architecture ranking. How well a particular architecture performs on the small proxy tasks might not always be indicative of its actual performance when applied to more extensive and state-of-the-art tasks, and thus, researchers seek to establish a robust link between the two domains. In attempts to address this pressing concern, the authors of the present paper have diligently employed a comprehensive set of measures. Their goal is to carefully select algorithms that can demonstrate a remarkable ability to generalize from small proxy tasks to larger and more complex real-world tasks, thereby enhancing the reliability and applicability of the AutoML approach. However, despite the earnest efforts made by the authors, there remains a notable omission in their work. They have not provided a quantitative analysis that establishes the explicit correlation between the performance achieved on the small proxy tasks and the corresponding performance on the final, real-world tasks. This apparent absence of a rigorous and quantitative examination is indeed regrettable, as such analytical insights are vital for instilling greater confidence in the effectiveness and reliability of the proposed AutoML techniques. In conclusion, while significant strides have been taken in addressing concerns regarding the search effectiveness in AutoML, the lack of a robust quantitative analysis linking the performance between the small proxy and the final tasks leaves a critical aspect unanswered. Efforts to bridge this analytical gap would undoubtedly enhance the overall understanding and trustworthiness of the AutoML methodologies employed.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive and insightful comments!\n\n> A more elegant solution would involve automating this process as an inherent aspect of defining the search space, further enhancing the overall methodology and augmenting its potential impact. \n\nWe would like to point out that a large part of the simplification procedure is already automated. For example, program 4 (12 statements) is a result of automatic simplification from program 8 (21 statements). This automatic simplification is done through pruning with abstract execution (More details in Section 2.2 and Appendix I). Note that we only perform the manual simplification step on the best algorithms from the funnel selection and such simplification follows naturally from our analysis of the programs to understand it better. \n\nWe agree with the reviewer that the simplification could be further automated by introducing heuristics for the mathematical equivalent transformation and ablation study, and we leave it for future work. \n\n> Sign updates have also been previously employed as optimizers in prior works.\n\nWe have discussed the previous works that share similarities, such as the use of sign operation, with the discovered algorithms (Section 3.2 and Appendix K). Note that none of the prior works that uses sign update showed superior results compared to Adam and the ablation study in Appendix L demonstrated the benefit from the two linear interpolations used in Lion. \n\n> As such, I eagerly look forward to any insights or findings that may shed light on the feasibility and practicality of applying the Lion algorithm to the challenging task of stable diffusion, as this would undoubtedly represent a valuable and impactful contribution to the scientific community.\n\nIn section 4.3, we evaluate Lion's performance in training diffusion models on ImageNet across three resolutions: 64, 128, and 256. While this approach differs from latent diffusion, it offers encouraging evidence of Lion's performance in stable diffusion.\n\n> More precisely, a central aspect of this concern is rooted in the correlation between the performance exhibited by programs during the search phase, specifically on the small proxy, and their subsequent performance on the final task. ... Efforts to bridge this analytical gap would undoubtedly enhance the overall understanding and trustworthiness of the AutoML methodologies employed.\n\nWe agree with the reviewer that the generalization gap between the proxy tasks and the final / target tasks poses a great challenge. We discussed this challenge in section 2.3. Specifically, Figure 11 in the Appendix demonstrates the meta-overfitting issue caused by the gap. Programs that perform well on proxy tasks don't necessarily perform well on larger tasks. We alleviate this challenge through the funnel selection process described in section 2.3 that uses progressively larger tasks to further select out the programs that are more likely to work on the final task. \n\n"}, {"review_id": "yai0ak3KAu", "reviewer": "Reviewer_Gizf", "summary": "In this paper, a new optimizer is designed by program search. The search process overcomes the problem of distribution differences for small data set proxies and large tasks. The new optimizer converges faster and at the same time has strong generalization over a wide range of models for different tasks.", "questions": "See weaknesses.", "limitations": "Limitations have been well discussed", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Apply the idea of program search of AutoML-Zero to the field of optimizer design. Efficient and highly generalized optimizers are designed.\n2. Solved the problem of complex and sparse search space by abstract execution pruning, warm-start and proxy task.\n3. solved the problem of distribution difference between proxy tasks and larger tasks by funnel selection.\n4. Sufficient experiments. The experiments include the validation of extensive models on different tasks.", "weaknesses": "1. The original algorithm obtained from the search is still more complex. It needs expert knowledge and labor cost for simplification. See the simplification of Program 4 to Program 1 for details.\n2. No specific validation experiments or theoretical analysis of the high memory efficiency of Lion compared to adam. It is only outlined at L233-L239. The authors can use some metrics, such as memory peaks, to compare the memory footprint of different optimizers.\n3. The design for the search space is heuristic. In some places, the authors just make statements about the approach, which makes it difficult to understand. For example, the last of the three forms of mutation. Why is the hyperparameter in the statement modified instead of replacing the statement itself?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive and insightful comments!\n\n> The original algorithm obtained from the search is still more complex. It needs expert knowledge and labor cost for simplification. See the simplification of Program 4 to Program 1 for details.\n\nWe would like to point out that a large part of the simplification procedure is already automated. For example, program 4 (12 statements) is a result of automatic simplification from program 8 (21 statements). This automatic simplification is done through pruning with abstract execution (More details in Section 2.2 and Appendix I). Note that we only perform the manual simplification step on the best algorithms from the funnel selection and such simplification follows naturally from our analysis of the programs to understand it better. \n\nWe agree with the reviewer that the simplification could be further automated by introducing heuristics for the mathematical equivalent transformation and ablation study, and leave it for future work.\n\n> No specific validation experiments or theoretical analysis of the high memory efficiency of Lion compared to adam. It is only outlined at L233-L239. The authors can use some metrics, such as memory peaks, to compare the memory footprint of different optimizers.\n\nLion keeps track of only the momentum while Adam keeps track of both momentum and second moments. So we believe the memory efficiency is evident since there are less floats to keep in the memory. We agree with the reviewer that metrics memory peaks can be helpful but it depends on specific implementations and hardwares, so we leave the comprehensive comparison across different codebases and hardwares for future work. \n\n> The design for the search space is heuristic. In some places, the authors just make statements about the approach, which makes it difficult to understand. For example, the last of the three forms of mutation. Why is the hyperparameter in the statement modified instead of replacing the statement itself?\n\nThe third type of mutation (modifying a statement by randomly changing its arguments) is introduced to make it easier to add small changes. Compared to replacing the existing statement with one randomly generated statement, i.e., a replacement mutation, it preserves most of the structure and allows \"fine-tuning\" the arguments. For example, consider a statement where most parts of the statements are good but one argument, which is a float constant, needs to be tuned. In this case, the modification mutation makes it easier to mutate the value of the constant based on the current value while a full replacement would make it harder to make such small changes. Additionally, the replacement mutation can already be composed with the first two types of mutations (insertion and deletion) so we didn't introduce it as a separate one. We will update the draft to make the motivation clearer. \n"}, {"review_id": "OxrXVV4gao", "reviewer": "Reviewer_weMp", "summary": "This paper proposes to formulate algorithm discovery as program search and discover optimization algorithms for deep neural network training. To bridge the large generalization gap between proxy and target tasks, the authors introduce program selection and simplification strategies. The proposed method discovers a simple and effective optimization algorithm called Lion, which is more memory-efficient than Adam as it only keeps track of the momentum. Lion is compared with widely used optimizers, such as Adam and Adafactor, for training a variety of models on different tasks. Empirical results demonstrate that Lion significantly outperforms these optimizers. The authors also examine the limitations of Lion and identify scenarios where its improvements are small or not statistically significant.", "questions": "Despite its sound empirical performances, I would also like to see some theoretical analysis on the searched Lion optimizer. How it can converge under which conditions? The authors may consider referring to [7] and see if similar convergence results can be obtained.", "limitations": "I suppose the key limitation of this work is its huge demand on computation power. For most research groups, the value of this work is solely the Lion optimizer (which is rather simple) instead of the search method, though it is certainly good to see some new optimizers no matter how they are found (by human or by machine). \n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is clear and easy to understand\n- Empirical results are very sufficient to demonstrate the superiority of proposed method\n", "weaknesses": "- The computation resources required by the proposed search method is unaffordable for most research groups\n- Theoretical analysis is missing in this work\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive and insightful comments!\n\n> The computation resources required by the proposed search method is unaffordable for most research groups\n\nWe acknowledge that the search cost of Lion (3K TPU V2 days) is still significant, and improving its efficiency is an important future work. However, the search process is a one-time task, while the benefits of Lion are persistent and sustainable. In our evaluation of image classification, we pre-train ViT-G/14 and CoAtNet-7 on JFT, each requires 512 TPU V4 chips for one week, equivalent to 3K TPU V4 days, or approximately 12K TPU V2 days. Lion's ability to save over 2x compute on this task compared to Adafactor already offsets the search cost. As more researchers use Lion on various tasks, these compute savings will continue to grow.\n\nFurthermore, Lion can also bring benefits to production models, such as ads, which require frequent training and operate at a large scale. In these cases, the advantages of Lion can be even greater and outweigh the one-time search cost by a significant margin.\n\n> Theoretical analysis is missing in this work. Despite its sound empirical performances, I would also like to see some theoretical analysis on the searched Lion optimizer. How it can converge under which conditions? The authors may consider referring to [7] and see if similar convergence results can be obtained.\n\nWe would like to emphasize that Lion is discovered through symbolic program search without any prior favoring sign operation or the decoupled momentum. While we acknowledge the importance of theoretical analysis, it is out of the scope of the current paper that focuses on the automatic discovery of the algorithm. Therefore, we leave it for future research endeavors.\n"}, {"review_id": "yGofyuBXkh", "reviewer": "Reviewer_Y4eN", "summary": "The paper introduces Lion, an automatic search strategy for deep learning optimizers. The paper claims that the resulting optimizer is more compute efficient than the baselines and achieves better accuracy on a variety of tasks from vision and language and model architectures. The search uses an evolutionary algorithm initialized wih AdamW that gradually mutates a population of optimizers and contains several additional strategies to increase search efficiency eg the removal of redundant instructions and a hash-based caching as well as funnel selection to mitigate overfitting on the proxy task. The standout discovery of the Lion optimizer is the use of a sign function in combination with a specific momentum strategy. Throughout their experiments, the authors demonstrate significant improvements, especially on larger architectures with large datasets. ", "questions": " - Program 4 is an automatic simplification of Program 8. However, the authors further do manual simplifications to arrive at the final optimizer in program 1. I understand that the sign simplification (from the code in red) might be hard to automatically find, however, the method also includes a step to remove statements with minimal effect. Why does this removal step not remove the arcsin and clip functions if the authors manually found that their removal does not cause a quality drop? \n - The authors discuss that Lion needs a smaller learning rate due to the sign operation compared to AdamW. How does this lr compare to the one used during the search? Did you use a small learning rate during the search which might bias the search towards the sign operation or does the search actually find an optimizer that is optimal for a learning rate that is orders of magnitude smaller than the one used during the search process? In the second case, do you think that searching with a smaller lr might influence the results? \n - Would it be possible to do a similar search for a second-order method, eg initialized with Newton's method or do you think the additional complexity of the hessian integration would make the search space too complex? ", "limitations": "On some tasks, the improvements are rather minimal. Also, the set of instructions is rather limited and does not allow for complex structures like for-loops. All significant limitations are addressed by the authors in the Appendix. ", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": " - The evaluation is very thorough. The authors include several vision and language models, contrastive learning as well as diffusion models. Throughout the experiments on Transformer models, I found the results to be a significant improvement over AdamW, especially since there has not really been a significant upgrade to AdamW in recent years. I am not fully convinced by the results on ResNets and I also think that it would have been better to have a larger selection of convolutional neural networks. But given the importance of transformers in ML, I think this is still a good result. \n - In general, I find the idea of the paper to automatically search for an optimizer to be relevant. \n - The presentation is clear overall. ", "weaknesses": " - I think the contribution is mostly technical and seems mostly like an engineering effort as the actual search strategy seems relatively simple.  For me, the results are still interesting enough for acceptance.\n - There is no code available and a lot of experiments are done on Google internal datasets such as JFT and the Imagen dataset which makes it impossible to reproduce a lot of the results in the paper. I could also not find any search hyperparameters in the appendix, eg learning rate used during the search. \n - The search method still produces quite a lot of unnecessary statements that need to be manually removed to arrive at the final program. While I understand the complexity of the task, I would still prefer a method that produces cleaner code automatically. Being able to produce simplified optimizers during training might even improve results, for example, what would happen if one would start the search again starting at the Lion optimizer instead of AdamW. \n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive and insightful comments!\n\n> Program 4 is an automatic simplification of Program 8. However, the authors further do manual simplifications to arrive at the final optimizer in program 1. I understand that the sign simplification (from the code in red) might be hard to automatically find, however, the method also includes a step to remove statements with minimal effect. Why does this removal step not remove the arcsin and clip functions if the authors manually found that their removal does not cause a quality drop?\n\nThis is a great question. We found two reasons why such statements remain in the program. (1) Some of the statements are added by evolution because it was helpful in some period of the search, for example, the clip function could help when the optimizers in the population are unstable and introduces very large gradients but becomes less helpful when the optimizer becomes inherently stable; (2) some statements are added because it helps on proxy tasks but makes no difference on larger target tasks, which can be seen as a type of meta-overfitting as discussed in section 2.3 and Figure 11; (3) some statements makes no difference in quality but are still preserved due to noise in evolution; note that since it doesn't make the quality worse, there is no pressure for evolution to remove them. \n\n\n> The authors discuss that Lion needs a smaller learning rate due to the sign operation compared to AdamW. How does this lr compare to the one used during the search? Did you use a small learning rate during the search which might bias the search towards the sign operation or does the search actually find an optimizer that is optimal for a learning rate that is orders of magnitude smaller than the one used during the search process? In the second case, do you think that searching with a smaller lr might influence the results?\n\nThe learning rate is represented as a searchable constant in the program and also mutated during the evolutionary search. The initial value for the learning rate is selected based on the initial program (AdamW) so there should be no bias to push it towards the sign operation. In fact, we are quite surprised by the use of the sign operation in the discovered program since prior works using sign operation usually achieve worse results than Adam. Note that, as a baseline, we also performed a hyperparameter search for AdamW on the proxy tasks with more compute and the final result is still worse than the discovered optimizer. \n\n> Would it be possible to do a similar search for a second-order method, e.g. initialized with Newton's method or do you think the additional complexity of the hessian integration would make the search space too complex?\n\nThanks for the suggestions! Since we use programs as the representation, it is definitely possible and quite natural to expand the search space to include more operations that can compose second order methods. Since this is our first attempt in this direction, we decided to keep it simple and only include operations for first order optimizers. One potential challenge with searching second order methods would be that introducing hessian matrices or operations like outer products can make the programs more expensive to evaluate, thus increasing the compute cost. \n"}]

Author rebuttal:

Meta-review: This paper proposes a symbolic search algorithm to discover novel optimizers. As a direct product of the algorithm, an novel optimizer is discovered, which is evaluated extensively on a wide range of large scale benchmarks, and show favorable performance than the standard AdamW optimizer. Reviewers are generally very positive about the findings, and one reviewer (hEv6) is particularly enthusiastic about it. The AC agrees that this is a solid work, and the fact that the optimizer demonstrating convincing performance gain over AdamW alone justifies acceptance. The AC agrees with Reviewer weMp that a deeper understanding of the discovered optimizer would be interesting, but also acknowledges that this may be beyond the scope of this work.
Decision: Accept (poster)