Original: Adam Can Converge with Proper β 1 and β
Revised: Adam Can Converge With Proper β 1 and β

Original: Abstract
Revised: Abstract

Original: Adam is one of the most popular algorithms in deep learning.
Revised: Ever since Reddi et al. (2018) pointed out the divergence issue of Adam, many new variantshave been designed to obtain convergence. However, vanilla Adam remains exceptionallypopular and it works well after proper tuning.

Original: Despite its practical popularity,the convergence of  Adam is far from being understood.
Revised: Why is the divergence not observed? Wepoint out there is a gap between theory and practice: Reddi et al. (2018) pick the samplesize n after picking the hyperparameters of Adam, i.e., ( β 1 , β 2 ) ; while practical applicationsoften fix n first and then tune ( β 1 , β 2 ) . Since Adam often works well in the latter setting, weconjecture it can still converge.

Original: In this work, we prove that whenthe 2nd-order momentum parameter β 2 is large and 1st-order momentum parameter β 1 <√ β 2 < 1 ,randomly shuffled Adam converges to the neighborhood of stationary points.
Revised: In this work, we confirm this conjecture under fixed n . Weprove that, when the 2nd-order momentum parameter β 2 is large and 1st-order momentumparameter β 1 < √ β 2 < 1 , Adam converges to the neighborhood of stationary points.

Original: Thesize of the neighborhood is propositional to thenoise variance of stochastic gradients.
Revised: The sizeof the neighborhood is propositional to the variance of stochastic gradients.

Original: Withextra strong growth condition, we prove that Adam can converge to stationary points.
Revised: Under an extracondition (strong growth condition), Adam converges to stationary points.

Original: Asβ 2 increases, our convergence result can cover any β 1 ∈ [0 , 1) including the default settingβ 1 = 0 .
Revised: As β 2 increases,our convergence result can cover any β 1 ∈ [0 , 1) including the default setting β 1 = 0 .

Original: To our best knowledge, we are the first to prove that Adam with any large β 1 canconverge without any form of bounded gradient assumption.
Revised: Toour best knowledge, we are the first to prove that Adam with arbitrarily large β 1 can convergewithout any form of bounded gradient assumption.

Original: We also investigate the divergence behavior of Adam.
Revised: This result shows that vanilla Adam,without any modification, is still theoretically justified to converge.

Original: When β 2 is small, we point out a largeregion of ( β 1 , β 2 ) combinations where Adam can diverge to infinity.
Revised: When β 2 is small, wefurther point out a large region of ( β 1 , β 2 ) combinations where Adam can diverge to infinity.

Original: We emphasize that bothour results above consider a quite different setting fromthe divergence result in Reddi et al.(2018): they consider picking ( β 1 , β 2 ) before picking the sample size n ; while we consider amore practical scenario with n fixed in advance. Our analysis shows that Adam’s behaviorchanges significantly with hyperparameters.
Revised: Our divergence result considers the same setting (fixing n ahead) as our convergence result,indicating that there is a phase transition from divergence to convergence when increasing β 2 .

Original: These results may providenew guidance towardsbetter hyperparameter tuning of Adam.
Revised: These results may provide guidance towards better hyperparameter tuning of Adam.

Original:  Introduction
Revised:  Introduction

Original: Modern machine learning tasks often aim to solve the following finite-sum problem.
Revised: Modern machine learning tasks often aim to solve the following finite-sum problem.

Original: [Equation]
Revised: [Equation]

Original: where n often stands for the number of samples or mini-batches and x denotes the trainable parameters.
Revised: where n is the number of samples or mini-batches and x denotes the trainable parameters.

Original: When combined with deep neural nets, Adam (Kingma & Ba, 2014) is one of the most dominantalgorithms for solving (1).
Revised: In deeplearning, Adam (Kingma & Ba, 2014) is one of the most popular algorithms for solving (1).

Original: In the past five years, Adam has reached tremendous success in differentmachine learning applications including natural language processing (NLP) (Vaswani et al., 2017; 
Revised: Ithas been applied to various machine learning domains such as natural language processing (NLP)(Vaswani et al., 2017; Brown et al., 2020;

Original: Brown et al., 2020; Devlin et al., 2018), generative adversarial networks (GANs) (Isola et al., 2017; Zhu et al., 2017; Radford et al., 2015) and computer vision (CV) (Dosovitskiy et al., 2021), etc..
Revised: Devlin et al., 2018), generative adversarial networks (GANs)(Radford et al., 2015; Isola et al., 2017; Zhu et al., 2017) and computer vision (CV) (Dosovitskiyet al., 2021).

Original: Despite its prevalence, Reddi et al.
Revised: Despite its prevalence, Reddi et al.

Original: (2018) point out that Adam can diverge under a wide range ofhyperparameters.
Revised: (2018) point out that Adam can diverge with a widerange of hyperparameters.

Original: A main result in (Reddi et al., 2018) conveys the message that 1 :
Revised: A main result in (Reddi et al., 2018) states that 1 :

Original: For any β 1 , β 2 s.t. {+
 ≤+} β 1 < √ β 2 < 1 , there exists a problem  that Adam will diverge.
Revised: For any β 1 , β 2 s.t. {+
 ≤+} β 1 < √ β 2 < 1 , there exists a problem such that Adam diverges.

Original: Here, β 1 is the hyperparameter to control Adam’s 1st-ordermomentum and β 2 control the weightedaverage of squared past gradients (a.k.a.
Revised: Here, β 1 and β 2 are the hyperparameter to control Adam’s 1st-order and 2nd-order momentum. More description of Adam can be seen in Algorithm 1 (presented later in Section 2.1).

Original: 2nd-order momentum).
Revised: Ever since(Reddi et al., 2018) pointed out the divergence issue, many new variants have been designed to obtainconvergence. For instance, (Zou et al., 2019) changed β 2 to iteration-dependent β 2 t ; AMSGrad Appendix D.2.

Original: More description can be seenin Algorithm 1 (presented later in Section 2.1).
Revised: (Reddi et al., 2018) enforced the adaptor v t (defined later in Algorithm 1) to be non-decreasing;

Original: Despite the divergence issue, Adam works wellin practice after proper tuning.
Revised: AdaBound (Luo et al., 2019) imposed constraint v t ∈ [ C l , C u ] to ensure the boundedness on effectivestepsize. These variants often reach good performance in practice. On the other hand, counter-intuitively, vanilla Adam remains exceptionally popular (see evidenceat (Scholar)). Without any modification on its update rules, Adam works well after proper tuning.

Original: Even more mysteriously, we find that the commonly reportedhyperparameters actually satisfy the divergence condition above.
Revised: Even more mysteriously, we find that the commonly reported hyperparameters actually satisfy thedivergence condition stated earlier.

Original: For instance, Kingma & Ba (2014) claim that ( β 1 , β 2 ) =
Revised: For instance, Kingma & Ba (2014) claimed that ( β 1 , β 2 ) =(0 .

Original: (0 . 9 , 0 . 999) is a “good choice for the tested machine learning problems" and itis indeed the default setting in practice.
Revised: 9 , 0 . 999) is a “good choice for the tested machine learning problems" and it is indeed the defaultsetting in PyTorch.

Original: Inrecent super-large models GPT-3 and Megatron (Brownet al., 2020; Smith et al., 2022), ( β 1 , β 2 ) = (0 . 9 , 0 . 95) is chosen as their final hyperparameters.
Revised: In super-large models GPT-3 and Megatron (Brown et al., 2020; Smith et al.,2022), ( β 1 , β 2 ) is chosen to be (0 . 9 , 0 . 95) .

Original: For GAN researchers (e.g. (Isola et al., 2017)), ( β 1 , β 2 ) =
Revised: GAN researchers (e.g. Radford et al. ; Isolaet al. (2017)) use ( β 1 , β 2 ) =

Original: (0 . 5 , 0 . 999) is often preferred.

Original: All thesehyperparameters satisfy β 1 < √ β 2 < 1 , and thus live in the “divergence region" by (Reddi et al.,2018).
Revised: All these hyperparameters live in the divergence regionβ 1 < √ β 2 .

Original: However, they all reach good performances and  the divergence issue is not observed.
Revised: Surprisingly, instead of observing the divergence issue, these hyperparameters achievegood performances and they actually show the sign of convergence.

Original: Why does Adam work well despite its  divergence issue?
Revised: Why does Adam work well despite its theoretical divergence issue?

Original: One might argue that neural nets trainingmight be somewhat “simpler".
Revised: Is there any mismatch betweendeep learning problems and the divergent example?

Original: We would like to point out it is not the crucial reason. Instead, wenotice a non-negligible gap between practical applications and the above divergence result: practicalapplications often fix the number ofsamples or mini-batches n in advance; while the divergenceresult picksn after picking ( β 1 , β 2 )  .This difference turns out to be crucial.
Revised: We take a closer look into the divergenceexample and find out the mismatch does exist. In particular, we notice an important (but oftenignored) characteristic of the divergence example: (Reddi et al., 2018) picks ( β 1 , β 2 ) before pickingthe sample size n .

Original: We explain as follows. To construct the divergence example, one needs to change n for different ( β 1 , β 2 ) .
Revised: Put in another way, to construct the divergence example, they change n fordifferent ( β 1 , β 2 ) .

Original: This is shown inthe proof of (Reddi et al., 2018) and is also restated in Appendix C.2.

Original: For instance, for ( β 1 , β 2 ) =(0 , 0 . 99) , we need one n to construct the divergent example; for ( β 1 , β 2 )=
Revised: For instance, for ( β 1 , β 2 ) = (0 , 0 . 99) , we use one n to construct the divergentexample; for ( β 1 , β 2 )

Original:  (0 , 0 . 9999) , we needanother n to construct another divergent example.
Revised: = (0 , 0 . 9999) , we use another n to construct another divergent example.

Original: However, in practical applications, sample size nis often given before picking ( β 1 , β 2 ) so the divergence result cannot be applied. Consideringthedifferent settings, the above divergence result might not provide proper guidance for Adam users.
Revised: This understanding reveals a gap between the divergence result and practice: In practical applicationsof Adam listed above, practitioners tune the hyperparameters ( β 1 , β 2 ) after the sample size n is given,while the divergence example considers the opposite order. Considering the good performance of Adam under fixed n , we conjecture that “the order of picking n or ( β 1 , β 2 ) " might be the major gapbetween theory and practice.

Original: So how would Adam behave whenthe sample size n is fixed?
Revised: So could Adam possibly converge when n is fixed?

Original: Can Adam possibly converge?
Revised: Unfortunately,the behavior of vanilla Adam is far less studied than its variants (perhaps due to the criticism ofdivergence).

Original: If so,what ( β 1 , β 2 ) pairs can ensure convergence? With these questions in mind, we run experiments ondifferent tasks.
Revised: To explore this question, we run experiments for different choices of ( β 1 , β 2 ) on afew tasks.

Original: First, we run Adam under different choices of ( β 1 , β 2 ) on a convex function (2) withfixed n (see the definition in Section 3.2).
Revised: First, we run Adam for a convex function (2) with fixed n (see the definition in Section3.2).

Original: Then, we run Adam  on real dataset MNIST and CIFAR-10with fixed batchsize.
Revised: Second, we run Adam for the classification problem on data MNIST and CIFAR-10 with fixedbatchsize.

Original: We observe some interesting phenomena in Figure 1 (a), (b) and (c).
Revised: We observe some interesting phenomena in Figure 1 (a), (b) and (c).

Original: First, when β 2 is large, the optimization error is small for almost all values of β 1 .
Revised: First, when β 2 is large, the optimization error is small for almost all values of β 1 .

Original: Second, whenβ 1 , β 2 are both small, there is a red region with relatively large error.
Revised: Second, whenβ 1 , β 2 are both small, there is a red region with relatively large error.

Original: On MNIST, CIFAR-10, theerror in the red region is around 0.4 times higher than that in the blue region.
Revised: On MNIST, CIFAR-10, theerror in the red region is around 0.4 times higher than that in the blue region.

Original: The situation is a lotworse on function (2): the error in the red region is 70 times higher than that in the blue region.
Revised: The situation is a lotworse on function (2): the error in the red region is 70 times higher than that in the blue region.

Original: [Figure]
Revised: [Figure]

Original: While Adam’s performances seem unstable in the red region, we find that Adam always performswell in the top blue region in Figure 1.
Revised: While Adam’s performances seem unstable in the red region, we find that Adam always performswell in the top blue region in Figure 1.

Revised: This seems to suggest that Adam can converge without anyalgorithmic modification, as long as β 1 and β 2 are chosen properly.

Original: We are led to the following question:
Revised: We ask the following question:

Original: When n is fixed ahead, does Adam provably converge in the top region of Figure 1?
Revised: When n is fixed, can vanilla Adam provably converge without any algorithmic modification?

Original: In this work, we theoretically explore this question.
Revised: In this work, we theoretically explore this question.

Original: When n is
Revised: Our contributions are summarized below andalso visualized in Figure 1 (d).

Original: fixed in advance, we prove the  results below.
Revised: We prove the following results when n is fixed.

Original: Our contributions are also visualized in Figure 1 (d) • When β 2 is large  and β 1 < √ β 2 ,we prove that Adam converges to the neighborhood of stationarypoints.
Revised: • We prove that when β 2 is large enough and β 1 < √ β 2 , Adam converges to the neighborhoodof stationary points.

Original: The size of the neighborhood is propositional to noise variance of stochastic gradients.
Revised: The size of the neighborhood is propositional to the variance of stochasticgradients.

Original: With  extra  strong growth condition, we prove that Adam can converge to stationary points.
Revised: With an extra condition (so-called strong growth condition), we prove that Adam canconverge to stationary points.

Original: Asβ 2 increases, these results can cover any momentum parameter β 1 ∈ [0 , 1) including the defaultsetting β 1 = 0 .
Revised: As β 2 increases, these results can cover any momentum parameterβ 1 ∈ [0 , 1) including the default setting β 1 = 0 .

Original: In particular, our analysis does not require bounded gradient assumption.•
Revised: In particular, our analysis does not requirebounded gradient assumption.•

Original: We point out thealgorithmic divergence in the red region in Figure 1 (d).
Revised: We study the divergence issue of Adam.

Original: We prove that: thereexists a function such that, Adam diverges to infinity when ( β 1 , β 2 ) is picked in the red region in Figure 1 (d). The size of the red region depends on n .
Revised: We prove that: for any fixed n , there exists a functionsuch that, Adam diverges to infinity when ( β 1 , β 2 ) is picked in the red region in Figure 1 (d).The size of the red region increases with n .

Original: The shape of the region follows thesolution to our analytic conditions.
Revised: The shape of the region follows the solution to ouranalytic conditions.•

Original: We emphasize that our  divergence result consider “fix n beforechoosing ( β 1 , β 2 ) ".
Revised: We emphasize a few characteristics of our results. (1) phase transition . The divergence resultconsiders the same setting (fixing n and then pick ( β 1 , β 2 ) ) as our convergence result, indicatingthat there is a phase transition from divergence to convergence. (2) problem-dependent bounds .Our

Original: This setting is quite different from (Reddi et al., 2018)mentioned in Section1 which considers a reverse ordering.
Revised: convergence and divergence regions of ( β 1 , β 2 ) are problem-dependent, which is drasticallydifferent from (Reddi et al., 2018) which established the problem-independent worst-case choiceof ( β 1 , β 2 ) . (3) non-asymptotic characterization . the “divergence region” of ( β 1 , β 2 ) expandsas n increases and converges to the whole region [0 , 1) 2 as n goes to infinity, which recovers(actually stronger than) the problem-independent divergence result of (Reddi et al., 2018) thatrequires β 1 < √ β 2 . In this sense, we can view the divergence result of (Reddi et al., 2018) as anasymptotic characterization of the divergence region (as n → ∞ ) and our divergence result as annon-asymptotic characterization (for any fixed n ).

Revised: We provide more discussion in Section 4.We believe

Original: We believe our results can boost new understandings for Adam.
Revised: our results can boost new understandings for Adam.

Original: While Reddi et al.
Revised: While Reddi et al.

Original: (2018) revealthat “Adam can diverge", our results show the other side of the coin: when n is fixed, Adam can stillwork if hyperparameters are chosen properly.
Revised: (2018) reveal that“Adam can diverge", our results show the other side of the coin: When n is fixed, Adam can still workwithout any algorithmic modification.

Original: Since “fixing n " is closer to practical scenarios, ourconvergence and divergence results may provide better guidance for hyperparameter tuning.
Revised: Since “fixing n " is closer to practical scenarios, our resultssuggests that Adam is still a theoretically justified algorithm and practitioners can use it confidently.

Original: We  emphasize two features in this work.
Revised: We further emphasize two features in this work.

Original: First, our convergence results can cover any β 1 ∈ [0 , 1) ,which can involve massive momentum.
Revised: First, our convergence results can cover anyβ 1 ∈ [0 , 1) , which can involve massive momentum.

Original: It turns out that large-momentum Adam is not easy to analyze.
Revised: It turns out that large-momentum Adam is noteasy to analyze.

Original: Even with stronger assumptions like bounded gradient ( ∥∇ f ( x ) ∥ < C , ∀ x ), its convergence is notwell understood (see related works in Section 2.2).
Revised: Even with stronger assumptions like bounded gradient ( ∥∇ f ( x ) ∥ < C , ∀ x ), itsconvergence is not well understood (see related works in Section 2.2).

Original: To our best knowledge, this is the first result thatproves Adam with any large β 1 can converge without any form of bounded gradient assumption.
Revised: To our best knowledge, this isthe first result that proves Adam with any large β 1 can converge without any form of bounded gradientassumption.

Revised: The proof contains a new method to handle unbounded momentum in the stochasticnon-linear dynamics system.

Original: Wewill highlight our technical novelty in Section 5.
Revised: We will highlight our technical novelties in Section 5.

Original: Second, when n is fixed, we point out Adam’s behavior  significantly changes under different ( β 1 , β 2 ) :it converges in certain region but it may diverge in others.
Revised: Second, when n is fixed, we point out the behavior of Adam significantly changes under different( β 1 , β 2 ) : It converges in certain region but it may diverge in others.

Revised: Pointing out these behaviorsin the ( β 1 , β 2 ) diagram may provide better guidance for tuning β 1 and β 2 .

Original: Removing boundedness condition isimportant to characterize such qualitative behavior.
Revised: Removing boundednesscondition is important to characterize such qualitative behaviors.

Original: With assumed boundedness on gradient oriterates, Adam never diverges, yet we have theoretical and numerical evidence showing that it can.
Revised: With assumed boundedness ongradient or iterates, Adam cannot diverge to infinity, yet we have theoretical and numerical evidenceshowing that it can.

Original:  Preliminaries
Revised:  Preliminaries

Original: 2.1 Problem Settings
Revised: 2.1 Problem Settings

Original: We consider finite-sum problem (1).
Revised: We consider finite-sum problem (1).

Original: We use x to denote the optimization variable.
Revised: We use x to denote the optimization variable.

Original: We denote ∇ f j asthe gradient of f j and let ◦ be the component-wise product.
Revised: We denote ∇ f j asthe gradient of f j and let ◦ be the component-wise product.

Original: The division and square-root operatorare component-wise as well.
Revised: The division and square-root operatorare component-wise as well.

Original: We present randamly shuffled Adam in Algorithm 1.Algorithm
Revised: We present randomly shuffled Adam in Algorithm 1.

Original: Initialize x 1 , 0 = x 0 , m 1 , − 1 = ∇ f ( x 0 ) and v 1 , − 1 = ∇ f ( x 0 ) ◦ ∇ f ( x 0 ) .

Original: for k = 1 → ∞ do

Original: [-Sample { τ k, 0 , τ k, 1 , · · · , τ k,n − 1 } as a random permutation of { 0 , 1 , 2 , · · · , n − 1 } for i = 
-]
Revised: [-Sample { τ k, 0 , τ k, 1 , · · · , τ k,n − 1 } as a random permutation of { 0 , 1 , 2 , · · · , n − 1 } for i = 
-]

Original: → n − 1 do m k,i =

Original: β 1 m k,i − 1 + (1 − β 1 ) ∇ f τ k,i ( x k,i ) v k,i =

Original: β 2 v k,i − 1 + (1 − β 2 ) ∇ f τ k,i ( x k,i ) ◦ ∇ f τ k,i ( x k,i ) x

Original: k,i +1 =

Original: [-x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
-]
Revised: [-x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
-]

Original: = x k,n v k

Original: [Equation]

Original: m

Original: [Equation]

Original: Break if certain exit condition is satisfied.

Original: end for In Algorithm 1, m denotes the 1st-order momentum and v denotes the 2nd-order momentum.
Revised: In Algorithm 1, m denotes the 1st-order momentum and v denotes the 2nd-order momentum.

Original: theyare weighted averaged by hyperparameter β 1 , β 2 , respectively.
Revised: theyare weighted averaged by hyperparameter β 1 , β 2 , respectively.

Original: Larger β 1 , β 2 will adopt more historyinformation.
Revised: Larger β 1 , β 2 will adopt more historyinformation.

Original: We denote x k,i , m k,i , v k,i ∈
Revised: We denote x k,i , m k,i , v k,i ∈

Original: R d as the value of x, m, v at the k -th outer loop (epoch)and i -th inner loop (batch), respectively.
Revised: R d as the value of x, m, v at the k -th outer loop (epoch)and i -th inner loop (batch), respectively.

Original: We choose η k = η 1 √ nk as the stepsize.
Revised: We choose η k = η 1 √ nk as the stepsize.

Original: In practice, ϵ isadopted for numerical stability and it is often chosen to be 10 − 8 .
Revised: In practice, ϵ isadopted for numerical stability and it is often chosen to be 10 − 8 .

Original: In our theory, we allow ϵ to be anarbitrary non-negative constant including 0 .
Revised: In our theory, we allow ϵ to be anarbitrary non-negative constant including 0 .

Original: In the vanilla version of Adam, it has an additional “bias correction” step (see (Kingma & Ba, 2014)).
Revised: In the original version of Adam in (Kingma & Ba, 2014), it has an additional “bias correction” step.

Original: This “bias correction” step can be implemented by changing the stepsize η k into ˆ η k = √ 1 − β k1 − β k 1 η kand using zero initialization.
Revised: This “bias correction” step can be implemented by changing the stepsize η k into ˆ η k = √ 1 − β k1 − β k 1 η kand using zero initialization.

Original: In Algorithm 1, the “bias correction” step is replaced by a specialinitialization, which also corrects the bias.
Revised: In Algorithm 1, the “bias correction” step is replaced by a specialinitialization, which also corrects the bias.

Original: Note that ˆ η k ∈ [ √ 1 − β 2 η k , 1 1 − β 1 η k ] is well-boundednear η k , so η k and ˆ η k brings the same convergence rate.
Revised: Note that ˆ η k ∈ [ √ 1 − β 2 η k , 1 1 − β 1 η k ] is well-bounded Algorithm Initialize x 1 , 0 = x 0 , m 1 , − 1 = ∇ f ( x 0 ) and v 1 , − 1 = max i ∇ f i ( x 0 ) ◦ ∇ f i ( x 0 ) . for k = 1 → ∞ do

Original: In addition, as the effect of initializationbecomes negligible when the training progresses, Adam with zero & our initialization will have thesame asymptotic behavior.{+{ 0 , 1 , 2 , · · · , n − 1 } for i = 
 → n − 1 do m k,i = β 1 m k,i − 1 + (1 − β 1 ) ∇ f τ k,i ( x k,i ) v k,i = β 2 v k,i − 1 + (1 − β 2 ) ∇ f τ k,i ( x k,i ) ◦ ∇ f τ k,i ( x k,i ) x+}
Revised: Sample { τ k, 0 , τ k, 1 , · · · , τ k,n − 1 } as a random permutation of {+{ 0 , 1 , 2 , · · · , n − 1 } for i = 
 → n − 1 do m k,i = β 1 m k,i − 1 + (1 − β 1 ) ∇ f τ k,i ( x k,i ) v k,i = β 2 v k,i − 1 + (1 − β 2 ) ∇ f τ k,i ( x k,i ) ◦ ∇ f τ k,i ( x k,i ) x+}

Revised: k,i +1 =

Original: {+x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
+}
Revised: {+x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
+}

Revised: = x k,n ; v k +1 , − 1 = v k,n − 1 ; m k +1 , − 1 = m k,n − 1 end for near η k , so η k and ˆ η k brings the same convergence rate.

Revised: In addition, as the effect of initializationbecomes negligible when the training progresses, Adam with zero & our initialization will have thesame asymptotic behavior.

