Adam with model exponential moving average is effective for nonconvex optimization

Authors: kwangjun  ahn
ashok  cutkosky
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 16:57:47

Abstract: In this work, we offer a theoretical analysis of two modern optimization techniques for training large and complex models: (i) adaptive optimization algorithms, such as Adam, and (ii) the model exponential moving average (EMA). Specifically, we demonstrate that a clipped version of Adam with model EMA achieves the optimal convergence rates in various nonconvex optimization settings, both smooth and nonsmooth. Moreover, when the scale varies significantly across different coordinates, we demonstrate that the coordinate-wise adaptivity of Adam is provably advantageous. Notably, unlike previous analyses of Adam, our analysis crucially relies on its core elements---momentum and discounting factors---as well as model EMA, motivating their wide applications in practice.

Main paper: 
 introduction: 
 In neural network training, the training loss F : R d → R is often optimized using an iterative optimization algorithm which starts with the initial iterate x 0 and then updates during each iteration t = 1, 2, . . . as follows:x t = x t-1 + z t , where z t denotes the increment chosen by the algorithm during the t-th iteration. One of the most popular optimization algorithms is Adam [Kingma and Ba, 2014] . Adam has gained significant attention due to its effectiveness in training Transformer-based language models [Zhang et al., 2020a , Kunstner et al., 2023 , Jiang et al., 2023 , Pan and Li, 2023 , Ahn et al., 2024a , Kunstner et al., 2024 , Zhang et al., 2024a] .The model exponential moving average (EMA) [Polyak and Juditsky, 1992, Ruppert, 1988 ] is an optimization technique that has gained popularity in conjunction with Adam for various recent applications. EMA maintains an exponential moving average of the model iterates, x t , which contributes to the stabilization of these iterates. There has been a resurgence of interest in this technique due to its effectiveness in training high-quality generative models [Yaz et al., 2018 , Karras et al., 2019 , Song et al., 2021b , Dhariwal and Nichol, 2021 , Nichol and Dhariwal, 2021 , Song et al., 2021a , Balaji et al., 2022 , Karras et al., 2022 , Rombach et al., 2022 , Kang et al., 2023 , Karras et al., 2023] . Moreover, a recent work by Block et al. [2024] demonstrates the effectiveness of EMA for both language modeling and imitation learning applications.In this work, we theoretically study the effectiveness of these two modern optimization techniques. Our main results can be informally summarized as follows.Theorem 1 (Informal). CLIPPED-ADAM with the EMA on its iterates achieves the optimal convergence rate for nonconvex optimization both for smooth and nonsmooth settings (Section 5). The coordinate-wise adaptivity of Adam is particularly effective when the scale varies across different coordinates (Section 6).Assumption 2. Let F : R d → R be a differentiable function with the following properties:• Let ∆ := F (x 0 ) -inf x F (x).• For any two points x and y, F (y) -F (x) = 1 0 ⟨∇F (x + t(y -x)), y -x⟩ dt.• Lipschitzness. F is G-Lipshitz, i.e., for any point x, ∥∇F (x)∥ ≤ G.• Stochastic gradient variance. For any point x, the stochastic gradient g ← STOGRAD(x, r) for randomness r ∈ Z satisfies E[g] = ∇F (x) and E ∥g -∇F (x)∥ 2 ≤ σ 2 .The Lipschitz continuity condition is a standard assumption in nonconvex nonsmooth settings. However, as we will discuss in Section 6, one of the key insights from our results is that Adam enables us to adapt to the Lipschitz constants coordinate-wise without requiring prior knowledge of these constants. Note that we almost certainly need some form of structural assumption on the "difficulty" of the loss function; thus, relaxing the Lipschitz assumption would likely come at the cost of another assumption, such as smoothness.The second condition, called well-behavedess in [Cutkosky et al., 2023, Definition 1] , is a mild regularity condition. For any locally Lipschitz function F , applying an arbitrarily small perturbation to the function is sufficient to ensure this condition [Cutkosky et al., 2023, Proposition 2] .For the notion of optimality, we follow Zhang and Cutkosky [2024] and consider the following notion of stationarity for nonconvex and nonsmooth functions. This notion is a slight relaxation of the notion of a Goldstein stationarity point [Goldstein, 1977] , which was further studied by recent works [Zhang et al., 2020a , Davis et al., 2022b , Tian et al., 2022 , Jordan et al., 2023] . Definition 3 ((λ, ε)-stationary point). Suppose F : R d → R is differentiable. We say x is a (λ, ε)-stationary point of F if ∥∇F (x)∥ [λ] ≤ ε, where∥∇F (x)∥ [λ] := inf p∈P(R d ), Ey∼p[y]=x ∥E[∇F (y)]∥ + λ • E ∥y -x∥ 2 .To further motivate this definition, we remark that (λ, ϵ)-stationary points retain the desirable properties of Goldstein stationary points. Specifically, the following result [Zhang and Cutkosky, 2024, Lemma 2.3] demonstrates that, akin to Goldstein stationary points, (λ, ϵ)-stationary points can be reduced to first-order stationary points with appropriate choices of λ when the objective function is smooth or second-order smooth. Lemma 4. If F is L-smooth, then an (L 2 ε -1 , ε)-stationary point x of F satisfies ∥∇F (x)∥ ≤ 2ε. Moreover, if F is H-second-order-smooth, then an (H/2, ε)-stationary point x of F satisfies ∥∇F (x)∥ ≤ 2ε.Moreover, as shown by [Zhang and Cutkosky, 2024, Lemma 2.4 ], (λ, ε)-stationary points can also be reduced to Goldstein stationary points whenF is Lipschitz. Lemma 5. Suppose F is G-Lipschitz. For any λ, ε, δ > 0, a (λ, ε)-stationary point is a (δ, ε ′ )- Goldstein stationary point, where ε ′ = (1 + 2G λδ 2 )• ε. Now we design algorithms that find (λ, ε)-stationary points efficiently.3 Discounted-to-nonconvex conversion: online learning of incrementsOur main results are built on the online-to-nonconvex conversion framework of Cutkosky et al. [2023] . At its core, this framework involves selecting the increment z t using an online learner, as discussed by Ahn et al. [2024b] . Specifically, we follow a variant developed by Zhang and Cutkosky [2024] , which carefully incorporates the discounting factor in the conversion process. Note that we make slight modifications to the version proposed by Zhang and Cutkosky [2024] as follows. Here Exp(1) denotes the exponential random variable with mean 1.Given Algorithm 1, it turns out we need to design an online learner that minimizes the discounted regret, formally defined below. It is worth noting that discounted regret has been recently studied with the goal of better adapting online learners to dynamic environments [Ahn et al., 2024b , Zhang et al., 2024b , Jacobsen and Cutkosky, 2024] .Algorithm 1 Discounted-to-nonconvex conversion (choosing increments via online learning) Input: Initial point x 0 , T ∈ N, online learning algorithm A, and discounting factor β ∈ (0, 1) for t = 1, 2 . . . , T do Receive z t from A // choose the increment using an online learner Update x t ← β-β t 1-β t x t-1 + 1-β 1-β t x t (Equivalently, x t ← 1-β 1-β t t s=1 β t-s x s ) end forDefinition 6 (Discounted regret). For a comparator u, the β-discounted regret is defined asRegret [β] t (u) := β t • t s=1 (ℓ [β] s (z s ) -ℓ [β] s (u)) = t s=1 β t-s ⟨g s , z s -u⟩ .The discounted regret of an online learner A can be used to upper bound the norm of averaged gradients, as shown in the following result.Lemma 7 (Discounted-to-nonconvex conversion). Suppose that F satisfies Assumption 2. Then for the comparator sequence chosen asu t := -D t s=1 β -s ∇F (xs) ∥ t s=1 β -s ∇F (xs)∥ , Algorithm 1 gives E t∼[T ] E E yt ∇F (y t ) ≤ ∆ DT + 2G + σ (1 -β)T + σ 1 -β + 1 DT β • E Regret [β] T (u T ) + (1 -β) • T t=1 E Regret [β] t (u t ) ,where y t is distributed over {x s } t s=1 as P(y t = x s ) = β t-s • 1-β 1-β t for s = 1, 2, . . . , t.The proof combines the techniques of [Cutkosky et al., 2023, Theorem 7] and [Zhang and Cutkosky, 2024, Theorem 3.3] . See Appendix A for details.We briefly explain how Lemma 7 can be used to find a (λ, ϵ)-stationary point (Definition 3). Recall that (λ, ϵ)-stationarity essentially requires producing a point x = E[y] such that both ∥E[∇F (y)]∥ and E∥y -x∥ 2 are small.Given this context, Lemma 7 states that as long as the discounted regret of the online learner A is small, we can ensure that the EMA iterates xt = E[y t ] serve as good candidates for (λ, ϵ)-stationarity, since the term E ∥E yt ∇F (y t )∥ can be kept small. The remaining task is to bound the variance term,E ∥y t -x t ∥2 , which will be addressed later in Lemma 10.Moreover, the comparator u t roughly models the update direction that an oracle algorithm with perfect knowledge of the loss would select. In the proof of Lemma 7, we demonstrate that moving along the u t direction effectively decreases the loss value, which forms the basis for establishing our convergence guarantee.Thanks to the discounted-to-nonconvex conversion, the task now reduces to designing an online learner that achieves low discounted regret.4 Scale-free Follow-the-Regularized-Leader (FTRL)In this section, we introduce an algorithmic component, called the Followed-The-Regularized-Leader (FTRL), a powerful online learning technique with various applications [Gordon, 1999 , Kalai and Vempala, 2005 , Shalev-Shwartz and Singer, 2006 , Abernethy et al., 2008 , Nesterov, 2009 , Hazan and Kale, 2010] .For the setting, consider the online linear optimization (OLO) setting, where during each round t = 1, . . . , T , and online learner chooses z t , and then the linear loss ℓ t (•) = ⟨v t , •⟩ is revealed by the environment. Here the goal of the online learner is to minimize the regret defined as t ⟨v t , z t -u⟩, where u is the comparator in hindsight. For this setting, FTRL is presented in Algorithm 2.Algorithm 2 Follow-the-Regularized-Leader (FTRL)Require: Regularizers {ψ t } : R d → R, the domain D ⊆ R d 1: for t = 1, 2, . . . , T do 2:Update z t ← argmin z∈D ψ t (z) +t-1 s=1 ℓ s (z)3:Receive the next loss ℓ t (•) = ⟨v t , •⟩ 4: end forThe key insight of Ahn et al. [2024b] and Zhang et al. [2024b] is that in order to design an online learner for discounted regret, it is important that the online learner is scale-free as described below. In particular, following Ahn et al. [2024b] , we consider a gradient adaptive scale-free FTRL algorithm called scale-free FTRL [Orabona and Pál, 2018 ].We will focus on the case where D = B D , the d-dimensional L 2 -ball of radius D > 0. Scale-free FTRL is given by Algorithm 2 with the following choie:ψ t (•) = 1 η t ∥•∥ 2 and η t = D t-1 s=1 ∥v s ∥ 2 .Then using the clipping operator clip D (x) := x min( D /∥x∥, 1), we can write down the update rule more explicitly as follows:z t ← argmin z∈B D 1 η t ∥z∥ 2 + t-1 s=1 ⟨v s , z⟩ = -clip D   D t-1 s=1 v s t-1 s=1 ∥v s ∥ 2   . (SCALE-FREE FTRL)Here, if the denominator is zero, i.e., v 1 = • • • = v t-1 = 0, then we set z t ← 0. Note that this algorithm is scale-free in the sense that when the loss sequence is scaled by a scalar c > 0, the updates remain the same.Let us now present the regret bound of SCALE-FREE FTRL.Lemma 8 (Gradient-adaptive regret bound). For any T > 0, loss sequence v 1:T and comparator u ∈ R d s.t. ∥u∥ ≤ D, SCALE-FREE FTRL guarantees the following regret bound:T t=1 ⟨v t , z t -u⟩ ≤ 4D T t=1 ∥v t ∥ 2 .We note that Lemma 8 follows (with a slightly worse constant) from [Orabona and Pál, 2018, Theorem 1] , and the version we invoke is here due to [Ahn et al., 2024b, Theorem A.1] .Recall from Lemma 7 that an online learner for the discounted-to-nonconvex conversion (Algorithm 1) needs to have a low discounted regret. To achieve this, following Ahn et al. [2024b] and Zhang et al. [2024b] , we simply substitute v t ← β -t g t into SCALE-FREE FTRL, resulting in the updatez t ← -clip D   D t-1 s=1 β -s g s t-1 s=1 β -2s ∥g s ∥ 2   . (β-FTRL)Here again, if the denominator is zero, i.e., g 1 = • • • = g t-1 = 0, then we set z t ← 0. Then, the following result characterizes the discounted regret guarantee of β-FTRL. Theorem 9 (Discounted regret bound). Let β ∈ (0, 1]. For any T > 0, loss sequence g 1:T and comparator u ∈ R d s.t. ∥u∥ ≤ D, β-FTRL guarantees the following static regret boundRegret [β] T (u) ≤ 4D T t=1 β 2(T -t) ∥g t ∥ 2 .We next use this result to design an algorithm for nonconvex optimization.5 Discounted-FTRL leads to adaptive nonconvex optimizationIn this section, as a warm-up, let us see the implications of choosing A = β-FTRL in Algorithm 1. First, let us obtain a bound on the expected discounted regret. By Theorem 9 together with Jensen's inequality, we have the following regret bound for any t = 1, 2, . . . , T :E Regret [β] t (u t ) ≤ 4D E T t=1 β 2(T -t) ∥g t ∥ 2 ≤ 4D T t=1 β 2(T -t) E ∥g t ∥ 2 . Since E ∥g t ∥ 2 ≤ G 2 + σ 2 and 1 √ 1-β 2 ≤ 1 √ 1-β , it follows that E Regret [β] t (u t ) ≤ 4D(G + σ) 1 -β 2 ≤ 4D(G + σ) √ 1 -β .(1)Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The abstract and introduction accurately reflect the paper's contributions regarding the analysis of Adam with EMA for nonconvex and nonsmooth optimization.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
