Title: ADAX: ADAPTIVE GRADIENT DESCENT WITH EXPO-NENTIAL LONG TERM MEMORY

Abstract: Adaptive optimization algorithms such as RMSProp and Adam have fast convergence and smooth learning process. Despite their successes, they are proven to have non-convergence issue even in convex optimization problems as well as weak performance compared with the first order gradient methods such as stochastic gradient descent (SGD). Several other algorithms, for example AMSGrad and AdaShift, have been proposed to alleviate these issues but only minor effect has been observed. This paper further analyzes the performance of such algorithms in a non-convex setting by extending their non-convergence issue into a simple nonconvex case and show that Adam's design of update steps would possibly lead the algorithm to local minimums. To address the above problems, we propose a novel adaptive gradient descent algorithm, named AdaX, which accumulates the long-term past gradient information exponentially. We prove the convergence of AdaX in both convex and non-convex settings. Extensive experiments show that AdaX outperforms Adam in various tasks of computer vision and natural language processing and can catch up with SGD.

1. INTRODUCTION:

In the era of deep learning, Stochastic Gradient Descent (SGD), though proposed in the last century, remains the most effective algorithm in training deep neural networks (Robbins & Monro, 1951) . Many methods have been created to accelerate the training process and boost the performance of SGD, such as momentum (Polyak, 1964 ) and Nesterov's acceleration (Nesterov, 1983) . Recently, adaptive optimization methods have become popular as they adjust parameters' learning rates in different scales, instead of directly controlling the overall step size. These algorithms schedule the learning rates using a weighted average of the past gradients. For example, AdaGrad (Duchi et al., 2011) chooses the square root of the global average of the past gradients as the denominator of the adaptive learning rates. It is shown that when the gradients are sparse or small, AdaGrad can converge faster than vanilla SGD. However, its performance is quite limited in non-sparse settings.

Other adaptive algorithms have been proposed to replace the global average in AdaGrad by using the exponential moving average of past gradients, such as RMSProp (Tieleman & Hinton, 2012) , AdaDelta (Zeiler, 2012) , and Adam (Kingma & Ba, 2015) . Among all these variants, Adam (Kingma & Ba, 2015) is the most popular yet controversial optimization algorithm since it has faster convergence rate than the others. However, Adam has worse performance (i.e. generalization ability in testing stage) compared with SGD. Recent theories (Wilson et al., 2017; Reddi et al., 2018) have also shown that Adam suffers from non-convergence issue and weak generalization ability. For example, Reddi et al. (2018) proposed that Adam's non-convergence problem originate from a mistake in their proof of convergence. They constructed a counterexample for Adam and thoroughly proved that Adam did not guarantee convergence even in a simple convex setting.

In the meantime, Zhou et al. (2019) theoretically proved that Adam could make a large update when gradients are small, and a small update when gradients are large, which probably lead the optimization process to wrong directions. Shazeer & Stern (2018) also empirically showed that Adam's parameter updates are not stable and its second moment could be out of date. Luo et al. (2019) examined the effective learning rate of Adam in training and found that Adam would produce too large or too small extreme learning rates. All the above analyses have suggested that Adam's unstable design of adaptive learning rate may impair the optimization process.

To address the above issues, we propose a new adaptive optimization method, termed AdaX, which guarantees convergence in both convex and non-convex settings. This is done by memorizing the long-term square of gradients exponentially as the second moment. To achieve this, we first extend the convex counterexample in Reddi et al. (2018) to a non-convex setting. We then analyze how the second moment instability in Adam and RMSProp could cause the optimization process to converge to a local minimum even without noisy gradients, revealing that Adam produces update steps that are too large in optimization. We also show how AMSGrad (Reddi et al., 2018) is unable to solve Adam's problem completely, because its effectiveness relies heavily on the magnitude of the maximal second moment, for instance, too large second moment would lead to insufficient training and thus sub-optimal solutions. To address the above problems, we introduce a novel AdaX algorithm and theoretically prove that it converges with a similar speed to Adam, but gets rid of the second moment instability. Extensive experiments show that AdaX outperforms Adam and is comparable to SGD with momentum in many tasks of computer vision and natural language processing.

2. BACKGROUND AND NOTATION:

Overview of Adaptive Methods. To compare AdaX with the previous optimization methods, we follow (Reddi et al., 2018) to present a generic framework of adaptive algorithms as shown in Algorithm 1. Let S + d be a set of positive definite matrices in R d×d , and F be the parameter domain. In line 1, we take φ t : F → R d and ψ t : F → S d + as input, which are unspecified moment functions that vary among different optimization algorithms. After obtaining the gradient at time t in line 3, we can calculate the corresponding first and second moment m t , V t . In line 5, α t = α/ √ t is chosen as the step size in each iteration for the convergence analysis of adaptive algorithms. The projection operation in line 6, Π F ,M (y) is defined as argmin x∈F √ M (x -y) , where M ∈ S + d and y ∈ R d .

Algorithm 1 Generic Adaptive Optimization Algorithm 1: Input: x ∈ F, step size α, sequence of functions {φ t , ψ t } T t=1 2: for t = 1 to T do 3:

g t = ∇f t (x t ) 4:

m t = φ t (g 1 , g 2 , . . . , g t ) and V t = ψ t (g 1 , g 2 , . . . , g t )

5:

α t = α/ √ t 6: x t+1 = Π F , √ Vt (x t -α t m t / √ V t ) 7: end for

The main differences between the adaptive methods and the conventional SGD are in line 4 and 6, where the matrix V t scales the overall step size α t element-wisely by 1/ √ V t , known as the adaptive learning rate. Using the general framework in Algorithm 1, we are able to summarize many adaptive optimization algorithms proposed recently. For example, AdaGrad (Duchi et al., 2011) designed its V t as the global average of past gradients, while Adam (Kingma & Ba, 2015) and RMSProp (Tieleman & Hinton, 2012) chose the exponential moving average as follows instead.

EQUATION

