AdaXod: a new adaptive and momental bound algorithm for training deep neural networks

Abstract

Adaptive algorithms are widely used in deep learning because of their fast convergence. Among them, Adam is the most widely used algorithm. However, studies have shown that Adam’s generalization ability is weak. AdaX is a variant of Adam, which introduces a novel second-order momentum, modifies the second-order moment of Adam, and has good generalization ability. However, these algorithms may fail to converge due to instability and extreme learning rates during training. In this paper, we propose a new adaptive and momental bound algorithm, called AdaXod, which characterizes of exponentially averaging the learning rate and is particularly useful for training deep neural networks. By setting an adaptively limited learning rate in the AdaX algorithm, the resultant AdaXod can effectively eliminate the problem of excessive learning rate in the later stage of neural networks training and thus results in stable training. We conduct extensive experiments on different datasets and verify the advantages of the AdaXod algorithm by comparing with other advanced adaptive optimization algorithms. AdaXod eliminates large learning rates during neural networks training and outperforms other optimizers, especially for some neural networks with complex structures, such as DenseNet.

Appendix A Details About Proof

Before proving Theorem  1, we need to establish the following result of Lemma 1.

Lemma 1

For the parameter settings and conditions assumed in Theorem  1, we have

$$\begin{aligned} \sum \limits _{t = 1}^T {\left[\frac{1}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}{m_t}} \right\| }^2} + \frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}{m_{t - 1}}} \right\| }^2}\right]} \le \frac{{{R_\infty }G_2^2}}{{1 - {\beta _1}}} \end{aligned}$$

(A1)

Proof

According to the definition of ${\eta _t}$, by the formula ${\eta _t}= {p_t}$, we have

$$\begin{aligned} {\left\| {{\eta _t}} \right\| _\infty } = {\left\| {{p_t}} \right\| _\infty } \le {R_\infty } \end{aligned}$$

(A2)

Hence,

$$\begin{aligned} \begin{aligned}&\sum \limits _{t = 1}^T {\left[\frac{1}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}{\varvec{m}_t}} \right\| }^2} + \frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}{\varvec{m}_{t - 1}}} \right\| }^2}\right]} \\&\quad \le \sum \limits _{t = 1}^T {\left[\frac{{{R_\infty }}}{{2(1 - {\beta _1})}}{{\left\| {{\varvec{m}_t}} \right\| }^2} - \frac{{{R_\infty }}}{{2(1 - {\beta _1})}}{{\left\| {{\varvec{m}_{t - 1}}} \right\| }^2}\right]} \\&\quad = \frac{{{R_\infty }}}{{2(1 - {\beta _1})}}\left[\sum \limits _{t = 1}^T {{{\left\| {{\varvec{m}_t}} \right\| }^2}} + \sum \limits _{t = 1}^T {{{\left\| {{\varvec{m}_{t - 1}}} \right\| }^2}} \right]\\&\quad \le \frac{{{R_\infty }}}{{2(1 - {\beta _1})}}\left[\frac{1}{T}{\left[\sum \limits _{t = 1}^T {\left\| {{\varvec{m}_t}} \right\| } \right]^2} + \frac{1}{T}{\left[\sum \limits _{t = 1}^T {\left\| {{\varvec{m}_{t - 1}}} \right\| } \right]^2}\right]\\&\quad \le \frac{{{R_\infty }G_2^2}}{{(1 - {\beta _1})}} \end{aligned} \end{aligned}$$

(A3)

The first inequality is due to ${\beta _{1t}} \le {\beta _1} < 1$. We complete the proof of this Lemma 1. $\square$

Proof of Theorem 1

Let ${\varvec{\theta }^*} = \arg {\min _{\theta \in F}}(\sum \nolimits _{t = 1}^T {{f_t}({\varvec{\theta }^*})} )$, which exists since F is closed and convex. From the definition of the projection operation, we get the observation.

$$\begin{aligned}{} & {} \begin{aligned} {\varvec{\theta }_{t + 1}} = \prod \nolimits _F {({\varvec{\theta }_t} - {\eta _t}{\varvec{m}_t})}= \min \left\| {{\eta _t}^{ - 1/2}(\varvec{\theta } - ({\varvec{\theta }_t} - {\eta _t}{\varvec{m}_t}))} \right\| \end{aligned} \end{aligned}$$

(A4)

