人工智能顶刊论文TPAMI2024-新型优化器Adan

本文发表于人工智能顶刊TPAMI（IEEE Transactions on Pattern Analysis and Machine Intelligence），原论文数学推导比较复杂，对于使用者来说不需要过于纠结数学细节，只需要了解算法的创新点在哪里以及如何使用到自己的项目中就行。本文根据个人理解进行阐述，并用一个例子说明如何使用文章提出的优化器。

论文摘要部分

摘要写得的比较清楚，我直接翻译原文

在深度学习中，不同类型的深度网络通常需要不同的优化器，而这些优化器往往需要经过多次试验才能选定，这使得训练过程效率低下。为缓解这一问题，并持续提升各类深度网络的模型训练速度，作者提出了自适应Nesterov动量算法（Adaptive Nesterov momentum algorithm，简称Adan）。Adan首先对传统的Nesterov加速进行了重新表述，开发出一种新的Nesterov动量估计（Nesterov momentum estimation，NME）方法，避免了在外部插值点计算梯度带来的额外开销。接着，Adan采用NME方法，在自适应梯度算法中估计梯度的一阶和二阶矩，以加速收敛。此外，作者证明了在非凸随机问题（如深度学习问题）上，Adan能在$O({ϵ}^{-3.5})$的随机梯度复杂度内找到一个$ϵ$近似的一阶驻点，这与已知的最优下界相匹配。大量实验结果表明，Adan在视觉、语言和强化学习任务上持续超越相应的当前最优（SoTA）优化器，并为许多流行的网络和框架（如ResNet、ConvNext、ViT、Swin、MAE、DETR、GPT-2、Transformer-XL和BERT）创造了新的当前最优记录。更令人惊讶的是，在ViT、GPT-2、MAE等模型上，Adan仅用当前最优优化器一半的训练成本（训练轮数）就能实现更高或相当的性能，并且在大范围的小批量大小（例如从1k到32k）下也表现出很强的适应性。

本文提出的方法

自适应梯度算法（如Adam和AdamW）已成为训练卷积神经网络（CNNs）和视觉Transformer（ViTs）的默认选择。与SGD对所有梯度坐标使用单一学习率不同，自适应算法根据目标函数当前的几何曲率为每个梯度坐标调整学习率，因此收敛速度更快。

Adan 的关键特性

• Nesterov 动量估计 (NME)： 重新定义了传统的 Nesterov 加速，开发了 NME 方法，避免了计算外推点梯度的额外开销。
• 自适应梯度算法： 使用 NME 估计自适应梯度算法中的梯度的一阶和二阶矩，以加速收敛。
• 收敛速度分析： 证明了 Adan 在非凸随机问题 (例如深度学习问题) 中找到$ϵ-$近似一阶驻点的随机梯度复杂度为 $O({ϵ}^{-3.5})$，与最佳已知下界相匹配。

Adan算法

步骤解释：

1. 初始化： 设置初始参数${\theta }_0$(算法开始时的模型参数值，通常是随机初始化)、学习率 $\eta$、动量系数 (${\beta }_1,{\beta }_2,{\beta }_3$)、稳定参数 $ϵ$(防止分母为零的情况)、权重衰减 ${\lambda }_k$和重启条件。

2. 计算梯度： 在当前参数 ${\theta }_k$处估计随机梯度$g_k$。

   新的动量值 ${\mathbf{m}}_k$ 是前一步动量值 ${\mathbf{m}}_{k-1}$和当前梯度$g_k$ 的加权和。

​ ${\mathbf{m}}_k=(1-{\beta }_1){\mathbf{m}}_{k-1}+{\beta }_1{\mathbf{g}}_k;$

​ 更新梯度差项的公式，梯度差系数${\beta }_2$控制了加权的比例，较大的${\beta }_2$ 表示更注重当前的梯度差，梯度差项 $v_k$可以帮助模型更好地捕捉梯度方向的变化趋势，从而提高收敛速度。

​ ${\mathbf{v}}_k=(1-{\beta }_2){\mathbf{v}}_{k-1}+{\beta }_2({\mathbf{g}}_k-{\mathbf{g}}_{k-1})$