where β 2 is the fixed second moment coefficient and g 2 i denotes the element-wise square of the gradients. The diagonal operation diag() perform the dimension transformation from R d to S d + . In order to improve the performance of Adam, Reddi et al. (2018) proposed AMSGrad, which took max operation on the second moment. Zhou et al. (2019) argued that we could replace g 2 t in V t with some past gradient squares g 2 t-n to temporarily remove the correlation between the first and second moment. Huang et al. (2019) constructed NosAdam, where a sequence of β 2t 's gave higher weights to past gradients. We provide a summary of different designs of adaptive learning rate in Table 1. It's noticeable that these algorithms, due to their exponential moving average design, still assign relatively high weights on recent gradients and past information is not emphasized. Besides, Loshchilov & Hutter (2019) noticed the difference between L 2 regularization and weight decay in adaptive algorithms and improved Adam with standard weight decay by proposing AdamW.

Optimization Framework. A commonly used framework for analyzing convex optimization algorithms was constructed by Zinkevich (2003) , named the online optimization problem. In this 

I diag( t i=1 g 2 i /t) (1 -β2)diag( t i=1 β t-i 2 g 2 i ) Adam AMSGrad AdaShift ψt ( 1-β 2 1-β t 2 )diag( t i=1 β t-i 2 g 2 i ) diag(max(vt-1, vt)) diag(β2vt-1 + (1 -β2)g 2 t-n ) NosAdam ... AdaX (ours) ψt diag(β2tvt-1 + (1 -β2t)g 2 t-n ) ... β 2 (1+β 2 ) t -1 diag( t i=1 (1 + β2) t-i g 2 i )

framework setting, the optimization algorithm chooses a parameter set θ t ∈ F and an unknown cost function f t (θ) evaluates its performance at θ t in each iteration. Suppose that there exists a best

parameter f t (θ * ) such that θ * = argmin θ∈F T t=1 f t (θ)

, then a metric used to show the algorithm's performance is the regret function R T = T t=1 f t (θ t ) -f t (θ * ) and we want to ensure that R T = o(T ) so that the algorithm will always converge to the optimal solution. (Zinkevich, 2003) .

Non-convergence of Adam. Reddi et al. (2018) proposed that the matrix Γ t defined as follows, was mistakenly assumed to be positive semi-definite in the original convergence proof of Adam.

EQUATION

where V t and α t are defined as in Algorithm 1. Based on such an observation, they constructed the following online convex optimization problem, in which Adam failed to converge to the optimal solution. Let C > 2 be a fixed constant and {f t } be the sequence of cost functions whose sum is to be minimized. Define f t as follows

EQUATION

where x ∈ F = [-1, 1]. In this problem, Adam could not distinguish between the true large gradient direction (C) and the noisy gradient directions (-1) because its √ V t scales the gradients to be of similar sizes, which forces the algorithm to reach a highly suboptimal solution x = 1 every three iterations. However, SGD and AdaGrad are both able to counteract the noisy gradients and converge to the optimum, which reveals the fact that Adam's design of adaptive learning rate is very unstable.

3. THE NONCONVERGENCE OF ADAM IN A NON-CONVEX SETTING:

In this section, we extend the non-convergence problem in (2) to the non-convex setting and explain why the fast convergence of Adam impairs its performance in the long term. Let C ∈ (1, +∞), λ ∈ (0, 1) be constants in R, consider the following simple sequence of non-convex functions f t .

f t (x) = Cλ t-1 x, for x ≥ 0 C 2 1-λ , for x < 0 ∀t ≥ 1 (3)

It can be easily observed that for the domain F = [-2, C/(1 -λ)], the minimum of each f t is obtained at x = 0. Suppose we start from x 0 > 0, then this problem simulates a situation where the gradient decreases exponentially as time increases, implying that the algorithm is approaching the global minimum and smaller step sizes are needed. Compared with the non-convergence problem proposed by Reddi et al. (2018) , no high-frequency noise exist in our gradients. However, next to the optimal solution, there is a local minimum trap where no gradients exist and thus no algorithm could escape. Let α 1 = α be the initial step size, we are able to show that SGD is capable of avoiding the trap, even without the learning rate decrease α t = α/ √ t, and will converge to x = 0 if initialized well. However, Adam ignores the gradient decrease information and always enters the trap regardless of initialization. We summarize the above results in the following lemma Lemma 3.1 In problem (3), with β 1 = 0, β 2 ∈ (0, λ 2 ) and α t ≥ α/t, Adam will always reach the local minimum, i.e.

T t=1 f (t)/T → C 2 1-λ , ∀x 1 , α 1 > 0.

We provide all the proofs in the Appendix. This parameter setting of Adam is the same as RMSProp except for the bias correction term, and the condition β 1 < √ β 2 mentioned by Kingma & Ba (2015) is automatically satisfied. α t ≥ α/t is a weak requirement for the step sizes and can be ensured with constant step sizes or α t = α/ √ t as in the convergence analysis. Intuitively, Adam would scale the decreasing gradient by 1/ √ V t , which approximately increases with the same speed. Therefore, its update steps would be larger than a fixed constant at any time step and would ultimately lead the parameters to the trap regardless of initialization. People may wonder whether the first moment design helps Adam in such a situation, but we can show that as long as the condition β 1 < √ β 2 is satisfied, Adam would always reach the local minimum. Hence, Adam converges faster than SGD because of its large updates, but it cannot slow down when approaching the global minimum.

AMSGrad is constructed to address the problem of Adam's large steps during optimization. However, it suffers from two major issues in a non-convex situation. 1) The never decreasing V t in AMSGrad could lead to early stops and therefore insufficient training during optimization, as revealed by Huang et al. (2019) ; 2) The time for achieving the maximum of V t is uncontrollable. We show that for certain cases in our problem, AMSGrad is unable to help Adam. Lemma 3.2 In problem (3), with β 1 = 0 and α t ≥ α/t, ∀β 2 ∈ (0, 1), ∃λ ∈ ( √ β 2 , 1), such that AMSGrad will always reach the local minimum, i.e.

T t=1 f (t)/T → C 2 1-λ , ∀x 1 , α 1 > 0.

