Symbolic Discovery of Optimization Algorithms — part 4/4

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

Limitations of Lion While we endeavour to evaluate Lion on as many tasks as possible, the assessment is limited to 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 of autoregressive language model trained on the large-scale internal dataset, which is arguably a more reliable metric the in-context learning benchmarks, and (3) masked language modeling on C4. These tasks have a common characteristic in that the datasets are massive and of high quality, which results in a 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 no better than AdamW if 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.

## 7 Related Work

Our work lies in the area of AutoML and meta-learning that includes learning to learn ( Andrychowicz et al., 2016 ; Ravi and Larochelle, 2017 ; Wichrowska et al., 2017 ; Bello et al., 2017 ; Xiong et al., 2022 ; Metz et al., 2019 ; Metz et al., 2022 ) , neural architecture search ( Real et al., 2019 ; Zoph and Le, 2017 ; Pham et al., 2018 ; Liu et al., 2019a ; Chen and Hsieh, 2020 ; Wang et al., 2021b ; So et al., 2019 ; Chen et al., 2021 ; Yang et al., 2022 ; Wang et al., 2021a ) and hyperparameter optimization ( Li et al., 2017 ; Jamieson and Talwalkar, 2016 ; Hutter et al., 2011 ; Dong et al., 2021 ) , etc. There is also a long history of using evolutionary methods to search for programs, i.e., genetic programming ( Koza, 1994 ; Brameier et al., 2007 ; Holland, 1992 ) . Our approach builds upon a symbolic search space similar to AutoML-Zero ( Real et al., 2020 ; Peng et al., 2020 ) . 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 ( Kingma and Ba, 2014 ; Bernstein et al., 2018 ; Duchi et al., 2011 ; Shazeer and Stern, 2018 ; Zhuang et al., 2020 ; Dozat, 2016 ; Anil et al., 2020 ; Liu et al., 2020 ; Reddi et al., 2018 ; Gupta et al., 2018 ; Riedmiller and Braun, 1993 ; Ma and Yarats, 2019 ) , which we discuss in Section 3.2 .

## 8 Conclusion

This paper proposes to discover optimization algorithms via program search. We propose techniques to address the challenges in searching an infinite and sparse search space, and large generalization gap between the proxy and target tasks. Our method discovers a simple and effective optimizer, Lion, that is memory-efficient and achieves strong generalization across architectures, datasets and tasks.

## Acknowledgements

We would like to thank (in alphabetical order) Angel Yu, Boqing 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.

## Appendix A Pseudocode for AdamW and Lion

## Appendix B Image Classification Tasks

Our evaluation covers various benchmarks: ImageNet, ImageNet ReaL ( Beyer et al., 2020 ) , ImageNet V2 ( Recht et al., 2019 ) , ImageNet A ( Hendrycks et al., 2021b ) , ImageNet R ( Hendrycks et al., 2021a ) , ImageNet Sketch ( Wang et al., 2019b ) , ObjectNet ( Barbu et al., 2019 ) , CIFAR-100 ( Krizhevsky, 2009 ) , and Oxford-IIIT Pet ( Parkhi et al., 2012 ) .

## Appendix 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 in Section 4.4 . Those tasks include Open-Domain Question Answering , Cloze and Completion Tasks , Winograd-Style Tasks , Common Sense Reasoning , In-Context Reading Comprehension , SuperGLUE , and Natural Language Inference .

• NLG: TriviaQA ( Joshi et al., 2017 ) , Natural Questions ( Kwiatkowski et al., 2019 ) , Web Questions ( Berant et al., 2013 ) .

• NLU: HellaSwag ( Zellers et al., 2019 ) , StoryCloze ( Mostafazadeh et al., 2016 ) , Winograd ( Levesque et al., 2012 ) , Winogrande ( Sakaguchi et al., 2020 ) , RACE ( Lai et al., 2017 ) , PIQA ( Bisk et al., 2020 ) , ARC ( Clark et al., 2018 ) , OpenbookQA ( Mihaylov et al., 2018 ) , BoolQ ( Clark et al., 2019 ) , Copa ( Gordon et al., 2012 ) , RTE ( Dagan et al., 2006 ) , WiC ( Pilehvar and Camacho-Collados, 2019 ) , Multirc ( Khashabi et al., 2018 ) , WSC ( Levesque et al., 2012 ) , ReCoRD ( Zhang et al., 2018 ) , CB ( de Marneffe et al., 2019 ) , Adversarial NLI ( Nie et al., 2020 ) .

