摘要

本文提出一种将算法发现表述为程序搜索的方法，并将其应用于发现用于深度神经网络训练的优化算法。利用高效的搜索技术来探索无限和稀疏的程序空间。为弥合代理任务和目标任务之间的巨大泛化差距，还引入了程序选择和简化策略。该方法发现了一种简单有效的优化算法Lion (evolutionary Sign Momentum)。它比Adam更节省内存，因为它只记录动量。与自适应优化器不同，通过符号操作计算的每个参数的更新具有相同的幅度。将Lion与广泛使用的优化器(如Adam和adfactor)进行了比较，用于在不同任务上训练各种模型。在图像分类方面，Lion将ViT在ImageNet上的准确率提高了2%，并在JFT上节省了高达5倍的预训练计算。在视觉-语言对比学习方面，在ImageNet上实现了88.3%的零样本和91.1%的微调精度，分别超过了之前的最佳结果2%和0.1%。在扩散模型上，Lion通过获得更好的FID分数并将训练计算量减少了2.3倍，超越了Adam。在自回归、掩码语言建模和微调方面，Lion表现出与Adam类似或更好的性能。对Lion的分析表明，其性能增益随着训练批大小的增加而增长。由于符号函数产生的更新范数更大，它还需要比Adam更小的学习率。此外，研究了Lion的局限性，并确定了其改进很小或没有统计意义的场景。Lion的实现是公开的。

1、简介

优化算法，即优化器，在训练神经网络中起着基本作用。近年来引入了大量手工优化器，主要是自适应优化器(Anil et al.， 2020;贝尔斯和亨尼格，2018;Bernstein等人，2018;Dozat, 2016;Liu等人，2020;庄等，2020)。然而，具有解耦权重衰减(Loshchilov and Hutter, 2019)的Adam (Kingma and Ba, 2014)，也称为AdamW，以及具有因子化二阶矩的adfactor (Shazeer and Stern, 2018)，仍然是训练大多数深度神经网络的事实上的标准优化器，特别是最近最先进的语言(Brown等人，2020;Devlin等人，2019;Vaswani等人，2017)，愿景(Dai等人，2021;Dosovitskiy等人，2021年;翟等人，2021)和多模态(Radford等人，2021;撒哈拉等人，2022年;Yu et al.， 2022)模型。

另一个方向是自动发现此类优化算法。学习优化(lmo)方法建议通过训练参数化模型(如神经网络)来发现优化器，以输出更新(Andrychowicz等人，2016;李和马利克，2017;Metz et al.， 2019, 2022)。然而，这些黑盒优化器通常在数量有限的小任务上进行训练，很难推广到最先进的设置，在这些设置中，训练更大的模型需要更多的训练步骤。另一行方法(Bello等人，2017;Wang et al.， 2022)应用强化学习或蒙特卡洛采样来发现新的优化器，其中搜索空间由预定义的操作数(如梯度和动量)和操作符(如一元和二进制数学操作)组成的树定义。然而，为了使搜索易于管理，它们通常通过使用固定操作数和限制树的大小来限制搜索空间，从而限制了发现的潜力。因此，所发现的算法尚未达到最先进水平。AutoML-Zero (Real et al.， 2020)是一项雄心勃勃的努力，它试图搜索机器学习管道的每个组件，同时评估玩具任务。这项工作遵循自动发现优化器的研究方向，特别是受到AutoML-Zero的启发，但旨在发现有效的优化算法，可以提高最先进的基准。

本文提出了一种将算法发现形式化为程序搜索的方法，并将其应用于发现优化算法。主要有两个挑战。第一个目标是在无限稀疏的程序空间中找到高质量的算法。第二种是进一步选择出可以从小型代理任务推广到更大的最先进的任务的算法。为应对这些挑战，采用了一系列技术，包括热启动和重启的进化搜索、抽象执行、漏斗选择和程序简化。

该方法发现了一种简单有效的优化算法:Lion，进化符号动量的缩写。与各种自适应算法不同的是，该算法仅跟踪动量，并利用符号操作来计算更新，从而降低了内存开销，并在所有维度上统一更新幅度。尽管很简单，Lion在一系列模型(Transformer、MLP、ResNet、U-Net和Hybrid)和任务(图像分类、视觉-语言对比学习、扩散、语言建模和微调)中表现出出色的性能。值得注意的是，通过在BASIC中将Adafactor替换为Lion，在ImageNet上实现了88.3%的零样本微调精度和91.1%的微调精度(Pham等人，2021)，分别超过了之前的最佳结果2%和0.1%。此外，Lion将JFT上的预训练计算减少了多达5倍，将扩散模型的训练效率提高了2.3倍，并实现了更好的FID分数，并在语言建模方面提供了类似或更好的性能，计算量节省了多达2倍。

