谷歌新出优化器内存小 效率高 训练GPT2果然快
发布时间：2023-02-21 08:56:29 所属栏目：动态 来源：
导读： 优化器即优化算法，在神经网络训练中起着关键作用。近年来，研究者引入了大量的手工优化器，其中大部分是自适应优化器。Adam 以及 Adafactor 优化器仍然占据训练神经网络的主流，尤其在语言、视觉和多模态领域更
|
优化器即优化算法，在神经网络训练中起着关键作用。近年来，研究者引入了大量的手工优化器，其中大部分是自适应优化器。Adam 以及 Adafactor 优化器仍然占据训练神经网络的主流，尤其在语言、视觉和多模态领域更是如此。 在本文中，来自谷歌、 UCLA 的研究者提出了一种通过程序搜索发现深度神经网络训练的优化算法的方法，进而发现 Lion（EvoLved Sign Momentum）优化器。实现这一目标面临两个挑战：首先是在无限稀疏的程序空间中寻找高质量的算法，其次是选择可以从小型任务泛化到更大、SOTA 任务的算法。为了应对这些挑战，该研究采用了一系列技术，包括具有热启动和重启的进化搜索、抽象执行、funnel 选择和程序简化。 Lion 在一系列模型（Transformer、MLP、ResNet、U-Net 和 Hybrid）和任务（图像分类、视觉 - 语言对比学习、扩散、语言建模和微调）中展示了出色的性能。值得注意的是，通过在 BASIC 中用 Lion 替换 Adafactor，该研究在 ImageNet 上实现了 88.3% 的零样本和 91.1% 的微调精度（Pham 等人，2021），分别超过之前的 SOTA 结果 2% 和 0.1%。 本文使用程序形式的符号表示具有以下优点：（1）它符合算法必须作为程序来执行的事实；(2) 与神经网络等参数化模型相比，程序等符号表示更易于分析、理解和迁移到新任务；(3) 程序长度可用于估计不同程序的复杂性，从而更容易选择更简单、通常更具通用性的程序。这项工作侧重于深度神经网络训练的优化器，但该方法通常适用于其他任务。 该研究采用以下技术来解决无限和稀疏搜索空间带来的挑战。首先应用正则化，因为它简单，可扩展，并在许多 AutoML 搜索任务中很成功；其次是简化程序空间中的冗余；最后，为了降低搜索成本，该研究通过减少模型大小、训练示例的数量和距离目标任务的 step 使成本降低。 搜索和 funnel 选择过程促成了程序 4，它通过从原始程序 8（附录）自动删除冗余语句获得。研究者进一步简化并得到程序 1 中的最终算法（Lion）。在简化过程中从程序 4 中删除了几个不必要的元素。其中 cosh 函数被删除，因为 m 将在下一次迭代中重新分配（第 3 行）。使用 arcsin 和 clip 的语句也被删除，因为研究者观察到没有它们质量不会下降。三个红色语句转换为一个符号函数。 动量跟踪。与 AdamW 和动量 SGD 中常用的 0.9 相比，Lion 中用于跟踪动量的默认 EMA 因子为 0.99 (β_2)。EMA 因子和插值的这种选择，使得 Lion 在记住动量梯度的 10 倍 history 和在更新中对当前梯度施加更多权重之间取得平衡。 内存和运行时优势。Lion 只保存动量，比 AdamW 等流行的自适应优化器占用的内存更小，这在训练大模型和 / 或使用大批量时非常有用。例如，AdamW 需要至少 16 个 TPU V4 芯片来训练图像分辨率为 224 且批大小为 4,096 的 ViT-B/16，而 Lion 只需要 8 个（均具有 bfloat16 动量）。 研究者在图像分类任务上执行涵盖各种数据集和架构的实验。除了在 ImageNet 上从头开始训练之外，他们还在两个更大的成熟数据集 ImageNet-21K 和 JFT 上进行预训练。图片大小默认为 224。 首先在 ImageNet 上从头开始训练。研究者训练 ResNet-50 达到 90 个 epoch，批大小为 1,024，其他模型训练 300 个 epoch，批大小为 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 显示了 JFT-300M 上不同预训练预算下三种 ViT 模型（ViT-B/16、ViT-L/16 和 ViT-H/14）的准确率。Lion 使 ViT-L/16 能够与 AdamW 在 ImageNet 和 ImageNet V2 上训练的 ViT-H/14 性能相匹配，但预训练成本降低了 3 倍。 本节重点介绍 CLIP 风格的视觉语言对比训练。研究者没有从头开始学习所有参数，而是使用强大的预训练模型初始化图像编码器。 对于锁定图像文本调整 (LiT) ，研究者通过使用相同的 frozen 预训练 ViT 以对比方式训练文本编码器，进而在 LiT 上对 Lion 和 AdamW 进行比较。下表 4 显示了 3 种模型尺度上的零样本图像分类结果，Lion 展示出了对 AdamW 的持续改进。 最近，扩散模型在图像生成方面取得了巨大的成功。鉴于其巨大的潜力，研究者测试了 Lion 在无条件图像合成和多模式文本到图像生成方面的表现。 对于 ImageNet 上的图像合成，研究者利用 2021 年论文《 Diffusion models beat gans on image synthesis 》中引入的改进 U-Net 架构，在 ImageNet 上执行 64×64、128×128 和 256×256 图像生成。如上图 5（中和右）所示，Lion 可以在 FID 分数上实现更好的质量和更快的收敛速度。 （编辑：汽车网） 【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容! |