Original: We follow the form of Algorithm 1, which makes results cleaner.
Revised: We follow the form of Algorithm 1, which makes results cleaner.

Original: In ouranalysis, we make the assumptions below.
Revised: In ouranalysis, we make the assumptions below.

Original: Assumption 2.1.
Revised: Assumption 2.1.

Original: we consider x ∈ R d and f i ( x ) satisfies gradient Lipschitz continuous with constant
Revised: we consider x ∈ R d and f i ( x ) satisfies gradient Lipschitz continuous with constant

Original: L .
Revised: L .

Original: We assume f ( x ) is lower bounded by a finite constant f ∗ .
Revised: We assume f ( x ) is lower bounded by a finite constant f ∗ .

Original: Assumption 2.2.
Revised: Assumption 2.2.

Original: f i ( x ) and f ( x ) satisfy: (cid:80) n
Revised: f i ( x ) and f ( x ) satisfy: (cid:80) n

Original: [Equation]
Revised: [Equation]

Original: [Equation] Assumption 2.2 is quite general, it covers some stronger conditions.
Revised: Assumption 2.2 is quite general.

Original: When D 1 = n , it becomes the“constant variance” with constant D 0 /n .
Revised: When D 1 = 1 /n , it becomes the “constant variance” with constant D 0 /n .

Original: “constant variance" condition is commonly used in both SGDand Adam analysis (e.g. (Ghadimi et al., 2016;
Revised: “constant variance" condition is commonly used in both SGD and Adam analysis (e.g.(Ghadimi et al., 2016;

Original: Zaheer et al., 2018; Huang et al., 2021)).
Revised: Zaheer et al., 2018; Huang et al., 2021)).

Original: Assumption2.2 allows more flexible choices of D 1 ̸ = n and thus it is weaker than “constant variance”.
Revised: Assumption 2.2 allows more flexiblechoices of D 1 ̸ = n and thus it is weaker than “constant variance”.

Original: When D 0 > 0 , the problem instance is sometimes called “non-realizable" (Shi et al., 2020).
Revised: When D 0 > 0 , the problem instance is sometimes called “non-realizable" (Shi et al., 2020).

Original: In thiscase, adaptive gradient methods are not guaranteed to reach the exact stationary point.
Revised: In thiscase, adaptive gradient methods are not guaranteed to reach the exact stationary point.

Original: Instead, theyonly converge to a bounded region (near stationary points)
Revised: Instead, theyonly converge to a bounded region (near stationary points)

Original: (Zaheer et al., 2018; Shi et al., 2020).
Revised: (Zaheer et al., 2018; Shi et al., 2020).

Original: Thisphenomenon indeed occurs for Adam in practice, even with diminishing stepsize (see Appendix D.2).
Revised: Thisphenomenon indeed occurs for Adam in experiments, even with diminishing stepsize (see Appendix B.1).

Original: The behavior of SGD is similar: constant stepsize SGD converges to a bounded region with its sizepropositional to the noise level D 

Revised: The behavior of SGD is similar: constant stepsize SGD converges to a bounded region with itssize propositional to the noise level D

Original: (Yan et al., 2018; Yu et al., 2019; Liu et al., 2020b).
Revised: (Yan et al., 2018; Yu et al., 2019; Liu et al., 2020b).

Original: When D 0 = 0 , Assumption 2.2 is often called “strong growth condition" (SGC) (Vaswani et al.,2019).
Revised: When D 0 = 0 , Assumption 2.2 is often called “strong growth condition" (SGC) (Vaswani et al.,2019).

Original: When ∥∇ f ( x ) ∥ = 0 , under SGC we have ∥∇ f j ( x ) ∥ = 0 for 
Revised: When ∥∇ f ( x ) ∥ = 0 , under SGC we have ∥∇ f j ( x ) ∥ = 0 for all j .

Original: all j . SGC is sometimes usedin SGD analysis e.g. (Schmidt & Roux, 2013; Vaswani et al., 2019).
Revised: SGC is increasinglypopular recently e.g.(Schmidt & Roux, 2013; Vaswani et al., 2019).

Original: This condition is reasonablein the overparameterized regime where neural networks are capable of interpolating all data points(Vaswani et al., 2019).
Revised: This condition is reasonablein the overparameterized regime where neural networks are capable of interpolating all data points: When every data point reaches zero training loss, the gradient of each point also converges to zero(Vaswani et al., 2019).

Original: We will show that Adam can converge  stationary points if SGC holds.
Revised: We will show that Adam can converge to stationary points if SGC holds.

Original: When n, f ∗ , L, D 0 , D 1 are fixed a priori, we use F n,f ∗ L,D 0 ,D 1 ( R d ) to denote the (sub) function classcontaining f ( x ) satisfying Assumption 2.1 and 2.2 with constant n, f ∗ , etc..
Revised: When n, f ∗ , L, D 0 , D 1 are fixed a priori, we use F n,f ∗ L,D 0 ,D 1 ( R d ) to denote the (sub) function classcontaining f ( x ) satisfying Assumption 2.1 and 2.2 with constant n, f ∗ , etc..

Original: We introduce thisnotation to clearly present the divergence result in Proposition 3.3.
Revised: We introduce thisnotation to clearly present the divergence result in Proposition 3.3.

Original: Without this pre-defined functionclass, the claim of divergence might bemisleading or confusing.
Revised: Without this pre-defined functionclass, the claim of divergence might be confusing.

Original: 2.2 Related Works
Revised: 2.2 Related Works

Original: Ever since Reddi et al.
Revised: Ever since Reddi et al.

Original: (2018) pointed out the divergence issue, there are many attempts on designingnew variants of Adam.
Revised: (2018) pointed out the divergence issue, there are many attempts on designingnew variants of Adam.

Original: For instance, Zou et al.

Original: and

Original: Chen et al. (2018) replace the constant( β 1 , β 2 ) by iterate-dependent ones e.g. β 1 t or β 2 t . These methods can address the divergence of theoriginal Adam. Many variants also achieve satisfactory performance in practice.
Revised: As introduced in Section 1, these methods can address the divergence ofthe original Adam and they often achieve satisfactory performance in practice.

Original: Since we focus onunderstanding Adam without modification, we introduce more variants later in Appendix C.1.
Revised: Since we focus onunderstanding Adam without modification , we introduce more variants later in Appendix D.1.

Original: There are several works trying to understand Adam’s convergence behavior, but they often requireextra conditions.
Revised: Compared with proposing new variants, the convergence of vanilla Adam is far less studied thanits variants (perhaps due to the criticism of divergence). We summarize a few works here. Zhouet al. (2018b) analyze the counter-example in (Reddi et al., 2018) and find certain hyperparametercan work.

Original: For instance, the analysis in Zaheer et al.
Revised: However, their analysis is restricted to the counter-example. Zaheer et al.

Original: (2018) reveals relation between mini-batchsizes and (non-)convergence of Adam.
Revised: (2018) revealsrelation between mini-batch sizes and (non-)convergence of Adam.

Original: Gadat & Gavra (2020) point out that Adam only converges tominimizers.

Original: However, these two works require β 1 = 0 and Adam is reduced to RMSProp (Hintonet al., 2012).
Revised: However, this work requireβ 1 = 0 and Adam is reduced to RMSProp (Hinton et al., 2012).

Original: De et al.
Revised: De et al.

Original: (2018)further analyze Adam with non-zero β 1 , but they assume the signof all stochastic gradients to keep the same.
Revised: (2018) analyze RMSPropand non-zero- β 1 Adam, but they assume the sign of all stochastic gradients to keep the same.

Original: It seems unclear how to check this condition a priori.
Revised: Itseems unclear how to check this condition a priori.

Original: Additionally, they require β 1 inversely related to the upper bound of gradient, which forces β 1 to besmall.
Revised: Additionally, they require β 1 to be inverselyrelated to the upper bound of gradient, which forces β 1 to be small (as a side note, this result only applies to full-batch Adam).

Original: Défossez et al.
Revised: Défossez et al.

Original: (2020) analyze Adam  β 1 < β 2 and provide some insights on the momentummechanisms.
Revised: (2020) analyze Adam with β 1 < β 2 and provide someinsights on the momentum mechanisms.

Original: However, they require ϵ (the hyperparameter for numerical stability)  to be large, inwhich case ϵ will be much larger than √ v .
Revised: However, their bound is inversely proportional to ϵ (thehyperparameter for numerical stability) and the bound goes to infinity when ϵ goes to 0.

Revised: This isdifferent from practical application since small ϵ such as 10 − 8 often works well.

Original: This is against the nature of adaptive gradient methodsbecause √ v no longer dominates in the choice of stepsize.
Revised: Further, using largeϵ is against the nature of adaptive gradient methods because √ v no longer dominates in the choiceof stepsize.

Revised: In this case, Adam is essentially transformed back to SGD.

Original: Twovery recent works (Huang et al.,2021) and (Guo et al., 2021) propose novel and simple frameworks to analyze Adam-family withlarge β 1 .
Revised: Two recent works (Huanget al., 2021) and (Guo et al., 2021) propose novel and simple frameworks to analyze Adam-familywith large β 1 .

Original: Yet, they require the effective stepsize of Adam will be bounded in certain interval, i.e.,1 √ v t + ϵ ∈ [ C l , C u ] 2 .
Revised: Yet, they require the effective stepsize of Adam to be bounded in certain interval, i.e.,1 √ v t + ϵ ∈ [ C l , C u ] 2 .

Original: This boundedness condition changes Adam into AdaBound (Luo et al.
Revised: This boundedness condition changes Adam into AdaBound (Luo et al.

Original: , 2019)and thus they cannot explain the observations on original Adam in Section 1.
Revised: , 2019)and thus they cannot explain the observations on original Adam in Section 1.

Original: Additionally, all ofthese works (including those for new variants) require bounded gradient assumptions.
Revised: Additionally, all ofthese works (including those for new variants) require bounded gradient assumptions.

Original: Among all these works, (Shi et al., 2020) is the first work that proves RMSProp (a non-momentumversion of Adam with β 1 = 0 ) can converge without bounded gradient assumption.
Revised: Among all these works, (Shi et al., 2020) is the first work that proves RMSProp (a non-momentumversion of Adam with β 1 = 0 ) can converge without bounded gradient assumption.

Original: When the2nd-order parameter β 2 is large, they show that RMSProp converges to a bounded region in general,and to stationary points under SGC.
Revised: When the2nd-order parameter β 2 is large, they show that RMSProp converges to a bounded region in general,and to stationary points under SGC.

Original: They also provide a convergence result for Adam with β 1 closeenough to 0.
Revised: They also provide a convergence result for Adam with β 1 closeenough to 0.

Original: However, a simple calculation by Zhang et al.
Revised: However, a simple calculation by Zhang et al.

Original: (2022) shows that they require β 1 < 10 − 7on real datasets.
Revised: (2022) shows that they require β 1 < 10 − 7on real datasets.

Original: Thus their result does not provide much extra information other than RMSProp.
Revised: Thus their result does not provide much extra information other than RMSProp.

Original: It is still unclear how β 1 affects the (non-)convergence of Adam.
Revised: Itis still unclear how β 1 affects the (non-)convergence of Adam.

Original: Proving convergence under largeβ 1 is non-trivial.
Revised: Proving convergence under large β 1is non-trivial.

Original: We point out there are at least two issues.
Revised: We point out there are at least two challenges.

Original: First, it seems unclear how to control they this Appendix the massive momentum m t .
Revised: First, it seems unclear how to controlthe massive momentum m t .

Original: Second, the momentum m t in Adam is multiplied by 1 / √ v t , causingnon-linear perturbation.
Revised: Second, the momentum m t in Adam is multiplied by 1 / √ v t , causingnon-linear perturbation.

Original: It turns out that we requires much different techniques to resolve these twoissues.
Revised: We use much different techniques to resolve these two issues.

Original: We highlight our technical novelties in Section 5.
Revised: We highlightour technical novelties in Section 5.

Original: 2.3 The Importance and Difficulties of Removing Bounded Gradient Assumptions
Revised: 2.3 The Importance and Difficulties of Removing Bounded Gradient Assumptions

Original: Here, we emphasize the importance to remove bounded gradient assumption.
Revised: Here, we emphasize the importance to remove bounded gradient assumption.

Original: First, unlike theassumptions in Section 2.1, bounded gradient is not common in SGD analysis.
Revised: First, unlike theassumptions in Section 2.1, bounded gradient is not common in SGD analysis.

Original: So it is of theoreticalinterests to remove this condition for Adam.
Revised: So it is of theoreticalinterests to remove this condition for Adam.

Original: Second, bounded gradient condition rules out thechances of gradient divergence a priori.
Revised: Second, bounded gradient condition rules out thechances of gradient divergence a priori.

Original: However, there are numerical evidences showing that Adam’sgradient can diverge (see Section 6 and Appendix D.2).
Revised: However, there are numerical evidences showing that Adam’sgradient can diverge (see Appendix B and Appendix B.1).

Original: Removing the boundedness assumption helpsus point out the divergence and convergence behavior of Adam in the ( β 1 , β 2 ) diagram.
Revised: Removing the boundedness assumptionhelps us point out the divergence and convergence behavior of Adam in the ( β 1 , β 2 ) diagram.

Original: However, it is often difficult to remove the boundedness condition, especially for stochastic algorithmswith momentum.
Revised: However, it is often difficult to remove the boundedness condition, especially for stochastic algorithmswith momentum.

Original: For instance, An early paper Bertsekas & Tsitsiklis (2000) analyzed SGD-typemethods without any boundedness condition.
Revised: For instance, An early paper Bertsekas & Tsitsiklis (2000) analyzed SGD-typemethods without any boundedness condition.

Original: But it is not until recently that Yu et al.
Revised: But it is not until recently that Yu et al.

Original: Liuet al. (2020b) prove SGDM (SGD with momentum) converges without bounded gradient assumption.
Revised: Liuet al. (2020b) prove SGDM (SGD with momentum) converges without bounded gradient assumption.

Original: Such attempts of removing boundedness assumptions are often appreciated for general optimizationproblems where “bounded-assumption-free" is considered as a major contribution.
Revised: Such attempts of removing boundedness assumptions are often appreciated for general optimizationproblems where “bounded-assumption-free" is considered as a major contribution.

Original: For Adam, the role of momentum  is even more intricate since it is multiplied by 1 / √ v t .
Revised: For Adam, the role of momentum m t is even more intricate since it is multiplied by 1 / √ v t .

Original: Combinedwith v t , the impact of previous signals not only affect the update direction, but also change the stepsizefor each component.
Revised: Combinedwith v t , the impact of previous signals not only affect the update direction, but also change the stepsizefor each component.

Revised: Further, both momentum m t and stepsize 1 / √ v t are random variables andthey are highly correlated.

Revised: Such statistical dependency causes trouble for analysis.

Original: As such, the role of momentum in Adam could be much different from that in SGDM or GDM.
Revised: In summary,the role of momentum in Adam could be much different from that in SGDM or GDM.

Original: Even with boundedness conditions, the convergence of large- β 1 Adam is still notwell understood (see related works in Section 2.2).
Revised: Even withboundedness conditions, the convergence of large- β 1 Adam is still not well understood (see relatedworks in Section 2.2).

Original: In this work, we propose new techniques to handle Adam’s momentum under any large β 1 , regardless of the gradient magnitude.
Revised: In this work, we propose new techniques to handle Adam’s momentum underany large β 1 , regardless of the gradient magnitude.

Original: These techniques arenot revealed in any existing works.
Revised: These techniques are not revealed in any existingworks.

Original: We introduce our technical contribution in Section 5.
Revised: We introduce our technical contribution in Section 5.

Original:  Main Results
Revised:  Main Results

Original: 3.1 Convergence Results
Revised: 3.1 Convergence Results

Original: [Figure]

Original: Here, we give the convergence results under large β 2 .
Revised: Here, we give the convergence results under large β 2 .

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: For any f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , when the hyperparameters in Algorithm 1 satisfies:β 1 < √ β 2 < 1 ; β 2 is greater or equal to a threshold γ 1 ( n ) ; and η k = η 1 √ nk , we will have: [Equation]
Revised: For any f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , we assume the hyperparameters in Algorithm 1satisfy: β 1 < √ β 2 < 1 ; β 2 is greater or equal to a threshold γ 1 ( n ) ; and η k = η 1 √ nk .

Original: where k m ∈
Revised: Let k m ∈

Original: N satisfies k m ≥ 4 and β ( k m − 1)n
Revised: Nsatisfies k m ≥ 4 and β ( k m − 1)

Revised: n≤ β n 1 √ k m − 1 , 4 we have the following results for any T > k m :

Original: [Equation]
Revised: [Table]

Original: Remark 1: the choice of β 2 .
Revised: Remark 1: the choice of β 2 .

Original: The threshold γ 1 ( n ) satisfies β 2 ≥ 1 − O (cid:16) 1 −
Revised: Our theory suggests that large β 2 should be used to ensure convergence.

Original: [Equation]
Revised: This message matches our experimental findings in Figure 1. We would like to point outthat the requirement of “large β 2 " is neccessary, because small β 2 will indeed lead to divergence(shown later in Section 3.2). We here comment a bit on the the threshold γ 1 ( n ) .

Original:  n 2 ρ (cid:17) (see inequality(32) and Remark G.8), where ρ is a constant that depends on the training trajectory.
Revised: γ 1 ( n ) satisfiesβ 2 ≥ 1 − O (cid:16) 1 − β n 1 n 2 ρ (cid:17) (see inequality (33) and Remark G.7), where ρ is a constant that depends on thetraining trajectory.

Original: In worst cases,ρ is upper bounded by n 2 . 5 , but we find the practical ρ to be much smaller.
Revised: In worst cases, ρ is upper bounded by n 2 . 5 , but we find the practical ρ to be muchsmaller.

Original: In Appendix D.2, weestimate ρ on MNIST and CIFAR-10.
Revised: In Appendix B.1, we estimate ρ on MNIST and CIFAR-10.

Original: In practical training process, we empirically observe thatρ ≈ O ( n ) , thus the required γ 1 ( n ) ≈ 1 − O (cid:0) n − 3 (cid:1) .
Revised: In practical training process, weempirically observe that ρ ≈ O ( n ) , thus the required γ 1 ( n ) ≈ 1 − O (cid:0) n − 3 (cid:1) .

Original: Note that our threshold of β 2 is a sufficientcondition for convergence, so there may be a gap between the practical choices and the theoreticalbound of β 2 .
Revised: Note that our thresholdof β 2 is a sufficient condition for convergence, so there may be a gap between the practical choicesand the theoretical bound of β 2 .

Original: Closing the gap will be an interesting future direction.
Revised: Closing the gap will be an interesting future direction.

Original: We further remark that γ 1 ( n ) increases with n .
Revised: We find that γ 1 ( n ) increases with n .

Original: This condition suggests that larger β 2 should be usedwhen the batchsize is small.
Revised: This property suggests that larger β 2 should be used when n islarge.

Original: This phenomenon is also verified by our experiments in Figure 4 (d), Section 6.
Revised: This phenomenon is also verified by our experiments in Figure 8 (d), Appendix B. We alsoremark that γ 1 ( n ) slowly increases with β 1 .

Original: Intuitively, this phenomenon is understandable since larger n brings larger noise in thestochastic gradient, and thus larger β 2 is required.
Revised: This property is visualized in Figure 1 (d) where thelower boundary of blue region slightly lifts up when β 1 increases.

Original: Remark 2: the choice of β 1 .
Revised: Remark 2: the choice of β 1 .

Original: β 1 < √ β 2 is also a requirement in (Kingma & Ba, 2014), buttheir convergence analysis requires bounded gradient and Reddi et al.
Revised: β 1 < √ β 2 is also a requirement in (Kingma & Ba, 2014), but theirconvergence analysis requires bounded gradient and Reddi et al.

Original: (2018) later pointed out thatthe proof has flaws.
Revised: (2018) later pointed out that theproof has flaws.

Original: Theorem 3.1 is the first result that proves the convergence of Adam with thesame choice of β 1 , but with less assumptions.
Revised: Theorem 3.1 is the first result that proves the convergence of Adam with the samechoice of β 1 , but with less assumptions.

Original: Since β 2 is suggested to be large, our convergenceresult can cover flexible choice of β 1 ∈ [0 , 1) .
Revised: Since β 2 is suggested to be large, our convergence result cancover flexible choice of β 1 ∈ [0 , 1) .

Original: For instance, β 2 = 0 . 999 brings the threshold ofβ 1 < 0 .
Revised: For instance, β 2 = 0 . 999 brings the threshold of β 1 < 0 .

Original: 9995 , which coversand basically all practical choices of β 1 reported in the literature (see 
Revised: 9995 ,which covers basically all practical choices of β 1 reported in the literature (see Section 1),

Original: Section 1), including the default setting β 1 = 0 . 
Revised: includingthe default setting β 1 = 0 . 9 .

Original: Remark 3: convergence to a bounded region.
Revised: Remark 3: convergence to a bounded region.

Original: When D 0 > 0 , Adam converges to a boundedregion near stationary points.
Revised: When D 0 > 0 , Adam converges to a boundedregion near stationary points.

Original: As discussed in Section 2.1, converging to bounded region is commonfor stochastic methods including constant-stepsize SGD (Yan et al., 2018;
Revised: As discussed in Section 2.1, converging to bounded region is commonfor stochastic methods including constant-stepsize SGD (Yan et al., 2018;

Original: Yu et al., 2019;
Revised: Yu et al., 2019;

Original: Liu et al.,2020b) and diminishing-stepsize RMSProp (Zaheer et al., 2018; Shi et al., 2020).
Revised: Liu et al.,2020b) and diminishing-stepsize RMSProp (Zaheer et al., 2018; Shi et al., 2020).

Original: This phenomenonis also observed in practice: even for convex quadratic function with D 0 > 0 , Adam with diminishingstepsize cannot reach exactly zero gradient (see Appendix D.2).
Revised: This phenomenonis also observed in practice: even for convex quadratic function with D 0 > 0 , Adam with diminishingstepsize cannot reach exactly zero gradient (see Appendix B.1).

Original: The good news is that, the constant
Revised: The good news is that, the constant

Original: O ( √ D 0 ) vanishes to 0 as β 2 goes to 1.
Revised: O ( √ D 0 ) vanishes to 0 as β 2 goes to 1.

Original: This can be seen in Remark G.17 in Appendix G.9.
Revised: This can be seen in Remark G.14 in Appendix G.9.

Original: As a corollary of Theorem 3.1, we have the following result under SGC (i.e., D 0 = 0 ).
Revised: As a corollary of Theorem 3.1, we have the following result under SGC (i.e., D 0 = 0 ).

Original: Corollary 3.2.
Revised: Corollary 3.2.

Original: Under the setting in Theorem 3.1.
Revised: Under the setting in Theorem 3.1.

Original: When D 0 = 0 for Assumption 2.2, we have
Revised: When D 0 = 0 for Assumption 2.2, we have

Original: [Equation]
Revised: [Equation]

Original: where k m ∈