$$\begin{aligned}{} & {} \begin{aligned} {\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| ^2}&\le {\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\eta _t}{\varvec{m}_t} - {\varvec{\theta }^*})} \right\| ^2}\\&= {\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| ^2} \\&\quad+ {\left\| {{\eta _t}^{1/2}{\varvec{m}_t}} \right\| ^2}- 2\left\langle {{\varvec{m}_t},{\varvec{\theta }_t} - {\varvec{\theta }^*}} \right\rangle \\&= {\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| ^2} + {\left\| {{\eta _t}^{1/2}{\varvec{m}_t}} \right\| ^2} \\&\quad -2\left\langle {{\beta _{1t}}{\varvec{m}_{t - 1}} + (1 - {\beta _{1t}}){\varvec{g}_t},{\varvec{\theta }_t} - {\varvec{\theta }^*}} \right\rangle \end{aligned} \end{aligned}$$

(A5)

Rearranging the above inequalities (A5), we get

$$\begin{aligned} \begin{aligned} \left\langle {{\varvec{g}_t},{\varvec{\theta }_t} - {\varvec{\theta }^*}} \right\rangle&\le \frac{1}{{2(1 - {\beta _{1t}})}}\left[{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| ^2} - {\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_{t + 1}} - {\varvec{\theta }^*})} \right\| ^2}\right]\\&\quad + \frac{1}{{2(1 - {\beta _{1t}})}}{\left\| {{\eta _t}^{1/2}{\varvec{m}_t}} \right\| ^2} - \frac{{{\beta _{1t}}}}{{1 - {\beta _{1t}}}}\left\langle {{\varvec{m}_{t - 1}},{\varvec{\theta }_t} - {\varvec{\theta }^*}} \right\rangle \\&\le \frac{1}{{2(1 - {\beta _{1t}})}}\left[{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| ^2} - {\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_{t + 1}} - {\varvec{\theta }^*})} \right\| ^2}\right] \\&\quad + \frac{1}{{2(1 - {\beta _{1t}})}}{\left\| {{\eta _t}^{1/2}{\varvec{m}_t}} \right\| ^2} + \frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{\left\| {{\eta _t}^{1/2}{\varvec{m}_{t - 1}}} \right\| ^2} \\&\quad + \frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{\left\| {{\eta _t}^{1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| ^2} \end{aligned} \end{aligned}$$

(A6)

Then by the convexity of the function ${{f_t}({\varvec{\theta }})}$ at each step, it yields

$$\begin{aligned} \begin{aligned} \sum \limits _{t = 1}^T {{f_t}({\varvec{\theta }_t})} - {f_t}({\varvec{\theta }^*})& \le \sum \limits _{t = 1}^T {\left\langle {{\varvec{g}_t},{\varvec{\theta }_t} - {\varvec{\theta }^*}} \right\rangle }\\& =\sum \limits _{t = 1}^T {\Big[\frac{1}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| }^2}\Big]}\\& \quad -\sum \limits _{t = 1}^T {\Big[\frac{1}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_{t + 1}} - {\varvec{\theta }^*})} \right\| }^2}}\Big]\\& \quad +\sum \limits _{t = 1}^T {\Big[\frac{1}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}{\varvec{m}_t}} \right\| }^2}} \Big]\\& \quad +\sum \limits _{t = 1}^T {\Big[\frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}{\varvec{m}_{t - 1}}} \right\| }^2}} \Big]\\& \quad + \sum \limits _{t = 1}^T {\frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| }^2}}\\ \end{aligned} \end{aligned}$$

(A7)

Then according to Lemma 1, equation (A7) can be further rearranged as

$$\begin{aligned}{} & {} \begin{aligned} \sum \limits _{t = 1}^T {{f_t}({\varvec{\theta }_t})} - {f_t}({\varvec{\theta }^*})&\le \sum \limits _{t = 1}^T {\Big[\frac{1}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| }^2}\Big]}\\&\quad -\sum \limits _{t = 1}^T {\Big[\frac{1}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_{t + 1}} - {\varvec{\theta }^*})} \right\| }^2}} \Big]\\&\quad + \sum \limits _{t = 1}^T {\frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| }^2}} + \frac{{{R_\infty }G_2^2}}{{(1 - {\beta _1})}}\\ \end{aligned} \end{aligned}$$

(A8)

