This article provides a brief introduction to an optimizer called AdaX, from the paper “AdaX: Adaptive Gradient Descent with Exponential Long Term Memory”. The reason for introducing this optimizer is that it once again confirms a conclusion mentioned in the previous article “Analysis of the AdaFactor Optimizer (with Open Source Implementation)”. These two articles can be read in comparison.
Adam & AdaX
The update format for AdaX is: \left\{\begin{aligned} &g_t = \nabla_{\theta} L(\theta_t)\\ &m_t = \beta_1 m_{t-1} + \left(1 - \beta_1\right) g_t\\ &v_t = (1 + \beta_2) v_{t-1} + \beta_2 g_t^2\\ &\hat{v}_t = v_t\left/\left(\left(1 + \beta_2\right)^t - 1\right)\right.\\ &\theta_t = \theta_{t-1} - \alpha_t m_t\left/\sqrt{\hat{v}_t + \epsilon}\right. \end{aligned}\right. where the default value of \beta_2 is 0.0001. By the way, here is my Keras implementation: https://github.com/bojone/adax.
For comparison, the update format for Adam is: \left\{\begin{aligned} &g_t = \nabla_{\theta} L(\theta_t)\\ &m_t = \beta_1 m_{t-1} + \left(1 - \beta_1\right) g_t\\ &v_t = \beta_2 v_{t-1} + \left(1 - \beta_2\right) g_t^2\\ &\hat{m}_t = m_t\left/\left(1 - \beta_1^t\right)\right.\\ &\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\ &\theta_t = \theta_{t-1} - \alpha_t \hat{m}_t\left/\sqrt{\hat{v}_t + \epsilon}\right. \end{aligned}\right. where the default value of \beta_2 is 0.999.
Equivalent Form Transformation
As we can see, the first difference between the two is that AdaX removes the bias correction for momentum (the step \hat{m}_t = m_t / (1 - \beta_1^t)), but this actually has little impact. The most significant change in AdaX is at v_t. Originally, v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 is in a moving average format, whereas v_t = (1 + \beta_2) v_{t-1} + \beta_2 g_t^2 does not look like a moving average. Furthermore, since 1 + \beta_2 > 1, does there seem to be a risk of exponential explosion? The original paper calls this “with Exponential Long Term Memory,” meaning that 1 + \beta_2 > 1 causes the weight of historical accumulated gradients not to decrease over time, but rather to increase. This is its long-term memory property.
In fact, the learning rate correction uses \hat{v}_t, so to determine if there is an explosion, we must observe \hat{v}_t. For Adam, we have: \begin{aligned} \hat{v}_t =& v_t\left/\left(1 - \beta_2^t\right)\right.\\ =& \frac{\beta_2 v_{t-1} + (1-\beta_2) g_t^2}{1 - \beta_2^t}\\ =& \frac{\beta_2 \hat{v}_{t-1}\left(1 - \beta_2^{t-1}\right) + (1-\beta_2) g_t^2}{1 - \beta_2^t}\\ =& \beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}\hat{v}_{t-1} + \left(1 - \beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}\right)g_t^2 \end{aligned} So if we set \hat{\beta}_{2,t}=\beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}, then the update formula is: \hat{v}_t =\hat{\beta}_{2,t}\hat{v}_{t-1} + \left(1 - \hat{\beta}_{2,t}\right)g_t^2 Based on the same logic, if we set \hat{\beta}_{2,t}=1 - \frac{\beta_2}{(1 + \beta_2)^t - 1}, then the update formula for \hat{v}_t in AdaX can also be written in the above form.
Comparison of Decay Strategies
Therefore, looking at \hat{v}_t which is actually used to correct the gradient, both Adam and AdaX follow a moving average format in their update formulas, except that the corresponding decay coefficients \hat{\beta}_{2,t} are different.
For Adam, when t=1, \hat{\beta}_{2,t}=0. At this point, \hat{v}_t is simply g_t^2, meaning the real-time gradient is used to correct the learning rate, which is when the correction intensity is at its maximum. When t \to \infty, \hat{\beta}_{2,t} \to \beta_2. At this stage, v_t is a weighted average of the accumulated squared gradients and the current squared gradient. Since \beta_2 < 1, it means the weight of the current gradient 1 - \beta_2 is non-zero. This may lead to training instability because, in the later stages of training, gradients become smaller and the training itself tends to stabilize, making the significance of learning rate correction smaller. Therefore, the intensity of learning rate correction should decrease, and as t \to \infty, the learning rate should ideally remain constant (equivalent to degenerating into SGD). This requires that when t \to \infty, \hat{\beta}_{2,t} \to 1.
For AdaX, when t=1, \hat{\beta}_{2,t}=0, and when t \to \infty, \hat{\beta}_{2,t} \to 1, which satisfies the ideal properties mentioned above. Therefore, from this perspective, AdaX is indeed an improvement over Adam. AdaFactor uses \hat{\beta}_{2,t} = 1 - \frac{1}{t^c}, which is also designed from this perspective. As for whether the strategy of AdaX or AdaFactor is superior, the author believes it is difficult to explain clearly from a theoretical standpoint and likely depends on experimental results.
Conclusion
And so, the article ends here. As stated at the beginning, this post is just a brief introduction to AdaX because it once again confirms a previous conclusion—that \hat{\beta}_{2,t} should satisfy the conditions “\hat{\beta}_{2,1}=0, \hat{\beta}_{2,\infty}=1.” This may become one of the basic conditions for the improvement of optimizers in the future.
When reposting, please include the original address of this article: https://kexue.fm/archives/7387
For more detailed reposting matters, please refer to: “Scientific Space FAQ”