MADGRAD 优化方法

一种用于随机优化的动量化、自适应、对偶平均梯度方法

文档可在 https://madgrad.readthedocs.io/en/latest/ 获取。

pip install madgrad

立即试用！这是一种兼具两全其美的优化器，它拥有与 SGD 相当的泛化性能，且收敛速度至少不逊于 Adam，通常甚至更快。我们提供了一个可直接替换 torch.optim 的实现 madgrad.MADGRAD，以及一个封装在 FairSeq 中的实例。对于 FairSeq，只需在项目文件中的任意位置导入 madgrad，并使用命令行选项 --optimizer madgrad，同时配合 --weight-decay、--momentum，以及可选的 --madgrad_eps。

如果您不想通过 pip 安装，也可以将包含该优化器的 madgrad.py 文件直接放入任何 PyTorch 项目中。如果您使用 FairSeq，则还需要配套的 fairseq_madgrad.py 文件。

注意事项：

• 您可能需要使用比平时更低的权重衰减值，很多时候设为 0 即可。
• 建议进行全面的学习率搜索，因为最佳学习率会与 SGD 或 Adam 不同。我们发现的最佳学习率包括：CIFAR10 数据集上 152 层 PreActResNet 的 2.5e-4，ImageNet 上 ResNet-50 的 0.001，使用 transformer_iwslt_de_en 模型的 IWSLT14 数据集上的 0.025，以及在 BookWiki 数据集上训练 RoBERTa 模型时使用 BERT_BASE 的 0.005。此外，在 NLP 模型中，梯度裁剪也有帮助。

镜像 MADGRAD

MADGRAD 的镜像下降版本也一并提供，名为 madgrad.MirrorMADGRAD。该版本在大规模 Transformer 训练中表现尤为出色，甚至优于 MADGRAD。对于数据集足够大、泛化差距不是主要问题的任务，推荐使用此版本。

由于镜像下降版本不会进行隐式正则化，因此您可以通常使用与其他优化器兼容的权重衰减值。

技术报告

无妥协的自适应性：一种用于随机优化的动量化、自适应、对偶平均梯度方法

我们提出了一种名为 MADGRAD 的新型优化方法，属于 AdaGrad 自适应梯度方法家族。MADGRAD 在多个领域的深度学习优化任务中表现出色，涵盖视觉领域的分类和图像到图像任务，以及自然语言处理中的循环神经网络和双向掩码模型。对于这些任务中的每一项，MADGRAD 在测试集性能上均能与 SGD 和 ADAM 竞争，甚至在传统上自适应方法表现不佳的问题上也能超越它们。

@misc{defazio2021adaptivity,
      title={无妥协的自适应性：一种用于随机优化的动量化、自适应、对偶平均梯度方法}, 
      author={Aaron Defazio 和 Samy Jelassi},
      year={2021},
      eprint={2101.11075},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

结果

许可证

MADGRAD 采用 MIT 许可证 许可。

MADGRAD 快速上手指南

MADGRAD 是一种结合了动量、自适应和双重平均梯度的随机优化方法。它兼具 SGD 的泛化性能和 Adam 的快速收敛速度（甚至更快），是深度学习训练的优秀选择。

环境准备

• 系统要求：支持 Linux、macOS 或 Windows。
• 前置依赖：
  • Python 3.6+
  • PyTorch 1.7+
  • (可选) FairSeq：若需在 FairSeq 框架中使用，请确保已安装 fairseq。

安装步骤

推荐使用 pip 进行安装。国内用户可使用清华源加速下载：

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

如果不希望通过 pip 安装，也可以直接将项目中的 madgrad.py 文件复制到你的 PyTorch 项目目录中。若使用 FairSeq，还需同时复制 fairseq_madgrad.py。

基本使用

1. 在原生 PyTorch 项目中使用

作为 torch.optim 的直接替代品，导入并实例化即可：

from madgrad import MADGRAD

# 假设 model 是你的神经网络模型
optimizer = MADGRAD(model.parameters(), lr=0.001, weight_decay=0, momentum=0.9)

# 后续训练循环与普通用法一致
for input, target in dataset:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

注意：

• 建议将 weight_decay 设置为较低值（通常为 0）。
• 需要重新搜索最佳学习率（LR）。参考值：CIFAR10 (ResNet) 约为 2.5e-4，ImageNet (ResNet-50) 约为 0.001。

2. 在 FairSeq 项目中使用

只需在项目的任意文件中导入 madgrad，然后在命令行启动训练时指定优化器参数：

fairseq-train ... --optimizer madgrad --weight-decay 0 --momentum 0.9 --madgrad-eps 1e-6

3. 进阶选项：Mirror MADGRAD

对于大规模 Transformer 训练（如数据集极大，无需担心泛化差距的情况），推荐使用镜像下降版本 MirrorMADGRAD，其表现通常优于标准版，且支持常规的 weight decay 设置。

from madgrad import MirrorMADGRAD

optimizer = MirrorMADGRAD(model.parameters(), lr=0.001, weight_decay=0.01)