当前位置: 首页 > 科技观察

彻底碾压AdamW!谷歌新优化器内存小效率高,网友:训练GPT2真快

时间:2023-03-20 12:02:35 科技观察

optimizer就是优化算法,在神经网络训练中起着关键作用。近年来,研究人员推出了大量手工制作的优化器,其中大部分是自适应优化器。Adam和Adafactor优化器仍然占据训练神经网络的主流,尤其是在语言、视觉和多模态领域。除了手动引入优化器,还有一个方向就是程序自动发现优化算法。之前有人提出L2O(learningtooptimize),通过训练神经网络来发现优化器。然而,这些黑盒优化器通常是在有限数量的小任务上训练的,因此很难推广到大型模型。已经尝试了其他方法,他们应用强化学习或蒙特卡洛采样来发现新的优化器。然而,为了简化搜索,这些方法通常会限制搜索空间,从而限制发现其他优化器的可能性。因此,目前的方法还没有达到最先进的水平。近年来,值得一提的是AutoML-Zero,它试图在评估任务时搜索机器学习管道的每个组件,对于优化器的发现非常有用。在这篇论文中,来自谷歌和加州大学洛杉矶分校的研究人员提出了一种方法,通过程序搜索来发现深度神经网络训练的优化算法,进而发现Lion(EvoLvedSignMomentum)优化器。实现这一目标面临两个挑战:首先,在无限稀疏的程序空间中寻找高质量的算法,其次,选择能够从小任务泛化到更大的SOTA任务的算法。为了应对这些挑战,该研究采用了一系列技术,包括具有热重启和重启的进化搜索、抽象执行、漏斗选择和程序简化。论文地址:https://arxiv.org/pdf/2302.06675.pdf项目地址:https://github.com/google/automl/tree/master/lion用AdamW和各种自适应优化器(这些优化器需要先保存两者和第二时刻),Lion只需要动量并利用符号操作来计算更新,并将额外的内存占用减少了一半。这在训练大型模型或大批量时很有用。例如,AdamW至少需要16个TPUV4来训练图像大小为224、批量大小为4096的ViT-B/16,而Lion只需要8个。另一个实际好处是,由于Lion的简单性,Lion在实验中具有更快的运行时间(步数/秒),通常比AdamW和Adafactor快2-15%,具体取决于任务、代码库和硬件。论文第一作者陈湘宁表示:我们的符号程序搜索找到了一个高效的只跟踪动量的优化器——Lion。与Adam相比,它实现了88.3%的零样本和91.1%的微调ImageNet精度,以及高达5倍(与ViT相比)、2.3倍(与扩散模型相比)和2倍(与LM相比)。)训练效率。Lion在一系列模型(Transformer、MLP、ResNet、U-Net和Hybrid)和任务(图像分类、对比视觉语言学习、扩散、语言建模和微调)中展示了出色的性能。值得注意的是,通过在BASIC中用Lion替换Adafactor,该研究在ImageNet上实现了88.3%的零样本和91.1%的微调精度(Pham等人,2021),分别超过之前的SOTA结果2%和0.1%。%。此外,Lion将JFT上的预训练计算减少了5倍,将扩散模型的训练效率提高了2.3倍,获得了更好的FID分数,并在语言建模方面提供了类似或更好的性能。性能,最多可节省2倍的计算量。Twitter用户crumb表示:我使用谷歌的Lion优化器训练了一个124M参数的GPT2模型,发现实现与Adam相同的损失所需的步数减少了37.5%。图源:https://twitter.com/aicrumb/status/1626053855329898496算法的符号发现本文采用程序形式的符号表示具有以下优点:(1)符合该算法必须作为程序执行的事实;(2)兼容程序等符号表示比神经网络等参数模型更容易分析、理解和迁移到新任务;(3)程序长度可以用来估计不同程序的复杂性,从而更容易选择更简单、通常更通用的程序。这项工作侧重于深度神经网络训练的优化器,但该方法通常适用于其他任务。下图中,Program2这个简化的代码片段使用了与AdamW相同的签名,以确保发现的算法具有更小或相等的内存占用;程序3中给出了AdamW的示例表示。本研究采用以下技术来解决无限和稀疏搜索空间带来的挑战。正则化首先被应用是因为它简单、可扩展并且在许多AutoML搜索任务中都很成功;其次,简化程序空间的冗余;最后,为了降低搜索成本,研究减少了模型大小,训练示例的数量和目标任务的步骤降低了成本。左图:显示了五个进化搜索实验的平均值和标准误差。右图:随着搜索的进行,冗余语句的百分比和缓存命中率都会增加。据Lion推导分析研究人员介绍,优化器Lion具有简单、内存效率高、搜索和元验证性能强等特点。推导搜索和漏斗选择过程导致程序4,它是通过从原始程序8(附录)中自动删除冗余语句获得的。研究人员在程序1中进一步简化得到了最终的算法(Lion)。程序4在简化过程中去掉了几个不需要的元素。其中删除了cosh函数,因为m将在下一次迭代中重新分配(第3行)。使用arcsin和clip的语句也被删除,因为研究人员观察到没有它们没有质量下降。三个红色语句转化为一个符号函数。尽管程序4中同时使用了m和v,但v仅改变了动量的更新方式(两个常数为~0.9和~1.1的插值函数等同于一个常数为~0.99的函数),不需要单独跟踪。请注意,不再需要偏差校正,因为它不会改变方向。下面的算法2显示了伪代码。分析符号更新和正则化。Lion算法通过符号运算在所有维度上产生幅度统一的更新,这在原理上不同于各种自适应优化器。直觉上,符号操作会增加更新的噪音,充当一种正则化形式并有助于泛化。下面的图11(右)显示了一个证明。动量跟踪。Lion中用于跟踪动量的默认EMA因子为0.99(β_2),而AdamW和动量SGD中通常为0.9。EMA因子和插值的这种选择允许Lion在记住10倍的动量梯度历史和在更新中将更多权重放在当前梯度之间取得平衡。超参数和批量大小选择。Lion比AdamW和Adafactor更简单并且超参数更少,因为它不需要ε和因式分解相关参数。Lion需要更小的学习率,这反过来需要更大的解耦权重衰减以达到类似的有效权重衰减强度(lr*λ)。内存和运行时优势。与AdamW等流行的自适应优化器相比,Lion仅节省动量并且内存占用更小,这在训练大型模型和/或使用大批量时非常有用。例如,AdamW需要至少16个TPUV4芯片来训练图像分辨率为224和批量大小为4,096的ViT-B/16,而Lion只需要8个(全部使用bfloat16动量)。Lion评估结果在实验部分,研究人员在各种基准测试中对Lion进行了评估,主要是将其与流行的AdamW(或当内存成为瓶颈时的Adafactor)进行比较。图像分类研究人员对涵盖各种数据集和架构的图像分类任务进行实验。除了在ImageNet上从头开始训练外,他们还在两个更大的成熟数据集ImageNet-21K和JFT上进行预训练。图像大小默认为224。首先在ImageNet上从头开始训练。研究人员对ResNet-50进行了90个时期的训练,批次大小为1,024,其他模型训练了300个时期,批次大小为4,096。如下表2所示,Lion在各种架构上明显优于AdamW。其次是在ImageNet-21K上进行预训练。研究人员在ImageNet-21K上对ViT-B/16和ViT-L/16进行了90个epoch的预训练,批量大小为4,096。下面的表2显示,即使训练集扩大10倍,Lion的表现仍然优于AdamW。最后在JFT上进行预训练。为了突破极限,研究人员对JFT进行了大量实验。下面的图4显示了三种ViT模型(ViT-B/16、ViT-L/16和ViT-H/14)在JFT-300M上不同预训练预算下的准确率。Lion使ViT-L/16能够与AdamW在ImageNet和ImageNetV2上训练的ViT-H/14的性能相媲美,但预训练成本要低3倍。下面的表3显示了具有更高分辨率和Polyak平均的微调结果。研究人员使用的ViT-L/16与AdamW之前训练的ViT-H/14的结果相匹配,同时将参数减少了2倍。在将预训练数据集扩展到JFT-3B后,Lion训练了ViT-g/14优于之前的ViT-G/14结果,参数减少了1.8倍。对比视觉语言学习本节重点介绍CLIP风格的对比视觉语言训练。研究人员没有从头开始学习所有参数,而是使用强大的预训练模型初始化图像编码器。对于LockedImageTextTuning(LiT),我们通过使用相同的冻结预训练ViT以对比方式训练文本编码器来比较LiT上的Lion和AdamW。下面的表4显示了3个模型尺度下的零样本图像分类结果,其中Lion显示了对AdamW的持续改进。下面的图5(左)显示了LiT-B/16-B的示例零样本学习曲线,在其他两个数据集上获得了类似的结果。扩散模型最近,扩散模型在图像生成方面取得了巨大成功。鉴于其巨大的潜力,研究人员测试了Lion在无条件图像合成和多模式文本到图像生成方面的性能。对于ImageNet上的图像合成,研究人员使用2021年论文《 Diffusion models beat gans on image synthesis 》中介绍的修改后的U-Net架构在ImageNet上进行了64×64、128×128和256×256图像生成。如上图5(中和右)所示,Lion在FID分数上实现了更好的质量和更快的收敛。对于文本到图像的生成,下面的图6显示了学习曲线。Lion在文本条件超分辨率模型上优于AdamW,尽管64×64基本模型没有显着提高。与AdamW相比,Lion获得了更高的CLIP分数并且FID指标的噪声更小。语言建模和微调本节重点介绍语言建模和微调。在纯语言任务上,研究人员发现调整β_1和β_2可以提高AdamW和Lion的质量。对于自回归语言建模,下图7显示了Wiki-40B的token-levelperplexity和PG-19的word-levelperplexity。Lion始终实现比AdamW更低的验证困惑。在Wiki-40B和PG-19上训练中型模型时,它分别实现了1.6倍和1.5倍的加速。当模型增加到大尺寸时,PG-19进一步实现了2倍的加速。对于掩码语言建模,研究人员还在C4数据集上进行了BERT训练,其中Lion在验证困惑度上的表现略好于AdamW。相关的学习曲线如下图11(左)所示。为了进行微调,研究人员在GLUE基准测试中对Base(220M)、Large(770M)和最大的11BT5模型进行了微调。下面的表6显示了GLUE开发集的结果,其中,平均而言,Lion在所有3种模型尺寸上都击败了AdamW。与其他流行优化器的比较该研究还使用四种流行的优化器RAdam、NAdam、AdaBelief和A??MSGrad在ImageNet上训练ViT-S/16和ViT-B/16(使用RandAug和Mixup)。Lion仍然是表现最好的,如下表7所示。有关更多技术细节,请参阅原始论文。