SDEs for Adaptive Methods: The Role of Noise

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 12:27:11

Abstract: Despite the vast empirical evidence supporting the efficacy of adaptive optimization methods in deep learning, their theoretical understanding is far from complete. In this work, we introduce novel SDEs for commonly used adaptive optimizers: SignSGD, RMSprop(W), and Adam(W). Our SDEs offer a quantitatively accurate description of these optimizers and help bring to light an intricate relationship between adaptivity, gradient noise, and curvature. Our novel analysis of SignSGD highlights a noteworthy and precise contrast to SGD in terms of convergence speed, stationary distribution, and robustness to heavy-tail noise. We extend this analysis to AdamW and RMSpropW, for which we observe that the role of noise is much more complex. Crucially, we support our theoretical analysis with experimental evidence by verifying our insights: this includes numerically integrating our SDEs using Euler-Maruyama discretization on various neural network architectures such as MLPs, CNNs, ResNets, and Transformers. Our SDEs accurately track the behavior of the respective optimizers, especially when compared to previous SDEs derived for Adam and RMSprop. We believe our approach can provide valuable insights into best training practices and novel scaling rules.

Main paper: 
 introduction: 
 Adaptive optimizers lay the foundation for effectively training of modern deep learning models.These methods are typically employed to optimize an objective function expressed as a sum across N individual data points: min x∈R d [f (x) := 1 coefficients: an approach which is theoretically sound in very limited settings (Kohatsu-Higa et al., 1997; Bishop and Del Moral, 2019) . Zhou et al. (2024) informally presented an SDE for (only) the parameters of AdamW: this is achieved under strong assumptions and various approximations, some of which are hard to motivate formally.Influence of noise on convergence. Several empirical papers demonstrate that adaptive algorithms adjust better to the noise during training. Specifically, (Zhang et al., 2020b ) noticed a consistent gap in the performance of SGD and Adam on language models and connected that phenomenon with heavy-tailed noise distributions. (Pascanu et al., 2013) suggests using gradient clipping to deal with heavy tail noise, and consequently several follow-up works analyzed clipped SGD under heavy-tailed noise (Zhang et al., 2020a; Mai and Johansson, 2021; Puchkin et al., 2024) . Kunstner et al. (2024) present thorough numerical experiments illustrating that a significant contributor to heavy-tailed noise during language model training is class imbalance, where certain words occur much more frequently than others. They demonstrate that adaptive optimization methods such as Adam and SignSGD can better adapt to such class imbalances. However, the theoretical understanding of the influence of noise in the context of adaptive algorithms is much more limited. The first convergence results on Adam and RMSprop were derived under bounded stochastic gradients assumption (De et al., 2018; Zaheer et al., 2018; Chen et al., 2019; Défossez et al., 2022) . Later, this noise model was relaxed to weak growth condition (Zhang et al., 2022; Wang et al., 2022) and its coordinate-wise version (Hong and Lin, 2023; Wang et al., 2024) and sub-gaussian noise (Li et al., 2023a) . SignSGD and its momentum version Signum were originally studied as a method for compressed communication (Bernstein et al., 2018) under bounded variance assumption, but with a requirement of large batches.Several works provided counterexamples where SignSGD fails to converge if stochastic and full gradients are not correlated enough (Karimireddy et al., 2019b; Safaryan and Richtarik, 2021) . In the case of AdamW, (Zhou et al., 2022 (Zhou et al., , 2024) ) provide convergence guarantees under restrictive assumptions such as bounded gradient and bounded noise. All aforementioned results only show that SignSGD, Adam, and RMSprop at least do not perform worse than vanilla SGD. None of them studied how noise affects the dynamics of the algorithm: In this work, we attempt to close this gap. 
limitations: 
 Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: See Section C.1.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. 
formal statements & insights: the sdes: 
 This section provides the general formulations of the SDEs of SignSGD (Theorem 3.2) and AdamW (Theorem 3.12). Due to the technical nature of the analysis, we refer the reader to the appendix for the complete formal statements and proofs.Assumptions and notation. In this section, we assume that ∇f γ (x) = ∇f (x)+Z(x), E[Z(x)] = 0 and, unless we study the cases where the gradient variance is unbounded, we write Cov(Z(x)) = Σ(x) where we omit the batch size unless relevant. To derive the stationary distribution around an optimum, we will approximate the loss function with a quadratic convex function f (x) = 1 2 x ⊤ Hx as commonly done in the literature (Ge et al., 2015; Levy, 2016; Jin et al., 2017; Poggio et al., 2017; Mandt et al., 2017; Compagnoni et al., 2023) . Regarding the notation, η > 0 is the step size, the mini-batches {γ k } are of size B ≥ 1 and modeled as i.i.d. random variables uniformly distributed on {1, . . . , N }. The β parameters refer to momentum parameters, γ > 0 is the (decoupled)L 2 -regularization parameter, and ϵ > 0 is a small scalar used for numerical stability.The following definition formalizes the idea that an SDE can be a "good model" to describe an optimizer. It is drawn from the field of numerical analysis of SDEs (see Mil'shtein (1986) ) and it quantifies the disparity between the discrete and the continuous processes.Definition 3.1 (Weak Approximation). A continuous-time stochastic process {X t } t∈[0,T ] is an order α weak approximation (or α-order SDE) of a discrete stochastic process {x k } ⌊T /η⌋ k=0 if for every polynomial growth function g, there exists a positive constant C, independent of the stepsize η, such that max k=0,...,⌊T /η⌋ |Eg (x k ) -Eg (X kη )| ≤ Cη α .In this section, we derive an SDE model for SignSGD, which we believe to be a novel addition to the existing literature. This derivation will reveal the unique manner in which noise influences the dynamics of SignSGD. First, we recall the update equation of SignSGD: The following theorem derives a formal continuous-time model for SignSGD.EQUATIONTheorem 3.2 (Informal Statement of Theorem C.5). Under sufficient regularity conditions, the solution of the following SDE is an order 1 weak approximation of the discrete update of SignSGD:EQUATIONwhere Σ(x) is the noise covariance Σ(x) = E[ξ γ (x)ξ γ (x) ⊤ ] and ξ γ (x) := sign(∇f γ (x)) -1 + 2P(∇f γ (x) < 0) the noise in the sample sign (∇f γ (x)).For didactic reasons, we next present a corollary of Theorem 3.2 that provides a more interpretable SDE. Figure 1 shows the empirical validation of this model for various neural network classes: All details are presented in Appendix F.Corollary 3.3 (Informal Statement of Corollary C.7). Under the assumptions of Theorem 3.2, and that the stochastic gradient is ∇f γ(x) = ∇f (x)+Z such that Z ∼ N (0, Σ), Σ = diag(σ 2 1 , • • • , σ 2 d ),the following SDE provides a 1 weak approximation of the discrete update of SignSGDEQUATIONwhere the error function Erf(x) and the square are applied component-wise.While Eq. (3) may appear intricate at first glance, it becomes apparent upon closer inspection that the properties of the Erf(•) function enable a detailed exploration of the dynamics of SignSGD. In particular, we demonstrate that the dynamics of SignSGD can be categorized into three distinct phases. The left of Figure 2 empirically verifies this result on a convex quadratic function.Lemma 3.4. Under the assumptions of Corollary 3.3 and signal-to-noise ratioY t := Σ -1 2 ∇f (Xt) √ 2 , 1. Phase 1: If |Y t | > 3 2, the SDE coincides with the ODE of SignGD:dX t = -sign(∇f (X t ))dt; (4) 2. Phase 2: If 1 < |Y t | < 3 2 : 1 (a) mY t + q -≤ dE[Xt] dt ≤ mY t + q + ;(b) For any a > 0, P ∥X t -E [X t ]∥ 2 2 > a ≤ η a d -∥mY t + q -∥ 2 2 ;3. Phase 3:If |Y t | < 1, the SDE is dX t = - 2 π Σ -1 2 ∇f (X t )dt + √ η I d - 2 π diag Σ -1 2 ∇f (X t ) 2 dW t .(5) 1 Let m and q1 are the slope and intercept of the line secant to the graph of Erf(x) between the points (1, Erf(1)) and 3 2 , Erf 3 2 , while q2 is the intercept of the line tangent to the graph of Erf(x) and slope m,(q + )i := q2 if ∂if (x) > 0 -q1 if ∂if (x) < 0 , (q -)i := q1 if ∂if (x) > 0 -q2 if ∂if (x) < 0, and q := max(q1, q2). The dynamics of the moments of X t predicted in Lemma 3.7 track the empirical ones (Center-Right); If the schedulers satisfy the condition in Lemma 3.9, the loss decays to 0 as prescribed. Otherwise, the loss does not converge to 0 (Right).The behavior of SignSGD depends on the size of the signal-to-noise ratio. In particular, the SDE itself shows that in Phase 3, the inverse of the scale of the noise Σ -12 premultiplies the gradient, thus affecting the rate of descent. This is not the case for SGD where Σ only influences the diffusion term. 2 To better understand the role of the noise, we need to study how it affects the dynamics of the loss and compare it with SGD.Lemma 3.5. Let f be µ-strongly convex, T r(∇ 2 f (x)) ≤ L τ , and S t := f (X t ) -f (X * ). Then, during 1. Phase 1, the loss will reach0 before t * = 2 S0 µ because S t ≤ 1 4 √ µt -2 √ S 0 2 ; 2. Phase 2 with ∆ := m √ 2σmax + ηµm 2 4σ 2 max : E[S t ] ≤ S 0 e -2µ∆t + η 2 (Lτ -µdq 2 ) 2µ∆ 1 -e -2µ∆t ; 3. Phase 3 with ∆ := 2 π 1 σmax + η π µ σ 2 max : E[S t ] ≤ S 0 e -2µ∆t + η 2 Lτ 2µ∆ 1 -e -2µ∆t .In Phase 1, the signal-to-noise ratio is large, meaning that SignSGD behaves like SignGD: Consistently with the analysis of SignGD in (Ma et al., 2022) , this explains the fast initial convergence of the optimizer as well as of RMSprop and Adam. In this phase, the loss undergoes a steady decrease which ensures the emergence of Phase 2 which in turn triggers that of Phase 3 which is characterized by an exponential decay to an asymptotic loss level: As a practical example, we verify the dynamics of the expected loss around a minimum in the center-left of Figure 2 .Lemma 3.6. For SGD, the expected loss satisfies:E[S t ] ≤ S 0 e -2µt + η 2 Lτ σ 2 max 2µ 1 -e -2µt .Remark: The two key observations are that:1. Both in Phase 2 and Phase 3, the noise level σ max inversely affects the exponential convergence speed, while this trend is not observed with SGD; 2. The asymptotic loss of SignSGD is (almost) linear in σ max while that of SGD is quadratic.Additionally, we characterize the stationary distribution of SignSGD around a minimum: Empirical validation is provided in the center-right of Figure 2 .Lemma 3.7. Let H = diag(λ 1 , . . . , λ d ) and M t := e -2 √ 2 π Σ -1 2 H+ η π Σ -1 2 H 2 t . Then, 1. E [X t ] = e - √ 2 π Σ -1 2 Ht X 0 t→∞ → 0; 2. Cov [X t ] = M t -e -2 √ 2 π Σ -1 2 Ht X 2 0 + η 2 2 π I d + η π H -1 H -1 Σ 1 2 (I d -M t ) , which as t → ∞ converges to η 2 2 π I d + η π H -1 H -1 Σ 1 2 .Lemma 3.8. Under the same assumptions as Lemma 3.7, the stationary distribution for SGD is:E [X t ] = e -Ht X 0 t→∞ → 0 and Cov [X t ] = η 2 H -1 Σ I d -e -2Ht t→∞ → η 2 H -1 Σ.As we observed above, the noise inversely affects the convergence rate of the iterates of SignSGD while it does not impact that of SGD. Additionally, while both covariance matrices essentially scale inversely to the hessian, that of SignSGD scales with Σ 1 2 while that of SGD scales with Σ.We conclude this section by presenting a condition on the step size scheduler that ensures the asymptotic convergence of the expected loss to 0 in Phase 3. For general schedulers, we characterize precisely the speed of convergence and the factors influencing it. Empirical validation is provided in the right of Figure 2 for a convex quadratic.Lemma 3.9. Under the assumptions of Lemma 3.5, any step size scheduler η t such that∞ 0 η s ds = ∞ and lim t→∞ η t = 0 =⇒ E[f (X t ) -f (X * )] t→∞ → ≲ L τ σ max 4µ π 2 η t t→∞ → 0. (6) Remark: Under the same conditions, SGD satisfies E[f (X t ) -f (X * )] t→∞ → ≲ Lτ σ 2 max 4µ η t t→∞ → 0.As noted in Bernstein et al. (2018) , the signal-to-noise ratio is key in determining the dynamics of SignSGD. Our SDEs help clarify the mechanisms underlying the dynamics of SignSGD: we show that the effect of noise is radically different from SGD: 1) It affects the rate of convergence of the iterates, of the covariance of the iterates, and of the expected loss; 2) The asymptotic loss value and covariance of the iterates scale in Σ 1 2 while for SGD it does so in Σ. On the one hand, low levels of noise will ensure a faster and steadier loss decrease close to minima for SignSGD than for SGD. On the other, SGD will converge to much lower loss values. A symmetric argument holds for high levels of noise, which suggests that SignSGD is more resilient to high levels of noise.Interestingly, we can replicate the efforts above also in case the noise structure is heavy-tailed as it is distributed according to a Student's t distribution. Notably, we derive the SDE for the case where the noise has infinite variance and show how little marginal effect this has on the dynamics of SignSGD.Lemma 3.10. Under the assumptions of Corollary 3.3 but the noise on the gradients U ∼ t ν (0, I d ) where ν ∈ Z + : The following SDE is a 1 weak approximation of the discrete update of SignSGDEQUATIONwhere Ξ(x) is defined asΞ(x) := x Γ( ν+1 2 ) √ πνΓ( ν 2 ) 2 F 1 1 2 , ν+1 2 ; 3 2 ; -x 2 νand 2 F 1 (a, b; c; x) is the hypergeometric function. Above, the function Ξ(x) and the square are applied component-wise.We now characterize the dynamics of SignSGD when the noise on the gradient has infinite variance.Corollary 3.11. Under the assumptions of Lemma 3.10 and ν = 2, the dynamics in Phase 3 is:EQUATIONWe observe that the dynamics of SignSGD when the noise is Gaussian (Eq. ( 5)) and when the noise is heavy-tailed with unbounded variance (Eq. ( 8)) are very similar: By comparing the constants in front of the drift terms Σ -1 2 ∇f (X t ), they are only ∼ 10% apart, and the diffusion coefficients are comparable. Not only do we once more showcase the resilience of SignSGD to high levels of noise, but in alignment with (Zhang et al., 2020b) , we provide theoretical support to the success of Adam in such a scenario where SGD would diverge.All the results derived above can be extended to this setting: this is left as an exercise for the reader.In the last subsection, we showcased how SDEs can serve as powerful tools to understand the dynamics of the simplest among coordinate-wise adaptive methods: SignSGD. Here, we extend the discussion to Adam with decoupled weight decay, i.e. AdamW:EQUATIONwhich, of course, covers Adam, RMSprop, and RMSpropW depending on the values of γ and β 1 .The following result proves the SDE of AdamW which we validate in Figure 3 for two simple landscapes and in Figure 4 for a Transformer and a ResNet.Theorem 3.12 (Informal Statement of Theorem C.31). Under sufficient regularity conditions,ρ 1 = O(η -ζ ) s.t. ζ ∈ (0, 1), and ρ 2 = O(1), the order 1 weak approximation of AdamW is:EQUATIONEQUATION)EQUATIONwhereβ i = 1 -ηρ i ∼ 1, γ i (t) = 1 -e -ρit, andP t = diag √ V t + ϵ γ 2 (t)I d .In contrast to Remark 4.3 of Malladi et al. (2022) , which suggests that an SDE for RMSprop and Adam is only viable if σ ≫ ∥∇f (x)∥ and σ ∼ 1 η , our derivation that does not need these assumptions:See Remark C.25 for a deeper discussion, the implications, and the experimental comparison.The following result demonstrates how the asymptotic expected loss of AdamW scales with the noise level. Notably, it introduces the first scaling rule for AdamW, extending the one proposed for Adam in (Malladi et al., 2022) to include weight decay scaling. It is crucial to understand that, unlike the typical approach in the literature (see (Jastrzebski et al., 2018; Malladi et al., 2022) ), our objective in deriving these rules is not to maintain the dynamics of the optimizers or the SDE unchanged. Instead, our goal is to offer a practical strategy for adjusting hyperparameters (e.g., from η to η) to retain certain performance metrics or optimizer properties as the batch size increases (e.g., from B to B).Therefore, in our upcoming analysis, we aim to derive scaling rules that preserve specific relevant aspects of the dynamics, such as the convergence bound on the loss or the speed. For a more detailed discussion motivating our approach, see Appendix E.Lemma 3.13. If f is µ-strongly convex and L-smooth, L τ := Tr(∇ 2 f (x)), and (∇f (x)) 2 = O(η), η = κη, B = Bδ, and ρi = α i ρ i , and γ = ξγ, AdamW satisfiesEQUATION)We derive the novel scaling rule by 1) Preserving the upper bound, which requires that κ = √ δ and ξ = κ; 2) Preserving the relative speed of M t , V t and X t , which requires that βi = 1 -κ 2 (1 -β i ).The left of Figure 5 shows the empirical verification of the predicted loss value and scaling rule on a convex quadratic function. 3 Interestingly, and consistently with Lemma 3.13, such a value is not Losses -AdamW 1 = 0.999, 2 = 0.998 1 = 0.999, 2 = 0.996 1 = 0.999, 2 = 0.992 1 = 0.9, 2 = 0.998 1 = 0.9, 2 = 0.996 1 = 0.9, 2 = 0.992 AdamW is run with regularization parameter γ = 1. AdamW R (AdamW Rescaled) is run as we apply the scaling rule with κ = 2. AdamW NR (AdamW Not Rescaled) is run as we apply the scaling rule with κ = 2 on all hyperparameters but γ, which is left unchanged: Our scaling rule holds, and failing to rescale γ leads the optimizer not to preserve the asymptotic loss level. The same happens for γ = 4 (Left); The same for RMSpropW (Center-Left); For AdamW, β 1 and β 2 influence which basin will attract the dynamics and how fast this will converge, but not the asymptotic loss level inside the basin (Center-Right). For both AdamW and RMSpropW, the variance at convergence predicted in Lemma 3.14 matches the experimental results (Right).0influenced by the choice of β i : We argue that β i do not impact the asymptotic level of the loss, but rather drive the selection of the basin and speed at which AdamW converges to it -The center-right of Figure 5 exemplifies this on a simple nonconvex landscape.We conclude this section with the stationary distribution of AdamW around a minimum which we empirically validate on the right of Figure 5 .Lemma 3.14. The stationary distribution of AdamW is(E[X ∞ ], Cov[X ∞ ]) = 0, η 2 I d + γH -1 Σ 1 2 -1 H -1 Σ 1 2.RMSpropW We derived the same results for RMSprop(W) and we reported them in Appendix C.4: importantly, we validate the SDE in Figure 3 for two simple landscapes and in Figure 4 for aTransformer and a ResNet. The results regarding the asymptotic loss level and stationary distributions are validated in the center-left and right of Figure 5 for a convex quadratic function.While for both SignSGD and Adam the asymptotic loss value and the covariance of the iterates scale linearly with Σ 1 2 , we observe for AdamW this is more intricate: The interaction between curvature, noise, and regularization implies that these two quantities are upper-bounded in Σ 1 2 and increasing Σ to infinity does not lead to their explosion: Weight decay plays a crucial stabilization role at high noise levels near the minimizer -See Figure 6 for a comparison across optimizers. Finally, we argue that β i play a key role in selecting the basin and the convergence speed to the asymptotic loss value rather than impacting the loss value itself.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]• 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: sde validation: 
 The point of our experiments is to validate the theoretical results derived from the SDEs. Therefore, we first show that our SDEs faithfully represent the dynamics of their respective optimizers. To do For each optimizer, we plot the loss value on a convex quadratic and compare its asymptotic value with the limits predicted by our theory. As we take Σ = σ 2 I d , we confirm that the loss of SGD scales quadratically in σ (Lemma 3.6), and linearly for SignSGD (Lemma 3.5) and Adam (Lemma 3.13 with γ = 0). For AdamW, the maximum asymptotic loss value is bounded in σ (Lemma 3.13 with γ > 0). In accordance with the experiments, our theory predicts that adaptive methods are more resilient to noise. so, we integrate the SDEs with Euler-Maruyama (Algorithm 1): This is particularly challenging and expensive as one needs to calculate the full gradients of the DNNs at each iteration. 4 We present the first set of validation experiments on a variety of architectures and datasets: An MLP on the Breast Cancer dataset, a CNN and a Transformer on MNIST, and a ResNet on CIFAR-10. All details are in Appendix F. 