分析了Lion的性质和局限性。用户应该知道，使用sign函数计算的uniform更新通常比SGD和自适应方法产生的范数更大。因此，Lion需要更小的学习率lr和更大的解耦权重衰减λ来保持有效的权重衰减强度。详情请参阅第5节。实验表明，Lion的增益随着批量大小的增加而增加，并且与AdamW相比，它对不同的超参数选择具有更强的鲁棒性。由于局限性，Lion和AdamW在一些大规模语言和图像-文本数据集上的差异不具有统计学意义。如果在训练过程中使用强增强或小批量大小(<64)， Lion的优势较小。详情参见第6节。

2、算法的符号发现

我们提出了一种将算法发现公式化为程序搜索的方法(Brameier等人，2007;Koza, 1994;Real等，2020)。我们使用程序形式的符号表示有以下优点:(1)它符合算法必须作为程序来执行的事实;(2)与神经网络等参数化模型相比，程序等符号表示更容易分析、理解并转移到新任务中;(3)程序长度可以用来估计不同程序的复杂程度，从而更容易选择较简单的，往往更具有一般性的程序。这项工作的重点是用于深度神经网络训练的优化器，但该方法一般适用于其他任务。

2.1、程序搜索空间

在设计程序搜索空间时，我们遵循以下三个标准:(1)搜索空间应足够灵活，以支持发现新算法;(2)程序应该易于分析并纳入机器学习工作流程;(3)程序设计应关注高层次的算法设计，而不是底层的实现细节。在命令式语言中，我们将程序定义为包含操作n维数组的函数，包括包含此类数组的列表和字典等结构。它们类似于使用NumPy / JAX的Python代码(Bradbury等人，2018;Harris等人，2020)以及优化算法的伪代码。下面概述了设计的细节，以及程序3中AdamW的示例表示。

输入/输出签名。该程序定义了一个训练函数，对正在搜索的优化算法进行编码，其中主要输入是当前训练步骤的模型权重(w)、梯度(g)和学习率调度值(lr)。主要输出是对权重的更新。该程序还包含了额外的初始化为零的变量，以在训练期间收集历史信息。例如，AdamW需要两个额外的变量来估计第一和第二矩。请注意，这些变量可以任意使用，我们在程序3中使用m和v的名称只是为了更好的可读性。程序2中的简化代码片段使用了与AdamW相同的签名，以确保所发现的算法具有更小或相等的内存占用。与之前的优化器搜索尝试相反(Bello等人，2017;Wang et al.， 2022)，我们的方法允许发现更新额外变量的更好方法。

构建块。train函数由一系列赋值语句组成，对语句或局部变量的数量没有限制。每个语句都使用常量或现有变量作为输入调用函数，结果存储在一个新的或现有的变量中。对于程序，我们选择了45个常用的数学函数，其中大多数对应于NumPy中的函数或线性代数中的操作。为了使程序更紧凑，还引入了一些函数，如线性插值函数interp(x, y, a)，它等价于(1 - a) * x + a * y.初步实验研究了包含更高级的特性，如条件语句和循环语句，以及定义和调用新函数，但这些并没有产生改进的结果，因此我们将它们排除。附录h总结了这些函数的详细描述。必要时，函数参数的类型和形状会自动校正，例如，将数组字典添加到标量时。

突变和冗余语句。进化搜索中使用的突变的设计与程序的表示紧密交织在一起。我们包含三种类型的突变:(1)在随机位置插入一条新语句，其中包含随机选择的函数和参数;(2)删除一条随机选择的语句;(3)通过随机改变其中一个函数参数来修改一条随机语句，该参数可以是变量或常量。为了改变参数，将其替换为一个现有的变量或一个新生成的常数，通过从正态分布X ~ N(0 1)中采样获得。此外，可以通过将现有的常数乘以一个随机因子2a来改变，其中a ~ N(0 1)。这些常数作为优化算法中的可调超参数，如AdamW中的峰值学习率和权重衰减。请注意，我们允许程序在搜索过程中包含冗余语句，即不影响程序最终输出的语句。这是必要的，因为修改仅限于影响一条语句。因此，多余的语句可以作为今后对prog进行重大修改的中间步骤。

无限稀疏搜索空间。考虑到语句和局部变量的数量是无限的，以及可变常量的存在，程序的搜索空间是无限的。即使我们忽略常量并限制程序长度和变量数量，潜在程序的数量仍然难以控制。对可能程序数量的粗略估计是np=nflnvna∗ln_{p}=n_{f}^{l} n_{v}^{n_{a} * l}np​=nfl​nvna​∗l​，其中nfn_{f}nf​是可能函数的数量，nvn_{v}nv​是局部变量的数量，nan_{a}na​是每个语句的平均参数数量，l是程序长度。更重要的是，这种挑战来自于搜索空间中高性能程序的稀疏性。为说明这一点，在一个低成本的代理任务上进行了随机搜索，评估了超过200万个程序。其中最好的程序仍显著低于AdamW。

2.2、高效搜索技术

本文采用以下技术来解决无限和稀疏搜索空间带来的挑战。