limitations: 
 Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The paper discuss the limitations of the work in Section 7.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: The assumptions and the proofs are provided both in the main text and the appendix.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experimental result reproducibility: 
 Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [NA] Justification: This paper is a theory paper and does not have any experiments. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
from gradient-adaptive regret to nonconvex optimization: 
 In order to obtain nonconvex optimization guarantees in terms of the (λ, ε)-stationarity (Definition 3), we need to handle the variance term. Following [Zhang and Cutkosky, 2024, Lemma 3.2] , the variance term can be bounded as follows.Lemma 10 (Variance bound). Using the notations of Lemma 7, for any t = 1, 2, . . . , T , β-FTRL satisfiesE t∼[T ] E ∥y t -x t ∥ 2 ≤ 12 D 2 (1 -β) 2 .Proof. From [Zhang and Cutkosky, 2024, Lemma 3.2 ], it follows that E T t=1 ∥y t -x t ∥ 2 ≤ 12 (1-β) 2 E T t=1 ∥z t ∥ 2. Now since ∥z t ∥ ≤ D for all t = 1, 2, . . . , T , after dividing each side by T , we get the desired inequality.Plugging the regret bound (1) into Lemma 7 and combining it with Lemma 10, we arrive at the following optimization guarantee in terms of the (λ, ε)-stationarity. See Section B.1 for a proof.Theorem 11. Suppose that F satisfies Assumption 2 and consider any λ > 0. For C > 0, choose A = β-FTRL in Algorithm 1 with the following parameters:β = 1 - ε 10C 2 , D = (1 -β)ε 1/2 4λ 1/2 , and T = 1 1 -β • max 4∆λ 1/2 ε 3/2 , 12C ε . Then we have Et∼[T ] E ∥∇F (x t )∥ [λ] ≤ (1 + G+σ C )ε.In other words, a randomly chosen model EMA x t is a (λ, (1 + G+σ C )ε)-stationary point, in expectation.Here, we discuss several notable aspects of the guarantee provided in Theorem 11.As shown in [Zhang and Cutkosky, 2024, Corollary 5 .1], the lower bound on the iteration complexity for finding a (λ, ε)-stationary point isΩ((G + σ) 2 ∆λ 1/2 ε -7/2 ), provided that λ ≤ G 4 ∆ 2 ε -1. Theorem 11 implies that setting C = G + σ achieves this optimal iteration complexity. Corollary 12. In Theorem 11, choosing C = G + σ leads to the following iteration complexity for finding a (λ, ε)-stationary point:O max (G + σ) 2 ∆λ 1/2 ε 7/2 , (G + σ) 3 ε 3 .In particular, treating G, σ, and ∆ as constants, as long as λ ≳ ε, this leads to the optimal complexity of O((G + σ) 2 ∆λ 1/2 ε -7/2 ).In light of Lemma 4, the above optimal complexity can be converted into the optimal complexities for smooth settings.Corollary 13 (Smooth settings). Corollary 12 implies the following optimal iteration complexity for smooth settings. Choosing λ = O(ε -1 ), it implies the optimal complexity of O(ε -4 ) for smooth loss functions [Arjevani et al., 2023] . Similarly, with λ = O(1), it achieves the optimal iteration complexitiy of O(ε -7/2 ) for second-order smooth loss functions [Arjevani et al., 2020] .We next discuss the benefits of using the gradient-adaptive regret bound (Theorem 9) by considering the case where we do not have knowledge of G, σ.A remarkable consequence of Theorem 11 is that, due to the gradient-adaptive regret bound of Theorem 9, the final convergence guarantee has a better dependence on G, σ in the case when we do not have knowledge of them. For concreteness, in the following discussion, we treat G, σ, ∆ as constants, and focus on the regime λ ≳ ε.First, our Theorem 11 with C = 1 (since we do not know G, σ) leads to the following iteration complexity for finding a (λ, ε)-stationary point:O (G + σ) 7/2 ∆λ 1/2 ε -7/2The price we pay for not knowing G, σ relative to the lower bound is a multiplicative factor of (G + σ) 3/2 . To see the benefit of this adaptive regret approach, let us consider the guarantees given by Zhang and Cutkosky [2024] . Their approach is based on choosing online gradient descent for A in Algorithm 1, when the learning rate is not properly tuned with the knowledge of G and σ, it would lead to the following (suboptimal) discounted regret bound:E Regret [β] t (u t ) ≤ O D(G + σ) 2 √ 1 -β .Then, the resulting iteration complexity becomes O(∆λ1/2 ( ε (G+σ) 2 ) -7/2 ), which is equal to O (G + σ) 7 ∆λ 1/2 ε -7/2. This is larger than the complexity due to our adaptive approach by a multiplicative factor of (G + σ) 7/2 . Next, we build on the results from this section and consider a better approach to design an adaptive nonconvex optimizer.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [NA] Justification: This paper is a theory paper and does not have any experiments. Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
