真实训练中Adam无法收敛到最佳位置？改进版 AdaX来拯救！

本文介绍了一种新的自适应步长优化器 AdaX，它简单而高效，能够较好地弥补 Adam 在高频梯度噪声时存在的缺陷，即在真实训练中无法收敛到最佳位置的问题。

作者 | 李文杰

编辑 | 丛 末

论文地址：https://arxiv.org/pdf/2004.09740.pdf

开源地址：https://github.com/switchablenorms/adax

1

故事背景

自从Reddi et al. [1] 在2018年指出了Adam 文章中的证明错误之后，对Adaptive算法的研究就成了一个热门，之后也出现了许多优秀的算法，如AdaShift[2], NosAdam [3], AdaBound [4],以及最新的Radam [5]，但是大多数算法都还难以在泛化性能上与SGD+Momentun打成平手。并且，Reddi et. al 文章中的反例也颇有一些令人费解之处，下面先简要回顾。

假设我们需要最小化函数ft(x),

x的定义域是[-1, 1]，并且起始值为 x=1，且C>2。在上图的函数中，每当时间t除以3余1时，正确的梯度（即C）会指示算法朝-1的方向走去，而不正确的梯度（即-1）则会让算法朝 1 走去。Reddi等人构造出这样一个拥有高频梯度噪声的函数，且证明了Adam会反复回到 x=1的位置难以动弹，而SGD和AdaGrad 这些算法却能够成功到达 x=-1。即Adam在高频梯度噪声时存在缺陷！

2

新的算法

然而，这么高频率的错误梯度，在实际的 DNN 训练中也许是寥寥无几。所以在真实训练中，Adam 依旧是可以收敛到不错的位置，只是可能并没有到最佳（比不上SGD所能达到的最佳位置）。所以在我们的论文中，我们重新设计了一个例子来模拟真实的训练场景。

这里 n 是一个小常数，可以简单设为1， 0< lambda < 1，C>1 都为常数。这个问题其实就是模拟在训练过程中，由于参数不断接近最优解，梯度的大小不断减小的过程。观察可以发现，f_t之和的最小值在0处达到。我们希望验证的是，在初始学习率（alpha_0)和初始值 (x_0)都可以随便取的情况下，各个算法是否都能到达 0 点。很显然，通过等比数列求和公式，SGD当然是可以到达0的。但神奇的是，无论取多小的alpha_0, 无论初始值靠近还是远离0，Adam都无法在0点停下来，都会到达负数解。

简单的来说，是因为Adam的分母V_t 对于小的梯度过于敏感，导致Adam的步长非常的大，大到了即便梯度在指数级减小，Adam的步长依旧大于一个常数。所以，虽然Adam收敛得非常快，但却有可能得不偿失，收敛到一个次优解，并且由于次优解处的梯度太小而难以重新回到最优解。

为了解决这个问题，我们重新设计了Adam的分母。由于梯度在训练后期减小非常快，甚至为0，无法提供更多有效信息。我们提出强调之前的v_t的重要性，而逐渐减弱对现在的梯度的适应。最简单的方式就是将Adam的分母公式反过来，将原来指数级遗忘的v_t改为指数级叠加，也就是我们的新AdaX算法如下：

简单的来说，就是将原本的0.999 v_{t-1} +0.001 g_t^2 改为了 1.0001 v_{t-1} + 0.0001 g_t^2。当然，第七行的系数也要随之变化。我们希望通过更合理地增速（因为分母的指数遗忘改为了指数积累），能够让Adaptive类算法达到和SGD一样的效果。在我们的实验中，AdaX的表现非常不错，甚至能在一些传统自适应算法效果很差的实验中表现得很好。

在CIFAR10， ImageNet和VOC Segmentation等等任务中, AdaX既能比SGD快，又能和SGD一样好，可以说是非常惊喜了。

我们在文章中还有其他更多的实验，欢迎大家移步arXiv看看，相关理论分析也在论文当中，也可以用代码试一试。从Adam改到AdaX变化非常的小，很容易实现。期待大家的结果。

参考文献：

[1]. Reddi, S. J., Kale, S., and Kumar., S. Onthe convergence of adam and beyond.Proceedings of the 6th InternationalConference on Learning Representations (ICLR), 2018. [2]. Zhou, Z., Zhang, Q., Lu, G., Wang, H., Zhang, W., and Yu, Y. Adashift:Decorrelation and convergence of adaptive learning rate methods. Proceedingsof 7th International Conference on Learning Representations (ICLR), 2019.

[3]. Huang, H., Wang, C., and Dong., B.Nostalgic adam: Weighting more of the past gradients when designing theadaptive learning rate. arXiv preprint arXiv: 1805.07557,2019. [4]. Luo, L., Xiong, Y., Liu, Y., and Sun, X. Adaptive gradi- ent methods withdynamic bound of learning rate. Proceedings of 7th InternationalConference on Learning Representations, 2019.

[5]. Liu L., Jiang, H., He, P., Chen, W., Liu,X., Gao, J., Han, J. On the Variance of the Adaptive Learning Rate andBeyond. Proceedings of 8th International Conference on LearningRepresentations, 2020.