SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 7: The zero-shot ImageNet accuracy curve of LiT-B/16-B (Left). FID comparison on  $64 \times 64$  (Middle) and  $128 \times 128$  (Right) image generation when training diffusion models.
Table 6: Architecture details for language modeling.
Model	#Params	\(n_{layers}\)	\(d_{model}\)	\(n_{heads}\)	\(d_{head}\)
Small-scale
Small	110M	12	768	12	64
Medium	336M	24	1024	16	64
Large	731M	24	1536	16	96
Large-scale
1.1B	1.07B	24	1536	16	96
2.1B	2.14B	32	2048	16	128
7.5B	7.49B	32	4096	32	128
Table 7: Training error  ${L}_{train}$  and landscape flatness  ${L}_{train}^{\mathcal{N}}$  of ViT-B/16 trained from scratch on ImageNet.
Optimizer	AdamW	Lion
ImageNet	75.48	77.44
ReaL	80.64	82.57
V2	61.87	64.81
\(L_{train}\)	0.61	0.75
\(L_{train}^{\mathcal{N}}\)	3.74	1.37
E Architecture Details for Language Modeling
Table 6 shows the Transformer architecture details for language modeling (Section 4.4). The dimension of the feed-forward layer is  $4 \times d_{model}$ . 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 \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 [13] 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.
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.
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_{train}^{\mathcal{N}} = \mathbb{E}_{\epsilon \sim \mathcal{N}}[L_{train}(w + \epsilon)]$  (average over 1K random noises) in Table 7. We observe that the ViT-B/16 trained by AdamW enjoys a smaller training error  $L_{train}$ . However, Lion can enable ViT to converge to flatter regions, as it helps the model retain comparably lower error against Gaussian perturbations.
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
Program 8: Raw program of Lion before removing redundant statements.
def train(w, g, m, v, lr):  
    g = clip(g, lr)  
    m = clip(m, lr)  
    v845 = sqrt(0.6270633339881897)  
    v968 = sign(v)  
    v968 = v - v  
    g = arcsin(g)  
    m = interp(g, v, 0.8999999761581421)  
    v1 = m * m  
    v = interp(g, m, 1.109133005142212)  
    v845 = tanh(v845)  
    lr = lr * 0.0002171761734643951  
    update = m * lr  
    v1 = sqrt(v1)  
    update = update / v1  
    wd = lr * 0.4601978361606598  
    v1 = square(v1)  
    wd = wd * w  
    m = cosupdate()  
    lr = tan(1.4572199583053589)  
    update = update + wd  
    lr = cos(v845)  
    return update, m, v