The lemma essentially states for any fixed β 2 , we can find a λ such that AMSGrad cannot help Adam because V t keeps increasing before stepping into the trap. Therefore, we still need an algorithm that can generate stable learning rates and control the update steps effectively. We emphasize that although the functions in (3) are not smooth, the problem does successfully provide some intuition on why Adam variants trains much faster than SGD, but cannot have comparable testing performance.

4. ALGORITHM AND CONVERGENCE:

Next, we introduce our novel optimization algorithm and present its special way of adjusting the adaptive learning rate. Based upon the above discussions that current gradients lead to unstable second moment and that long-term memory algorithms are preferred, we design our algorithm AdaX by weighting exponentially more on the history gradients and less on the current gradients, as shown in Algorithm 2. The most important differences between AdaX and Adam are in line 6 and 7, where instead of an exponential moving average, we change β 2 to 1+β 2 and accumulate the past gradients. Such a design guarantees that noisy and extreme gradients cannot greatly influence the update steps, and vt would gradually become stable. Similar to Kingma & Ba (2015)'s derivation, in order to

Algorithm 2 AdaX Algorithm 1: Input: x ∈ F, step size {α t } T t=1 , β 1 , β 2 , β 3 2: Initialize m 0 = 0, v 0 = 0 3: for t = 1 to T do 4: g t = ∇f t (x t ) 5: m t = β 1 m t-1 + (1 -β 3 )g t 6: v t = (1 + β 2 )v t-1 + β 2 g 2 t 7: vt = v t /[(1 + β 2 ) t -1] and V t = diag(v t ) 8: x t+1 = Π F , √ Vt (x t -α t m t / √ vt )

9: end for achieve an unbiased estimate of second moment, we obtain our bias correction term as follows. Let g t be the gradient at timestep t and further suppose g t 's are drawn from a stationary distribution g t ∼ p(g t ). Take expectation on both sides of line 6 in Algorithm 2, we get

EQUATION

To maintain an accurate second moment, we would naturally divide v t by (1 + β 2 ) t -1 in line 7. However, it's worth mentioning that we do not include a first moment correction term (1 -β t 1 ) as Kingma & Ba (2015) did for the following reason. Consider the Stochastic Gradient Descent with momentum(SGDM) algorithm and Adam's first moment, SGDM:

m t = γm t-1 + g t = t i=1 γ t-i g i Adam: m t = β 1 m t-1 + (1 -β 1 )g t = (1 -β 1 ) t i=1 β t-i 1 g i

It can be observed that they have the same form except for the scaling constant 1 -β 1 , and therefore the first order bias correction term is counter intuitive. We change the scaling constant from 1 -β 1 to 1 -β 3 in line 5 in our algorithm to obtain a more general form of the moment expression, and when β 3 = β 1 , it helps to scale the hyper-parameters (such as step size, weight decay) of adaptive algorithms to the same size as SGD. Next we show that our algorithm ensures the positive semi-definiteness of Γ t and hence does not have the non-convergence issue of Adam. Consider the following lemma which leads to the conclusion that Γ t is positive semi-definite in our algortihm, Lemma 4.1 Algorithm 2 ensures that the matrix Vt

EQUATION

Finally, we provide the convergence analysis of our algorithm in both convex and non-convex settings. Using the analysis framework by Zinkevich (2003) , the following theorem states that we are able to obtain a regret bound of O( √ T ), which is the same as the results of Reddi et al. (2018) . A domain

F is said to have bounded diameter if x -y ∞ ≤ D ∞ , ∀x, y ∈ F for some D ∞ ∈ R.

Theorem 4.1 Let {x t } and {v t } be the sequences obtained from Algorithm 2,

α t = α/ √ t, β 1,1 = β 1 , β 1,t ≤ β 1 , for all t ∈ [T ] and β 2t = β 2 /t, β 3t = 1-1/ √ t.

Assume that F has bounded diameter D ∞ and ∇f t (x) ≤ G ∞ for all t ∈ [T ] and x ∈ F. Then for x t generated using Algorithm 2, we have the following bound on the regret.

R T ≤ D 2 ∞ 2α T (1 -β 1 ) d i=1 v1/2 T,i + D 2 ∞ 2(1 -β 1 ) T t=1 d i=1 β 1t v1/2 t,i α t + αC √ 1 + log T (1 -β 1 ) 3 √ β 2 d i=1 g 1:T,i 2 (5)

The following corollary follows naturally from the above theorem.

Corollary 4.1 Suppose

EQUATION

Analyzing optimization algorithms in a non-convex setting is slightly different from the convex case, where instead of the average regret, stationarity in gradient is utilized to show convergence in time. Following Chen et al. (2019) , suppose we are minimizing a cost function f that satisfies the following three assumptions A1. f is differentiable and has L-Lipschitz gradient, i.e. ∀x, y, ∇f (x) -∇f (y) ≤ L x -y . and f (x * ) > ∞ where x * is the optimal solution. A2. Suppose that the true gradients and the noisy gradients are bounded i.e.

∇f (x t ) ≤ G ∞ , g t ≤ G ∞ , ∀t ≥ 1. Also, α t mt √ vt ≤ G for some G > 0 A3.

The noisy gradient is unbiased and the noise is independent, i.e.

g t = ∇f (x t ) + η t , E[η t ] = 0 and η i is independent of η j if i = j.

then we would obtain the following theorem and corollary, which prove that AdaX converges with a speed close to AMSGrad as mentioned by Chen et al. (2019) .

Theorem 4.2 Let {x t } and {v t } be the sequences obtained from Algorithm 2,

α t = α/ √ t, β 1,1 = β 1 , β 1,t ≤ β 1 , for all t ∈ [T ] and β 2t = β 2 /t, β 3t = β 1 . Assume that ∇f t (x) ≤ G ∞ for all t ∈ [

T ] and x ∈ F. Then for x t generated using Algorithm 2, we have the following bound.

EQUATION

where C 1 , C 2 , C 3 , C 4 are constants independent of T Corollary 4.2 Suppose β 2t = β 2 , with the other assumptions same as in Theorem 4.2, we have