Original: N satisfies k m ≥ 4 and β ( k m −

Original: [Equation] n≤ [Equation] Under SGC (i.e. D 0 = 0 ), Corollary 3.2 states that Adam can converge to stationary points.
Revised: Under SGC (i.e. D 0 = 0 ), Corollary 3.2 states that Adam can converge to stationary points.

Original: This isindeed the case in practice.
Revised: This isindeed the case in practice.

Original: For instance, function (2) satisfies SGC and we observe 0 gradient normafter Adam converges (see Appendix D.2).
Revised: For instance, function (2) satisfies SGC and we observe 0 gradient normafter Adam converges (see Appendix B and Appendix B.1).

Original: The convergence rate in Corollary 3.2 is comparable tothat of SGD under the same condition in (Vaswani et al., 2019).
Revised: The convergence rate in Corollary 3.2 iscomparable to that of SGD under the same condition in (Vaswani et al., 2019).

Original: 3.2 Divergence Results
Revised: 3.2 Divergence Results

Original: Theorem 3.1 shows that when β 2 is large, any β 1 < √ β 2 ensures convergence.
Revised: Theorem 3.1 shows that when β 2 is large, any β 1 < √ β 2 ensures convergence.

Original: Now we consider  the222 case where β 2 is small.
Revised: Now we consider 253 the case where β 2 is small.

Original: We will show that in this case, a wide range of β 1 is facing the risk  of divergingto infinity.
Revised: We will show that in this case, a wide range of β 1 is facing the risk 254 of diverging to infinity.

Revised: The divergence of small- β 2 Adam suggests that “large β 2 " is necessary inthe convergence result Theorem 3.1.

Original: We construct a counter-example in F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: We construct a counter-example in F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: Consider f ( x ) = (cid:80) n − 1 i =

Revised: Considerf ( x ) = (cid:80) n − 1 i =

Original: f i ( x ) forx ∈ R , we define f i ( x ) as:
Revised: f i ( x ) for x ∈ R , we define f i ( x ) as:

Original: [Equation]
Revised: [Equation]

Original: Summing up all the f i ( x ) , we can see that
Revised: Summing up all the f i ( x ) , we can see that

Original: [Equation]
Revised: [Equation]

Original: is a lower bounded convex smooth function with optimal solution x ∗ = − 2 .
Revised: is a lower bounded convex smooth function with optimal solution x ∗ = − 2 .

Original: Function (2) allows bothiterates and gradients to diverge to infinity.
Revised: Function (2) allowsboth iterates and gradients to diverge to infinity.

Original: As shown in Figure 1 (a), when running Adam on (2), there exists a red large-error region.
Revised: As shown in Figure 1 (a), when running Adam on(2), there exists a red large-error region.

Original: Thisshows the sign of divergence.
Revised: This shows the sign of divergence.

Original: We further theoretically verify the conjecture.
Revised: We further theoreticallyverify the conjecture in Proposition 3.3.

Original: Proposition 3.3.
Revised: Proposition 3.3.

Original: For any function class F n,f ∗ L,D 0 ,D 1 ( R d ) , there exists a f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , s.t.when ( β 1 , β 2 ) satisfies analytic condition (9) , (10) , (11) in Appendix E, Adam’s iterates and functionvalues diverge to infinity.
Revised: For any function class F n,f ∗ L,D 0 ,D 1 ( R d ) , there exists a f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , s.t.when ( β 1 , β 2 ) satisfies analytic condition (11) , (12) , (13) in Appendix E, Adam’s iterates and functionvalues diverge to infinity.

Original: By solving these conditions in NumPy , we plot the orange region in 2.
Revised: By solving these conditions in NumPy , we plot the orange region in Figure2.

Original: Thesize of the region depends on n and it expands to the whole region when n goes to infinity. [Equation] b e t a
Revised: The size of the region depends on n and it expands to the whole region when n goes to infinity.0.

Original: We present the proof  in Appendix E.
Revised: The proof can be seen in Appendix E.

Original: We find the “divergence region"always stays below the “convergence threshold" 
Revised: We find the “divergenceregion" always stays below the “convergence threshold" γ 1 ( n ) in

Original: γ 1 ( n ) in Theorem3.1, so the two results are self-consistent (see the remark in Appendix E).
Revised: Theorem 3.1, so the two results are self-consistent (see the remarkin Appendix E).

Original: Proposition 3.3 states the divergence of iterates and functionvalues.
Revised: Proposition 3.3 states the divergence of iteratesand function values.

Original: Our experiments also show the divergence of gradient (see Section 6 and Appendix D.2).
Revised: Consistently, our experiments also show thedivergence of gradient (see Appendix B and Appendix B.1).

Original: These results characterize Adam’sdivergence behavior both numerically and theoretically.
Revised: Theseresults characterize Adam’s divergence behavior both numericallyand theoretically.

Original: We emphasize that the orange region is not discussed in (Reddi et al.,2018) because we consider n fixed while they allow n changing.
Revised: We emphasize that the orange region is not discussed in (Reddi et al.,2018) because we consider n fixed while they allow n changing.

Original: When n is allowed to increase, our orange region will expand tothe whole region and thus we can derive a similar result as (Reddiet al., 2018).
Revised: When n is allowed to increase, our orange region will expand to thewhole region and thus we can derive a similar (actually stronger)result as (Reddi et al., 2018).

Original: We provide more explanation in Section 4.
Revised: We provide more explanation in Section 4.

Original: Combining Theorem 3.1 and Proposition 3.3, we establish a clearer image onthe relation between ( β 1 , β 2 ) and qualitative behavior of Adam.
Revised: Combining Theorem 3.1 and Proposition 3.3, we establisha clearer image on the relation between ( β 1 , β 2 ) and qualitativebehavior of Adam.

Original:  Reconciling Our Results with (Reddi et al., 2018)
Revised:  Reconciling Our Results with (Reddi et al., 2018)

Original: We discuss more on the relation between (Reddi et al., 2018) and our results.
Revised: We discuss more on the relation between (Reddi et al., 2018) and our results.

Original: Recall Reddi et al.(2018) show that “for any β 1 < √ β 2 < 1 , Adam can diverge".
Revised: Recall Reddi et al.(2018) show that “for any β 1 < √ β 2 < 1 , Adam can diverge".

Original: This does not contradict with ourconvergence results in Theorem 3.1.
Revised: This does not contradict with ourconvergence results in Theorem 3.1.

Original: Further, it is different from our divergence result in Proposition3.3.
Revised: Further, it is different from our divergence result in Proposition3.3.

Original: As discussed in Section 1, the key difference lies in whether ( β 1 , β 2 ) is picked before or afterpicking n , which often stands for sample size or the number of mini-batch.
Revised: As discussed in Section 1, the key difference lies in whether ( β 1 , β 2 ) is picked before or afterpicking n , which often stands for sample size or the number of mini-batch.

Original: Since n is fixed when the (sub-)function class F n,f ∗ L,D 0 ,D 1 ( R d ) is given, this key difference can also beexpressed as whether ( β 1 , β 2 ) is picked before or after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: Since n is fixed when the (sub-)function class F n,f ∗ L,D 0 ,D 1 ( R d ) is given, the above key difference canalso be expressed as whether ( β 1 , β 2 ) is picked before or after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: Now we discuss thefollowing two cases.
Revised: Now wediscuss the following two cases.

Original: Case I: When ( β 1 , β 2 ) is picked before picking F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: Case I: When ( β 1 , β 2 ) is picked before picking F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: As discussed in Section 1, the divergence result requires differentn for different ( β 1 , β 2 ) .
Revised: As discussed in Section 1, the divergence result requires differentn for different ( β 1 , β 2 ) .

Original: In this sense, the considered function classis constantly changing.
Revised: In this sense, the considered function classis constantly changing.

Original: It does not contradict with our Theorem3.1 which considers a fixed function class with fixed n .
Revised: It does not contradict with our Theorem 3.1which considers a fixed function class with fixed n .

Original: For Case I , we illustrate Adam’s behavior in Figure 3.
Revised: For Case I , weillustrate Adam’s behavior in Figure 3.

Original: The red region withis proved by (Reddi et al., 2018).
Revised: The red region is proved by(Reddi et al., 2018).

Original: For completeness, we remove thecondition “ β 1 < √ β 2 " and further prove that Adam will divergeto infinity for any ( β 1 , β 2 ) ∈ [0 , 1) 2 .
Revised: For completeness, we remove the condition“ β 1 < √ β 2 " and further prove that Adam will diverge to infinityfor any ( β 1 , β 2 ) ∈ [0 , 1) 2 .

Original: The result is shown in thefollowing Corollary 4.1.
Revised: The result is shown in the following Corollary 4.1.

Original: Corollary 4.1.
Revised: Corollary 4.1.

Original: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , there exists a functionsatisfying Assumption 2.1 and 2.2 that the Adam’s iterates andfunction values diverge to infinity.
Revised: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , there exists a functionsatisfying Assumption 2.1 and 2.2 that the Adam’s iterates andfunction values diverge to infinity.

Original: Rigorously speaking, Reddi et al.

Original: (2018) consider constrained problems.

Original: So their claimed “divergence" actually means the iterates willstay in a huge region with the size of whole feasible set.

Original: We callit “non-convergence" to distinguish from our result of “diverge toinfinity" in Figure 3.

Original: Proof of Corollary 4.1 can be seen in the final paragraph in Appendix
Revised: Proof of Corollary 4.1 can be seen in the final paragraph in Appendix

Original: E.
Revised: E.

Original: In the proof, we also require different n to cause divergence for different ( β 1 , β 2 ) .
Revised: In the proof, we also require different n to cause divergence for different ( β 1 , β 2 ) .

Original: So the functionclass is constantly changing.
Revised: So the functionclass is constantly changing.

Original: As a result, in Case I , we cannot prove any convergence result.
Revised: As a result, in Case I , we cannot prove any convergence result.

Original: [Table]
Revised: [Table]

Original: The orange region in Figure

Original: [Table] Case II: When ( β 1 , β 2 ) is picked after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: Case II: When ( β 1 , β 2 ) is picked after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: When the function class is pickedin advance, sample size n will also be fixed.
Revised: When the function class is pickedin advance, sample size n will also be fixed.

Original: This case is closer to most practical applications.
Revised: This case is closer to most practical applications.

Original: Inthis case, we find that Adam’s behavior changes significantly in the different region of Figure 3.
Revised: Inthis case, we find that Adam’s behavior changes significantly in the different region of Figure 3.

Original: First, ∀ f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) will converge when β 1 < √ β 2 and β 2 is large.
Revised: First, ∀ f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) will converge when β 1 < √ β 2 and β 2 is large.

Original: Second, ∃ f ( x ) ∈
Revised: Second, ∃ f ( x ) ∈

Original: F n,f [Equation] L,D 0 ,D 1 ( R d ) will diverge to infinity when ( β 1 , β 2 ) are in the orange region in Figure 2.
Revised: F n,f ∗ L,D 0 ,D 1 ( R d ) will diverge to infinity when ( β 1 , β 2 ) are in the orange region in Figure 2.

Original: Since Case
Revised: Since Case

Original: II is closer to practical scenarios, these results can provide better guidance for hyperparameter tuningfor Adam users.
Revised: II is closer to practical scenarios, these results can provide better guidance for hyperparameter tuningfor Adam users.

Original: We provide some suggestions  in Appendix B.
Revised: We provide some suggestions for practitioners in Appendix C.

Original: For Case II , we summarize the possible behaviors of Adam in Table 1.
Revised: For Case II , we summarize the possible behaviors of Adam in Table 1.

Original: We also illustrate ourconvergence and divergence results in Figure 1 (d).
Revised: We also illustrate ourconvergence and divergence results in Figure 1 (d).

Original: Note that there are some blanket areas where
Revised: Note that there are some blanket areas where

Original: Adam’s behavior remains unknown, this part will be left as interesting future work.
Revised: Adam’s behavior remains unknown, this part will be left as interesting future work.

Original:  Proof Ideas for the Convergence Result
Revised:  Proof Ideas for the Convergence Result

Original: We nowbriefly introduce our proof ideas for the convergence result in Theorem 3.1.
Revised: We now introduce our proof ideas for the convergence result in Theorem 3.1.

Original: Simply put, we 289 need to  control the update direction m k,i / √ v k,i near the negative gradient direction.
Revised: Simply put, we want to 318 control the update direction m k,i / √ v k,i inside the dual cone of gradient direction.

Original: Namely, we want to show
Revised: Namely: E ⟨∇ f ( x k, 0 ) ,

Original: [Equation]
Revised: [Equation]

Original: However, directly proving (3) could be difficult because both m k,i and v k,i distort the trajectory.
Revised: However, directly proving (3) could be difficult because both m k,i and v k,i distort the trajectory.

Original: To start with, we try to control the movement of v k,i .
Revised: To start with, we try to control the movement of v k,i .

Original: We use a standard approach for Adam-type algorithms: use large or increasing β 2 (e.g. (Shi et al., 2020;
Revised: We use a standard approach for Adam-type algorithms: use large or increasing β 2 (e.g. (Shi et al., 2020;

Original: Zou et al., 2019; Chen et al., 2018)).
Revised: Zou et al., 2019; Chen et al., 2018b)).

Original: Recall v k,i = (1 − β 2 ) (cid:80) ij =1 β i − j 2 ∇ f τ k,j ( x k,j ) ◦ ∇ f τ k,j ( x k,j ) +
Revised: Recall v k,i = (1 − β 2 ) (cid:80) ij =1 β i − j 2 ∇ f τ k,j ( x k,j ) ◦ ∇ f τ k,j ( x k,j ) +

Original: β
Revised: β

Original: i 2 v k, 0 , we have v k,i ≈ v k, 0 when β 2 is large.
Revised: i 2 v k, 0 , we have v k,i ≈ v k, 0 when β 2 is large.

Original: In this case, we have:
Revised: In this case, we have:

Original: [Equation]
Revised: [Equation]

Revised: where the first “ ≈ " is due to the large β 2 and the second “ ≈ " is our goal.

Original: Now we need to show:
Revised: Now we need to show:

Original: [Equation]
Revised: [Equation]

Original:  where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) , similarly for m l,k, 

Revised: √ [Equation] where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) , similarly for m l,k,

Original: and v l,k, 0 .
Revised: and v l,k, 0 .

Original: ( ∗ ) is due to thefinite-sum structure.
Revised: ( ∗ ) is due to thefinite-sum structure.

Original: However, it is not easy to prove (4).
Revised: However, it is not easy to prove (4).

Original: We point out some technical issues below.
Revised: We point out some technical issues below.

Original: Issue I: massive momentum.
Revised: Issue I: massive momentum.

Original: Directly proving (4) is still not easy.
Revised: Directly proving (4) is still not easy.

Original: We need to first consider asimplified problem: for every l ∈ [ d ] , assume we treat ∂ l f ( x k, 0 ) / √ v l,k, 0 as a constant, how to bound 
Revised: We need to first consider asimplified problem: for every l ∈ [ d ] , assume we treat ∂ l f ( x k, 0 ) / √ v l,k, 0 as a constant, how to bound E [Equation]

Original: E (cid:80) n − 1 i =0 (cid:0) m l,k,i − ∂ l f τ k,i ( x k, 0 ) (cid:1) ?
Revised: i [Equation] (cid:0) [Equation]

Original: It turns out that this simplified problem is still non-trivial.
Revised: It turns out that this simplified problem is still non-trivial.

Original: When β 1 is large, m l,k,i contains heavyhistorical signals which significantly distort the trajectory from gradient direction.
Revised: When β 1 is large, m l,k,i contains heavyhistorical signals which significantly distort the trajectory from gradient direction.

Original: Existing literatures(Zaheer et al., 2018; De et al., 2018; Shi et al., 2020) take a naive approach: they set β 1 ≈ 

Revised: Existing literatures(Zaheer et al., 2018; De et al., 2018; Shi et al., 2020) take a naive approach: they set β 1 ≈

Original: so thatm l,k,i ≈ ∂
Revised: so thatm l,k,i ≈ ∂

Original: l f τ k,i ( x k,i ) .
Revised: l f τ k,i ( x k,i ) .

Original: Then we get (4) ≈ 0 .
Revised: Then we get (4) ≈ 0 .

Original: However, this method cannot be applied here since weare interested in practical cases where β 1 is large in [0 , 1) .
Revised: However, this method cannot be applied here since weare interested in practical cases where β 1 is large in [0 , 1) .

Original: Issue II:  non-linear dynamics.
Revised: Issue II: stochastic non-linear dynamics.

Original: Even if we solve Issue I , it is still unclear how to prove (4).
Revised: Even if we solve Issue I , it is still unclear how to prove(4).

Original: This isbecause: for every l ∈ [ d ] , ∂ l f ( x k, 0 ) / √ v l,k, 0 is a r.v. instead of a constant.
Revised: This is because: for every l ∈ [ d ] , ∂ l f ( x k, 0 ) / √ v l,k, 0 is a r.v. instead of a constant.

Original: With this term involved,we are facing with a stochastic non-linear dynamics, which could be difficult to analyze.
Revised: With this terminvolved, we are facing with a stochastic non-linear dynamics, which could be difficult to analyze.

Revised: Further, ∂ l f ( x k, 0 ) / √ v l,k, 0 is statistically dependent with (cid:0) m l,k,i − ∂ l f τ k,i ( x k, 0 ) (cid:1) , so we are notallowed to handle the expectation E ( ∂ l f ( x k, 0 ) / √ v l,k, 0 ) separately .

Original: Unfortunately, even with additional assumption like bounded gradient, there is no general approachto tackle the above issues.
Revised: Unfortunately, even with additional assumptions like bounded gradient, there is no general approachto tackle the above issues.

Original: In this work, we propose a solution regardless of gradient magnitude.
Revised: In this work, we propose solutions regardless of gradient magnitude.

Original: Step 1.
Revised: Solution to Issue I.

Original: We prove the following Lemma to resolve Issue I .
Revised: We prove the following Lemma to resolve Issue I .

Original: Lemma 5.1.
Revised: Lemma 5.1.

Original: (Informal) Consider Algorithm 1.
Revised: (Informal) Consider Algorithm 1.

Original: For every l ∈ [ d ] and any β 1 ∈ [0 , 1) , we have the 309 following result under Assumption 2.1.δ [Equation] β 1 ) : =
Revised: For every l ∈ [ d ] and any β 1 ∈ [0 , 1) , we have the 339 following result under Assumption 2.1.

Original: E

Original: [Equation]
Revised: [Equation]

Original: where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) ; m l,k,i
Revised: where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) ; m l,k,i

Original: = (1 − β 1 ) ∂ l f  ( x k,i ) + β 1 m l,k,i − 1 .
Revised: = (1 − β 1 ) ∂ l f τ k,i ( x k,i ) + β 1 m l,k,i − 1 .

Original: We present the proof idea in Appendix A. Simply put, we construct a simple toy example called“color-ball" model, which shows a special property of δ ( β 1 ) .
Revised: We present the proof idea in Appendix A. Simply put, we construct a simple toy example called“color-ball" model (of the 1st kind). This toy model shows a special property of δ ( β 1 ) .

Original: We find out: for Algorithm 1, error termsfrom successive epochs can be canceled, which keeps the momentum roughly in the descent direction.
Revised: We find out:for Algorithm 1, error terms from successive epochs can be canceled, which keeps the momentumroughly in the descent direction.

Original: This important property is not revealed in any existing work.
Revised: This important property is not revealed in any existing work.

Original: In Appendix A, we demonstrate the“color-ball” model in simple language where only basic probability knowledge is required.

Original: Remark 4: When assuming bounded gradient ∥∇ f ( x ) ∥ ≤ G , a naive upper bound would beδ ( β 1 ) =
Revised: Remark 4: When assuming bounded gradient ∥∇ f ( x ) ∥ ≤ G , a naive upper bound would beδ ( β 1 ) =

Original: O ( G ) .
Revised: O ( G ) .

Original: However, such constant upper bound does not imply δ ( β 1 ) is close to 0.
Revised: However, such constant upper bound does not imply δ ( β 1 ) is close to 0.

Original: It will nothelp prove the convergence.
Revised: It will nothelp prove the convergence.

Original: This might be partially the reason why large- β 1 Adam is hard to analyzeeven under bounded gradient (see related works in Section 2.2).
Revised: This might be partially the reason why large- β 1 Adam is hard to analyzeeven under bounded gradient (see related works in Section 2.2).

Original: We emphasize Lemma 5.1 holds trueregardless of gradient norm, so it could be deployed in both bounded or unbounded gradient analysis.
Revised: We emphasize Lemma 5.1 holds trueregardless of gradient norm, so it could be deployed in both bounded or unbounded gradient analysis.

Original: Step 2.
Revised: Solution to Issue II.

Original: We try to show (4) by adopting Lemma 5.1.
Revised: We try to show (4) by adopting Lemma 5.1.

Original: However, the direct application cannotwork since ∂ l f ( x k, 0 ) √ v l,k, 0 is not constant.
Revised: However, the direct applicationcannot work since ∂ l f ( x k, 0 ) √ v l,k, 0 is not constant.

Original: Despite its randomness, we find out that when β 2 is large, (cid:110) [Equation]
Revised: Despite its randomness, we find out that when β 2 islarge, (cid:110) ∂

Original: l [Equation] ( x [Equation]
Revised: l f ( x k, 0 ) √ v l,k, 0 (cid:111) ∞

Original: [Equation] ,k, [Equation] ∞ k =1 behaves like Cauchy sequence which has unique limit point.
Revised: k =1 behaves like Cauchy sequence which has unique limit point.

Original: As such, although∂ l f ( x k, 0 ) √ v l,k, 0 brings extra perturbation, the quantity in (4) share the similar asymptotic behavior as δ ( β 1 ) .
Revised: As such, although∂ l f ( x k, 0 ) √ v l,k, 0 brings extra perturbation, the quantity in (4) share the similar asymptotic behavior as δ ( β 1 ) .

Original: Therefore, with the help of Step 1 , we prove the following Lemma 5.2.
Revised: We prove the following Lemma 5.2.

Original: Lemma 5.2.
Revised: Lemma 5.2.

Original: (Informal) Under Assumption 2.1 and 2.2, consider Algorithm 1 with large β 2 and β 1 < √ β 2 .
Revised: (Informal) Under Assumption 2.1 and 2.2, consider Algorithm 1 with large β 2 andβ 1 < √ β 2 .

Original: For those l with gradient component larger than certain threshold, we have:
Revised: For those l with gradient component larger than certain threshold, we have:

Original: [Equation]
Revised: [Equation]

Original: We conclude the proof of (4) by some additional analysis on “those l with small gradient component".
Revised: [Equation] In Appendix A, we introduce how to derive (6) from (5).

Original: This case is a bit easier since it reduce to bounded gradient case. Since the whole proof is quite long,we provide a proof sketch in Appendix G.1. The whole proof is presented in Appendix G.
Revised: To do so, we introduce a new type of“color-ball" model (we call it color-ball of the 2nd kind) which adopts the random perturbation of∂ l f ( x k, 0 ) √ v l,k, 0 .

Original: To support our theory, we provide more simulations and real-data experiments. All the experimentalsettings and hyperparameters are presented in Appendix D. We aim to show:(I). When β 2 is large, a large range of β 1 gives good performance, including all β 1 < √ β 2 .(II).
Revised: Understanding color-ball model of the 2nd kind is crucial for proving Lemma 5.2.

Original: When β 2 is small, a large range of β 1 performs relatively badly

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: (d) MNIST(batchsizev.s. β 2 ) Counter-example (2) . We run more experiments on (2) with initialization x = − 5 and n =5 , 10 , 15 , 20 .
Revised: We conclude the proof of (4) by some additional analysis on “those l with small gradient component".

Original: We show the result of n = 20 in Figure 4 (a), (b); the rest are shown in Appendix D.2.

Original: We find that: first, when β 2 is small, gradient norms increase dramatically. This emphasizes theimportance of discarding bounded gradient assumptions. Second, when β 2 is large, Adam convergesto stationary points for β 1 < √ β 2 . This matches our theory since (2) satisfies D 0 = 0 .
Revised: This case is a bit easier since it reduces to bounded gradient case.

Original: For functionswith D 0 > 0 , we further show that Adam can only converge toa bounded region in Appendix D.2. MNIST and CIFAR-10. As shown in Figure 1 (b)& (c) in Section 1, the training results matchboth claim (I) and (II).
Revised: For readers who wants to learn moreabout the idea of tackling Issue I and II , please refer to Appendix A where we formally introduce the1st and 2nd kind of color-ball models.

Original: In addition, there is a convex-shaped boundary on the transition from low lossto higher loss, this boundary roughly matches the condition in Theorem 3.1. NLP. We use Adam to train Transformer XL (Dai et al., 2019) on the WikiText-103 dataset (Merityet al., 2016). This architecture and dataset is widely used in NLP tasks (e.g. (Howard & Ruder, 2018;
Revised: Since the whole proof is quite long, we provide a proof sketchin Appendix G.1.

Original: See et al., 2017)). As shown in Figure 4 (c), the training results match both claim (I) and (II). Batchsize and β 2 . As shown in Figure 4 (d): on MNIST, smaller batchsize requires larger β 2 toreach small loss.
Revised: The whole proof is presented in Appendix G.

Original: This matches the message by Theorem 3.1: the threshold of β 2 increases with n .

Original:  Conclusions
Revised:  Conclusions

Original: In this work, we explore the (non-)convergence of Adam.
Revised: In this work, we explore the (non-)convergence of Adam.

Original: When β 2 is large, we prove that Adamcan converge with any β 1 < √ β 2 .
Revised: When β 2 is large, we prove that Adamcan converge with any β 1 < √ β 2 .

Original: When β 2 is small, we further show that Adam might diverge toinfinity for a wide range of β 1 .
Revised: When β 2 is small, we further show that Adam might diverge toinfinity for a wide range of β 1 .

Original: One interesting question is to verify the advantage of Adam over
Revised: One interesting question is to verify the advantage of Adam over

Original: SGD.
Revised: SGD.

Original: In this work, we focus on the fundamental issue of convergence.
Revised: In this work, we focus on the fundamental issue of convergence.

Original: Proving faster convergence of
Revised: Proving faster convergence of

Original: Adam would be our future work.
Revised: Adam would be our future work.

Original: Checklist

Original: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [Yes]

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?

Original: [Yes] If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results? [Yes]
Revised: Appendix Organization

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes]

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?

