A Brief Look at the AdaX Optimizer (with Open-Source Implementation)

AdaX优化器浅析（附开源实现） 苏剑林 · 2020-05-11 · Machine Learning optimizer

This post gives a brief introduction to an optimizer called AdaX, from AdaX: Adaptive Gradient Descent with Exponential Long Term Memory. The reason for introducing this optimizer is that it once again confirms a conclusion mentioned earlier in A Brief Look at the AdaFactor Optimizer (with Open-Source Implementation); the two posts are worth reading side by side.

Adam & AdaX

The update rule for AdaX is

\begin{equation}\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.\end{equation}

where the default value of $\beta_2$ is $0.0001$. By the way, here's my Keras implementation: https://github.com/bojone/adaxmore

For comparison, the update rule for Adam is

\begin{equation}\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.\end{equation}

where the default value of $\beta_2$ is $0.999$.

Transforming into an Equivalent Form

As we can see, the first difference between the two is that AdaX drops the bias-correction step for the momentum term (the step $\hat{m}_t = m_t\left/\left(1 - \beta_1^t\right)\right.$), but that actually doesn't matter much. AdaX's biggest change is at $v_t$: whereas $v_t = \beta_2 v_{t-1} + \left(1 - \beta_2\right) g_t^2$ is a moving-average form, $v_t = (1 + \beta_2) v_{t-1} + \beta_2 g_t^2$ doesn't look like a moving average at all, and moreover $1 + \beta_2 > 1$ — which seems to carry a risk of exponential blow-up? The original paper calls this "Exponential Long Term Memory," referring to the fact that $1 + \beta_2 > 1$ causes the weight of accumulated historical gradients not to shrink over time but instead to grow larger and larger — that's its long-term memory property.

In fact, the term used to correct the learning rate is $\hat{v}_t$, so to determine whether there's really any blow-up, what we should actually look at is $\hat{v}_t$. For Adam, we have

\begin{equation}\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}\end{equation}

so if we set $\hat{\beta}_{2,t}=\beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}$, then the update formula becomes

\begin{equation}\hat{v}_t =\hat{\beta}_{2,t}\hat{v}_{t-1} + \left(1 - \hat{\beta}_{2,t}\right)g_t^2\end{equation}

By the same reasoning, 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.

Comparing Decay Strategies

So, looking at $\hat{v}_t$ — the quantity actually used to correct the gradient — both Adam and AdaX have update formulas that take the form of a moving average; the only difference lies in the corresponding decay coefficient $\hat{\beta}_{2,t}$.

For Adam, when $t=1$, $\hat{\beta}_{2,t}=0$, in which case $\hat{v}_t$ is simply $g_t^2$ — that is, the learning rate is corrected using the instantaneous gradient, giving the strongest correction. When $t\to\infty$, $\hat{\beta}_{2,t}\to \beta_2$, and in that case $v_t$ is a weighted average of the accumulated squared gradient and the current squared gradient; since $\beta_2 < 1$, this means the weight $1 - \beta_2$ of the current gradient is never zero. This can lead to training instability, because later in training the gradient becomes small and training itself tends toward stability, at which point correcting the learning rate no longer serves much purpose — so the strength of the correction should decrease. Moreover, $t\to\infty$: ideally the learning rate should settle to a constant (at which point it effectively degenerates to SGD), which requires that as $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$ — satisfying the desired property above. So, from this angle, AdaX is indeed an improvement over Adam. AdaFactor, for its part, uses $\hat{\beta}_{2,t} =1 - \frac{1}{t^c}$, which is also designed with this same principle in mind. As for which strategy — AdaX's or AdaFactor's — is actually better, I think that's hard to settle purely on theoretical grounds; it probably has to be determined experimentally.

And That's It

Well, that's where this post ends. As I said at the start, this is just meant to be a brief introduction to AdaX, because it once again confirms an earlier conclusion — that $\hat{\beta}_{2,t}$ should satisfy the condition "$\hat{\beta}_{2,1}=0,\hat{\beta}_{2,\infty}=1$" — which may well become one of the basic requirements for future optimizer improvements.