conclusion: 
 We derived the first formal SDE for SignSGD, enabling us to demonstrate its dynamics traversing three discernible phases. We characterize how the signal-to-noise ratio drives the dynamics of the loss in each of these phases, and we derive the asymptotic value of the loss function, as well as the stationary distribution. Regarding the role of noise, we draw a straightforward comparison with SGD. For SignSGD, the noise level 
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: [Yes] Justification: As we state in Section F, we run our experiments on an NVIDIA GeForce RTX 3090.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: All we do is derive some convergence bounds and similar results.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: [Yes] Justification: It can have a positive impact as it helps understand adaptive optimizers better.Possibly, it might help reduce the cost of fine-tuning thanks to our novel scaling law.Guidelines:• The answer NA means that there is no societal impact of the work performed. 
 
 Appendix: Σ has an inverse linear effect on the convergence speed of the loss and the iterates. However, it linearly affects the asymptotic expected loss and the asymptotic variance of the iterates. In contrast, for SGD, noise does not influence the convergence speed but has a quadratic impact on the loss level and variance. We also examine the scenario where the noise has infinite variance and demonstrate the resilience of SignSGD, showing that its performance is only marginally affected. Finally, we generalize the analysis to include AdamW and RMSpropW.Specifically, we leverage our novel SDEs to derive the asymptotic value of the loss function, their stationary distribution on a convex quadratic, and a novel scaling rule. The key insight is that, similarly to SignSGD, the loss level and covariance matrix of the iterates of Adam and RMSprop scale linearly in the noise level Σ 1 2 . For AdamW and RMSpropW, the complex interaction of noise, curvature, and regularization implies that these two quantities are bounded in terms of Σ 1 2 , showing that weight decay plays a crucial stabilization role at high noise levels near the minimizer. Interestingly, the SDEs for Adam and RMSprop are straightforward corollary of our general results and were derived under much less restrictive and more realistic assumptions than those in the literature. Finally, we thoroughly validate all our theoretical results: We compare the dynamics of the various optimizers with the respective SDEs and find good agreement on simple landscapes and deep neural networks.For Adam and RMSprop, our SDEs track them better than those derived in (Malladi et al., 2022) .We believe that our results can be extended to other optimizers commonly used in practice such as Signum, AdaGrad, AdaMax, and Nadam. Additionally, inspired by the insights from our SDE analysis, there is potential for designing new optimization algorithms that combine the strengths of existing methods while mitigating their weaknesses. For example, developing hybrid optimizers that adaptively switch between different strategies based on the training phase or current state of the optimization process could offer superior performance. In this section, we list some papers that derived or used SDEs to model optimizers. In particular, we focus on the aspect of empirically verifying the validity of such SDEs in the sense that they indeed track the respective optimizers. We divide these into three categories: Those that did not carry out any type of validation, those that did it on simple landscapes (quadratic functions et similia), and those that did small experiments or neural networks.None of the following papers carried out any experimental validation of the approximating power of the SDEs they derived. Many of them did not even validate the insights derived from the SDEs: (Liu et al., 2021; Hu et al., 2019; Bercher et al., 2020; Zhu and Ying, 2021; Cui et al., 2020; Maulén Soto, 2021; Wang and Wu, 2020; Lanconelli and Lauria, 2022; Ayadi and Turinici, 2021; Soto et al., 2022; Li and Wang, 2022; Wang and Mao, 2022; Bardi and Kouhkouh, 2022; Chen et al., 2022; Kunin et al., 2023; Zhang et al., 2023; Sun et al., 2023; Li et al., 2023b; Gess et al., 2024; Dambrine et al., 2024; Maulen-Soto et al., 2024) .The following ones carried out validation experiments on artificial landscapes, e.g. quadratic or quartic function, or easy regression tasks: (Li et al., 2017 (Li et al., , 2019;; Zhou et al., 2020b; An et al., 2020; Fontaine et al., 2021; Gu et al., 2021; Su and Lau, 2023; Ankirchner and Perko, 2024) .The following papers carried out some experiments which include neural networks: (Paquette et al., 2021; Compagnoni et al., 2023) . In particular, they both simulate the SDEs with a numerical integrator and compare them with the respective optimizers: The first validates the SDE on a shallow MLP while the second does so on a shallow and a deep MLP. Regarding (Li et al., 2021; Malladi et al., 2022) , they do not validate their SDEs: Rather, their approach conceptually proceeds as follows:1. Derive an SDE for an optimizer which we now dub "A";2. Notice that simulating the SDE is too expensive;3. Define another discrete-time algorithm called SVAG which also has the same SDE as "A" but does not numerically integrate the SDE as it does not even require access to it: It does not need access neither to the drift nor to the diffusion term;4. Simulate SVAG and show that it tracks "A" successfully;5. Conclude that the SDE is a good approximation for "A".However, they never validated that the SDE is a good approximation for "A" or for SVAG either.With the same logic, they could have done the following:1. Derive an SDE for "A";2. Notice that simulating the SDE is too expensive;3. Define another discrete-time algorithm called "B" which coincides with "A" and thus of course shares the same SDE; 4. Simulate "B" and show that it tracks "A" perfectly;5. Conclude that the SDE is a good approximation for "A".In particular, the only fact they prove is that SVAG is a discrete-time optimizer that shares the same SDE as "A" because it describes a discrete trajectory that is a 1st-order approximation of the SDE of "A". Technically speaking, "A" also does the same. One cannot conclude that the SDE derived for "A" is a good model for "A" by simply comparing two algorithms "A" and "B" that share the same SDE.Otherwise, simply comparing an optimizer "A" with itself would do the trick. An SDE's empirical validation can only occur if the SDE is simulated with a numerical integrator that requires access to the drift and diffusion terms (Higham, 2001; Milstein, 2013) .In this section, we summarize some important results in the analysis of Stochastic DifferentialEquations Mao (2007) ; Øksendal (1990) . The notation and the results in this section will be used extensively in all proofs in this paper. We assume the reader to have some familiarity with Brownian motion and with the definition of stochastic integral (Ch. 1.4 and 1.5 in Mao (2007) ).We start with some notation: Let (Ω, F, {F t } t≥0 , P) be a filtered probability space. We say that an event E ∈ F holds almost surely (a.s.) in this space ifP(E) = 1. We call L p ([a, b], R d ), with p > 0, the family of R d -valued F t -adapted processes {f t } a≤t≤b such that b a ∥f t ∥ p dt ≤ ∞.Moreover, we denote byM p ([a, b], R d ), with p > 0, the family of R d -valued processes {f t } a≤t≤b in L([a, b], R d ) such that E b a ∥f t ∥ p dt ≤ ∞. We will write h ∈ L p R + , R d , with p > 0, if h ∈ L p [0, T ], R d for every T > 0.Similar definitions hold for matrix-valued functions using theFrobenius norm ∥A∥ := ij |A ij | 2 .Let W = {W t } t≥0 be a one-dimensional Brownian motion defined on our probability space and letX = {X t } t≥0 be an F t -adapted process taking values on R d . Definition B.1. Let the drift be b ∈ L 1 R + , R d and the diffusion term be σ ∈ L 2 R + , R d×m .X t is an Itô process if it takes the formX t = x 0 + t 0 b s ds + t 0 σ s dW s .We shall say that X t has the stochastic differentialEQUATION)Theorem B.2 (Itô's Lemma). Let X t be an Itô process with stochastic differential dX t = b t dt + σ t dW t . Let f (x, t) be twice continuously differentiable in x and continuously differentiable in t, taking values in R. Then f (X t , t) is again an Itô process with stochastic differentialEQUATIONStochastic Differential Equations (SDEs) are equations of the formdX t = b(X t , t)dt + σ(X t , t)dW t .First of all, we need to define what it means for a stochastic process X = {X t } t≥0 with values in R dto solve an SDE.Definition B.3. Let X t be as above with deterministic initial condition X 0 = x 0 . Assume b :R d × [0, T ] → R d and σ : R d × [0, T ] → R d×m are Borel measurable; X t is called a solution to the corresponding SDE if 1. X t is continuous and F t -adapted; 2. b ∈ L 1 [0, T ], R d ; 3. σ ∈ L 2 [0, T ], R d×m ;4. For every t ∈ [0, T ]X t = x 0 + t 0 b(X s , s)ds + t 0 σ(X s , s)dW (s) a.s.Moreover, the solution X t is said to be unique if any other solution X ⋆ t is such thatP {X t = X ⋆ t , for all 0 ≤ t ≤ T } = 1.Notice that since the solution to an SDE is an Itô process, we can use Itô's Lemma. The following theorem gives a sufficient condition on b and σ for the existence of a solution to the corresponding SDE.Theorem B.4. Assume that there exist two positive constants K and K such that1. (Global Lipschitz condition) for all x, y ∈ R d and t ∈ [0, T ] max{∥b(x, t) -b(y, t)∥ 2 , ∥σ(x, t) -σ(y, t)∥ 2 } ≤ K∥x -y∥ 2 ; 2. (Linear growth condition) for all x ∈ R d and t ∈ [0, T ] max{∥b(x, t)∥ 2 , ∥σ(x, t)∥ 2 } ≤ K(1 + ∥x∥ 2 ).Then, there exists a unique solution X t to the corresponding SDE, andX t ∈ M 2 ([0, T ], R d ).Numerical approximation. Often, SDEs are solved numerically. The simplest algorithm to provide a sample path (x k ) k≥0 for X t , so that X k∆t ≊ xk for some small ∆t and for all k∆t ≤ M is called Euler-Maruyama (Algorithm 1). For more details on this integration method and its approximation properties, the reader can check Mao (2007) .Algorithm 1 Euler-Maruyama Integration Method for SDEs input The drift b, the volatility σ, and the initial condition x 0 . Fix a stepsize ∆t;Initialize x0 = x 0 ; k = 0; while k ≤ T ∆t do Sample some d-dimensional Gaussian noise Z k ∼ N (0, I d ); Compute xk+1 = xk + ∆t b(x k , k∆t) + √ ∆t σ(x k , k∆t)Z k ; k = k + 1; end while output The approximated sample path (x k ) 0≤k≤⌊ T ∆t ⌋ .In this section, we introduce the theoretical framework used in the paper, together with its assumptions and notations.First of all, many proofs will use Taylor expansions in powers of η. For ease of notation, we introduce the shorthand that whenever we write O (η α ), we mean that there exists a function K(x) ∈ G such that the error terms are bounded by K(x)η α . For example, we writeb(x + η) = b 0 (x) + ηb 1 (x) + O η 2 to mean: there exists K ∈ G such that |b(x + η) -b 0 (x) -ηb 1 (x)| ≤ K(x)η 2 .Additionally, we introduce the following shorthand:• A multi-index is α = (α 1 , α 2 , . . . , α n ) such that α j ∈ {0, 1, 2, . . .}; • |α| := α 1 + α 2 + • • • + α n ; • α! := α 1 !α 2 ! • • • α n !; • For x = (x 1 , x 2 , . . . , x n ) ∈ R n , we define x α := x α1 1 x α2 2 • • • x αn n ; • For a multi-index β, ∂ |β| β f (x) := ∂ |β| ∂ β 1 x 1 ∂ β 2 x 2 •••∂ βn xn f (x);• We also denote the partial derivative with respect to x i by ∂ ei .Definition C.1 (G Set). Let G denote the set of continuous functions R d → R of at most polynomial growth, i.e. g ∈ G if there exists positive integers ν 1 , ν 2 > 0 such that |g(x)| ≤ ν 1 1 + |x| 2ν2 , for all z ∈ R d .The next results are inspired by Theorem 1 of Li et al. (2017) and are derived under some regularity assumption on the function f . Assumption C.2. Assume that the following conditions on f, f i , and their gradients are satisfied:• ∇f, ∇f i satisfy a Lipschitz condition: there exists L > 0 such that|∇f (u) -∇f (v)| + n i=1 |∇f i (u) -∇f i (v)| ≤ L|u -v|;• f, f i and its partial derivatives up to order 7 belong to G;• ∇f, ∇f i satisfy a growth condition: there exists M > 0 such that|∇f (x)| + n i=1 |∇f i (x)| ≤ M (1 + |x|). Lemma C.3 (Lemma 1 Li et al. (2017)). Let 0 < η < 1. Consider a stochastic process X t , t ≥ 0 satisfying the SDE dX t = b (X t ) dt + √ ησ (X t ) dW twith X 0 = x ∈ R d and b, σ together with their derivatives belong to G. Define the one-step difference ∆ = X η -x, and indicate the i-th component of ∆ with ∆ i . Then we have1. E∆ i = b i η + 1 2 d j=1 b j ∂ ej b i η 2 + O η 3 ∀i = 1, . . . , d; 2. E∆ i ∆ j = b i b j + σσ T (ij) η 2 + O η 3 ∀i, j = 1, . . . , d; 3. E s j=1 ∆ (ij ) = O η 3 for all s ≥ 3, i j = 1, . . . , d.All functions above are evaluated at x.Theorem C.4 (Theorem 2 and Lemma 5, Mil'shtein (1986) ). Let Assumption C.2 hold and let us define ∆ = x 1 -x to be the increment in the discrete-time algorithm, and indicate the i-th component of ∆ with ∆i . If in addition there existsK 1 , K 2 , K 3 , K 4 ∈ G so that 1. E∆ i -E ∆i ≤ K 1 (x)η 2 , ∀i = 1, . . . , d; 2. E∆ i ∆ j -E ∆i ∆j ≤ K 2 (x)η 2 , ∀i, j = 1, . . . , d; 3. E s j=1 ∆ ij -E s j=1 ∆ij ≤ K 3 (x)η 2 , ∀s ≥ 3, ∀i j ∈ {1, . . . , d}; 4. E 3 j=1 ∆ij ≤ K 4 (x)η 2 , ∀i j ∈ {1, . . . , d}.Then, there exists a constant C so that for all k = 0, 1, . . . , N we have|Eg (X kη ) -Eg (x k )| ≤ Cη.Modeling of discrete-time algorithms using SDEs relies on Assumption C.2. As noted by Li et al. (2021) , the approximation can fail when the stepsize η is large or if certain conditions on ∇f and the noise covariance matrix are not met. Although these issues can be addressed by increasing the order of the weak approximation, we believe that the primary purpose of SDEs is to serve as simplification tools that enhance our intuition: We would not benefit significantly from added complexity.In this subsection, we provide the first formal derivation of an SDE model for SignSGD. Let us consider the stochastic process X t ∈ R d defined as the solution ofEQUATIONwhereEQUATIONand ξ γ (x) := sign(∇f γ (x)) -1 + 2P(∇f γ (x) < 0) the noise in the sample sign (∇f γ (x)). The following theorem guarantees that such a process is a 1-order SDE of the discrete-time algorithm of SignSGD Consider the stochastic process X t defined in Eq. ( 16) and fix some test function g ∈ G and suppose that g and its partial derivatives up to order 6 belong to G. Then, under Assumption C.2, there exists a constant C > 0 independent of η such that for all k = 0, 1, . . . , N , we haveEQUATIONwith x 0 ∈ R d , η ∈ R >0|Eg (X kη ) -Eg (x k )| ≤ Cη.That is, the SDE ( 16) is an order 1 weak approximation of the SignSGD iterations (18).Lemma C.6. Under the assumptions of Theorem C.5, let 0 < η < 1 and consider x k , k ≥ 0 satisfying the SignSGD iterationsx k+1 = x k -ηsign (∇f γ k (x k )) with x 0 ∈ R d .From the definition the one-step difference ∆ = x 1 -x, then we have1. E ∆i = -(1 -2P (∂ i f γ < 0)) η ∀i = 1, . . . , d; 2. E ∆i ∆j = (1 -2P (∂ i f γ < 0)) (1 -2P (∂ j f γ < 0)) + Σ(ij) η 2 ∀i, j = 1, . . . , d; 3. E s j=1 ∆ij = O η 3 ∀s ≥ 3, i j ∈ {1, . . . , d}.All the functions above are evaluated at x.Proof of Lemma C.6. First of all, we have that by definitionEQUATIONwhich impliesEQUATIONSecond, we have that by definitionEQUATIONwhich implies thatEQUATIONFinally, by definitionEQUATIONwhich concludes our proof.Proof of Theorem C.5. To prove this result, all we need to do is check the conditions in Theorem C.4.As we apply Lemma C.3, we make the following choices:• b(x) = -(1 -2P (∇f γ (x) < 0)); • σ(x) = Σ(x).First of all, we notice that ∀i = 1, . . . , d, it holds that• E ∆i 1. Lemma C.6 = -(1 -2P (∂ i f γ (x) < 0)) η; • E∆ i 1. Lemma C.3 = -(1 -2P (∂ i f γ (x) < 0)) η + O η 2 .Therefore, we have that for some K 1 (x) ∈ G,EQUATIONAdditionally, we notice that ∀i, j = 1, . . . , d, it holds that• E ∆i ∆j 2. Lemma C.6 = (1 -2P (∂ i f γ (x) < 0)) (1 -2P (∂ j f γ (x) < 0)) η 2 + Σ(ij) (x)η 2 ; • E∆ i ∆ j 2. Lemma C.3 = (1 -2P (∂ i f γ (x) < 0)) (1 -2P (∂ j f γ (x) < 0)) + Σ(ij) (x) η 2 + O η 3 .Therefore, we have that for some K 2 (x) ∈ G,EQUATIONAdditionally, we notice that ∀s ≥ 3, ∀i j ∈ {1, . . . , d}, it holds that• E s j=1 ∆ij 3. Lemma C.6 = O η 3 ; • E s j=1 ∆ ij 3. Lemma C.3 = O η 3 .Therefore, we have that for someEQUATION)Additionally, for some K 4 (x) ∈ G, ∀i j ∈ {1, . . . , d},EQUATION)To conclude, Eq. ( 25), Eq. ( 26), Eq. ( 27), and Eq. ( 28) allow us to conclude the proof.Corollary C.7. Let us take the same assumptions of Theorem C.5, and that the stochastic gradient is ∇f γ (x) = ∇f (x) + U such that U ∼ N (0, Σ) that does not depend on x. Then, the following SDE provides a 1 weak approximation of the discrete update of SignSGDEQUATION)where the error function Erf(x) and the square are applied component-wise, andΣ = diag σ 2 1 , • • • , σ 2 d .Proof of Corollary C.7. First of all, we observe thatEQUATION)where Φ is the cumulative distribution function of the standardized normal distribution. Remembering thatEQUATIONwe have thatEQUATIONSimilarly, one can prove that Σ defined in (17) becomesEQUATION)Corollary C.8. Let us take the same assumptions of Theorem C.5, and that the stochastic gradient is ∇f γ (x) = ∇f (x) + √ ΣU such that U ∼ t ν (0, I d ) that does not depend on x and ν is a positive integer number. Then, the following SDE provides a 1 weak approximation of the discrete update of SignSGDEQUATION)where Ξ(x) is defined asEQUATIONand 2 F 1 (a, b; c; x) is the hypergeometric function. Above, function Ξ(x) and the square are applied component-wise, andΣ = diag σ 2 1 , • • • , σ 2 d .Proof of Corollary C.8. First of all, we observe thatEQUATIONwhere F ν (x) is the cumulative function of a t distribution with ν degrees of freedom. Remembering thatEQUATIONwe have thatEQUATIONSimilarly, one can prove that Σ defined in (17) becomesEQUATION)Lemma C.9. Under the assumptions of Corollary C.7 and signal-to-noise ratioY t := Σ -1 2 ∇f (Xt) √ 2 , 1. Phase 1: If |Y t | > 3 2, the SDE coincides with the ODE of SignGD:EQUATIONProof of Lemma C.9. Exploiting the regularity of the Erf function, we approximate the SDE in (29) in three different regions:1. Phase 1:If |x| > 3 2 , Erf(x) ∼ sign(x). Therefore, if Σ -1 2 ∇f (Xt) √ 2 > 3 2 , (a) Erf Σ -1 2 ∇f (Xt) √ 2 ∼ sign Σ -1 2 ∇f (Xt) √ 2 = sign (∇f (X t )); (b) Erf Σ -1 2 ∇f (Xt) √ 2 2 ∼ sign Σ -1 2 ∇f (Xt) √ 2 2 = (1, . . . , 1).Therefore,EQUATION2. Phase 2: Let m and q 1 are the slope and intercept of the line secant to the graph of Erf(x)between the points (1, Erf(1)) and 3 2 , Erf 3 2 , while q 2 is the intercept of the line tangent to the graph of Erf(x) and slope m. If 1 < x < 3 2 , we have thatEQUATIONEQUATIONTherefore, we have that ifEQUATIONwhereEQUATIONandEQUATIONTherefore,m √ 2 Σ -1 2 ∇f (X t ) + q -≤ dE [X t ] dt ≤ m √ 2 Σ -1 2 ∇f (X t ) + q + ; (48) (b) Similar to the above, m √ 2 Σ -1 2 ∇f (X t ) + q - 2 ≤ Erf Σ -1 2 ∇f (X t ) √ 2 2 ≤ m √ 2 Σ -1 2 ∇f (X t ) + q + 2Therefore,EQUATION) 3. Phase 3: If |x| < 1, Erf(x) ∼ 2 √ π . Therefore, if Σ -1 2 ∇f (Xt) √ 2 < 1, (a) Erf Σ -1 2 ∇f (Xt) √ 2 ∼ 2 π Σ -1 2 ∇f (X t ); (b) Erf Σ -1 2 ∇f (Xt) √ 2 2 ∼ 2 π Σ -1 2 ∇f (X t ) 2 .Therefore,EQUATIONLemma C.10 (Dynamics of Expected Loss). Let f be µ-strongly convex, T r(∇ 2 f (x)) ≤ L τ , andS t := f (X t ) -f (X * ).Then, during 1. Phase 1, the dynamics will stop beforet * = 2 S0 µ because S t ≤ 1 4 √ µt -2 √ S 0 2 ; 2. Phase 2 with ∆ := m √ 2σmax + ηµm 2 4σ 2 max : E[S t ] ≤ S 0 e -2µ∆t + η 2 (Lτ -µdq 2 ) 2µ∆ 1 -e -2µ∆t ; 3. Phase 3 with ∆ := 2 π 1 σmax + η π µ σ 2 max : E[S t ] ≤ S 0 e -2µ∆t + η 2 Lτ 2µ∆ 1 -e -2µ∆t .Proof of Lemma C.10. We prove each point by leveraging the shape of the law of X t derived in Lemma C.9:1. Phase 1:EQUATION)meaning that the dynamics will stop before t * = 2 f (X0)-f (X * ) µ ;2. Phase 2: By applying the Itô Lemma to f (X t ) -f (X * ) and thatEQUATIONwe have that if q := max(q 1 , q 2 ),EQUATIONEQUATION3. Phase 3: By applying the Itô Lemma to f (X t ) -f (X * ), we have that:EQUATIONEQUATIONEQUATIONSince f is µ-Strongly Convex, f is also µ-PL. Therefore, we haveEQUATIONTherefore,EQUATIONwhich implies that if k := 2µEQUATIONLemma C.11. Under the assumptions of Lemma 3.5, for any step size scheduler η t such thatEQUATION)Proof of Lemma C.11. For any scheduler η k used inEQUATIONthe SDE of Phase 3 isEQUATIONTherefore, analogously to the calculations in Lemma C.10, we have thatEQUATION)Therefore, using l'Hôpital's rule we have thatEQUATIONLemma C.12. Let H = diag(λ 1 , . . . , λ d ) andM t := e -2 √ 2 π Σ -1 2 H+ η π Σ -1 2 H 2 t . Then, 1. E [X t ] = e - √ 2 π Σ -1 2 Ht X 0 ; 2. V ar [X t ] = M t -e -2 √ 2 π Σ -1 2 Ht X 2 0 + η 2 2 π I d + η π H -1 H -1 Σ 1 2 (I d -M t ).Proof of Lemma C.12. The proof is banal: The expected value derivation leverages the martingale property of the Brownian motion while that of the variance uses the Ito Isomerty.Lemma C.13. Let H = diag(λ 1 , . . . , λ d ). Then, EEQUATION)Proof of Lemma C.13. Since the matrix H is diagonal, we focus on a single component. We apply the Ito Lemma to λi(X i t ) 2 2 :EQUATIONwhich implies thatEQUATIONTherefore,EQUATION)Lemma C.14. Under the assumptions of Corollary C.8, where ∇f γ (x) = ∇f (x) + √ ΣU , we have that the dynamics of SignSGD in Phase 3 is:EQUATION)Proof of lemma C.14. We apply Eq. ( 34) with ν = 2 and linearly approximate Ξ(x) as |x| < 1, In this subsection, we provide our formal derivation of an SDE model for RMSprop. Let us consider the stochastic process L t := (X t , V t ) ∈ R d × R d defined as the solution ofwhere 2Ξ(x) ∼ x √ 2 .EQUATION)EQUATIONwhere β = 1 -ηρ, ρ = O(1), andP t := diag (V t ) 1 2 + ϵI d .Remark C.15. We observe that the term in blue is the only difference w.r.t. the SDE derived in (Malladi et al., 2022) (see Theorem D.2): This is extremely relevant when the gradient size is not negligible. Figure 7 shows the comparison between our SDE, the one derived in (Malladi et al., 2022) , and RMSprop itself: It is clear that even on simple landscapes, our SDE matches the algorithm much better. Importantly, one can observe that the SDE derived in (Malladi et al., 2022) is only slightly worse than ours at the end of the dynamics: As we show in Lemma C.17 The following theorem guarantees that such a process is a 1-order SDE of the discrete-time algorithm of RMSprop EQUATION)v k+1 = βv k + (1 -β) (∇f γ k (x k )) 2 (89) with (x 0 , v 0 ) ∈ R d × R d , η ∈ R >0k := (x k , v k ) ∈ R d × R d , 0 ≤ k ≤ N denotea sequence of RMSprop iterations defined by Eq. (88). Consider the stochastic process L t defined in Eq. ( 86) and fix some test function g ∈ G and suppose that g and its partial derivatives up to order 6 belong to G. Then, under Assumption C.2 and ρ = O(1) there exists a constant C > 0 independent of η such that for all k = 0, 1, . . . , N , we have|Eg (L kη ) -Eg (l k )| ≤ Cη.That is, the SDE (86) is an order 1 weak approximation of the RMSprop iterations (88).Proof. The proof is virtually identical to that of Theorem C.5. Therefore, we only report the key steps necessary to conclude the thesis. First of all, we observe that sinceEQUATION)Then,EQUATIONTherefore, we work with the following algorithm as all the approximations below only carry an additional error of order O(η 2 ), which we can ignore. Therefore, we have thatEQUATION)Therefore, if ∇f γj (x j ) = ∇f (x j ) + Z j (x j ), E[Z j (x j )] = 0, and Cov(Z j (x j )) = Σ(x j )1. E[x k+1 -x k ] = -η diag(v k + ϵI d ) -1 2 ∇f (x k ) ; 2. E[v k -v k-1 ] = ηρ (∇f (x k-1 )) 2 + diag(Σ(x k )) -v k-1 .Then, we have that ifΦ k := ∇f (x k ) √ v k +ϵI d - ∇fγ k (x k ) √ v k +ϵI d 1. E[(x k+1 -x k )(x k+1 -x k ) ⊤ ] = E[(x k+1 -x k )]E[(x k+1 -x k )] ⊤ (94) + η 2 E (Φ k ) (Φ k ) ⊤ (95) = E[(x k+1 -x k )]E[(x k+1 -x k )] ⊤ (96) + η 2 (diag(v k ) + ϵI d ) -1 Σ(x k ); (97) 2. E[(v k -v k-1 )(v k -v k-1 ) ⊤ ] = E[(v k -v k-1 )]E[(v k -v k-1 )] ⊤ + O(ρη 2 ); 3. E[(x k+1 -x k )(v k -v k-1 ) ⊤ ] = E[(x k+1 -x k )]E[(v k -v k-1 ) ⊤ ] + 0.ThereforeEQUATION)EQUATION)Lemma C.17. If (∇f (x)) 2 = O(η), Theorem D.2 is a Corollary of Theorem C.16.Proof. In the proof of Theorem C.16, one drops the term η(∇f (x)) 2 as it is of order η 2 .Corollary C.18. Under the assumptions of Theorem C.16 with Σ(x) = σ 2 I d , η = κη, B = Bδ, and EQUATION)EQUATIONEQUATION)By enforcing that the speed of V t matches that of X t , one needs ρ = κ 2 ρ, which implies β = 1 -κ 2 (1 -β).Proof of Lemma C.19. In order to recover the scaling of β, we enforce that the rate at which V t converges to its limit matches the speed of X t : We need ρ = κ 2 ρ, which recovers the classic scalingEQUATION)EQUATIONTherefore, V t t→∞ → σ 2 Bδ 1, meaning that under these conditions:EQUATIONwhich satisfies the following for µ-strongly convex functionsEQUATIONmeaning that E[f (X t ) -f (X * )] t→∞ ≤ ησLτ 4µ √ B κ √ δ .Since the asymptotic the loss is η Remark: Under the same conditions, SGD satisfiesdX t = -κ∇f (X t )dt + κ 1 √ δ η B σI d dW t (107)and thereforeEQUATIONmeaning that asymptotically the loss is η 2 Lτ σ 2 2µB κ δ which does not depend on κ andδ if κ δ = 1. Lemma C.20. For f (x) := x ⊤ Hx 2 , the stationary distribution of RMSprop is (E[X ∞ ]], Cov(X ∞ )) = 0, η 2 Σ 1 2 H -1 .Proof. As (∇f (x)) 2 = O(η) and t → ∞, we havedX t = -Σ -1 2 HX t dt + √ ηI d dW t (109)which implies thatEQUATIONThe thesis follows from the martingale property of Brownian motion and the Itô isometry. The other two do the same for RMSpropW. In both cases, the first is an MLP on the Breast Cancer Dataset and the second a CNN on MNIST: Our SDEs match the respective optimizers.In this subsection, we derive the SDE of RMSpropW defined as Theorem C.21. Under the same assumptions as Theorem C.16, the SDE of RMSpropW isx k+1 = x k -η ∇f γ k (x k ) √ v k+1 + ϵI d -ηγx k (111) v k+1 = βv k + (1 -β) (∇f γ k (x k )) 2 (112) with (x 0 , v 0 ) ∈ R d ×R d , η ∈ R >0EQUATION)EQUATION)where β = 1 -ηρ, ρ = O(1), γ > 0, andP t := diag (V t ) 1 2 + ϵI d .Proof. The proof is the same as the of Theorem C.16 and the only difference is that ηγx k is approximated with γX t dt. EQUATION)EQUATIONEQUATION)By enforcing that the speed of V t matches that of X t , one needs ρ = κ 2 ρ, which implies β =1 -κ 2 (1 -β).Proof of Lemma C.23. In order to recover the scaling of β, we enforce that the rate at which V t converges to its limit matches the speed of X t : We need ρ = κ 2 ρ, which recovers the classic scalingEQUATION)EQUATIONTherefore, V t t→∞ → σ 2 Bδ 1, meaning that under these conditions:EQUATIONwhich satisfies the following for µ-strongly convex and L-smooth functionsdE[f (X t ) -f (X * )] ≤ κ 2µ √ Bδ σ + ξγ 1 + µ L E[f (X t ) -f (X * )]dt + κ 2 ηL τ 2 dt, (121) meaning that E[f (X t ) -f (X * )] t→∞ ≤ ηLτ σL 2 κ 2µ √ BδL+σξγ(L+µ) .Since the asymptotic the loss ηLτ σL Lemma C.24. For f (x) :=x ⊤ Hx 2 , the stationary distribution of RMSpropW is(E[X ∞ ]], Cov(X ∞ )) = 0, η 2 (HΣ -1 2 + γI d ) -1 .Proof. As (∇f (x)) 2 = O(η) and t → ∞, we havedX t = -Σ -1 2 HX t dt + √ ηI d dW t -γX t dt (122)which implies thatEQUATION)The thesis follows from the martingale property of Brownian motion and the Itô isometry. In this subsection, we provide our formal derivation of an SDE model for Adam. Let us consider the stochastic process LEQUATION)EQUATION)EQUATIONwhere 1), andβ i = 1 -ηρ i , γ i (t) = 1 -e -ρit , ρ 1 = O(η -ζ ) s.t. ζ ∈ (0, 1), ρ 2 = O(P t = diag √ V t + ϵ γ 2 (t)I d .Remark C.25. The terms in purple and in blue are the two differences w.r.t. that of (Malladi et al., 2022) which is reported in Theorem D.5. The first appears because we assume realistic values of β 1while the second appears because we allow the gradient size to be non-negligible. For two simple landscapes, Figure 10 compares our SDE and that of Malladi et al. (2022) with Adam: In both cases, the first part of the dynamics is perfectly represented only by our SDE. While the discrepancy between the SDE of (Malladi et al., 2022) and Adam is asymptotically negligible in the convex setting, we observe that in the nonconvex case, it converges to a different local minimum than ours and of Adam. Finally, Theorem D.5 is a corollary of ours when (∇f (x)) 2 = O(η) and ρ 1 = O(1):It only describes the dynamics where the gradient to noise ratio is vanishing and only for unrealistic values of β 1 = 1 -ηρ 1 . In Figure 11 , we compare the dynamics of our SDE, that of Malladi et al. (2022) , and Adam on an MLP, a CNN, a ResNet, and a Transformer. One can clearly see that our SDE more accurately captures the dynamics. Details on these experiments are available in Appendix F.The following theorem guarantees that such a process is a 1-order SDE of the discrete-time algorithm of Adamv k+1 = β 2 v k + (1 -β 2 ) (∇f γ k (x k )) 2 (127) m k+1 = β 1 m k + (1 -β 1 )∇f γ k (x k ) (128) mk = m k 1 -β k 1 -1 (129) vk = v k 1 -β k 2 -1(130)EQUATIONwith Let(x 0 , m 0 , v 0 ) ∈ R d × R d × R d , η ∈ R >0 is the step size, β i = 1 -ρ i η for ρ 1 = O(η -ζ ) s.t.l k := (x k , m k , v k ) ∈ R d × R d × R d , 0 ≤ k ≤ N denotea sequence of Adam iterations defined by Eq. ( 127). Consider the stochastic process L t defined in Eq. ( 124) and fix some test function g ∈ G and suppose that g and its partial derivatives up to order 6 belong to G.Then, under Assumption C.2 ρ 1 = O(η -ζ ) s.t. ζ ∈ (0, 1), while ρ 2 = O(1), there exists a constant C > 0 independent of η such that for all k = 0, 1, . . . , N , we have|Eg (L kη ) -Eg (l k )| ≤ Cη.That is, the SDE (124) is an order 1 weak approximation of the Adam iterations (127).Proof. The proof is virtually identical to that of Theorem C.5. Therefore, we only report the key steps necessary to conclude the thesis. First of all, we observe that sinceEQUATIONThen,EQUATIONTherefore, we work with the following algorithm as all approximations only carry an additional error of order O(η 2 ), which we can ignore. Therefore, we have thatEQUATIONTherefore, if ∇f γj (x j ) = ∇f (x j ) + Z j (x j ) and E[Z j (x j )] = 0, and Cov(Z j (x j )) = Σ(x j ), we have that1. E[v k -v k-1 ] = ηρ 2 (∇f (x k-1 )) 2 + diag(Σ(x k )) -v k-1 ; 2. E[m k+1 -m k ] = ηρ 1 [∇f (x k ) -m k ] ; 3. E[x k+1 -x k ] = - η √ v k +ϵI d √ 1-(1-ηρ2) k 1-(1-ηρ1) k+1 (m k + ηρ 1 (∇f (x k ) -m k )) . Then, we have 1. E[(x k+1 -x k )(x k+1 -x k ) ⊤ ] = E[(x k+1 -x k )]E[(x k+1 -x k )] ⊤ + O(η 4 ρ 2 1 ); 2. E[(x k+1 -x k )(m k -m k-1 ) ⊤ ] = E[(x k+1 -x k )]E[(m k -m k-1 )] ⊤ + 0; 3. E[(x k+1 -x k )(v k -v k-1 ) ⊤ ] = E[(x k+1 -x k )]E[(v k -v k-1 )] ⊤ + 0; 4. E[(v k -v k-1 )(v k -v k-1 ) ⊤ ] = E[(v k -v k-1 )]E[(v k -v k-1 )] ⊤ + O(η 2 ρ 2 2 ); 5. E[(m k -m k-1 )(m k -m k-1 ) ⊤ ] = E[(m k -m k-1 )]E[(m k -m k-1 )] ⊤ + η 2 ρ 2 1 Σ(x k-1 ); 6. E[(v k -v k-1 )(m k -m k-1 ) ⊤ ] = E[(v k -v k-1 )]E[(m k -m k-1 )] ⊤ + O(η 2 ρ 1 ρ 2 ).Since in real-world applications,EQUATION)EQUATION)EQUATIONwhereβ i = 1 -ηρ i , γ i (t) = 1 -e -ρit, andP t = diag √ V t + ϵ γ 2 (t)I d .Corollary C.27. Under the assumptions of Theorem C.26 withEQUATIONEQUATION)EQUATIONLemma C.28. Under the assumptions of Corollary C.27, f is µ-strongly convex, L τ := Tr(∇ 2 f (x)), and (∇f (x)) 2 = O(η), the asymptotic dynamics of the iterates of Adam satisfies the classic scalingrule κ = √ δ because E[f (X t )] t→∞ ≤ ησLτ 4 √ B κ √ δ .To enforce that the speed of M t and V t match that ofX t , one needs ρi = κ 2 ρ i , which implies βi = 1 -κ 2 (1 -β i ).Proof. First of all, we need to ensure that the relative speeds of X t , M t , and V t match. Therefore, we select α i = κ 2 , which recovers the scaling rules for βi = 1 -κ 2 (1 -β i ). Then, recalling that(∇f (x)) 2 = O(η), we have that as t → ∞, V t → σ 2Bδ , and M t → ∇f (X t ) with high probability.Therefore,EQUATIONEQUATION)dV t = 0. (147) Therefore, if H(X t , V t ) := f (X t ) + Lτ δB ρ 2 σ 2 ∥Mt∥ 2 2 2and ξ ∈ (0, 1) we have that by Itô's lemma,EQUATIONLet us now select ξ such thatEQUATIONTherefore,EQUATIONwhich implies thatEQUATIONAnalogously,EQUATIONwhich gives the square root scaling rule.Lemma C.29. Under the assumptions of Corollary C.27,f (x) = x ⊤ Hx 2 s.t. H = diag(λ 1 , • • • , λ d ) and (∇f (x)) 2 = O(η), the dynamics of Adam implies that f (X t ) → ησd 4 √ B κ √ δ . Proof. Recalling that (∇f (x)) 2 = O(η), we have that as t → ∞, V t → σ 2Bδ , and M t → λX t with high probability. Therefore, in the one-dimensional caseEQUATIONEQUATION)dV t = 0. (159) Therefore, if H(X t , V t ) := λX 2 t 2 + λδB ρ 2 σ 2 M 2 t 2,5 we have that by Itô's lemma,EQUATION)Once again, since M t → λX t , we have thatH(X t , V t ) = λX 2 t 2 + λδB ρ 2 σ 2 M 2 t 2 → λX 2 t 2 + λ λδB ρ 2 σ 2 λX 2 t 2 = 1 + λ λδB ρ 2 σ 2 λX 2 t 2 =: Kf (X t ).(164) Therefore,EQUATIONwhich implies that E[f (X t )] → ησ 4 √ B κ √δ , which also gives the square root scaling rule. The generalization to d dimension is analogous and one needs to sum across all the dimensions.Lemma C.30. Let f (x) := x ⊤ Hx 2 where H = diag(λ 1 , . . . , λ d ). The stationary distribution of Adam is (E[X ∞ ]], Cov(X ∞ )) = 0, η 2 Σ 1 2 H -1 .Proof. The expected value follows immediately from the fact thatEQUATIONFor the covariance, we focus on the one-dimensional case. We define H(X t , V t ) :=X 2 t 2 + λ 2 2σ 2 ρ 2 M 2 t 2 .With the same arguments as Lemma C.29, we haveEQUATIONwhich implies thatEQUATIONThe thesis follows by applying the same logic to multiple dimensions.In this subsection, we derive the SDE of AdamW defined as defined as Theorem C.31. Under the same assumptions as Theorem C.26, the SDE of AdamW isEQUATIONx k+1 = x k -η mk+1 vk+1 + ϵI d -ηγx k (173) with (x 0 , m 0 , v 0 ) ∈ R d × R d × R d , η ∈ R >0 is the step size, β i = 1 -ρ i η for ρ 1 = O(η -ζ ) s.t. ζ ∈ (0, 1), ρ 2 = O(1), γ > 0,EQUATION)EQUATION)whereβ i = 1 -ηρ i , γ > 0, γ i (t) = 1 -e -ρit, andP t = diag √ V t + ϵ γ 2 (t)I d .Proof. The proof is the same as the of Theorem C.26 and the only difference is that ηγx k is approximated with γX t dt.Figure 4 and Figure 9 validate this result on a variety of architectures and datasets.Corollary C.32. Under the assumptions of Theorem C.31 withEQUATIONEQUATION)EQUATIONLemma C.33 (Scaling Rule at Convergence). Under the assumptions of Corollary C.32, f is µstrongly convex and L-smooth, L τ := Tr(∇ 2 f (x)), and (∇f (x)) 2 = O(η), the asymptotic dynamics of the iterates of AdamW satisfies the novel scaling rule if κ = √ δ and ξ = κ becauseEQUATIONBy enforcing that the speed of V t matches that of X t , one needs ρ = κ 2 ρ, which implies βi =1 -κ 2 (1 -β i ).Proof. The proof is the same as Lemma C.28 where we also use L-smoothness as in Lemma C.23.Lemma C.34. For f (x) := x ⊤ Hx 2 , the stationary distribution of AdamW is (E[X ∞ ]], Cov(X ∞ )) = 0, η 2 (HΣ -1 2 + γI d ) -1 .Proof. The proof is the same as Lemma C.30.Theorem D.1 (Original Malladi's Statement). Let σ 0 := ση, ϵ 0 := ϵη, and c 2 := 1-β η 2 . Define the state of the SDE as L t = (X t , u t ) and the dynamics asEQUATION)EQUATION)whereP t := σ 0 diag (u t ) 1/2 + ϵ 0 I d .Theorem D.2 (Informal Statement of Theorem C.2 Malladi et al. (2022) ). Under sufficient regularity conditions and ∇f (x) = O( √ η), the following SDE is an order 1 weak approximation of RMSprop:EQUATIONwhere β = 1 -ηρ, ρ = O(1), and P t := diag (V t )1 2 + ϵI d .Lemma D.3. Theorem D.1 and Theorem D.2 are equivalent.Proof. It follows applying time rescaling t := ηξ and observing that W t = W ηξ = √ ηW ξ .Theorem D.4 (Original Malladi's Statement). Let c 1 := (1 -β 1 ) /η 2 , c 2 := (1 -β 2 ) /η 2 and define σ 0 , ϵ 0 in Theorem D.1. Let γ 1 (t) := 1 -exp (-c 1 t) and γ 2 (t) := 1 -exp (-c 2 t). Define the state of the SDE as L t = (X t , m t , u t ) and the dynamics asEQUATION)EQUATIONEQUATIONwhereP t := σ 0 diag (u t ) 1/2 + ϵ 0 γ 2 (t)I d .Theorem D.5 (Informal Statement of Theorem D.2 Malladi et al. (2022) ). Under sufficient regularity conditions and ∇f (x) = O( √ η), the following SDE is an order 1 weak approximation of Adam:EQUATION)EQUATION)whereβ i = 1 -ηρ i , γ i (t) = 1 -e -ρit , ρ i = O(1), andP t = diag √ V t + ϵ γ 2 (t)I d .Lemma D.6. Theorem D.4 and Theorem D.5 are equivalent.Proof. It follows applying time rescaling t := ηξ and observing that W t = W ηξ = √ ηW ξ .In this section, we provide a gentle introduction to the meaning of deriving an SDE model for an optimizer and discuss how SDEs have been used to derive scaling rules. To aid the intuition of the reader, we informally derive an SDE for SGD with learning rate η, mini-batches γ B of size B, and starting point x 0 = x, which we dub SGD (η,B) . The iterates are given by:x k+1 = x k -η∇f γ B k (x k ) (191) which for U k := √ η(∇f (x k ) -∇f γ B k (x k)), we rewrite asEQUATIONwhereE[U k ] = 0 and Cov(U k ) = η B Σ(x k ) = η B 1 n B i=0 (∇f (x k )-∇f i (x k ))(∇f (x)-∇f i (x k )) ⊤ .If we now consider the SDEEQUATIONits Euler-Maruyama discretization with pace ∆t = η and Z k ∼ N (0, I d ) isEQUATION)Since the Eq. ( 191) and Eq. ( 194) share the first two moments, it is reasonable that by identifying t = kη, the SDE in Eq. ( 193) is a good model to describe the iterates of SGD in Eq. ( 191).Informally, we need a "good model", which is an SDE that is close to the real optimizer. This is formalized in the following definition which comes from the field of numerical analysis of SDEs (see Mil'shtein (1986) ) and bounds the disparity between the the discrete and the continuous process. To see if an SDE satisfies such a definition, one has to check that for ∆ = x 1 -x and ∆ = X η -x,1. E∆ i -E ∆i = O(η 2 ), ∀i = 1, . . . , d;2. E∆ i ∆ j -E ∆i ∆j = O(η 2 ), ∀i, j = 1, . . . , d.Example: Let us prove that the SDE in Eq. ( 193) is a valid approximation of SGD (η,B) : The first condition is easily verified. Coming to the second condition we have that1. E∆ i ∆ j = η 2 ∂ i f (x)∂ j f (x) + η 2 B Σ(x); 2. E ∆i ∆j = η 2 ∂ i f (x)∂ j f (x) + η 2 B Σ(x) + O(η 3 );whose difference is of order η 3 and thus satisfies the condition. However, we observe that if the scale of the noise is too small w.r.t η, i.e. Σ(x) = O(η α ) for α ≥ 0, then the simplest SDE model describing SGD (η,B) is the ODE dX t = -∇f (X t )dt as in that case1. E∆ i ∆ j = η 2 ∂ i f (x)∂ j f (x) + O(η 2+α ); 2. E ∆i ∆j = η 2 ∂ i f (x)∂ j f (x) + O(η 2 ),whose difference is also of order η 2 . Much differently, if Σ(x) = O(η -α ) for α > 0, the simplest model is the SDE in Eq. ( 193). We highlight that simplest does not mean best: The SDE is more accurate than the ODE even in a regime with low noise, but this observation serves as a provocation.One has to pay attention when deriving SDEs: Some models are more realistic than others.Let us dig deeper into this thought as we derive two SDEs for SGD with learning rate η := κη and batch size B := δB for κ > 1 and δ > 1, which we dub SGD (η, B) . The first is derived considering that the learning rate is η and carries an error of order O(η) w.r.t. SGD (η, B)EQUATIONThe second one instead is derived considering η as the learning rate and κ as a constant "scheduler".Consistently with (Li et al., 2017) , the SDE which carries an error of order O(η) w.r.t SGD (η, B) isEQUATIONWhile they both are valid models, there are three reasons why one should prefer the latter:1. It fully reflects the fact that a larger learning rate results in a faster and noisier dynamics 2. It has intrinsically less error than the other;3. It is consistent with the optimizer in that there is no combination of κ and δ that can ever leave the dynamics unchanged.E.1 Deriving scaling rules Jastrzebski et al. (2018) observed that only the ratio between η and B matters in determining the dynamics of Eq. ( 194). Therefore, they argue that for κ = δ the SDE for SGD (κη,δB) coincides with that of SGD (η,B) and that this implies that the path properties of the optimizers are the same. On the contrary, the path of SGD (η,B) strongly depends on the hyperparameters: The speed and volatility of the dynamics are driven by η, and no choice of B can undo this. We remind the reader that the goal of these rules is not to keep the dynamics of the optimizers unaltered, but rather to give a practical way to change a hyperparameter, e.g. η, and have a principled way to adjust the others, e.g. B, such that the performance of the optimizer is preserved. Therefore, we propose deriving scaling rules as we preserve certain relevant quantities of the dynamics such as the convergence bound on the expected loss or the speed. To show this quantitative, we use this rationale to derive the scaling rule of SGD as we aim at preserving the asymptotic loss level.Lemma E.2. If f is a µ strongly convex function, L τ ≤ Tr(∇ 2 f (x)) and Σ(x) = σ 2 I d , then:1. Under the dynamics of Eq. (193) we have: (197) 2. Under the dynamics of Eq. (195) we have: (198) 3. Under the dynamics of Eq. (196) we have:E[f (X t ) -f (X * )] ≤ (f (X 0 ) -f (X * ))e -2µt + η 2 L τ σ 2 2µB 1 -e -2µt ;E[f (X t ) -f (X * )] ≤ (f (X 0 ) -f (X * ))e -2µt + η 2 L τ σ 2 2µB κ δ 1 -e -2µt ;EQUATIONThe first bound implies that the asymptotic limit of the expected loss for SGD (η,B) isη 2 Lτ σ 22µB . The last two bounds predict that the asymptotic loss level for SGD (η, B) isη 2 Lτ σ 2 2µB κ δ .Since the objective of the scaling rule is to find κ and δ such that SGD (η, B) achieves the same loss level as SGD (η,B) , we recover the linear scaling rule setting κ = δ. However, only the last bound can correctly capture the fact that the dynamics of SGD (η, B) is κ times faster than that of SGD (η,B) .We conclude the discussion with a simple sample of how deriving a scaling rule from the SDE itself inevitably leads to the wrong conclusion. We define the following algorithm which is inspired byAdamW and which we dub SGDW:EQUATIONTherefore, one would naively deduce that to keep the SDE unchanged, one can simply use the linear scaling rule of SGD and leave γ unaltered. However, one can easily derive the upper bound on the expected loss for a convex quadratic function and observe that to preserve that, it is imperative to scale γ by κ as well.We thus conclude that:1. Eq. ( 196) is a better model for SGD (η, B) as it represents the dynamics more accurately;2. Maintaining the shape of the SDE does not preserve the path properties of the optimizer;3. Deriving a scaling rule uniquely from the SDE might lead to the wrong conclusions in the general case.Remark E.4. We highlight that Theorem 5.3 of Malladi et al. (2022) claimed to have formally derived one for RMSprop: In line with (Jastrzebski et al., 2018) , they argue that if they were to find a scaling rule that would leave their SDE unchanged, this would imply that even the dynamics of the iterates of RMSprop itself would be unchanged. First, we remind the reader that an SDE is formally defined as an equation that drives the dynamics plus an initial condition (See (Karatzas and Shreve, 2014), Section 5). While their scaling rule does leave the equation unchanged, it alters the initial condition, thus changing the SDE itself: This invalidates their claim and proof. Second, contrary to their claim, the rule is only valid near convergence as their SDE is only valid there. Third, Lemma E.2 offers a shred of concrete evidence that keeping the SDE unchanged does not imply that the path properties of the optimizers are preserved. Fourth, Lemma E.3 is a piece of concrete evidence that deriving scaling rules directly and naively from the SDE might lead to the wrong conclusions.In this section, we provide the modeling choices and instructions to replicate our experiments. All experiments we run on one NVIDIA GeForce RTX 3090 GPU. The code is implemented in Python 3 (Van Rossum and Drake, 2009) mainly using Numpy (Harris et al., 2020) , scikit-learn (Pedregosa et al., 2011) , and JAX (Bradbury et al., 2018) .F.1 SignSGD: SDE validation (Figure 1 )In this subsection, we describe the experiments we run to produce Figure 1 : The loss dynamics of SignSGD and that of our SDE match on average.DNN on Breast Cancer Dataset (Dua and Graff, 2017) This paragraph refers to the left of Figure 1 . The DNN has 10 dense layers with 20 neurons each activated with a ReLu. We minimize the binary cross-entropy loss. We run SignSGD for 50000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ∼ N (0, σ 2 I d ) where σ = 1. The learning rate is η = 0.001. Similarly, we integrate the SignSGD SDE (Eq. ( 7)) with Euler-Maruyama (Algorithm 1) with ∆t = η. Results are averaged over 3 runs and the shaded areas are the average ± the standard deviation.CNN on MNIST (Deng, 2012) This paragraph refers to the center-left of Figure 1 . The CNN has a (3, 3, 32) convolutional layer with stride 1, followed by a ReLu activation, a (2, 2) max pool layer with stride (2, 2), a (3, 3, 32) convolutional layer with stride 1, a ReLu activation, a (2, 2) max pool layer with stride (2, 2). Then the activations are flattened and passed through a dense layer that compresses them into 128 dimensions, a final ReLu activation, and a final dense layer into the output dimension 10. The output finally goes through a softmax as we minimize the cross-entropy loss. We run SignSGD for 40000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ∼ N (0, σ 2 I d ) where σ = 1. The learning rate is η = 0.001. Similarly, we integrate the SignSGD SDE (Eq. ( 7)) with Euler-Maruyama (Algorithm 1) with ∆t = η. Results are averaged over 3 run and the shaded areas are the average ± the standard deviation.Transformer on MNIST This paragraph refers to the center-right of Figure 1 . The Architecture is a scaled-down version of (Dosovitskiy et al., 2021) , where the hyperparameters are patch size=28, out features=10, width=48, depth=3, num heads=6, and dim ffn=192. We minimize the cross-entropy loss as we run SignSGD for 5000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ∼ N (0, σ 2 I d ) where σ = 1. The learning rate is η = 0.001. Similarly, we integrate the SignSGD SDE (Eq. ( 7)) with Euler-Maruyama (Algorithm 1) with ∆t = η. Results are averaged over 3 runs and the shaded areas are the average ± the standard deviation.ResNet on CIFAR-10 ( Krizhevsky et al., 2009) This paragraph refers to the right of Figure 1 .The ResNet has a (3, 3, 128) convolutional layer with stride 1, followed by a ReLu activation, a second (3, 3, 64) convolutional layer with stride 1, followed by a residual connection from the first convolutional layer, then a (2, 2) max pool layer with stride (2, 2). Then the activations are flattened and passed through a dense layer that compresses them into 128 dimensions, a final ReLu activation, and a final dense layer into the output dimension 10. The output finally goes through a softmax as we minimize the cross-entropy loss. We run SignSGD for 5000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ∼ N (0, σ 2 I d ) where σ = 1. The learning rate is η = 0.001.Similarly, we integrate the SignSGD SDE (Eq. ( 7)) with Euler-Maruyama (Algorithm 1) with ∆t = η.Results are averaged over 3 runs and the shaded areas are the average ± the standard deviation.F.2 SignSGD: insights validation (Figure 2 )In this subsection, we describe the experiments we run to produce Figure 2 : We successfully validate them all.Phases: Lemma 3.4 and Lemma 3.5 In this paragraph, we describe how we validated the existence of the phases of SignSGD as predicted in Lemma 3.4 and Lemma 3.5. To produce the left of Figure 2 ), we simulated the full SDE (Eq. ( 16)) and the one describing Phase 3 (Eq. ( 5)). The optimized function is f (x) = x ⊤ Hx 2 for H = diag(1, 2), x 0 drawn (and fixed for all runs) from a normal distribution N (0, 0.01), η = 0.001, and Σ = σ 2 I d where σ = 0.1. We integrate the SDEs with Euler-Maruyama (Algorithm 1) with ∆t = η and for 3000 iterations. Results are averaged over 500 runs and the shaded areas are the average ± the standard deviation. Clearly, the two SDEs share the same dynamics.To produce the center-left of Figure 2 , we repeat the above as x 0 drawn (and fixed for all runs) from a normal distribution N (0, 1). Then, we plot the average loss values together with the theoretical prediction of Phase 1 and Phase 3: They perfectly overlap.F.5 RMSpropW & AdamW: SDE validation (Figure 3 , Figure 4 )The settings are exactly the same as those for RMSprop and Adam. The regularization parameter used is always γ = 0.01. We observe that our SDEs match the respective algorithm with a good agreement.F.6 RMSpropW & AdamW: insights validation (Figure 5 )In this subsection, we describe the experiments we run to produce where H = diag(1, 3). We run Adam for 20000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ∼ N (0, σ 2 I d ) where σ = 0.01. The learning rate is η = 0.001, γ = 4, β = 0.999, β 1 = 0.9, and β 2 = 0.999. We plot the evolution of the average variances with the theoretical predictions of Lemma C.24 and Lemma 3.14: Results are averaged over 100 runs.F.7 Effect of noise -validation (Figure 6 )In this subsection, we describe the experiments run to produce Figure 6 : All bounds on the asymptotic expected loss value for SGD, SignSGD, Adam, and AdamW are perfectly verified.We optimize the loss f (x) = x ⊤ Hx 2 where H = diag(1, 1) as we run each optimizer for 100000 iterations with η = 0.01. We repeat this procedure five times, one for each σ ∈ {0.01, 0.1, 1, 10, 100}.As we train, we inject noise on the gradient as distributed as N (0, σ 2 I d ). We plot the average loss together with the respective limits predicted by our Lemmas. For each optimizer and each σ, the average asymptotic loss matches the predicted limit. Therefore, we verify that the loss of SGD scales quadratically in σ, that of Adam and SignSGD scales linearly, and that of AdamW is limited in σ.The analysis of Malladi et al. (2022) suggests that, when scaling batch size B by a factor κ one has to scale up (↑) the learning rate η by a factor √ κ and scale down (↓) β 2 to the value 1 -κ(1 -β 2 ).Our SDE analysis confirms similar rules (Lemma 3.13) but additionally suggests scaling up the decoupled weight decay parameter γ by a factor √ κ. We test this in two settings: VGG11 and ResNet34 (convolutional networks) on CIFAR-10 classification. We select a base batch size of 256, and run AdamW with η = 0.001, β 2 = 0.99, and γ = 0.1. We consider scaling the batch by a factor 4: In Table 1 , we show the effect of updating each hyperparameter with the proposed rule and we denote by a "•" the model parameters of the base run with B = 256. We train for 150 epochs the model with B = 256, and 150 × 4 the model with B = 4 × 256. Experiments are repeated 3 times.We find that, while improvements are marginal, they are consistent with our theoretical results. 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.Ths SDE of SGD is dXt = -∇f (Xt)dt + √ ηΣ 1 2 dWt.Table 1 in Appendix F.8 shows that our scaling rule works on DNNs: it confirms that failing to rescale the weight decay parameter is suboptimal.Many papers derived SDEs to model optimizers: most of them do not validate them, some do so on quadratic functions, andPaquette et al. (2021);Compagnoni et al. (2023) do it on NNs: See Appendix A for details.Inspired by(Barakat and Bianchi, 2021) 


Decision: This paper introduces novel Stochastic Differential Equations (SDEs) for adaptive optimization methods such as SignSGD, RMSProp(W), and Adam(W). The authors provide a detailed analysis of these methods, focusing on their dynamics, convergence behavior, and robustness to noise. The theoretical findings are supported by experiments on various neural network architectures, confirming the accuracy of the derived SDEs in modeling the behavior of these optimizers, which are appreciated by the reviewers.

However, I have concerns about the correctness of one of the main theoretical result in this paper, Theorem 3.2 (or its formal version, Theorem C.5 in appendix). Theorem C.5 mainly says that SignSGD (equation 18) and SDE (equation 16) are order-1 weak approximation. However, for the setting studied in this paper, i.e., batches are sampled i.i.d. from the uniform distribution of a training set of size $N$, both the drift term and diffusion matrix in equation 16 are not continuous, because $\mathrm{Pr}(\nabla f_\gamma(x)<0)$ only takes finite set of discrete values of {$0,\frac{1}{N},\frac{2}{N},\ldots,1$}. Therefore even the uniqueness and existence of solution of (16) is not obvious and the proof framework introduced by [Li et al., 17] does not apply here. This mistake is also reflected in the proof of Theorem 3.5, which applies Lemma 3.3 without satisfies its condition (drift $b$ and diffusion matrix $\Sigma$ needs to be lipschitz continuous).

With that being said, this issue could probably be fixed by only considering infinite size training dataset, or equivalently, noise distribution with continuous/smooth densities, like gaussian distribution. However, even in that case, the proof framework of [Li et al., 17] (which is further refined in [Li et al., 19]) does not apply for the SDE approximation of SignSGD directly. Therefore I think the current version is not ready for being published and the amount of modification is significant, which needs another round of review. I encourage the authors to fix the statement and the proof and resubmit to another venue.

- Li, Q., Tai, C. and Weinan, E., 2017, July. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning (pp. 2101-2110). PMLR.

- Li, Q., Tai, C. and Weinan, E., 2019. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40), pp.1-47.