​ 更新二阶矩的公式。其中$n_k$表示 第$ k$ 步的二阶矩值，是一个向量，表示模型参数在训练过程中经历的平均平方变化。${\beta }_3$ 是二阶矩系数，控制二阶矩项的权重。通常取值在 (0, 1) 之间。$g_k$是第 $k $步的随机梯度，它是一个向量，表示模型参数在第 $k$ 步的瞬时变化。二阶矩项$n_k$可以帮助模型更好地理解梯度方向的变化趋势，并计算学习率，从而提高收敛速度和泛化能力。二阶矩系数 ${\beta }_3$可以调整二阶矩项的权重，以平衡模型对过去和当前梯度值平方信息的利用。

​ ${\mathbf{n}}_k=(1-{\beta }_3){\mathbf{n}}_{k-1}+{\beta }_3[{\mathbf{g}}_k+(1-{\beta }_2)({\mathbf{g}}_k-{\mathbf{g}}_{k-1}){]}^2$

​ 计算学习率 ${\mathbf{\eta }}_k$ 的公式学习率公式可以根据二阶矩值自动调整学习率的大小，从而实现自适应学习率调整。

​ ${\mathbf{\eta }}_k=\eta /(\sqrt{{\mathbf{n}}_k}+\epsilon )$

​ 更新模型参数 ${\mathbf{\theta }}_{k+1}$ 的公式。${\lambda }_k$是权重衰减系数，用于控制权重衰减的程度。通常取值在 (0, 1) 之间。新的模型参数值${\mathbf{\theta }}_{k+1}$ 是前一步模型参数值${\mathbf{\theta }}_k$ 减去学习率 ηk 乘以动量项${\mathbf{m}}_k$加上梯度差项 $(1-{\beta }_2){\mathbf{v}}_k$ 的商。模型参数更新公式综合考虑了动量项、梯度差项和权重衰减项，从而实现自适应参数更新。

​ ${\mathbf{\theta }}_{k+1}=(1+{\lambda }_k\eta {)}^{-1}[{\mathbf{\theta }}_k-{\mathbf{\eta }}_k\circ ({\mathbf{m}}_k+(1-{\beta }_2){\mathbf{v}}_k)];$

Adan 算法中的重启条件

Adan 算法中的重启条件用于在训练过程中重新初始化动量项和参数，以稳定优化过程并提高收敛速度。具体来说，当满足以下条件时，算法将触发重启：

$(k+1){\sum }_{t=0}^k{\Vert {\mathbf{y}}_{t+1}-{\mathbf{y}}_t\Vert }_{\sqrt{{\mathbf{n}}_t}}^2>R^2$，其中$y_t$ 是第 $t+1$ 步的模型参数，即$y_t=x_k-\beta {\eta }_k[g_k+(1-{\beta }_2)(g_k-g_{k-1})]$,$n_t$是第 $t$ 步的二阶矩值，$R$是重启阈值，控制重启的频率。较大的 $R$ 表示较少的重启，较小的$ R$ 表示较多的重启。

这个条件简单说来表示当模型参数在历史迭代中经历的总变化超过阈值 $R^2$ 时，触发重启。总变化可以看作是模型探索新解空间的能力，如果总变化过小，模型可能陷入局部最优解，而重启可以使其重新探索解空间。重启条件可以帮助模型在训练过程中更好地平衡探索和利用，从而提高收敛速度和泛化能力。重启频率可以通过调整 R 的值来控制，以适应不同的训练任务和数据集。

优化器中常用的动量解释

动量是深度学习优化器中的一个标配组件，从SGDM开始就存在。动量项是优化算法中用于加速收敛的一种技术，它借鉴了物理学中动量的概念。在优化算法中，动量项用于累积梯度信息，帮助模型更好地捕捉梯度方向的变化趋势，从而提高收敛速度。其工作原理如下：