second input that is required to be a scalar; (6) ternary function interpolate (interp) that uses the third argument  $\mathbf{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$  and  $\epsilon = 10^{-8}$  respectively.
Figure 10: Left: Ablation for the effect of batch size. Lion prefers a larger batch than AdamW. ImageNet accuracy of ViT-B/16 trained from scratch when we vary  $lr$  and  $\lambda$  for AdamW (Middle) and Lion (Right). Lion is more robust to different hyperparameter choices.
Figure 11: Left: The meta-validation (defined in Section 2.3) curves of two search runs measured on a  $\sim 500\mathrm{x}$  larger meta-validation task compared to the proxy. The blue one meta-overfits at  $\sim 15\%$  of the search progress, while the orange one meta-overfits at  $\sim 90\%$  and achieves a better metric.
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.
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
Figure 12: Log perplexity of the small (Left), medium (Middle), and large (Right) size Transformer on PG-19. Since  $\beta_{1} = 0.95$ ,  $\beta_{2} = 0.98$  in Lion when performing language modeling, we compare to Ablation<sub>0.95</sub> and Ablation<sub>0.98</sub> with  $\beta = 0.95$  and  $\beta = 0.98$ , respectively (see Section L for the definition). Lion is still the best-performing one.
Table 8: Accuracy on ImageNet, ImageNet ReaL, and ImageNet V2. Numbers in  $(\cdot)$  are from Dai et al. [21], Dosovitskiy et al. [26]. Results are averaged from three runs.
Model	#Params	Optimizer	RandAug + Mixup	ImageNet	ReaL	V2
Train from scratch on ImageNet
ResNet-50	25.56M	SGD	76.22	82.39	63.93
AdamW	X	76.34	82.72	64.24
Lion	76.45	82.72	64.02
Mixer-S/16	18.53M	AdamW	X	69.26	75.71	55.01
Lion	69.92	76.19	55.75
Mixer-B/16	59.88M	AdamW	X	68.12	73.92	53.37
Lion	70.11	76.60	55.94
ViT-S/16	22.05M	AdamW	X	76.12	81.94	63.09
Lion	76.70	82.64	64.14
AdamW	✓	78.89	84.61	66.73
Lion	79.46	85.25	67.68
ViT-B/16	86.57M	AdamW	X	75.48	80.64	61.87
Lion	77.44	82.57	64.81
AdamW	✓	80.12	85.46	68.14
Lion	80.77	86.15	69.19
CoAtNet-1	42.23M	AdamW	✓	83.36 (83.3)	-	-
Lion	84.07	-	-
CoAtNet-3	166.97M	AdamW	✓	84.45 (84.5)	-	-
Lion	84.87	-	-
Pre-train on ImageNet-21K then fine-tune on ImageNet
ViT-B/16384	86.86M	AdamW	X	84.12 (83.97)	88.61 (88.35)	73.81
Lion	84.45	88.84	74.06
ViT-L/16384	304.72M	AdamW	X	85.07 (85.15)	88.78 (88.40)	75.10
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 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 10\mathrm{x}$ , 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)
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.
Model	Optimizer	CoLA	SST-2	MRPC	STS-B	QQP	MNLI -m	MNLI -mm	QNLI	RTE	Avg
Base	AdamW	60.87	95.18	92.39 / 89.22	90.70 / 90.51	89.23 / 92.00	86.77	86.91	93.70	81.59	87.42
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
Large	AdamW	63.89	96.10	93.50 / 90.93	91.69 / 91.56	90.08 / 92.57	89.69	89.92	94.45	89.17	89.46
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
11B	AdamW	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.65	93.04 / 93.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 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).
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 9 (Left) in the Appendix for the learning curves.
Fine-tuning We fine-tune Base (220M), Large (770M), and the largest 11B T5 [77] on the GLUE benchmark [91]. Every model is fine-tuned for 500K steps with a batch size of 128 and a constant learning rate. Table 9 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.
K Comparison with Other Popular Optimizers
We also employ four popular handcrafted optimizers: RAdam [58], NAdam [27], AdaBelief [105], AMSGrad [82] and two optimizers discovered by AutoML: PowerSign [5] and AddSign [5] to train ViT-S/16 and ViT-B/16 on ImageNet (with RandAug and Mixup). We thoroughly tune the peak learning rate  $lr$  and decoupled weight decay  $\lambda$  [60] 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 to real-world learning tasks. Figure 8 further shows that the learning curves of the five adaptive optimizers are pretty similar, whereas Lion has a unique one that learns faster.
L Ablations
Momentum tracking To ablate the effects of both  $\beta_{1}$  and  $\beta_{2}$ , we compare to a simple update rule:  $\mathfrak{m} = \operatorname{interp}(\mathfrak{g}, \mathfrak{m}, \beta)$ ; update  $= \operatorname{sign}(\mathfrak{m})$ . Two optimizers, Ablation<sub>0.9</sub> and Ablation<sub>0.99</sub>, are created with  $\beta$  values of 0.9 and 0.99 respectively. Illustrated by Table 10, the two ablated optimization algorithms perform worse than all five compared baselines, let alone our Lion. Further
$^{2}$ https://github.com/deepmind/optax
Table 10: The performance of various optimizers to train ViT-S/16 and ViT-B/16 on ImageNet (with RandAug and Mixup). Lion is still the best performing one.
Model	Task	AdamW	RAdam	NAdam	Ada-Belief	AMSGrad	Power-Sign	Add-Sign	Ablation0,9	Ablation0,99	Lion
ViT-S/16	ImageNet	78.89	78.59	78.91	78.71	79.01	77.36	77.37	78.23	78.19	79.46
RealL	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.35	66.82	65.17	64.52	66.13	65.96	67.68
ViT-B/16	ImageNet	80.12	80.26	80.32	80.29	79.85	78.95	78.50	79.54	79.90	80.77
RealL	85.46	85.45	85.44	85.48	85.16	84.76	84.49	85.10	85.36	86.15
V2	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 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 10 (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.
M Hyperparameter Tuning
To ensure a fair comparison, we tune the peak learning rate  $lr$  and decoupled weight decay  $\lambda$  for both AdamW (Adafactor) and our Lion using a logarithmic scale. The default values for  $\beta_{1}$  and  $\beta_{2}$  in AdamW are set as 0.9 and 0.999, respectively, with an  $\epsilon$  of  $1e - 8$ , while in Lion, the default values for  $\beta_{1}$  and  $\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  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.99$  in AdamW, and  $\beta_{1} = 0.95$ ,  $\beta_{2} = 0.98$  in Lion. In our experience, reducing  $\beta_{2}$  results in shorter memorization of historical information and enhanced training stability. Additionally, the  $\epsilon$  in AdamW is set as  $1e - 6$  instead of the default  $1e - 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  $\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  $1r * \lambda$ : update += w *  $\lambda$ ; update *= 1r, the value of  $\lambda$  used for Lion is  $3 - 10x$  larger than that for AdamW in order to maintain a similar strength. For instance,
-  $lr = 1e - 4$ ,  $\lambda = 10.0$  in Lion and  $lr = 1e - 3$ ,  $\lambda = 1.0$  in AdamW when training ViT-B/16 on ImageNet with strong augmentations,
-  $lr = 3e - 5$ ,  $\lambda = 0.1$  in Lion and  $lr = 3e - 4$ ,  $\lambda = 0.01$  in AdamW for diffusion models,
-  $lr = 1e - 4$ ,  $\lambda = 0.01$  in Lion and  $lr = 1e - 3$ ,  $\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 10 (Middle and Right), we alter both  $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.
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 [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, 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.
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) theImagen 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.
Table 11: One-shot evaluation on English NLP tasks. TriviaQA, NQs, and WebQs are NLG tasks and the rest are NLU tasks. This corresponds to Table 3 in the main text.
Task	1.1B
Adafactor	Lion	2.1B
Adafactor	Lion	7.5B
Adafactor	Lion	6.7B
GPT-3	8B
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	50.7	50.3	59.4	59.3	68.2	68.3	66.5	68.2
StoryCloze	74.8	74.4	78.2	78.3	81.2	81.5	78.7	78.7
Winograd	75.1	80.2	81.3	82.1	85.3	84.2	84.6	85.3
Winogrande	59.7	60.5	64.8	65.7	71.4	71.0	65.8	68.3
RACE-m	52.0	50.8	55.1	53.8	59.1	61.3	54.7	57.7
RACE-h	36.8	35.4	40.3	40.7	44.5	43.9	44.3	41.6
PIQA	69.4	69.9	71.3	72.1	75.5	74.5	76.3	76.1
ARC-e	64.3	62.0	69.5	68.9	72.4	72.7	62.6	71.3
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	57.8
WiC	47.6	47.3	46.1	48.1	50.9	48.1	50.3	47.3
Multirc (F1a)	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	86.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 R1	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