[论文解读] Symbolic Discovery of Optimization Algorithms

Xiangning Chen, Liang Chen+10 more|arXiv (Cornell University)|Feb 13, 2023

Machine Learning and Data Classification被引用 167

一句话总结

论文将优化算法发现表述为程序搜索，并识别出一个简单、内存高效的优化器 Lion，它通过动量跟踪和基于符号的更新，在视觉、语言和扩 diffusion 任务中提升性能。

ABSTRACT

We present a method to formulate algorithm discovery as program search, and apply it to discover optimization algorithms for deep neural network training. We leverage efficient search techniques to explore an infinite and sparse program space. To bridge the large generalization gap between proxy and target tasks, we also introduce program selection and simplification strategies. Our method discovers a simple and effective optimization algorithm, $ extbf{Lion}$ ($ extit{Evo$ extbf{L}$ved S$ extbf{i}$gn M$ extbf{o}$me$ extbf{n}$tum}$). It is more memory-efficient than Adam as it only keeps track of the momentum. Different from adaptive optimizers, its update has the same magnitude for each parameter calculated through the sign operation. We compare Lion with widely used optimizers, such as Adam and Adafactor, for training a variety of models on different tasks. On image classification, Lion boosts the accuracy of ViT by up to 2% on ImageNet and saves up to 5x the pre-training compute on JFT. On vision-language contrastive learning, we achieve 88.3% $ extit{zero-shot}$ and 91.1% $ extit{fine-tuning}$ accuracy on ImageNet, surpassing the previous best results by 2% and 0.1%, respectively. On diffusion models, Lion outperforms Adam by achieving a better FID score and reducing the training compute by up to 2.3x. For autoregressive, masked language modeling, and fine-tuning, Lion exhibits a similar or better performance compared to Adam. Our analysis of Lion reveals that its performance gain grows with the training batch size. It also requires a smaller learning rate than Adam due to the larger norm of the update produced by the sign function. Additionally, we examine the limitations of Lion and identify scenarios where its improvements are small or not statistically significant. Lion is also successfully deployed in production systems such as Google search ads CTR model.

研究动机与目标

• 促使发现新的优化器，超越手工设计，旨在提升深度神经网络的训练效率和泛化能力。
• 将优化器发现表述为可执行的程序搜索问题，以探索新颖的算法设计。
• 开发在无限、稀疏搜索空间中导航并筛选在大规模任务上具有泛化能力的解决方案的技术。
• 在视觉、语言和扩散模型上展示所发现优化器的实际性能。

提出的方法

• 将优化器表示为对权重、梯度和辅助状态执行的命令式程序，具有固定的训练签名。
• 使用来自45个数学函数集合的一组语句的巨大搜索空间来编码候选算法。
• 应用带有热启动（来自 AdamW）和重新启动的进化搜索来探索该空间。
• 采用抽象执行和缓存来修剪无效/语义等价的程序并加速评估。
• 使用漏斗筛选和元验证，在逐步增大的代理任务上选择对目标任务具有泛化能力的算法。
• 通过消除冗余并使更新对齐为符号化动量更新，将发现的程序简化为 Lion。

实验结果

研究问题

• RQ1 是否可以通过程序搜索在无限、稀疏的空间中有效发现优化算法？
• RQ2 自动发现的优化器是否能够从代理任务泛化到大规模、前沿的训练设置？
• RQ3 在架构和任务跨越的情况下，所发现的优化器具有什么特性和实际极限？
• RQ4 在现实世界的训练中，与 AdamW 和 Adafactor 相比，简单的带符号更新与动量有何差异？

主要发现

• Lion 是一个简单、内存高效的优化器，能够跟踪动量并使用符号更新，在某些设置下使 ImageNet 零样本准确率比先前基线高出最多 2 个百分点，微调准确率高出 0.1 个百分点。
• Lion 将 JFT 的预训练计算量降低多达 5 倍，且可将扩散模型训练加速约 2.3 倍。
• 在 ImageNet 的多种模型上，Lion 的表现优于 AdamW，在容量更大模型和使用更大批量时获得更大提升。
• 在视觉-语言学习（LiT/BASIC 设置）中，Lion 将零-shot ImageNet 准确率相比 AdamW 提高约 1.0–1.7 个百分点，并获得更好的检索结果。
• Lion 的性能优势随批量大小增加而增强，并且需要更小的学习率和更大的权重衰减来维持有效的正则化。
• 搜索过程在其他优化器中也能识别出 Lion，并且在代理任务搜索后期的元过拟合倾向于更好地泛化到目标任务。