Original: [No] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: The Appendix is organized as follows.

Original: [Yes] (c) Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?{+• Appendix A introduces two “color-ball" models and their applications in proving Lemma 5.
+}
Revised: {+• Appendix A introduces two “color-ball" models and their applications in proving Lemma 5.
+}

Original: [No] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [No]
Revised: andLemma 5.2.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators? [Yes] (b) Did you mention the license of the assets?
Revised: This part is important for proving Theorem 3.1.• Appendix B introduces experiments to support our theory.• Appendix C provide some suggestions for hyperparameter tuning of Adam.• Appendix D.1 provide more discussions on some recent related works.

Original: [No]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [No] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No]
Revised: Appendix D.2 re-state thenon-convergence results in (Reddi et al., 2018).

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A]

Original: Negative Social Impact

Original: This script may provide better guidance for deep neural network training. It would have certainnegative social impact if neural nets are deployed for illegal usage. Appendix Organization The Appendix is organized as follows.
Revised: • Appendix E provides detailed proof for Proposition 3.3.• Appendix F provides some more notations and technical lemmas that serve for the proof of

Original: • Appendix A discusses the ideas of “color-ball" model and its application in proving Lemma5.1. This part is important for proving Theorem 3.1.• Appendix B provide some suggestions for hyperparameter tuning of Adam.• Appendix C.1 provide more discussions on some recent related works. Appendix C.2re-state the non-convergence results in (Reddi et al., 2018).• Appendix D provides more detailed setting of the experiments, as well as some additionalexperiments mentioned in the full script.• Appendix E provides detailed proof for Proposition 3.3.
Revised: Theorem 3.1.

Original: • AppendixF provides more some notations and technical lemmas that serve for the proof ofTheorem 3.1.• Appendix G provides detailed proof for Theorem 3.1.
Revised: • Appendix G provides detailed proof for Theorem 3.1.

Original: Especially, Appendix G.1 provide aproof roadmap.
Revised: Especially, Appendix G.1 provide a proofroadmap.

Original: AAn Introduction to the “Color-Ball” Model
Revised: A Introduction to the “Color-Ball” Models: the Key Ingredients to ProveTheorem 3.1

Original: We now introduce ideas of the “color-ball" model and its application in bounding δ ( β 1 ) in Lemma5.1.
Revised: We now introduce two “color-ball" models and their applications in tackling Issue I and Issue IImentioned in Section 5.

Original: This part is important for proving Theorem 3.1. 
Revised: These two color-ball models are important for proving Theorem 3.1. Solution to Issue I.

Original: As discussed in Section 5, we wish to show that δ ( β 1 ) vanishes with k .
Revised: As discussed in Section 5, we wish to show that δ ( β 1 ) vanishes with k .

Original: Formally, we wish to get thefollowing equation (5) for every l ∈ [ d ] .
Revised: Formally, we wish to get the following equation (7) for every l ∈ [ d ] .

Original: [Equation]
Revised: [Equation]

Original: When k is large, O ( β k 1 ) vanishes faster than O ( 1 √ k ) and thus Lemma 5.1 can be proved.
Revised: When k is large, O ( β nk 1 ) vanishes faster than O ( 1 √ k ) and thus Lemma 5.1 can be proved.

Original: In thefollowing context, we will carefully quantify the mismatch between m l,k,i and the ∂ l f ( x k, 0 ) .
Revised: In thefollowing context, we will carefully quantify the mismatch between m l,k,i and the ∂ l f ( x k, 0 ) .

Original: Wefind out that the error terms from successive epochs can be cancelled, which keeps the momentumroughly in the descent direction.
Revised: Wefind out that the error terms from successive epochs can be cancelled, which keeps the momentumroughly in the descent direction.

Original: To help readers understand our idea, we consider a toy model calledthe “color-ball " model as follows.
Revised: To help readers understand our idea, we introduce the “color-ball "model (of the 1st kind) as follows.

Original: The “color-ball " model. Consider a box containing two balls labeled with constant c 0 , c 1 ∈ R ,respectively.
Revised: The color-ball model of the 1st kind. Consider a box containing two balls labeled with constantc 0 , c 1 ∈ R , respectively.

Original: In each round (epoch), we randomly sample balls from the box without replacement,then we put them back.
Revised: In each round (epoch), we randomly sample balls from the box withoutreplacement, then we put them back.

Original: We denote the 1st sampled label in the k -th epoch as a k and the 2nd sampledone as b k ; a k , b k ∈ { c 0 , c 1 } .
Revised: We denote the 1st sampled label in the k -th epoch as a k and the2nd sampled one as b k ; a k , b k ∈ { c 0 , c 1 } .

Original: We define two random variables m 0 and m 1 as follows (assumeβ ∈ 
Revised: We define two random variables m 0 and m 1 as follows(assume β ∈ [0 , 1) ):

Original: [Equation]
Revised: [Equation]

Original: where m 0 ,k denotes the summand of m 0 in k -th epoch, similarly for m 1 ,k .
Revised: where m 0 ,k denotes the summand of m 0 in k -th epoch, similarly for m 1 ,k .

Original: Note that in each epoch,m 0 and m 1 share the same sample order.
Revised: Note that in each epoch,m 0 and m 1 share the same sample order.

Original: Further, we introduce the following deterministic constants.
Revised: Further, we introduce the following deterministic constants.

Original: f 1 =

Original: [Equation]

Original: f 1 ,k + β 2 + β(cid:124) (cid:123)(cid:122) (cid:125) f 1 ,k −· · · + β 2( k − 1) + β 2( k − 1)+

Original: cid:124) (cid:123)(cid:122) (cid:125) f 1 ,);

Original: [Equation]
Revised: [Equation]

Original: f

Original: [Equation]
Revised: [Equation]

Original: k

Original: [Equation]

Original: f

Original: [Equation]

Original: k −· · ·

Original: [Equation]

Original: [Equation] ( k [Equation] 1)+ (cid:124) (cid:123)(cid:122) (cid:125) f 0 ,); where f 0 ,k denotes the summand of f 0 in k -th epoch, similarly for f 1 ,k .
Revised: where f 0 ,k denotes the summand of f 0 in k -th epoch, similarly for f 1 ,k .

Original: Now we prove the following
Revised: Now we prove the following

Original: Lemma A.1. Lemma A.1.
Revised: Lemma A.1. Lemma A.1.

Original: In this color-ball example, we have
Revised: In the color-ball model of the 1st kind, we have(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) E (cid:34)

Original: [Equation]
Revised: [Equation]

Original: where the expectation is taken on all the possible draws. For the color-ball example with n ≥ 2 balls,we have
Revised: (cid:88) i =

Original: [Equation]
Revised: [Equation]

Original: i
Revised: (cid:88)

Original: [Equation]
Revised: [Equation]

Original: m
Revised: where the expectation is taken on all the possible draws.

Original: [Equation] i= [Equation] (cid:35)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
Revised: For the color-ball example with n ≥ 2 balls,we have(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) E (cid:34) n − 1 (cid:88) i

Revised: = m i − n − 1 (cid:88)

Original: [Equation]
Revised: [Equation]

Original: How is Lemma A.1 related to (5) ?
Revised: How is Lemma A.1 related to (7) ?

Original: In this color-ball toy example, m i mimics the possiblerealization of momentum up to the i -th inner loop in k -th epoch.
Revised: In this color-ball toy example, m i mimics the possiblerealization of momentum up to the i -th inner loop in k -th epoch.

Original: f i mimics the stochastic gradient∇
Revised: f i mimics the stochastic gradient∇

Original: f i ( x k, 0 ) .
Revised: f i ( x k, 0 ) .

Original: This is because we can expand ∇ f i ( x k, 0 ) into an infinite-sum sequence ∇ f i ( x k, 0 ) =(1 − β 1 ) ∇ f i ( x k, 0 ) (cid:80) ∞
Revised: This is because we can expand ∇ f i ( x k, 0 ) into an infinite-sum sequence ∇ f i ( x k, 0 ) =(1 − β 1 ) ∇ f i ( x k, 0 ) (cid:80) ∞

Original: j =0 β
Revised: j =0 β

Original: j 1 , which shares a similar structure as f i .
Revised: j 1 , which shares a similar structure as f i .

Original: In this sense, Lemma A.1 mayprovide ideas to prove (5).
Revised: In this sense, Lemma A.1 mayprovide ideas to prove (7).

Original: Nevertheless, there are still gap between these two, we will explain thegap later.
Revised: Nevertheless, there are still gap between these two, we will explain thegap later.

Original: Proof.
Revised: Proof.

Original: We use E k [ · ] to denote the conditional expectation given all the history up to the beginningof k -th epoch.
Revised: We use E k [ · ] to denote the conditional expectation given all the history up to the beginningof k -th epoch.

Original: Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 1 i =

Revised: Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 1 i =

Original: m i (cid:105) .
Revised: m i (cid:105) .

Original: Since all the history beforek -th epoch is fixed, we relegate this part to later discussion and first focus on the expectation of 536 (cid:80)i =

Revised: Since all the history beforek -th epoch is fixed, we relegate this part to later discussion and first focus on the expectation of 528 (cid:80)i =

Original: m i,k .
Revised: m i,k .

Original: As shown in Figure 5 (upper part), there are 2 possible realization.
Revised: As shown in Figure 4 (upper part), there are 2 possible realization.

Original: With the help of Figure 5 (upper part), we have the following result.
Revised: With the help of Figure 4 (upper part), we have the following result.

Original: [Equation]
Revised: [Equation]

Original: Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 1 i =

Revised: Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 1 i =

Original: m i − (cid:80) 1 i =

Revised: m i − (cid:80) 1 i =

Original: f i (cid:105) .
Revised: f i (cid:105) .

Original: Using the similarstrategy as (6), we have the following result.
Revised: Using the similarstrategy as (8), we have the following result.

Original: The calculation is illustrated in Figure 5 (lower part)and Figure 6.
Revised: The calculation is illustrated in Figure 4 (lower part)and Figure 5. [Figure]

Original: [Equation]
Revised: [Equation]

Original: i =
Revised: [Figure] [Equation] [Figure]

Original: [Equation]
Revised: [Equation]

Original: i
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: i = f i
Revised: [Figure] [Equation] [Figure]

Original: [Equation]
Revised: [Equation]

Original: k + 1 (cid:88)i
Revised: [Figure] (cid:88)

Original: = m i,k − 1 − 1 (cid:88) i = f
Revised: [Figure]

Original: i,k −(cid:35)(cid:41)
Revised: (cid:35)(cid:41)

Original: [Equation]
Revised: [Equation]

Original: (cid:40)
Revised: [Figure] [Figure] [Equation] [Figure]

Original: [Equation]
Revised: [Equation]

Original: i = m i,k −(cid:41) − 1 (cid:88) i = f i,k −
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: We observe that only the highest order term remains in the calculation.
Revised: We observe that only the highest order term remains in the calculation.

Original: Repeat this process untilk = 1 , we will get the results in Lemma A.1.
Revised: Repeat this process untilk = 1 , we will get the results in Lemma A.1.

Original: The above analysis also holds for general n ≥ 2 .
Revised: The above analysis also holds for general n ≥ 2 .

Original: Formore evidence, please refer to Appendix G.1 where we further provide a derivation for n = 3 .

Original: The gap between Lemma A.1 and Equation (5) .
Revised: The gap between Lemma A.1 and Equation (7) .

Original: Lemma A.1 shows the key idea of provingequation (5).
Revised: Lemma A.1 shows the key idea of provingequation (7).

Original: However, due to its idealized setting, the color-ball toy example is still far from our realgoal (5).
Revised: However, due to its idealized setting, the color-ball toy example is still far from our realgoal (

Original: We list some of the gaps here.
Revised: We list some of the gaps here.

Original: (i) In each possible trajectory: x k,i is changing with k and i , while the balls are fixed in the color-ball example.
Revised: • In each possible trajectory: x k,i is changing with k and i , while the balls are fixed in the color-ballexample.

Original: (ii) When taking expectation: x is changing in different trajectory, while the balls is fixed in thecolor-ball example. (iii) x is a vector in R d while the label of the balls are constant in R .
Revised: • When taking expectation: x is changing in different trajectory, while the balls is fixed in thecolor-ball example.• x is a vector in R d while the label of the balls are constant in R .

Original: It requires extra technical lemmas to handle these differences.
Revised: It requires extra technical lemmas to handle these differences.

Original: Briefly speaking, the error caused by(i) and (ii) can be neutralized by diminishing stepsize.
Revised: We provide more discussions in Appendix G Solution to Issue II. We now discuss how to resolve Issue II mentioned in Section 5.

Revised: The solutioncontains two parts.

Revised: First, we need to prove (5).

Revised: Second, we derive (6) from (5).

Revised: Due to the limitedspace, we relegate the first part to Appendix G.1 (related to Lemma G.2).

Revised: Now, we discuss the secondpart: Assume we have (5), how do we use it to prove (6)?

Revised: To answer this question, we introduce thecolor-ball model of the 2nd kind.

Revised: The color-ball model of the 2nd kind. Consider the same setting as the color-ball model in Step 1 .

Revised: We define a sequence of real random variable { r j } kj =1 with the following relation.

Revised: [Equation]

Revised: Further, we assume r j is fixed when fixing the history up to j -th round.

Revised: The sequence { r j } kj =1mimics the sequence { ∂ l f ( x k, 0 ) √ v l,k, 0 } ∞ k =1 in (5).

Revised: Now

Original: Issue (iii) can be handled by component-wiseanalysis of x .
Revised: [Equation] we consider the following quantities. [Equation] [Equation]

Revised: [Equation]

Revised: k −

Revised: [Equation]

Revised: [Equation]

Revised: ( k −

Revised: [Equation]

Revised: β

Revised: [Equation]

Revised: ( k − 1)+ (cid:124) (cid:123)(cid:122) (cid:125) f 1 ,) 

Revised: 

Revised: [Equation]

Revised: [Equation]

Revised: (cid:125)

Revised: [Equation]

Revised: ( k −

Revised: [Equation]

Revised: ( k − 1)

Revised: [Equation]

Revised: (cid:124) (cid:123)(cid:122) (cid:125) f

Revised: [Equation]

Revised: )   ; .

Revised: We now prove the following Lemma A.2.

Revised: Lemma A.2.

Revised: Consider the color-ball model of the 2nd kind, we have

Revised: [Equation]

Revised: For general n = 1 , 2 , 3 , · · · , we have

Revised: [Equation]

Revised: Proof.

Revised: We only describe the proof idea here.

Revised: The proof contains the following 4 steps.

Revised: Step 2.1.

Revised: We firstly take E k [ · ] and thus r k can be viewed as a constant.

Revised: We use the color-ballprocedure as in Figure 6 to calculate E k (cid:104)

Revised: [Figure]

Revised: Step 2.2.

Original: {+We change E k (cid:104) r k (cid:80) 1 i =
+}
Revised: {+We change E k (cid:104) r k (cid:80) 1 i =
+}

Original: {+m i,k − 1 (cid:105) into E k (cid:104) r k − 1 (cid:80) 1 i =
+}
Revised: {+m i,k − 1 (cid:105) into E k (cid:104) r k − 1 (cid:80) 1 i =
+}

Revised: m i,k − 1 (cid:105) + Error; where Error =

Revised: O (1 / √ k ) .

Revised: Step 2.3.

Original: {+We calculate E k − 1 E k (cid:104) r k − 1 (cid:80) 1 i =
+}
Revised: {+We calculate E k − 1 E k (cid:104) r k − 1 (cid:80) 1 i =
+}

Revised: m i,k − 1 (cid:105) =

Original: {+E k − 1 (cid:104) r k − 1 (cid:80) 1 i =
+}
Revised: {+E k − 1 (cid:104) r k − 1 (cid:80) 1 i =
+}

Revised: m i,k − 1 (cid:105) .

Original: {+This part isillustrated in the top row in Figure 
+}
Revised: {+This part isillustrated in the top row in Figure 
+}

Revised: Step 2.4.

Revised: For the leftovers in Step 2.1 , we change all r k into r k − 1 .

Revised: Then we do the cancellation tocalculate E k − 1 E k (cid:104)

Original: {+r k − 1 (cid:80) 1 i =
+}
Revised: {+r k − 1 (cid:80) 1 i =
+}

Original: {+m i,k − 1 + r k − 1 (cid:80) 1 i =
+}
Revised: {+m i,k − 1 + r k − 1 (cid:80) 1 i =
+}

Original: {+m i,k − r k − 1 (cid:80) 1 i =
+}
Revised: {+m i,k − r k − 1 (cid:80) 1 i =
+}

Revised: f i (cid:105) .

Revised: This step is shown inthe second and third row in Figure 7.

Revised: Repeat this process to the 1st epoch.

Revised: We can prove the Lemma A.2.

Revised: We emphasize that here are still some gap between Lemma A.2 and our goal in (6) in Lemma 5.2.

Revised: First, we have the similar gap as discussed at the end of the Solution to Issue I .

Revised: Second, the conditionin Lemma 5.2 has requirement on the gradient norm, while this requirement is temporarily ignored inthe color-ball method of the 2nd kind.

Revised: It requires some technical lemmas to handle these gaps.

Original: For more details, please refer to the complete proof in Appendix G.
Revised: For more details, please refer to the complete proof in Appendix G.

Revised: B Experiments

Revised: To support our theory, we provide more simulations and real-data experiments.

Revised: All the experimentalsettings and hyperparameters are presented in Appendix B.2.

Revised: We aim to show:(I).

Revised: When β 2 is large, a large range of β 1 gives good performance, including all β 1 < √ β 2 .(II).

Original: B Some Potential Implications for Practical Users
Revised: When β 2 is small, a large range of β 1 performs relatively badly. [Figure] [Figure] [Figure]

Revised: [Figure]

Revised: g r a d i e n t n o r m

Revised: [Figure]

Revised: Counter-example (2) .

Revised: We run more experiments on (2) with initialization x = − 5 and n =5 , 10 , 15 , 20 .

Revised: We show the result of n = 20 in Figure 8 (a), (b); the rest are shown in Appendix B.1.

Revised: We find that: first, when β 2 is small, gradient norms increase dramatically.

Revised: This emphasizes the importance of discarding bounded gradient assumptions.

Revised: Second, when β 2 is large, Adam convergesto stationary points for β 1 < √ β 2 .

Revised: This matches our theory since (2) satisfies D 0 = 0 .

Revised: For functionswith D 0 > 0 , we further show that Adam can only converge to a bounded region in Appendix B.1.

Revised: MNIST and CIFAR-10.

Revised: As shown in Figure 1 (b)& (c) in Section 1, the training results matchboth claim (I) and (II).

Revised: In addition, there is a convex-shaped boundary on the transition from low lossto higher loss, this boundary roughly matches the condition in Theorem 3.1.

Revised: NLP.

Revised: We use Adam to train Transformer XL (Dai et al., 2019) on the WikiText-103 dataset (Merityet al., 2016).

Revised: This architecture and dataset is widely used in NLP tasks (e.g. (Howard & Ruder, 2018;

Revised: See et al., 2017)).

Revised: As shown in Figure 8 (c), the training results match both claim (I) and (II).

Revised: Batchsize and β 2 .

Revised: As shown in Figure 8 (d): on MNIST, smaller batchsize requires larger β 2 toreach small loss.

Revised: Since batchsize equals to (number of total sample)/(number of batches).

Revised: In thecontext of finite-sum setting with n summand, n usually stands for the number of batches (e.g., Inthe extreme case when batchsize = 1 , n equals to the number of total samples).

Revised: Therefore, smallerbatchsize brings larger n .

Revised: As such, Figure 8 (d) matches the message by Theorem 3.1: the thresholdof β 2 increases with n .

Revised: B.1 More experiments

Revised: Estimation on ρ in Theorem 3.1.

Revised: To ensure convergence, Theorem 3.1 requires β 2 ≥ γ 1 ( n ) =1 − O ((1 − β n 1 ) /n 2 ρ ) .

Revised: Now we estimate the constant ρ .

Revised: According to our definition in Remark G.7in Appendix G.1, ρ = ρ 1 ρ 2 ρ 3 , where ρ 1 , ρ 2 , ρ 3 are defined as follows.

Revised: [Equation]

Revised: These definitions can also be seen in Appendix F.1.

Revised: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .

Revised: However, ρ is highly dependent on the problem instance f ( x ) and training process.

Revised: We now estimatehow ρ changes with Adam’s trajectory on MNIST and CIFAR-10.

Revised: We use β 1 = 0 .

Revised: 9 , β 2 = 0 .

Revised: On both datasets, we set batchsize to be 64, which brings n = 937 on MNIST and n = 781 on

Revised: CIFAR-10.

Revised: We collect ρ 1 , ρ 2 , ρ 3 along the training process and estimate their distribution density.

Revised: The results are shown in Figure 9 and 10.

Revised: [Figure]

Original: {+On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
+}
Revised: {+On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
+}

Revised: ≈ O ( n ) ,ρ 3 < 0 .

Revised: 1 ≈ O (1 / √ n ) .

Revised: Therefore, ρ = ρ 1 ρ 2 ρ 3 ≈ O ( n ) .

Revised: [Figure]

Revised: More experiments on function (2) .

Revised: Figure 1 (a) shows the optimality gap after 50k iterationswhen n = 10 and initialization x = 1 .

Revised: Here, we provide more relevant experiments.

Revised: First, wheninitialized at x = 1 , we run experiments with n = 5 , 10 , 15 , 20 .

Revised: The results are shown in Figure 11.

Revised: We observe that the blue region shrinks as n increases.

Revised: This matches our conclusion in Theorem 3.1:when n increases, the convergent threshold of β 2 increases, which means we need larger β 2 to ensureconvergence.

Revised: It also matches the conclusion in Theorem 3.3: when n increases, the divergence regionwill expand (more evidence can be seen in Appendix E).

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: n

Revised: [Figure]

Revised: When initializing at x = − 5 , we further demonstrate that the gradient norm of f ( x ) can dramaticallyincrease.

Revised: All the setting is the same as that in Figure 11 except for the change of initialization.

Revised: Starting at x = − 5 , the algorithm can touch the “quadratic" side of function (2) where the gradient isunbounded.

Revised: The results are shown in Figure 12.

Revised: We observe a similar pattern as that in Figure 11.

Revised: Asa result, the gradient norm of f ( x ) is large in the left bottom corner.

Revised: We further plot the change of gradient norm along the iterations.

Revised: We pick β 1 = 0 .

Revised: 1 and β 2 =0 , 0 . 1 , 0 . 3 , 0 . 5 , 0 .

Revised: 99 to see the phase transition when increasing β 2 .

Revised: The result is shown in

Revised: Figure 13.

Revised: When β 2 is small, the gradient norm of f ( x ) increases rapidly along the iteration.

Revised: Mostof them are even much larger than the upper bound of color bar in Figure 12.

Revised: As a result, there is aphase transition from diverge to converge when increasing β 2 from 0 to 1.

Revised: [Figure]

Revised: Adam converges to a exact stationary point under SGC where D 0 = 0 .

Revised: Since function (2)satisfies SGC, we further provide empirical evidence that Adam converges to a exact stationary pointunder SGC.

Revised: We use function (2) n = 20 and initialization x = − 5 .

Original: {+We choose some large enough β 
+}
Revised: {+We choose some large enough β 
+}

Revised: [Figure]

Revised: to ensure the convergence.

Revised: As shown in Figure 14, We observe 0 gradient norm after Adam converges.

Revised: All the hyperparameter setting is the same as before.

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: Adam converges to a bounded region when D 0 > 0 .

Revised: Now we show that Adam may not convergeto an exact stationary point when D 0 > 0 .

Revised: Instead, it converges to a bounded region near thestationary point.

Revised: For this part, we re-state the example from (Shi et al., 2020).

Revised: Consider the followingfunction.

Revised: [Equation]

Revised: Summing up

Revised: f

Revised: j x ) we get

Revised: [Equation]

Revised: This f ( x ) is lower bounded by − a 2 9 with optimal solution x ∗ = 0 .

Original: {+At the optimal point x ∗ = 0 ,∇ f j ( x ∗ ) ̸ = 
+}
Revised: {+At the optimal point x ∗ = 0 ,∇ f j ( x ∗ ) ̸ = 
+}