1. 累积梯度信息： 在训练过程中，模型会根据当前的梯度信息更新参数。动量项会记录过去几步的梯度信息，形成一个累积的动量向量。
2. 平滑梯度方向： 当梯度方向发生较缓慢的变化时，动量项会平滑梯度方向，使模型能够更好地跟随梯度变化趋势，从而避免陷入局部最优解。
3. 加速收敛： 当梯度方向发生较快速的变化时，动量项会保留部分过去动量，帮助模型更快地响应梯度变化，从而加速收敛速度。

二阶矩值和动量的不同用途

在 Adan 算法中，二阶矩和动量是两个不同的概念，它们在优化过程中起着不同的作用：

二阶矩：一个向量，用于估计模型参数在训练过程中经历的平均平方变化。它可以帮助模型更好地理解梯度方向的变化趋势，并计算学习率，从而提高收敛速度和泛化能力。

**动量：**一个向量，用于累积梯度信息，帮助模型更好地捕捉梯度方向的变化趋势，从而提高收敛速度。动量可以帮助模型更快地捕捉梯度方向的变化趋势，从而加速收敛速度。同时，动量可以平滑梯度方向，使模型能够更好地跟随梯度变化趋势，从而避免陷入局部最优解。

二阶矩和动量在优化过程中起着不同的作用：二阶矩主要用于计算学习率，而动量主要用于加速收敛速度。二阶矩和动量都是优化算法中的重要概念，它们可以帮助模型更好地适应训练过程中的变化，提高收敛速度和泛化能力。

对比Adam和RMSProp优化器的改进：

Adan优化器与Adam和RMSProp相比，主要有以下几方面的区别和改进：

1. Nesterov动量：Adan的主要创新之一是引入了Nesterov动量，它在梯度更新时使用了一个外推点的梯度，这样可以更准确地预测参数的未来变化轨迹，从而加速收敛。相比之下，Adam和RMSProp都采用了较为简单的动量计算，Adam使用的是梯度的移动平均，RMSProp则是使用梯度平方的平均值。它们并没有利用Nesterov加速技术，这使得它们在一些优化问题上收敛速度较慢。
2. 自适应梯度计算：Adan在自适应梯度计算方面做出了改进。RMSProp根据梯度平方来调整学习率，Adam则使用梯度的一阶和二阶矩来调整学习率，而Adan进一步通过引入梯度差异（即连续两次梯度之间的变化）来优化动量估计。这个改进使得Adan能够更好地处理动态变化的优化问题，提高了稳定性和效率，尤其是在大规模模型训练中。
3. 收敛速度：理论上，Adan相比Adam和RMSProp具有更快的收敛速度，特别是在深度神经网络（DNN）中。Adan通过结合Nesterov加速和更精确的动量估计，显著加快了收敛速度。例如，Adan的复杂度为 $O({ϵ}^{-3.5})$，而Adam的复杂度为 $O({ϵ}^{-4})$，这意味着Adan能够以更少的训练步数达到更优的结果。
4. 实证表现：在实际实验中，Adan在多种任务上都表现出了优于Adam和RMSProp的性能，尤其是在视觉任务（如ViT）、语言模型（如LSTM、GPT-2）以及Transformer-XL等大规模模型的训练中，Adan能够在更少的训练轮次内达到更高的准确率

论文实验结果

在计算机视觉、自然语言处理模型和任务上都有不同程度的涨点，训练曲线也相对稳定。

本地实验测试

下载github源码：https://github.com/sail-sg/Adan

解压缩后安装：

cd Adan
Python setup.py install

测试一下：

from adan import Adan

通过两步使用Adan

第一步，添加Adan需要的超参数：

parser.add_argument('--max-grad-norm', type=float, default=0.0, help='if the l2 norm is large than this hyper-parameter, then we clip the gradient  (default: 0.0, no gradient clip)')
parser.add_argument('--weight-decay', type=float, default=0.02,  help='weight decay, similar one used in AdamW (default: 0.02)')
parser.add_argument('--opt-eps', default=None, type=float, metavar='EPSILON', help='optimizer epsilon to avoid the bad case where second-order moment is zero (default: None, use opt default 1e-8 in adan)')
parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA', help='optimizer betas in Adan (default: None, use opt default [0.98, 0.92, 0.99] in Adan)')
parser.add_argument('--no-prox', action='store_true', default=False, help='whether perform weight decay like AdamW (default=False)')

