Symbolic Discovery of Optimization Algorithms
Xiangning Chen1 2 §§\mathsection Chen Liang1 §§\mathsection Da Huang1 Esteban Real1 Kaiyuan Wang1 Yao Liu1 Hieu Pham1 Xuanyi Dong1 Thang Luong1 Cho-Jui Hsieh2 Yifeng Lu1 Quoc V. Le1 §§\mathsectionEqual、Core Contribution 1Google 2UCLA Work done as a student researcher at Google Brain.Work done while at Google.
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, Lion (EvoLved Sign Momentum). 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% zero-shot and 91.1% 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 a
中文速览
深度神经网络的训练长期依赖Adam、AdamW等手工设计的优化器，而自动发现新优化器的方法要么泛化性差、要么搜索空间受限，始终未能超越这些经典算法。研究团队将优化器发现问题转化为程序搜索（program search），用进化算法在一个由数学运算语句构成的无限程序空间中探索，并通过抽象执行剪枝、漏斗式筛选和程序化简等策略，解决搜索空间稀疏、代理任务与真实任务存在泛化鸿沟这两大核心难题。最终发现了一个极简高效的优化算法——Lion（EvoLved Sign Momentum）：它只需维护一阶动量，用符号函数（sign）统一每个参数的更新幅度，内存占用比Adam更低。在图像分类、视觉-语言对比学习、扩散模型、语言模型等多项任务上，Lion全面持平或超越AdamW和Adafactor，ImageNet零样本准确率达88.3%，JFT预训练算力最多节省5倍，已在谷歌搜索广告点击率预估系统中落地部署。这项工作证明了通过程序搜索自动发现的优化算法完全可以在实际生产规模的任务上击败精心调校的人工设计算法，为自动机器学习研究开辟了新方向。
原文 arXiv:2302.06675;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2302.06675v4