coordinate-wise adaptivity via (clipped-)adam: 
 In this section, we consider the setting where the Lipschitzness constants vary across different coordinates, which is empirically observed to be reflective of practical neural network training (see, e.g. [Crawshaw et al., 2022 , Zhuang et al., 2022] ). Formally, we consider the following setting. Assumption 14. Under the same setting as Assumption 2, we replace the last two conditions with the following coordinate-wise version:• For each coordinate i = 1, 2, . . . , d, there is a Lipschitzness constant G i > 0 and a variance constant σ i > 0 such that ∀x, |∂ i F (x)| ≤ G i and the stochastic gradient g ← STOGRAD(x, r) satisfies E[g[i]] = ∂ i F (x i ) and E |g[i] -∂ i F (x)| 2 ≤ σ 2 i . (Here, ∂ i F denotes the partial deriva- tive of F w.r.t. the i-th coordinate.)Let G := (G 1 , . . . , G d ) and σ := (σ 1 , . . . , σ d ). Then, the above condition implies the last two conditions in Assumption 2 with G = ∥G∥ 2 and σ = ∥σ∥ 2 .As we mentioned before, the previous approaches [Cutkosky et al., 2023, Zhang and Cutkosky, 2024] choose the online learner A to be online gradient descent, and hence choosing the learning rate requires the knowledge of G i , σ i for all i. However, for neural network training, d is equal to the number of parameters in the network, so tuning them individually is computationally infeasible. We instead consider running β-FTRL coordinate-wise in Algorithm 1, which will automatically adapt to each coordinate. We begin with an important observation that such an approach in fact leads to a popular optimizer widely used in practice.For notational simplicity, fix a coordinate among i = 1, 2, . . . , d, and let us denote the iterate by x t , the stochastic gradient by g t , and the update by z t . Then the resulting optimizer becomes:z t+1 = -clip D   D t s=1 β t-s 1 g s t s=1 β t-s 2 g 2 s   , (CLIPPED-ADAM)where β 1 = β and β 2 = β 2 . Here, again if the denominator is zero, i.e., if g 1 = • • • = g t = 0, then we set the update to be zero, i.e., z t+1 = 0. Note that this is almost exactly the Adam optimizer [Kingma and Ba, 2014] , except that now we add clipping to control the variance of the iterates relative to their EMA. Notably, CLIPPED-ADAM retains one of the most important properties of Adam: it is scale-invariant. The scale invariance causes the optimizer to make updates of the same magnitude on each coordinate even when the scale differs across different coordinates.In practice, we expect that the clipping operation will effectively be a no-op. This is because, when the algorithm is converging (even if the convergence is somewhat slow), the gradients are likely to behave as approximately mean-zero random variables (due to factors such as stochastic noise, unstable training trajectories, etc.). In such cases, standard concentration inequalities imply thatt s=1 β t-s g s ≲ t s=1 (β t-s g s ) 2, and hence, the clipping has no effect. We also remark that CLIPPED-ADAM does not consider the "bias correction" terms in the original updates of Adam [Kingma and Ba, 2014] . However, note that the bias correction terms are coordinateindependent, and they can be merged into the scalar D.We next discuss the theoretical guarantees of CLIPPED-ADAM for nonconvex and nonsmooth optimizaton. Inspired by [Duchi et al., 2010, McMahan and Streeter, 2010] , where the coordinate-wise online learners lead to regret bounds with respect to the L 1 norms of stochastic gradients, we consider the following variant of Definition 3, in the same vein as [Cutkosky et al., 2023, Section 4] .Definition 15 ((λ, ε)-L 1 -stationary point). Suppose F : R d → R is differentiable. We say x is a (λ, ε)-L 1 -stationary point of F if ∥∇F (x)∥ [λ] 1 ≤ ε, where ∥∇F (x)∥ [λ] 1 := inf p∈P(R d ), Ey∼p[y]=x ∥E[∇F (y)]∥ 1 + λ • E ∥y -x∥ 2 2 . Using the fact ∥•∥ 1 ≤ √ d ∥•∥ 2 , one can connect the two notions of (λ, ε)-stationary points. Lemma 16. A (λ/ √ d, ε/ √ d)-stationary point is a (λ, ε)-L 1 -stationary point.In order to obtain the guarantee in terms of L 1 -norm, we consider the coordinate-wise version of discounted-to-online conversion, in the same vein as [Cutkosky et al., 2023, Appendix G] . See Section A.1 for details.Lemma 17 (L 1 -variant of Lemma 7). Suppose that F satisfies Assumption 14. Consider the comparator sequence chosen as u t defined asu t [i] := -D t s=1 β -s ∂iF (xs) | t s=1 β -s ∂iF (xs)| for i = 1, 2, . . . , d. Then, Algorithm 1 gives E t∼[T ] E E yt ∇F (y t ) 1 ≤ ∆ DT + 2 ∥G + σ∥ 1 (1 -β)T + ∥σ∥ 1 1 -β + 1 DT β • E Regret [β] T (u T ) + (1 -β) • T t=1 E Regret [β] t (u t ) ,where y t is distributed over {x s } t s=1 as P(y t = x s ) = β t-s • 1-β 1-β t for s = 1, 2, . . . , t.Next, let us consider the (expected) regret bound. Fix a coordinate i = 1, . . . , d. Then, by the one-dimensional version of Theorem 9 together with Jensen's inequality, we have the following regret bound for any t = 1, 2, . . . , T :E Regret [β] t (u t [i]) ≤ 4D T t=1 β 2(T -t) E |g t [i]| 2 ≤ 4D(G i + σ i ) √ 1 -β .Hence, taking the sum over all coordinates i = 1, . . . , d, we obtainEQUATION)Combining these together, one get the following guarantee in terms of the L 1 norm. See Section B.2 for a proof.Theorem 18. Suppose that F satisfies Assumption 14 and consider any λ > 0. For C > 0, choose the coordinate-wise optimizer CLIPPED-ADAM in Algorithm 1 with the following parameters:β = 1 - ε 10C 2 , D = (1 -β)ε 1/2 4d 1/2 λ 1/2 , and T = 1 1 -β • max 4∆d 1/2 λ 1/2 ε 3/2 , 12C ε . Then we have Et∼[T ] E ∥∇F (x t )∥ [λ] 1 ≤ (1 + ∥G+σ∥ 1 C)ε. In other words, a randomly chosen modelEMA x t is a (λ, (1 + ∥G+σ∥ 1 C)ε)-L 1 -stationary point, in expectation.In this section, we discuss the benefits of coordinate-wise adaptivity by examining the guarantee from Theorem 18 and compare it with that of Theorem 11. We begin with the (λ, ε)-L 1 -stationary point guarantee due to Theorem 18. We consider the scenario where β is carefully tuned by making the optimal choice of C. Corollary 19. In Theorem 18, choosing C = ∥G + σ∥ 1 leads to the following iteration complexity for finding a (λ, ε)-L 1 -stationary point:O max ∥G + σ∥ 2 1 ∆d 1/2 λ 1/2 ε 7/2 , ∥G + σ∥ 3 1 ε 3 .In order to better appreciate the benefits of coordinate-wise adaptivity, let us compare the above iteration complexity with that of Theorem 11.For concreteness, we treat G = ∥G∥ 2 and σ = ∥σ∥ 2 as constants throughout, and more importantly, we assume that the coordinates are heterogeneous in the sense thatEQUATION)The assumption (3) roughly says that a few coordinates of G + σ take much larger values than the rest; if all the coordinates of G + σ have similar magnitudes, then ∥G + σ∥ 1 ≈ √ d ∥G + σ∥ 2 . In the case λ ≳ ε, Corollary 19 implies that the iteration complexity isEQUATIONNext, let us consider the counterpart that does not adapt to each coordinate separately. In this case, we apply Lemma 16, which tells us that to find a (λ, ε)-L 1 stationary point it suffices to find a (λ/ √ d, ε/ √ d)-stationary point. Then, from Corollary 12, the iteration complexity isEQUATIONHence, when (3) holds, (4) can be lower than (5) by a multiplicative factor of d, showing the benefits of coordinate-wise adaptivity.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [NA] Justification: This paper is a theory paper and does not have any experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