热启动和重启的进化。我们应用正则化演化，因为它简单、可扩展，并在许多AutoML搜索任务上显示出成功(Holland, 1992;Real等人，2019,2020;So等人，2019;Ying et al.， 2019)。它保持P个算法的种群，并通过循环逐步改进。每个循环随机选择T<P算法，并选择表现最好的作为父算法，即锦标赛选择(Goldberg and Deb, 1991)。然后复制和变异这个父算法以产生一个子算法，并将其添加到种群中，而最古老的算法则被删除。通常，进化搜索从随机候选开始，但我们将初始种群作为AdamW进行热启动，以加速搜索。默认情况下，我们使用2的锦标赛大小和1K的人口大小。为了进一步提高搜索效率，采用了两种重启方式:(1)从初始程序重新启动，由于进化过程中的随机性，这可以导致不同的局部最优值，并鼓励探索。这可以通过并行运行多个搜索来实现。(2)从目前找到的最好的算法重新启动，以进一步优化它，鼓励利用。图2(左)显示了五个进化搜索实验的均值和标准误差。我们通过只允许进化中常量的突变来运行基于AdamW的超参数调优，并通过随机程序采样来运行随机搜索，两者的计算量都增加了4倍。我们的搜索明显优于两个基线取得的最佳结果，如图中的两条虚线所示。搜索适应度的高方差需要通过从初始程序重新启动进行多次重复。当搜索适应度在~ 300K进展后停滞时，从目前找到的最佳程序重新启动，可以进一步提高橙色曲线所示的适应度。

通过抽象执行进行剪枝。本文提出从3个方面消除程序空间中的冗余:具有语法或类型/形状错误的程序、功能等价的程序和程序中的冗余语句。在程序实际执行之前，我们执行一个抽象执行步骤，该步骤(1)推断变量类型和形状以检测有错误的程序，并不断改变父程序，直到生成有效的子程序;(2)生成一个哈希，唯一地标识如何从输入计算输出，允许我们缓存和查找语义重复的程序(Gillard等人，2023);(3)识别在实际执行和分析过程中可以忽略的冗余语句。例如，去掉程序8中所有冗余语句后，得到程序4。与实际执行相比，抽象执行的成本可以忽略不计，因为每个输入和函数都被自定义值替换，例如哈希。抽象执行的细节见附录一。初步实验表明，搜索过程可能会被无效程序淹没，如果不过滤掉无效程序，就无法取得进展。如图2(右)所示，随着搜索的进行，冗余语句的百分比和缓存命中率都在增加。基于5次搜索，每次覆盖300K个程序，最终有69:8±1:9%的冗余语句，这意味着删除冗余语句平均使程序缩短3倍，从而更容易分析。缓存命中率为89:1±0:6%，表明使用哈希表作为缓存可以减少约10倍的查找开销。

代理任务和搜索成本。为了减少搜索成本，通过减少模型大小、训练样本数量和目标任务的步数来创建低成本的代理。在一个TPU V2芯片上，可以在20min内完成对代理的评估。我们使用验证集上的准确率或复杂度作为适应度。每个搜索实验使用100个TPU V2芯片，运行时间约为72小时。在每次搜索实验中，总共产生了20 - 300k个程序。然而，由于通过抽象执行使用了缓存，实际计算的程序数量约为20-30K。为了结合重新启动，我们开始5次重复的搜索实验，然后从迄今找到的最佳算法开始另一轮搜索初始化。这导致了总成本为~ 3K TPU V2天。有关代理任务的详细信息，请参阅附录F。

2.3、泛化:方案选择与简化

搜索实验可以在代理任务上发现有前途的程序。我们通过增加模型大小和训练步骤来使用比代理任务更大的元验证任务的性能，以选择泛化超出代理任务的程序，然后进一步简化它们。当搜索适合度不断增长，但元验证指标下降时，就会出现元过拟合现象，表明所发现的算法已经过度拟合代理任务。图3(左)显示了两个例子，其中蓝色曲线表示早期元过拟合，橙色曲线表示后期元过拟合。

泛化差距大。由于搜索过程中的代理任务与目标任务之间存在很大的差距，所发现的算法面临着重大挑战。代理任务通常可以在一个TPU V2芯片上在20分钟内完成，而目标任务可以>大104倍，需要在512 TPU V4芯片上进行几天的培训。此外，我们希望优化器在不同的架构、数据集甚至不同的领域都能表现良好，因此发现的算法需要表现出强大的分布外泛化。演化过程中的稀疏搜索空间和固有噪声进一步加剧了这一挑战，导致不同运行之间的泛化属性不一致。我们的观察表明，后来进行元过拟合的进化搜索实验往往会发现泛化效果更好的优化算法。参见图3(右)。

漏斗的选择。为了缩小泛化差距，我们根据搜索适合度收集有前景的程序，并使用一系列元验证任务添加额外的选择步骤，以选择泛化效果更好的程序。为了节省计算，我们应用了一个漏斗选择过程，逐渐增加元验证任务的规模。例如，从代理任务A开始，我们通过增加模型大小和训练步骤来创建一个10倍大的任务B。只有在任务B上超过基线的算法才会在任务C上进行评估，任务C是任务B的100倍。这种方法允许我们逐渐过滤出泛化性能较差的算法，最终导致选择泛化性能较好的算法。