$$\begin{aligned}{} & {} \begin{aligned} \sum \limits _{t = 1}^T {{f_t}({\varvec{\theta }_t})} - {f_t}({\varvec{\theta }^*})& = \frac{1}{{2(1 - {\beta _1})}}\sum \limits _{t = 2}^T {{{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta } _t} - {\varvec{\theta } ^*})} \right\| }^2}} \\& \quad - \frac{1}{{2(1 - {\beta _1})}}\sum \limits _{t = 2}^T {{{\left\| {{\eta _{t - 1}}^{ - 1/2}({\varvec{\theta } _t} - {\varvec{\theta } ^*})} \right\| }^2}} \\& \quad + \frac{1}{{2(1 - {\beta _1})}} \Big[{\left\| {{\eta _t}^{ - 1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| ^2}\Big] + \frac{{{R_\infty }G_2^2}}{{(1 - {\beta _1})}} \\& \quad + \sum \limits _{t = 1}^T {\frac{{{\beta _{1t}}}}{{2(1 - {\beta _{1t}})}}{{\left\| {{\eta _t}^{1/2}({\varvec{\theta }_t} - {\varvec{\theta }^*})} \right\| }^2}} \\& =\frac{1}{{2(1 - {\beta _1})}} \Big[\sum \limits _{i = 1}^d {{\eta _{t,i}}^{ - 1}({\varvec{\theta }_{t,i}} - \varvec{\theta }_i^*)} \Big]\\& \quad + \frac{1}{{2(1 - {\beta _1})}} \Big[\sum \limits _{t = 2}^T {\sum \limits _{i = 1}^d {{\eta _{t,i}}^{ - 1}{{({\varvec{\theta }_{t,i}} - \varvec{\theta }_i^*)}^2}} } \Big]\\& \quad - \frac{1}{{2(1 - {\beta _1})}} \Big[\sum \limits _{t = 2}^T {\sum \limits _{i = 1}^d {{\eta _{t - 1,i}}^{ - 1}{{({\varvec{\theta }_{t,i}} - \varvec{\theta }_i^*)}^2}} } \Big]\\& \quad + \frac{1}{{2(1 - {\beta _1})}} \Big[\sum \limits _{t = 2}^T {\sum \limits _{i = 1}^d {{\beta _1}_t{\eta _{t,i}}^{ - 1}{{({\varvec{\theta }_{t,i}} - \varvec{\theta }_i^*)}^2}} \Big]} +\frac{{{R_\infty }G_2^2}}{{(1 - {\beta _1})}}\\& \le \frac{{D_\infty ^2}}{{2(1 - {\beta _1})}} \Big[\sum \limits _{i = 1}^d {{\eta _{t,i}}^{ - 1}} + \sum \limits _{t = 2}^T {\sum \limits _{i = 1}^d { \Big[{\eta _{t,i}}^{ - 1} - {\eta _{t - 1,i}}^{ - 1}} } \Big]\Big] \\& \quad + \frac{{D_\infty ^2}}{{2(1 - {\beta _1})}}\sum \limits _{t = 1}^T {\sum \limits _{i = 1}^d {{\beta _1}_t{\eta _{t,i}}^{ - 1}} } + \frac{{{R_\infty }G_2^2}}{{(1 - {\beta _1})}}\\& \le \frac{{D_\infty ^2\sqrt{T} }}{{2(1 - {\beta _1})}}\sum \limits _{i = 1}^d {{{{\hat{\eta }} }_{t,i}}^{ - 1}} + \frac{{{R_\infty }G_2^2}}{{(1 - {\beta _1})}} \\& \quad +\frac{{D_\infty ^2}}{{2(1 - {\beta _1})}}\sum \limits _{t = 1}^T {\sum \limits _{i = 1}^d {{\beta _1}_t{\eta _{t,i}}^{ - 1}} } \\ \end{aligned} \end{aligned}$$

(A9)

Equation (A8) uses inequality ${\beta _{1t}} \le {\beta _1} < 1$ to get equation (A9). According to the definition of ${\eta _t}$, we have $\eta _{t,i}^{ - 1} \ge \eta _{t - 1,i}^{ - 1}$. Using the ${D_\infty }$ bound on the feasible region, we get the equation (A9). From the above formula, the regret obtained by AdaXod enjoys an upper bound of O(T). The proof is finished. $\square$