Reviews: review: paper_summary: paper_summary: This paper derives SDEs for SignSGD, RMSprop, and Adam.
The analysis offers insights into the convergence speed, stationary distribution, and robustness to heavy-tail noise of adaptive methods.
strength_weakness: strength_weakness: Refer to Questions and Limitations.
questions: questions: - What learning rate (lr) do the experiments in Figure 4 use? Within what range of lr does this SDE align well with the original algorithm (experimentally)?

- Could the authors intuitively explain why the asymptotic expected loss of SignSGD is proportional to $\sigma_{\max}$ instead of $\sigma_{\max}^2$?

- How can the derived SDE explain the loss spike phenomenon of SignSGD/AdamW?

- Many works about SGD noise ([1][2][3]) admit the noise structure $\mathbb{E}(g_i-g)(g_i-g)^{\top}\sim \mathcal{L} H$. What conclusions (such as those related to the training phases) can be derived from the SDE if we change the noise assumption in Corollary 3.3 to $\mathbb{E}(g_i-g)(g_i-g)^{\top}\sim\mathcal{L} H$?


[1] Ziyin et al. Strength of minibatch noise in SGD.

[2] Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part II: Continuous time analysis.

[3] Wu et al. The alignment property of SGD noise and how it helps select flat minima: A stability analysis.
limitations: limitations: The SDE for AdamW is limited to quadratic functions.
paper_summary: paper_summary: The authors derive SDE for signSGD and Adam(W). The experiments show that the algorithm will converge toward the limit of the theorem indicates.
strength_weakness: strength_weakness: 1. In Remark after Lemma 3.6, the authors claim that Sign-SGD is (almost) linear in $\sigma_{max}$. However, with $\Delta$ either in Phase 2 or Phase 3, there should be $\sigma_{max}^2$ in the final bound.