min t∈[T ] E ∇f (x t ) 2 ≤ G ∞ α √ T ( C 1 G 2 ∞ α 2 c 2 + C 2 dα c + C 3 dα 2 c 2 + C 4 ) (8)

5. EXPERIMENTS:

In this section, we evaluate the performance of AdaX on various tasks in comparison with SGD with Nesterov momentum (SGDM), Adam(W), and many other common optimizers. The implementation of AdaX consists of two parts, AdaX and AdaX-W, representing using L 2 regularization and standard weight decay in the algorithm respectively as discussed in Loshchilov & Hutter (2019) . We relegate the detailed implementation of AdaX to the Appendix. We show that AdaX, combined with a proper weight decay, is capable of performing better than Adam and SGDM in many tasks. L 2 Regularization. We used L 2 regularization of strength 5e-4 for all the optimizers. As shown in Figure 1a and 1b, although AdaX was relatively slow at the beginning compared with other optimizers, its testing accuracy quickly caught up with the others after the first learning rate decrease and became the highest (91.6) at last.

Weight Decay. The baseline was trained with SGDM with weight decay 5e-4. Adam with weight decay, named AdamW (Loshchilov & Hutter, 2019) was also trained for comparisons. Although AdamW with 1e-4 weight decay converged much faster than SGDM and AdaX-W, its final accuracy could not catch up with the other two (see Figure 1c & 1d ).A higher weight decay could potentially help AdamW achieve a better performance (92.1), but it was as slow as SGDM. On the other hand, AdaX-W with step size rate 0.5 and 5e-4 weight decay converged fast and yielded the same performance as SGDM(92.32). The best result we obtained was AdaX-W with step size 0.25 and 1e-3 weight decay, which resulted in 92.51 Top-1 accuracy, even slightly higher than SGDM.

5.2. CONVOLUTIONAL NEURAL NETWORK ON IMAGENET:

We also conducted experiments to examine the performance of AdaX-W on ImageNet (Deng et al., 2009) . SGDM, AdaX-W, and AdamW were used to train a ResNet-18 model on ImageNet, with a standard 1e-4 weight decay and 0.1, 0.5, 1e-3 step sizes as in CIFAR-10 respectively. A warm up scheme was applied in the initial 25k iterations (Goyal et al., 2017) , and then the step size was multiplied by 0.1 at the 150k, 300k and 450k-th iteration steps. As observed from Figure 2 , although AdamW was fast at the beginning of training, its test accuracy stagnated after the second learning rate decrease. AdaX-W, on the other hand, converged faster than SGDM without loss of testing accuracy (69.87), as shown in Table 3a . AdaX has also been validated on One Billion Word (Chelba et al., 2013) dataset of language modeling task. For the One Billion Word, we used a two-layer LSTMs with 2048 hidden states and sampled softmax. The global experiment settings in the released public code Rdspring1 was adopted in this study. For both vanilla Adam and AdaX, the LSTMs were trained using for 5 epochs, with learning rate decaying to 1e-8 linearly. Similarly, the weight decay for AdaX was set to 0. Note that the regular SGD is not suitable in this task, so it is not included in the comparison.

The training loss and the validation perplexity is shown in Figure 3a and Table 3b . We can see that the AdaX outperforms the Adam baseline by a significant margin (35.22 vs. 36.90) . Moreover, similar to the effect on image classification tasks described above, AdaX starts a little slower at the early stage, but it soon surpasses Adam on both training and validation performance.

5.4. TRANSFER LEARNING:

To further examine the robustness of AdaX in transfer learnings such as semantic segmentation, we evaluated its performance on the PASCAL VOC2012 augmented dataset (Everingham et al., 2014) (Hariharan et al., 2011) . The classic Deeplab-ASPP model (Chen et al., 2016 ) was adopted with a ResNet-101 backbone pretrained on MS-COCO dataset (Lin et al., 2014) . Adaptive methods are seldom used in semantic segmentation tasks because of their bad performances, but AdaX can surprisingly be applied to train these models as well. The initial step sizes for different algortithms was recorded in section A.7 and all other parameters stuck to the default setting in Chen et al. (2016) . We evaluated the algorithms' performances at the 5k, 10k, 15k and 20k iterations using intersection over union (IoU). As can be observed in Figure 3c , AdaX-W trained faster than SGDM and obtained a higher IoU (76.5) at the same time. On the other hand, AdamW was not capable of obtaining comparable results.

5.5. COMPARISON OF SECOND MOMENT DESIGN:

Finally, we compared the second moment design of Adam and AdaX empirically and showed the instability of Adam's second moment could have a large impact on its performance while our design was stable and robust. We first evaluated the performance of different algorithms in our synthetic example (3). The problem parameters were set to be C = 1e -3, λ = 0.9999, x 0 = 1. To ensure fair comparisons, default hyperparameters were chosen for all the algorithms, specifically α 0 = 0.1, γ = 0.9 for SGDM, α 0 = 1e -3, β 1 = 0.9, β 2 = 0.999 for Adam, and α 0 = 0.15, β 1 = 0.9, β 2 = 1e -4, β 3 = 0.999 for AdaX. As shown in Figure 4a and 4b, SGD and AdaX quickly converged under the strong gradient decrease information and could potentially reach the global minimum if initialized well. However, the update steps of Adam decreased with a much slower rate, which resulted in substantial changes in x and ultimately lead the algorithm to the local minimum. Zaheer et al. (2018) found that Adam's performance could be affected by different values of , which was originally designed to avoid zeros in the denominator. Such a finding revealed the instability ∈ {1e -8, 1e -5, 1e -3} on the CIFAR-10 dataset. We found that larger values of , since it helped to stabilize very small second moment, did improve AdamW's performance by around 0.35 percent accuracy on CIFAR-10. However, AdaX-W's performance was not affected by different choices of as shown in Figure 4d , which again proved our claim that a long term memory design was more stable than the exponential moving average in Adam.

