论文算法发现

Symbolic Discovery of Optimization Algorithms (Lion)

符号化程序搜索发现优化器 Lion

Xiangning Chen, Chen Liang, Da Huang, Esteban Real, et al. (Google) · Google · NeurIPS 2023 · 2023-02 · 被引 698

一句话Google 把「发现优化器」形式化为程序空间上的进化搜索(约 3K TPU V2 天),从 AdamW 出发演化出只需一个 momentum 状态的 sign-based 优化器 Lion:ViT 在 ImageNet 上最高 +2%,JFT 预训练省最多 5x 计算,且被 optax/timm/Keras/bitsandbytes 等主流库实际采纳——机器发现的训练算法进入生产的少数真实案例。

这是什么

自动发现优化算法(learning to optimize)是个老方向,但此前的尝试要么用神经网络参数化优化器(难迁移、难分析),要么在固定模板里搜超参组合(搜不出新结构)。这篇工作把优化器表示成一段命令式程序:一个 train(w, g, lr) 函数,由一串对 n 维数组的赋值语句组成,词表是 45 个 NumPy 风格的数学函数,语句数和局部变量数不设限。程序空间无限大且高性能程序极其稀疏——作者随机采样 2M 个程序,最好的一个仍明显不如 AdamW。

在这个空间上跑 regularized evolution(锦标赛选择 + 三种变异:插入/删除/修改语句),从 AdamW 热启动,配合 restart、abstract execution 剪枝和低成本 proxy task,最终搜出的程序经人工化简后就是 Lion(EvoLved Sign Momentum):只维护一个 momentum,更新方向取 sign,幅度对所有参数一致。比 Adam 少一个二阶矩状态,内存减半,实现只有几行。

发表于 NeurIPS 2023,S2 引用 698。它的意义有两层:作为优化器,Lion 是 AdamW 之后少数被社区大规模实际使用的新优化器;作为方法论,它演示了「proxy task 搜索 + meta-validation 漏斗筛选 + 化简」这套让机器发现的算法跨 10^4 倍规模 gap 泛化的流程。

JFT 预训练的 ViT 在 ImageNet 上的精度 vs 预训练计算量:同等精度下 Lion(蓝)比 AdamW(橙)省约 3x 计算,同等计算下精度高约 0.7-1%。这是论文主打的 compute saving 结果。

机制与做法

搜索空间与效率技巧

程序保持和 AdamW 相同的输入输出签名(权重、梯度、lr,外加若干初始化为零的历史状态变量),保证搜出的算法内存不超过 AdamW;但与之前的 optimizer search 不同,历史变量怎么更新也是可搜的。变异只改一条语句,所以允许冗余语句存在——它们是后续大改动的中间脚手架,搜索后期约 69.8% 的语句是冗余的。

关键效率技巧是 abstract execution:真正跑程序之前先做一遍抽象执行,推断类型/形状过滤非法程序、给「输入到输出的计算方式」算 hash 用于去重缓存、标记冗余语句。缓存命中率到后期约 89%,相当于把搜索成本降了约 10x。每个 proxy task 在单块 TPU V2 上 20 分钟内跑完,一次搜索用 100 块 TPU V2 跑约 72 小时、生成 20-30 万个程序(实际执行约 2-3 万个),全流程(5 次重启 + 1 次从最优程序续搜)约 3K TPU V2 天。

ImageNet 精度随 batch size 的变化:bs=64 时 Lion 与 AdamW 打平,bs 越大 Lion 优势越明显(32K 时 77.9 vs 75.4)。这张图同时展示了 Lion 的适用边界——小 batch 场景没有收益。

跨越 proxy 到目标任务的泛化 gap

proxy task 和目标任务之间有超过 10^4 倍的规模差,直接拿搜索 fitness 最高的程序去大任务上多半失败——存在 meta-overfitting:搜索 fitness 还在涨,但在约 500x 大的 meta-validation 任务上指标开始掉。50 次搜索的统计显示一半的 run 很早就 meta-overfit,而 overfit 得越晚的 run 泛化越好。

对策是漏斗式筛选:候选程序先在 10x 大的任务上过筛,赢过 baseline 才上 100x 大的任务,逐级过滤。选出的程序再做化简:先自动删冗余语句,再删掉去掉后几乎无影响的语句,最后人工重排、重命名、做数学等价变形。原始搜出的程序里 cosh/arcsin/clip 等都是可删的噪声,三条红色语句合起来等价于一个 sign——Lion 就是这么从一坨演化产物里提炼出来的。

Lion 本体:sign momentum 为什么行

Lion 的更新:u = sign(interp(g, m, β1)),然后 m = interp(g, m, β2),默认 β1=0.9、β2=0.99。和 signSGD 的区别在于 momentum 的追踪(β2=0.99,记住约 10x 长的梯度历史)和使用(更新时以 β1=0.9 更偏重当前梯度)被解耦。sign 让每个维度的更新幅度一致,自带类似梯度噪声的正则效应——ViT-B/16 用 Lion 训练误差比 AdamW 高但验证精度高 2%,且收敛到更平坦的区域。

实用后果:更新的 norm 比 AdamW 大,所以 lr 要缩小 3-10x、weight decay 相应放大以保持有效正则强度;没有 ε 和二阶矩,超参更少;只存一个 bfloat16 momentum,训 ViT-B/16(bs 4096)AdamW 需要至少 16 块 TPU V4 而 Lion 只要 8 块,单步还快 2-15%。论文自己的 Limitations 也写得很老实:batch size < 64 时优势消失;在 Imagen base、大规模内部数据的自回归 LM perplexity、C4 masked LM 上与 AdamW 相当——数据又大又干净时优化器差距会收窄。