2. All the stationarity holds when Hessian is the same from $X_0$ to $X_t$ and convergence holds for strongly convex. However, the hessian changes a lot during network training.
questions: questions: 1. Notation of $W_t$ is not defined.  What is $W_t$?

2. How can we extend Lemma 3.13 to convex setting (or even nonconvex case)?
limitations: limitations: N/A
paper_summary: paper_summary: This work derives SDEs for adaptive gradient methods and study the role of gradient noise. The analysis starts from theoretically driving the SDE for SignSGD and highlight its significant difference from SGD. The work further generalize the SDE analysis for AdamW and RMSpropW, two popular adaptive optimizers with decoupled weight decay and reveal key properties of weight decay. Finally, the work integrates the derived SDEs with Euler-Maruyama to confirm that the SDEs faithfully track their respective optimizers with various modern neural networks.
strength_weakness: strength_weakness: -It seems that the reported theoretical results and insights cannot directly lead to some theory-inspired and improved methods. This raise a question on the significance of this work.

-While this work did literature review, some important references are still missing, such as [1] on analyzing Adam using SDEs. As weight decay plays a key role in the results, it may be helpful to review recent papers analyzing novel or overlooked properties of weight decay.


Reference:

[1] Xie, Z., Wang, X., Zhang, H., Sato, I., & Sugiyama, M. (2022, June). Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International conference on machine learning (pp. 24430-24459). PMLR.
questions: questions: - Please see the weaknesses.