Revised: so we have D 0 > 0 .

Revised: When running Algorithm 1 on this function, we observe that Adamwith diminishing stepsize does not converge to 0 gradient norm.

Revised: Instead, it converge to a boundedregion.

Revised: Further, the size of the region shrinks when β 2 increases.

Revised: These phenomena matches ourdiscussion in Remark 3.

Revised: The result is shown in Figure 15.

Revised: In the experiment, we choose β 1 = 0 .

Original: {+9 ,a = 3 , x 
+}
Revised: {+9 ,a = 3 , x 
+}

Revised: = − 2 and diminishing stepsize η k = 0 .

Revised: 1 √ k .

Revised: B.2 Experimental Settings

Revised: Here, we introduce our experimental settings.

Revised: • Experiments on function (2) .

Revised: We use Algorithm 1 with cyclic order f 0 , f 1 , f 2 and so on.

Revised: Wereport the optimality gap x − x ∗ after 50k iteration, or equivalently 50000 /n epochs

Revised: We useϵ =

Revised: 10 − 8 for numerical stability.

Revised: We use diminishing stepsize η k = 0 .

Revised: 1 / √ k , where k is the indexof epoch.

Revised: Unless otherwise stated, this setting applies to all the other experiments on function (2).

Revised: In Figure 1 (a), we use n = 10 and initialization x = 1 .

Revised: We will report more results with differentn and different initialization in Appendix B.1.

Revised: Instead, it only converge to a bounded region near the

Revised: point

Revised: • MNIST (Deng, 2012).

Revised: We use one-hidden-layer neural network with width =16.

Revised: We set batchsize

Revised: =1, weight decay =0, stepsize =0.0001 and train for 20 epochs.

Revised: We use ϵ = 10 − 8 for numericalstability.• CIFAR-10 (Krizhevsky et al., 2009).

Revised: We use ResNet-18 (He et al., 2016) as the architecture.

Revised: Wechoose batchsize =16, weight decay =5e-4 , initial stepsize=1e-3.

Revised: We use a stage-wise constantlearning rate scheduling with a multiplicative factor of 0.1 on epoch 30, 60 and 90.

Revised: We useϵ =

Revised: 10 − 8 for numerical stability.

Revised: For MNIST and CIFAR-10, larger batchsize will bring similar pattern as that in Figure 1, but thephase transition will occur at some smaller β 2 .

Revised: • NLP.

Revised: The WikiText-103 dataset is a collection of over 100 million tokens extracted from the setof verified ‘Good’ and ‘Featured’ articles on Wikipedia.

Revised: The base model of Transformer XLcontains 16 self-attention layers.

Revised: In each self-attention layer, there are 10 heads.

Revised: The encodingdimension of each head is set to be 41.

Revised: We set batchsize = 60, number of iteration =

Revised: 200k, andinitial stepsize = 0.00025.

Revised: We use cosine learning rate scheduler, which is a popular choice fortraining Transformer.

Revised: We use ϵ = 10 − 8 for numerical stability.

Revised: C Some Potential Implications for Practical Users

Original: Theorem 3.1 and Proposition 3.3 establish a clearer image on the relation between ( β 1 , β 2 ) and qualitative behavior of Adam, which may have certain implication for practical users.
Revised: Theorem 3.1 and Proposition 3.3 establish a clearer image on the relation between ( β 1 , β 2 ) and qualitative behavior of Adam, which may have certain implication for practical users.

Original: Many practitioners are still doing grid or random search over ( β 1 , β 2 ) , which could be costly.
Revised: Many practitionersare still doing grid or random search over ( β 1 , β 2 ) , which could be costly.

Original: The following advice maycut down a large portion of search space.
Revised: The following advice maycut down a large portion of search space.

Original: Suppose we start with some random point ( β ∗ 1 , β ∗ 2 ) , weprovide the following suggestions for tuning β 1 and β 2 .
Revised: Suppose we start with some random point ( β ∗ 1 , β ∗ 2 ) , weprovide the following suggestions for tuning β 1 and β 2 .

Original: Case 1: If ( β ∗ 1 , β ∗ 2 ) lies above the blue curve in Figure 2.
Revised: Case 1: If ( β ∗ 1 , β ∗ 2 ) lies above the blue curve in Figure 2.

Original: We point out two sub-cases.
Revised: We point out two sub-cases.

Original: First, if
Revised: First, if

Original: Adam with ( β ∗ 1 , β ∗ 2 ) diverges, then any points with “ β 1 < √ β 2 " and “ β 2 < β ∗ 2 " shall not be trieddue to the risk of divergence.
Revised: Adam with ( β ∗ 1 , β ∗ 2 ) diverges, then any points with “ β 1 < √ β 2 " and “ β 2 < β ∗ 2 " shall not be trieddue to the risk of divergence.

Original: Second, if Adam with ( β ∗ 1 , β ∗ 2 ) converges, then any points above thispoint will converge.
Revised: Second, if Adam with ( β ∗ 1 , β ∗ 2 ) converges, then any points above thispoint will converge.

Original: You may either fix β 1 = β ∗ 1 and increase β 2 , or fix β 2 = β ∗ 2 and try smaller β 1 .
Revised: You may either fix β 1 = β ∗ 1 and increase β 2 , or fix β 2 = β ∗ 2 and try smaller β 1 .

Original: Both ways have convergence guarantee.
Revised: Both ways have convergence guarantee.

Original: Case 2: If ( β ∗ 1 , β ∗ 2 ) lies below the blue curve in Figure 2.
Revised: Case 2: If ( β ∗ 1 , β ∗ 2 ) lies below the blue curve in Figure 2.

Original: We still discuss two sub-cases.
Revised: We still discuss two sub-cases.

Original: First, ifyou observe divergence at ( β ∗ 1 , β ∗ 2 ) , then do not further try any point on the left.
Revised: First, ifyou observe divergence at ( β ∗ 1 , β ∗ 2 ) , then do not further try any point on the left.

Original: We do not suggestexploring the points in below either, since the majority of them will still face the risk of divergence.
Revised: We do not suggestexploring the points in below either, since the majority of them will still face the risk of divergence.

Original: Instead, we suggest fix β 1 = β ∗ 1 and increase β 2 .
Revised: Instead, we suggest fix β 1 = β ∗ 1 and increase β 2 .

Original: Since our Theorem 3.1 applies for any β 1 < √ β 2 ,algorithm will converge when β 2 is large enough.
Revised: Since our Theorem 3.1 applies for any β 1 < √ β 2 ,algorithm will converge when β 2 is large enough.

Original: Second, if you observe convergence at ( β ∗ 1 , β ∗ 2 ) (which is also possible according to Figure 1 (a)),then we suggest: (i) either fix β 1 = β ∗ 1 and increase β 2 , or (ii) fix β 2 = β ∗ 2 and try smaller β 1 .
Revised: Second, if you observe convergence at ( β ∗ 1 , β ∗ 2 ) (which is also possible according to Figure 1 (a)),then we suggest: (i) either fix β 1 = β ∗ 1 and increase β 2 , or (ii) fix β 2 = β ∗ 2 and try smaller β 1 .

Original: Bothways push ( β 1 , β 2 ) into the region of Theorem 3.1 with convergence guarantee.
Revised: Bothways push ( β 1 , β 2 ) into the region of Theorem 3.1 with convergence guarantee.

Original: C More Discussion on Related Works.
Revised: D More Discussion on Related Works.D.1 More Related Works on the Convergence Analysis of Adam Family

Original: C.1 More Related Works on the Convergence Analysis of Adam Family

Original: New variants of Adam.
Revised: New variants of Adam.

Original: Ever since Reddi et al.
Revised: Ever since Reddi et al.

Original: (2018) pointed out the non-convergence issueof Adam, one active line of work has tried to design new variants of Adam that can be proved toconverge.
Revised: (2018) pointed out the non-convergence issueof Adam, one active line of work has tried to design new variants of Adam that can be proved toconverge.

Original: For instance, Zou et al.
Revised: For instance, Zou et al.

Original: (2019); Gadat & Gavra (2020)
Revised: (2019); Gadat & Gavra (2020)

Original: ; Chen et al.
Revised: ; Chen et al.

Revised: (2018b, 2021)

Original: replacethe constant hyperparameters by iterate-dependent ones e.g. β 1 t or β 2 t .
Revised: replacethe constant hyperparameters by iterate-dependent ones e.g. β 1 t or β 2 t .

Original: AMSGrad (Reddi et al.,2018) and AdaFom (Chen et al., 2018) modify { v t } to be an non-decreasing sequence.
Revised: AMSGrad (Reddi et al.,2018) and AdaFom (Chen et al., 2018b) modify { v t } to be an non-decreasing sequence.

Original: Similarly,
Revised: Similarly,

Original: AdaBound (Luo et al., 2019) impose lower and upper bounds on { v t } to prevent the effective stepsizefrom vanishing or exploding.
Revised: AdaBound (Luo et al., 2019) impose lower and upper bounds on { v t } to prevent the effective stepsizefrom vanishing or exploding.

Original: Zhou et al.
Revised: Zhou et al.

Original: (2018) also adopt a new estimate of v t to correct thebias.
Revised: (2018b) also adopt a new estimate of v t to correct thebias.

Original: There are also attempts to combine Adam with Nesterov momentum (Dozat, 2016)
Revised: There are also attempts to combine Adam with Nesterov momentum (Dozat, 2016)

Original: as well aswarm-up techniques (Liu et al., 2020a).
Revised: as well aswarm-up techniques (Liu et al., 2020a).

Revised: Padam (Chen et al., 2018a) also introduce a partial adaptiveparameter to improve the generalization performance.

Original: As reported in the literature, these variants often reach goodperformance in many applications.
Revised: As reported in the literature, these variantsoften reach good performance in many applications.

Original: Some more discussions on (Guo et al., 2021) and (Huang et al., 2021).
Revised: There are also some works providing theoreticalanalysis on the variants of Adam. For instance, Zhou et al. (2018a) study the convergence of AdaGradand AMSGrad with zero and non-zero β 1 under bounded gradient condition. Gadat & Gavra (2020)study the asymptotic behavior of a subclass of adaptive gradient methods from landscape point ofview.

Revised: Their analysis applies to Adam-variants with β 1 = 0 and β 2 increasing along the iterates (itcould also be understood as RMSProp with increasing β 2 ).

Revised: When this script is under review, a newwork (Iiduka, 2022) appear on line.

Revised: Iiduka (2022) analyze the convergence of AMSGrad by relaxingthe Lipschitz-gradient condition.

Revised: However, their analysis requires extra conditions on both boundedgradient and bounded domain.

Revised: Some more discussions on (Guo et al., 2021) and (Huang et al., 2021).

Original: Here, we discuss more ontwo recent works Guo et al.
Revised: Here, we discuss more ontwo recent works Guo et al.

Original: and Huang et al.
Revised: and Huang et al.

Original: As mentioned in Section 2.2, they bothrequires some extra conditions.
Revised: As mentioned in Section 2.2, they bothrequire some extra conditions.

Original: First, both Guo et al.
Revised: First, both Guo et al.

Original: and Huang et al.
Revised: and Huang et al.

Original: (2021) requires boundedgradient assumption.
Revised: (2021) requires boundedgradient assumption.

Original: This can be seen in Assumption 2 in (Guo et al., 2021).
Revised: This can be seen in Assumption 2 in (Guo et al., 2021).

Original: In (Huang et al.
Revised: In (Huang et al.

Original: , 2021),they require bounded iterates ( their Theorem 1) or change Adam into AdaBound (Luo et al.
Revised: , 2021),they require bounded iterates ( their Theorem 1) or change Adam into AdaBound (Luo et al.

Original: , 2019)by clipping (their Remark 2, Corollary 1).
Revised: , 2019)by clipping (their Remark 2, Corollary 1).

Original: Both settings inherent boundedness on gradient.
Revised: Both settings inherent boundedness on gradient.

Original: Besides bounded gradient, both (Huang et al., 2021) and (Guo et al., 2021) requires 1 / ( √ v t + ϵ ) ≤ C u .
Revised: Besides bounded gradient, both (Huang et al., 2021) and (Guo et al., 2021) requires 1 / ( √ v t + ϵ ) ≤ C u .

Original: This condition is stated in Assumption 2 in (Guo et al., 2021) and Assumption 3 in (Huang et al.,2021)
Revised: This condition is stated in Assumption 2 in (Guo et al., 2021) and Assumption 3 in (Huang et al.,2021)

Original: (they presented it as H t ⪰ ρI d ≻ 0 , where matrix H t = diag ( √ v t + ϵ ) ).
Revised: (they presented it as H t ⪰ ρI d ≻ 0 , where matrix H t = diag ( √ v t + ϵ ) ).

Original: Combining these twoconditions, the effective stepsize of Adam will be bounded in certain interval 1 √ v t + ϵ ∈
Revised: Combining these twoconditions, the effective stepsize of Adam will be bounded in certain interval 1 √ v t + ϵ ∈

Original: [ C l , C u ] .
Revised: [ C l , C u ] .

Original: Suchboundedness condition changes Adam into AdaBound (Luo et al., 2019) and thus cannot explain theobservations of Adam in Section 1.
Revised: Suchboundedness condition changes Adam into AdaBound (Luo et al., 2019) and thus cannot explain theobservations of Adam in Section 1.

Original: C.
Revised: D.2 A Brief Introduction to (Reddi et al., 2018)

Original: An Brief Introduction to (Reddi et al., 2018)

Original: Here, we re-state two counter examples by (Reddi et al., 2018).
Revised: Here, we re-state two counter examples by (Reddi et al., 2018).

Original: For the consistence of notation, wewill re-state their results under our notation in the full script.
Revised: For the consistence of notation, wewill re-state their results under our notation in the full script.

Original: They consider the convex problem((Reddi et al., 2018)): min (cid:80) n − 1 i =

Revised: They consider the convex problem((Reddi et al., 2018)): min (cid:80) n − 1 i =

Original: f i ( x ) where x ∈ [ − 1 , 1] , n ≥ 3 :
Revised: f i ( x ) where x ∈ [ − 1 , 1] , n ≥ 3 :f i ( x ) = (cid:26) nx, for i =− x, otherwise, (9)

Original: [Equation]
Revised: Note that (9) satisfy both Assumption 2.1 and 2.2 (with D 1 = n 2 + n − 1 and D 0 = 0 ), so ourassumptions do not rule out this counter-example a priori.

Original: Since it is a constrained problem with feasible set x ∈ [ − 1 , 1] , the optimal solution is x ∗ = − 1 .
Revised: This is a constrained problem with feasibleset x ∈ [ − 1 , 1] , the optimal solution is x ∗ = − 1 .

Original: In(Reddi et al., 2018), they use notation C rather than n .
Revised: Since they consider constrained problems, theirclaimed “divergence" actually means the iterates will stay in a huge region with the size of wholefeasible set. Here, we call it “non-convergence" to distinguish from our result of “diverge to infinity"in Proposition 3.3.

Original: They consider sampling f i in the cyclic order: f 0 , f 1 , f 2 .
Revised: They consider sampling f i in the cyclic order: f 0 , f 1 , f 2 .

Original: In (Reddi et al., 2018), Function (7) ispresented as an online optimization problem.
Revised: In (Reddi et al., 2018), Function (9) ispresented as an “online optimization problem with non-zero average regret".

Original: We choose to use the form of (7) since it is moreconsistent with our notation in Algorithm 1.
Revised: We choose to use theform of (9) since it is more consistent with our notation in Algorithm 1.

Original: We re-state their results as follows.
Revised: We re-state their results asfollows.

Original: Theorem C.1 (Theorem 2 in (Reddi et al., 2018)) .
Revised: Theorem D.1 (Theorem 2 in (Reddi et al., 2018)) .

Original: For any fixed ( β 1 , β 2 ) s.t.
Revised: For any fixed ( β 1 , β 2 ) s.t.

Original: β 1 < √ β 2 , there existsfunction (7) with large enough n , s.t.
Revised: β 1 < √ β 2 , there existsfunction (9) with large enough n , s.t.

Original: Adam will converge to highly sub-optimal solution x = 1 .
Revised: Adam will converge to highly sub-optimal solution x = 1 .

Original: We briefly re-state the non-convergent condition for this Theorem.
Revised: We briefly re-state the non-convergent condition for this Theorem.

Original: As stated in Equation (7), Appendix
Revised: As stated in Equation (7), Appendix

Original: B in (Reddi et al., 2018), for every fixed ( β 1 , β 2 ) , they need a “constant n that depends on β 1 and β 2 ".
Revised: B in (Reddi et al., 2018), for every fixed ( β 1 , β 2 ) , they need a “constant n that depends on β 1 and β 2 ".

Original: As such, they requires different n to cause non-convergence on different ( β 1 , β 2 ) .
Revised: As such, they require different n to cause non-convergence on different ( β 1 , β 2 ) .

Original: So the consideredfunction class is constantly changing.
Revised: So the consideredfunction class is constantly changing.

Original: For completeness, we further re-state Theorem 1 in (Reddi et al., 2018).
Revised: For completeness, we further re-state Theorem 1 in (Reddi et al., 2018).

Original: Theorem C.2 (Theorem 1 in (Reddi et al., 2018)) .
Revised: Theorem D.2 (Theorem 1 in (Reddi et al., 2018)) .

Original: For function (7) , when β 1 = 0 and β 2 = 1 / ( n 2 +1) ,
Revised: For function (9) , when β 1 = 0 and β 2 = 1 / ( n 2 +1) ,

Original: Adam will converge to highly sub-optimal solution x = 1 .
Revised: Adam will converge to highly sub-optimal solution x = 1 .

Original: This theorem considers choosing ( β 1 , β 2 ) after n .
Revised: This theorem considers choosing ( β 1 , β 2 ) after n .

Original: However, this result only show divergence on asingle point ( β 1 , β 2 ) =
Revised: However, this result only show non-convergence ona single point ( β 1 , β 2 ) =

Original: (0 , 1 / ( n 2 + 1)) .
Revised: (0 , 1 / ( n 2 + 1)) .

Original: This point lies somewhere on the left boundary of Figure 3.
Revised: This point lies somewhere on the left boundary of Figure3.

Original: It seems unclear how Adam’s behavior would change as we change the ( β 1 , β 2 ) to anywhere else.
Revised: It seems unclear how Adam’s behavior would change as we change the ( β 1 , β 2 ) to anywhere else.

Original: D Experimental Settings and More Experiments

Original: D.1 Experimental Settings

Original: Here, we introduce our experimental settings that are omitted in the full script.

Original: • Experiments on function (2) .

Original: We use Algorithm 1 with cyclic order f 0 , f 1 , f 2 and so on.

Original: We report the optimality gap x − x ∗ after 50k iteration, or equivalently 50000 /n epochs.

Original: Weuse ϵ

Original: = 10 − 8 for numerical stability.

Original: We use diminishing stepsize η k = 0 .

Original: 1 / √ k , where k is

Original: the index of epoch.

Original: Unless otherwise stated, this setting applies to all the other experiments

Original: on function (2).

Original: In Figure 1 (a), we use n = 10 and initialization x = 1 .

Original: We will report more results withdifferent n and different initialization in Appendix D.2.

Original: • MNIST (Deng, 2012).

Original: We use one-hidden-layer neural network with width =16.

Original: We setbatchsize =1, weight decay =0, stepsize =0.0001 and train for 20 epochs.

Original: We use ϵ = 10 − 8for numerical stability.

Original: • CIFAR-10 (Krizhevsky et al., 2009).

Original: We use ResNet-18 (He et al., 2016) as the architecture.

Original: We choose batchsize =16, weight decay =5e-4 , initial stepsize=1e-3.

Original: We use a stage-wise

Original: constant learning rate scheduling with a multiplicative factor of 0.1 on epoch 30, 60 and 90.We use ϵ = 10 − 8 for numerical stability.

Original: For MNIST and CIFAR-10, larger batchsize will bring similar pattern as that in Figure 1,but the phase transition will occur at some smaller β

Original: • NLP. The WikiText-103 dataset is a collection of over 100 million tokens extracted from theset of verified ‘Good’ and ‘Featured’ articles on Wikipedia. The base model of Transformer XL contains 16 self-attention layers.
Revised: E Proof of Proposition 3.3

Original: In each self-attention layer, there are 10 heads.

Original: Theencoding dimension of each head is set to be 41.

Original: We set batchsize = 60, number of iteration= 200k, and initial stepsize = 0.00025. We use cosine learning rate scheduler, which is apopular choice for training Transformer. We use ϵ = 10 − 8 for numerical stability. D.2 More experiments
Revised: We restate our counter-example here.

Original: Estimation on ρ in Theorem 3.1.

Original: To ensure convergence, Theorem 3.1 requires β 2 ≥ γ 1 ( n ) =1 − O ((1 − β n 1 ) /n 2 ρ ) .

Original: Now we estimate the constant ρ .

Original: According to our definition in Remark G.8in Appendix G.1, ρ = ρ 1 ρ 2 ρ 3 , where ρ 1 , ρ 2 , ρ 3 as defined as follows.

Original: [Equation]

Original: ρ 2 ≥ | max i ∂ l f

Original: [Equation]

Original:  x k [Equation] [Equation] |1 n (cid:80) {+n − 1 i =
+}
Revised: Consider f ( x ) = (cid:80) {+n − 1 i =
+}

Original: [Equation]

Original: f [Equation] 2 ; ρ 3 ≥ [Equation]
Revised: f i ( x ) for x ∈ R , we define f i ( x ) asfollows:

Original: =

Original: [Equation]

Original: f

Original: [Equation]

Original: [Equation]

Original: (cid:113) 1 n (cid:80) ni =

Original: [Equation]

Original: ∂ l f

Original: [Equation]
Revised: [Equation]

Original: x k
Revised: x <

Original: [Equation]

Original: 0 ) |.

Original: These definitions can also be seen in Appendix F.1.

Original: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .

Original: However, ρ is highly dependent on the problem instance f ( x ) and training process.

Original: We now estimatehow ρ changes with Adam’s trajectory on MNIST and CIFAR-10.

Original: We use β 1 = 0 .

Original: 9 , β 2 = 0 .

Original: On both datasets, we set batchsize to be 64, which brings n = 937 on MNIST and n = 781 on

Original: CIFAR-10.

Original: We collect ρ 1 , ρ 2 , ρ 3 along the training process and estimate their distribution density.

Original: The results are shown in Figure 7 and 8.

Original: [Figure]

Original: [Figure]

Original: [-On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
-]
Revised: [-On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
-]

Original: ≈ O ( n ) ,ρ 3 < 0 .

Original: 1 ≈ O (1 / √ n ) .

Original: Therefore, ρ = ρ 1 ρ 2 ρ 3 ≈ O ( n ) .

Original: More experiments on function (2) .

Original: Figure 1 (a) shows the optimality gap after 50k iterationswhen n = 10 and initialization x = 1 .

Original: Here, we provide more relevant experiments.

Original: First, wheninitialized at x = 1 , we run experiments with n = 5 , 10 , 15 , 20 .

Original: The results are shown in Figure 9.

Original: We observe that the blue region shrinks as n increases.

Original: This matches our conclusion in Theorem 3.1:when n increases, the convergent threshold of β 2 increases, which means we need larger β 2 to ensureconvergence.

Original: It also matches the conclusion in Theorem 3.3: when n increases, the divergence regionwill expand (more evidence can be seen in Appendix E).

Original: When initializing at x = − 5 , we further demonstrate that the gradient norm of f ( x ) can dramaticallyincrease.

Original: All the setting is the same as that in Figure 9 except for the change of initialization.

Original: Starting at x = − 5 , the algorithm can touch the “quadratic" side of function (2) where the gradient isunbounded.

Original: The results are shown in Figure 10.

Original: We observe an similar pattern as that in Figure 9.

Original: Asa result, the gradient norm of f ( x ) is large in the left bottom corner.

Original: We further plot the change of gradient norm along the iterations.

Original: We pick β 1 = 0 .

Original: 1 and β 2 =0 , 0 . 1 , 0 . 3 , 0 . 5 , 0 .

Original: 99 to see the phase transition when increasing β 2 .

Original: The result is shown in

Original: Figure 11.

Original: When β 2 is small, the gradient norm of f ( x ) increases rapidly along the iteration.

Original: Mostof them are even much larger than the upper bound of color bar in Figure 10.

Original: As a result, there is aphase transition from diverge to converge when increasing β 2 from 0 to 1.

Original: [Figure]

Original: [Figure]

Original: Adam converges to a exact stationary point under SGC where D 0 = 0 .