The experiments shown above verify the effectiveness of AdaX, showing that the accumulated longterm past gradient information can enhance the model performance, by getting rid of the second moment instability in vanilla Adam. It is also worth noticing that the computational cost for each step of AdaX and Adam are approximately the same, as they both memorize the first and second momentum in the past. Therefore AdaX enables one to get higher performance than Adam in various tasks using the same training budget.

6. CONCLUSION:

In this paper, we present a novel optimization algorithm named AdaX in order to improve the performance of traditional adaptive methods. We first extend the non-convergence issue of Adam to a non-convex case, and show that Adam's fast convergence impairs its convergence. We then propose our variant of Adam, analyze its convergence, and evaluate its performance on various learning tasks. Our theoretical analysis and experimental results both show that AdaX is more stable and performs better than Adam in various tasks. In the future, more experiments still need to be performed to evaluate the overall performance of AdaX and AdaX-W. Moreover, our paper is a first step into designing adaptive learning rates in ways different from simple and exponential average methods. Other new and interesting designs should also be examined. We believe that new adaptive algorithms that outperform AdaX in convergence rate and performance still exist and remain to explore.

A APPENDIX:

A.1 PROOFS OF LEMMA 3.1, LEMMA 3.2

We consider a one dimensional non-convex case where {f t } are a sequence of linear functions that have decreasing gradients in the long term. We want to show that because Adam trusts its current gradient as the second moment, its step sizes are too large and the algorithm would converge to a suboptimal solution. Let constant C be the initial gradient, define cost function f t as follows:

f t (x) =    Cλ t-1 x, for x ≥ 0 C 2 1 -λ , for x < 0 ∀t ≥ 1 (9)

where λ is the decreasing factor of gradient. Consider F = [-2, C/(1 -λ)], then it's obvious that the minimum regret is obtained at x = 0. Let the initial step size be α 1 = α, we then consider the performances of different algorithms in this setting.

(SGD). We first show that without the momentum, vanilla SGD is able to converge to the optimum solution and avoid the trap. Take derivative with respect to x, we obtain that

∇f t (x) = Cλ t-1 , for x ≥ 0 ∞ t=1 ∇f t (x) = ∞ t=1 Cλ t-1 = C 1 -λ (10)

Therefore, even if we set α t = α, ∀t ≥ 1, as long as the initial point x 0 ≥ αC 1-λ , SGD is able to avoid the trap. If α t = α/ √ t, then the condition can be even less strict:

x 0 ≥ ∞ t=1 αCλ t-1 √ t

. SGD is able to converge to the optimum if the equal signs are true.

(Adam). We consider the Adam algorithm with the following parameter setting:

β 1 = 0, 0 < β 2 < λ < 1, and α t = α √ t (11)

Note that this parameter setting of Adam is the same as RMSProp, but we can further show that even if β 1 = 0, we still obtain similar results. Consider how v t changes in time, before it reaches the negative region, the gradients are positive and

v t = β 2 v t-1 + (1 -β 2 )(Cλ t-1 ) 2 = t i=1 β t-i 2 (1 -β 2 )(Cλ i-1 ) 2 = (1 -β 2 )C 2 λ 2 β t 2 t i=1 ( λ 2 β 2 ) i = (1 -β 2 )C 2 (λ 2t -β t 2 ) λ 2 -β 2 (12)

Note that λ > √ β 2 , therefore the update rule is:

x t+1 = x t -α t g t √ v t = x t -α t λ 2 -β 2 λ t-1 (1 -β 2 )(λ 2t -β t 2 ) = x t -α t λ 2 -β 2 λ -1 (1 -β 2 )(1 -( β2 λ 2 ) t ) ≤ x t - α √ t λ 2 -β 2 λ 2 (1 -β 2 ) (13)

Note that the series ∞ t=1 1 √ t diverges, hence Adam would always reach the negative region. Same argument applies as long as α t ≥ α/t. We would emphasize here that the bias correction term in Adam does not change the final result as 1 -β t 2 ≥ 1 -β 2 and therefore the update steps are still bounded. We could further show that when β 1 = 0, β 1 < √ β 2 , Adam will still go to the negative region. Since

m t = β 1 m t-1 + (1 -β 1 )g t = t i=1 (1 -β 1 )β t-i 1 g i , therefore m t √ v t = (1 -β 1 )(λ t -β t 1 ) λ -β 1 • λ 2 -β 2 (1 -β 2 )(λ 2t -β t 2 ) = (1 -β 1 ) λ 2 -β 2 (λ -β 1 ) (1 -β 2 ) • λ t -β t 1 λ 2t -β t 2 = (1 -β 1 ) λ 2 -β 2 (λ -β 1 ) √ 1 -β 2 • 1 -(β 1 /λ) t 1 -(β 2 /λ 2 ) t ≥ (1 -β 1 ) λ 2 -β 2 (λ -β 1 ) √ 1 -β 2 • (1 - β 1 λ ) (14)

Since the update steps are lower bounded, the algorithm would still go to the negative region.

(AMSGrad). We now evaluate the performance of AMSGrad in our formulated problem. Note that v t in AMSGrad would take the same form as Adam, and vt = max{v t } t i=1 . We suppose that the maximum is obtained at v 1 as an example, then

EQUATION

As we can see, AMSGrad partially solves the problem of Adam and restores the gradient decrease information as its v t is lower bounded. If the maximum of v t is obtained before the parameters enter the trap, AMSGrad could possibly have a better performance in this problem as it prevents the update steps from being too large. However, one important determining factor is the time when the maximum value is obtained. If v t in fact keeps increasing before a very large number T , then AMSGrad would have the same performance as Adam. We explain the above intuition as follows.

Let

h(t) = λ 2t -β t 2 , then dh(t) dt = ln λ 2 • λ 2t -ln β 2 • β t 2 (16) If dh(t) dt ≥ 0, we have ln λ 2 • λ 2t -ln β 2 • β t 2 ≥ 0 ( λ 2 β 2 ) t ≤ ln β 2 ln λ 2

When the equal sign is true, we have