- Could you please explain more how L2 regularization and decoupled weight decay behaves differently in your results?
limitations: limitations: This work discussed the limitations in the appendix.

Comments: title: Thanks!
Thanks for the rebuttal.
Thanks!
Reviewer's curiosity: Noise Structure and Conjecture on Spiking Phenomena
comment: Dear Reviewer,

Thank you for your trust and the updated score: We truly appreciate it.

Best regards,

The Authors
Thanks for the authors' response. I have no further questions. Since the authors claim to add clarification in the final version, I raise my score to 6.
Dear Reviewer,

Thank you for your response.

We are glad to know that some of your concerns have been resolved: Could you please share any remaining issues or suggestions you might have? Your feedback is invaluable and will assist us in refining our manuscript further.

We appreciate your time and consideration.

Best regards,

The Authors
Thanks for the rebuttal and addressing some of the concerns.

I will keep the rating as 6. I tend to accept this work.
Dear Reviewer,

Thank you for your kind words and the updated score: We truly appreciate it.

Best regards,

The Authors
Many thanks to the authors for your careful explanation and detailed rebuttal. I feel that this paper is of great help in understanding signSGD/Adam. I have raised my score.
**_Given the length limit for the Rebuttal, we decided to include these minor results in an Official Comment._**

**Continuation of Answer to Q4 - The Third Noise Structure:**

