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

这是 Adan 的官方 PyTorch 实现。请在此处查看论文 here。如果你觉得我们的 Adan 对你的项目有帮助或启发，请引用这篇论文并为本仓库点亮星星。谢谢！

@article{xie2024adan,
  title={Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models},
  author={Xie, Xingyu and Zhou, Pan and Li, Huan and Lin, Zhouchen and Yan, Shuicheng},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
  year={2024},
  publisher={IEEE}
}

支持的项目

• Adan 的黎曼版本（Radan），包含在 Manify 中 —— 一个非欧几里得机器学习库。
• Adan 在 NVIDIA 提供的框架 NeMo 中得到支持。
• Adan 是高保真文本到3D生成项目的默认优化器。更多详情请参见 Consistent3D。
• Adan 是 Masked Diffusion Transformer V2 的默认优化器。更多详情请参见 MDT V2。
• Adan 在 Meta AI 提供的项目 D-Adaptation 中得到支持。
• Adan 将在 百度飞桨 提供的项目 Paddle 中得到支持。
• Adan 在 Huggingface 提供的 Timm 中得到支持。
• Adan 是文本到3D项目 DreamFusion Project 的默认优化器。更多结果请参见 这里。
• 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 中的优化器 betas (默认值: 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 的 GPU 内存消耗略高于 Adam/AdamW。然而，这个问题可以通过使用 ZeroRedundancyOptimizer 来解决，该工具会在分布式数据并行进程中共享优化器状态以减少每进程的内存占用。具体来说，当在两个以上的 GPU 上使用 ZeroRedundancyOptimizer 时，Adan 和 Adam 几乎消耗相同的内存量。

3) 更多详细步骤和结果

请参考以下链接获取详细步骤。在这些详细步骤中，我们甚至包括了用于可重复性的 docker 镜像。

• ViTs、ResNets 和 ConvNext 的 说明。
• MAE 的 说明。
• BERT 的 说明。
• Transformer-XL 的 说明。
• GPT2 的 说明。
• 文本到3D DreamFusion 的 结果。

各项任务的结果

大型语言模型的结果

专家混合模型 (MoE)

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

模型规模	8x0.1B	8x0.1B	8x0.1B	8x0.3B	8x0.3B	8x0.3B	8x0.6B
Token 数量	10B	30B	100B	30B	100B	300B	300B
AdamW	2.722	2.550	2.427	2.362	2.218	2.070	2.023
Adan	2.697	2.513	2.404	2.349	2.206	2.045	2.010

GPT2-345m

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

	步数	pass@1	pass@10	pass@100	下载
GPT2-345m (Adam)	300k	0.0840	0.209	0.360	log&config
GPT2-345m (Adan)	150k	0.0843	0.221	0.377	log&config

Adan 在仅一半成本的情况下获得了可比的结果。

视觉任务的结果

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

模型	轮次	训练设置	准确率 (%)	配置	批量大小	下载
ViT-S	150	I	80.1	config	2048	log/model
ViT-S	150	II	79.6	config	2048	log/model
ViT-S	300	I	81.1	config	2048	log/model
ViT-S	300	II	80.7	config	2048	log/model
ViT-B	150	II	81.7	config	2048	log/model
ViT-B	300	II	82.6	config	2048	log/model
ResNet-50	100	I	78.1	config	2048	log/model
ResNet-50	200	I	79.7	config	2048	log/model
ResNet-50	300	I	80.2	config	2048	log/model
ResNet-101	100	I	80.0	config	2048	log/model
ResNet-101	200	I	81.6	config	2048	log/model
ResNet-101	300	I	81.9	config	2048	log/model
ConvNext-tiny	150	II	81.7	config	2048	log//model
ConvNext-tiny	300	II	82.4	config	2048	log/model
MAE-small	800+100	---	83.8	config	4096/2048	log-pretrain/log-finetune/model
MAE-Large	800+50	---	85.9	config	4096/2048	log-pretrain/log-finetune/model

NLP 任务结果

BERT-base

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

预训练	配置	批量大小	日志	模型
Adan	config	256	log	model

在 GLUE 任务上的微调	指标	结果	配置
CoLA	Matthew's 相关系数	64.6	config
SST-2	准确率	93.2	config
STS-B	Pearson 相关系数	89.3	config
QQP	准确率	91.2	config
MNLI	匹配准确率/不匹配准确率	85.7/85.6	config
QNLI	准确率	91.3	config
RTE	准确率	73.3	config

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

Transformer-XL-base

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

	步数	测试困惑度	下载
基线 (Adam)	200k	24.2	log&config
Transformer-XL-base	50k	26.2	log&config
Transformer-XL-base	100k	24.2	log&config
Transformer-XL-base	200k	23.5	log&config

大型语言模型的结果

GPT2-345m

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

	步数	pass@1	pass@10	pass@100	下载
GPT2-345m (Adam)	300k	0.0840	0.209	0.360	log&config
GPT2-345m (Adan)	150k	0.0843	0.221	0.377	log&config

Adan 以一半的成本获得了相当的结果。

扩散模型的结果

我们展示了由 DreamFusion 项目 支持的文本到 3D 任务的结果。更多可视化结果可以在这里找到 here。 使用 Adam 和 Adan 从文本提示 Sydney opera house, aerial view 生成的示例：

内存和效率

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

模型	模型大小 (MB)	Adam 峰值 (MB)	Adan 峰值 (MB)	FusedAdan 峰值 (MB)	Adam 时间 (ms)	Adan 时间 (ms)	FusedAdan 时间 (ms)
ResNet-50	25	7142	7195	7176	9.0	4.2	1.9
ResNet-101	44	10055	10215	10160	17.5	7.0	3.4
ViT-B	86	9755	9758	9758	8.9	12.3	4.3
Swin-B	87	16118	16202	16173	17.9	12.8	4.9
ConvNext-B	88	17353	17389	17377	19.1	15.6	5.0
Swin-L	196	24299	24316	24310	17.5	28.1	10.1
ConvNext-L	197	26025	26055	26044	18.6	31.1	10.2
ViT-L	304	25652	25658	25656	18.0	43.2	15.1
GPT-2	758	25096	25406	25100	49.9	107.7	37.4
GPT-2	1313	34357	38595	34363	81.8	186.0	64.4

Adan 快速上手指南

Adan 是一种基于自适应 Nesterov 动量算法的优化器，专为加速深度模型训练而设计。以下是快速上手指南。

---

环境准备

系统要求

• Python 3.7 或更高版本
• PyTorch 1.10 或更高版本
• 支持 Linux、macOS 和 Windows（推荐在 Linux 上使用）

前置依赖

确保已安装以下依赖：

pip install torch

如果需要国内镜像源加速安装，可以使用以下命令：

pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple

---

安装步骤

通过以下命令安装 Adan：

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

如果需要使用原始版本的 Adan（非 Fused 版本），请执行以下步骤：

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

---

基本使用

最简单的使用示例

步骤 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 中的优化器 betas (默认: None，使用 Adan 默认值 [0.98, 0.92, 0.99])')
parser.add_argument('--no-prox', action='store_true', default=False, help='是否以 AdamW 的方式执行权重衰减 (默认: False)')

步骤 2：创建 Adan 优化器

将现有的优化器替换为 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
)

---

以上是 Adan 的基本使用方法。更多详细配置和实验结果，请参考官方文档或相关链接。