t = log λ 2 β 2 ln β 2 ln λ 2 = ln ln β2 ln λ 2 ln λ 2 β2 lim β2→λ 2 t = lim β2→λ 2 1 β2 ln(β2) -1 β2 = - 1 ln(λ 2 ) lim λ→1 -t = ∞ (17)

The first equal sign in the first limit is due to L'Hospital's rule. Therefore, the value of T where v T = max{v t } depends on the difference between β 2 and λ 2 , and the value of λ. If β 2 is close to λ 2 or λ is close to 1, then v t needs a large number of steps to obtain the maximum. In such cases, AMSGrad may not able to help Adam. Specifically, for a fixed β 2 , since lim λ→1 -t = ∞. and

g t / √ v t = λ 2 -β 2 λ 2 (1 -β 2 ) = 1 √ 1 -β 2 1 - β 2 λ 2 (18)

we know that larger λ 2 will lead to both larger update steps and larger T when the maximum is obtained, hence ∃λ ∈ (0, 1), such that AMSGrad cannot avoid the trap.

(AdaX). We provide the performance of AdaX in this problem for completeness. We only show for the case when β 1 = β 3 = 0, but the same results hold when the first order momentum is used.

EQUATION

As we can see, AdaX successfully restores the gradient decrease information and controls the decrease speed by an almost fixed parameter instead of trusting the gradient at a certain moment, and is therefore expected to perform better than AMSGrad since its step sizes are not affected by extreme gradients. With a suitable initial step size and starting point, AdaX is able to avoid the designed trap.

A.2 PROOF OF LEMMA 4.1

EQUATION

where in the first inequality we utilize the fact that

(1 + β 2 ) t ≥ 1 + tβ 2 and hence t (1+β2) t -1 ≥ t-1 (1+β2) t -(1+β2)

. Intuitively, it is easier to see this inequality if we simply let β 2 to be a small number such as 1e-4 in our implementation, then the denominator doesn't change much while the numerator decreases.

A.3 AUXILLARY LEMMAS FOR CONVERGENCE ANALYSIS:

Lemma A.1 Assume that β 21 = β 2 , β 2t = β 2 /t, with β 2 ∈ (0, 1) and vt = [(1 + β 2t )v t-1 + β 2t g 2 t ]/[(1 + β 2t ) t -1], V t = diag(v t ), then we have Vt α 2 t Vt-1 α 2 t-1

Proof: Similar to Lemma 4.1 in the algorithm section, we have

EQUATION

The first inequality comes from deleting the last term β2 t g 2 t and second one comes from the following fact:

EQUATION

Therefore the positive semi-definiteness is satisfied.

Lemma A.2 For the parameter settings and conditions assumed in Theorem 4.1, we have

EQUATION

Proof: We first analyze with the following process directly from the update rules, note that

EQUATION

T t=1

EQUATION

where the first inequality is due to an application of Cauchy-Schwarz inequality. The second inequality is due to the fact that β 1t ≤ β 1 , ∀t. The third inequality follows from

T j=1 β T -j 1 ≤ 1/(1 -β 1

) and the fact that 1 -β 3j ≤ 1. The fourth one comes from only keeping one of the positive terms in the denominator. The final one is from the fact that

(1 + β 2T ) T -1 ≤ C for a constant C = √ e β2 -1

. By using induction on all the terms in the equation, we are able to further bound it.

T t=1

EQUATION

The second equality follows from a re-arrange of sum order. The third inequality comes from the fact that

T j=t β j-t 1 ≤ 1/(1-β 1 ) ≤ 1/(1-β 1 )

. The fourth inequality is again from the Cauchy-Schwarz inequality and the final one is from the harmonic sum T t=1 ≤ (1 + log T ). By denoting g 1:t,i to be the vector of the past graidents from time 1 to t in the i-th dimension, i.e. g 1:t,i = [g 1,i , g 2,i , ..., g t,i ],

we complete the proof of the lemma.

Lemma A.3 For the parameter settings and conditions assumed in Theorem 4.1, we have

EQUATION

Proof: Using the definition of L2 norm, by Lemma A.1, since

EQUATION

where the first inequality is from separating the first term and getting rid of the last negative term in the summation. The last inequality is from a telescopic summation and the diameter bound that 

EQUATION

Using Lemma 4 in Reddi et al. (2018) proved by Mcmahan & Streeter (2010) with a direct substitute of

z 1 = (x t -α t V -1/2 t m t ), Q = V 1/2

and z 2 = x * for x * ∈ F, the following inequality holds:

V 1/4 t (u 1 -u 2 ) 2 = V 1/4 t (x t+1 -x * ) 2 ≤ V 1/4 t (x t -α t V -1/2 t m t -x * ) 2 = V 1/4 t (x t -x * ) 2 + α 2 t V -1/4 t m t 2 -2α t m t , (x t -x * ) = V 1/4 t (x t -x * ) 2 + α 2 t V -1/4 t m t 2 -2α t β 1t m t-1 + (1 -β 1t )g t , (x t -x * ) (31)

where the first equality is due to the fact that Π F ,

√:

Vt (x * ) = x * . Rearrange the last inequality, we obtain

EQUATION

The second inequality comes from applications of Cauchy-Schwarz and Young's inequality. We now make use of the approach of bounding the regret using convexify of f t as in Kingma & Ba (2015) . Following Lemma A.2 and Lemma A.3, we have

EQUATION

A.4.2 PROOF OF COROLLARY 4.1

Proof. We first take a look at the size of v1/2 t,i , note that ∇f

t (θ) ∞ ≤ G ∞ vt,i = 1 (1 + β 2t ) t -1 t j=1 β 2j Π t-j k=1 (1 + β 2(t-k+1) )g 2 j,i ≤ G 2 ∞ β 2 t j=1 β 2 j Π t-j k=1 (1 + β 2 t -k + 1 ) ≤ G 2 ∞ t j=1 1 j Π t-j k=1 (1 + 1 t -k + 1 ) = G 2 ∞ t j=1 t + 1 j(j + 1) = tG 2 ∞ (34)

The first inequality is due to the fact that (1 + β 2t ) t ≥ (1 + β 2 ) and the gradient bound. The second inequality follows from β 2 < 1. The last inequality is from the telescopic sum. Then we have the following inequality,