Original: Since function (2)satisfies SGC, we further provide empirical evidence that Adam converges to a exact stationary pointunder SGC.

Original: We use function (2) n = 20 and initialization x = − 5 .

Original: We choose some large enough β 2to ensure the convergence.

Original: As shown in Figure 12, We observe 0 gradient norm after Adam converges.

Original: All the hyperparameter setting is the same as before.

Original: Adam converges to a bounded region when D 0 > 0 .

Original: Now we show that Adam may not convergeto an exact stationary point when D 0 > 0 .

Original: Instead, it converge to a bounded region near the stationarypoint.

Original: For this part, we re-state the example from (Shi et al., 2020).

Original: Consider the following function.

Original: [Equation]

Original: Summing up f j ( x ) we get

Original: [Equation]

Original: This f ( x ) is lower bounded by − a 2 9 with optimal solution x ∗ = 0 .

Original: In (Shi et al., 2020), they try

Original: SGD, RMSProp and AMSGrad on this function, the result is shown in Figure 13.

Original: They observe that

Original: [Figure]

Original: (a

Original: [Equation]

Original: [Figure]

Original: [Figure]

Original: [Equation]

Original: [Figure]

Original: n

Original: [Figure]

Original: e n t n o r m [0.1,0.99][0.1,0.9][0.1,0.7][0.1,0.5][0.1,0.3] [0.1,0.1][0.1,0] (

Original: [Figure]

Original: ) n =

Original: [Figure]

Original: RMSProp with diminishing stepsize may not converge to the optimal solution when D 0 > 0 , even ifthey set the initialization at x = x ∗ = 0 .

Original: See more description in (Shi et al., 2020, Appendix A.4) [Figure] E Proof of Proposition 3.  We restate our counter-example here.
Revised: [Equation]

Original: [-Consider f ( x ) = (cid:80) n − 1 i =
 f i ( x )-] forx ∈ R , we define f i ( x ) asfollows:
Revised: [-Consider f ( x ) = (cid:80) n − 1 i =
 f i ( x )-] for

Original: [Equation]
Revised: [Equation]

Original: Summing up all the f i ( x ) , we can see that
Revised: Summing up all the f i ( x ) , we can see thatf ( x ) =

Original: [Equation] is a convex smooth function with optimal solution x ∗ = − 2 and optimal value f ( x ∗ ) = − 3 / 2 .
Revised: (cid:26) x, x ≥ −12 ( x + 2) 2 − 32 , x < −is a convex smooth function with optimal solution x ∗ = − 2 and optimal value f ( x ∗ ) = − 3 / 2 .

Original: However, we are going to show that, for any fixed n > 2 , there exists an orange region shown in
Revised: However, we are going to show that, for any fixed n > 2 , there exists an orange region shown in

Original: Figure 2, s.t., Adam with any  β 1 - β 2  combination in the yellow region converges to x = ∞ ratherthan the optimal solution x = − 2 , causing the divergence.
Revised: Figure 2, s.t., Adam with any ( β 1 , β 2 ) combination in the yellow region diverge to x = ∞ ratherthan the optimal solution x = − 2 , causing the divergence.

Original: Now we introduce the formal statement of
Revised: Now we introduce the formal statement of

Original: Proposition 3.3.
Revised: Proposition 3.3.

Original: Proposition E.1.
Revised: Proposition E.1.

Original: (Formal statement of Proposition 3.3.) Consider the convex function (8) for afixed n .
Revised: (Formal statement of Proposition 3.3.) Consider the convex function (10) for afixed n .

Original: Starting at the initialization x = 1 and initial stepsize η 1 , the iterates of Adam diverge toinfinity if the following holds:
Revised: Starting at the initialization x = 1 and initial stepsize η 1 , the iterates of Adam diverge toinfinity if the following holds:

Original: [Equation]
Revised: [Equation]

Original: The analysis is motivated by that in (Reddi et al., 2018, Theorem 1).
Revised: The analysis is motivated by that in (Reddi et al., 2018, Theorem 1).

Original: However, (Reddi et al., 2018,
Revised: However, (Reddi et al., 2018,

Original: Theorem 1) considers a simplified case with β 1 = 0 .
Revised: Theorem 1) considers a simplified case with β 1 = 0 .

Original: Here, we consider non-zero β 1 , especially forthose β 1 > √ β 2 .
Revised: Here, we consider non-zero β 1 , especially forthose β 1 > √ β 2 .

Original: To show the divergence, we aim to prove the following claim: (we denote x k,i asthe value of x at the k -th outer loop and i -th inner loop )
Revised: To show the divergence, we aim to prove the following claim: (we denote x k,i asthe value of x at the k -th outer loop and i -th inner loop )

Original: Claim: for any fixed n > 2 , there exists an orange region shown in Figure 2 s.t., Adam with anyβ 1 - β 2 combination in the orange region givesx k +1 , 0 > 1 as long as x k, 0 = 1 .
Revised: Claim: for any fixed n > 2 , there exists an orange region shown in Figure 2 s.t., Adam with anyβ 1 - β 2 combination in the orange region gives

Original: Since the gradient stays constant when x > 1 , so x will go to infinity if the claim holds.
Revised: Since the gradient stays constant when x > 1 , so x will go to infinity if the claim holds.

Original: To prove thisclaim, we only need to analyze the trajectory of Adam within one particular outer loop, e.g., the k -thouter loop.
Revised: To prove thisclaim, we only need to analyze the trajectory of Adam within one particular outer loop, e.g., the k -thouter loop.

Original: We will show that x k +1 , 0 > 1 if this outer loop is initialized with x k, 0 = 1 .
Revised: We will show that x k +1 , 0 > 1 if this outer loop is initialized with x k, 0 = 1 .

Original: Similarly as(Reddi et al., 2018), we assume f i ( x ) are sampled in the order of f 0 ( x ) , f 1 ( x ) , · · · , f n − 1 ( x ) withinthe k -th outer loop.
Revised: Similarly as(Reddi et al., 2018), we assume f i ( x ) are sampled in the order of f 0 ( x ) , f 1 ( x ) , · · · , f n − 1 ( x ) withinthe k -th outer loop.

Original: Now let us prove the claim.
Revised: Now let us prove the claim.

Original: For function (8), the update rule of Adam is shown as follows.
Revised: For function (10), the update rule of Adam is shown as follows.

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: We decompose the total movement (cid:80) n − 1 i =

Revised: We decompose the total movement (cid:80) n − 1 i =

Original: δ k,i into three terms as follows.
Revised: δ k,i into three terms as follows.

Original: [Equation]
Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Original:  We will show that for some β 1 and β 2 : ( I ) , ( II ) > 0 and ( III ) < 0 .
Revised: [Equation] [Table] [Equation] We will show that for some β 1 and β 2 : ( I ) , ( II ) > 0 and ( III ) < 0 .

Original: However, ( I ) and ( II )outweigh ( III ) , causing the divergence.
Revised: However, ( I ) and ( II )outweigh ( III ) , causing the divergence.

Original: First, we show that m k − 1 ,n − 1 < 0 when β 1 is small.
Revised: First, we show that m k − 1 ,n − 1 < 0 when β 1 is small.

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: when β 1 is small,we have (1 − β n − 1) > (1 − β 1 ) β n − 1n , which implies − m k − 1 ,n − 1 > 0 .
Revised: when β 1 is small,we have (1 − β n − 1) > (1 − β 1 ) β n − 1n , which implies − m k − 1 ,n − 1 > 0 .

Original: For thesechoices of β 1 , we have ( I ) > 0 .
Revised: For thesechoices of β 1 , we have ( I ) > 0 .

Original: Now we derive a lower bound for ( II ) .
Revised: Now we derive a lower bound for ( II ) .

Original: [Equation]
Revised: [Equation]

Original: The inequality is due to the fact that v k, 0 ≤ n 2 .
Revised: The inequality is due to the fact that v k, 0 ≤ n 2 .

Original: Since β j 2 n 2 is small when β 2 is small and j is closeto n , there exists some small β 2 such that β j 2 n 2 ≤ 0 .
Revised: Since β j 2 n 2 is small when β 2 is small and j is closeto n , there exists some small β 2 such that β j 2 n 2 ≤ 0 .

Original: 1 for at least one j < n .
Revised: 1 for at least one j < n .

Original: For these small enoughβ 2 , we keep the summand with j ≥ log β 2 (0 .
Revised: For these small enough β 2 , we keep the summand with j ≥ log β 2 (0 .

Original: 1 /n 2 ) and drop the rest.
Revised: 1 /n 2 ) and drop the rest.

Original: We have the following lowerbound for ( II ) .
Revised: We have the following lowerbound for ( II ) .

Original: [Equation]
Revised: [Equation]

Original: However, these lower bound only holds for the small β 2 .
Revised: However, this lower bound only holds for the small β 2 .

Original: With simple modification, we derive anuniversal lower bound of ( II ) for any β 2 ∈ (0 , 1) .
Revised: With simple modification, we derive auniversal lower bound of ( II ) for any β 2 ∈ (0 , 1) .

Original: [Equation]
Revised: [Equation]

Original: Now we derive a lower bound for ( III ) .
Revised: Now we derive a lower bound for ( III ) .

Original: [Equation]
Revised: [Equation]

Original: The remaining step is to show for small enough step size η 1 , the iterates will stay in the linear region,thus the above gradient holds for all iterates in the trajectory.
Revised: The remaining step is to show for small enough step size η 1 , the iterates will stay in the linear region,thus the above gradient holds for all iterates in the trajectory.

Original: As we initial x as x 0 = 1 , if for all m < n and k > 0 , (cid:80) mi [-=
-] δ i,k ≥ − 2 , we can conclude all iteratesin  trajectory stay 
Revised: As we initial x as x 0 = 1 , if for all m < n and k > 0 , (cid:80) mi [-=
-]=0 δ k,i ≥ − 2 , we can conclude all iteratesin the trajectory stay in the linear region.

Original: Because it holds that m k − 1 ,n − 1 ≤ 0 and v k − 1 ,n − 1 ≥ 0 , we have 
Revised: Because it holds that m k − 1 ,n − 1 ≤ 0 and v k − 1 ,n − 1 ≥ 0 , we have the following result after droppingmost negative terms in the definition of δ k,i :

Original: [Equation]
Revised: [Equation]

Original: Therefore, to make (cid:80) mi =

Revised: Therefore, to make (cid:80) m i =

Revised: δ k,i ≥

Revised: [Equation]

Original: δ i,k ≥ (cid:80) mi =0 − η  1(1 − β 1 ) β i 1 √ (1 − β 2 ) β n 2 ≥ − 2 , we have
Revised: mi =0 − η [Equation] 1 − β [Equation]

Revised: β

Revised: [Equation]

Revised: 1 √

Revised: [Equation]

Revised: 1 −

Revised: [Equation]

Revised: β n 2 ≥ − 2 , we have

Original: [Equation]
Revised: [Equation]

Original: To show the divergence, we want to show that there exists some β 1 and β 2 s.t.
Revised: To show the divergence, we want to show that there exists some β 1 and β 2 s.t.

Original: the both of thefollowing conditions hold:
Revised: the both of thefollowing conditions hold:

Original: [Equation]
Revised: [Equation]

Original: The proof of Theorem 3.3 is completed.
Revised: The proof of Theorem 3.3 is completed.

Original: With the help of Python, we visualize the region where( C1 ) and ( C2 ) hold.
Revised: With the help of Python, we visualize the region where( C1 ) and ( C2 ) hold.

Original: The results are shown in Figure 14.
Revised: The results are shown in Figure 16.

Original: We use orange color to indicate the regionwhere ( C1 ) holds.
Revised: We use orange color to indicate the regionwhere ( C1 ) holds.

Original: White color is used for the counter part.
Revised: White color is used for the counter part.

Original: As for ( C2 ) , we use the gray verticalline to indicate the line where (1 − β n − 1) =
Revised: As for ( C2 ) , we use the gray verticalline to indicate the line where (1 − β n − 1) =

Original: (1 − β 1 ) β n − 1n .
Revised: (1 − β 1 ) β n − 1n .

Original: Note that there are two solutions tothis equation: one solution is β 1 = 1 and the other solution lies in 0 < β 1 < 1 , this is why there aretwo vertical lines in the figure.
Revised: Note that there are two solutions tothis equation: one solution is β 1 = 1 and the other solution lies in 0 < β 1 < 1 , this is why there aretwo vertical lines in the figure.

Original: ( C2 ) holds on the left hand side of the left gray vertical line.
Revised: ( C2 ) holds on the left hand side of the left gray vertical line.0.

Revised: 0.8 beta 0.

Revised: b e t a(a) n =0.

Revised: 0.8 beta 0.

Revised: b e t a(b) n =0.

Revised: 0.8 beta 0.

Original:  The intersection of two regions will be the region where Adam diverges, which is actually the orangeregion in Figure 14.
Revised: b e t a(c) n =0. 0.8 beta 0. b e t a(d) n = The intersection of two regions will be the region where Adam diverges, which is actually the orangeregion in Figure 16.

Original: As n increases, ( C2 ) holds for a wide range of β 1 , so the grey vertical linesmove towards β 1 = 1 and finally get overlapped.
Revised: As n increases, ( C2 ) holds for a wide range of β 1 , so the grey vertical linesmove towards β 1 = 1 and finally get overlapped.

Original: In addition, the size of divergence region increaseswith n
Revised: In addition, the size of divergence region increaseswith n

Original: Relation with γ 1 ( n ) in Theorem 3.1 According to Theorem 3.1, γ 1 ( n ) is at least in the order of1 − O ( n − 2 ) .
Revised: Relation with γ 1 ( n ) in Theorem 3.1 According to Theorem 3.1, γ 1 ( n ) is at least in the order of1 − O ( n − 2 ) .

Original: Combining with Figure 14.
Revised: Combining with Figure 16.

Original: It is not hard to see that γ 1 ( n ) is always larger than theupper boundary of the orange region, so there is no contradiction.
Revised: It is not hard to see that γ 1 ( n ) is always larger than theupper boundary of the orange region, so there is no contradiction.

Original: Proof of Corollary 4.1.
Revised: Proof of Corollary 4.1.

Original: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , condition (C1), (C2) and (C3) can be satisfiedby some sufficiently large n .
Revised: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , condition (C1), (C2) and (C3) can be satisfiedby some sufficiently large n .

Original: Therefore, Adam will diverge and the proof is concluded.
Revised: Therefore, Adam will diverge and the proof is concluded.

Original: F Some More Notations and Useful Lemmas for Convergence AnalysisF.1 More notations
Revised: F Some More Notations and Useful Lemmas for Convergence Analysis

Revised: F.1 More notations

Original: • We denote x k,i , m k,i , v k,i ∈
Revised: • We denote x k,i , m k,i , v k,i ∈

Original: R
Revised: R

Original: d as the value of x, m, v at the k -th outer loop and i -th inner loop.
Revised: d as the value of x, m, v at the k -th outer loop and i -th inner loop.

Original: Further, we denote x l,k,i , m l,k,i , v l,k,i ∈ R as the l -th component of x k,i , m k,i , v k,i .
Revised: Further, we denote x l,k,i , m l,k,i , v l,k,i ∈ R as the l -th component of x k,i , m k,i , v k,i .•

Original: • We denote η k as the stepsize at the k -th epoch (outer loop).
Revised: We denote η k as the stepsize at the k -th epoch (outer loop).

Original: We will focus mainly on diminishing step size, especially η k = η 1 √ nk , where n is the number of batches (inner loop).
Revised: We will focus mainly on diminishingstep size, especially η k = η 1 √ nk , where n is the number of batches (inner loop).

Original: • We denote ∂ l f ( x ) =
Revised: • We denote ∂ l f ( x ) =

Original: ∂∂x l f ( x ) , ∂ l f j ( x ) =
Revised: ∂∂x l f ( x ) , ∂ l f j ( x ) =

Original: ∂∂x l f j ( x ) .
Revised: ∂∂x l f j ( x ) .

Original: Further, we will use τ k,i to index[Table] the i -th randomly chosen batch in the k -th epoch.
Revised: Further, we will use τ k,i to index the i -thrandomly chosen batch in the k -th epoch.

Original: In this sense, we denote ∂ l f τ k,i ( x ) as[Table] ∂∂x l f τ k,i ( x ) .
Revised: In this sense, we denote ∂ l f τ k,i ( x ) as ∂∂x l f τ k,i ( x ) .

Original: [Table]
Revised: [Table]

Original: • Given an epoch k , we denote α as the index of the coordinate with the greatest gradient:
Revised: • Given an epoch k , we denote α as the index of the coordinate with the greatest gradient:

Original: [Equation]
Revised: [Equation]

Original: 0.8 beta 0.

Original: b e t a(a) n =0.

Original: 0.8 beta 0.

Original: b e t a(b) n =0.

Original: 0.8 beta 0.

Original: b e t a(

Original: [Figure]

Original: 0.8 beta 0. b e t a [Figure] • We define ρ 1 , ρ 2 , ρ 3 as constants satisfying the following condition for any l = 1 , · · · , d :
Revised: • We define ρ 1 , ρ 2 , ρ 3 as constants satisfying the following condition for any l = 1 , · · · , d :

Original: [Equation]
Revised: [Table]

Original: ρ 1 , ρ 2 , ρ 3 are problem-dependent constants.
Revised: ρ 1 , ρ 2 , ρ 3 are problem-dependent constants.

Original: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .
Revised: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .•

Original: ρ 3 is larger when ∂ l f i ( x k, 0 ) are more aligned. • We define the constant △ x : [Equation]
Revised: We define the constant △ x := η 0 √ x L √ d √ 1 − β 2 1 − β 1 1 − β 1 √ β .

Original: This constant will be used repeatedly, especially in Lemma F.3.
Revised: This constant will be used repeatedly, especiallyin Lemma F.2.•

Original: • We define Q k : [Equation]
Revised: We define Q k := △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 .

Original: This constant will be used repeatedly.
Revised: This constant will be used repeatedly.•

Original: • We define
Revised: We define

Original: [Equation]
Revised: [Equation]

Original: This constant will be used repeatedly.
Revised: This constant will be used repeatedly.

Original: Note that δ 1 → 0 when β 2 → 1 . • E ( · ) means taking expectation over the whole trajectory. • We will use E k [ · ] as a shorthand for E [ · (cid:12)(cid:12) x k, 0 , x k − 1 ,n − 1 , · · · , x 1 , 0 ] , i.e. the conditionalexpectation given all the history up to x k, 0 .
Revised: Note that δ 1 → 0 when β 2 → 1 .• E ( · ) means taking expectation over the whole trajectory.• We will use E k [ · ] as a shorthand for E [ · (cid:12)(cid:12) x k, 0 , x k − 1 ,n − 1 , · · · , x 1 , 0 ] , i.e. the conditional expectation given all the history up to x k, 0 .

Original: Similarly, E k − 1 stands for the conditionalexpectation given all the history up to x k − 1 , 0 .
Revised: Similarly, E k − 1 stands for the conditional expectation givenall the history up to x k − 1 , 0 .

Original: We further introduce some ‘index’ notation of the stochastic gradients, these notations will be usefulin the proof.
Revised: We further introduce some ‘index’ notation of the stochastic gradients, these notations will be usefulin the proof.

Original: For all the stochastic gradient, we have
Revised: For all the stochastic gradient, we have

Original: • True index: every stochastic gradient has its own true index: f 1 ( · ) , f 2 ( · ) , · · · , f n − 1 ( · ) .
Revised: • True index: every stochastic gradient has its own true index: f 1 ( · ) , f 2 ( · ) , · · · , f n − 1 ( · ) .

Original: Allthese { f i ( · ) } n − 1 i =0 are fixed once the optimization problem is formulated.
Revised: All these{ f i ( · ) } n − 1 i =0 are fixed once the optimization problem is formulated.

Original: Note that for a fixed x , we have the following relationship between f i and f :
Revised: Note that for a fixed x , we have the following relationship between f i and f :

Original: [Equation]
Revised: [Equation]

Original: • Random-shuffle index: At the k -th epoch, all the stochastic gradients are sampled in theorder of f τ k, 0 ( · ) , f τ k, 1 ( · ) , · · · , f τ k,n − 1 ( · ) .
Revised: • Random-shuffle index: At the k -th epoch, all the stochastic gradients are sampled in the order off τ k, 0 ( · ) , f τ k, 1 ( · ) , · · · , f τ k,n − 1 ( · ) .

Original: Since Algorithm 1 is sampling without replacement, there is an implicit bijective mappingbetween { f i ( · ) } n − 1 i ={+=
+} and { f τ k,i ( · ) } n − 1 i =0 .
Revised: Since Algorithm 1 is sampling without replacement, there is an implicit bijective mapping between{ f i ( · ) } n − 1 i {+=
+} and { f τ k,i ( · ) } n − 1 i =0 .

Original: Further, for a fixed x , we have the following usefulproperty:
Revised: Further, for a fixed x , we have the following useful property:n

Original: [Equation]
Revised: [Equation]

Original: • Index for m and v : As shown in Algorithm 1, we denote :
Revised: • Index for m and v : As shown in Algorithm 1, we denote :

Original: [Equation]
Revised: [Equation]

Original: m l

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: F.2 Some useful lemmas

Original: We begin with proving some useful lemmas.
Revised: We begin with proving some useful lemmas.

Original: Lemma F.1.
Revised: Lemma F.1.

Original: Let M and k be 2 integers with k > M , we have

Original: [Equation]

Original: Proof.

Original: For the lower bound,

Original: [Equation]

Original: For the upper bound,

Original: [Equation]

Original: Lemma F.2.

Original: For any β ̸ = 1 , we have
Revised: For any β ̸ = 1 , we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Proof only involves basic calculation, we omit the proof here.
Revised: Proof only involves basic calculation, we omit the proof here.

Original: Lemma F.3.
Revised: Lemma F.2.

Original: Under Assumption 2.1, if β 1 < √ β 2 , we have
Revised: Under Assumption 2.1, if β 1 < √ β 2 , we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We start with bounding | x l,k,i +1 − x l,k,i | .
Revised: We start with bounding | x l,k,i +1 − x l,k,i | .

Original: By the update formula of Adam, we have
Revised: By the update formula of Adam, we have| x l

Original: [Equation]
Revised: [Equation]

Original: Now, by the Lipschitz assumption of ∂ l f τ k,i ( x k,i ) (Assumption 2.1), we have(cid:12)(cid:12)
Revised: Now, by Assumption 2.1, we have

Original: [Equation]
Revised: [Equation]

Original: Proof is completed.
Revised: Proof is completed.

Original: Lemma F.4.
Revised: Lemma F.3.

Original: Under Assumption 2.1 and 2.2, for any integer i ∈ [0 , n − 1] , we have the followingtwo results,[Equation]
Revised: Under Assumption 2.1 and 2.2, for any integer i ∈ [0 , n − 1] , we have the following two results,

Original: [Equation]
Revised: [Equation]

Original: where | ∂ α f ( x k, 0 ) | = max l ∈ [1 ,d ] | ∂ l f ( x k, 0 ) | .
Revised: where | ∂ α f ( x k, 0 ) | = max l ∈ [1 ,d ] | ∂ l f ( x k, 0 ) | .

Original: Proof.
Revised: Proof.

Original: | ∂ l f τ k,j ( x k,j ) |

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table] [Equation] [Table] [Equation]
Revised: where ( ∗ ) is because of the definition (see Appendix F.1): ρ 1 is a constant satisfying ρ 1 ≥

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table] [Equation] [Table] [Equation]
Revised: ni =

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table] [Equation] [Table] [Equation]
Revised: f i

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table] [Equation] where ( ∗ ) is because of the definition (see Appendix F.1): ρ 1 is a constant satisfying ρ 1 ≥(cid:80) ni =1 | ∂ l f i ( x k, 0 ) | √ (cid:80) ni
Revised: x k

Original: =1 | ∂ l f i ( x k, 0 ) | 2 ; ( ∗∗ ) is due to ∥∇ f ( x k, 0 )
Revised: [Equation]

Original:  ∥ 22 ≤ d | ∂ α f ( x k, 0 ) | 2; ( ∗∗ ) is because √ x + y ≤ 793 √ x + √ y , for any x, y ≥ 0 .
Revised: 0 ) | 2 ; ( ∗∗ ) is due to ∥∇ f ( x k, 0 ) ∥ 22 ≤ d | ∂ α f ( x k, 0 ) | 2 .

Original: The proof of (21) follows the same procedure as above.
Revised: The proof of (23) follows the same procedure as above.