[2] discusses two possible assumptions on $\Sigma$: $\|\Sigma(x)\| \leq C f(x)$ and $\|\Sigma(x)\| \leq C f(x)\left[1+|x|^2\right]$. Even though none was in line with the prescription of the Reviewer, we still thought that the one they used, i.e. $\Sigma = C f(x) I_d$ as per Section 2.4, is interesting. Therefore, we take $\Sigma := \sigma^2 f(x) I_d$, where we changed the constant to $\sigma^2$ to maintain consistency with the rest of our paper.

Under this assumption, we have that for $Y_t := \frac{\nabla f(X_t)}{\sqrt{2  f(x)} \sigma}$, Corollary 3.3 becomes:

\begin{align}
d X_t = - Erf \left( Y_t \right) dt + \sqrt{\eta} \sqrt{I_d - diag(Erf \left(Y_t \right))^2} d W_t.
\end{align}

As a consequence, Lemma 3.5 becomes:

Let $f$ be $\mu$-strongly convex, $S_t:=f(X_t) - f(x_*)$, and $Tr(\nabla^2 f(x)) \leq \mathcal{L}_{\tau}$. Then, during
1. Phase 1, the loss will reach $0$ before $t_* = 2 \sqrt{\frac{S_0}{\mu}}$ because $S_t \leq \frac{1}{4} \left( \sqrt{\mu}t - 2 \sqrt{S_0}\right)^2$;
2. Phase 2 with $\beta := \frac{\eta}{2} \left( \mathcal{L}_{\tau} - \mu d \hat{q}^2 - \frac{m^2 \mu^2}{\sigma^2}\right)$ and $\alpha:=  \frac{\sqrt{2} m \mu}{\sigma}$, 
\begin{equation}
\mathbb{E}[S_t] \leq \frac{\beta^2 \left( \mathcal{W}\left( \frac{(\beta + \sqrt{S_0} \alpha)}{\beta} \exp\left(-\frac{\alpha^2 t - 2 \sqrt{S_0} \alpha}{2 \beta} - 1 \right) \right) + 1 \right)^2}{\alpha^2} \overset{t \rightarrow \infty}{\rightarrow} \frac{\beta^2}{\alpha^2},\end{equation}
where $\mathcal{W}$ is the Lambert $\mathcal{W}$ function.
3. Phase 3 it is the same as Phase 2 but with $\beta := \eta \left( \frac{\mathcal{L}_{\tau}}{2} - \frac{2 \mu^2}{\pi \sigma^2}\right)$ and $\alpha:= 2 \sqrt{\frac{2}{\pi}} \frac{\mu}{\sigma}$.

