ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

Lion 优化器完全指南:Google AutoML 中经符号程序搜索发现的轻量级优化算法及其多框架实现

2026/10/4 13:42:45 拓冰建站 浏览量
Lion 优化器完全指南:Google AutoML 中经符号程序搜索发现的轻量级优化算法及其多框架实现 人工智能深度学习机器学习计算机视觉AutoML【免费下载链接】automlGoogle Brain AutoML项目地址https://gitcode.com/gh_mirrors/au/automl点击查看免费下载本文以 Google Brain AutoML 仓库 lion 目录中的官方 README 与 lion_pytorch.py、lion_optax.py、lion_tf1.py、lion_tf2.py 等实现源码为依托系统讲解 Lion 优化器的更新规则、内存与运行效率优势、跨任务实证表现以及最关键的学习率 / 权重衰减 / 批大小调参方法论。读完本文你将理解 Lion 与 AdamW 的本质差异掌握把现有训练脚本迁移到 Lion 时学习率缩小 3-10 倍、权重衰减放大 3-10 倍的实操规则并能直接使用仓库中的 PyTorch、JAX (Optax)、TensorFlow 1/2 四套官方实现。背景从符号程序搜索中发现而非手工设计的优化器Lion 优化器来自论文Symbolic Discovery of Optimization AlgorithmsarXiv:2302.06675由 Xiangning Chen 等 Google 研究者发表于 2023 年。与历史上靠理论推导和手工试错设计的 Adam、AdamW 不同Lion 的更新规则是通过符号程序搜索symbolic program search自动发现的搜索算法在给定搜索空间内组合出候选更新规则在真实训练任务上评估并迭代进化最终收敛到 Lion 这一简洁而高效的规则。这一搜索过程的底层基础设施在本仓库中以代号 Hero 开源即 hero 目录。从 hero/core.py 的模块文档可以看出它提供了支撑搜索空间的四类核心能力从字符串解析程序、在给定命名空间中执行程序支持具体执行与抽象执行、对程序进行变异、将程序序列化为字符串用于可视化。核心类包括Program可解析、执行与变异的程序、Namespace程序执行与变异所处的名字-值映射环境、Expression构成程序的各类表达式如函数调用、字面量、条件、循环、函数定义。Lion 的默认超参数$\beta_10.9$、$\beta_20.99$正是通过程序搜索过程发现的结果这一点在 README 的超参数一节有明确说明。本仓库同时开源了 Lion 在四种深度学习框架下的实现lion_pytorch.pyPyTorchLion优化器类lion_optax.py基于 Optax 的 JAX 实现按optax.chain组合scale_by_lion、add_decayed_weights与学习率缩放lion_tf2.pyTF2 / Keras 风格实现继承tf.keras.optimizers.legacy.Optimizerlion_tf1.pyTF1 实现继承tf.compat.v1.train.Optimizer。此外README 提到 Lion 已被集成进多个主流生态包括 Praxis、Optax、Keras、Timm、T5X 以及社区流行的 lucidrains PyTorch 实现方便在不同框架中直接使用。Lion 的更新规则与算法本质更新规则Lion 的每次参数更新只需维护一个动量momentum状态核心规则如下以 lion_pytorch.py 的step实现为准# 1. 解耦权重衰减与学习率相乘AdamW 风格 p.data.mul_(1 - group[lr] * group[weight_decay]) # 2. 组合动量与当前梯度update beta1 * m (1 - beta1) * grad update exp_avg * beta1 grad * (1 - beta1) # 3. 对 update 取符号sign按学习率施加到参数上 p.add_(update.sign_(), alpha-group[lr]) # 4. 更新动量m beta2 * m (1 - beta2) * grad exp_avg.mul_(beta2).add_(grad, alpha1 - beta2)即每个时间步的更新量是sign(β1·m (1−β1)·g)这一逐元素二值的 ±1 向量参数以学习率为幅度沿该方向移动。注意到步骤 1 的权重衰减与学习率相乘因此有效权重衰减为 $lr \times \lambda$——这正是 README 调参指南中$\lambda$ 需要随 $lr$ 等比例反向调节这一结论的根源。JAX/Optax 实现用optax.chain把三个变换串联起来scale_by_lion负责动量与符号运算optax.add_decayed_weights(weight_decay, mask)负责可掩码的权重衰减_scale_by_learning_rate负责学习率缩放其中scale_by_lionlion_optax.py的更新函数实现为updates sign((1 - b1) * g b1 * m)并支持mu_dtype指定动量累加器的 dtype、mask指定跳过权重衰减的参数子集。TF2 实现lion_tf2.py则对稠密与稀疏更新分别用tf.function(jit_compileTrue)做了 XLA 编译且提供了稀疏梯度_resource_apply_sparse的专门路径。简单、省内存、运行更快README 强调 Lion 相比 AdamW 及各类自适应优化器的三个直接收益内存占用减半AdamW 这类自适应优化器需要同时保存一阶矩first moment与二阶矩second moment两份状态而 Lion 只需要一份动量。对于大模型或大批量训练这是可观的显存 / 内存节省。硬件门槛降低README 给出的实例是用 ViT-B/16输入 224×224、batch size 4,096训练时AdamW 至少需要 16 个 TPU V4 芯片而 Lion 只需 8 个。运行更快由于运算简单README 报告在实验中 Lion 通常比 AdamW 与 Adafactor 快 2%–15%取决于任务、代码库与硬件。Lion 优化算法示意图跨架构、跨任务的实证表现README 汇总了论文中 Lion 在四类任务域上的结果均为论文/README 报告数据具体超参数见论文图像分类在 ImageNet 上从头训练、或在 ImageNet-21K 上预训练的各种架构上Lion 均优于 AdamW。在 JFT-300M 上Lion 最多可节省 5 倍预训练成本。在更高分辨率微调与 Polyak 平均之后得到的 ViT-L/16 追平了此前 AdamW 训练的 ViT-H/14 的结果而模型体积小 2 倍。视觉-语言对比学习在 LiT 上Lion 在多个零样本图像分类与图文检索基准上优于 AdamW。在 BASIC-L 上Lion 达到 88.3% 零样本与 91.1% 微调 ImageNet 准确率分别超过此前 SOTA 2% 与 0.1%。扩散模型在 ImageNet 上 64×64、128×128、256×256 三个分辨率的图像生成实验中Lion 的 FID 优于 AdamW且最多节省 2.3 倍训练计算量。语言建模在 Wiki-40B 与 PG-19 上的语言建模中Lion 在验证困惑度perplexity指标上最多节省 2 倍计算量且 Transformer 越大收益越明显。相比 AdafactorLion 训练出的 LM 具备更好的平均上下文学习in-context learning能力。在 GLUE 上微调 T5 时Lion 优于 AdamW。超参数调优指南核心实操章节README 中Instructions for hyperparameter tuning and batch size choices一节是迁移到 Lion 时最需要遵循的实操依据本节逐条展开。更少的超参数不需要 εLion 相比 AdamW 与 Adafactor 更简单、超参数更少它不包含 $\epsilon$也没有 Adafactor 那类与矩阵分解相关的超参数。论文为保证公平对比对 AdamWAdafactor与 Lion 都在对数刻度上只调峰学习率 $lr$ 与解耦权重衰减 $\lambda$ 两个量。默认值对照如下优化器β1β2ε其他AdamW 默认0.90.9991e-8—Lion 默认0.90.99无—论文第 4.4 节调参结果AdamW0.90.991e-6—论文第 4.4 节调参结果Lion0.950.98无—仓库四个实现文件的默认超参数与此一致PyTorch 版为lr1e-4, betas(0.9, 0.99), weight_decay0.0Optax 版为b10.9, b20.99, weight_decay0.0TF1/TF2 版为learning_rate0.0001, beta1/beta_10.9, beta2/beta_20.99, wd0.0。调参经验README 原文要点调小 β2 增强稳定性经验表明减小 $\beta_2$ 会让模型对历史信息的记忆更短从而提升训练稳定性。AdamW 的 ε 建议调大论文实验中 AdamW 的 $\epsilon$ 设为 1e-6 而非默认 1e-8 时稳定性更好这与 RoBERTa 中的观测类似。学习率通常比 AdamW 小 3–10 倍由于 Lion 的更新量是 sign 运算产生的逐元素二值 ±1 向量其范数天然大于其他优化器生成的更新量因此Lion 合适的学习率通常是 AdamW 的 1/3 到 1/10。并且注意学习率的初始值、峰值、结束值需要按同一比例同时调整不能只改峰值。除学习率外训练配置如学习率调度、梯度裁剪与更新裁剪等不需要改动。权重衰减通常比 AdamW 大 3–10 倍由于有效权重衰减是 $lr \times \lambda$而 $lr$ 缩小了 3–10 倍为了维持相近的衰减强度Lion 的 $\lambda$ 应比 AdamW 大 3–10 倍。三个可直接复用的配置实例README 给出了三组经过验证的配置对照可直接作为迁移起点任务/模型Lion lrLion λ对照优化器 lr对照 λViT-B/16 ImageNet强增强1e-410.0AdamW 1e-31.0扩散模型3e-50.1AdamW 3e-40.017.5B 语言建模1e-40.01Adafactor 1e-30.001论文中所有任务的完整超参数详见原文。对超参数更稳健除了峰值性能超参数敏感度与调参难度同样决定优化器能否被实际采用。README 报告在 ImageNet 上从头训练 ViT-B/16同时扰动 $lr$ 与 $\lambda$ 时Lion 的热力图显示其对不同超参数组合比 AdamW更稳健即更不容易因调参不当而大幅掉点。批大小消融与超参数热力图批大小选择Lion 偏好大 batch小 batch 也稳健一个常见的顾虑是sign 操作会引入噪声Lion 是否必须用很大的 batch size 才能确定可靠的更新方向README 用一组实验回答了这个问题——在 ImageNet 上从头训练 ViT-B/16固定总训练轮数为 300并采用 RandAug 与 Mixup 增强变化 batch size最优 batch sizeAdamW 为 256Lion 为 4,096说明 Lion 确实更偏好大 batch。小 batch 依然稳健即使 batch size 小到 64Lion 的性能依然保持稳健。超大 batch 收益显著当 batch size 放大到 32K训练步数仅 11K 步时Lion 相比 AdamW 获得显著的 2.5 个百分点准确率增益77.9% vs. 75.4%展示了它在超大批量训练场景下的有效性。仓库实现快速上手PyTorch从 lion_pytorch.py 导入即可使用接口与标准 PyTorch 优化器一致from lion_pytorch import Lion optimizer Lion(model.parameters(), lr1e-4, betas(0.9, 0.99), weight_decay10.0)JAX / Optaxlion_optax.py 暴露了与 Optax 生态一致的lion(...)变换可与optax.chain、optax.adam等自由组合from lion_optax import lion tx lion(learning_rate1e-4, b10.9, b20.99, weight_decay10.0)其中weight_decay会与学习率相乘与 PyTorch 语义一致且可通过mask对部分参数跳过权重衰减mu_dtype可指定一阶累加器的数值精度。TensorFlow 2 / Keraslion_tf2.py 提供 Keras 风格优化器可直接传入model.compilefrom lion_tf2 import Lion optimizer Lion(learning_rate0.0001, beta_10.9, beta_20.99, wd0.0)TensorFlow 1lion_tf1.py 提供tf.compat.v1.train.Optimizer子类适合 TF1 训练图或 Estimator 场景参数名为beta1、beta2、wd。引用如果本文内容对你的工作有帮助可按照 README 中的说明引用该工作misc{chen2023symbolic, title{Symbolic Discovery of Optimization Algorithms}, author{Xiangning Chen and Chen Liang and Da Huang and Esteban Real and Kaiyuan Wang and Yao Liu and Hieu Pham and Xuanyi Dong and Thang Luong and Cho-Jui Hsieh and Yifeng Lu and Quoc V. Le}, year{2023}, eprint{2302.06675}, archivePrefix{arXiv}, primaryClass{cs.LG} }进一步深入可阅读本仓库中 lion/README.md 原始文档、四个框架的完整实现源码lion_pytorch.py、lion_optax.py、lion_tf1.py、lion_tf2.py以及符号搜索基础设施 hero/core.py含Program、Namespace、Expression核心类与其配套的 hero/fn_lib.py、hero/core_test.py 测试。赞分享人工智能深度学习机器学习计算机视觉AutoML【免费下载链接】automlGoogle Brain AutoML项目地址https://gitcode.com/gh_mirrors/au/automl点击查看免费下载相关推荐革命性优化器Lion符号程序搜索发现的新一代深度学习优化算法革命性优化器Lion符号程序搜索发现的新一代深度学习优化算法 还在为AdamW优化器的内存占用和训练速度而烦恼吗Google Brain团队通过符号程序搜索人工智能深度学习机器学习计算机视觉AutoML深度解析Lion优化器符号程序搜索如何发现更高效的深度学习优化算法深度解析Lion优化器符号程序搜索如何发现更高效的深度学习优化算法 Lion优化器是Google Brain通过符号程序搜索技术发现的全新深度学习优化器相比人工智能深度学习机器学习计算机视觉AutoMLLion优化器革命AdamW-Lion混合算法实战指南Lion优化器革命AdamW Lion混合算法实战指南 还在为深度学习训练中的优化器选择而烦恼吗传统的AdamW虽然稳定但内存占用大而新兴的Lion优化器人工智能深度学习机器学习计算机视觉AutoML上一篇HS2-HF Patch终极指南5分钟解决HoneySelect2汉化与MOD整合难题下一篇HoneySelect2终极汉化与MOD整合完整指南HS2-HF Patch快速安装教程创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考