简化。简单的程序更容易理解，我们的直觉是它们更容易泛化，所以我们用以下步骤简化程序。首先，我们删除冗余语句，这些语句对通过抽象执行确定的最终输出没有贡献。其次，我们删除非冗余的语句，但删除时产生的差异最小。这一步也可以通过在突变过程中禁用新语句的插入来实现。最后，我们手动重新排列语句，为变量分配清晰和描述性的名称，并将程序转换为更简单的、数学上等效的形式。

3、Lion的推导与分析

我们之所以使用Lion优化器，是因为它的简单性、内存效率以及在搜索和元验证方面的强大性能。请注意，搜索还发现了附录D中所示的其他现有或新算法，例如，一些具有更好的正则化，一些类似AdaBelief (Zhuang等人，2020)和AdaGrad (Duchi等人，2011)。

3.1、求导

搜索和漏斗选择过程导致程序4，它是通过自动从原始程序8中删除冗余语句来获得的(在附录中)。我们进一步简化它，得到程序1中的最终算法(Lion)。在简化过程中，从程序4中删除了几个不必要的元素。cosh函数被删除，因为m将在下一次迭代中被重新赋值(第3行)。使用arcsin和clip的语句也被删除，因为没有它们，我们观察到质量没有下降。三个红色语句转换为一个符号函数。虽然在程序4中都使用了m和v，但v只是改变了动量的更新方式(两个常数为∼0.9和∼1.1的interp函数相当于一个常数为∼0.99的interp函数)，不需要单独跟踪。注意，偏差修正不再需要，因为它不会改变方向。算法2显示了伪代码。

3.2、分析

信号更新和正则化。Lion算法通过符号操作在所有维度上生成大小一致的更新，这在原则上不同于各种自适应优化器。直观地说，符号操作为更新添加了噪声，这作为一种形式的正则化，并有助于泛化(Chen等人，2022;Foret等人，2021;Neelakantan等人，2017)。一个证据如附录中的图11(右)所示，Lion在ImageNet上训练的ViT-B/16相比AdamW有更高的训练误差，但在验证集上的准确率高2%(如表2所示)。此外，附录G的结果表明，Lion在更平滑的区域会导致收敛，这通常会导致更好的泛化。

动量跟踪。在Lion中用于跟踪动量的默认EMA因子是0.99 (β2)，而AdamW和动量SGD中常用的EMA因子是0.9。在应用符号运算之前，电流梯度和动量用0.9 (β1)的因子插值。这种EMA因子和插值的选择允许Lion在记住动量梯度的~ 10倍长的历史和在更新中对当前梯度施加更多的权重之间进行平衡。β1和β2的必要性将在第4.6节中进一步讨论。

超参数和批处理大小选择。与AdamW和Adafactor相比，Lion更简单，超参数更少，因为它不需要$ϵ$和因式分解相关的参数。如果我们忽略权重衰减项，更新是一个按元素计算的二进制$\pm 1$，比其他优化器(如SGD和自适应算法)产生的范数更大。因此，Lion需要更小的学习率和更大的去耦权值衰减来实现类似的有效权值衰减强度($\mathrm{lr}\ast \lambda$)。关于调优Lion的详细信息可以在第5节中找到。此外，Lion相对于AdamW的优势随着批大小的增加而增大，这符合通过数据并行扩大模型训练的常见做法(章节4.6)。

内存和运行时优势。Lion只保存动量，因此比AdamW等流行的自适应优化器占用的内存更小，这在训练大型模型和/或使用大批量时是有益的。例如，AdamW需要至少16个TPU V4芯片来训练图像分辨率为224、批处理大小为4096的ViT-B/16，而Lion只需要8个(都有bfloat16动量)。另一个实际的好处是，在我们的实验中，由于Lion的简单性，它的运行时间(步数/秒)更快，通常比AdamW和Adafactor快2-15%，这取决于任务、代码库和硬件。

与现有优化器的关系。符号操作已经在之前的优化器中进行了探索(Bernstein et al.， 2018;Riedmiller和Braun, 1993)。最接近我们的是手工优化器signSGD (Bernstein et al.， 2018)(及其动量变体)，它也利用符号操作来计算更新，但与Lion具有不同的动量更新规则。他们的重点是减轻分布式训练中代理之间的通信成本，他们在训练ConvNets进行图像分类任务时观察到性能较差。另一方面，NAdam (Dozat, 2016)结合更新的第一矩和梯度来计算更新，但Lion解耦了动量跟踪以及它如何通过β2应用于更新。在第4.5节中可以找到Lion与相关优化器的比较。

4、Lion评估