_Please, find an empirical validation of these bounds in Figure 1.c of the attached .pdf file._

**Continuation of Answer to Q3 - Reviewer's curiosity: Conjecture on Spiking Phenomena:**

This is a very interesting question that we do not address in this paper. While this is not a fundamental element for the flow and contribution of our paper, we gladly try to answer it, both for our and the Reviewer's curiosity.

Although we can not answer this in the general case, we offer the following conjecture to provide an intuition of how one could explain the spiking behavior of the mentioned optimizers.

Since the SDE of RMSprop is less complex and less complicated to work with, we restrict ourselves to this case: Generalizing is only a matter of technicalities.

Let us remind that the SDE of RMSprop is 
$$
d X_t = - V_t^{-\frac{1}{2}} (\nabla f(X_t) dt + \sqrt{\eta} \Sigma(X_t)^{\frac{1}{2}} d W_t)
$$
$$
d V_t = \rho( (\nabla f(X_t))^2 +  diag(\Sigma(X_t)) - V_t)) dt,
$$
where $\beta = 1 - \eta \rho$.

Intuitively, the dynamics of the parameters $X_t$ is a preconditioned version of SGD. Much differently, $V_t$ is a process that tracks the squared gradient and its noise.

This implies that the expected iterates follow the dynamics:
\begin{equation}
    d \mathbb{E}[X_t] = -  \mathbb{E}\left[\frac{\nabla f(X_t)}{\sqrt{V_t}}\right] dt.
\end{equation}
Consistently with the noise structure proposed by the Reviewer and used in many papers (see [3] and references therein), let us assume that the covariance of the noise $\Sigma(x)$ scales proportionally to the loss function, e.g. $\Sigma \sim f(x)$.

Spikes seem to happen when the loss is essentially $0$, meaning that $\nabla f(X_t) \sim 0$, $\Sigma \sim 0$, and $V_t \sim 0$. However, if we now draw a minibatch of data for which the gradient is not $0$, e.g. some data points that are outliers, $V_t$ might not have the time to "catch up" with this anomaly. Therefore, the numerator $\nabla f(X_t)$ is non-$0$ while the denominator $\sqrt{V_t}$ is still essentially $0$, meaning that the ratio $\frac{\nabla f(X_t)}{\sqrt{V_t}}$ spikes to infinity, drastically disturbing the dynamics of the iterates and in turn that of the loss function which might spike.