discussion: 
 Our analyses of Adam based on the discounted-to-online conversion is quite different than the previous ones. As discussed in Section 1.1, the previous analyses often result in guarantees that are not quite reflective of practice-e.g., the rates get better without momentum and the rates are no better than that of non-adaptive methods. In contrast, our analyses and results highlight the role of the practical components as highlighted below.• Momentum. In order to obtain a low discounted regret, any sensible online learner should integrate the past history of stochastic gradients g 1:t to make the decision z t+1 . Such online learners under the discounted-to-online conversion lead to momentum methods that integrate g 1:t to obtain the next increment z t+1 . In particular, non-momentum methods would correspond to aggressive online learners that only use the last gradient g t to make the decision z t+1 . This perspective provides new insights into understanding the role of momentum, as echoed by Ahn et al. [2024b] .• Adaptive learning rates. The adaptive learning rate due to β-FTRL leads to a gradient-adaptive regret bound (Theorem 9), which is important to obtain a better Lipshitzness dependence (Section 5.2) as well as the coordinate-wise adaptivity for high-dimension settings (Section 6.3). Our analysis offers theoretical benefits of adaptive learning rate from a discounted regret perspective.• Model EMA. Lastly, the discounted-to-nonconvex conversion (Algorithm 1) naturally leads to guarantees in terms of the model EMA, x t . At a high level (see Appendix A precise details), this is because for a dynamic environment, it is important to discount the losses such that online learners adapt to changing environments. The appearance of model EMA in the discounted-to-nonconvex conversion provides a new perspective on its role.Our analyses and results have several limitations and raise several interesting questions. Firstly, CLIPPED-ADAM does not precisely match the original Adam algorithm, warranting further investigation into the original Adam update. Specifically, our analysis suggests choosing β 1 = β and β 2 = β 2 , which does not align with the commonly used practical choices. Understanding the exact roles of these practical choices for β 1 and β 2 would be valuable.In Section 5.2, we observed that our iteration complexity for finding a (λ, ε)-stationary point is O(∆(G + σ) 7/2 λ 1/2 ε -7/2 ) when G and σ are unknown. Investigating whether this complexity is optimal presents another intriguing direction for future research.Lastly, from a practical standpoint, developing a more advanced online learner for discounted regret and designing an algorithm that surpasses Adam in practicality would have significant practical implications.where the last line follows from the fact E[⟨∇F (x t ) -g t , z t ⟩] = 0. More specifically, note that the randomness in the stochastic gradient oracle is independent of the randomness due to α t . Since E[∇F (x t ) -g t ] = 0, it follows thatE[⟨∇F (x t ) -g t , z t ⟩] = E[⟨E[∇F (x t ) -g t ], z t ⟩] = 0,where the inner expectation is with respect to the randomness in the stochastic gradient oracle and the outer is with respect to all other quantities. Now let us handle each term.1 : Note that using the definition of y n , we haveE n t=1 β n-t (1 -β) ⟨∇F (x t ), u n ⟩ = (1 -β) E n t=1 β n-t ∇F (x t ), -D n t=1 β n-t ∇F (x t ) ∥ n t=1 β n-t ∇F (x t )∥ = (1 -β n ) E n t=1 1 -β 1 -β n β n-t ∇F (x t ), -D n t=1 1-β 1-β n β n-t ∇F (x t ) n t=1 1-β 1-β n β n-t ∇F (x t ) = -D(1 -β n ) E E yn ∇F (y n ) ≤ -D E E yn ∇F (y n ) + DGβ n .Therefore, summing over n = 1, . . . , T , we obtain:E T n=1 n t=1 β n-t (1 -β) ⟨∇F (x t ), u n ⟩ ≤ -D E T t=1 E yt ∇F (y t ) + DG 1 -β .2 : For the second term, using Cauchy-Schwartz inequality, we haveE n t=1 β n-t ⟨∇F (x t ) -g t , -u n ⟩ ≤ E n t=1 β n-t (∇F (x t ) -g t ) 2 E ∥u n ∥ 2 .Using the bounded variance assumption on the stochastic gradient oracle, we haveE n t=1 β n-t (∇F (x t ) -g t ) 2 = E n t=1 β 2(n-t) ∥∇F (x t ) -g t ∥ 2 ≤ σ 2 1 -β 2 .Therefore, summing over n = 1, . . . , T , and using the fact that 1 1-β 2 ≤ 1 1-β , we get the following bound on the second term:E T n=1 n t=1 β n-t (1 -β) ⟨∇F (x t ) -g t , -u n ⟩ ≤ T n=1 (1 -β) • σD 1 -β 2 ≤ σDT 1 -β .3 : Lastly, for the third term, we haveE T n=1 n t=1 β n-t (1 -β) ⟨g t , z t -u n ⟩ = (1 -β) E T n=1 n t=1 E β n-t g t , z t -u n = (1 -β) E T t=1 Regret [β]t (u t ) .Note that for each t, since x t ← x t-1 + α t z t for α t ∼Exp(1), Lemma 20 yieldsE[F (x t ) -F (x t-1 )] = E ⟨∇F (x t ), z t ⟩ = E ⟨g t , z t ⟩ = E ⟨g t , z t -u T ⟩ + E ⟨g t , u T ⟩ ≤ E ⟨g t , z t -u T ⟩ + D(G + σ) .Thus,E T t=1 β T -t+1 (F (x t ) -F (x t-1 )) = β E T t=1 β T -t g t , z t -u T + β T -t D(G + σ) ≤ β E[Regret [β] T (u T )] + D(G + σ) 1 -β .Combining A and BCombining the above analyses and rearranging, it follows thatD E T t=1 E yt ∇F (y t ) ≤ ∆ + DG 1 -β + σDT 1 -β + (1 -β) E T t=1 Regret [β] t (u t ) + β E[Regret [β] T (u T )] + D(G + σ) 1 -β .Dividing both sides by DT , we get the desired result.The proof closely follows that of Lemma 7. In particular, with ∆ := F (x 0 ) -inf x F (x), we have-∆ ≤ E T n=1 n t=1 β n-t (1 -β)(F (x t ) -F (x t-1 )) A + E T t=1 β T -t+1 (F (x t ) -F (x t-1 )) B .We begin with the term B . Using the same decomposition as before, we haveE[F (x t ) -F (x t-1 )] = E ⟨g t , z t -u T ⟩ + E ⟨g t , u T ⟩ = E ⟨g t , z t -u T ⟩ + d i=1 E g t [i]u T [i] ≤ E ⟨g t , z t -u T ⟩ + D d i=1 (G i + σ i ) .Thus,E T t=1 β T -t+1 (F (x t ) -F (x t-1 )) = β E T t=1 β T -t g t , z t -u T + β T -t d i=1 D(G i + σ i ) ≤ β E[Regret [β] T (u T )] + D d i=1 (G i + σ i ) 1 -β .Moving onto the term A , we again use the same decomposition:E[F (x t ) -F (x t-1 )] = E ⟨∇F (x t ), u n ⟩ 1 + E ⟨∇F (x t ) -g t , -u n ⟩ 2 + E ⟨g t , z t -u n ⟩ 3 .As before, let us handle each term one by one separately.1 : Note that using the definition of y n , for each coordinate i = 1, . . . , d, we haveE n t=1 β n-t (1 -β)∂ i F (x t )u n [i] = (1 -β) E n t=1 β n-t ∇F (x t ) -D n t=1 β n-t ∂ i F (x t ) | n t=1 β s-t ∂ i F (x t )| = (1 -β n ) E   n t=1 1 -β 1 -β n β n-t ∇F (x t )   -D n t=1 1-β 1-β n β n-t ∇F (x t ) n t=1 1-β 1-β n β n-t ∇F (x t )     = -D(1 -β n ) E E yn ∂ i F (y n ) ≤ -D E E yn ∂ i F (y n ) + DG i β n .Therefore, summing over i = 1, . . . , d and then n = 1, . . . , T , we obtain:E T n=1 n t=1 β n-t (1 -β) ⟨∇F (x t ), u n ⟩ ≤ -D E T t=1 E yt ∇F (y t ) 1 + D i=1 G i 1 -β .2 : For each coordinate i = 1, . . . , d, we haveE n t=1 β n-t (∂ i F (x t ) -g t [i])(-u n [i]) ≤ E n t=1 β n-t (∂ i F (x t ) -g t [i]) 2 E |u n [i]| 2 .Using the coordinate-wise bounded variance assumption on the stochastic gradient oracle,E n t=1 β n-t (∂ i F (x t ) -g t [i]) 2 = E n t=1 β 2(n-t) |∂ i F (x t ) -g t [i]| 2 ≤ σ 2 i 1 -β 2 .Therefore, summing over n = 1, . . . , T , and using the fact that 1 1-β 2 ≤ 1 1-β , we get the following bound on the second term:E T n=1 n t=1 β n-t (1 -β) ⟨∇F (x t ) -g t , -u n ⟩ ≤ T n=1 (1 -β) • D d i=1 σ i 1 -β 2 ≤ DT d i=1 σ i 1 -β .3 : We use the same manipulation as before:E T n=1 n t=1 β n-t (1 -β) ⟨g t , z t -u n ⟩ = (1 -β) E T t=1 Regret [β] t (u t ) .Combining the above, we get the desired result in Lemma 17.By Definition 3, since E[yt] = x t , it holds thatE t∼[T ] ∥∇F (x t )∥ [λ] ≤ E t∼[T ] E yt ∇F (y t ) + λ E yt ∥y t -x t ∥ 2 .We begin with the second term (the variance term). By Lemma 10, we haveλ E t∼[T ] E yt ∥y t -x t ∥ 2 ≤ 12 λD 2 (1 -β) 2 .D = (1-β)ε 1/2 4λ 1/2 , it follows that λ • Et∼[T ] Ey t ∥y t -x t ∥ 2 ≤ ε.Next consider the first term (the norm of the averaged gradients). Plugging the regret bound (1) into Lemma 7, we getE t∼[T ] E yt ∇F (y t ) ≤ ∆ DT + 2G + σ (1 -β)T + σ 1 -β + 4(G + σ) T √ 1 -β + 4(G + σ) 1 -β ≤ 4∆λ 1/2 (1 -β)ε 1/2 T + 6G + 5σ (1 -β)T + (4G + 5σ) 1 -β ,where the last line follows since1 √ 1-β ≤ 1 1-β and D = (1-β)ε 1/2 4λ 1/2 . Choosing β = 1 -( ε 10C ) 2 , the last term is bounded by G+σ 2C ε. Moreover, choosing T = (1 -β) -1 • max 4∆λ 1/2 ε -3/2 , 12Cε -1, the first and second terms are bounded by ε and G+σ 2C ε, respectively. This concludes the proof.By Definition 15, since E[yt] = x t , it holds thatE t∼[T ] ∥∇F (x t )∥ [λ] 1 ≤ E t∼[T ] E yt ∇F (y t ) 1 + λ E yt ∥y t -x t ∥ 2 2 .We begin with the second term (the variance term). This time, given that now each coordinate of update z t is bounded by D, i.e., |z t [i]| ≤ D, applying the variance bound due to Lemma 10 coordinate-wise implies:λ E t∼[T ] E yt ∥y t -x t ∥ 2 2 ≤ 12 λdD 2 (1 -β) 2 .D = (1-β)ε 1/2 4d 1/2 λ 1/2 , it follows that λ Et∼[T ] Ey t ∥y t -x t ∥ 2 ≤ ε.Next consider the first term (the L 1 -norm of the averaged gradients). Plugging the regret bound (2) into Lemma 17, and doing similar manipulations as the proof of Theorem 11, we getE t∼[T ] E yt ∇F (y t ) 1 ≤ ∆ DT + ∥6G + 5σ∥ 1 (1 -β)T + ∥4G + 5σ∥ 1 1 -β = 4∆d 1/2 λ 1/2 (1 -β)ε 1/2 T + ∥6G + 5σ∥ 1 (1 -β)T + ∥4G + 5σ∥ 1 1 -β ,where the last line follows sinceD = (1-β)ε 1/2 4d 1/2 λ 1/2 . Choosing β = 1 -( ε 10C ) 2 , the last term is bounded by ∥G+σ∥ 1 2C • ε. Moreover, choosing T = (1 -β) -1 • max 4∆d 1/2 λ 1/2 ε -3/2 , 12Cε -1 , the firstand second terms are bounded by ε and ∥G+σ∥ 1 2C• ε, respectively. This concludes the proof.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [NA] Justification: This paper is a theory paper and does not have any experiments.• The answer NA means that the paper does not include experiments. • The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [NA] Justification: This paper is a theory paper and does not have any experiments.Guidelines:• The answer NA means that the paper does not include experiments. Justification: The research adheres to the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This paper is a theory paper, and we do not forsee any direct societal implications arising from the research.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper is a theory paper and does not have any experiments.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The paper accurately credits the existing results by citing the papers that inspired the techniques used in this work.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: The paper does not introduce any new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: This paper examines the benefit of model exponential moving average when the model is trained with Adam. In particular, the paper uses a discount-to-nonconvex conversion framework with a scale-free FTRL learner to obtain non convex optimization guarantees. This resultant algorithm resembles a clipped variant of Adam with model EMA, thus providing optimization guarantees for it. While there are differences with the version of Adam used in practice, the results are nevertheless insightful. The reviews for the paper were mostly positive except for one reviewer with low-confidence providing a lower score. While some of their concerns were valid, I believe the rebuttal mostly addresses them. I think this paper will be valuable addition to the conference. I recommend acceptance.