在本节中，我们介绍在各种基准上对Lion的评估。我们主要将它与AdamW(或Adafactor，当内存是瓶颈时)进行比较，因为它非常流行，并且在大多数学习任务中是事实上的标准优化器。momentum SGD的结果只包含在ResNet中，因为它在其他地方的表现比AdamW更差。我们还在4.5节中对其他流行的优化器进行了基准测试，它们的性能和学习曲线与AdamW类似。我们确保每个优化器都针对每个任务进行了优化(有关优化细节，请参阅第5节)。默认情况下，学习率调度是带有10K步预热的余弦衰减，动量保存为bfloat16以减少内存占用。

4.1、图像分类

在图像分类任务上进行了包括各种数据集和架构的实验(有关数据集的详细信息，请参见附录B)。除了在ImageNet上从头开始训练外，还在两个更大的成熟数据集ImageNet- 21k和JFT上进行预训练(Sun等人，2017)。图像大小默认为2242，否则由下标指定。

在ImageNet上从头开始训练。以下是之前的作品(Dosovitskiy等人，2021;He等人，2016)，我们以1024个批次大小训练ResNet-50 90个批次，以及以4096个批次大小训练300个批次的其他模型。如表2所示，Lion在各种体系结构上的性能明显优于AdamW。实验结果表明，在容量较大的模型上，改进效果更明显，ViT-B/16和ViT-S/16的准确率分别提高了1.96%和0.58%。归纳偏差越少，性能差距也会越大。当应用强增强时，Lion相对于AdamW的增益会缩小，但尽管在训练过程中进行了强正则化，但它在CoAtNet-3上的性能仍比AdamW高出0.42% (Dai等人，2021)。

在ImageNet-21K上进行预训练。在ImageNet-21K上对ViT-B/16和ViT-L/16进行了90次预训练，批处理大小为4096。表2显示，即使训练集扩大10倍，Lion仍然超过AdamW。更大的模型上的差距始终更大，ViT-L/16和ViT-B/16分别为+0.52%和+0.33% (ImageNet)， +0.57%和+0.23%(真实)，+0.74%和+0.25% (V2)。

对JFT进行预训练。为了突破极限，我们在JFT上进行了广泛的实验。我们遵循Dosovitskiy等人(2021)和Zhai等人(2021)的设置来进行预训练和微调。图1(左)和4展示了三种ViT模型(ViT- b /16、ViT- l /16和ViT- h /14)在JFT-300M上不同预训练预算下的准确性。Lion使ViT-L/16与AdamW在ImageNet和ImageNet V2上训练的ViT-H/14的性能相匹配，但预训练成本减少了3倍。在ImageNet ReaL上，计算量进一步减少了5倍。另一个证据是，即使由Zhai et al.(2021)的AdamW训练ViT-L/16 400万步，其性能仍然落后于Lion训练的100万步相同模型。

表3显示了微调结果，具有更高的分辨率和Polyak平均。所提出的ViT-L/16与AdamW训练的之前的ViT-H/14结果相匹配，同时缩小了2倍。在更具挑战性的基准上，例如ViT-L/16的+1.33% (V2)， +6.08% (A)， +5.54% ®，优势更大。将预训练数据集扩展到JFT-3B后，Lion训练的ViT-g/14优于之前的ViT-g/14结果(Zhai等人，2021)，参数减少了1.8倍。所提出的ViT-G/14在ImageNet上进一步取得了90.71%的准确率。

4.2、视觉-语言对比学习

本节重点介绍CLIP风格的视觉-语言对比训练(Radford et al.， 2021)。我们不是从头开始学习所有参数，而是使用一个强预训练的模型初始化图像编码器，因为它被认为更有效(Zhai et al.， 2022)。

锁定图像文本调整(LiT)。通过使用相同的冻结预训练ViT，以对比的方式训练文本编码器(Zhai et al.， 2022)，我们在LiT上对Lion和AdamW进行比较(Zhai et al.， 2022)。所有模型训练1B个图像-文本对，批大小为16384个。表4给出了三种模型尺度下的零拍图像分类结果，名称表示大小，如:LiT-B/16- b表示ViT-B/16，文本编码器为基本尺寸Transformer。我们的Lion方法在LiT-B/32-B、LiT-B/16-B和LiT-g/14288-L的零镜头ImageNet精度上分别提高了+1.10%、+1.13%和+0.66%，与AdamW相比，得到了持续的改善。图5(左)描述了LiT-B/16-B的零次射击学习曲线示例。在其他两个数据集上也得到了类似的结果。在MSCOCO (Lin et al.， 2014)和Flickr30K (Plummer et al.， 2015)上的零镜头图像-文本检索结果见图9(附录)。评估指标是Recall@K，根据查询的ground truth标签是否出现在检索到的前k个示例中计算。Lion在两个数据集上都优于AdamW，在Flicker30K上，Recall@1的增益比Recall@10更大，这意味着更准确的检索结果:图像+1.70% vs +0.60% !文本+2.14% vs.文本+0.20% !的形象。

