本文介绍了一种新的自适应步长优化器 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.