T t=1 d i=1 β 1t v1/2 t,i α t ≤ dG ∞ β 1 α T t=1 λ t-1 t ≤ dG ∞ β 1 α(1 -λ) 2 (35)

The second inequality is due to the arithmetic geometric series sum

EQUATION

Therefore we have the following regret bound

EQUATION

A.5 PROOF OF NON-CONVEX CONVERGENCE A.5.1 PROOF OF THEOREM 4.2

Proof. We first directly refer to the original paper and obtain the following bound (Chen et al., 2019 ).

EQUATION

where C 1 , C 2 , C 3 are constants independent of d and T , C 4 is a constant independent of T . For the first term, assume that min j∈

EQUATION

where the first inequality follows from Lemma A.1 as vt

α 2 t ≥ vt-1 α 2 t

. The second inequality is from the gradient bound ∇f (x t ) ≤ G ∞ . For the second term with C 2 , similarly by the positive semidefiniteness in Lemma A.1, we have

EQUATION

The second equality is from the telescope sum and for the third term

E C 3 T -1 t=2 α t √ vt - α t-1 vt-1 2 ≤ E C 3 T -1 t=2 α c α t √ vt - α t-1 vt-1 1 ≤ C 3 dα 2 c 2 (42)

where the first inequality is because

| αt √ vt -αt-1 √ vt-1 | ≤ αt-1 √ vt-1 ≤ α c

and the last one is due to the previous inequality with second term. Hence in summary, we have

EQUATION

Now by the proof of Corollary 4.1, we have an upper bound ( vt ) j ≤ tG 2 ∞ , therefore we can bound the effective step sizes,

EQUATION

And thus we have

EQUATION

The last inequality is due to the fact that

T t=1 1/t ≥ T +1

t=1 1/tdt = log(1 + T ), therefore we have

EQUATION

We would emphasize that the assumption α t m t / √ vt ≤ G in the theorem is automatically satisfied as αt

√ vt ≤ α1 √ v1 = α c . Hence α t m t / √ vt ≤ αG∞ c . A.5.2 PROOF OF COROLLARY 4.2

Similar to the proof of Theorem 4.2, we still have

EQUATION

Note that (v t ) j has the following upper bound as ∇f

EQUATION

)

And thus we have

E T t=1 α t ∇f (x t ), ∇f (x t )/ vt ≥ E T t=1 α √ tG ∞ ∇f (x t ) 2 ≥ α G ∞ min t∈[T ] E ∇f (x t ) 2 T t=1 1 √ t ≥ α G ∞ min t∈[T ] E ∇f (x t ) 2 √ T (49)

where the last inequality is by the fact that T t=1

1 √ t ≥ √

T , therefore we have

EQUATION

A.6 IMPLEMENTATION OF ADAX AND ADAX-W

The detailed implementations of AdaX and AdaX-W are as follows. The performance of AdaX is robust with respect to the value of β 2 , but we recommend smaller values such as 1e -4, 1e -5 to reduce computational cost. The value of β 3 can be set as 0.9 to match the first moment design of Adam, consequently we need 1 100 step sizes. The size of weight decay should also be adjusted respectively.

Algorithm 3 AdaX Algorithm with L 2 Regularization and Weight Decay 1: Input: x ∈ F, step size {α t } T t=1 , (β 1 , β 2 , β 3 ) = (0.9, 1e -4, 0.999), weight decay λ = 5e -4, = 1e -12 2: Initialize m 0 = 0, v 0 = 0 3: for t = 1 to T do 4: g t = ∇f t (x t ) + λx t 5: m t = β 1 m t-1 + (1 -β 3 )g t 6:

v t = (1 + β 2 )v t-1 + β 2 g 2 t 7: vt = v t /[(1 + β 2 ) t -1] and V t = diag(v t ) 8:

x t+1 = Π F , √ Vt (x t -α t (m t /( √ vt + ) + λx t ))

9: end for

A.7 HYPER-PARAMETER TUNING IN THE EXPERIMENTS:

The hyperparameters in different algorithms have a huge impact on their performances in the experiments. To efficiently find the best step sizes, we follow Wilson et al. (2017) to perform a logarithmically-spaced grid search of the optimal step sizes and we list the step sizes we tried in the following tables, where the step sizes in bold are the ones with best performances and used in the experiments section.

Step size: Image Classification

• SGD(M) {10, 1, 1e-1, 1e-2, 1e-3}

• Adam/AdamW {1e-2, 3e-3, 1e-3, 3e-4, 1e-4}

• AdaGrad {5e-2, 1e-2, 5e-3, 1e-3, 5e-4}

• AMSGrad {1e-2, 3e-3, 1e-3, 3e-4, 1e-4}

• RMSProp {1e-2, 5e-3,1e-3, 5e-4, 1e-4}

• AdaX(ours) {5e-1, 2.5e-1, 2e-1, 1.5e-1, 1e-1, 5e-2 1e-2, 5e-3, 1e-3}

• AdaX-W(ours) {1,5e-1, 4e-1, 3e-1, 2.5e-1(CIFAR), 1e-1, 1e-2, 5e-3, 1e-3}

Step size: VOC2012 Segmentation

• SGD(M) {1e-3, 5e-4, 2.5e-4, 1e-4, 5e-5}

• AdamW {5e-4, 1e-5, 5e-5, 1e-6, 5e-7}

• AMSGrad {5e-4, 1e-5, 5e-5, 1e-6, 5e-7}

• RMSProp {1e-4, 5e-5, 1e-5, 5e-6, 1e-6}

• AdaX-W(ours) {1e-2, 5e-3, 1e-3, 5e-4, 1e-4}

Step size: Billionwords

• SGD {30, 20, 10, 5, 1, 1e-1, 1e-2 }

• Adam {5e-3, 2e-3, 1e-3, 5e-4, 1e-4 }

• AdaX(ours) {5e-3, 2e-3, 1e-3, 5e-4, 1e-4} × {0.5, 1, 5, 15, 25, 50, 100}

Momentum parameters. For the momentum parameters of Adam, RMSProp and AMSGrad, we tuned over (β 1 , β 2 ) = {(0.9, 0.999), (0.99, 0.999), (0.99, 0.9999)} and found that the default values (0.9, 0.999) as in Kingma & Ba (2015) yield the best result. For the momentum parameters (β 1 , β 2 , β 3 ) in AdaX, we directly applied β 1 = 0.9 as in Adam and we tuned β 2 over {1e -3, 1e -4, 1e -5}. We found that the value of β 2 didn't affect the general performance of AdaX. Tuning β 3 is the same as tunning the step sizes, as we will show in the next subsection.

For the other hyperparameters such as weight decay and batch size, we directly apply the same settings as in the baselines (He et al., 2016) (Chen et al., 2016 )(Rdspring1) (Loshchilov & Hutter, 2019) .

A.8 THE EFFECT OF β 3

As we mentioned in section 4 and 5, adding a new hyper-parameter β 3 does not influence the first momentum sum of AdaX as it is just a constant common factor that can be extracted out as follows.

m t = β 1 m t-1 + (1 -β 3 )g t = (1 -β 3 ) t i=1 β t-i 1 g i (51)

Therefore, tuning the hyperparameter β 3 will have exactly the same effect as tuning the step size α.

As long as (1 -β 3 )α remains the same, the performance of AdaX doesn't change. We also show this fact empirically using different combinations of (α, β 3 ) on the CIFAR-10 dataset. As shown above, we trained the ResNet-20 model (He et al., 2016) on the CIFAR-10 dataset with different combinations of α and β 3 from {(0.5, 0.999), (0.05, 0.99), (0.005, 0.9)}. However, the

:



Figures and Tables:

Figure 1: Training and Testing Accuracy on CIFAR-10

5.1 CONVOLUTIONAL NEURAL NETWORK ON CIFAR-10Using ResNet-20 proposed byHe et al. (2016), we verified the performance of AdaX on CIFAR-10 (Krizhevsky et al., 2009) image classification task. In our experiments, we utilized a learning rate Testing Top-1 Accuracy on ImageNet

Figure 2: Training and Testing Results on ImageNet.



Figure 3: (a) Traning loss curves for Adam and AdaX on One Billion Word. (b, c) Training Loss and Testing Results on VOC2012 Segmentation task. In (c), dashed lines are mean accuracy values and solid lines are Intersection over Union (IoU) values

Figure 4: Second momentum stability comparisons. (a),(b). Training updates and x-positions on Problem (3). (c),(d) Training and testing results of AdamW and AdaX-W with different 's

x -x * ≤ D ∞ A.4 PROOF OF REGRET BOUND A.4.1 PROOF OF THEOREM 4.1 Proof. Following the proof given by Reddi et al. (2018), we provide the proof of regret bound in Theorem 4.1. Beginning with the definition of the projection operation Π F , √ Vt , we have the observation

Figure 5: Training and Testing Results on CIFAR-10 with Different combinations of (α, β 3 ).

Comparisons of different designs of the second moment

Table content: <table><tr><td>SGDM</td><td>AdaGrad</td><td>RMSProp</td></tr><tr><td>ψt</td><td/><td/></tr></table>

Validation accuracy on CIFAR-10.

Table content: <table><tr><td/><td>Method</td><td>Top 1 Acc</td><td>Method</td><td>Top 1 Acc</td></tr><tr><td/><td>Adam</td><td>90.96 ± 0.03</td><td>AdamW, wd=1e-4</td><td>91.86± 0.04</td></tr><tr><td>-</td><td>AdaGrad AMSGrad</td><td>86.56 ± 0.07 90.98 ± 0.04</td><td>AdamW, wd=5e-4 SGDM</td><td>92.12 ± 0.05 92.30 ± 0.09</td></tr><tr><td/><td>RMSProp</td><td colspan="3">89.64 ± 0.05 AdaX-W, wd=5e-4(ours) 92.32 ± 0.04</td></tr><tr><td/><td colspan="4">AdaX(ours) 91.60 ± 0.10 AdaX-W, wd=1e-3(ours) 92.51 ± 0.07</td></tr></table>



Table content: <table/>

Peer review (iclr2020):
In this paper, the authors propose a new adaptive gradient algorithm AdaX, which the authors claim results in better convergence and generalization properties compared to previous adaptive gradient methods. The paper overall is fairly clear, although the writing can be improved in places. Section 3 is interesting, although calling the section the *nonconvergence of Adam* is a bit misleading, since the algorithm does converge to a local minimum. I have some concerns about the rest of the paper however. I am a bit confused about the changes proposed to Adam that gives rise to the AdaX algorithm. 1. Won*t replacing beta2 with 1+beta2 keep increasing the magnitude of the denominator of the algorithm just like AdaGrad does? In that case, is the algorithm expected to work better when having sparse gradients? 2. I also do not quite understand how to interpret using a separate hyperparameter for the momentum term (ie the beta3 hyperparameter that is introduced). How is beta1 and beta3 related? The numerator loses the interpretation of momentum, i.e., averaged past gradients, when using a separate beta3 parameter, and this does not feel like a principled change. I have a number of questions about the experiments as well, which makes it hard for me to interpret the significance of the empirical results presented: 1. What is the minibatch size used? Do any of the conclusions presented change if the minibatch size is changed? 2. Is the learning rate tuned? The authors mention the initial learning rate used for the experiments, but it is not clear why those values are used? Was the same initial learning rate used for all algorithms? 3. Were the beta1, beta2 and beta3 values tuned for AdaX? What about beta1 and beta2 for Adam? What are the optimal values of these parameters that were observed? 4. How sensitive is performance to the values of these hyperparameters? Overall I think this work requires quite a bit of work before it is ready for publication, and would benefit from a much more thorough empirical evaluation of the algorithm. ================================== Edit after rebutall: I thank the authors for their response. While the paper has definitely improved in the newer draft, after having read the other reviews and the updated draft, I believe the paper still requires a bit of work before being ready for publication. I am sticking to my score.