BASIC。Pham等人(2021)提出同时扩大批处理规模、数据集和模型规模，实现了对CLIP的大幅改进。它使用在JFT-5B上预训练的复杂CoAtNet (Dai et al.， 2021)作为图像编码器。在6.6B图像-文本对上进行对比训练，批大小为65,536。为了突破极限，我们只在最大的BASIC-L上进行实验，并在图像编码器预训练和对比学习阶段使用Lion。如表1所示，我们在基线上获得了显著的2.6%的增益，在零样本 ImageNet分类上达到了88.3%的精度。请注意，这一结果比之前的最佳结果(Yu et al.， 2022)高2.0%。在其他五个健壮性基准测试中，性能增益是一致的。在Lion获得的BASIC-L中对图像编码器(CoAtNet-7)进行微调后，我们进一步在ImageNet上实现了91.1%的top-1精度，比之前的SOTA提高了0.1%。

4.3、扩散模型

最近，扩散模型在图像生成方面取得了巨大的成功(Dhariwal和Nichol, 2021;Ho和Salimans, 2022;Ho等人，2020;撒哈拉等人，2022;Song等人，2021)。鉴于其巨大的潜力，我们测试了Lion在无条件图像合成和多模态文本到图像生成方面的性能。

ImageNet上的图像合成。我们利用Dhariwal和Nichol(2021)中介绍的改进的U-Net架构，并在ImageNet上执行64×64, 128×128和256×256图像生成。批大小设置为2048，在整个训练过程中学习率保持不变。对于解码，我们在没有分类器引导的情况下应用DDPM (Ho等人，2020)进行1K采样步骤。评估指标是标准FID评分。如图1(右)和图5(中、右)所示，Lion在FID评分上实现了更好的质量和更快的收敛。请注意，Lion和AdamW之间的差距随着图像分辨率的增加而增加，此时生成任务变得更具挑战性。在生成256 × 256张图像时，Lion在440K步长达到AdamW的最终性能，减少2.3x迭代。最终的FID得分为4.1 (Lion) vs. 4.7 (AdamW)，作为参考，ADM (Dhariwal and Nichol, 2021)的FID为10.94。

Text-to-image 生成。我们遵循Imagen (sahara et al.， 2022)设置来训练一个基本的64 × 64文本到图像模型和一个64 × 64 !256 × 256超分辨率模型。所有模型都在一个高质量的内部图像-文本数据集上进行训练，该数据集的批处理大小为2048，并且具有恒定的学习率。由于计算的限制，我们的基本U-Net的宽度为192，而原始2B模型的宽度为512，而600M超分辨率模型与原始Imagen设置相同。在训练的同时，从MSCOCO (Lin et al.， 2014)验证集中采样2K图像用于实时评估。我们使用CLIP评分来衡量图像-文本对齐，使用零镜头FID-30K来衡量图像保真度。应用了权重为5.0的无分类器引导(Ho和Salimans, 2022)，因为它已被证明可以改善图像-文本对齐。图6描述了学习曲线。虽然在64 × 64的基础模型上没有明显的改进，但Lion在文本条件超分辨率模型上优于AdamW。与AdamW相比，它实现了更高的CLIP评分，并且具有更少的噪声FID指标。

4.4、语言建模与微调

本节主要讨论语言建模和调优。在只有语言的任务中，我们发现调整β1和β2可以提高AdamW和Lion的质量。有关调优的详细信息，请参阅第5节。

自回归语言建模。在Hua等人(2022)之后，我们首先在两个较小规模的学术数据集Wiki- 40B (Guo等人，2020)和PG-19 (Rae等人，2020)上进行实验。变压器采用小型(110M)、中型(336M)、大型(731M)三个尺度。架构细节可以在附录e中找到。所有模型每批训练125K步，使用218个令牌，学习速率计划为10K步热身，然后是线性衰减。对于Wiki-40B，上下文长度设置为512,PG-19设置为1024。图7说明了Wiki-40B的令牌级困惑度和PG-19的单词级困惑度。Lion始终比AdamW实现更低的验证复杂度。当在Wiki-40B和PG-19上训练中等大小的模型时，它分别实现了1.6倍和1.5倍的加速。当模型增加到大尺寸时，PG-19上的加速进一步增加到2倍。

扩大语言模型和预训练数据集的规模已经彻底改变了NLP领域。所以我们进一步进行更大规模的实验。我们的预训练数据集类似于GLaM中使用的数据集(Du et al.， 2022)，由跨越广泛自然语言用例的1.6万亿令牌组成。遵循GPT-3 (Brown et al.， 2020)，我们训练了三个模型，参数范围从1.1B到7.5B，用于300B令牌，批大小为3M令牌，上下文长度为1K。我们在三个自然语言生成(NLG)和21个自然语言理解(NLU)任务上对它们进行了评估(参见附录C的任务细节)。在这个庞大的数据集上，我们在整个训练过程中没有观察到困惑度的差异。然而，Lion在平均上下文内学习能力方面优于Adafactor，如表5所示。我们的7.5B基线模型，训练300B令牌，优于8B PaLM，训练780B令牌，证明了我们的设置的强度。Lion在NLG和NLU任务上都优于Adafactor，特别是在NLG任务上，1.1B、2.1B和7.5B模型的精确匹配性能分别提高了+1.0、+0.9和+0.6。

