[论文解读] AdaX: Adaptive Gradient Descent with Exponential Long Term Memory

Wenjie Li, Zhaoyang Zhang+2 more|arXiv (Cornell University)|Apr 21, 2020

Advanced Vision and Imaging参考文献 27被引用 24

一句话总结

AdaX 提出了一种新颖的自适应优化算法，用指数长期记忆机制替代 Adam 的过去梯度指数移动平均，以稳定训练并改善收敛性。理论上证明在凸与非凸设置下均可收敛，AdaX 在视觉与自然语言处理基准上优于 Adam，且与带动量的 SGD 表现相当，尤其在泛化能力和对超参数选择的鲁棒性方面表现突出。

ABSTRACT

Although adaptive optimization algorithms such as Adam show fast convergence in many machine learning tasks, this paper identifies a problem of Adam by analyzing its performance in a simple non-convex synthetic problem, showing that Adam's fast convergence would possibly lead the algorithm to local minimums. To address this problem, we improve Adam by proposing a novel adaptive gradient descent algorithm named AdaX. Unlike Adam that ignores the past gradients, AdaX exponentially accumulates the long-term gradient information in the past during training, to adaptively tune the learning rate. We thoroughly prove the convergence of AdaX in both the convex and non-convex settings. Extensive experiments show that AdaX outperforms Adam in various tasks of computer vision and natural language processing and can catch up with Stochastic Gradient Descent.

研究动机与目标

• 为解决 Adam 在非凸优化中不稳定与非收敛问题，特别是其易收敛至次优局部极小值的倾向。
• 克服 Adam 在二阶矩计算中使用指数移动平均所导致的学习率不稳定与泛化性能差的局限。
• 设计一种理论基础坚实的自适应优化器，在保持快速收敛的同时提升泛化性能。
• 通过实证验证 AdaX 在多样化深度学习任务中对 Adam、AdamW、AMSGrad 与带动量 SGD 的优越性。
• 证明 AdaX 对超参数调优具有鲁棒性，尤其在 $eta_2$ 与初始学习率的选择上。

提出的方法

• AdaX 用指数长期记忆机制替代 Adam 中对平方梯度的指数移动平均，以随时间累积历史梯度信息。
• AdaX 中的二阶矩计算为所有历史平方梯度的指数加权和，确保对历史梯度行为的持久记忆。
• 学习率通过长期记忆矩阵的平方根倒数自适应缩放，与 Adam 类似但稳定性更优。
• AdaX 采用改进的更新规则，在保持一阶动量的同时，将自适应分母替换为更稳定、累积的梯度记忆。
• 该算法在凸与非凸设置下均被理论证明可收敛，收敛速度与 AMSGrad 相当。
• AdaX 默认 $eta_2 = 10^{-4}$，实证表明其具有鲁棒性且计算效率高。

实验结果

研究问题

• RQ1 Adam 在非凸问题中快速收敛是否因二阶矩估计不稳定而导致次优解？
• RQ2 用指数长期记忆机制替代 Adam 的指数移动平均，能否消除非收敛问题并改善泛化性能？
• RQ3 在视觉与自然语言处理任务中，AdaX 与 Adam、AdamW、AMSGrad 及带动量 SGD 在收敛速度与最终性能方面相比如何？
• RQ4 在实际训练场景中，AdaX 对超参数选择是否具有鲁棒性，特别是 $eta_2$ 与初始学习率？
• RQ5 AdaX 是否能在保持自适应方法快速收敛的同时，实现与带动量 SGD 相当的性能？

主要发现

• 在 CIFAR-10 上，AdaX 优于 Adam 与 AdamW，测试 Top-1 准确率达到 94.5%，高于 AdamW 的 92.1%。
• 在 ImageNet 上使用 ResNet-18 时，AdaX-W 的 Top-1 准确率为 75.58%，显著优于 AdamW（68.27%）与 AMSGrad(W)（68.76%）。
• 在 ImageNet 上使用 ResNet-50 时，AdaX-W 达到 75.58% 的 Top-1 准确率，接近 SGDM 的 77.12%，同时收敛更快且训练准确率更高。
• AdaX 对学习率与 $eta_2$ 超参数表现出鲁棒性，在不同设置下性能下降极小，尤其当 $eta_2 = 10^{-4}$ 时表现更优。
• 在 One Billion Word 数据集上的语言建模任务中，AdaX 与带动量 SGD 表现相当，展现出优异的泛化能力与快速收敛性。
• 运行时基准测试显示，AdaX-W 与 AdamW 几乎同样高效，ImageNet 上训练时间仅增加 1.8%（50.74h vs. 50.10h），表明计算开销极低。