关键结果

ViT 系列在 ImageNet 从头训练最高 +2%(ViT-B/16 无增广:75.48 → 77.44);JFT 预训练达到同等精度省最多 5x 计算(图上 ViT-L/16 约 3x)。
视觉-语言对比学习:zero-shot ImageNet 88.3%、fine-tune 91.1%,分别刷新当时纪录 2% 和 0.1%。
扩散模型上 FID 更好且省最多 2.3x 训练计算;自回归/masked LM 与 Adam 持平或略好。
Lion 相对 AdamW 的优势随 batch size 增大而扩大(bs 32K 时 ImageNet 77.9 vs 75.4),但 bs=64 时两者打平——适合大规模数据并行、不适合小 batch 场景。
搜索方法本身的对照:同等 4x 计算下,进化搜索的 fitness 显著超过 AdamW 超参调优和随机搜索两条 baseline;abstract execution 的缓存带来约 10x 搜索加速。
内存:只存 momentum(可 bfloat16),状态量是 AdamW 的一半;单步速度快 2-15%。
论文明确列出无效场景:强数据增广下增益缩水、ResNet 上与 SGD/AdamW 差距很小、大规模高质量语料的 LM 预训练上与 AdamW 相当。

实证核查

扎实视觉任务上的增益被社区广泛复现,Lion 被 optax、timm、Keras、bitsandbytes 等主流库正式收录,是机器发现算法进入生产的真实案例;但在 LLM 预训练上,后续独立基准显示精调过的 AdamW 不输 Lion——这点论文 Limitations 里其实自己承认了。另外开源的只有优化器本身,发现流程(搜索代码)没有开源。

Lion 是简单有效的优化器,比 Adam 省内存,在视觉任务上有实打实的提升。

被生态大规模采纳:官方实现在 google/automl/lion(含 optax/PyTorch/TF1/TF2 四个版本,gh api 确认);第三方 lucidrains/lion-pytorch 2.2k stars 且 2026 年仍在维护;Keras 内置 keras.optimizers.Lion,optax、timm、bitsandbytes(8-bit Lion)均收录。一个新优化器被这么多框架正式集成,在 AdamW 之后属于少数。

在自回归、masked LM 和 fine-tuning 上,Lion 与 Adam 相当或更好。

这条在后续独立基准里打了折:Semenov et al. 的 LLM 预训练优化器基准(openreview.net/forum?id=fL9qDVnMJF,被引 58)显示 Lion 需要配大 weight decay 才能与 AdamW 竞争;arXiv 2507.08472(预算受限 LLM 预训练对比)结论是 AdamW 下游精度稳定优于 Lion 和 Sophia;Stanford 2025 年的优化器横评也认为 Lion 这类 scalar 优化器相对精调 AdamW 加速有限,真正的加速来自 Muon 等矩阵型优化器。公平地说,论文 Limitations 自己写了大规模高质量语料上与 AdamW 相当,声称本身不算夸大。

方法把算法发现形式化为程序搜索,搜索成本约 3K TPU V2 天,并已部署到 Google 搜索广告 CTR 模型等生产系统。

google/automl/lion 目录里只有 4 个优化器实现文件 + README,进化搜索、abstract execution、漏斗筛选的代码均未开源(gh api repos/google/automl/contents/lion 确认),所以「发现流程」本身无法复现,只能复现产物;生产部署(ads CTR)是内部系统,无法外部核实。仓库 google/automl 已于 2025-03 archive,但这不影响 Lion 的使用——各框架都有自己的实现。

Lion 更新方向取 sign 是搜出来的,不是拍脑袋设计。

后续理论工作给了背书:Chen et al.《Lion Secretly Solves a Constrained Optimization Problem》(arXiv 2310.05898)证明 Lion 隐式在解一个带 ℓ∞ 约束的优化问题,说明这个演化产物在数学上是有结构的,不是过拟合 proxy task 的偶然产物。这也是「机器发现 → 人类事后给出理论解释」的一个好例子。

与我们方向的关系

对 ai4ai/算法发现方向,这篇是 pre-LLM 时代符号搜索路线的标杆:不用 LLM 生成代码,纯靠进化 + 变异也能在无限程序空间里找到可部署的算法。它解决泛化 gap 的三件套——低成本 proxy、meta-validation 漏斗筛选、自动+人工化简——对任何「在小任务上搜、指望大任务上用」的算法发现工作(包括 LLM 驱动的 FunSearch/AlphaEvolve 一路)都直接可借鉴,meta-overfitting 的量化分析(50 次 run 的统计)尤其值得参考。

另一个可借鉴的点是 abstract execution:类型/形状检查 + 语义 hash 去重 + 冗余语句标记,成本几乎为零却带来约 10x 搜索加速。用 LLM 做算法搜索时同样需要这类廉价的程序等价性判定和去重机制,否则采样预算会被重复候选吃掉。

阅读笔记

搜索代码未开源是复现「方法」的硬伤,读方法部分时只能信论文的消融。用 Lion 时注意:lr 要比 AdamW 缩 3-10x、weight decay 放大,小 batch(<64)别用;LLM 预训练场景对它期望别太高。google/automl 仓库已 archive,PyTorch 用户直接用 lucidrains/lion-pytorch 或 bitsandbytes 的 8-bit 版即可。