Original: G Proof of Theorem 3.1G.1 A roadmap of the proof
Revised: G Proof of Theorem 3.1G.1 A roadmap of the proof

Original: Our proof is based on the Descent Lemma:
Revised: Our proof is based on the Descent Lemma:

Original: [Equation]
Revised: [Equation]

Original: The expectation E ( · ) is taken on the whole trajectory.
Revised: The expectation E ( · ) is taken on the whole trajectory.

Original: Summing both sides from the initializationk =
Revised: Summing both sides from the initializationk =

Original: t 0 to k = T , we have the following re-arranged inequality: (usually we set t 0 = 1 .)
Revised: t 0 to k = T , we have the following re-arranged inequality: (usually we set t 0 = 1 .)

Original: [Equation]
Revised: [Equation]

Original: To prove the convergence, we need an upper bound for L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 , as well as a lower boundfor E ⟨∇ f
Revised: To prove the convergence, we need an upper bound for L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 , as well as a lower boundfor E ⟨∇ f

Original: ( x k, 0 ) , x k, 0 − x k +1 , 0 ⟩ (and such a lower bound should be in the order of 1 √ k E ∥∇ f
Revised: ( x k, 0 ) , x k, 0 − x k +1 , 0 ⟩ (and such a lower bound should be in the order of 1 √ k E ∥∇ f

Original: ( x k, 0 ) ∥ ).
Revised: ( x k, 0 ) ∥ ).

Original: The upper bound for L 2 E ∥ x k +1 , 0 − x k, 

Revised: The upper bound for L 2 E ∥ x k +1 , 0 − x k,

Original: ∥ 22 is relatively easy to get: according to the update rule of
Revised: ∥ 22 is relatively easy to get: according to the update rule of

Original: Adam, this term is in the order of O ( η 2 k m l,k,i /v l,k,i )
Revised: Adam, this term is in the order of O ( η 2 k m l,k,i /v l,k,i )

Original: =
Revised: =

Original: O ( η 2 k ∥∇ f
Revised: O ( η 2 k ∥∇ f

Original: ( x k, 0 ) ∥ / ∥∇ f ( x k, 0 ) ∥ )
Revised: ( x k, 0 ) ∥ / ∥∇ f ( x k, 0 ) ∥ )

Original: =
Revised: =

Original: O ( 1 k ) .
Revised: O ( 1 k ) .

Original: Further, recall (cid:80) Tk = t 0 1 k ≤ log T +1 t 0 , so L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 contributes to the log term in Theorem3.1.
Revised: Further, recall (cid:80) Tk = t 0 1 k ≤ log T +1 t 0 , so L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 contributes to the log term in Theorem3.1.

Original: The proof is shown in Lemma F.3.
Revised: The proof is shown in Lemma F.2.

Original: However, the lower bound for E ⟨∇ f
Revised: However, the lower bound for E ⟨∇ f

Original: ( x k, 0 ) , x k+1 , 0 − x k, 0 ⟩ = E (cid:104) (cid:80) dl
Revised: ( x k, 0 ) , x k

Original: [-=1 (cid:80) n − 1 i =
-]
Revised: [-=1 (cid:80) n − 1 i =
-][Equation]

Original: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i (cid:105)requires sophisticated derivation, we explain as follows.
Revised: requires sophisticated derivation, we explain as follows.

Original: Before taking the expectation, we firstwork on every possible realization of (cid:80) dl
Revised: Before taking the expectation, we firstwork on every possible realization of (cid:80) dl

Original: =1 (cid:80) n − 1 i =

Revised: =1 (cid:80) n − 1 i =

Original: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i .
Revised: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i .

Original: We perform the followingdecomposition for every l ∈ [1 , d ] .
Revised: We perform the followingdecomposition for every l ∈ [1 , d ] .

Original: n − 1 (cid:88) i =

Original: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i = n − 1 (cid:88)

Original: i = ∂ l f ( x k, 0 ) √ v l,k,i ( ∂ l f i ( x k, 0 ) + m l,k,i − ∂ l f i ( x k, 0 )) = (cid:34) n − 1 (cid:88) i =
Revised: [Equation]

Original: ∂ l f ( x k, 0 ) √ v l,k,i ∂ l f i ( x k, 0 ) (cid:35) (cid:124) (cid:123)(cid:122) (cid:125)

Original: ( a ) + (cid:34) ∂

Original: l f ( x k, 0 ) n − 1 (cid:88) i = 1 √ v l,k,i ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b ) .

Original: First of all, we introduce the following Lemma G.1 to further decompose ( a ) and ( b ) .
Revised: First of all, we introduce the following Lemma G.1 to further decompose ( a ) and ( b ) .

Original: The intuitionof this decomposition is as follows: by increasing β 2 , we can control the movement of the movingaverage factor v l,k,i , so the trajectory of Adam is easier to control.
Revised: The intuitionof this decomposition is as follows: by increasing β 2 , we can control the movement of the movingaverage factor v l,k,i , so the trajectory of Adam is easier to control.

Original: Lemma G.1.
Revised: Lemma G.1.

Original: Under Assumption 2.1, for those l satisfying max i | ∂ l f i ( x k, 0 )
Revised: Under Assumption 2.1, for those l satisfying max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k :=
Revised: | ≥ Q k :=

Original: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , we have the following lower bound for ( a ) in (23) :
Revised: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , we have the following lower bound for ( a ) in (25) :

Original: [Equation]
Revised: [Equation]

Original: (cid:12)(cid:12)(cid:12)(cid:12)

Original: [Equation]
Revised: [Equation]

Revised: where δ 1 = (1 − β 2 )4 nρ 2 β n

Original: [Equation]
Revised: [Equation]

Original: where δ 1 = (1 − β 2  ) 
Revised: ρ 2 , ρ 3 are constants satisfying ρ 2 ≥ | max i ∂ l f i ( x k, 0 ) | 2 1 n (cid:80) ni =1 | ∂ l f i ( x k, 0 ) | 2 ,ρ 3 ≥ |

Original: [Equation]
Revised: [Equation]

Original: ρ 2 , ρ 3 are constants satisfying ρ 2 ≥
Revised: l f i

Original: [Equation]
Revised: [Equation]

Original: δ
Revised: k

Original: [Equation]
Revised: [Equation]

Original: 0 when β
Revised: n [Equation] l f i

Original: [Equation]
Revised: [Equation]

Revised: k

Revised: [Equation]

Revised: [Equation]

Revised: | 2 .

Revised: Note that δ 1 → 0 when β 2 → 1 .

Original: Under the same condition, we also have a lower bound for ( b ) in (23) :
Revised: Under the same condition, we also have a lower bound for ( b ) in (25) :

Original: [Equation]
Revised: [Equation]

Original: The proof can be seen in Appendix G.2.
Revised: The proof can be seen in Appendix G.2.

Original: Lemma G.1 decomposes both ( a ) and ( b ) : in the denominator,we approximate v l,k,i (which is changing with i ) by v l,k, 0 (which is fixed).
Revised: Lemma G.1 decomposes both ( a ) and ( b ) : in the denominator,we approximate v l,k,i (which is changing with i ) by v l,k, 0 (which is fixed).

Original: Accordingly, theapproximation error can be controlled by increasing β 2 (since δ 1 → 0 when β 2 → 1 ).
Revised: Accordingly, theapproximation error can be controlled by increasing β 2 (since δ 1 → 0 when β 2 → 1 ).

Original: However,
Revised: However,

Original: Lemma G.1 can only be applied to those l with max i | ∂ l f i ( x k, 0 )
Revised: Lemma G.1 can only be applied to those l with max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k , so we need to discuss thefollowing two cases to complete the whole picture.
Revised: | ≥ Q k , so we need to discuss thefollowing two cases to complete the whole picture.

Original: Case 1: unbounded gradient.
Revised: Case 1: unbounded gradient.

Original: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 )
Revised: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k , we havethe following decomposition:
Revised: | ≥ Q k , we havethe following decomposition:

Original: [Table]
Revised: [Equation]

Revised: ( a ) + (cid:34) ∂

Revised: l f ( x k, 0 ) n − 1 (cid:88) i = 1 √ v l,k,i ( m l,k,i − ∂ l f i ( x k, 0 ))

Original:  (cid:123)(cid:122) (cid:125) ( a 1 ) [-− (cid:34) δ(cid:115) 
-]
Revised: (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b ) [-− (cid:34) δ(cid:115) 
-].

Original: ρβ nn − 1 (cid:88) i = | ∂ l f i ( x k, 0 ) |(cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( a 2 ) +(cid:34) ∂ l f ( x k, 0 ) √ v l,k,n − 1 (cid:88) i =
Revised: LemmaG.1 ≥ ∂ l f ( x k, 0 ) 2 √ v l,k,− δ 1 (cid:12)(cid:12)(cid:12)(cid:12) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:12)(cid:12)(cid:12)(cid:12) n − 1 (cid:88) i = | ∂ l f i ( x k, 0 ) | + ∂ l f ( x k, 0 ) √ v l,k,n − 1 (cid:88) i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 )) 
Revised: ( m l,k,i − ∂ l f i ( x k, 0 )) − δ 1 (cid:12)(cid:12)(cid:12)(cid:12) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:12)(cid:12)(cid:12)(cid:12) n − 1 (cid:88) i = | m l,k,i − ∂ l f i ( x k, 0 ) | LemmaG.10 ≥ (cid:20) ∂

Original: (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b 1 ) − (cid:34) δ(cid:115) [-
-] ρβ nn − 1 (cid:88) i =| m l,k
Revised: l f ( x k, 0 ) 2 √ v l,k,(cid:21) (cid:124) (cid:123)(cid:122) (cid:125) ( a 1 ) − (cid:34) δ(cid:115) [-
-]2 ρβ nn − 1 (cid:88) i =

Original: i − ∂ l f i ( x k, 0 ) | (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b 2 ) , (26)
Revised: | ∂ l f i ( x k, 0 ) | (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( a 2 ) + [Equation]

Original: Case 2: bounded gradient.
Revised: Case 2: bounded gradient.

Original: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 ) | ≤ Q k , the analysisdegenerates to the “bounded gradient” scenario, we have the following lower bound:
Revised: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 ) | ≤ Q k , the analysisdegenerates to the “bounded gradient” scenario, we have the following lower bound:

Original: [Equation]
Revised: [Equation]

Original: where F:= △ 1 n 2 √ n 32 √ [Equation] [Equation] − β
Revised: where F 1 :=

Original: [Equation]

Original: ) n

Original: [Equation]
Revised: [Equation]

Original: n
Revised: n

Original: [Equation]
Revised: [Equation]

Original: − β 1 1 − β 1 √ β 2 n , ( ∗ ) is due to the fact that | ∂ l f ( x k, 0 ) | ≤n max i | ∂ l f i ( x k, 0 ) | ≤ nQ k .
Revised: n , ( ∗ ) is due to the fact that | ∂ l f ( x k, 0 ) | ≤n max i | ∂ l f i ( x k, 0 ) | ≤ nQ k .

Original: Combining Case 1 & 2 together, we have the following result (note that ( a 1 ) , ( a 2 ) , ( b 1 ) , ( b 2 ) varywith l .
Revised: Combining Case 1 & 2 together, we have the following result (note that ( a 1 ) , ( a 2 ) , ( b 1 ) , ( b 2 ) varywith l .

Original: A more precise notation will be ( a 1 ) l , etc.
Revised: A more precise notation will be ( a 1 ) l , etc.

Original: We drop the subscript for brevity).
Revised: We drop the subscript for brevity).

Original: d

Original: [Equation]
Revised: [Equation]

Original: l large
Revised: l large

Original: [Equation]
Revised: [Equation]

Original: (cid:125) unbounded gradient − (cid:88) l
Revised: (cid:125) unboundedgradient [Equation]

Original:  small
Revised: l small

Original: F 1 1 √ k
Revised: F [Equation] k

Original: (cid:124) (cid:123)(cid:122) (cid:125) bounded gradient ≥(cid:88) l large
Revised: [Equation] (cid:125) boundedgradient ≥

Original: [Equation]
Revised: [Equation]

Revised: l large

Original: [Equation]
Revised: [Equation]

Original:  (cid:124) (cid:123)(cid:122) (cid:125) unbounded gradient − d(cid:88)
Revised: ) } (cid:124) (cid:123)(cid:122) (cid:125) unboundedgradient − d

Original: l = (cid:26) 
Revised: [Equation] (cid:26) F

Original: F 1 1 √ k (cid:27) (cid:124) (cid:123)(cid:122) (cid:125) bounded gradient ≥   (cid:88)
Revised: [Equation] √ k (cid:27) (cid:124) (cid:123)(cid:122) (cid:125) boundedgradient ≥   (cid:88)

Original: l large ( a 1 )   +   (cid:88)
Revised: l large ( a 1 )   +   (cid:88)

Original: l large ( b 1 )   −   (cid:88) l large { ( a 2 ) + ( b 2 ) }   − dF 1 1 √ k , (27){+
 [Equation]+} where “ l large” stands for the gradient component in Case 1 and “ l small” indicates those in Case2 .
Revised: l large ( b 1 )   −   (cid:88) l large { ( a 2 ) + ( b 2 ) }   − dF {+
 [Equation]+} where “ l large” stands for the gradient component in Case 1 and “ l small” indicates those in Case2 .

Original: We assume “ l large” is not an empty set, otherwise the analysis degenerates to the proof withbounded gradient assumption.
Revised: We assume “ l large” is not an empty set, otherwise the analysis degenerates to the proof withbounded gradient assumption.

Original: Now we take expectation on (27).
Revised: Now we take expectation on (29).

Original: The expectation is taken on all the possible trajectories up to thei -th iteration in k -th epoch.
Revised: The expectation is taken on all the possible trajectories up to thei -th iteration in k -th epoch.

Original: [Equation]
Revised: [Equation]

Original: In the following context, we will discuss how to bound all terms in (28), respectively.
Revised: In the following context, we will discuss how to bound all terms in (30), respectively.

Original: First andforemost, we derive a lower bound for E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .
Revised: First andforemost, we derive a lower bound for E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .

Original: Since E (cid:104)(cid:80) l large ( b 1 ) (cid:105) contains all thehistorical gradient information, it involves great effort to handle it.
Revised: Since E (cid:104)(cid:80) l large ( b 1 ) (cid:105) contains all thehistorical gradient information, it involves great effort to handle it.

Original: We will show that the lower boundof E (cid:104)(cid:80) l large ( b 1 ) (cid:105) will vanish when β 2 is large and k goes to infinity.
Revised: We will show that the lower boundof E (cid:104)(cid:80) l large ( b 1 ) (cid:105) will vanish when β 2 is large and k goes to infinity.

Original: Recall E (cid:104)(cid:80) [Equation]
Revised: Recall E (cid:104)(cid:80) l large ( b 1 ) (cid:105) =

Original: = [Equation] ( m [Equation]{+x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
+}
Revised: E (cid:104)(cid:80) l large ∂ l f ( {+x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
+}

Original:  − ∂ [Equation] f i ( x k [Equation]
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) .
Revised: (cid:105) .

Original: When β 1 is large,m l,k,i contains heavy historical signals.
Revised: When β 1 is large,m l,k,i contains heavy historical signals.

Original: It seems unclear how large E (cid:104)(cid:80) l large ( b 1 ) (cid:105) would bewhen β 1 goes to 1.
Revised: It seems unclear how large E (cid:104)(cid:80) l large ( b 1 ) (cid:105) would bewhen β 1 goes to 1.

Original: Existing literatures (Zaheer et al., 2018; De et al., 2018; Shi et al., 2020)
Revised: Existing literatures (Zaheer et al., 2018; De et al., 2018; Shi et al., 2020)

Original: takea naive approach: they set β 1 ≈ 0 so that m k,i ≈ ∇ f τ k,i ( x k,i ) .
Revised: takea naive approach: they set β 1 ≈ 0 so that m k,i ≈ ∇ f τ k,i ( x k,i ) .

Original: Then we get δ ( β 1 ) ≈ 0 .
Revised: Then we get δ ( β 1 ) ≈ 0 .

Original: However,this naive method cannot be applied here since we are interested in practical cases where β 1 is largein [0 , 1) .
Revised: However,this naive method cannot be applied here since we are interested in practical cases where β 1 is largein [0 , 1) .

Original: We emphasize the following technical difficulties in bounding E (cid:104) (cid:80) l large ( b 1 ) (cid:105) for anyβ 1 ∈ [0 , 1) :
Revised: We emphasize the following technical difficulties in bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) for anyβ 1 ∈ [0 , 1) :

Original:  Issue (i)
Revised: • Issue (i)

Original: In order to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) , we need to first know how to bound its simpler version: E (cid:104)(cid:80) n − 1 i =

Revised: In order to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) , we need to first know how to bound its simplerversion: E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) .
Revised: (cid:105) .

Original: This term measures the difference between the currentgradient and weighted historical gradients.
Revised: This term measures the difference between the currentgradient and weighted historical gradients.

Original: It seems unclear that how large this term couldbe when β 1 goes to 1. Issue (ii) Even if we can bound E (cid:104)(cid:80) n − 1 i =

Revised: It seems unclear that how large this term could bewhen β 1 goes to 1.• Issue (ii) Even if we can bound E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , it is still unclear how to bound
Revised: (cid:105) , it is still unclear how to bound

Original: E (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Revised: E (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , which is further multiplied by a random variable∂ l f ( x k, 0 ) [Table] l,k, 0 .
Revised: (cid:105) , which is further multiplied by a random variable∂ l f ( x k, 0 ) √ v l,k, 0 .

Original:  Issue (iii)
Revised: • Issue (iii)

Original: Even if we can bound E (cid:104) ∂
Revised: Even if we can bound E (cid:104) ∂

Original: l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Revised: l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , it is still different from E (cid:104)(cid:80) l large ( b 1 ) (cid:105) which contains additional operation “ (cid:80) l large " inside the expectation.
Revised: (cid:105) , it is still different from E (cid:104)(cid:80) l large ( b 1 ) (cid:105) which contains additional operation “ (cid:80) l large " inside the expectation.

Original: Note that the set “ l large" is a random variable which changes along different trajectories, so there is still non-negligible gap between bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105)and
Revised: Note that the set “ l large" is a random variable which changes along different trajectories, so there is still non-negligible gap between bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105)

Original: [Equation]{+andE (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
 ( m l,k,i − ∂ l f i ( x k, 0 )) (cid:105) .+}
Revised: {+andE (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
 ( m l,k,i − ∂ l f i ( x k, 0 )) (cid:105) .+}

Original: To our best knowledge, there is no general approach to tackle the above issues.
Revised: To our best knowledge, there is no general approach to tackle the above issues.

Original: In the followingcontent, we will overcome difficulties (i), (ii) and (iii) in Step 1, 2 and 3 respectively.
Revised: In the followingcontent, we will overcome difficulties (i), (ii) and (iii) in Step 1, 2 and 3 respectively.

Original: In Step 1, wewill discuss how to bound E (cid:104)(cid:80) n − 1 i =

Revised: In Step 1, wewill discuss how to bound E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , which is a simplified version of E [( b 1 )] .
Revised: (cid:105) , which is a simplified version of E [( b 1 )] .

Original: Bounding this term will shed light on bounding the whole term E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .
Revised: Bounding this term will shed light on bounding the whole term E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .

Original: Then in Step 2 and3, we will prove several technical lemmas to handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 and “ (cid:80) l large ", by whichwe can tackle (ii) and (iii) .
Revised: Then in Step 2 and3, we will prove several technical lemmas to handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 and “ (cid:80) l large ", by whichwe can tackle (ii) and (iii) .

Original: Combining all together we can bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) . BoundingE (cid:104)(cid:80)
Revised: Combining all together we can bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) . Bounding

Original: [Equation]
Revised: [Equation]

Original: Step 1.
Revised: Step 1.

Original: We now introduce the key idea of bounding
Revised: We now introduce the key idea of bounding

Original: [Equation] As discussed in Appendix A, we analyze{+E (cid:104)(cid:80) n − 1 i =
+}
Revised: {+E (cid:104)(cid:80) n − 1 i =
+}

Original: [Figure] [Equation] f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) for every epoch and find out some pattern across the successive epochs. In Appendix A, we distill our idea into a toy example with two balls in a box. Here,to better convey our message, we further describe a similar case with three balls. The “color-ball " model with three balls. Consider a box containing 3 balls labeled a 0 , a 1 , a 2 ∈ R(could be negative).
Revised: (cid:105) .

Original: In each round (epoch), we randomly sample all the balls from the box withoutreplacement, then we put them all back.

Original: We denote the 1st sampled label in the k -th epoch as a τ k, 0 ,and similarly for a τ k, 1 , a τ k, 2 . We define the following random variables: ( β ∈ [0 , 1] ) [Equation] [Equation]
Revised: In Appendix A, we distill our idea into a toy example called "thecolor-ball model of the 1st kind" and thus we prove Lemma A.1.

Original: Further

Original: [Equation] we introduce the following deterministic constants. [Equation] For the simplicity of analysis
Revised: This lemma is crucial for bounding

Original: [Equation]

Original: we denote

Original: [Equation]

Original: [Equation]

Original: [-2 i =
-]
Revised: [-2 i =
-]

Original: [-f i (cid:17) k − 1 + · · · + (cid:16)(cid:80) 2 i =
-]
Revised: [-f i (cid:17) k − 1 + · · · + (cid:16)(cid:80) 2 i =
-]

Original: [-f i (cid:17) 1 ,where (cid:16)(cid:80) 2 i =
-]
Revised: [-f i (cid:17) 1 ,where (cid:16)(cid:80) 2 i =
-]

Original: [-m i (cid:17) k and (cid:16)(cid:80) 
-]
Revised: [-m i (cid:17) k and (cid:16)(cid:80) 
-]

Original: [-i =
-]
Revised: [-i =
-]

Original: f i (cid:17) k contain the summand in the k -th epoch, so on so for.

Original: [-Now, wetry to calculate E (cid:104)(cid:80) 2 i =
-]
Revised: [-Now, wetry to calculate E (cid:104)(cid:80) 2 i =
-]

Original: [-m i − (cid:80) 2 i =
-]
Revised: [-m i − (cid:80) 2 i =
-]

Original: f i (cid:105) , where the expectation is taken on all the possible drawsuntil the end of k -th round (epoch).

Original: Lemma G.2.

Original: In the colorb-ball example above, we have

Original: E

Original: [Equation]

Original: For general n = 1 , 2 , 3 , · · · , we have

Original: E

Original: [Equation]

Original: Proof.

Original: [-Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 2 i =
-]
Revised: [-Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 2 i =
-]

Original: [-m i − (cid:80) 2 i =
-]
Revised: [-m i − (cid:80) 2 i =
-]

Original: f i (cid:105) .

Original: Since all the history before k -th epoch is fixed, we delegate the history part to later analysis.

Original: Now, we focus on 890 (cid:16)(cid:80)

Original: [Equation]

Original:  − [-(cid:16)(cid:80) 
-] i [Equation] f{+=
+}
Revised: E (cid:104)(cid:80) n − [-(cid:16)(cid:80) 
-]1 i {+=
+}

Original: [Equation] (cid:17) k . As shown in Figure 15, there are (3)! =
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: 6 possible order in the k -thepoch.

Original: To calculate E k ( · ) , we need to sum up all 6 possible shuffle order.

Original: Figure 15 tells us: when all thehistory sampling orders are fixed, we have: (we use

Original: = all 6 possible shuffle )(cid:88) s (cid:40)(cid:32) 2 (cid:88)

Original: [Equation]

Original: Since all the sampling order occur with the same probability 1 / 6 , we have:

Original: [Equation]

Original: [-Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 2 i =
-]
Revised: [-Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 2 i =
-]

Original: [-m i − (cid:80) 2 i =
-]
Revised: [-m i − (cid:80) 2 i =
-]

Original: f i (cid:105) .
Revised: (cid:105) .

Original: Since all the historybefore ( k − 1) -th epoch is fixed, we do not discuss them for now.

Original: In particular, we need to calculate

Original: E k −(cid:40)

Original: E k (cid:34)(cid:32) 2 (cid:88) i

Original: = m i (cid:33) k − (cid:32) 2 (cid:88) i =

Original: f i (cid:33) k

Original: + (cid:32) 2 (cid:88) i

Original: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Original: f i (cid:33) k −(cid:35)(cid:41) =

Original: E k −(cid:40)

Original: [Equation]

Original: (cid:88) i

Original: = a i ( [Equation] 1 − [Equation]
Revised: We refer the readers to Appendix A for more explanation.

Original: 2 ) + (cid:32) 2 (cid:88) i

Original: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Original: f i (cid:33) k −(cid:41) .

Original: We highlight the blue term since it is “the residue from the k -th epoch”, this term will be canceledout later.

Original: [-Similarly as Figure 15, we sum up all possible draws (cid:80) s (cid:16)(cid:80) 2 i =
-]
Revised: [-Similarly as Figure 15, we sum up all possible draws (cid:80) s (cid:16)(cid:80) 2 i =
-]

Original: m

Original: i (cid:17) k − 1 in Figure 16.

Original: As shown in Figure 16, we have(cid:88) s   2 (cid:88)

Original: i = a i ( − 2 6 β 1 − 4 6 β 2 ) + (cid:32) 2 (cid:88) i

Original: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Original: f i (cid:33) k

Original: −  = 2 (cid:88)

Original: i = a i ( − 2 β 1 − 4 β 2 ) + 2 (cid:88)

Original: i = a i (2 β 1 + 4 β 2 ) (cid:124) (cid:123)(cid:122) (cid:125)

Original: canceloutthe k -thresidue + 2 (cid:88)

Original: i = a i (6 β 3 + 4 β 4 + 2 β 5 ) − 2 (cid:88)

Original: i = a i (6 β 3 + 6 β 4 + 6 β 5 ) (cid:124) (cid:123)(cid:122) (cid:125) createthe ( k − 1) -thresidue

Original: = 2 (cid:88)

Original: i = a i ( − 2 β 4 − 4 β 5 ) .

Original: Equivalently,

Original:  E k −  E k   (cid:32) 2 (cid:88) i = m i (cid:33) k − (cid:32) 2 (cid:88) i = f i (cid:33) k{+=
+}
Revised: Lemma A.1 can provide insights in bounding E (cid:104)(cid:80) n − 1 i {+=
+}

Original: + (cid:32) 2 (cid:88) i

Original: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Original: f i (cid:33) k −   

Original: Figure i = a i ( − 2 6 β 4 − 4 6 β 5 )

Original: From (29) we find out that: the residue from the k -th epoch will be cancelled out by part of the redterms in the ( k − 1) -th epoch; meanwhile, new residue term will be created.

Original: Moreover, we observethe following pattern in the residue : 1) the order of β involves the largest & second largest order thatpossibly occur within the epoch.