掩码语言建模。我们还在C4数据集上执行BERT训练(Raffel et al.， 2020)。它要求语言模型在输入序列中重新构造随机屏蔽的标记。我们使用与小规模自回归实验相同的架构和训练设置。在验证复杂度方面，Lion的表现略好于AdamW: 4.18 vs. 4.25(小)，3.42 vs. 3.54(中等)，3.18 vs. 3.25(大)。学习曲线参见附录中的图11(左)。

微调。我们在GLUE基准(Wang et al.， 2019a)上对Base (220M)、Large (770M)和最大的11B T5 (Raffel et al.， 2020)模型进行了微调。每个模型都经过了500K步的微调，批处理大小为128，学习率恒定。表6显示了GLUE开发集中的结果。对于MRPC和QQP，我们报告F1 / Accuracy分数，对于STS-B，我们报告Pearson / Spearman相关性，对于其他数据集，我们报告它们的默认度量。平均而言，Lion在所有三个模型量表上都击败了AdamW。它在T5基础、大型和11B模型的12分中分别获得10分、12分和10分。

4.5、与其他流行优化器的比较

我们还使用了四种流行的优化器RAdam (Liu et al.， 2020)、NAdam (Dozat, 2016)、AdaBelief (Zhuang et al.， 2020)和AMSGrad (Reddi et al.， 2018)在ImageNet上训练vt - s /16和vt - b /16(使用RandAug和Mixup)。我们彻底调整每个优化器的峰值学习率lr和解耦权重衰减λ (Loshchilov和Hutter, 2019)，而其他超参数则被设置为Optax的默认值——一个用于JAX的梯度处理和优化库。从表7可以看出，Lion仍然是性能最好的。我们注意到，基线之间没有明显的赢家。AMSGrad在vv - s /16上表现最好，但在vv - b /16上表现最差。图10(附录)进一步显示了五种自适应优化器的学习曲线非常相似，而Lion有一个独特的学习速度更快的优化器。

4.6、消融

动量跟踪。为了消除β1和β2的影响，我们比较了一个简单的更新规则:m = interp(g, m， β);Update =符号(m)。两个优化器，Ablation0.9和Ablation0.99，其β值分别为0.9和0.99。从表7中可以看出，两种烧蚀优化算法的性能都不如所有五个比较基线，更不用说我们的Lion了。对语言建模任务的进一步研究(如附录中的图12所示)得出了类似的结论。这些结果验证了使用两个线性插值函数的有效性和必要性，让Lion记住更长的梯度历史，同时为当前梯度分配更高的权重。

批次大小的影响。有些人可能会质疑Lion是否需要大批量才能准确确定方向，因为标识操作会增加噪音。为了解决这个问题，我们在ImageNet上使用各种批处理大小训练vvi - b /16模型，同时保持总训练周期为300，并结合RandAug和Mixup技术。如图8(左)所示，AdamW的最佳批大小为256,Lion的最佳批大小为4096。这表明Lion确实更喜欢更大的批处理大小，但即使是较小的64批处理大小，它的性能仍然很健壮。此外，当批量大小扩大到32K，导致只有11K的训练步骤时，Lion比AdamW获得了2.5%的显著精度提高(77.9%对75.4%)，证明了它在大批量训练设置中的有效性。

5、超参数调优

为了确保公平的比较，我们使用对数刻度调整AdamW (Adafactor)和Lion的峰值学习率l r和解耦权重衰减\\lambda。AdamW中\\beta ＿1{和}\\beta ＿2{的默认值分别设置为0.9和0.999,}\\epsilon为$1e-8$，而Lion中β1\\beta_1β1​和β2\\beta_2β2​的默认值通过程序搜索过程发现，分别设置为0.9和0.99。我们只在语言任务(章节4.4)中调优这些超参数，其中AdamW中的β1=0.9,β2=0.99\\beta_1{=0.9, }\\beta_2=0.99β1​=0.9,β2​=0.99, Lion中的β1=0.95,β2=0.98\\beta_1{=0.95, }\\beta_2=0.98β1​=0.95,β2​=0.98。此外，AdamW中的}\\epsilon被设置为1 e-6，而不是默认的$1e-8$，因为它提高了我们实验中的稳定性，类似于RoBERTa中的观察结果(Liu et al.， 2019b)。