其中：

opt-betas:为了与其他优化器使用习惯保持一致, 论文中的$\beta$参数 在代码中实际是$(1-\beta )$。

foreach (bool):如果为 True, Adan 会使用 torch._foreach 实现，运行速度更快但会耗费更多存储.

no-prox: 确定了带有权重衰减的参数更新规则。默认情况下，Adan以论文中算法所示的方式更新参数:

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

但是也可以按照AdamW:的方式更新参数：

$${\mathbf{\theta }}_{k+1}=(1-\lambda \eta ){\mathbf{\theta }}_k-{\mathbf{\eta }}_k\circ ({\mathbf{m}}_k+(1-{\beta }_2){\mathbf{v}}_k).$$

第二步：创建 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)

使用中的一些帮助：

1.为了使 Adan 更加简单，在论文中除表 12 外的所有实验中，均未使用 Adan 的重启策略。但论文中表 12 显示，重启策略可以进一步略微提升 Adan 的性能。

2.Adan 通常允许使用较大的峰值学习率，而这种学习率通常会使其他优化器（例如 Adam 和 AdamW）失效。例如，在所有实验中（MAE 预训练和 LSTM 除外），Adan 使用的学习率是 Adam/AdamW 的 5-10 倍。

3.Adan 相对来说对 ${\beta }_1$、${\beta }_2$ 和 ${\beta }_3$ 更为鲁棒，尤其是 ${\beta }_2$。如果希望获得更好的性能，可以先调整 ${\beta }_3$，然后再调整 ${\beta }_1$。

4.在单节点上，Adan 的 GPU 内存成本略高于 Adam/AdamW。然而，这一问题可以通过使用 ZeroRedundancyOptimizer 解决。该优化器通过在分布式数据并行进程中共享优化器状态来减少每个进程的内存占用。具体来说，当在两个以上的 GPU 上使用 ZeroRedundancyOptimizer 时，Adan 和 Adam 的内存占用几乎相同。

下面我写了一个在Resnet18上训练测试CiFAR10的代码测试Adan、Adam和RMSprop，Adan的超参数设定按照论文中的来，我没有修改，同样的训练10个epoch:

import time
import torch
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
from torch import nn
from torch.utils.data import DataLoader
from torchvision import models
from adan import Adan  # 假设你已安装Adan并正确导入

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# CIFAR10数据集的转换和加载
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

trainloader = DataLoader(trainset, batch_size=128, shuffle=True)
testloader = DataLoader(testset, batch_size=100, shuffle=False)

# ResNet18模型加载
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 10)  # 修改最后一层以适应CIFAR10
model = model.to(device)

# 定义损失函数
criterion = nn.CrossEntropyLoss()


# 训练函数
def train(model, optimizer, trainloader, criterion, epochs=10):
    model.train()
    start_time = time.time()
    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        for inputs, labels in trainloader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()

            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        print(f"Epoch {epoch + 1}, Loss: {running_loss / len(trainloader)}, Accuracy: {100 * correct / total}%")

    end_time = time.time()
    return end_time - start_time