Reviews: review: paper_summary: paper_summary: This paper proposes a variant of Adam involving per-iteration clipping of the updates (which look like Adam's updates) and EMA-based weight averaging. More specifically, two versions of this variant are proposed -- one is a global (w.r.t. all the coordinates) version and one is a per-coordinate version. The two main technical tools used in developing the proposed algorithm are (i) the online-to-nonconvex framework of references Cutkosky et al. [2023] and Zhang and Cutkosky [2024] wherein the update directions are chosen with an online learner, and (ii) choosing the online learner to be scale-free Follow-the-Regularized-Leader (FTRL) based on reference Ahn et al. [2024b]. The proposed algorithm attains the optimal convergence rate for both smooth and non-smooth nonconvex settings (although w.r.t. a different notion of stationary points for the non-smooth case). In addition, the benefit of the per-coordinate version is discussed.
strength_weakness: strength_weakness: **1.** It appears to me that this paper's main convergence results (Theorems 11 and 19) can hold even with $\mathbf{y}_t$ as defined in Lemma 7. Since Theorems 11 and 19 are *not* w.r.t. the last *averaged* iterate (which is what is used in practice), it seems that EMA-based averaging is not strictly necessary in the context of the proposed algorithms, i.e., the results hold even with the unaveraged iterates appropriately sampled (as described in Lemma 7).

**2.** I understand that this is a theoretical paper but because a new algorithm has been proposed involving clipping and EMA-based averaging, it would have been nice to show some empirical results at least in some simple settings. This would have made the proposed algorithm more appealing and convincing to practitioners.
questions: questions: **1.** The purpose of Lemma 10 is not immediately clear to me. Is it to replace $\mathbf{y}_t$ in Lemma 7 with $\overline{\mathbf{w}}_t$?

**2.** It seems to me that clipping is always activated in $\beta$-FTRL and CLIPPED-ADAM(?) This is because with the choice of $\beta_1 = \beta$ and $\beta_2 = \beta^2$ in CLIPPED-ADAM for instance, $\sum_{s=1}^t \beta_1^{t-s} g_s \geq \sqrt{\sum_{s=1}^t \beta_2^{t-s} g_s^2}$. 

**3.** Is the Lipschitz assumption necessary?
limitations: limitations: Discussed.
paper_summary: paper_summary: The authors show that with clipping and model exponential average, (Clipping) Adam can perform better than SGD.
strength_weakness: strength_weakness: 1.  The paper is hard to follow.

    (i) In Algorithm 1, it would be better to add "output: $\bar{w}_T$"
    
    (ii) it is unclear what $v_s$ is in Definition 6. 

2. All of the following results are based on Lemma 7, I am not sure whether Lemma 7 is suitable for SGD.
questions: questions: If change Lemma 7 with a different formulation will the conclusion change?
limitations: limitations: N/A
paper_summary: paper_summary: In this paper, the authors analyze the online learning framework 'discounted-to-nonconvex conversion', and propose a bound for the expectation of the exponential average of the output gradient. Moreover, they find that if the online learner follows a specific loss structure(FTRL), then the output of this online learner would follow a clip-Adam rules.
strength_weakness: strength_weakness: Since this paper is built upon the 'discounted-to-nonconvex conversion' framework proposed in [1], focusing on the perspective of online learning, many definitions, concepts, and details require further elaborations and interpretations, especially for readers who are unfamiliar with this topic. I will list all my confusions and questions in the 'Questions' part and update my rating according to the response from the authors.

[1] Qinzi Zhang and Ashok Cutkosky. Random scaling and momentum for non-smooth non-convex optimization. In International Conference on Machine Learning. PMLR, 2024.
questions: questions: 1. In line 109, the description of the Algorithm 1, could the author explain the reason they introduce the notations $\mathbf{x_t}$ and $\mathbf{w_t}$, since from my perspective, $\mathbf{x_t}$ represent the iterations of coefficients of the loss function. Then the stochastic gradient $\mathbf{g_t}$ should also be computed by $\mathbf{x_t}$, instead of $\mathbf{w_t}$. I also wonder about the implication of $\alpha_t$ in the definition of $\mathbf{w_t}$, since it only benefits the further proof. 

2. I might be wrong, but there might be typos in the formula of Definition 6. According to the definition of $\ell_t^{[\beta]}(z)$ in Algorithm 1,  the RHS of formula might be $\beta^t\sum_{s=1}^t \langle \beta^{-s} \mathbf{v_s},  \mathbf{z_s} -\mathbf{\mu}\rangle$. Meanwhile, I also wonder whether $\mathbf{v_s}$ is actually \mathbf{g_s} defined in Algorithm 1 since I never find the definition of $\mathbf{v_s}$. Moreover, could the author explain the meanings and implications of such a definition?

3. Could the author explain the meaning of choosing such a comparator $u_t$ in Lemma 7.

4. I might be wrong, but as I checked the Lagrangian function of constrained optimization in line 137, I guess the RHS should be $-\text{clip}_D(\eta \sum \mathbf{v_s})$. 

5. Could the authors explain the details of how they get Corollary 12 from the results of the definition $T$ in Theorem 11. 

6. Could the authors explain why line 466 holds since I'm not sure about the independence among $\mathbf{z_t}$ and $\mathbf{g_t}$. Moreover, could the author explain about the randomness in their formulas and which variables they take expectation w.r.t during the proof?

7. Could the author further discuss the implications of their main result? Since from my perspective, it looks like they use the exponential discounted loss and choose FTRL learner in Discounted-to-nonconvex conversion, then they could get an updating rules share a similar structure with Adam. However, this result is far-fetched, and I have no idea what it implies.
limitations: limitations: No limitations.

Comments: title: Thank you
Thank you
comment: You're right of course  that if all the  $\mathbf{g_t}$s have the same sign, then the clipping will occur. Our observation is that this is extremely unlikely to happen. The reason is the following: imagine that the algorithm is converging (even if it is converging somewhat slowly). In such a case, we would expect the gradients to be roughly mean-zero random variables (e.g. https://arxiv.org/pdf/2111.06328 or https://arxiv.org/pdf/1508.00882), in which case the clipping should be expected to not occur.

  In fact, the case in which clipping does occur frequently results in "too good to be true" results that should provide further intuitive evidence of it being unlikely. For example, https://jmlr.org/papers/v18/17-079.html shows that if the prediction of FTRL is clipped on every round then in fact we would suffer only *logarithmic* regret rather than the standar $\sqrt{T}$ regret bound, and http://proceedings.mlr.press/v97/cutkosky19b.html shows that in the case that $\sum_{t=1}^T \mathbf{g_t}\ge \Omega(\sqrt{T})$ (even if clipping doesn't occur every single round), we can obtain preconditioning "for free" without any extra computation cost.
- To clarify the new scheme, the update is now $\mathbf{x_t}=  \mathbf{x_{t-1}} + \alpha_t \mathbf{z_t}$, where $\alpha_t$ is sampled from an exponential distribution (which is a very light-tailed distribution with mean 1 and variance 1). So $\alpha_t$ is still used.

  Via an elementary argument involving  fundamental theorem of calculus  and the pdf $p(\alpha) = \exp(-\alpha)$,  we obtain $\mathbb{E}[F(\mathbf{x_t}) - F(\mathbf{x_{t-1}})] = \mathbb{E}[\langle \mathbf{g_t}, \mathbf{x_t} -  \mathbf{x_{t-1}}\rangle]$.

  This identity says that the function gap is *exactly equal* to the linearization of the function gap. In a sense, the randomization makes the first-order Taylor approximation perfectly correct. In the current submission, we sample $\mathbf{g_t}$ at $\mathbf{w_t}$ to ensure this identity. However, changing $\alpha_t$ to be exponentially distributed and multiplying the update  $\mathbf{z_t}$ by $\alpha_t$ means that  this is no longer necessary. 


  In short, although $\mathbf{w_t}$ plays a significant role, the entirety of that role is contained in establishing the identity $\mathbb{E}[F(\mathbf{x_t}) - F(\mathbf{x_{t-1}})] = \mathbb{E}[\langle \mathbf{g_t}, \mathbf{x_t} -  \mathbf{x_{t-1}}\rangle]$. This other option keeps the same identity and so works just as well.

  That said, if the reviewers prefer to keep the results  the same as in the submitted version for which all of the analysis is completely available, we are of course happy to do so.

- Regarding why $\mathbb{E}[ F(\mathbf{w_t}) - \mathbf{g_t}|\mathbf{z_t}]=0$: this is is because $\mathbf{g_t}$ is a standard stochastic gradient oracle evaluated at $\mathbf{w_t}$ and as such $\mathbb{E}[\mathbf{g_t} | \mathbf{w_t}]=F(\mathbf{w_t})$. As a concrete illustration, suppose that $F=\mathbb{E}[f(x,r)]$, where $x$ is the model parameter, $r$ is a randomly sampled data point, and $f$ is the loss of the model on the data point. Then  $\mathbf{g_t}$ can be taken to be $\nabla  f(\mathbf{w_t}, r_t)$  for some i.i.d. sample $r$.  This formalism exactly captures the standard approach used in training. In this case, since $r_t$ is independent of $\mathbf{w_t}$ and $\mathbf{z_t}$, we have $\mathbb{E}[ F(\mathbf{w_t}) - \mathbf{g_t} | \mathbf{z_t}]=0$.

- Regarding the clipping: You're right that this is not too common in current practice. However, we feel this is more an analytical artifact than a real change in the algorithm because the clipping should not be active in most iterations. The reason for this is that intuitively as the algorithm convergence we expect the gradients to be dominated by noise so that they look roughly like mean-zero random variables. In this case, the clipping will not occur. On the other hand, if the clipping *does* occur frequently, then an alternative analysis based upon https://jmlr.org/papers/v18/17-079.html would actually show that we converge incredibly fast. This in turn would put us in the regime in which gradients are dominated by noise, and so clipping would stop occurring.
I thank the authors for their response.
I still have some resolved questions. Since the author claimed they had dropped the notation of $w_t$, which plays a significant role in their original proof, I can not check the correctness of their results in their new version. Despite this significant problem, I also feel confused about their explanations for the following points:
- Since they drop the variable $w_t$, why they still need $\alpha_t$ for their further proof? At least in their original version,  $\alpha_t$ was introduced in the definition of $w_t$, and now, I even don't know where is this $\alpha_t$ first introduced.
- The authors might misunderstand my concerns about line 466. Actually, I want to know why $E[\nabla F(w_t)-g_t|z_t]=0$.
- The authors's explanation of implications and motivations does not convince me, several definitions or settings, seems too far-fetched and specific. While the authors consider clipping a minor modification, I have never used it or seen it used by others in practice. Similar questions also exist for the definition of comparator $u_t$ and learning rate $\eta$. As I stated in my review, it seems only given such specific forms, the authors get an updating rules share a similar structure with Adam. However, I can not see any direct connections between these frameworks and Adam. 

Due to these questions, I have to keep my original rating score. However, I want to stress that I'm not familiar with online optimization. Specifically, my comments about these intermediate variables might be unfair and it's welcome to point it out with more convincing evidence. Moreover, I keep my confidence at 2 and if other reviewers feel more confident in their judgment, I would defer to them.
Thanks for the rebuttal; I'll keep my score. Regarding the direction of the inequality, I believe my direction is correct (with exact inequality and not inequality up to constants) if all the $g_s$'s are non-negative (this can be seen by squaring both sides).
Thank you for the response I have raised my score.