Lion生成的更新是一个元素二进制±1，由于符号操作的结果，因此它比其他优化器生成的更新具有更大的范数。根据我们的经验，Lion的合适学习率通常比AdamW小3-10倍。由于有效权值衰减为lr * λ: update += w * λ;update *= lr, Lion的λ值比AdamW的λ值大3-10倍，以保持相似的强度。例如:(1)在强增强的ImageNet上训练ViT-B/16时，Lion中lr = 1e−4，λ = 10:0; AdamW中lr = 1e−3，λ = 1:0;(2)在Lion模型中lr = 3e−5，λ = 0:1;在AdamW模型中lr = 3e−4，λ = 0:01;(3)对于7.5B语言建模，Lion中lr = 1e−4，λ = 0:01, Adafactor中lr = 1e−3，λ = 0:001。所有超参数请参见表12(附录)。

除了峰值性能之外，对超参数的敏感性和调优它们的难度对于在实践中采用优化器也是至关重要的。在图8(中右)中，我们在ImageNet上从头训练vvi - b /16时改变lr和λ。从热图中可以看出，与AdamW相比，Lion对于不同的超参数选择更加稳健。

6、限制

搜索的限制。尽管努力减少搜索空间的限制，但它仍然受到流行的一阶优化算法的启发，导致对类似算法的偏见。它还缺乏构建高级二阶算法所需的功能(Anil等人，2020;Gupta等人，2018;马丁斯和格罗斯，2015)。搜索成本仍然很大，算法简化需要人工干预。进一步减小搜索空间中的偏差，发现更多新颖的算法，提高搜索效率是未来的重要方向。当前的程序结构非常简单，因为我们没有发现更高级的程序构造(如条件语句、循环语句和定义新函数)的良好用法。探索如何将这些元素结合起来，就有可能开启新的可能性。

Lion的局限性。虽然我们努力在尽可能多的任务上评估Lion，但评估仅限于选定的任务。在视觉任务上，当使用强增强时，Lion和AdamW之间的性能差距缩小。Lion还有几个任务的执行情况与AdamW类似，包括:(1)Imagen文本到图像的基础模型，(2)在大规模内部数据集上训练的自回归语言模型的复杂度，这可以说是一个更可靠的上下文学习基准指标，以及(3)C4上的掩码语言建模。这些任务有一个共同的特征，即数据集是大量的和高质量的，这导致了优化器之间的差异减小。另一个潜在的限制是批量大小。尽管人们经常扩大批处理大小以实现更多的并行性，但如果批处理大小较小(<64)， Lion的性能可能并不比AdamW好。此外，Lion仍然需要在bfloat16中跟踪动量，这对于训练大型模型来说可能是昂贵的。一个潜在的解决方案是分解动量以节省内存。

7、相关工作

我们的工作在于AutoML和元学习领域，包括学习如何学习(Andrychowicz等人，2016;贝罗等人，2017;梅茨等人，2019,2022;拉维和拉罗谢尔，2017年;Wichrowska等人，2017;Xiong等人，2022)，神经结构搜索(Chen和Hsieh, 2020;Chen等人，2021;刘等，2019a;Pham等人，2018;Real等人，2019年;So et al.， 2019;Wang et al.， 2021;Yang等，2022;Zoph和Le, 2017)和超参数优化(Dong等人，2021;Hutter et al.， 2011;贾米森和塔尔沃卡，2016年;Li et al.， 2017)等。使用进化方法来搜索程序也有很长的历史，即遗传编程(Brameier et al.， 2007;荷兰,1992;Koza, 1994)。我们的方法建立在一个类似于AutoML-Zero的符号搜索空间上(Peng等人，2020;Real等，2020)。然而，我们的目标不是为玩具任务发现具有固定维度矩阵、向量和标量的程序，而是开发在n维数组上操作的程序，并可以推广到最先进的任务。其他相关工作包括大量手工优化器(Anil等人，2020;伯恩斯坦等人，2018;Dozat, 2016;Duchi等人，2011;Gupta等人，2018;金玛与巴，2014;Liu et al.， 2020;Reddi等人，2018;Riedmiller和Braun, 1993;Shazeer和Stern, 2018;Zhuang等人，2020)，我们在第3.2节中讨论。

8、结论

本文提出通过程序搜索来发现优化算法。我们提出了一些技术来解决搜索无限和稀疏的搜索空间的挑战，以及代理任务和目标任务之间的大泛化差距。我们的方法发现了一个简单而有效的优化器Lion，它是内存高效的，并实现了跨架构、数据集和任务的强泛化。

致谢

我们要感谢(按字母顺序)Angel Yu, Gong Boqing, Chen Cheng, Chitwan sarah, Daiyi Peng, David So, Hanxiao Liu, Hanzhao Lin, Jeff Lund, Jiahui Yu, jingu Xu, Julian Grady, Junyang Shen, Kevin Regan, Li Sheng, Liu Yang, Martin Wicke, Mingxing Tan, Mohammad Norouzi, Qiqi Yan, Rakesh Shivanna, Rohan Anil，高瑞琪，Steve Li, Vlad Feinberg，张文博，William Chan，小王，翟晓华，李亚光，Yang Li, zhushu Li, Zihang Dai, zirrui Wang的帮助讨论，谷歌Brain团队提供了一个支持性的研究环境。

唱吧歌曲