# 测试函数
def test(model, testloader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in testloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Test Accuracy: {100 * correct / total}%')


# 定义优化器并进行训练
def run_experiment(optimizer_name, optimizer_params):
    print(f"Running experiment with {optimizer_name} optimizer:")

    # 初始化模型
    model = models.resnet18(pretrained=True)
    model.fc = nn.Linear(model.fc.in_features, 10)  # 修改最后一层以适应CIFAR10
    model = model.to(device)

    optimizer = optimizer_name(model.parameters(), **optimizer_params)
    train_time = train(model, optimizer, trainloader, criterion, epochs=10)

    print(f"Training time for {optimizer_name}: {train_time:.2f} seconds")
    test(model, testloader)
    print("\n" + "=" * 50 + "\n")


# RMSProp实验
rmsprop_params = {'lr': 0.01}
run_experiment(optim.RMSprop, rmsprop_params)

# Adam实验
adam_params = {'lr': 0.001}
run_experiment(optim.Adam, adam_params)

# Adan实验
adan_params = {
    'lr': 0.001,
    'weight_decay': 0.02,
    'max_grad_norm': 0.0,
    'no_prox': False
}
run_experiment(Adan, adan_params)

运行输出：

Running experiment with <class 'torch.optim.rmsprop.RMSprop'> optimizer:
Epoch 1, Loss: 2.4968030233212444, Accuracy: 23.17%
Epoch 2, Loss: 1.525029681222823, Accuracy: 45.332%
Epoch 3, Loss: 1.1440054220921547, Accuracy: 59.082%
Epoch 4, Loss: 0.9422917885853507, Accuracy: 66.778%
Epoch 5, Loss: 0.8193090390366362, Accuracy: 71.25%
Epoch 6, Loss: 0.7436501835008411, Accuracy: 74.384%
Epoch 7, Loss: 0.664197140032678, Accuracy: 76.768%
Epoch 8, Loss: 0.5913919725686388, Accuracy: 79.518%
Epoch 9, Loss: 0.5181573262757353, Accuracy: 81.964%
Epoch 10, Loss: 0.4775610526123315, Accuracy: 83.598%
Training time for <class 'torch.optim.rmsprop.RMSprop'>: 40.26 seconds
Test Accuracy: 70.16%

==================================================

Running experiment with <class 'torch.optim.adam.Adam'> optimizer:
Epoch 1, Loss: 0.8910103748216653, Accuracy: 70.016%
Epoch 2, Loss: 0.5801464280356532, Accuracy: 80.438%
Epoch 3, Loss: 0.4486024776078246, Accuracy: 84.732%
Epoch 4, Loss: 0.35622213636060507, Accuracy: 88.022%
Epoch 5, Loss: 0.27932105611657243, Accuracy: 90.442%
Epoch 6, Loss: 0.2237478634890388, Accuracy: 92.424%
Epoch 7, Loss: 0.18003986158486826, Accuracy: 93.848%
Epoch 8, Loss: 0.1396366751817105, Accuracy: 95.262%
Epoch 9, Loss: 0.12056110491094839, Accuracy: 95.826%
Epoch 10, Loss: 0.10910908536761618, Accuracy: 96.322%
Training time for <class 'torch.optim.adam.Adam'>: 39.90 seconds
Test Accuracy: 80.64%

==================================================

Running experiment with <class 'adan.Adan'> optimizer:
Epoch 1, Loss: 0.843513864537944, Accuracy: 71.122%
Epoch 2, Loss: 0.4961339098870602, Accuracy: 83.222%
Epoch 3, Loss: 0.3540309919115832, Accuracy: 87.93%
Epoch 4, Loss: 0.2544422761901565, Accuracy: 91.348%
Epoch 5, Loss: 0.18906724664485058, Accuracy: 93.608%
Epoch 6, Loss: 0.1524442125140401, Accuracy: 94.898%
Epoch 7, Loss: 0.12659881578858398, Accuracy: 95.828%
Epoch 8, Loss: 0.10798529118223263, Accuracy: 96.334%
Epoch 9, Loss: 0.09582332747719248, Accuracy: 96.712%
Epoch 10, Loss: 0.08477451177814123, Accuracy: 97.124%
Training time for <class 'adan.Adan'>: 47.10 seconds
Test Accuracy: 82.86%

这里没有调参，从测试集准确率上，确实要更好一些。

这里再单独试试Adan和Adam在更大学习率上的效果，设置学习率为0.01，训练50个epoch

Running experiment with <class 'torch.optim.adam.Adam'> optimizer:
Epoch 50, Loss: 0.04389519359001799, Accuracy: 98.49%
Training time for <class 'torch.optim.adam.Adam'>: 201.90 seconds
Test Accuracy: 75.3%

Running experiment with <class 'adan.Adan'> optimizer:
Epoch 50, Loss: 0.14560543577117688, Accuracy: 94.894%
Training time for <class 'adan.Adan'>: 237.20 seconds
Test Accuracy: 76.17%

从测试集效果更好一点，但是也更耗时。