Algorithm with a better regularization. It dynamically calculates the dot product between the weight and gradient, before computing the weight decay.

Algorithm that tracks the second moment without EMA decay, which is the same as AdaGrad.

Algorithm uses the difference between gradient and momentum to track the second moment, resembling AdaBelief.

## Appendix 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 in the proxy task, we are more likely to discover algorithms with better regularization (Program ), and if we reduce the search progress, we are likely to find simple variants of AdamW (Program and ). Future work can explore this potential to discover optimizers specialized for different tasks.

## Appendix E Architecture Details for Language Modeling

Table 10 shows the Transformer architecture details for language modeling (Section 4.4 ). The dimension of the feed-forward layer is 4 × d m ​ o ​ d ​ e ​ l 4\times d_{model} . We use vocabulary size 32K for small-scale and 256K for large-scale models.

## Appendix 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 64\times 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 ( Chelba et al., 2013 ) for 20K steps with batch size 64, sequence length 32 and vocabulary size 3K. The evaluation time may vary for different programs, but typically a evaluation can be done on one TPU V2 chip within 20min. The validation accuracy or perplexity is used as the fitness.

## Appendix 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 ( Foret et al., 2021 ; Keskar et al., 2017 ; Chen et al., 2022 ; Chen and Hsieh, 2020 ) . Following Chen et al. (2022) , we measure the landscape flatness at convergence by L t ​ r ​ a ​ i ​ n 𝒩 = 𝔼 ϵ ∼ 𝒩 ​ [ L t ​ r ​ a ​ i ​ n ​ ( w + ϵ ) ] L^{\mathcal{N}}_{train}=\mathbb{E}_{\epsilon\sim\mathcal{N}}[L_{train}(w+\epsilon)] (average over 1K random noises) in Table 10 . We observe that the ViT-B/16 trained by AdamW enjoys a smaller training error L t ​ r ​ a ​ i ​ n L_{train} . However, Lion can enable ViT to converge to flatter regions, as it helps the model retain comparably lower error against Gaussian perturbations.

Raw program of Lion before removing redundent statements.

## Appendix H Available Functions

We include 43 available functions that can be used in 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.

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 , log1p , square , sqrt , cube , cbrt , sign , reciprocal and binary functions like + , - , * , / , power , maximum , minimum with the same semantic as the corresponding function in NumPy / JAX.

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_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_sim that treats the two inputs as two vectors and computes their cosine similarity; (5) binary clip_by_global_norm ( clip ) that clips the global norm of the first input to the value of the second input that is required to be a scalar; (6) ternary function interpolate ( interp ) that uses the third argument a , required to be a scalar, to compute a linear interpolation of the first two arguments x and y with (1 - a) * x + a * y .

Functions producing commonly used constants This includes get_pi , get_e , get_eps that generates π \pi , e e and ϵ = 10 − 8 \epsilon=10^{-8} respectively.

## Appendix I Abstract Execution

We propose to prune the large search space with abstract execution. Our approach is motivated by the fact that a large number of programs are invalid, functionally equivalent, or 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 a hash for each unique computation from inputs to outputs to detect redundant statements. The abstract execution can be seen as a static analysis of 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 of the program.

Detecting errors with type / shape inference To detect programs containing errors, we infer the type and shape of each variable in the program through the following steps: (1) replace each input with an abstract object that only contains type and shape information, and replace each statement with a type and shape inference function; (2) iterate through all statements. Instead of executing the 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 a static analysis of the program, exposing errors caused by type and shape mismatch. Note that there are still run-time errors, such as 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 it 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 a functional hash value for every unique computation from the inputs to the outputs as follows: (1) a 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 of all outputs. We then build a hash table that maps each unique functional hash value to the fitness of the corresponding program. When a new program is generated, we first look up its hash value and only perform evaluation if it is not found or if we want to evaluate it multiple times to reduce measurement noise. In our experiments, this technique reduces the search cost by ∼ \sim 10x, as depicted in Figure 2 (Right).

Identifying redundant statements by tracking dependencies In program evolution, redundant statements are included to 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 a recursive manner using the following steps: (1) 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 a 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., a set of all depending statements; (3) compute the union of all statements that each output depends on, which contains all non-redundant statements. By filtering out redundant statements, we obtain a simplified version of the program that is cheaper to execute and easier to analyze. In our experiments, this reduces the program length by ∼ \sim 3x on average, as shown in Figure 2 (Right).

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .