Adan: 自适应 Nesterov 动量算法，用于更快地优化深度模型

这是 Adan 的官方 PyTorch 实现。论文可以在这里查看。如果您发现我们的 adan 对您的项目有帮助或启发，请引用这篇论文并为这个仓库点星。谢谢！

@article{xie2024adan,
  title={Adan: 自适应 Nesterov 动量算法，用于更快地优化深度模型},
  author={Xie, Xingyu and Zhou, Pan and Li, Huan and Lin, Zhouchen and Yan, Shuicheng},
  journal={IEEE 模式分析与机器智能汇刊},
  year={2024},
  publisher={IEEE}
}

支持的项目

• Adan 在 NVIDIA 的框架 NeMo 中得到支持。
• Adan 是高保真文本到3D生成项目的默认优化器。更多详情请参见 Consistent3D。
• Adan 是掩码扩散变换器 V2 的默认优化器。更多详情请参见 MDT V2。
• Adan 在 Meta AI 的项目 D-Adaptation 中得到支持。
• Adan 将在 百度飞桨 的项目 Paddle 中得到支持。
• Adan 在 Huggingface 的 Timm 中得到支持。
• Adan 是文本到3D DreamFusion 项目的默认优化器。更多结果请参见这里。
• Adan 在 OpenMMLab 项目的 MMClassification 中得到支持。使用 Adan 训练 ViT-B 的日志和示例在这里。
• TF 的实现（第三方）请参考 DenisVorotyntsev/Adan。
• JAX 版本（第三方）在 Deepmind/optax 中实现并得到支持。

新闻

• :fire::fire::fire: 在大型语言模型（如 MoE 和 GPT2）上的结果已发布。
• 内存占用更少的 FusedAdan 已发布。

安装

python3 -m pip install git+https://github.com/sail-sg/Adan.git

默认安装 FusedAdan。如果您想使用原始的 Adan，请通过以下方式安装：

git clone https://github.com/sail-sg/Adan.git
cd Adan
python3 setup.py install --unfused

使用方法

为了方便您使用 Adan，我们简要提供了一些直观的说明，然后提供一些通用的实验技巧，最后为论文中的每个实验提供更多细节（如具体命令和超参数）。

1) 使用 Adan 的两个步骤

步骤 1. 通过在配置中添加以下超参数来添加 Adan 依赖的超参数：

parser.add_argument('--max-grad-norm', type=float, default=0.0, help='如果 l2 范数大于此超参数，则裁剪梯度（默认：0.0，不进行梯度裁剪）')
parser.add_argument('--weight-decay', type=float, default=0.02,  help='权重衰减，类似于 AdamW 中使用的值（默认：0.02）')
parser.add_argument('--opt-eps', default=None, type=float, metavar='EPSILON', help='优化器 epsilon，用于避免二阶矩为零的不良情况（默认：None，在 adan 中使用优化器默认值 1e-8）')
parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA', help='Adan 中的优化器 beta 值（默认：None，在 Adan 中使用优化器默认值 [0.98, 0.92, 0.99]）')
parser.add_argument('--no-prox', action='store_true', default=False, help='是否像 AdamW 那样执行权重衰减（默认为 False）')

opt-betas：为保持与我们的使用习惯一致，论文中的 $\beta$ 实际上是代码中的 $(1-\beta)$。

foreach (bool)：如果为 True，Adan 将使用 torch._foreach 实现。它更快但使用略多的内存。

no-prox：它决定了带权重衰减的参数的更新规则。默认情况下，Adan 按照论文中算法 1 所示的方式更新参数：

$$\boldsymbol{\theta}_{k+1} = ( 1+\lambda \eta)^{-1} \left[\boldsymbol{\theta}_k - \boldsymbol{\eta}_k \circ (\mathbf{m}_k+(1-{\color{blue}\beta_2})\mathbf{v}_k)\right]$$

但也可以像 Adamw 那样更新参数：

$$\boldsymbol{\theta}_{k+1} = ( 1-\lambda \eta)\boldsymbol{\theta}_k - \boldsymbol{\eta}_k \circ (\mathbf{m}_k+(1-{\color{blue}\beta_2})\mathbf{v}_k).$$

步骤 2. 创建 Adan 优化器如下。在这一步中，我们可以直接使用以下命令替换原始优化器：

from adan import Adan
optimizer = Adan(param, lr=args.lr, weight_decay=args.weight_decay, betas=args.opt_betas, eps = args.opt_eps, max_grad_norm=args.max_grad_norm, no_prox=args.no_prox)

2) 实验技巧

• 为了简化Adan，除论文中的表12外的所有实验中，我们没有使用Adan的重启策略。但表12显示重启策略可以进一步略微提高Adan的性能。
• Adan通常允许使用更大的峰值学习率，这往往会导致其他优化器(如Adam和AdamW)失败。例如，除了MAE预训练和LSTM外的所有实验中，Adan使用的学习率是Adam/AdamW的5-10倍。
• Adan对beta1、beta2和beta3相对鲁棒，尤其是对beta2。如果你想要更好的性能，可以先调整beta3，然后再调整beta1。
• 在单节点上，Adan比Adam/AdamW略微增加了GPU内存消耗。然而，这个问题可以通过使用ZeroRedundancyOptimizer来解决，它在分布式数据并行进程中共享优化器状态以减少每个进程的内存占用。具体来说，当在两个以上的GPU上使用ZeroRedundancyOptimizer时，Adan和Adam消耗的内存几乎相同。

3) 更多详细步骤和结果

请参考以下链接获取详细步骤。在这些详细步骤中，我们甚至包含了docker镜像以便复现。

• ViTs、ResNets和ConvNext的说明
• MAE的说明
• BERT的说明
• Transformer-XL的说明
• GPT2的说明
• text-to-3D DreamFusion的结果

各种任务的结果

大型语言模型的结果

专家混合系统（MoE）

为了研究Adan优化器对LLM的有效性，我们使用MoE模型进行了预训练实验。这些实验使用了RedPajama-v2数据集，有三种配置，每种包含8个专家：8x0.1B（共0.5B可训练参数）、8x0.3B（2B可训练参数）和8x0.6B（4B可训练参数）。这些模型分别用10B、30B、100B和300B个token的采样数据进行训练。

GPT2-345m

我们提供了在BigCode数据集上预训练的GPT2-345m的配置和日志，并通过零样本学习在HumanEval数据集上进行评估。HumanEval用于衡量从文档字符串合成程序的功能正确性。它包含164个原创编程问题，评估语言理解、算法和简单数学，其中一些可与简单的软件面试问题相媲美。我们在评估时设置Temperature = 0.8。

Adan仅用一半的成本就获得了可比的结果。

视觉任务的结果

为了方便您使用Adan，我们提供了ImageNet-1k实验的配置和日志文件。

自然语言处理任务结果

BERT-base

我们提供了在Bookcorpus和Wikipedia数据集上预训练并在GLUE任务上微调的BERT-base模型的配置和日志文件。请注意，我们在./NLP/BERT文件夹中提供了BERT-base的配置、日志文件和详细说明。

关于GLUE任务的微调，请参阅相应配置文件中的总批量大小。

Transformer-XL-base

我们提供了在WikiText-103数据集上训练的Transformer-XL-base的配置和日志。此实验的总批量大小为60*4。

大型语言模型的结果

GPT2-345m

我们提供了在来自BigCode的数据集上预训练的GPT2-345m的配置和日志，并通过零样本学习在HumanEval数据集上进行评估。HumanEval用于衡量从文档字符串合成程序的功能正确性。它包含164个原始编程问题，评估语言理解、算法和简单数学，其中一些可与简单的软件面试问题相媲美。我们在评估时设置温度 = 0.8。

Adan仅用一半的成本就获得了可比的结果。

扩散模型的结果

我们展示了由DreamFusion项目支持的文本到3D任务的结果。更多可视化结果可以在这里找到。 以下是使用Adam和Adan从文本提示"悉尼歌剧院，鸟瞰图"生成的示例：

内存和效率

以下是优化器峰值内存和运行时间的简要比较。持续时间是200次optimizer.step()的总时间。我们进一步详细比较了GPT-2上的Adam和FusedAdan。更多结果请参见这里。