Symbolic Discovery of Optimization Algorithms — part 3/4

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

Locked-image text Tuning (LiT) We perform a comparison between Lion and AdamW on LiT ( Zhai et al., 2022 ) by training the text encoder ( Zhai et al., 2022 ) in a contrastive manner using the same frozen pre-trained ViT. All models are trained for 1B image-text pairs with a batch size of 16,384. Table 4 shows the zero-shot image classification results on three model scales, with the name specifies the size, e.g., LiT-B/16-B denotes a ViT-B/16 and a base size Transformer as the text encoder. Our method, Lion, demonstrates consistent improvement over AdamW with gains of +1.10%, +1.13%, and +0.66% on zero-shot ImageNet accuracy for LiT-B/32-B, LiT-B/16-B, and LiT-g/14 288 -L, respectively. Figure 5 (Left) depicts an example zero-shot learning curve of LiT-B/16-B. Similar results are obtained on the other two datasets. The zero-shot image-text retrieval results on MSCOCO ( Lin et al., 2014 ) and Flickr30K ( Plummer et al., 2015 ) can be found in Figure 9 (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 AdamW on both datasets, with a larger gain in Recall@1 than Recall@10 on Flicker30K, implying more accurate retrieval results: +1.70% vs. +0.60% for image → \rightarrow text and +2.14% vs. +0.20% for text → \rightarrow image.

BASIC Pham et al. (2021) propose to scale up batch size, dataset, and model size simultaneously, achieving drastic improvements over CLIP. It uses a sophisticated CoAtNet ( Dai et al., 2021 ) pre-trained on JFT-5B as the image encoder. Furthermore, the contrastive training is performed on 6.6B image-text pairs with a 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 1 , we achieve a significant 2.6% gain over the baseline, striking a 88.3% accuracy on zero-shot ImageNet classification. Note that this result is 2.0% higher than the previous best result ( Yu et al., 2022 ) . 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 a 91.1% top-1 accuracy on ImageNet, which is 0.1% better than the previous SOTA.

### 4.3 Diffusion Model

Recently, diffusion models achieve a huge success on image generation ( Ho et al., 2020 ; Song et al., 2021 ; Dhariwal and Nichol, 2021 ; Ho and Salimans, 2022 ; Saharia et al., 2022 ) . Given its enormous potential, we test the performance of Lion on unconditional image synthesis and multimodal text-to-image generation.

Image synthesis on ImageNet We utilize the improved U-Net architecture introduced in Dhariwal and Nichol (2021) and perform 64 × 64 64\times 64 , 128 × 128 128\times 128 , and 256 × 256 256\times 256 image generation on ImageNet. The batch size is set as 2,048 and the learning rate remains constant throughout training. For decoding, we apply DDPM ( Ho et al., 2020 ) for 1K sampling steps without classifier-free guidance.The evaluation metric is the standard FID score. Illustrated by Figure 1 (Right) and 5 (Middle and Right), Lion enables both better quality and faster convergence on the FID score. Note that the gap between Lion and AdamW tends to increase with the image resolution, where the generation task becomes more challenging. When generating 256 × 256 256\times 256 images, Lion achieves the final performance of AdamW at 440K steps, reducing 2.3x iterations. The final FID scores are 4.1 (Lion) vs. 4.7 (AdamW), and for reference, the FID of ADM ( Dhariwal and Nichol, 2021 ) is 10.94.

Text-to-image generation We follow the Imagen ( Saharia et al., 2022 ) setup to train a base 64 × 64 64\times 64 text-to-image model and a 64 × 64 → 256 × 256 64\times 64\rightarrow 256\times 256 super-resolution model. All models are trained on a high-quality internal image-text dataset with a batch size of 2,048 and a 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 is identical to the original Imagen setup. Along with the training, 2K images are sampled from the MSCOCO ( Lin et al., 2014 ) 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 ( Ho and Salimans, 2022 ) with a weight of 5.0 is applied as it has been shown to improve image-text alignment. Figure 7 depicts the learning curve. While there is no clear improvement on the base 64 × 64 64\times 64 model, Lion outperforms AdamW on the text-conditional super-resolution model. It achieves a higher CLIP score and has a less noisy FID metric compared to AdamW.

### 4.4 Language Modeling and Fine-tuning

This section focuses on language modeling and fine-tuning. On language-only tasks, we find that tuning β 1 \beta_{1} and β 2 \beta_{2} can improve the quality for both AdamW and Lion. See Section 5 for tuning details.

Autoregressive language modeling We first experiment on two smaller-scale academic datasets Wiki-40B ( Guo et al., 2020 ) and PG-19 ( Rae et al., 2020 ) following Hua et al. (2022) . The employed Transformer spans three scales: small (110M), medium (336M), and large (731M). The architecture details can be found in Appendix E . All models are trained with 2 18 2^{18} tokens per batch for 125K steps, with a learning rate schedule of 10K steps warmup followed by linear decay. The context length is set to 512 for Wiki-40B and 1,024 for PG-19. Figure 7 illustrates the token-level perplexity for Wiki-40B and word-level perplexity for PG-19. Lion consistently achieves lower validation perplexity than AdamW. It achieves 1.6x and 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.

Scaling up the scale of language models and pre-training datasets has revolutionized the field of NLP. So we further perform larger-scale experiments. Our pre-training dataset, similar to that used in GLaM ( Du et al., 2022 ) , consists of 1.6 trillion tokens spanning a wide range of natural language use cases. Following GPT-3 ( Brown et al., 2020 ) , we train three models, ranging from 1.1B to 7.5B parameters, for 300B tokens with a batch size of 3M tokens and a context length of 1K. We evaluate them on three natural language generative (NLG) and 21 natural language understanding (NLU) tasks (see Appendix C for task details). On this massive dataset, we observe no perplexity difference throughout training. Nevertheless, Lion outperforms Adafactor on the average in-context learning ability, as shown in Table 5 . Our 7.5B baseline model, trained for 300B tokens, outperforms the 8B PaLM, trained for 780B tokens, demonstrating the strength of our setup. Lion outperforms Adafactor on both NLG and NLU tasks, particularly on the NLG tasks, with an exact match improvement of +1.0, +0.9, and +0.6 for the 1.1B, 2.1B, and 7.5B models, respectively.

Masked language modeling We also perform BERT training on the C4 dataset ( Raffel et al., 2020 ) . 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 11 (Left) in the Appendix for the learning curves.

Fine-tuning We fine-tune Base (220M), Large (770M), and the largest 11B T5 ( Raffel et al., 2020 ) models on the GLUE benchmark ( Wang et al., 2019a ) . Every model is fine-tuned for 500K steps with a batch size of 128 and a constant learning rate. Table 6 shows the results on 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. It achieves 10, 12, and 10 wins out of 12 scores for T5 Base, Large, and 11B models, respectively.

### 4.5 Comparison with Other Popular Optimizers

We also employ four popular handcrafted optimizers: RAdam ( Liu et al., 2020 ) , NAdam ( Dozat, 2016 ) , AdaBelief ( Zhuang et al., 2020 ) , AMSGrad ( Reddi et al., 2018 ) and two optimizers discovered by AutoML: PowerSign ( Bello et al., 2017 ) and AddSign ( Bello et al., 2017 ) to train ViT-S/16 and ViT-B/16 on ImageNet (with RandAug and Mixup). We thoroughly tune the peak learning rate l ​ r lr and decoupled weight decay λ \lambda ( Loshchilov and Hutter, 2019 ) of every optimizer, while other hyperparameters are set as the default values in Optax. 2 2 2 https://github.com/deepmind/optax As shown in Table 7 , 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 to real-world learning tasks. Figure 10 (in the Appendix) further shows that the learning curves of the five adaptive optimizers are pretty similar, whereas Lion has a unique one that learns faster.

### 4.6 Ablations

Momentum tracking To ablate the effects of both β 1 \beta_{1} and β 2 \beta_{2} , we compare to a simple update rule: m = interp(g, m, β \beta ); update = sign(m) . Two optimizers, Ablation 0.9 and Ablation 0.99 , are created with β \beta values of 0.9 and 0.99 respectively. Illustrated by Table 7 , the two ablated optimization algorithms perform worse than all five compared baselines, let alone our Lion. Further ablation studies on the language modeling task (as depicted in Figure 12 in 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 a large batch size to 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 8 (Left), the optimal batch size for AdamW is 256, while for Lion is 4,096. This indicates that Lion indeed prefers a larger batch size, but its performance remains robust even with a small 64 batch size. Furthermore, when the batch size enlarges to 32K, leading to only 11K training steps, Lion achieves a significant 2.5% accuracy gain over AdamW (77.9% vs. 75.4%), demonstrating its effectiveness in the large batch training setting.

## 5 Hyperparameter Tuning

To ensure a fair comparison, we tune the peak learning rate l ​ r lr and decoupled weight decay λ \lambda for both AdamW (Adafactor) and our Lion using a logarithmic scale. The default values for β 1 \beta_{1} and β 2 \beta_{2} in AdamW are set as 0.9 and 0.999, respectively, with an ϵ \epsilon of 1 ​ e − 8 1e-8 , while in Lion, the default values for β 1 \beta_{1} and β 2 \beta_{2} 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 β 1 = 0.9 \beta_{1}=0.9 , β 2 = 0.99 \beta_{2}=0.99 in AdamW, and β 1 = 0.95 \beta_{1}=0.95 , β 2 = 0.98 \beta_{2}=0.98 in Lion. In our experience, reducing β 2 \beta_{2} results in shorter memorization of historical information and enhanced training stability . Additionally, the ϵ \epsilon in AdamW is set as 1 ​ e − 6 1e-6 instead of the default 1 ​ e − 8 1e-8 as it improves stability in our experiments, similar to the observations in RoBERTa ( Liu et al., 2019b ) .

The update generated by Lion is an element-wise binary ± 1 \pm 1 , as a result of the sign operation, therefore it has a larger norm than those generated by other optimizers. Based on our experience, a suitable learning rate for Lion is typically 3-10x smaller than that for AdamW. 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 lr * λ \lambda : update += w * λ \lambda ; update *= lr , the value of λ \lambda used for Lion is 3-10x larger than that for AdamW in order to maintain a similar strength. For instance, • l ​ r = 1 ​ e − 4 lr=1e-4 , λ = 10.0 \lambda=10.0 in Lion and l ​ r = 1 ​ e − 3 lr=1e-3 , λ = 1.0 \lambda=1.0 in AdamW when training ViT-B/16 on ImageNet with strong augmentations,

• l ​ r = 3 ​ e − 5 lr=3e-5 , λ = 0.1 \lambda=0.1 in Lion and l ​ r = 3 ​ e − 4 lr=3e-4 , λ = 0.01 \lambda=0.01 in AdamW for diffusion models,

• l ​ r = 1 ​ e − 4 lr=1e-4 , λ = 0.01 \lambda=0.01 in Lion and l ​ r = 1 ​ e − 3 lr=1e-3 , λ = 0.001 \lambda=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 8 (Middle and Right), we alter both l ​ r lr and λ \lambda when training ViT-B/16 from scratch on ImageNet. Suggested by the heatmaps, Lion is more robust to different hyperparameter choices compared to AdamW.

## 6 Limitations

Limitations of search Despite the efforts to 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 ( Anil et al., 2020 ; Gupta et al., 2018 ; Martens and Grosse, 2015 ) . 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, as we do not find a 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.