Original: 2) the coefficients are fixed to be − 2 6 and − 4 6 .

Original: This is actually anuniversal pattern for every epoch from 1 to k .

Original: We illustrate this pattern in the following Figure 17.

Original: [-Combining the above analysis and Figure 17, we observe the following pattern: when iterativelyplugging in E k − j , the residue term remaining in the ( k − j + 1) -th epoch will be erased (i.e. 908 (cid:80)i =
-]
Revised: [-Combining the above analysis and Figure 17, we observe the following pattern: when iterativelyplugging in E k − j , the residue term remaining in the ( k − j + 1) -th epoch will be erased (i.e. 908 (cid:80)i =
-]

Original: [-a i β ( j − 1) ( − 2 β 4 − 4 β 5 ) ), whereas a new residue term will be produced (i.e. (cid:80) 2 i =
-]
Revised: [-a i β ( j − 1) ( − 2 β 4 − 4 β 5 ) ), whereas a new residue term will be produced (i.e. (cid:80) 2 i =
-]

Original: a i β j ( − 2 β 4 −

Original: 4 β 5 ) ).

Original: Since E [ · ] = E [ E 1 ( E k − 1 ( · · · E k ( · )))] , we derive the following result by repeatedly applyingthe steps in (29):

Original: [Equation]

Original: The above analysis also holds for general n = 1 , 2 , 3 , · · · , we just need to change the constants in thefollowing form:

Original: [Equation]

Original: Since the analysis for general “ n balls” is exactly the same as that of n = 3 , we omit the proof forbrevity.

Original: [-The color-ball toy example above can provide insights in bounding E (cid:104) (cid:80) n − 1 i =
-]
Revised: [-The color-ball toy example above can provide insights in bounding E (cid:104) (cid:80) n − 1 i =
-]

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) .
Revised: (cid:105) .

Original: However, there is still some gap between E (cid:104)(cid:80) n − 1 i =

Revised: However, there isstill certain gap between E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) and the quantities in Lemma G.2.
Revised: (cid:105) and the quantities in Lemma A.1.

Original: We elaborate as follows:
Revised: Weelaborate as follows:

Original: • To mimic the color-ball example, we need to expand (cid:80) n − 1 i =

Revised: • To mimic the color-ball example, we need to expand (cid:80) n − 1 i =

Original: ∂ l f i ( x k, 0 ) into an infinite sumsequence: (cid:80) n − 1 i =

Revised: ∂ l f i ( x k, 0 ) into an infinite sumsequence: (cid:80) n − 1 i =

Original: ∂
Revised: ∂

Original: l f i ( x k, 0 )(1 − β 1 )(1 + β 1 + · · · β ∞ 1 ) .
Revised: l f i ( x k, 0 )(1 − β 1 )(1 + β 1 + · · · β ∞ 1 ) .

Original: However, (cid:80) n − 1 i =

Revised: However, (cid:80) n − 1 i =

Original: m l,k,i is a finitesum sequence up to the order of β kn .
Revised: m l,k,i is a finite sumsequence up to the order of β kn .

Original: In contrast, both sequences in the color-ball example 
Revised: In contrast, both sequences in the color-ball example are “finite

Original: are “finite sum”.
Revised: sum”.

Original: As such, there is an error term caused by “finite sum v.s. infinite sum”.
Revised: As such, there is an error term caused by “finite sum v.s. infinite sum”.

Original: • When taking the expectation, the variable x in each possible trajectory is different.
Revised: • When taking the expectation, the variable x in each possible trajectory is different.

Original: In contrast, in the color-ball example, { a i } 2 i =0 are fixed in all shuffling order (so it is much 
Revised: In contrast, inthe color-ball example, { a i } 2 i =0 are fixed in all shuffling order ( so it is much easier

Original: easier to calculate the expectation by summing them up).
Revised: to calculate the expectation by summing them up

Original: To mimic the color-ball example, we repeatedly take conditional expectation at the beginningof each k -th epoch.
Revised: To mimic the color-ball example, we repeatedly take conditional expectation at the beginning ofeach k -th epoch.

Original: In this way, x k, 0 will be fix.
Revised: In this way, x k, 0 will be fix.

Original: Despite x k,i is still changing across thetrajectory, we can transform x k,i into x k, 0 by using Lipschitz property.
Revised: Despite x k,i is still changing across the trajectory,we can transform x k, i into x k, 0 by using Lipschitz property.

Original: • In each trajectory of (cid:80) n − 1 i =

Revised: • In each trajectory of (cid:80) n − 1 i =

Original: m l,k,i , the variable x in the summand of m l,k,i varies with k and 
Revised: m l,k,i , the variable x in the summand of m l,k,i varies with k and i ;

Original: i ; while the variable in ∂ l f i ( x k, 0 ) is fixed to be x k, 0 .
Revised: while the variable in ∂ l f i ( x k, 0 ) is fixed to be x k, 0 .

Original: In contrast, in the color-ball example,{ a i } 2 i =0 are the same across the epoch.
Revised: In contrast, in the color-ball example, { a i } 2 i =0are the same across the epoch.

Original: To mimic the final step in the color-ball example (Figure 17), at each step of conditionalexpectation, we need to simultaneously move the variables in m l,k,i and ∂ l f i ( x k, 0 ) (using Lipschitz property) so that they can match and cancel out with each other.
Revised: To mimic the final step in the color-ball example (Figure 5), at each step of conditional expectation,we need to simultaneously move the variables in m l,k,i and ∂ l f i ( x k, 0 ) (using Lipschitz property)so that they can match and cancel out with each other.

Original: This operationwill introduce new error terms and it is our duty to put them under control.
Revised: This operation will introduce new errorterms and it is our duty to put them under control.

Original: We omit the proof for bounding E (cid:104)(cid:80) n − 1 i =

Revised: We omit the proof for bounding E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) since this is not our actual goal.
Revised: (cid:105) since this is not our actual goal.

Original: Instead, we will directly use the above ideas to bound E [( b 1 )] .
Revised: Instead, we will directly use the above ideas to bound E [( b 1 )] .

Original: To do so, we need to further tacklethe issue (ii) and (iii) mentioned before.
Revised: To do so, we need to further tacklethe issue (ii) and (iii) mentioned before.

Original: We explain as follows.
Revised: We explain as follows.

Original: Bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) :
Revised: Bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) :

Original: Step 2.
Revised: Step 2.

Original: We now resolve (ii), i.e., handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 .
Revised: We now resolve issue (ii) , i.e., handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 .

Original: Thekey idea is as follows: when we calculate E [( b 1 )] , we sequentially take conditional expectation 942 E k ( · ) , E k − 1 ( · ) , etc.. When taking E k ( · ) , we will fix all the historical information up to k -th epoch,so ∂ l f ( x k, 0 ) √ v l,k, 0 can be regarded as a constant.
Revised: The key idea is as follows: when we calculate E [( b 1 )] , we sequentially take conditional expectation 953 E k ( · ) , E k − 1 ( · ) , etc.. When taking E k ( · ) , we will fix all the historical information up to k -th epoch, so ∂ l f ( x k, 0 ) √ v l,k, 0 can be regarded as a constant.

Original: In this sense, E k (cid:104)
Revised: In this sense, E k (cid:104)

Original: [Table]

Original: i − ∂

Original: [Equation]
Revised: [Equation]

Original: f i [Equation] x k, 0 )) (cid:105)can be calculated following the same idea as color-ball toy example.
Revised: can be calculated following the same idea as color-ball toy example.

Original: However, when taking E k − 1 ( · ) , ∂ l f ( x k, 0 ) √ v l,k, 0 will become a random variable which changes with differenttrajectories.
Revised: However, when taking E k − 1 ( · ) , ∂ l f ( x k, 0 ) √ v l,k, 0 will become a random variable which changes with differenttrajectories.

Original: In this case, the color-ball method cannot be applied.
Revised: In this case, the color-ball method cannot be applied.

Original: To fix this issue, we introducethe following lemma to change ∂ l f ( x k, 0 ) √ v l,k, 0 into ∂ l f ( x k − 1 , 0 ) √ v k − 1 , 0 , which can again be regarded as a fixedconstant when taking E k − 1 ( · ) .
Revised: To fix this issue, we introducethe following lemma to change ∂ l f ( x k, 0 ) √ v l,k, 0 into ∂ l f ( x k − 1 , 0 ) √ v k − 1 , 0 , which can again be regarded as a fixedconstant when taking E k − 1 ( · ) .

Original: Lemma G.3.
Revised: Lemma G.2.

Original: Suppose Assumption 2.1 holds and β 1 < √ β 2 .
Revised: Suppose Assumption 2.1 holds and β 1 < √ β 2 .

Original: For any integer j ∈ [0 , k ] , if
Revised: For any integer j ∈ [0 , k ] , ifmax

Original: [Equation]
Revised: [Equation]

Original:  Q 
Revised: [Equation] | ≥ Q k

Original: [Equation]
Revised: [Equation]

Original:  (where
Revised: j (where

Original: Q
Revised: Q

Original: [Equation]
Revised: [Equation]

Original: β
Revised: :=

Original: [Equation]
Revised: [Equation]

Original: 2 ), then we have the following result:
Revised: 2 ), then we have the following result:

Revised: [Equation]

Revised: (cid:16)

Revised: [Equation]

Revised: [Equation]

Revised: n(cid:17)

Original: [Equation]
Revised: [Equation]

Revised: 1 , where δ 1 is defined in (16) .

Original: The proof of Lemma G.3 can be seen in Appendix G.3.
Revised: The proof of Lemma G.2 can be seen in Appendix G.3.

Revised: To proceed, we combine Lemma G.2 and the“color-ball method of the 2nd kind", and thus we can prove Lemma A.2.

Original: Bounding E(cid:104)(cid:80) l large ( b 1 ) (cid:105) :
Revised: This lemma is crucial forour current goal: bounding E ( b 1 ) . We refer the readers to Appendix A for more information. We emphasize that here are still the following gap between Lemma A.2 and our goal E ( b 1 ) . • We have the similar gap as discussed at the end of the Step 1 .

Revised: • The condition in Lemma G.2 has requirement on the gradient norm, while this requirement is

Revised: temporarily ignored in the color-ball method of the 2nd kind.

Revised: • The result in Lemma G.2 has additional error terms other than O (1 / √ k ) .

Revised: This is slightly different

Revised: from the setting in Lemma A.2.

Revised: It requires some technical lemmas to handle these gaps.

Revised: We fill in these gaps in Lemma G.4.

Revised: Thetechnical details can be seen Appendix G.4.

Revised: Bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) :

Original: Step 3.
Revised: Step 3.

Original: Now we shift gear to tackle (iii): handling the random variable“ l large".
Revised: Now we shift gear to tackle (iii): handling the random variable“ l large".

Original: We rewrite “ l large" into the indicator function as follows:
Revised: We rewrite “ l large" into the indicator function as follows:

Original: E  (cid:88) l large ( b1 )
Revised: E  (cid:88) l large [Equation] b

Original: [Equation]
Revised: [Equation]

Original: , where I k := I (cid:16) max i | ∂ l f i ( x k, 0 ) 
Revised: )  [Equation] [Equation] E (cid:34)

Revised: [Equation]

Original: | ≥ Q k :=
Revised: l = I [Equation] l f (

Revised: [Equation]

Original:  △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 (cid:17) is the indicator function ( I ( A ) =
Revised: 0 ) √ v l,k,n − 1 (cid:88) i = ( m l,k,i − ∂ l f i ( x k, 0 )) (cid:35) , where I k : = I (cid:16) max i | ∂ l f i ( x k, 0 ) | ≥ Q k := △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 (cid:17) is the indicator function ( I ( A ) =

Original: 1when event A holds and I ( A )
Revised: 1when event A holds and I ( A )

Original: = 

Revised: =

Original: otherwise.)
Revised: otherwise.)

Original: Similarly as before, when taking E k ( · ) , “ I k " can beregarded as a constant index.
Revised: Similarly as before, when taking E k ( · ) , “ I k " canbe regarded as a constant index.

Original: Therefore, E k (cid:104)(cid:80) dl
Revised: Therefore, E k (cid:104)(cid:80) dl

Original: =

Revised: =

Original: I k ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Revised: I k ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) can becalculated following the same idea as color-ball toy example.
Revised: (cid:105) canbe calculated following the same idea as color-ball model of the 1st kind.

Original: However, when taking  E k − 1 ( · ) , I kwill become a random variable which changes with different trajectories.
Revised: However, when taking 981 E k − 1 ( · ) , I k will become a random variable which changes with different trajectories.

Original: In this case, the color-ballmethod cannot be applied.
Revised: In this case,the color-ball method cannot be applied.

Original: Similarly as before, we introduce the following lemma to change I k into961 I k − 1 (defined later), which can again be regarded as a fixed when taking E k − 1 ( · ) .
Revised: Similarly as in Step 2 , we introduce the following lemma tochange I k into I k − 1 (defined later), which can again be regarded as a fixed when taking E k − 1 ( · ) .

Original: Lemma G.4.
Revised: Lemma G.3.

Original: Suppose Assumption 2.1 holds and β 1 < √ β 2 .
Revised: Suppose Assumption 2.1 holds and β 1 < √ β 2 .

Original: For 0 ≤ j ≤ k , we define
Revised: For 0 ≤ j ≤ k , we define

Original: I k − j := I (cid:16) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp = k − j Q p (cid:17) , where Q k :=
Revised: I k − j := I (cid:16) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp = k − j Q p (cid:17) , where Q k :=

Original: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , then we havethe following results.
Revised: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , then we havethe following results.

Original: I k
Revised: I k

Original: = I   max i | ∂ l f i ( x k, 0 )
Revised: = I   max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )
Revised: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )

Original: | ≥ k (cid:88) p =
Revised: | ≥ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original: 
Revised: 

Original: + I   max i | ∂ l f i ( x k, 0 ) | ≥ Q k and max i | ∂ l f i ( x k − j, 0 ) |
Revised: + I   max i | ∂ l f i ( x k, 0 ) | ≥ Q k and max i | ∂ l f i ( x k − j, 0 ) |

Original: ≤ k (cid:88) p =
Revised: ≤ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original:  , (30)
Revised:  , (31)

Original: I   max i | ∂ l f i ( x k, 0 )
Revised: I   max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )
Revised: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )

Original: | ≥ k (cid:88) p =
Revised: | ≥ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original:  =
Revised:  =

Original: I   max i | ∂ l f i ( x k − j, 0 )
Revised: I   max i | ∂ l f i ( x k − j, 0 )

Original: | ≥ k (cid:88) p =
Revised: | ≥ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original:  =
Revised:  =

Original: I k − j .
Revised: I k − j .

Original: Proof.
Revised: Proof.

Original: Equation (30) is straightforward, we only prove (31) here.
Revised: Equation (31) is straightforward, we only prove (32) here.

Original: Under Assumption 2.1, we have| ∂ l f i ( x k, 0 ) − ∂ l f i ( x k − j, 0 ) | ≤ n △ n ( k − 1) + · · · + n △ n ( k − j ) ≤ Q k − 1 + · · · + Q k − j .
Revised: Under Assumption 2.1, we have| ∂ l f i ( x k, 0 ) − ∂ l f i ( x k − j, 0 ) | ≤ n △ n ( k − 1) + · · · + n △ n ( k − j ) ≤ Q k − 1 + · · · + Q k − j .

Original: To show the 967 second inequality, it requires comparing the value between Q k and n △ n ( k ) , which are both problemdependent constants.
Revised: To show the 989 second inequality, it requires comparing the value between Q k and n △ n ( k ) , which are both problemdependent constants.

Original: Here, the inequality holds when n △ n ( k ) ≤ Q k .
Revised: Here, the inequality holds when n △ n ( k ) ≤ Q k .

Original: If otherwise, we can alwaysdefine ˜ Q k :=
Revised: If otherwise, we can alwaysdefine ˜ Q k :=

Original: max { Q k , n △ n ( k ) } and the inequality still holds by changing all the Q k into ˜ Q k .
Revised: max { Q k , n △ n ( k ) } and the inequality still holds by changing all the Q k into ˜ Q k .

Original: Wetemporarily omit this step for now.
Revised: Wetemporarily omit this step for now.

Original: Since | ∂ l f i ( x k, 0 ) − ∂ l f i ( x k − j, 0 ) | ≤
Revised: Since | ∂ l f i ( x [Equation] , 0 ) − ∂ l f i ( [Equation]

Revised: | ≤

Original: Q k −+ · · · +
Revised: Q k −+ · · · +

Original: Q k − j , the event 972 (cid:110) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp
Revised: Q k − j , the event 994 (cid:110) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp

Original: = k − j Q p (cid:111) implies the event { max i | ∂ l f i ( x k, 0 )
Revised: = k − j Q p (cid:111) implies the event { max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k } , so theproof is completed.
Revised: | ≥ Q k } , so theproof is completed.

Original: Now we are ready to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .
Revised: Now we are ready to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .

Original: Combining Step 1, 2 and 3 together, we prove thefollowing Lemma G.5.
Revised: Combining Step 1, 2 and 3 together, we prove thefollowing Lemma G.4.

Original: Lemma G.5.
Revised: Lemma G.4.

Original: Under Assumption 2.1, consider β 1 < √ β 2 , when k is large such that: k ≥ 4 ;β ( k − 1)
Revised: Under Assumption 2.1, consider β 1 < √ β 2 , when k is large such that: k ≥ 4 ;β ( k − 1)

Original: n≤ β n 1 √ k − 1 , we have the following result:
Revised: n≤ β n 1 √ k − 1 , we have the following result:

Original: [Equation]
Revised: [Equation]

Original: where δ 1 = (1 − β 2 ) [Equation] , which goes to 0 when β 2 goes to 1.
Revised: where δ 1 = (1 − β 2 )4 nρ 2 β n+ (cid:18) 1 √ β n 2 −(cid:19) , which goes to 0 when β 2 goes to 1.

Original: Proof can be seen in Appendix G.4.
Revised: Proof can be seen in Appendix G.4.

Original: Now we bound the error terms: E (cid:104)(cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .
Revised: Now we bound the error terms: E (cid:104)(cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .

Original: Since they are multiplied by δ 1 , these two terms vanish when β 2 → 1 .
Revised: Since they are multiplied by δ 1 , these two terms vanish when β 2 → 1 .

Original: Bounding E (cid:104) (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .
Revised: Bounding E (cid:104)(cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .

Original: We bound these two terms in the following Lemma G.6.
Revised: We bound these two terms in the following Lemma G.5.

Original: Lemma G.6.
Revised: Lemma G.5.

Original: Given all the history up to x k, 0 , we denote α as the index of the coordinate withthe greatest gradient: α = arg max l =1 , 2 , ··· ,d | ∂ l f ( x k, 0 ) | .
Revised: Given all the history up to x k, 0 , we denote α as the index of the coordinate withthe greatest gradient: α = arg max l =1 , 2 , ··· ,d | ∂ l f ( x k, 0 ) | .

Original: Under Assumption 2.1 and 2.2, considerβ 1 < √ β 2 , when k is large such that: k ≥ 4 ; β ( k − 1) n
Revised: Under Assumption 2.1 and 2.2, considerβ 1 < √ β 2 , when k is large such that: k ≥ 4 ; β ( k − 1) n

Original: [Equation]
Revised: [Equation]

Original: k − 1 , we have the following results:
Revised: k − 1 , we have the following results:

Original: [Equation]
Revised: [Equation]

Original: E 
Revised: l

Original: [Equation]
Revised: [Equation]

Original: where δ 1 is defined in (14) , D 0 is defined in Assumption 2.2.
Revised: where δ 1 is defined in (16) , D 0 is defined in Assumption 2.2.

Original: Detailed proof of Lemma G.6 is shown in Appendix G.8.
Revised: Detailed proof of Lemma G.5 is shown in Appendix G.8.

Original: Bounding E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .
Revised: Bounding E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .

Original: With the help of Lemma G.5and G.6, we can bound E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .
Revised: With the help of Lemma G.4and G.5, we can bound E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .

Original: The intuition is asfollows:
Revised: The intuition is asfollows:

Original: • ( a 1 ) is in the order of | ∂ l f ( x k, 0 ) | , multiplied by some positive constant;
Revised: • ( a 1 ) is in the order of | ∂ l f ( x k, 0 ) | , multiplied by some positive constant;• (cid:80) l large { ( a 2 )

Original: • (cid:80) l large { ( a 2 ) + ( b 2 ) } vanishes when β 2 → 1 .
Revised: + ( b 2 ) } vanishes when β 2 → 1 .• E (cid:110)(cid:80) l large ( b 1 ) (cid:111) vanishes when β 2 → 1 and k goes to infinity

Original: • E (cid:110)(cid:80) l large ( b 1 ) (cid:111) vanishes when β 2 → 1 and k goes to infinity.
Revised: Therefore, E (cid:110)(cid:80) [Equation] [Equation] +

Original: Therefore, E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b2 ) } (cid:111) is still in the order of | ∂ l f ( x k, 0 ) |(multiplied by a positive constant when β 2 is large).
Revised: [Equation] l large [Equation] + ( b

Revised: [Equation]

Revised: } (cid:111) is still in the order of | ∂ l f ( x k, 0 ) |(multiplied by a positive constant when β 2 is large).

Original: More formal results are shown in Lemma G.7.
Revised: More formal results are shown in Lemma G.6.

Original: Lemma G.7.
Revised: Lemma G.6.

Original: Under Assumption 2.1 and 2.2, when the hyperparameters satisfy: i) β 1 < √ β 2 < 1 ,ii) β 2 is large enough such that A ( β 2 ) is small enough to satisfy (32) , where A ( β 2 ) is a non-negativeconstant that approaches 0 when β 2 approaches 1 .
Revised: Under Assumption 2.1 and 2.2, when the hyperparameters satisfy: i) β 1 < √ β 2 < 1 ,ii) β 2 is large enough such that A ( β 2 ) is small enough to satisfy (33) , where A ( β 2 ) is a non-negativeconstant that approaches 0 when β 2 approaches 1 .

Original: More Specifically, A ( β 2 ) needs to satisfy ( ρ 1 , ρ 2and ρ 3 are defined in Appendix F.1). 
Revised: More Specifically, A ( β 2 ) needs to satisfy ( ρ 1 , ρ 2and ρ 3 are defined in Appendix F.1). [Figure]

Original: [Equation]
Revised: [Equation]

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]