Adaptive Methods through the Lens of SDEs: Theoretical Insights on the Role of Noise
Authors: Enea Monzio Compagnoni
Tianlin Liu
Rustem Islamov
Frank Norbert Proske
Antonio Orvieto
Aurelien Lucchi
Venue: ICLR 2025 Poster (2025)
Keywords: Stochastic Differential Equations
Stochastic Optimization
Adaptive Methods
Primary area: optimization
TL;DR: We derive novel SDEs for SignSGD, RMSprop(W), and Adam(W), providing a more accurate theoretical and understanding of their dynamics, convergence, and robustness. We validate our findings with experiments on various neural network architectures.

Abstract: Despite the vast empirical evidence supporting the efficacy of adaptive optimization methods in deep learning, their theoretical understanding is far from complete. This work introduces novel SDEs for commonly used adaptive optimizers: SignSGD, RMSprop(W), and Adam(W). These SDEs offer a quantitatively accurate description of these optimizers and help illuminate 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.

Published as a conference paper at ICLR 2025
ADAPTIVE METHODS THROUGH THE LENS OF SDES:
THEORETICAL INSIGHTS ON THE ROLE OF NOISE
Enea Monzio Compagnoni', Tianlin Liu', Rustem Islamov', Frank Norbert Proske", Antonio Orvieto 9.1.5, Aurelien I Lucchi UUniversity of Basel, Switzerland; Uniiveriity of Oslo, Norway; Max Planck Institute for Intelligent Systems, Germany
'ELLIS Institute Tübingen, Germany
Tübingen Al Center, Germany (enea. monzi coopmpagnnoaaannono.....o tiarlin. liv, rustem. rustem.islamov/ @unibas.ddbbb...
proske backee@mahh.uuke uio. no, antonio@tue.ellis.. aurelien. lucchi @unibas ch
ABSTRACT
Despite the vast empirical evidence supporting the efficacy of adaptive optimization methods in deep learning, their theoretical understanding is far from complete. This work introduces novel SDEs for commonly used adaptive optimizers: SignSGD, RMSprop(W) (W), and AW). These SDEs offer a quantitatively accurate descrip- tion of these optimizers and help illuminate intricate relationship between adaptivity, gradient noise, and curvature. Our novel analysis of SignSGD high- lights a noteworthy and precise contrast to SGD terms of convergence speed, stationary distribution, and robustness to heavy-tail noise. We extend this analysis to AdamW W and RMSpropW, for which we observe that the role of noise 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 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.
INTRODUCTION
Adaptive optimizers lay the foundation for effective training of modern deep learning models. These methods are typically employed to optimize an objective function expressed as sum of losses across N individual data points: min (fx) + Σ f((z)], where f. fi _d + IR, i= 1,..., N. Due to the practical difficulties of selecting the learning rate of stochastic gradient descent, adaptive methods have grown in popularity over the past decade. At high level, these optimizers adjust the learning rate for each parameter based on the historical gradients. Popular optimizers that belong to this family R.Sprrp (Tieleman and Hinton, 2012), Adam (Kingma and Ba, 2015), SignSGI (Bernstein et al.. 2018), AdamW (Loshchilov and Hutter, 2019), and many other variants. SignSGD is often used for compressing gradients in distributed machine learning (Karimireddy al., 2019a), but it also has gained popularity due to its connection to RMSprop and Adam (Balles and Hennig, 2018). The latter algorithms have emerged as the standard methods for training modern large language models, partly because of enhancements in signal propagation (Noci et al., 2022).
Although adaptive methods are widely favored in practice, their theoretical foundations remain enig- matic. Recent research has illuminated some of their advantages: Zhang et al. (2020b) demonstrated how gradient clipping addresses heavy-tailed gradient noise, Pan and (2022) related the success of Adam over SGD sharpness, and Yang et al. (2024) showed that adaptive methods more resilient to poor learning rate tuning than SGD. At the same time, many optimization studies focus on worst-case convergence rates: These rates (e.g., Défossez et al. (2022)) are valuable, yet they provide an incomplete depiction of algorithm behavior, showing no quantifiable advantage over standard SGD. One particular aspect still lacking clarity is the precise role of noise in the algorithm trajectory.
Our investigation aims to study how gradient noise influences the dynamics of adaptive optimizers and how it impacts their asymptotic behaviors in terms of expected loss and stationary distribution. In

---

Published as a conference paper at ICLR 2025
particular, we want to understand which algorithms more resilient to high (possibly heavy-tailed) gradient noise levels. To do this, we rely on stochastic differential equations (SDEs) which have become popular in the literature to study the behavior of optimization algorithms (Li et al., 2017; Jastrzebski et al., 2018). These continuous-time models unlock powerful tools from Itô calculus, enabling to establish convergence bounds, determine stationary distributions, unveil implicit regularization, and elucidate the intricate interplay between landscape and noise. Notably, SDEs facilitate direct comparisons between optimizers by explicitly illustrating how each hyperparameter and certain landscape features influence their dynamics (Orvieto and Lucchi, 2019; Malladi et al., 2022: Compagnoni et al., 2023; 2024; 2025).
We begin by analyzing SignSGD, showing how the ratio between the gradient and the level of gradient noise affects its dynamies and elucidating the impact of noise at convergence. After examining the case where the gradient noise has an infinite variance, we extend our analysis to Adam and RMSprop with decoupled weight decay (Loshchilov and Hutter, 2019) - i.e. AdamW and RMSpropW: for both, we refine batch size scaling rules and compare the role of noise to SignSGD. Our analysis provides some theoretical grounding for the resilience of these adaptive methods to high noise levels. Importantly, we highlight that Adam and RMSprop byproducts of our analysis and that novel SDEs are derived under weaker assumptions than those in the literature (Zhou et al., 2020a; Malladi al., 2022).
Contributions
We identify our key contributions as follows:
1. We derive the first' SDE for SignSGD under very general assumptions: We show that SignSGD exhibits three different phases of the dynamics and characterize the loss behavior in these phases, including the stationary distribution and asymptotic loss value;
2. We prove that for SignSGD, noise inversely affects the convergence rate of both the loss and the iterates. Differently, it has a linear impact on the asymptotic expected loss and the asymptotic variance of the iterates. This is in contrast to SGD, where noise does not influence the convergence speed, but it has a quadratic effect on the loss and variance of the iterates. Finally, we show that, even if the noise has infinite variance, SignSGD is resilient: its performance is only marginally impacted. In the same conditions, SGD diverges;
3. We derive new, improved, SDEs for A and RMSpropW and use them to (i) show novel batch size scaling rule and (ii) inspect the stationary distribution and stationary loss value in convex quadratics. In particular, we dive into the properties of weight decay: while for vanilla Adam and RMSprop the effect of noise convergence mimics SignSGD, something different happens in Adam W and RMSpropW - Due to an intricate interac- tion between noise, curvature, and regularization, decoupled weight decay plays a crucial stabilization role at high noise levels near the minimizer;
4. We empirically verify every theoretical insight we derive. Importantly, we integrate our SDEs with Euler-Maruyama to confirm that our SDEs faithfully track their respective optimizers. We do so on MLP, CNN, a ResNet, and a Transformer. For RMSprop and Adam, our SDEs exhibit superior modeling power than the SDEs already in the literature. We emphasize that while our results rely certain regularity assumptions for loss functions and gradient noise, their applicability extends beyond these. For example, we validate our novel scaling rule for AdamW Pythia-like 160M LLM (Biderman al., 2023) trained on B/10 tokens from the SlimPajama dataset (Soboleva et al., 2023).
RELATED WORK
SDE approximations and applications. (Li et al., 2017) introduced a formal theoretical framework aimed at deriving SDEs that effectively model the inherent stochastic nature of optimizers. Ever since, SDEs have found several applications in the field of machine learning, for instance in connection with stochastic optimal control to select the stepsize (Li et al., 2017: 2019) and batch size (Zhao et al., 2022), the derivation of convergence bounds and stationary distributions (Compagnoni et al., 2023: 2024), implicit regularization (Smith et al., 2021), and scaling rules (Jastrzebski et al.. 2018). Previous work by Malladi et al. (2022) has already made strides in deriving SDE models for RMSprop and Adam, albeit under somewhat restrictive assumptions. They establish a scaling rule that they
In concurrent work, Xiao et al. (2024) derived an SDE for SignSGD in the high dimensional setting for linear regression task: See Appendix F in Xiao al. (2024) for a comparison with SDE.
2

---

Published as a conference paper at ICLR 2025
assert remains valid throughout the entirety of the dynamies. While their derivation builds on the approach of Jastrzebski et al. (2018), which may be problematic in the general setting (see Appendix analysis suggests that the SDEs proposed in Malladi et al. (2022)
E for further discussion), provide accurate approximations primarily near minima, implying that the corresponding scaling rules might not hold globally. (Zhou et al., 2020a) derived a Lévy SDE for Adam; however, the approximation relies on random coefficients, technique that is theoretically sound only under very specific conditions (see Kohatsu-Higa et al. (1997); Bishop and Del Moral (2019)). Xie et al. (2022) modeled AdamW W with an SDE to investigate the roles of learning rate adaptivity and momentum in saddle-point escaping and flat minima selection. However, since their derivation does not follow a formal framework, it currently lacks comprehensive approximation guarantees. Finally, Zhou et al. (2024) presented an informal SDE for the iterates of AdamW: Their derivation relies on several strong assumptions and approximations which may benefit from further formal justification.
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 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 significant contributor 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 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 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: 2024) provided convergence guarantees under restrictive assumptions such as bounded gradient and bounded noise. All aforementioned results only show that SignSGI,, Adam, and MSSrrop 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.
FORMAL STATEMENTS & INSIGHTS: THE SDES
This section provides the general formulations of the SDEs of SignSGD (Theorem 3.2) and AdamW (Theorem 3.) Due to the technical nature of the analysis, we refer the reader to the appendix for the complete formal statements and proofs and only provide a sketch of the proof of key results.
Assumptions and notation. In this section, we collect most of the notation and assumptions used in the paper. All our analysis take place on filtered probability space (S2, (Fr (Fe)c20.33 P). The batches are of size B 2 1 and modeled as i.i.d. random variables uniformly distributed on (1....,A0 jj. We assume that the stochastic gradient (xx()) == 1B Eiey V(f((z)) can be decomposed as V fz) + Z(z), where fz(x) is the full gradient and Z (xx) the batch noise. We assume that E[Z(z)] = O and unless we study the cases where the gradient variance is unbounded, we write (o)( = (x) (we omit the size of the batch y unless relevant.) s.t. v (xr) is bounded, Lipschitz, satisfies affine growth, and together with its derivatives, it grows at most polynomially fast (Definition 2.5 in Malladi al. (2022)). Importantly, we assume that Z(x) has a bounded and smooth probability density function whose derivatives are all integrable: A common assumption in the literature is for (1r) to be Gaussian² (Ahn et al., 2012 2: Chen et al., 2014; Mandt et al., 2016; Stephan et al., 2017: Zhu et al., 2019; Wu et al., 2020; Xie et al., 2021), while our assumption allows for heavy-tailed distributions such the Student's Specifically, Li al. (2017); Mertikopoulos and Staudigl (2018); Raginsky and Bouvrie (2012); Zhu et al. (2019): Mandt et al. (2016); Ahn et al. (2012); Jastrzebski et al. (2018) use a Gaussian noise with constant covariance matrix to model batch noise. To derive the stationary distribution around an optimum, we approximate the loss function
"ee Jastrzchski et al. (2018) for the justification why this might be the case.
3

---

Published as a conference paper at ICLR 2025
Losses DMN
Losses - CNN
Lasses Transorrmrr
Loases Restict
62-
2.4
3
0.1
LOGK
1.4
02
30
0.
2000a Iterations
5900
Iterations
Iterations
Iterations
Figure Comparison of SignSGD and its SDE in terms of f(z): Our SDE successfully tracks the dynamics of SignSGD on several architectures, datasets, and hyperparameters: DNN on the Breast Cancer dataset (Left); CNN on MNIST (Center-L.eft); Transformer on MNIST (Center-Right); ResNet on CIFAR-10 (Right).
Har as commonly done in the literature (Ge et al., with a quadratic convex function f(x) = 1 2015; Levy, 2016; Jin al., 2017; Poggio et al., 2017; Mandt et al., 2017: Compagnoni et al., 2023). Finally, > 0 is the step size, the Bs refer to momentum parameters. 0 > 0 is the (decoupled) I, I'-regulatization parameter, and c is a scalar used for numerical stability, Finally, we use W to indicate a Brownian motion.
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' Mil'shtein (1986)) and it quantifies the disparity between the discrete and the continuous processes.
Definition 3.1 (Weak Approximation), A continuous
stochastic process {Xi (Xelre(o.T) is an order
a weak approximation (or a-order SDE) of a discrete stochastic process {ax } [T/nl for every polynomial growth function g. there exists a positive constant C, independent of the stepsize 7, such that maxk=o....1I// [Eg (2xx) - Eg (Xx,)) < Cr7a.
3.1 SIGNSGD SDE
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:
Tk+1
Tx 7 sign (*)))))))
(I)
The following theorem derives formal continuous time model for SignSGD.
Theorem 3.2 (Informal Statement of Theorem C.16). Under sufficient regularity conditions, the solution of the following SDE is an order 1 weak approximation of the discrete update of SignSGD:
dX, = -1 2P
0)) + ((X.)dW.
(2)
where E(x) is the noise covariance E(x) E(5, (()) and (x) == sign( ( (x) 2P Ja (x) < 0) the noise of sigu (x))..
-
Proof idea. One needs to prove that the first and second moments of the increments of the discretiza- tion of the SDE match those of SignSGD up to an error of order ((1) and (17²), respectively. O
For didactic reasons, we next present a corollary of Theorem 3.2 that provides a more interpretable SDE. To do so, we model the batch noise with Gaussian distribution with constant covariance matrix,3 which is a common approach in the literature (Li et al., 2017: Mertikopoulos and Staudigl, 2018; Raginsky and Bouvrie, 2012; Zhu et al., 2019; Mandt et al., 2016; Ahn et al., 2012; Jastrzebski et al., 2018). 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.19). Under the assumptions of Theorem 3.2, and that the stochastic grudient is V (x) (x(x) + Z such that Z ~ (0, E), Σ = diag(of, (a ... ,0%) the following SDE provides 1 weak approximation of the discrete update of SignSGD
Σ vV((e)
Σ 1-11)(()))
dx. = -Erf
di + vTn Ia diag [Erf
dW.
(3)
V2
22
where the error function Frf(x) = v So - -² dt and the square are applied component-wise.
See Scction C.5 for more realistic noise structures.
4

---

Published as a conference paper at ICLR 2025
While Eq. may appear intricate 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 dynamies of SignSGD can be categorized into three distinct phases. The left of Figure 2 empirically verifies this result on a convex quadratic function.
applied element-wiss and Lemma 3.4. Under the assumptions of Coroll. 3.3, Y;: := --
V2
some constants E Rt, qt E R q R., the dynamics of SignSGD exhibits three phases:
1. Phase 1: If > 2' the SDE coincides with the ODE of SignGD:
dX, = - ( (X))))
(4)
2. Phase 2: 1f1 < [V]] < 2:
(a) -mY; q* < dE[ < -my - -
(b) For any a > 0, P [||X, Ex > a] S (d - |mY, + +
3. Phase 3: Yf the SDE is
d.x,- = - V
2 2 y ((xx)dd Vny la diag (>x)))(X))) dW.
(5)
For ease of reading, will informally refer to the gradient fxx the "signal" and to E as the "noise". Then, Lemma 3.4 tells us that the behavior of SignSGI) 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 E premultiplies (z). affecting the rate of descent. This is not the case for SGD where Σ only influences the diffusion term. To better understand the role of the noise, we study how it affects the dynamics of the loss on strongly convex functions and compare it with SGD. The dynamics of E [[|V ((x,) [2]] for general non-convex smooth functions is presented in Lemma C.24.
Remark:
Lemma 3.5. Let f be p-strronly convex, Tr(V²f(a) Lr, max be the maximum eigenvalue of Σ, and S. := f(X,) - f(X.). Then, during
1. Phase 1, S; S 4 (Vut - 2/50)²
2 so SignSGD stays in this phase for at most t., = 2/550 So:
2. Phase as A =
V2o
m
+
mum??
4rr
(Cr pdq²)
[[S S She -2ndt
--2-0)))
2 2113
3. Phase as A = (VZ + 14
17 L, (1-e² 22-0-)))))
[[S.] S Soc Ap
2 2
Proof idea. For each Phase, we use the respective SDE of SignSGD from Lemma 3.4 derive the SDE of S. via Itô's lemma. Then, we take its expectation to obtain the ODE of E [S.] and leverage the assumptions to establish a bound.
As per Eq. 4, during Phase SignSGD behaves like SignGD: Lemma 3.5 shows that, consistently with the analysis of SignGD in (Ma et al., 2022), such strong decrease in the loss value explains the fast initial convergence of the optimizer as well as of R.Sprrp and Adam. undergoes a 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, verify the dynamics of the expected loss around a minimum in the center-left of Figure 2.
this phase, the loss
*Soe Lemma C.21 for their definitions.
"The SDE of SGD is dX,- -X.((x)dt + vns.taw.
5

---

Published as a conference paper at ICLR 2025
Tejectories
x..00
0.00
0.10
0.12
bo 2)
00035
00033
G0005
Tiee
Verlks
Lesses
0.30
0.34
00005
0.32
Trw Red
0.0
Berations
FEX| 1
VariAs pedd
Iteraions
Figure 2: Phases of SignSGD: The ODE of Phase and the SDE of Phase 3 overlap with the "Full" SDE as per Lemma 3.4. In Phase 2, the dynamics satisfies the prescribed bounds (Left); Phases of the Loss: The bounds derived in Lemma 3.5 for the loss during the different phases correctly track the loss evolution (Center-Lefit The dynamics of the moments of X, predicted in Lemma 3.7 track the empirical ones (Center-Right)
the schedulers satisfy the condition in Lemma 3.9, the loss decays to O as prescribed. Otherwise, the loss does not converge to O (Right). For each figure, (ax) Mz for H diag 2), 7) 0.001, and E o² I2 where 0.1.
Lemma 3.6. For SGD, the expected loss satisfies: B S Soe 2ut + Loomn 24
(1 e 2ut).
Remark:
The two key observations are that:
1. Both in Phase 2 and Phase 3, the noise level max inversely affects the exponential conver- gence speed, while this trend is not observed with SGD;
2. The asymptotic loss of SignSGD (almost) linear in maa while that of SGD is quadratic.
W.nss Indeed, the asymptotic value of E[S.) in Phase 3 scales with á <2 5p
when the
noise max dominates the learning rate 17 and/or the minimum eigenvalue of the Hessian, or in general when nM ~ O, we can conclude that the scaling is (almost) linear in mmx
Additionally, we characterize the stationary distribution of SignSGD around minimum. To do this, study the behavior of SignSGD the literature (Ge et al., 2015; Levy, 2016; Jin et al., 2017: Poggio et al., 2017; Mandt et al., 2017; Compagnoni et al., 2023). Empirical validation is provided in the center-right of Figure 2.
quadratic loss function, which is common approach in
Lemma 3.7. Let H- diag(A (Ai... Ad) and M, == 2 x 2 H4 23. 11p): Then,
1. E = = V2E b AHtXo;
2. Cou [Xo] (M- X3
x3+ +
H H-1 (V !HH
(Ia Mi).
Therefore, we have that the stationary distribution of SignSGD is:
(B[X=]) = 0, 2
2 La 2 TH
H -
Proof idea. For the E [2). we take the expected value of the SDE of Phase 3 from Lemma 3.4 and integrate the resulting ODE, For CovlX, we derive the SDE of X,X, via Itô's lemma, take the expectation, and integrate the resulting ODE. Then, we subtract E[X,JE(X,)"
Lemma 3.8. Under the same assumptions as Lemma 3.7, the stationary distribution for SGD is:
e = - 5=20 0 and Coe [X.] - - (12-1 1:: 2 3H-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 D 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 convex quadratic as we use ny (1+1)) for & (10 2.2 ).

---

Published as a conference paper at ICLR 2025
Lemma 3.9. Under the assumptions of Lemma 3.5, any step size scheduler nu such that
nads = oo and lim Th = = [l)( - f(X.)] ' C 400
Tt
0.
(6)
4
Remark:
Under the same conditions, SGD satisfies [E)) f f(x.)] -
4M
1/1
$ 0.
Conclusion: As noted in Bernstein et (2018), the Signal-to-noisse 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 Σ while for SGD it does 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 much lower loss values. symmetric argument holds for high levels of noise, which suggests that SignSGD is more resilient to high levels of noise.
3.1.1 HEAVY-TAILED NOISE
Interestingly, we can replicate the efforts above also in case the noise (xx) is heavy-tailed as is distributed according Student's 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 where the noise on the gradients Z (0, Ia) and v E zt, the following SDE is a 1 weak approximation of the discrete update of SignSGD
-2 (=(EEvVI(X,)) di + viv/ la - 4 diag (E(E-bVS(X))
aw,
: : - is is the gamma function, and 2 F1 is the hypergeometric function. Above, the E(r) and the square are applied component-wise.
T(442) where E(x) defined as 3(x) VFrl( 2
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 v = 2, the dynamics in Phase is:
dx, dx,= = -1=-V -X--VV3-- HVvx)x( + vnV la diag (2:0))00) (2=)v)(X)) dW,.
(8)
Conclusion: We observe that the dynamics of SignSGD when the noise is Gaussian (Eq. 5) w.r.t. when it is heavy-tailed with unbounded variance (Eq. 8) are very similar: By comparing the constants (v1/2 and / 22/m) in front of the drift terms Σ + V f(X)), 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 scenario where SGD would diverge.
All the results derived above can be extended to this heavy-tailed noise setting: See Compagnoni et al. (2025) for detailed discussion in the distributed setting.
3.2 ADAMW SDE
In the last subsection, showeased how SDEs can serve as powerful tools to understand the dynamics of the simplest among coordinate-wist adaptive methods: SignSGD. Here, we extend the discussion to Adam with decoupled weight decay, i.e. AdamW:
0k+1 = 32th + (1 (1-32) (T
m 31m+ + (1 B1) (xx),
m
m
2
Ik+1 = I&
Vūu+1
700x my
=
3
0g = 1 3/
(9)
which, of course, covers Adam, RMSprop, and RMSpropW depending on the values of and 3
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 ResNet.
7

---

Published as a conference paper at ICLR 2025
Trajeloriis
Losses
Trajectories
Losses
Adamnl
1.3
1.2
10
Assart
38
10
-3.6
11-1
10600
3206 2049
40699 $8000
5090 12492 1569 2000
2000
X1
Rerations
Iterations
Figure 3: The two images on the left compare the SDEs of AdamW and RMSpropW with the respective optimizers in terms of trajectories and f((:) for convex quadratic function while the other two provide a comparison for an embedded saddle. In all cases, we observe good agreements.
Theorem 3.12 (Informal Statement of Theorem C.53). Under sufficient regularity conditions, P ()r s.t. < (0, 1), and P2 = (1), the order weak approximation of Adam is:
dX,, Vi2(t)
Pr (M² + nPh (V) - M,)) 0X.dl
(10)
11(()
dM, PI (Vf (X.) M)) dt + Vi V EX.,jaWW
(11)
dV; L' ((V.((x,)) + diag (2 (X.)) - Vi) dt,
(12)
where B, =1 np; ~ 1, 44()= = - p lu and P: = diag VV. + eever())la
M, and Vi the exponential moving averages of the gradient and the squared gradient, respectively, P. acts as an adaptive preconditioner, scaling the parameter updates X, based on the accumulated squared gradients in V,- While M, is the momentum term and captures the history of gradients to smooth out the updates, 711 (X,) - Mc) adjusts M, towards the current gradient, ensuring responsiveness to recent changes. Finally, - O.X, dt applies regularization by shrinking the parameters.
We highlight that in contrast to Remark 4.3 of Malladi et al. (2022), which suggests that an SDE for Adam is only viable if >> [(x]] and σ ~ 1, our derivation that does not need these assumptions: 17* See Remark C.46 for a deeper discussion, the implications, and the experimental comparisons.
The following result demonstrates how the asymptotic expected loss of Adam scales with the noise level. Notably, it introduces the first scaling rule for 1, proposing an alternative to the one proposed for Adam in (Malladi et al., 2022) and extending it 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 (c.g., from 17 to 1j) to retain certain performance metrics optimizer properties as the batch size increases (e.g., from B to B). Therefore, in our upcoming analysis, we aim derive scaling rules that preserve specific relevant aspects of the dynamics, such as the convergence bound on the loss the speed. See Appendix E for more detailed discussion motivating our approach.
Lemma 3.13. If f is p-strongly convex and L-smooth, TH(V2 f(x)) S L,, X. 0, E(x) = o² and f(x))² = O()), ń = kn, B = B8, and p. a4Pb, and 0 = E0, AdamW satisfies
(x,) f(x.) 60 nLyaT.
(13)
2 2jVBBL. (EO(T.+
We derive the novel scaling rule by 1) Preserving the upper bound, which requires that v8 and E = k; 2) Preserving the relative speed of Mr V, and X., which requires that [3 1- Ba).
x(1
The left of Figure 5 shows the empirical verification of the predicted loss value and scaling rule on a convex quadratic function: Consistently with Lemma 3.13, such a value is bounded w.r.t. a, meaning that the loss of AdamW does not diverge to infinity even if there is an infinite level of gradient noise: See the right of Figure for an experimental validation. Interestingly, the asymptotic loss value not influenced by the choice of 35: We argue that B, 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 Fig. 5 exemplifies this simple non-convex landscape. Finally, we observe that the scaling rule informally derived in Malladi et al. (2022) prescribes 3. 1 x² (1 B.): This can be recovered by preserving other quantities. While there is no reason to prefer one over the other a priori, our analysis on L.LMs in Appendix F.8 shows that rescaling might be preferable in practice.
8

---

Published as a conference paper at ICLR 2025
Losses Tansforme
ResNet
Losses Transformer
Losses ResNet
40
Marn
304
233
LOGG
LosS
15
0.1
181
Reratiors
lterations
Iterations
Rterall ons
Figure The two images on the left represent the comparison between AdamW and its SDE in terms of f (z). The two on the right do the same for RMSpropW. In both cases, the first is Transformer on MNIST and the second a ResNet on CIFAR-10: Our SDEs match the respective optimizers.
Losses
Losses
Losses
Manet ki ge
21
Lh-4
VartX
19
920 12393 10000 2000
4200
12G0
200
2003 20000 2086
4900 300
20000
2200
Iterations
Iterations
Iterations
Figure 5: The loss predicted in Lemma 3.13 matches the experimental results on a convex quadratic function. AdamW run with regularization parameter = Adam W R (AdamW Rescaled) is run as we apply the scaling rule with x = 2. Adam W NR (AdamW Not Rescaled) is run as apply the scaling rule with holds, and failing rescale 0 leads the optimizer not to preserve the asymptotic Ioss level. The same happens for 0 = 4 (Left); The same for RMSpropW (Center-Leftt) For AdamW, 31 and 32 influence which basin will attract the dynamics and how fast this will converge, but not the asymptotic loss level inside the basin (Center- (Center-Right.. For both AdamW and RMSpropW, the variance at convergence predicted in Lemma 14 matches the experimental results (Right).
on all hyperparameters but 0, which left unchanged: Our scaling rule
Interestingly, the fact that the weight decay is decoupled is key to determining the dependency of the asymptotic loss of AdamW w.r.t. the noise level a. While the asymptotic loss of Adam
upper-bounded in a, the same does not hold if we Adam the L'-regularized loss f(x) + 0||=|1
Under the same assumptions of Lemma 3.13, the dynamics of Adam on f(x) + A||=l1 implies that
to00 nL,0
L
Elf(x.) f(x.)) f(x.)]
(14)
2 2μL + 0(L + n))
meaning that the asymptotic loss level grows linearly in a: See Figure 14 for empirical validation.
We conclude this section with the stationary distribution of AdamW around minimum which empirically validate on the right of Figure 5.
Lemma 3.14. If (xx) E, the stationary distribution of AdamW W is (0. ((+ -1
[B[Xx,)
We derived the analogous results for RMSprop(W) and we reported them in Appendix C.7: importantly, we validate the SDE in Figure 3 for simple landscapes and in Figure 4 for Transformer and a ResNet. The results regarding the asymptotic loss level and stationary distributions are validated in the center-left and right of Figure for convex quadratic function.
RMSpropW
While for both SignSGD and Adam the asymptotic loss value and the covariance of the iterates scale linearly with Σi, we observe for AdamW this is more intricate: The interaction between curvature, noise, and regularization implies that these two quantities are upper-bounded in E and increasing Σ to infinity does not lead to their explosion: decoupled weight decay plays a crucial stabilization role at high noise levels near the minimizer - See Figure 6 for a comparison across optimizers. Finally, argue that 33 play key role in selecting the basin and the convergence speed to the asymptotic loss value rather than impacting the loss value itself.
Conclusion:

---

Published as a conference paper at ICLR 2025
5G0
SignSGD
Adam
AdamWy
122
bete (m-18-1
102
bete (g=10)
Oete
32-1
Unit langgi
Uiit ta-20
12-
3371
13"
10-
23:1
25450 20000 13600 10000
25332 50000 73600 100000
2330 56000 25008 330038 Reretions
25000 Sidde 7::00 200303 Iterations
Iterations
Rteraions
Figure 6: For SGD (Left), SignSGD (Center-Left), Adam (Center-Rights and AdamW: For each optimizer, we plot the loss value on convex quadratic and compare its asymptotic value with the limits predicted by our theory. As we take o²Id, Id, 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 = 0). For AdamW, the maximum asymptotic loss value is bounded in a (Lemma 3.13 with 0 0). In accordance with the experiments, our theory predicts that adaptive methods are more resilient to noise.
4 EXPERIMENTS: SDE VALIDATION
The point of our experiments 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 so, we integrate the SDEs with Buler-Maruyama (Algorithm 1): This particularly challenging and expensive as one needs to calculate the full gradients of the DNNs at each iteration. Ours is the first set of validation experiments on various architectures and datasets: An MLP on the Breast Cancer dataset, CNN and a Transformer on MNIST, and a ResNet on CIFAR-10. Details in Appendix F.
5 CONCLUSION
We derived the first formal SDE for SignSGD, enabling 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 straightforward comparison with SGD. For SignSGD, the noise level VE 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 Adam W 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 R.Spprop scale linearly in the noise level EJ. For AdamW and RMSpropW, the complex interaction of noise, curvature, and regularization implies that these two quantities are bounded in terms of Σ showing that decoupled weight decay plays a crucial stabilization role at high noise levels near the minimizer. Interestingly, the SDEs for Adam and RMSprop are a 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 the various optimizers with the respective SDEs and find good agreement on simple landscapes and DNNs. For Adam and RMSprop, our SDEs track them more faithfully than those derived in (Malladi et al., 2022).
Future work We helieve 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 and preserve 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.
MMany papers derive SDFs to model optimizers, but most lack validation. Some use toy landscapes, while only Paquette et al. (2021); Compagnoni al. (2023) validate on simple DNNs. See Appendix A for details.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGMENTS
Enea Monzio Compagnoni, Rustem Islamov, and Aurelien Lucchi acknowledge the financial support of the Swiss National Foundation, SNF grant No 207392. Tianlin Liu acknowledges the financial support of the European Research Council Starting Grant knowledges the financial support of the Hector Foundation. Frank Norbert Proske acknowledges the financial support of the Norwegian Research Council (project No 274410) and MSCA4Ukraine (project No 101101))3).
Antonio Orvieto ac-
REFERENCES
Ahn, S., Korattikara, A., and Welling, M. (2012). Bayesian posterior sampling via stochastic gradient fisher scoring. arXiv preprini arXiv: 1206.6380..
An, J., Lu, J., and Ying, L. (2020). Stochastic modified equations for the asynchronous stochastic gradient descent. Information and Inference: Journal of the IMA, 9(4:85518733
Ankirchner, S. and Perko, S. (2024). A comparison of continuous approximations to stochastic gradient descent. Journal of Machine Learning Research, 25(13):1-55.
Ayadi, I. and Turinici, G. (2021). Stochastic runge-kutes methods and adaptive sgd-g2 stochastic
gradient descent. In 2020 25th International Conference on Pattern Recognition (ICPR), pages $220-8227. IEEE
Balles, and Hennig, (2018). Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference on Machine Learning, pages 404-413. PMLR.
Barakat, A. and Bianchi, P. (2021). Convergence and dynamical behavior of the adam algorithm for nonconvex stochastic optimization. SIAM Journal on Optimization, 3 31(1):444-774.
Bardi, M. and Kouhkouh, H. (2022). Deep relaxation of controlled stochastic gradient descent via singular perturbations. arXiv preprint arXiv:2209. 055664.
Bercher, A., Gonon, L., Jentzen, A., and Salimova, D. (2020). Weak error analysis for stochastic gradient descent optimization algorithms. arXiv preprint arXiv:2007.02723.
Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anandkumar, A. (2018). signSGD: Compressed optimisation for non non-convex problems. In Proceedings of the 35th International Conference on Machine Learning.
Biderman, S., Schoelkopf, H., Anthony, Q., Bradley, H., O'Brien, K., Hallahan, E., Khan, M. A.,
Purohit, S., Prashanth, U. S., Raff, E., Skowron, A., Sutawika, L., and Van Der Wal, O. (2023). Pythia: suite for analyzing large language models across training and scaling. In Proceedings of the 40th International Conference on Machine Learning, ICML'22 '2. JMLR.org.
Bishop, A. N. and Del Moral, P. (2019). Stability properties of systems of linear stochastic differential equations with random coefficients. SIAM Journal on Control and Optimization,
Bock, S. and WeiB, M. G. (2021). Local convergence of adaptive gradient descent optimizers. arXiv preprint
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs.
Chen, P., Lu, J., and Xu, (2022). Approximation stochastic variance reduced gradient langevin dy- namics by stochastic delay differential equations. Applied Mathematics & Optimization, 85(2):15.
Chen, T., Fox, E., and Guestrin, C. (2014). Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pages 1683-1691, PMLR
Chen, X., Liu, S., Sun, R., and Hong, M. (2019). On the convergence of class of adam-type algorithms for non-convex optimization. In International Conference Learning Representations.

---

Published as a conference paper at ICLR 2025
Compagnoni, E. M., Biggio, L., Orvieto, A., Proske, N., Kersting, H.,, and Lucchi, A. (2023). An sde for modeling sam: Theory and insights. In International Conference on Machine Learning, pages
PMLR.
Compagnoni, E. M., Islamov, R., Proske, N.N., and Lucchi, A. (2025). Unbiased and sign compression in distributed learning: Comparing noise resilience via SDEs. In The 28th International Conference Artificial Intelligence and Statistics.
Compagnoni, E. M., Orvieto, A., Kersting, H., Proske, F., and Lucchi, A. (2024). Sdes for minimax optimization. International Conference on Artificial Intelligence and Statistics, pages 4834 4842. PMLR.
Cui, 7.-X., Fan, Q. and Jia, C. (2020). Momentum methods for stochastic optimization time-varying directed networks. Signal Processing, 174:107614.
Dambrine, M., Dossal, C., Puig, B., and Rondepierre, A. (2024). Stochastic differential equations for modeling first order optimization methods. SIAM Journal Optimization, (2(2):44021444
De, S., Mukherjee, A., and Ullah, E. (2018). Convergence guarantees for rmsprop and adam in non-convex optimization and an empirical comparison to nesterov acceleration. arXiv preprins arXiv:1807.06766
Défosse7, A., Bottou, L.., Bach, F., and Usunier, N. (2022). A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research.
Del Moral, P. and Niclas, A. (2018). A taylor expansion of the square root matrix function. Journal of Mathematical Analysis and Applications, 465(1):259-266
Deng, L. (2012). The mnist database of handwritten digit images for machine learning research. IEFE Signal Processing Magazine, 29(6):1411142
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. (2021). An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations.
Dua, D. and Graff, C. (2017). UCI machine learning repository.
Fontaine, X., De Bortoli, V., and Durmus, A. (2021). Convergence rates and approximation results for sgd and its continuous time counterpart. In Conference on Learning Theory, pages 1965-2058. PMLR.
Ge, R., Huang, F., Jin, C., and Yuan, Y. (2015). Escaping from saddle points online stochastic gradient for tensor decomposition. Conference on Learning Theory, pages 797-842.
Gess, B., Kassing, S., and Konarovskyi, V. (2024). Stochastic modified flows, mean-fiee -ield limits and dynamics of stochastic gradient descent. Journal of Machine Learning Research, 25(3)):
Gu, H., Guo, X., and Li, X. (2021). Adversarial training for gradient descent: Analysis through its continuous time approximation. arXiv preprint arXiv:2105.0800
Guo, Z., Xu, Y., Yin, W., Jin, R., and Yang, T. (2021). A novel convergence analysis for algorithms of the adam family. arXiv preprint ar arXiv:2112.0344
Harris, C. R., Millman, K. J., der Walt, S. J., Gommers, R., Virtanen, P., Coumapeau,, D... Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, J. F., Wiebe, M., Peterson, P., Gérard-Marchamn P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C., and Oliphant, T. E. (2020). Array programming with NumPy. Nature, (85(7825)3357-362
Higham, D. J. (2001). An algorithmic introduction numerical simulation of stochastic differential equations. SIAM review, 43(3):525-546.
Hong, Y. and Lin, J. (2023). High probability convergence of adam under unbounded gradients and affine variance noise. arXiv preprint 0ttii:2311.02000
12

---

Published as a conference paper at ICLR 2025
Hu, W., Li, J., and Zhou, X. (2019). On the global convergence of continuous- time stochastic heavy- ball method for nonconvex optimization. In 2019 IEEE International Conference on Big Data (Big Data), pages 94-104. IEEE.
Huang. F., Li, J., and Huang, H. (2021). Super-adam: faster and universal framework of adaptive gradients. Advances in Neural Information Processing Systems, 34:9074-9085.
Ikeda, N. and Watanabe, S. (2014), Stochastic differential equations and diffusson processes. Elsevier.
Jastrzebski, S., Kenton, Z., Arpit, D., Ballas, N.. Fischer, A., Bengio, Y., and Storkey, A. (2018). Three factors influencing minima in sgd. ICANN 2018.
Jin, C., Ge, R., Netrapalli, P., Kakade, S. M., and Jordan, M. I. (2017). How to escape saddle points efficiently, In International Conference on Machine Learning, pages 1724-1732. PMLR.
Karatzas, I. and Shreve, S. (2014). Brownian motion and stochastic calculus, volume 113. springer.
Karimireddy, S. P., Rebjock, Q., Stich, S., and Jaggi, M. (2019a). Error feedback fixes signsgd and other gradient compression schemes. International Conference on Machine Learning, pages 3252-3261. PMLR.
Karimireddy, S. P., Rebjock, Q., Stich, S., and Jaggi, M. (2019b). Error feedback fixes SignSGD and other gradient compression schemes. In Proceedings of the 36th International Conference on Machine Learning.
Kingma, D. P. and Ba, (2015). Adam: A method for stochastic optimization. In International Conference on Learning Representations
Kohatsu-Higa, A., León, J. A., and Nualart, D. (1997). Stochastic differential equations with random coefficients. Bernoulli, pages 233-245.
Krizhevsky, A., Hinton, G., et al. (2009). Learning multiple layers of features from tiny images. Toronto, ON, Canada.
Kunin, D., Sagastuy-Bras Breena, J., Gillespie, L., Margalit, E., Tanaka, H., Ganguli, S., and Yamins, D. (2023). The limiting dynamics of sgd: Modified loss, phase-space oscillations, and anomalous diffusion. Neural Computation, 36(1): 36(1):511-744
Kunstner, F., Yadav. R., Milligan, A., Schmidt, M., and Bietti, A. (2024). Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402.1944 19449.
Lanconelli, A. and Lauria, C. S. (2022). A note on diffusion limits for stochastic gradient descent. arXiv preprint arXiv:2210.11255
Levy, K. Y. (2016). The power of normalization: Faster evasion of saddle points. arXiv preprint arXiv:1611.04831
Li, H., Rakhlin, A., and Jadbabaie, A. (2023a). Convergence of adam under relaxed assumptions. In Thirty-sevntth Conference on Neural Information Processing Systems.
Li, and Wang, (2022). On mniform-in-time diffusion approximation for stochastic gradient descent. arXiv preprint vrXiv:2207.04922
Li Q., Tai, C., and Weinan, E. (2017). Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pages 2101-2110. PMLR.
Li, Q. Tai, C., and Weinan, E. (2019). Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. The Journal of Machine Learning Research. 1041):144-1520.
Li, Z., Malladi, S., and Arora. S. (2021). On the validity of modeling SGD with stochastic differential equations (SDEs). In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems.
13

---

Published as a conference paper at ICLR 2025
Li, Z., Wang. Y., and Wang, Z. (2023b). Fast equilibrium of sgd in generic situations. In The Twelfth International Conference on Learning Representations.
Liu, T., Chen, Z., Zhou, E., and Zhao, T. (2021). A diffusion approximation theory of momentum stochastic gradient descent in nonconvex optimization. Stochastic Systems.
Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. In International Conference on Learning Representations.
Luo, L., Xiong, Y., and Liu, Y. (2019). Adaptive gradient methods with dynamic bound of learning rate. In International Conference on Learning Representations.
Ma, C., Wu, L., and Weinan, E. (2022). A qualitative study of the dynamic behavior for adaptive gradient algorithms. In Mathematical and Scientific Machine Learning, pages 671-692. PMLR.
Mai, V. V. and Johansson, M. (2021). Stability and convergence of stochastic gradient clipping: Beyond lipschitz continuity and smoothness. In International Conference on Machine Learning.
Malladi, S., Lyu, K., Panigrahi, A., and Arora, S. (2022). On the SDEs and scaling rules for adaptive gradient algorithms. In Advances in Neural Information Processing Systems.
Mandt, S., Hoffmann M., and Blei, D. (2016). A variational analysis of stochastic gradient algorithms. In International conference machine learning, pages 354-363. PMLR.
Mandt, S., Hoffman, M. D., and Blei, D. M. (2017). Stochastic gradient descent approximate bayesian inference. JMLR 2017.
Mao, X. (2007). Stochastic differential equations and applications. Elsevier.
Maulen-Soto, R., Fadili, J., Attouch, H., and Ochs, P. (2024). Stochastic inertial dynamics via time scaling and averaging. arXiv preprint arXiv:2403.16775.
Maulén Soto, R. I. (2021). A continuous-time model of stochastic gradient descent: convergence rates and complexities under lojasiewicz inequality. Universidad de Chile.
Mertikopoulos, and Staudigl, M. (2018). On the convergence of gradient-like flows with noisy gradient input. SIAM Journal on Optimization, 28(1):663-197
Milstein, G. N. (2013). Numerical integration of stochastic differential equations, volume 313. Springer Science Business Media.
Mil'shtein, G. (1986). Weak approximation of solutions of systems of stochastic differential equations. Theory of Probability & Its Applications, 30(4):750-766.
Noci, L., Anagnostidis, S., Biggio, L.,, Orvieto, A., Singh, S. P., and Lucchi, A. (2022). Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27198-27211.
Øksendal, B. (1990). When is stochastic integral a time change of a diffusion? Journal of theoretical probability, 3(2):207-226.
Orvieto, A. and Lucchi, A. (2019). Continuous-tims models for stochastic optimization algorithms. Advances in Neural Information Processing Systems, 32.
Pan, Y. and Li, Y. (2022). Toward understanding why adam converges faster than SGD for transform- ers. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop).
Paquette, C., Lee, K., Pedregosa, F., and Paquette, E. (2021). Sgd in the large: Average-case analysis, asymptotics, and stepsize criticality. In Conference on Learning Theory, pages 3548 3548-3626. PMLR.
Paquette, E., Paquette. C., Xiao, L., and Pennington, J. (2024). 4+ 3 phases of compute-optimal neural scaling laws. arXiv preprint arXiv:2405.15074,
Pascanu, R., Mikolov, T., and Bengio, (2013). On the difficulty of training recurrent neural networks. In International conference on machine learning.
14

---

Published as a conference paper at ICLR 2025
Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg. V., Vanderplas, J., Passos, A., Couraapeau, D., Brucher, M., Perrot, M., and Duchesmay, E. (2011). Scikit-lcam: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830.
Poggio, T., Kawaguchi, K., Liao, Q. Miranda, B., Rosasco, L., Boix, X., Hidary, J., and Mhaskar, H. (2017). Theory of deep learning ii:: explaining the non-overfitting puzzle. arXiv preprint arXiv: 1801.0017.
Puchkin, N., Gorbunov, E., Kutuzov, N., and Gasnikov, A. (2024). Breaking the heavy-tailed noise barrier in stochastic optimization problems. In International Conference on Artificial Intelligence and Statistics.
Raginsky, M. and Bouvrie, J. (2012). Continuous-time stochastic mirror descent on network: Variance reduction, consensus, convergence. In 2012 IEEE 51st IEEE Conference on Decision and Control (CDC), pages 6793-6800. IEEE.
Reddi, S. Kale, S., and Kumar, S. (2018). On the convergence of adam and beyond. In International Conference Learning Representations
Safaryan, M. and Richtarik, P. (2021). Stochastic sign descent methods: New algorithms and better theory. In Proceedings of the 38th International Conference on Machine Learning-
Shi, N. and Li, D. (2021). Rmsprop converges with proper hyperparameter. In International conference on learning representation.
Smith, S. L., Dherin, B., Barrett, D. G. T., and De, (2021). On the origin of implicit regularization in stochastic gradient descent. ArXiv, 12212101.12116..
Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, J. R., Hestness, J., and Dey, N. (2023). SlimPajame A 627B token cleaned and deduplicated version of RedPajama.
Soto, R. M., Fadili, J., and Attouch, H. (2022). An sde perspective on stochastic convex optimization. arXiv preprint arXiv:2207.0273 22.00000.
Stephan, M., Hoffman, M. D., Blei, D. M., et al. (2017). Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 18(134):1-35
Su, I. and Lau, V. K. (2023). Accelerated federated learning over wireless fading channels with adaptive stochastic momentum. IEEE Internet of Things Journal.
Sun, J., Yang, Y., Xun, G., and Zhang, A. (2023). Scheduling hyperparameters to improve generaliza- tion: From centralized sgd to asynchronous sgd. ACM Transactions on Knowledge Discovery from Data, 17(2):1337.
Tieleman, T. and Hinton, G. (2012). Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude.
Van Rossum, G. and Drake, F. L. (2009). Python 3 Reference Manual. CreateSpace, Scotts Valley, CA.
Wang, B., Fu, J., Zhang, H., Zheng, N., and Chen, W. (2024). Closing the gap between the upper bound and lower bound of adam's 's iteration complexity. Advances in Neural Information Processing Systems, 36.
Wang. B., Zhang, Y., Zhang, H. Meng Q., Ma, Z.-.., Liu, T.-Y., and Chen, W. (2022). Provable adaptivity in adam. arXiv preprint arXiv:2208. 09900.
Wang, Y. and Wu, S. (2020). Asymptotic analysis via stochastic differential equations of gradient descent algorithms in statistical and computational paradigms. Journal of machine learning research, 111
Wang, Z. and Mao, Y. (2022). Two facets of sde under an information ceeemation-thorettic lens: Generalization of sgd via training trajectories and via terminal states. arXiv preprint arXiv:2211.1069)
15

---

Published as a conference paper at ICLR 2025
Wojtowytsch, S. (2024). Stochastic gradient descent with noise of machine learning type part ii: Continuous time analysis. Journal of Nonlinear Science, 34(1):16.
Wolkowicz, H. and Styan, G. P. (1980). Bounds for eigenvalues using traces. Linear algebra and its applications, 29:471-506.
Wu, J., Hu, W., Xiong, H., Huan, J., Braverman, V., and Zhu, Z. (2020). On the noisy gradient descent that generalizes as sgd. In International Conference on Machine Learning, pages 10367-10376.
PMLR.
Wu, L., Wang, M., and Su, W. (2022). The alignment property of sgd noise and how helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems, 35:4680-4693.
Xiao, K. L.,, Marshall, N., Agarwala, A., and Paquette, E. (2024). Exact risk curves of signsgd in high- digh-dimensions: Quantifying preconditioning and noise-compression effects. arXiv preprint arXiv:2411.12133.
Xie, Z., Wang. X., Zhang, H., Sato, I., and Sugiyama, M. (2022). Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International conference machine learning, pages 24430 24459. PMLR.
Xie, Z., Yuan, L.,, Zhu, Z., and Sugiyama, M. (2021). Positive-negativy momentum: Manipulating stochastic gradient noise to improve generalization. In Proceedings of the 38th International Conference on Machine Leurning, volume 139 of Proceedings of Machine Learning Research, pages 11448-114 11458. PMLR.
Yang, J., Li, X., Fatkhullin, I., and He, N. (2024). Two sides of one coin: the limits of untuned sgd and the power of adaptive methods. Advances in Neural Information Processing Systems, 36.
Zaheer, M., Reddi, S., Sachan, D., Kale, S., and Kumar, S. (2018). Adaptive methods for nonconvex optimization. Advances in neural information processing systems, 31.
Zhang. J., He, T., Sra, S., and Jadbabaie, A. (2020a). Why gradient clipping accelerates training: A theoretical justification for adaptivity. International Conference Learning Representations,
Zhang. J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra. S. (2020b). Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems.
Zhang, Y., Chen, C., Shi, N., Sun, R., and Luo, Z.-Q. (2022). Adam can converge without any modification on update rules. Advances in neural information processing systems.
Zhang, Z., Li, Y.. Luo, T., and Xu, Z.-Q. (2023), Stochastic modified equations and dynamics of dropout algorithm. arXiv preprint arXiv:2305.1588
Zhao, J., Lucchi, A., Proske, F. N., Orvieto, A., and Kersting, H. (2022). Batch size selection by stochastic optimal control. In Has it Trained Yet? NewrIPS 2022 Workshop.
Zhou, P., Feng. J., Ma, C., Xiong, C., Hoi. S. C. H., al. (2020a). Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems,
Zhou, P., Xie, X., Lin, Z., and Yan, (2024). Towards understanding convergence and generalization of adamw. IEEE Transactions on Pattern Analysis and Machine Intelligence.
Zhou, P. Xie, X.. and Shuicheng, Y. (2022). Win: Weight-decayy-nteege nesterov accelera- tion for adaptive gradient algorithms. In The Eleventh International Conference on Learning Representations.
Zhou, X., Yuan, H., Li, J., and Sun. Q. (2020b). Stochastic modified equations for continuous limit of stochastic admm. arXiv preprint arXiv:2003.03532.
Zhu, Y. and Ying, L. (2021). A sharp convergence rate for a model equation of the asynchronous stochastic gradient descent. Communications in Mathematical Sciences.
16

---

Published as a conference paper at ICLR 2025
Zhu, Z... Wu, J., Yu, B., Wu, L., and Ma, J. (2019). The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. JCML.
Ziyin, L.,, Liu, K., Mori, T., and Ueda, M. (2021). Strength of minibatch noise in sgd. arXiv preprint arXiv:2102.05375.
A ADDITIONAL RELATED WORKS
In this section, we list some papers that derived or used SDEs to model optimizers. In particular, we focus on the aspect 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 neural networks.
While the following works contribute to the theoretical understanding. they do not provide direct experimental validation of their SDE approximations or the derived insights: (Liu 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 al., 2023; Sun al., 2023; Li et al., 2023b; Gess al., 2024; Dambrine et al., 2024; Maulen-Soto et al., 2024).
The following ones carried validation experiments on artificial landscapes, e.g. quadratic or quartic function, or easy regression tasks: (Li et al., 2017; 2019; Zhou et al., 2020b; An et al., 2020; Fontaine al., 2021; Gu al., 2021; Su and Lau, 2023; Ankirchner and Perko, 2024).
Some recent papers have performed experiments involving neural networks (see. e.g., (Paquette al., 2021; Compagnoni et al., 2023)). In both works, the authors simulate the corresponding SDEs using numerical integrators and then compare the results with those of the associated optimizers. Specifically, (Paquette et al., 2021) validates the SDE on a shallow MI.P, while (Compagnoni et al., 2023) performs the validation on both shallow and deep MLP.
In contrast, the studies by (Li et al., 2021; Malladi et al., 2022) do not empirically validate their derived SDEs through numerical integration. While their analysis of SVAG provides valuable insights, we believe there are conceptual aspects worth discussing:
1. After deriving an SDE for an optimizer which we refer to as "Optimizer A", one observes that simulating these SDEs can computationally expensive;
2. To mitigate this computational challenge, (Li et al., 2021; Malladi et al., 2022) introduce a discrete-time algorithm called SVAG, which shares the same SDE as "Optimizer A". Importantly, SVAG does not perform a numerical integration the original SDE, as it does not require access to either the drift or diffusion term;
3. (Li et al., 2021; Malladi et al., 2022) simulate SVAG and observe that it tracks "Optimizer A" closely, which they interpret as evidence that the SDE is a suitable approximation for "Optimizer A".
However, it is important to note that (Li et al., 2021; Malladi al., 2022) do not directly validate the SDE by simulating it. To illustrate potential concern with this approach, consider the following hypothetical scenario:
1. Derive an SDE for "Optimizer A";
2. Recognize that simulating the SDE is computationally costly;
3. Define another discrete-time algorithm called "Optimizer B", which coincides with "Opti- mizer A" and therefore, by definition, shares the same SDE;
4. Simulate "Optimizer B" and observe that it tracks " "Optimizer A" perfectly, since they are identical by construction:
5. Conclude that the SDE is good approximation for "Optimizer A".
17

---

Published as a conference paper at ICLR 2025
While the reasoning in prior work (Li et al., 2021; Malladi et al., 2022) demonstrates that SVAG is discrete-time optimizer that shares the same SDE as "Optimizer A", this alone does not establish that the SDE is the correct most suitable model for "Optimizer A". Simply comparing two algorithms that share an SDE does not necessarily confirm the validity of the SDE itself. Otherwise, an optimizer compared with itself would trivially satisfy this criterion. A more direct validation of the SDE would require numerical integration using method that explicitly incorporates the drift and diffusion terms (Higham, 2001; Milstein, 2013).
B STOCHASTIC CALCULUS
In this section, summarize some important results in the analysis of Stochastic Differential Equations 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)).
B.I ITÔ'S LEMMA
start with some notation: Let (R, (Fr)e20, P) be a filtered probability space. We say that an event F holds almost surely (a.s.) in this space if P(F) = We call (" b), Rd), with p > the family of -aluud Fa-adapted processes {fiJases such that
fell 80.
Moreover, we denote by MP((a, b), Rd), with p > the family of RRd processes {feassesh
in C(aa, b), Rd) such that E [12 [| S xo. We will write h L2 (R,,,R.) with p > 0, E C" (I0, 2], 1R for every T 0. Similar definitions hold for matrix- functions using the Frobenius norm ||||| = /E., A:yl2
Let W' = {Web2o X = (X be an Fi-adapted process taking values on Rd
one-dimensional Brownian motion defined on our probability space and let
Definition B.1. Let the drift be I E L² RR,, and the diffusion term be a E L² (R+,, Rdxm).. X, is Itô process if takes the form
x, xo + I bads (' o,dW,,
We shall say that X, has the stochastic differential
dX, b.dt + o,dW,
(15)
Theorem B.2 (Itô's Lemma). Let X, be an Itô process with stochastic differential dX, = b, + d. Let (x,t) 1) be twice continuously differentiable in and continuously differentiable in t, taking values in R. Then [(X,, 1) is again an hō process with stochastic differential
1 Tr (ex(x()))) di (V (16)
tf) (x, (V)
B.2 STOCHASTIC DIFFERENTIAL EQUATIONS
Stochastic Differential Equations (SDEs) are equations the form
dX, (x, aXx, t)dW,
First all, we need to define what it means for a stochastic process X = (X.h22 with values d to solve SDE.
18

---

Published as a conference paper at ICLR 2025
Definition B.3. Let X2 be above with deterministic initial condition Xu zu- Assume b
R' x [0 - R4 and a R x [0, + Raxm are Borel measurable; X, is called a solution to the corresponding SDE if
1. X, is continuous and Fr-adapted;
2. b L ([0, IR):
3. L ((0, T) Rd
4. For every / E [0,
X, = 3o J 5)d S( oX2,, a.s.
Moreover, the solution X, is said to be unique if any other solution Xī is such that
P X;, for all 0 s T) 1.
Notice that since the solution to an SDE is an Itô process, we can use Itō's lemma. The following theorem gives sufficient condition on and for the existence of solution to the corresponding SDE.
Theorem B.4. Assume that there exist two positive constants K and K such that
1. (Global Lipschitz for all x, y E Rd and e [0, T]
max[1b(x,, t) bby, t)1,, | t) - oy, 1) } s K|ls yll::
2. (Linear growth condition) for all x E Rd and E [0,
max (1b(a,, 1)1 (x, t) s K( + ||
Then, there exists unique solution X; the corresponding SDE, and X, (M R*).
Numerical approximation. Often, SDEs are solved numerically, The simplest algorithm to provide a sample path (ax) for Xx, so that XeAt 2 2g for some small At and for all & s M is called Fuler-Maruyama (Algorithm 1). For more details on this integration method and its approximation properties, the reader can check Mao (2007).
Algorithm Euler- Euler-Maruyann Maruyama Integration Method for SDEs
input
The drift b, the volatility a, and the initial condition Zg. Fix stepsize At; Initialize 0o = Tp:
0;
while <
T do
Ai
Sample some d-dimensional Gaussian noise Z₁ ~ NNO. Compute + tt + At b(Kx, kAt) + VAl a (fx, 7p; k end while
output The approximated sample path ($x) T
C THEORETICAL FRAMEWORK WEAK APPROXIMATION
In this section, we introduce the thcoretical framework used in the paper, together with its assumptions and notations.
19

---

Published as a conference paper at ICLR 2025
First of all, many proofs will use Taylor expansions in powers of 17. For ease of notation, we introduce the shorthand that whenever we write O (17%). we mean that there exists a function (x) E such write
that the error terms are bounded by K( (z))n For example,
b(x 7) (o(x)) + 191 (x) (
to mean: there exists E G such that
1b(z 77) bo(x) na4 (z)) S K())))
Additionally, we introduce the following shorthand:
A multi-inde index is = (A1. o.,,..., an) such that aj
a := 0 + 02 + ..... ani
:= a lollll! an!
For (21, 22 Inn RR, we define x = 2i x2
002
For multi-index 3, [ f(x) =
@IPl
fz))
B. 2 8
We also denote the partial derivative with respect to x by de:.
Definition C.1 (G Set). Let G denote the of continuous functions R - R of at most polynomial
growth, i.e. E G if there exists positive integers v1, 1/ > such that [g(x)] S v (1 + | for
all x E Rdd.
Definition C.2 (C) [RR,,O R)). C (R", R) denotes the space functions whose k-th derivatives are bounded.
The next results are inspired by Theorem of Li et al. (2017) and are derived under some regularity assumption on the function
C.1 ASSUMPTIONS
In general, we assume some regularity in the loss function.
Assumption C.3, Assume that the following conditions on f. E CØ (IR", (&",R), R), and their gradients are satisfied:
Vf, Vfr satisfy a Lipschitz condition: there exists L > 0 such that
VV(u)) - Vf(()) + Σ Vf((u)) f(() & L ul:
f. and its partial derivatives up to order belong to G;
Vf, V.f satisfy a growth condition: there exists M > 0 such that
[(((()) < M(1 [=]). [V(()) +
Regarding the gradient noise, each optimizer has its mild assumptions which weaker or in line with the literature.
SignSGD
1. The gradient noise Z(x) admits strictly positive density function gz for all æ and require that g RR x R + [0, xc) s.t. (x, + 9z (y) in C*(RR x R") such that all partial derivatives of up to order are integrable with respect to y and s.t. their L'-rmsss uniformly bounded æ. This assumption covers Gaussian and Student's t, thus being more general than the literature. Indeed, the Gaussianity of the noise commonly assumed:
20

---

Published as a conference paper at ICLR 2025
Among others, see Ahn et al. (2012); Chen et al. (2014); Mandt al. (2016); Stephan et al. (2017); Zhu et al. (2019); Wu et al. (2020); Xie et al. (2021), while Jastrzebski et al. (2018) offers intuitive justification well;
2. For all compact sets K
sup 19(2,-)) E L (R)), =ek
which of course covers the Gaussian case, thus being more general than the literature.
3. The functions in Eq. 18 to be in G, which, we show below, covers Gaussian and Student's I, thus being more general than the literature.
Adam(W) and RMSprop(WW W)
1. In line with Malladi et al. (2022), we assume that vI(x) is: G' together with its derivatives, Lipschitz, bounded, and satisfy Affine Growth;
2. The term (V fxx))² to be Lipschitz and of affine growth, which is consequence of assuming bounded gradients as often done in the literature on the convergence of R.Mprrpp and Adam: Among many, (Luo et al., 2019; Défossez et al., 2022; Guo et al., 2021; Huang et al., 2021) together with the discussion in Section 2.1 of Shi and Li (2021).
Remark All the assumptions above are in line with or more general than those commonly found in the literature. In line with Remurk 11 of the seminal paper Li et al. (2019), we observe that while some of these assumptions might seem strong, loss functions in applications have inward pointing gradients for sufficiently large Therefore, we could simply modify the loss to satisfy the assumptions above.
Regarding the drift and diffusion coefficients, we highlight that many papers in the literature following this framework do not check for their regularity before applying the approximation theorems Hu al. (2019); An et al. (2020): Zhu and Ying (2021); Cui al. (2020); Maulén Soto (2021); Wang and Mao (2022); Compagnoni et al. (2023; 2024); Li et al. (2017). At first sight, it would seem that not even the seminal paper Li et al. (2019) checks these conditions carefully. However, a deeper investigation shows that they are restricting their analysis to compact sets to leverage the regularity and convergence properties of mollifiers: The assumption regarding the compactness of the domain is not highlighted nor assumed in any part of the paper. Therefore, we conclude that, willingly or not, most papers implicitly make these assumptions.
C.2 TECHNICAL RESULTS
In this subsection, we provide some results that will instrumental in the derivation of the SDEs.
Lemma C.4. Assume the existence of probability density 9u of the gradient noise Z(x) for all and require that R" x R2" + [0, oo); (I, y) + (3) is in C (RR" x R²) such that all partial derivatives of 9 up to order are integrable with respect to y and such that their I norms are uniformly bounded in x. Further, let f E C* (R²) and h IR" + I. be a bounded Borel measurable function. Define the function & by
k(x) E[((())))
Then there exists version Te of with 1: E (R")
Proof. Let y be smooth and compactly supported. Then for all multi indices 8 with | substitution, Fubini's 's theorem, and integration by parts imply that
/ k(x))))) p(x)de =
/ E (M((((())
a p(x)dx
J JR.
h(y)g= vf()))dyo) p(z)dz
(()))
(-1)181 / 3
h(u)) (g-(v ( (())))dyp(()ddd
21

---

Published as a conference paper at ICLR 2025
So
J h(w)ag (0 - V dy
a weak derivative @ & of le on any bounded open set. For compact sets K we obtain that
L/ M(u)olg) (e
(v- - VVf())) dæ
K RM
<
|M (K(
sup
(gm (VV dy
"ER JR
for all p 2 2 because of our assumptions on and f and substitution (A" Lebesgue measure). So follows from Sobolev embeddings with respect to Hölder spaces that for all bounded and open sets f there exists a version & of l such that k (2)). The latter version can extended to 9 R". which we also denote by k. Since 0 is bounded for [3] < 8, we conclude that E (RR). O
Lemma C.5. Assuming that for all compact sets K
sup zeK
19(a, I. (TR",,
and the positivity of the density functions, we have that for m 1...., that
| .p....0.... 0r. A) & Cln((),
(17)
where the function lm(x) is defined as
(m(x)
1 m(x) + (())(n 1)1/2
(1+ 2s(x)( - 1)2/2
(1 m(x) s(x)(n - 1) 11/2
((+11/2)
g/ A(x))
r+1
max mm
(18)
Proof. To prove this, we need the fact that the Fréchet derivatives of the square root function yp can be represented as follows (see Theorem 1.1 in Del Moral and Niclas (2018)):
-b((A) He -tp(A)
and higher derivatives of order m 2 are given by
m!
V" V'"s(A)[HH.. H -VP(A)
Σ (p+ 1)1(q 1)!
[V+1 ( (AA)HH ..... H])
L x ( (())H,...... H)]
(19)
for all A e S and symmetric matrices H. Moreover, have the following estimate for 2 0:
|| {(A|| s (v + 2 Amin (A) (A)-(m+1/22)
(20)
where Amii (A) is the smallest eigenvalue of A and C,. := m = (3m) We find that 0xA7/2(x) V=((())) and and
0,0,A1?(2)) = +444
Thus, it follows from Eq. (20) that
| CAmin (1m
and
|
5 (A)))) [() [(A(())
+ +C2Amin 11
22

---

Published as a conference paper at ICLR 2025
More generally, for l.....
|0 03. (z)
Sz.
1 C.
Σ
Amin (A(x:) ((xx))-(r+((2)
(r+1122)
| A(x))
1+1
max 8
(21)
Let us now provide a lower bound for Amin terms of (r(((x)) and tr((A()))²). Define
(x)) = -n=- n- ((((((((())))) -tt(((=)))) m(x) = tr(A(z))
n
Then, from Corollary 2.1, Corollary 2.2, and Theorem 2.1 Wolkowicz and Styan (1980), we obtain
2n(x) 1)1/2
(1+
Amin (A(x))
Amax (A(x))
m(x) s(x) - 1)-1/2
1
2s(x)(n 1)
1 m(x) - 8(x) - 1)-1/3
m(x) + (nx)
Therefore, from Eq. (21), we have for = 1..... that
Pzm, | (12))))))) S Cl,()),
(22)
where the function 4m (2) is defined
T2
2s(x)(n 1)2/2
-x+11/2)
(x(x)
1+ m(x) - s(x)(1 - 1
m(x) + (nz)(n 1)1/2
x max |/.81 A(x) r+1
1
(23)
|8 5m
The following results key to guarantee that SDE is weak approximation an optimizer.
Lemma C.6 (Lemma I Li et al. (2017)). Let < n 1. Consider a stochastic process X,,t 2 0 satisfying the SDE
dx; = b
+ Vio (X:)
with Xo E R and b, a together with their derivatives belong G. Define the one-step difference A = X, x, and indicate the i-th component of A with D, Then we have
1. EA,, byn ES= b,0,, 1/ + ( Vi 1, d;
2. KA,A, = [b,b, + Too (m)) 1/ 0
Vi, = 1.....d: d; for all 2 3,
3. IE I11;; AG, = O for
1,...,d.
All functions above are evaluated x.
23

---

Published as a conference paper at ICLR 2025
Theorem C.7 (Theorem and Lemma Mil'sh (1986)), Let Assumption C.3 hold and let us define Á I to be the increment in the discrete-time algorithm, and indicate the i-th component of with A. If in addition there exists K1, K2, K3, K4 G so that
1. EA: | & Ki (x)),, Vi -.....,,d;
2. EA,4, E& ر S K2(x)n²
:..=.....,.6
3. IIT A,, E IT Aiy $ <z()))
Vs Vs Vi, (I.....,d) d]:
4. -1 ( Vi, E fl,...,,. d}.
Then, there exists constant C' so that for all k = 0, l..... N we have
[Eg (Kkn) - Eg (xx) S Cn.
C.3 LIMITATIONS
Modeling discrete-time algorithms using SDEs relies on Assumption C.3. As noted by Li et al. (2021), the approximation can fail when the stepsize is large or if certain conditions on Vf 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.
C.4 FORMAL DERIVATION SIGNSGD
In this subsection, provide the first formal derivation of an SDE model for SignSGD. Let us consider the stochastic process X, E IRd defined as the solution of
dX, -1-
-11-2P1 (L1-P(V55(X)) < 0t) + Viy DdX
(24)
where
E(x) = (x)e
(25)
and E, (x) == sign(V - + 2 < 0) the noise in the sample sign (v))))). The following theorem guarantees that such process 1-order SDE of the discrete-time algorithm of SignSGD
Fa+1 2 7 sign (Srr (xx)),
(26)
with zo E Rd, 77 E R>0 the step size, the bathees }nw} are modelled i.i.d. random variables uniformly distributed on {1. .N, and of size B 2 1.
Before proceeding, we ensure that the SDE admits unique solution and that its coefficients are sufficiently regular.
Lemma C.8. The drift term b(x) - (1 2P((f.(x) < 0)) is Lipschitz, satisfies affine growth, and belongs to the space G together with its derivatives.
Proof. Since we are assuming that the gradient noise has smooth and bounded probability den- sity function," ? the drift can be rewritten in terms of the CDF F2(x) of the noise as b(x) (F2()v)(z) - 1, whose derivative (-F the Hessian of f are bounded, we conclude that the derivative is bounded. Therefore, the drift is Lipschitz and as regular as ff, meaning that each entry is in G, together with its derivatives. Finally, since it is bounded, it has affine growth.
Since the density function and
Lemma C.9. The diffiusion coefficient VE satisfies the affine growth condition.
Proof. Since it bounded, the result follows immediately.
"This is commonly assumed in the literature. Among others, Ahn et al. (2012); Chen et al. (2014); Mandt al. (2016); Stephan al. (2017); Zhu et al. (2019); Wu et al. (2020); Xie et al. (2021) assume that it is Gaussian, while Jastrzehski et al. (2018) offers intuitive justification.
24

---

Published as a conference paper at ICLR 2025
Lemma C.10. Let assume the same assumptions Lemmna C.4. Additionally, assume that [g ((x,))] -) E (I")))
sup =EK
for all compact sets K. Then the entries of E. in Eg. 25 are in C% (RR)).
Proof. By the definition of Σ in terms of the sign-function and dominated convergence, from the additional assumption on it follows that Σ is continuous. So Lemma C.4 entails that the entries of E. are Cg (TR*).
Lemma C.11. Under the assumption that
g(z,8) v) > 0,
(27)
the covariance matrix E is positive definite.
Proof. For y = (y1,...,,1) yn) observe that
n S (z) (
(E()),,)) Σ y.E (E ys =
Using the definition of &- and the positivity of the density J. we can argue by contradiction and see that for y & O, the right-hand side of the equation must strictly greater than zero for all r. Therefore, (x) E S for all where & denotes the open set of positive definite matrices in the space of symmetric n X n matrices.
Corollary C.12. Since E is positive definite and its entries are in C7 (R"), VEE is Lipschitz.
Proof. The function
-:S+S, A
has Fréchet derivatives of any order on 8 (see egg. Del Moral and Niclas (2018)). Therefore, E2/² E (*R())), and since E E C7 (R"),, 122 is Lipschitz continuous (see Proposition 6.2 in Ikeda and Watanabe (2014)).
Proposition C.13. Assume the conditions of Lemma C.5 and assume that the functions Im (x) for m == l......, 7 in Eq. (18) are of polynomial growth. Then 1//2 E G together with its derivatives. Corollary C.14. If the noise Z(x) -N(0,, ((x) L, (0, then I./2 G together with its derivatives.
Proof. With the definition of [xx) given in Lemma 3.10, the function K(x) == 4 is in G' together with its derivative: It is easy to verify that all the derivatives of K ((x) are bounded even in the case v = 1, which the most pathological one. Therefore, vE(()) is G together with its derivatives.
Remark C.15. Based on the above results, we have that under mild assumptions on the noise structures (see Sec. C.1) that cover and generalize the well-accepted Gaussianity, e-g. covering Student's as well, the SDE of SignSGD admits a unique solution and its coefficients are regular enough to apply Lemma C.6 and Thm. C.7.
Theorem C.16 (Stochastic modified equations). Let 0 < <<<<,, 0 and set N = [T/n]- Let It Rdd 0 < < N denote sequence of SignSGD iterations defined by Eq. 26. Consider the stochastic process X; defined in Eg. 24 and fix some test function E G and suppose that g and its partial derivatives up to order 6 belong G. Then, under Assumption C.3, there exists a constant C 0 independent of such that for all 0.1.....NN we have
[Eg (Xx,) - Eg (2x)) S Cn-
That is, the SDE 24 is an order 1 weak approximation of the SignSGD iterations 26.
25

---

Published as a conference paper at ICLR 2025
Lemma C.17. Under the assumptions of Theorem C.16, ler 0 r 1 and consider Ik. k 2 O satisfying the SignSGD iterations
Tk+1 = Ik siiign (V (2x))
with 20 E IRd, From the definition the one-step difference Á = 21 then we have
1. EA,-- (1-2P < 0)) Vi ,..... d;
2. EA = ((1 ((1-2IP 2IP (0.fa < 0)) (1 2P (0, < 0)) + E(n)) n/ Vi.j =
1,...,,h d;
3. E =1; (n²) 3, in {l.....d]
All the finctions above are evaluated at a.
Proof of Lemma C.17. First of all, we have that by definition
E
x²] -1 (x)))
(28)
which implics
EA, = - (1 2P
< (0))
vi= 1....,d. d.
(29)
Second, we have that by definition
E (xn 2)) (x1 =)=]
=E [1x - x)) E (11 x;) x
(30)
(sign (Vf((z)) 2IP (x) <0)
(31)
(sign (v/(()) 1+ 2P (V/(2) <0)7T(??
(32)
which implics that
KA,A, = (1- 0)) (1 2P (0, < < + Vi,j I,...,d.
(33)
Finally, by definition
E
A, = 011²)
V, i; (1......
(34)
which concludes our proof.
Proof of Theorem C.16. To prove this result, all we need to do is check the conditions in Theorem C.7. As we apply Lemma C.6, we make the following choices:
b(z) (11 2P (z) < 0)):
(x) = (E()).
First all, we notice that Vi 1,...,d. d, holds that
IE Lemma C.17
(1 - 21P (0.f4 < <)) m
EA Lemma C.6
(1 2P (z. (
26

---

Published as a conference paper at ICLR 2025
Therefore, have that for some K1(x) G.
IEA; EA - Ki()) =i
(35)
Additionally, notice that Vi,j 1,..., d. it holds that
E Lemma C.17
-1-23 2IP (1-2IP
< < (1- -1- (0)()))))))))))))))))))))))))))) <0)) -(t) +
J, O (17²)
((1 20-2P <0) (1-2)))))))111))1111111))11))))1 (P) <0) [(n)(t)) o²
Therefore, have that for some K2(x) G.
EA,A, EA,A, & K2(z)n²,
= = L...., d.
(36)
Additionally, we notice that Vs 2 Vi E {1.... d), it holds that
E [[; D, Lemma 0 (1²):
E [T 1T==14 Ai
C.6
0 (17²).
Therefore, we have that for some Kz(x)= E G,
A
E
A.
I K3(x))
(37)
j=1
j=1
Additionally, for some (1)x) E G, Viz E (L.....d}.
E
A A
Ka(z)n²
(38)
Remark C.18. Remembering Remark C.15, and thanks to Eq. 35, Eq. 36, Eq. 37, and Eq. 38, the thesis follows from Lemma C.6 and Thm. C.7.
In all the following results, the reader will notice that all the drifts, diffusion terms, and noise assumptions are selected to guarantee that the SDE we derived for SignSGD is indeed a weak approximation for SignSGD even without the mollification argument used Li et al. (2019) to handle the regularity issues.
Corollary C.19, Let take the same assumptions of Theorem C.16, and that the stochastic gradient is V fz(()) = fxx) + Z such that (0, Σ) that does not depend on x. Then, the following SDE provides 1 weak approximation of the discrete update of SignSGD
Σ 1 (x()))
Σ
Vf ))
dx. = -Frf
di V La diag
Erf
dW,, (39)
/2
V2
where the error function Erf(z) and the square applied component-wise, and Σ diag ..)).
Proof of Corollary C919. First of all, observe that
1 1 2P (V
0) = 1 - 2P
+
Σ 0))==))=110 -1-24 (-E-AVV(()). (40)
27

---

Published as a conference paper at ICLR 2025
where @ is the cumulative distribution function of the standardized normal distribution. Remembering that
(1
Φ(x) =
1 Erf
(41)
V2
we have that
1
Σ 3 V
))
1 - 2P (V < 0) = 11 2/ 1 Erf
= Erf
(42)
V2
v2
Similarly, one can prove that defined in 25 becomes
Σ x
Σ Ia dddia
Erf
(43)
V2
Corollary C.20. Let us take the same assumptions of Theorem C. 16, and that the stochastic gradient (z) = V(()) + VEZ such that Z (0, Ia) that does not depend on x and a positive integer number. Then, the following SDE provides 1 weak approximation of the discrete update of SignSGD
dx, -2 22 (=+vV((x) di + vil la diag (s (S(2-AV/(X:)) dW,, (44)
where E(x) is defined as
2+ ( Vavl
1 V v
x²
F
(45)
2
2
and 2 (a, b; x; ) is the hypergeometric function. funttiiiottiiccinioonn. Above, function E(x) and the square are applied component-wise, and Σ diag (ot, ... ,
Proof.
First of all, we observe that
+ 2P (z) < 0) = (Vf()) < -1 -
(46)
where F, (x) is the cumulative function a distribution with degrees of freedom. Remembering that
F. (x) = (xx()),
(47)
we have that
1-2P <0) = = 1-2 { E C- $ {(())) = (V
(48)
Similarly, one can prove that becomes
= Ia 4diag E, (
(49)
x- (-44)((,)
Lemma C.21. Under the assumptions of Corollary C. 19 and signal-to signal-to-noise ratio Y;:=
v2
1. Phase 1: 1f > 2. the SDE coincides with the ODE of SignGD:
dX, = Si (X)())))d;
(50)
2. Phase 2: 1f1 < [Y] < 3.
28

---

Published as a conference paper at ICLR 2025
(a) my q-= dE [
I mY; + qt:
(b) P [1]X - E [ 2 > a] S (d |mY + q l1):
3. Phase 3: | the SDE is
2 Mー (xx,)d + voy la diag [ fx))) dW, (51)
2
dX,, = -
Proof of Lemma C.21. Exploiting the regularity of the Erf function, we approximate the SDE in Eq. 39 in three different regions:
x: (x)(
1. Phase 1: If
[o|| > 3. Brf(x)) sagn(ar). Therefore, if
v2
((x)))
(a) Erf
sign
= sign (Vf(X,));
V2
v2
Σ ((x.)
(Vx)))))
(b) Erf
sign
(1,...;1). 1).
V2
V2
Therefore,
2
Σ Σ (xxx))
Σ 2
dX, Errf
dt + vT la - diag
Erf
dW,
V2
V2
sign ((S(())))))
(52)
2. Phase 2:
Let m and qı are the slope and intercept of the line secant to the graph of Erf(z)
between the points (1, Erf()) and ( Erf (2)). while Q2 is the intercept of the line tangent the graph of Er)( and slope m. If
S we have that
ma
Erf( (z) + 92-
(53)
Analogously, if - <
mæ 02 Erf(x) < m. Q1.
(54)
x- x-(b)(X,)
Therefore, we have that if 1
// then
V2
(a)
Σ fVX))
m Σ
fXx,
Erf
- V22
m M
+ +
(55)
2
V2
where
(q+);
92
if0.f(z)>0
(56)
1
if0.f(z) <0,
and
91 ifa,(x))
(qT); =
(57)
-92 if <0,
Therefore,
f((,) qt <
dE [X]
m M V - --: (58)
V2
dl
/2
(b) Similar to the above,
< Erf (=())(()))
m x x
+ +
m 5 1 ((), + +a*)
<
2
22
2
29

---

Published as a conference paper at ICLR 2025
Therefore,
P [I][X - BE > a] S IP s.t.
xx; - [X]] | > a]
(59)
x - B [ > va
Σ x,,0,((11)
7
Σ
1 Erf
(60)
v2
7
m
D
(xx)) +
12
(61)
V2
3. Phase 3: If |x| < 1, Brf(x)
2 Therefore, if
E ((x))) v2
<1.
Va
[v((x))
(a)
Erf
ト fXx,;;
V2
( v2
< (I-AvS(x.)²
(b)
Erf
2/4
Therefore,
Σ V
(=-VV((X))
Σ
dx, Erf
dp + vil - diag
-dddiag
Erf
V2
V2
S(vV)((X)d + wTy la
2
diag
(S=AV)((.))
2
aW.
(62)
π
Lemma C.22 (Dynamics of Expected Loss). Let be p-strongly convex, Tr(V2f(x)) s C., and S = f(X.) - f(X.). Then, during
1. Phase 1, the dynamics will stop before l. = 2/ So because S. S 1 3 (V - 2//50)² ?
2. Phase 2 with A : 12omm
mum 10
E(Si < Sye Buat + (Er-udi")) (1 c 2w))))):
241
(V²
3. Phase with A :=
2
E[S
Soe 2pAt
Le (1 - E 200r).
2
Proof of Lemma C.22. We prove each point by leveraging the shape of the law of X₂ derived in Lemma C.21:
1. Phase 1:
d(f(X))- - f(x,)) = (ixX = -VV S -IV
(63)
Since f is p. - PL, we have that - -VV < 22(((((X)) - f(X.)), which implies that
f(x.) - f(x.) < 4 -V
2/ - -((x)) -(X))
(64)
meaning that the dynamics will stop hefore l., = 2/ -(x))-(X.).
2. Phase 2: By applying the Itô Lemma to f(X,) f(x.) and that
m D - v-x,)( + - < Erf Σ
(X()))
m -ー
+
(65)
V2
22
22
30

---

Published as a conference paper at ICLR 2025
we have that if ;= max(g1:
a(f(xi) - f(x))) -
m
E +Vf(X + V + O(Noiee)
(66)
V2
2
+
2
! Tr
v2s(x))
la
diag m (=- + +
V2 (/2-
di
(67)
m 1
+ n dt
(68)
V2 o max
77 I1 m
E 2 Vf(( + q 12 + (Noise)
(69)
2 V2
m
[F(())))d (x))(ll + n
dt
(70)
22 max
2
num
nudg² dt
22m [f(X)) (71)
Vf(x))l
2
2
max
(Noise)
(72)
num 2 (f(X)) - f(X.))dt
m
S 2/
(73)
m0
40 max
+ 2
(C, - uud"") dt + ONoise),
(74)
which implies that if k == 2p
+ num=
(2mmm
1o?
( - S ((x.))e + (Cr udgz))
(1 kk-
(75)
2k
3. Phase 3: By applying the Itô Lemma to f(Xe) f(X.), we have that:
d()(x.) f(x.)) = V 22
(2
V + O
(76)
"1 Tr
2
Σ )ン)V)())) ") >)
La diag
dt (77)
2 IV + (NNiie)
(78)
m
" Tr
? 1 (V2 f(x)) dt
(V
(79)
π o
2
+ 7 P
[V
(80)
T Tmax
π o 2
! ((()))((x))dd O(Noise)
(81)
Since f is u-Strongly Convex, f is also p-PL. Therefore, we have
(v
2 7 P
(U(X))-((X)) S 2p
(((x))- I(X)))d
(82)
π o
(V2 f(X.)) ONoise).
(83)
Therefore,
- [E] S -2p
+
(E(F(X)) - f(x.)))dt 17
2/2 C,
T Tmax
(84)
31

---

Published as a conference paper at ICLR 2025
which implies that k := 2p
(1222..
+ 1 24
1TCr
BIf(X) (
2k (1 (1-E* c.).
(85)
We can weaken the regularity of f from p-strongly convex to μ-PL: This results in less tight bounds as expected.
Lemma C.23 (Dynamies of Expected Loss). Let f be p-Pl, I-smooth, and S. == (xx) ((X.)) Then, during
1. Phase 1, the dynamics will stop before = =2 2 because Se S 1 (Viit 2255))
2. Phase with A
T3. [8.) S Soe 2pdr
2nde
yld + TnB
(1 e. 222d)):
3. Phase with A
:: Taass
B[S S SOE
-2
(1 E -2201).. 4pA
yLd
Proof of Lemma C.22. We prove each point by leveraging the shape of the law of X, derived in
Lemma C.21:
1. Phase 1:
-110)///1//// d(f(X,) f(X)) V/ -IV - S - laat..
(86)
Since f is F1 PL, we have that - (V < -2((((X)) f(X.)), which implies that
1 s(x,) [(x.) - 4 (Vil (V - 1(x.))
- 2
(87)
meaning that the dynamics will stop before to = 2// [(Xo)-
2. Phase 2: By applying the Itô Lemma to f(X)) s(X.) and that
v2 Σ (x.()) + q < Erf Σ \
m
--iv +
(88)
v2
22
we have that if q ; max 42),
Σ 1 (V((x))=a =a-) T
- 5
(/22
Vf(())dtt OO(oiis))
(89)
m Σ --=vV + +)) dt
+ ?l Tr V2j(Xi)
La Ia-diag
2
-
(90)
s m
[Vf(X,) - fxx, || dt + ?Ld
2 dt
(91)
max
< 2p V 20 mux
m
(f(X,) - f(x.)) + yLd
dt ONoise),
(92)
which implies that if
m
V20mn
Elf(x,) f(x.)) E (f(Xo) - ((x.))) nLd (1- c 220d).
(93)
4nA
32

---

Published as a conference paper at ICLR 2025
3. Phase 3: By applying the Itô Lemma to f(Xi) f(X,), have that:
2 vs(x,) Σ fX O(Noise) + 7Ld
( -
2 dt
(94)
2 [ 12at + O(Noise) + nLd
2 dt
(95)
π max
Since f is p-PL, we have
d(s(x)) - f(X.)) - - 2 __(x) (J(X,) - f(X.))dt + yLd
dt + (Noise). (96)
0max
2
Therefore, for A : v 2
Omn
(lf(X,) f(x.)l (f(Xo) f(x.))) 2pddr 7
4nA (1-e
(97)
We can weaken the regularity of f from μ-PL to L-Smooth: Of course, we can only bound the expected norm of the gradient.
Lemma C.24 (Dynamics of Expected Gradient Norm). Let f be L-smooth, Ti be learning rate yds. scheduler such that lom oi 0 and of = oc, where of (n)) ds. Then, during
1. Phase 1. V
((X))-((X.) t 0:
2. Phase 2,
m V2
V f X X 112 mir E V (Xga.a)) )
S 0mar [((o)) - -f(x))
+
nyd oi
2330 O;
6
2 0
3. Phase 3. E V (Xp) / Ld of of / max f(Xo) 10:00 0:
where i, g0,2), (222), and f are random times with distribution
Proof of Lemma C.22. We prove each point by leveraging the shape of the law of X, derived in Lemma C.21:
Phase 1: ([((x))-(X..) T V/VVV((X) o1 1/
IVX((())
(98)
(XX)) di
(99)
Therefore, by integrating over time and using the law of the unconscious statistician
f(Xo) - f(x)) 20 0:
VV (X₂))li
(100)
2. Phase 2: By applying the Itô Lemma to f(Xi) f(X.) and that
fX())
m Σ vXXi))
. Erf
m Σー v-(()(
+ +
(101)
2
V2
Similar to what we have shown above, we have that
d(f(X.) f(x.))
mell
(102)
/2 max
Ld
+1 dl O(Noise)
(103)
33

---

Published as a conference paper at ICLR 2025
Therefore, by integrating over time and using the law of the unconscious statistician we have
v2 B V (x B (X S mmax
((X))-
f(x.) + ?Idd?
(0 O;
Φ
2
(104)
3. Phase 3: By applying the Itô Lemma to f(Xi) f(X,), ), have that:
- - f(x.)) - 221
mp ((X, Eat + O(Noise) + nn 2
dt (105)
π max
Therefore, by integrating over time and using the law of the unconscious statistician we have
T max oi
f(Xo) -f(x))
(X V
+ 2 max
(106)
Lemma C.25. Under the assumptions of Lemma 3.5, for any step size scheduler ni such that
20
nyds = oc and lim 1): = - E[f(x,) - f(x.) '230 0.
(107)
t
Proof of Lemma C.25. For any scheduler Ti used in
2k+1 It
mys sign (frx (zx))
(108)
the SDE of Phase 3 is
dX. 2 ST-D)((X))mdl + Vime V la diag ( (Σ-tv)(X.)) dW.
2
2 π
(109)
Therefore, analogously to the calculations Lemma C.22, we have that
20
di
f(Xo) s(x.) xc= So
mm
(110)
So So
onna no + n
Therefore, using l'Hôpital's rule we have that
n.ds = oo and lim me
=0 = B[f(X) f(x.)) O.
(111)
+000
Lemma C.26. Let H = diag()) .A.,,., and Mj == -2(VEs - - Then,
1. E [1] e - Xu:
x + (V H-1 (Ia - Mc).
2. Var [Xi] = -
Proof of Lemma C.26. The proof is banal: The expected value derivation leverages the martingale property the Brownian motion while that of the variance uses the Itō Isomerty.
x. nx.
Lemma C.27. Let H diag(Ar..... Aa). Then, E
is equal to
(x
-2x (V2L 1/1++1 -m7 h1z)
-2h, (VZL
(V=++00))
77
+
1 1-c
(112)
2
4
2
+
A.n
xo
34

---

Published as a conference paper at ICLR 2025
Proof of Lemma C.27. Since the matrix H is diagonal, focus single component. We apply the Itō Lemma to Ax(x;)
(x 2
(xx)²²
nA dl
2 2 xx))2
2
di +
di + O(Nois))
(113)
2
2
which implies that
(x
(x
-2
V
232
11
-2 -2(V
xo?
APu
)
1 1-e
4
V
2 o:
+
xo?
A.n
(114)
Therefore,
X H
xxx)))
22
Aed ma
7
1-e -2 ( +
2
4
2 1 +
Am
Ro
(115)
Lemma C.28. Under the assumptions of Corollary C.20, where V.z(z)) V[(())=V\(()) + VZ, we have that the dynamics of SignSGD in Phase 3 is:
dx, -X 2 v((x)dd vay La diag - (2=tv)(Xi)) dW,.
(116)
Proof of lemma C.28. We apply Eq. 44 with v = 2 and linearly approximate E(x) |ar| < 1, where 23
V2
C.5 ALTERNATIVE NOISE ASSUMPTIONS
In this subsection, report the consequences of assuming different noise structures. We do provide the proofs as they mimic those of Corollary C.19 and Lemma C.21. We validate our results Figure 7.
Assumption from (Ziyin et al., 2021) As per Eq. (16) in Corollary 2 of (Ziyin al., 2021), take Σ = o² fr. V (xx,), where we added the constant a² as parameter to control the scale of the noise and (*..) > 0. Under this assumption, have that for Y, v f(x.) +vx(x) and
V2V
Sa(X,) = E, (signn V f, (Xx) J,, Corollary C.19 becomes:
dX, -rf (Yt) dt + viv Sx((;)) Erf (r) (Y.) dW.
(117)
As consequence, Lemma C.21 becomes:
Lemma C.29, Let f be u-strongly convex, Tr(V2f(x)) < C., Amax the largest eigenvalue of fxx),, and S. = f(X,) - f(x.). Then, during
1. Phase 1, the loss will reach before t, = =2 So because S. S (Vrt (/tt-2//o))²
2. Phase 2 with A
mpmm 4f Yo Areu
E[S.] < -0e-2ppr
/2f(r= A
y
(C.-udg") 2n1
(1 e1---2us1) -2ud));
3. Phase 3 with A = (Vz
B[S.] < -ge-2pd:
)o Nave fix. Jo
3 2p1
(1 22DD)).
35

---

Published as a conference paper at ICLR 2025
Assumption from (Wojtowytsch, 2024) (Wojtowytsch, 2024) discusses two possible assumptions
on Σ: [[x)]| & Cf(2) and ||[()]] Σ Cf(x) + As Section 2.4, they ultimately use
== Ca) Therefore, we take Σ == (=)) where we changed the constant to o² to maintain
f((x))
consistency with the rest of our paper. Under this assumption, we have that for Y. /2f(2)o*
Corollary C.19 becomes:
dx, Erf (Y,) dt + vIv La (Eag (Y,))
(118)
As consequence, Lemma C.21 becomes:
Lemma C.30. Let f be p-strongly convex, (r((2f(x)) & Cr, and S. == f(X,) - f(X.). Then, during
1. Phase 1, the loss will reach 0 before t, =2 So "
because St S (ViitVVt55 -2v5o)
2. Phase with B = z/2 (Er - pudla² m² and a= = 22mm
32 (
(- aa1--2//ua - -1) +1)
B
exp
23
3²
E [
(119)
a²
a
2n² 3. Phase 3 with B = 17 (5r
and a = = 1 2
xa²
3² (
(3-v50a) exp ( a -
- -1)) -11)+11
+1)
t to g²
S
(120)
a22
a
where W is the Lambert W function.
Assumption from (Wu et al., 2022) (Wu et al., 2022) proposes novel structure of as being aligned with the Fisher Information Matrix and proportional to the loss function. Consistently with this, we take > = o² V²= where we changed the constants to o² to maintain consistency
with the rest of our paper. Under this assumption, we have that for Y, == ((')(x))) tVX(x) and
f2V
Sa(X,) = [(sign V (i((X J, Corollary C.19 becomes:
dXx = -Erf (Y)) di + viV (x.)) Erf (Y,) Erf (Y.) dW,
(121)
As consequence, Lemma C.21 becomes:
Lemma C.31. Let be p-strongly convex, L-smooth, Tr(V2f(x)) < C,, and S = ((Xi))f(X.)) Then, during
1. Phase 1, the loss will reach 0 before t, = =2 Sa because S, S 1 (Viit - 22/50)
2. Phase 2 with
= (c+ pulq²
m²p² 0²1
and
v2ma Via
32 (w (5-)) exp
a2412/55aa - 11) +1) -1))+1)
t to g²
[
(122)
a²
a
3. Phase with B = (
24² x
and a = =2/² VIo'
82 (w) (8-vFoa) exp
(- 0 -2 25 - 1) +1)
B
20
g2
E[S <
(123)
a2
where Wi is the Lambert W function.
36

---

Published as a conference paper at ICLR 2025
Losses
Losses
Losses
107
10
10
10"
10r
Loss
10"1
30
10:2
Bouns 27,
Bzond 171.
10-2
Srone 77h
5ge500
Bound
Bound
7hh
10-4
10
10
13
707
10
10
13)
10'
10
lterations
Iterations
Iteratoons
(a) E(x) = oR JZ.
(b) [
(c) E(x) := o² f(x)aa.
Figure 7: Empirical validation of the bounds for SignSGD derived for the noise structure:
[(x) = (=. (left), E(z) := o² f(x)V² (center), M(x) := x² f(x) (right). Each empirical validation has been carried out on strongly convex quadratic loss functions. The experiments are averaged over 500 runs.
Assumption from (Paquette al., 2024) In this section, we adopt the notation of (Paquette et al., 2024). As per Eq. 5 of (Paquette et al., 2024), the loss function can be rewritten as
f(0) マ (D(WO- 6), (WO b)), where D diag (j (j-20) E 1x"
(124)
Without loss of generality, we define @ == W0 b, which implies that
Φ ToDo
Dá 2
f(0)
where diag (j-2o) E R where 115
(125)
The stochastic gradient is unbiased and its covariance is the well-known BE(Ø) = (0 T D D
DØo D = 2f( V(()) where B is the batch size. Under this assumption, we have that for Y, : )B fvlle) and S(ø.) (L Corollary
V2
C.19 becomes:
dΦ. -Ers + viv S(i+) Ers (Yr) Eri dw.
(126)
As consequence, Lemma C.21 becomes:
Lemma C.32. Let f be as above. fo := f(t)), and C, == Tr(D). Let p be the mininum eigenvalue of D. and L be its maximum one. Then, during
1. Phase I, the loss will reach 0 before 22
Sa
because Sa S 1 (Vrl 2/To)²
a
2. Phase 2 with B = Cr and a MHVB V2L
32 (
(f+ /ooa)
a²t-2v/na - -1)) + +1)² 23
exp
g²
S
to
(127)
a²
a
3. Phase with B = Cr and a: ==:/:"//00 N2
VI.
32 (
(I- (oa) exp
(- - 25 Sha
1)) +1)
S
28
t
g2
(128)
a.
a
where W is the Lambert W function.
See Figure 8 for empirical validation.
37

---

Published as a conference paper at ICLR 2025
Losses
Losses
100
109
10-1
10
Loss
Loss
10
10-2
10-3
S SgD
10-²
SignSGD
Bound (Ph.
Bound (Ph.
Bound (Ph,
Bound (Ph. 2)
10
Bound (Ph.
100
16
102
10a
10
10
10a
10a
Iterations
Iterations
(a) 5, d 50.
(b) 50, d 500.
Figure 8: Empirical validation of the bounds for SignSGD derived in Lemma C.32: In both experi- ments, a 0.25, 3 2, 17 = 0.001, B 256, N 10000, and trajectories are averaged over 500
runs.
C.6 FORMAL DERIVATION RMSPROP
Trajectories
Trajectories
1.50
LC-
6.6
10-1
1-h
-0.0
11-1
-15
50C Oer
200 DONTER
x2 100 100 Iterations
Figure 9: The first two subfigures on the left compare our SDE, that from Malladi et al. (2022), and RMSprop in terms of trajectories and (ar), respectively, for convex quadratic function. The other subfigures do the same for an embedded saddle, and one can observe that our SDE tracks R.Mprrpp more faithfully.
In this subsection, we provide our formal derivation of an SDE model for RMSprop. Let us consider the stochastic process I, == (X,, Vi) Rdd x Rdd defined as the solution of
ax, - PP (Vf(X,) + (TE(X))a aV. = (x())(xx)) + diag E(X))) - dt))
(129)
(130)
where 3 = 1- npp, p= 011), and := diag (V.)² + eTa-
Remark C.33. 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 extremely relevant when the gradient size is not negligible. Figure 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 tracks the algorithm more faithfully. Importantly, one can observe that the SDE derived in (Malladi et al., 2022) is only slightly less accurate than ours at the end of the dynamics: As we show in Lemma C.37, Theorem D.2 is a corollary of Theorem C.34 when fz(x) O( n e.g. it only describes the dynamics where the gradient is vanishing. In Figure 10, we compare the two SDEs in question with RMSprop on an MLP, CNN, ResNet, and Transformer: Our SDE exhibits superior description of the dynamics.
38

---

Published as a conference paper at ICLR 2025
DNN
CNN
LosseS- RosNet
0.7
10:80
(han)
201
(0mm)
0.2
0.2
0.
250
:5e U000 1530.750 tteraions
Rerations
Iterations
Iterations
Figure 10: We compare our SDE, that from Malladi et al. (2022), and RMSprop in terms of f(x): The first is an MIP on the Breast Cancer dataset, the second a CNN on MNIST, the third a Transformer on MNIST, and the last ResNet on CIFAR-10: One can observe that SDE tracks RMSprop more faithfully.
The following theorem guarantees that such process is 1-order SDE of the discrete-time algorithm of RMSprop
V Sn (xx)
Ik+1 IK 17 /ut+1 + eld
(131)
Uk+1 Bor + (1 - 3) (V (xx))²
(132)
with (zo, vo) E Rd Rd, R*0 is the step size, = ln for p 0(),, the mini-batches (7%} are modelled as i.i.d. random variables uniformly distributed on fl.... .N., and of size 2 1.
Theorem C.34 (Stochastic modified equations). Let 0 < < 1,T> 0 and set N [T/nl-
Let 15 == (Ik, U&) Rdd Rd O S k S N denote a sequence of RMSprop iterations defined by Eq. 131. Consider the stochastic process Lt defined in Eq. 129 and fix some test function g E G and suppose that g and its partial derivatives up to order belong to Then, under Assumption C.3 and p = O(1) there exists constant C > 0 independent of such that for all = 0, 0,1,..., N, we have
[gg (Lx,) - Eg (I+) S Cn--
That is, the SDE 129 is an order weak approximation of the RMSprop iterations 131.
Proof: The proof is virtually identical to that of Theorem C. Therefore, we only report the key steps necessary to conclude the thesis. First of all, we observe that since 3 = Tp
uk+1 2k -np (vx
(1a)))
(133)
Then,
Uk 1
Uk+1 O(n)
O uk+1
- Uk
~
(1)))
(134)
uk++1
Uk++1
Uk
uk+1
Uk
Therefore, we work with the following algorithm as all the approximations below only carry an additional error of order On?2), which we can ignore. Therefore, we have that
V ffz T (Ix)
1k+1 Ih =
(135)
Vuh
ela
1s uk- -n Uk=1 (V (Tx-
(136)
Therefore, if VS, (xy) = Vf(x,) (T,)) ( = 0, and Cou Zz (11)) E(x,)
1. [[ - xz,] 1リ diag(u, (y) (raa -
2. EE (Vk vR- =1 [V 1)) diag(E(xx)) ) Uk
39

---

Published as a conference paper at ICLR 2025
Then, we have that if 4x = Vuutela
V((x))
V. (8x) Vok+ela
1.
2x)² [[(x++1 [xx) xx)]
(137)
+ n (0x) (1%)
(138)
B[ - mx)) - xx)1
-
(139)
+ rf (diag((v)) + ela) [xxxx)
(140)
2. El( (k-1) 1k-1) = B((u - 0k-1) - uk-1)] + (pp²²)
3. [[( - xx)(1E - U/= 1 B[(x+++ - x&) [x)] - Uk ok- 1
Remark C.35, Let us remember that by assumption, (x(()) and vE(x) are Lipschitz, grow at most affinely, and are in G together with their derivative. Therefore, the drift and diffusion terms of the SDE governing X, are the ratio between regular functions and a uniformly lower bounded process. Therefore, they are turn regular, modulo dividing by Vv. + &f + 6 s.t. & ~ O rather than by VV. + (See Bock and Weiß (2021) as they experimentally verify that this has no impact on the performance of the optimizer). Regarding the ODE governing V,, Σ(X,) is Lipschitz because (Xx()) bounded and Lipschitz. Additionally, is smooth, and with affine growth. On top of this, we need the term (x))² to be Lipschitz and of affine growth, which is consequence of assuming bounded gradients as often done in the literature on the convergence of RMSprop and Adam: Among many, see (Luo et al., 2019; Défossez et al., 2022: Guo et al., 2021: Huang et al., 202 I) together with the discussion Section 2.1 of Shi and Li (2021). Alternatively, exactly as done in Theorem al. (2019), one regularize the drifts and the diffusion terms with mollifiers on sufficiently large compact Reddi et al. (2018), which automatically implies that drift and diffusion coefficients satisfy all necessary regularity conditions. Importantly, needs to then send the mollification parameter conclude our statement. Therefore, we the SDE of RMSprop for Pa diag (V.) 3 + ela is
Li
dx,- -.---- =T.((.))dd))) dV. = p(((VF((x))) + diag(E(Xx)) dt)
(141)
(142)
L
Remark C.36. In all the following results, the reader will notice that all the drifts, diffusion terms, and noise assumptions are selected to guarantee that the SDE we derived for RMSprop is indeed 1 weak approximation for RMSprop even without the mollification argument. Importantly, our analysis of RMSprop focuses on its behavior convergence, i.c. (Vf(z))² O(1). Therefore, there is no need to assume bounded gradients or compact domain.
Lemma C.37. If (V)))) = (T)) Theorem D.2 is a Corollary of Theorem C.34.
Proof. In the proof of Theorem C.34, one drops the term 17 (=))) is of order 172.
Corollary C.38. Under the assumptions of Theorem C.34 with (xx) = d"La. ij = kn, B B8, and
~= = ap.
-VV((X))dt (-V)))))>
1
dx, diigg(V))
VoV B TadW;
(143)
dV. ap (1FF(X)))
0² - Vi
dt.
(144)
B8
Lemma C.39 (Scaling Rule at Convergence). Under the assumptions of Corollary C.38, f is p- strongly convex, Tr Tr(V2 f(x)) & L., and (x))) O(1), the asymptotic dynamics of the iterates of RMSprop satisfies the classic scaling rule K. J because
C, K
B[f(X,) - f(x.)
(145)
4pyBB V8
40

---

Published as a conference paper at ICLR 2025
By enforcing that the speed of Vi matches that of Xt. needs kp, which implies 3 1 x(1 3).
Proof of Lemma C.39. In order recover the scaling of we enforce that the rate at which V, converges to its limit matches the speed of Xx: We need = kp, which recovers the novel scaling = 111 x(1 - 3). Additionally, since (V fxx)) = O()) we have that
1 7 V8 Vo B oladWi)
dX; diaag(V))
-Vf(())dtt
(146)
dV. xp (
02
V
dt.
(147)
B8
Therefore, V, 20 d² 1, meaning that under these conditions:
B5
dX,
d( + n/nlaa
(148)
which satisfies the following for j-strongly convex functions
BB dE[f(X.) - f(X.)(() - 2xp B[f(X.) - f(X.))dt + xo
dt,
(149)
2
14000 no
meaning that (lf(X,) f(X.))
4uvB V8.
Since the asymptotic the loss is y L.a does not depend on and ō if v8 1, we recover the classic scaling rule.
2nvB v8
Remark:
Under the same conditions, SGD satisfies
7 V8 ooldd.
dX. -
(150)
and therefore
1 02 1
B[f(X)) f(x,) & (f(X)) - f(X.))e 2unt
+
2 2uB 8 (1- -unn)),
(151)
meaning that asymptotically the loss is n C. is which does not depend on k. and 8 oh = 2 2yB
Lemma C.40. For f(z) = xJ He the stationary distribution of RMSprop is (E]Xoe]l, (Xo)) (o.grtn11)
2-
Proof. As (Vf(z))² (()) and f co, have
dX, - HX,tt + AnladW,
(152)
which implics that
X, = e bm Xoo v L' x 5 brtHadd
(153)
The thesis follows from the martingale property of Brownian motion and the Itô isometry.
41

---

Published as a conference paper at ICLR 2025
C.7 RMSPROPW
Lasses DNN
Losses CMN
Losses DNN
Losses- CNN
0.7
0.5
Lo00
Loss
0.-
0.2
220
830
125
1928
5000
1000 2038
2000
1283 1033 Iterations
230
20033 aaze score 8000 1065S)
Figure 11: The first two represent the comparison between Adam W and its SDE in terms of f(x). The other two do the same for RMSpropp.. In both cases, the first is 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
(TR) 700x
Fk++1 Ti 1) /uk+1 cTa
(154)
1k Buk + (1 3) (V
(155)
with (zo, vo) E IRd x Idd, 17 E IR is the step size, 8 = 1 - for p ( > 0, the mini-batches } are modelled as i.i.d. random variables uniformly distributed on {1, (1,-... .N), and of size B
Theorem C.41. Under the same assumptions as Theorem C.34, the SDE of RMSpropW is
dx, = - -P- + (7x((X))dW)) - dtt
(156)
dV; ((())(x))) + diag(>((X.)) Vi))
(157)
where 3 1pp
Tpp, p = (1), 0 > 0, and P. diag (V)) ela-
Proof. The proof is the same the of Theorem C.34 and the only difference is that 110r is approximated with d.tdtt.
Figure and Figure validate this result on a variety architectures and datasets. Remark C.42. See Remark C.35 and Remark C.36 for discussion on the regularity of the SDE derived in Theorem C.41.
Corollary C.43. Under the assumptions of Theorem C.41 with (xx) = o² 1j = km, B Bö, and P = ap, and 0 = g0,
dx, <diig(V.)) (-VS(X,)d 1 x² dV. = ap ((TF(X))) + Bö" 1-1-V Vo
V8 B Boladw.) - 50nx,d
(158)
-
dt.
(159)
Lemma C.44 (Scaling Rule at Convergence). Under the assumptions of Corollary C.43, f is p- strongly convex and L-smooth, f())) < L,, X. = 0, and (V)(z))² = O(1), the asymptotic dynamics of the iterates of RMSpropW satisfies the novel scaling rule if k = V8 and & = k because
Elf(X.) f(X.)) 1+000 nL,LL
(160)
2
2μv + (E
By enforcing that the speed of V, matches that of X,. one needs p = ki, which implies 3 1 x(1 B).
42

---

Published as a conference paper at ICLR 2025
Trajectories ries
Trajectoris
04-
Hald
Dud
07
X1
Iterations
X;
Rerations
Figure 12: The first two the left compare our SDE, that from Malladi al. (2022), and Adam in terms of trajectories and [(xr,, respectively, for a convex quadratic function. The others do the same for embedded saddle: Our SDE is shown here to track Adam more faithfully.
Proof of Lemma C.44. In order to recover the scaling of B3, we enforce that the rate at which Vi converges its limit matches the speed of Xx: We need " = kp, which recovers the novel scaling x(1 - 8). Additionally, since (V (V)(x))² = =O(1) we have that
V8 V KE0X,dt
1 /
dX, diag(V,)
-Vf(X
(161)
dV. up ( 0² Vi
[B8 1 dt.
(162)
Therefore, V, 20 x² B5 1, meaning that under these conditions:
BBB V((x)) + kv/iladW, kç0X,dl,
dx,
(163)
a
which satisfies the following for j-strongly convex and L-smooth functions
VB
Eo (1+ (1) ((S(X)) ((x.))dt + x2nR, dl, (164)
dE| - S(X))) < x
2μ
2
meaning that [()(X,) -(cs))-J((XX. f(x.))
1400
S
yC, gb
2py P
Since the asymptotic the loss NC-oL
does not depend and and & if E
2pv BoL
V8, we recover the novel scaling rule.
A similar result be derived for general X.: The final expression is very convoluted and brings marginally negligible added value.
Lemma C.45. For f(x) = x7 Hx the stationary distribution of RMSprop is (E[[ Cou(Xoo)) = (0, 3H2 +11a)-1)
Proof. As (Vf(z))² O()) and + 00, have
dX, = -ーー HX,dt + ATladWe - 0x,dt
(165)
which implies that
X, x, = -E Httriae (Xo+ vn v c +
(166)
The thesis follows from the martingale property of Brownian motion and the Itô isometry.
43

---

Published as a conference paper at ICLR 2025
C.8 FORMAL DERIVATION ADAM
In this subsection, we provide our formal derivation of an SDE model for Adam. Let us consider the stochastic process b == (X,, Mc. Vi)= E IR x IRd x IR defined as the solution of
dX, Vi2(t) a((1)
PT (M, NPI (X, - M.))
(167)
dM. 01 (Vf - Mt) v (X,)
(168)
dV. P2 ((VI(X))) + +diag (X)). VV) dl,
(169)
where B, 1 7ps,, (t) P. diag v cv/2(t/)
ptt, P1 Onn s.t. 5 (0, 1), p2 O(1). to. and
Remark C.46. The terms in purple and in blue are the two differences w.r.t. that of (Malladi et al., 2022) which reported in Theorem D.5. The first appears because we assume realistic values of 31 while the second appears because we allow the gradient size to be non-negligible, For two simple landscapes, Figure 12 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 non-convex case, it converges to different local minimum than ours and of Adam. Finally, Theorem D.5 is a corollary of ours when (Vf()))² ( and #1 = O(1): It only describes the dynamics where the gradient to noise ratio is vanishing and only for unrealistic values of 31 = 7/1- In Figure 13, we compare the dynamics of our SDE, that of Malladi et al. (2022), and Adam an MLP, a CNN, ResNet, and Transformer. One can observe that SDE captures the Adam's 's dynamics more accurately. Details on these experiments are in Appendix F.
The following theorem guarantees that such process is 1-order SDE of the discrete-time algorithm of Adam
Mk++ = 820k + (1 (2) V Sr (rx))²
(170)
m = Bimx + (1 (1 (Fx)
(171)
mg m& 1)a))* -1
(172)
dy Uj (/2
(173)
m
Ik+1 Ts "I V /btl + ela
(174)
with (won mp, vo) E Rdd x Rdd x Rd, E R >00 is the step size, B, = pid for (1)-5) s.t.
<E (0,1), 1), P2 (1), the mini-batche {Tk} are modelled as i..d. random variables uniformly distributed on {1, (1.-.. NN, and of size B 2 1.
Theorem C.47 (Stochastic modified equations). Let < <,TT,, 0 and set N = [T/nl- Let la == (Ix, my, Un) Rd x Rd x Rd, 2, k N denote a sequence of Adam iterations defined by Eq. 170. Consider the stochastic process La defined in Eq. 167 and fix some test function E and suppose that g and its partial derivatives up to order 6 belong to G. Then, under Assumption C.3 p1 O(n 5) s.t. (0, 1), while P2 O1). there exists constant > 0 independent of 1) such that for all k = 0, .....,,,, N, we have
[Eg (Læn) - Eg (I+) S Cn- That is, the SDE 167 an order 1 weak approximation of the Adam iterations 170 for > to-
Proof. The proof is virtually identical to that of Theorem 1. Therefore, we only report the key steps necessary to conclude the thesis. First of all, observe that since 31 1711
Uk+ Uh
-MP1 (V (xx))))
(175)
44

---

Published as a conference paper at ICLR 2025
DNN
Loscas CWN
Losses Transformer
0.2
sie slind
Locs
Loss
0.4
Rerations
Iterations
[eer lterations
Iterations
Figure 13: We compare our SDE, that from Malladi et al. (2022), and Adam in terms of f(z): The first is an MLP on the Breast Cancer dataset, the second a CNN on MNIST, the third a Transformer on MNIST, and the last ResNet CIFAR-10: One can observe that our SDE matches the algorithm more accurately.
Then,
Uk 1
Uk+1 + O(n)
O Ck+1
I - (1 V
0()))
(176)
Vuk+1
0k++1
Uk
Uk++1
Therefore, we work with the following algorithm as all approximations only earry an additional error of order (v)²), which we can ignore. Therefore, we have that
t/ Uk = np2 (uk-1 (x (xx-1))
(177)
mg+11 me 1111 (m4 Vir (xx))
(178)
my = my (1 - 34)
(179)
6x = Uk (1 81)
(180)
v 1122)
Tk+1 - J4
Ukk + eTa (1 7p1 )a+1 (m² + 7pp (V (xx) - mx)).
(181)
Therefore, if (Z;) = Vf(xx)) + (z1z, and [x₁)( = 0, and Coo (zz())) Ma,)), we have that
1. EE [1 Uk- = 1122 (V (V diag(F(xx)) - 2k
2. Emm++1 mi] = np1 f - mp]
\- (m² + npI (V -ma)).
3. E[ - xxk] = -or 1-
Then, have
1. [x((+11 - Tx TT)T]] [x ax) ax) [[))1 - xx)]
Onn1pp);
2. E[(xx+1 - [m)(m² - -m 1) [] = E[(xx+1 xx)] - mm=1)] +0;
3. E[(Fe+1 [Fx Uk [x(++1 2e [B[ 1x-1)) 0;
4. [[(v 1 (k- || B(( - Uk 1) B[ - T D(mrp2))
5. E[(m: m (m - m [ = El - m m--1) - mm T)] + 7 p
6. [ vk-1) mi T| E v - 017 p)))))))))
45

---

Published as a conference paper at ICLR 2025
Since in real-world applications, P1 () s.t. < (0, while P2 (1), have
dX, Viz(t)
(t(t) Pt (M, + np1 (Vf - Mi)
(182)
dM, Pr (Vf (X,) - M,) dt v (X,) dW.
(183)
dV. p2 ((V)(x,))² diag (2 (Xx)) Vi)+ dt.
(184)
where 3; = 7ps, 4(t) 1 - e- to. and P. diag vV, + (2())aa
Remark C.48. See Remark C.35 and Remark C.36 for discussion on the regularity of the SDE derived in Theorem C.47.
Corollary C.49. Under the assumptions of Theorem C.47 with E(x) = oa n = 6J,, B = Bó, Ω1 01P1, and P2 a2p2
dX,-- = -K (2(() (1)
P. (M, (Vf Mt))dt
(185)
dM. = 0n (V) (X,) Mj) dt + vna a TadW,
(186)
VB8
dV. = a2p2 (1V1(X)))² + d² - Vi dt.
(187)
Lemma C.50. Under the assumptions of Corollary C.49, is p-strrongly convex, (v(2ff(x) S C.. and (xx)) = O1)), the asymptolic dynamics of the iterates of Adam satisfies the classic scaling rule vb because [x.)
t-000
LoLr
To enforce that the speed of Mr and Ve match that of 4vBB v8"
Xt, one needs P: kPt, which implies B x(1 Bt).
Proof. First of all, we need to ensure that the relative speeds of X,, M,, and V2 match. Therefore, we select a; = K, which recovers the scaling rules for 3r = 1 x(1 - B.). Then, recalling that (V fx))) = O()), we have that as t + cc, Vi + o² and Me + fx(X,) with high probability. Therefore.
BB
B
dX, = K ((X))dt
(188)
dM, k//PP
dW.
(189)
VB8
dV. = O.
(190)
Therefore, if HX,, V.) = f(X.)
2.5B | and & E (0, 1) we have that by Itô's lemma, p7
B5 vs(x)) di + (L,8BMM,
Me Me)
dH(X,, = -
k
2=O2
vB8
dW.
(191)
C,8B
2² npp o²
+
dt
(192)
2
0i0²
B8
k B8
[VXx + Noise + 1&nC
dt
(193)
2
VB
+1 (X))) di + Noise + 12nCC
2 dt
(194)
22p & (I(X.) 1- 1VV((
B8
*2 dt.
dt Noise
(195)
2
2
46

---

Published as a conference paper at ICLR 2025
Let us now select such that 1-$ ue
Lor this means that { dzpi+uCroB (0, 1) and { =
d'p
of
1 + p Lro Since M, + (X()), we have that pio=
VB8 dH Vi) S 2mp {H V)) + R2nG,
dt + Noise.
(196)
2
Therefore.
[X,(
(1 (1+ peptr
C
t 1 110 LT n
0¹O02 E Vo)l
S
(197)
& 4uvB B v8
which implies that
1000 noL, x 4uvBB v8
B[f(X]]
&
(198)
Analogously,
1 11 L k
-(X.))
5 ApVB v8
(199)
which gives the square root scaling rule.
Lemma C.51. Under the assumptions of Corollary C.49, f(x) = 2² Hz s.t. H = diag(A1, ... ),))
2
and (V {(x))² = On)), the dynamics of Adam implies that f(X,)
yod
4v B v8
Recalling that (Vf(z)²² (1)) we have that as + xc, Vv + B3 and Mr - AX2 with high probability. Therefore, in the one-dimensional case
Proof.
dX, BB
σ XX,dt
(200)
dM, = nP dW.
(201)
VB8
dV. = O.
(202)
Therefore, if (X,, Vi) == XX2 + pf
AB M2, 8 we have that by Itô's lemma,
B6
ABB
dH(x,, V.) = -AXx)
AX,
di +
(pfozM,
Mr
KVTPI dW,
(203)
VB8
1 AOB
mp o²
dt
(204)
2 pi02
B8
VBA
R x2mpf0² 2OB
= -2RA
f(X,)d:
2B8 pio² dt + Noise.
(205)
V Bd
R²11 dt + Noise.
-KAA f(X,)dt
(206)
Once again, since M, + XXe, we have that
AX? 2
XBB
M²
AX?
XBB AX?
XX?
H(X,, V.))
+
1
= KJ(X)).
2
2
0 2
1i02
2
(207)
Therefore,
VBB
x2nd2
KXElf(X))= - 2nA
E[f
dt.
(208)
2
which implies that [Xf( - 11
A/B v which also gives the square root scaling rule. The general- ization to d dimension is analogous and one needs to sum across all the dimensions.
Inspired by (Barakat and Bianchi, 2021)
47

---

Published as a conference paper at ICLR 2025
Lemma C.52. Let (() x Mr where H diag(A Ad). The stationary distribution of Adam (E[Xx]],, (X (0.323=11
Proof. The expected value follows immediately from the fact that
dX, -
(209)
For the covariance, we focus on the one dimeesiioio case. We define H(,,, Vi) == 2 With the same arguments as Lemma C.51, we have
X²
+
xa M?
20 2
A X²dt dt + " dt + Noise,
d(x,)² =
(210)
2
which implies that
B[X;]
10 2X
(211)
The thesis follows by applying the same logic to multiple dimensions.
C.9 ADAMW
In this subsection, derive the SDE of Adam W defined as defined as
0k+1 = 320k + (1 32) V Sr (xx))²
(212)
m 8ime + (1 - B1 V fr (xx)
(213)
rit ma (1 1- -1)
(214)
in U3 (1 (/) -1
(215)
mg+1
Ik+1 Ih Votti ela
n0xx
(216)
with (zo, mo, vo) Rd x RRd Rd 7 R20 is the step size, 3r 1- pn for p1 0
s.t. < (0, 1). P2 = O(1. 0 0, the mini-batches {mk are modelled as i.idd random variables uniformly distributed on {1.... .N, and of size B Z
Theorem C.53. Under the same assumptions as Theorem C.47, the SDE of AdamW is
dx, Vi2(t) P-- 3M,, + no1 (X) - M:)) 0X,dl (tt)
(217)
d.M PI (Vf (X.) M,) dt + //1p/>// (X,) dW,
(218)
dV. P2 ((V)(X))) + diag (E (X))) Vi) dt. where B, = 11 1p, O. (t) = 1 - e t to. and P. = diag VV + ee/12(()
(219)
Proof. The proof is the same as the of Theorem C.47 and the only difference is that 1,0 is approximated with dxtdt.
Figure and Figure 1l validate this result on a variety of architectures and datasets. Remark C.54, See Remark C.35 and Remark C.36 for a discussion on the regularity of the SDE derived in Theorem C.53.
Corollary C.55. Under the assumptions of Theorem C.53 with E(x) = =02la. 7 = kJ,, B = B8,
# 01P11 0 = E0, and 22 = 0222
(2((t) P. (M, + 701P1 (V) - Mi)ddt - KE0X,dt
dX1, = K (1(1)
(220)
dM, = apPi (V. (X,) - M,) dt + whapp IadW,
(221)
VBd
AV. = 422o (VS(x)))3
B5" Ia V dt.
(222)
48

---

Published as a conference paper at ICLR 2025
Lemma C.56 (Scaling Rule at Convergence). Under the assumptions of Corollary C.55, f is p- strongly convex and Lmoothh f(x)) & C. X, = 0, and (V fxx)) = O(T), the asymptotic dynamics of the iterates of AdamW W satisfies the novel scaling rule if R = v8 and { = K. because
60 nL,OL I
K
((f(X)) - f(x.))
(223)
2
2pv BšL o +
By enforcing that the speed of Vi matches that of Xt, one needs X rip. which implies B 1 - x(1 B.).
Proof.
The proof is the same Lemma C.50 where we also use L-smoothness in Lemma C.44.
similar result be derived for general X.: The final expression very convoluted and brings marginally negligible added value.
Lemma C.57. For f(z) H= the stationary distribution of AdamW is (E[Xx Coo((xo)) = (o, (0 (HE-)++1a)))
2
Proof. The proof is the same as Lemma C.52.
Finally, we prove generalization of Lemma C.56 the L-smooth case.
Lemma C.58. Let f L-smooth, te be a learning rate scheduler such that lim
di
0 and
of
$ 30 oo, where of [o(n) Then
E V (x+)) < (((X)) -(x.)) + 2, 1
Pin
Lr
1 1 0, (224)
0902
2
AVōB
o
where is random time with distribution $1
Proof. The proof is the same as Lemma C.24.
D SDES FROM THE LITERATURE
Theorem D.1 (Original Malladi's Statement), Let au :- o7, == en), and 62 = 1.8 Define the
state of the SDE L, (Xe, uz) and the dynamics
dX, = PP P. (Vf (Xx) dt + 1//2 (X,) dW.
(225)
du: = C2 (diag ( (X,)) mx) dt
(226)
where Pt = og diag (w)) 1/2 + Eg
Theorem D.2 (Informal Statement of Theorem C.2 Malladi et al. (2022)). Under sufficient regularity
conditions and Vf(x) = ((vr), the following SDE is an order weak upproximation of RMSprop:
dX, - P (V + vnD((x)+aw.)
(227)
dV, p(diag(E(X)) - di)
(228)
where 33 = 1- np. -- O(1),, and P == diag (v)) 2 + cId.
Lemma D.3. Theorem D.I and Theorem D.2 are equivalent.
Proof. It follows applying time rescaling ( = 1/ and observing that W, Wye WWW
n
49

---

Published as a conference paper at ICLR 2025
Theorem D.4 (Original Malladi's Statement), Let := == 31) :=:====...11111 (1- B2) /nf and define Gu, Eo in Theorem D.I. Let (1(t) := 1 - 1 (-c1t) and (2(t) == exp(pppp (-o2t). Define the state of
==:=(1-
me, U4) the SDE as L, = me,
and the dynamics as
dX, vial()
Pt m,dt
(229)
(1(1)
dm, c1 (Vf (Xt, mm,) dt + aoo 2092 (X₂) .,,
(230)
du, 02 (diag (2 (Xi)) uz) dt,
(231)
where Pt = og diag (u.) + /2 (t)
Theorem D.5 (Informal Statement of Theorem D.2 Malladi et al. (2022)), Under sufficient regularity conditions and (z) (V))), the following SDE is an order weak approximation of Adam:
dx. vi2(I)
P. P. M,dt
P
(232)
(t)
dM, 0I (VS (Xx) MM) dt /nnii 1// dW,
(233)
dV: P2 (diag (Xx)) - V()) dt.
(234)
where B; =1 nps, (:t)= - e Pi 0(1), and P. diag Vv. + e212(1)Ta
Lemma D.6. Theorem D.4 and Theorem D.5 are equivalent.
Proof. It follows applying time rescaling t == T7C and observing that W; = Wn = VnW
E SDE CANNOT BE DERIVED NOR USED NAIVELY
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 7/ of size B, and starting point øo = x, which we dub SGD (n.B), (n.( The iterates are given by:
Fk+1 Fx 1y SyV.g (2k)
(235)
which for U = (V()) (xx)))))) we rewrite as
23 7/ fxxx) + VTJU..
(236)
where E[U+] and Cov(Us) If we now consider the SDE
2 [x+x) 1 -//
dX, = - (f( + 7l )
(237)
its Fuler-Maruyama discretization with pace Al = and Z₁ N Ia) is
x(x) Zn. = - (X)x))++++++++) - + VTV B
7
(238)
Since the Eq. 235 and Eq. 238 share the first moments, is reasonable that by identifying (. knj. the SDE in Eq. 237 is a good model to describe the iterates of SGD in Eq. 235.
Informally, we need "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.
Definition E. (Weak Approximation). A continuous-time stochastic process (X1 is order [T/m) if for every
weak approximation (or a-orrer SDE) of discrete stochastic process {ax k== polynomial growth function g, there exists a positive constant C. independent the stepsize 7, such that maxx==0...,122/
Eg (xx) Eg (Xx,)) < Ct
To if SDE satisfies such definition, has check that for A x1 and A X,, F.
50

---

Published as a conference paper at ICLR 2025
1. EA; - EA = O Vi = l.....dd 2. EA,A, - EA - O()), Vi,j = ,....,d. d.
Example: Let us prove that the SDE in Eq. 237 is valid approximation of (GD(1.)); The first condition is easily verified. Coming to the second condition we have that
1. EA.A, - -00_0_0000 (
2. = 1 =(,))++++
whose difference is of order v/² and thus satisfies the condition. However, we observe that the scale of the noise is too small w.r.t 17, i.. [(x) = O17")) for 0, then the simplest SDE model describing (G B) the ODE dX, = VV(()) as in that case
1. EA,A,
2. EA,A, =
whose difference is also of order 1j². Much differently, if E(x) 0() for a > O, the simplest model is the SDE in Eq. 237. We highlight that simplest does not mean best: The SDE more accurate than the ODE even in a regime with low noise, but this observation serves 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 ñ Kr) and batch size := 8B for and > 1, which we dub (..B, The first is derived considering that the learning rate is and carries error of order (ī)) w.r.t. (G.
dx, = - (VX + V = - -V
+ v (X,))
(239)
The second one instead is derived considering y as the learning rate and k. as a constant "scheduler". Consistently with (Li et al., 2017), the SDE which carries error of order (γ) w.r.t. SGD is
TI x(X))
dx, - xVX + "V B8
(240)
While they both are valid models, there are three reasons why one should prefer the latter:
It fully reflects the fact that larger learning rate results 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 x and that can ever leave the dynamics unchanged.
E.1 DERIVING SCALING RULES
Jastrzebski et al. (2018) observed that only the ratio between 7 and B matters in determining the dynamics of Eq. 238. Therefore, they argue that for ô the SDE for SGD (87,8B coincides with that of SGD v. and that this implies that the path properties of the optimizers are the same. On the contrary, the path of (GD( strongly depends on the hyperparameters: The speed and volatility of the dynamics are driven by 1J. 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 hyperparameter, e.g. 77, and have 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 the convergence bound the expected loss. To show this quantitatively, we use this rationale to derive the scaling rule of SGD as we aim at preserving the asymptotic loss level.
Lemma F.2. If f is a μ strongly convex function, Tr Tr(V)[(>)) < L, and E(x) T²Ta. then:
51

---

Published as a conference paper at ICLR 2025
1. Under the dynamics of Eq. 237 we have:
11 0
-(Xx,))((x.)) (f(X)) x(
(1- e 2u0);
(241)
2 2pB
2. Under the dynamics of Eq. 239 we have:
E(f(X,) ((X.)- - (f(Xo) f(X.))e 2ut o² (1-e-2ut)
(242)
B 8
3. Under the dynamics of Eq. 240 we have:
-2pntt B(f(X.) f(xx)) S (f(Xo) f(x.))e 2pnt + o² x.
(1 e--nnt. 2pB 8
(243)
The first bound implies that the asymptotic limit of the expected loss for SGD( C. The C. x last two bounds predict that the asymptotic loss level for (. is 2 2nd 3 Since the objective of the scaling rule is to find and å such that (f. achieves the same loss level as SGD (n. we recover the linear scaling rule setting K However, only the last bound can correctly capture the fact that the dynamics of SGD We thus conclude that:
2μB
times faster than that of SGD
1. Eq. 240 is a better model for (G. 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.3. In Malladi et al. (2022), Theorem 5.3 proposes formal derivation of scaling rule for RMSprop. Following the perspective of Jastrzebski et al. (2018), 8), the authors suggest that a scaling rule that leaves their SDE unchanged would also preserve the dynamics of the RMSprop iterates. We note, however, that an SDE defined not only by the equation that governs the dynamics but also by its initial condition (see Karatzas and Shreve (2014). Section 5). Although the scaling rule in question leaves the differential equation unchanged. it modifies the initial condition of the process u,, and hence, the overall SDE is altered. This observation suggests that the claim and proof in Malladi et al. (2022) may need to be revisited.
Furthermore, the rule appears to be valid only in the vicinity of convergence, as the corresponding SDE is applicable primarily in that regime. Additionally, Lemma E.2 provides concrete evidence that maintaining the form of the SDE does not necessarily imply that the path properties of the optimizer preserved. We offer these clarifications with the hope of contributing constructively to the discussion.
F EXPERIMENTS
In this section, provide the modeling choices and instructions replicate our experiments. The code is implemented in Python (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 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 1. The learning rate is 77 0.001. Similarly, integrate the SignSGD SDE (Eq. 7) with Euler-Maruyama (Algorithm 1) with At = 17. Results are averaged over 3 runs and the shaded areas are the average + the standard deviation.
with Gaussian noise Z (000 o d²Ix) Ia) where
52

---

Published as a conference paper at ICLR 2025
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). (3, 3, 32) convolutional layer with stride 1, ReLu activation, (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, 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 60000 epochs we calculate the full gradient and inject with Gaussian noise Z ~ (0, La where = 0.4. The learning rate is = 0.001. Similarly, we integrate the SignSGD SDE (Eq. 7) with Suler-Maruyanu (Algorithm 1) with At n7. 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 feature=100, width=488 depth=3, num heads=6, and dim ffr=192. We minimize the cross- cross-entropy loss as we run SignSGD for 5000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ~ (0, o² o²Ta) where
1. The learning rate is v 0.001. Similarly, we integrate the SignSGD SDE (Eq. 7) with Fuler-Maruyama (Algorithm 1) with At = Results are averaged over 3 runs and the shaded areas 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 (3, 3, 32) convolutional layer with stride 1, followed by a ReLu activation, second (3, 3, 32) convolutional layer with stride followed by residual connection from the first convolutional layer, then (2, 2) max pool layer with stride (2, 2). Then the activations are flattened and passed through dense layer that compresses them into 128 dimensions, final ReLu activation, and final dense layer into the output dimension 10. The output finally goes through a softmax 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 (0, o²Ia) where = 1. The learning rate is 17 0.001. Similarly, we integrate the SignSGD SDE (Eq. 7) with Marlyamm (Algorithm 1) with At = 17. Results are averaged over 3 runs and the shaded areas 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 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. 24) and the one describing Phase 3 (Eq. 5). The optimized function fxx
ddiag(1, 2), zo drawn (and fixed for all runs) from a normal distribution (00, 0.01), 7 = 0.001, and = La where a = 0.1. We integrate the SDEs with (Algorithm 1) with At 17 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.
Ha for 2
To produce the center-left of Figure 2, we repeat the above as xo drawn (and fixed for all runs) from normal distribution (0, 1). Then, plot the average loss values together with the theoretical prediction of Phase 1 and Phase 3: They perfectly overlap.
Stationary distribution: Lemma 3.7 In this paragraph, we describe how we validated the conver-
gence behavior predicted in Lemma 3.7. To produce the center-right of Figure 2), we run SignSGD on f(x) = 2 We run this for 5000 times and report the evolution of the moments. Then, we add lines representing the theoretical predictions derived in Lemma 3.7: They match.
He for H = diag (1, 2), xig = (0.001, 0.001), = 0.001 and Σ = I²I, where a = 0.1.
Schedulers: Lemma 3.9 In this paragraph, we describe how we validated the convergence hehavior predicted in Lemma 3.9. To produce the right of Figure 2, we run SignSGD f(z) = =T for H diag(1, 1, 2), zo (0.01, 0.01), 17 = 0.01 and Σ T where 0.1. We used the scheduler
2
53

---

Published as a conference paper at ICLR 2025
(++1ja nh" -++1a for & (0.1,0.5,11 1.5). For the first two choices of 1, Vit satisfies our sufficient
condition for the convergence of SignSGD: In the figure, we observe that indeed SignSGD converges to with the same speed as the one predicted in the Lemma. For = 1.5, we observe that SignSGD does not converge following the theoretical curve because it does not satisfy our sufficient condition. Results are averaged over 500 runs.
F.3 RMSPROP: SDE VALIDATION (FIGURE 9 AND FIGURE 10)
In this subsection, we describe the experiments we run to produce Figure 9 and Figure 10: The dynamics of our SDE matches that of RMSprop more accurately than the SDE derived in (Malladi et al., 2022).
This paragraph refers to the left and center-left of Figure 9. We optimize the function (x) HH where H = (iaag(1,,)))) 2). We run RMSprop for 2000 epochs as calculate the full gradient and inject it with Gaussian noise Z ~ (N00, IaIa) where 0.1. The learning rate is 17 = 0.01, 3 = 0.99. Similarly, we integrate our RMSprop SDE (Eq. 129) and that of Malladi (Eq. 227) with Euler-Maruyama (Algorithm 1) with AL 1/. Results averaged over 500 runs and the shaded areas are the average + the standard deviation: Our SDE matches RMSprop more accurately.
Quadratic convex function
Embedded saddle This paragraph refers to the center-right and right of Figure 9. We optimize the 2 1 - x $ E 2 where H diag(-
À = and { = 0.1. We run RMSprop for 1600 epochs as we calculate the full gradient and inject it with Gaussian noise Z ~ (00 L2 where a = 0.01. The learning rate is = 0.01, 3 = 0.99. Similarly, we integrate our RMSprop SDE (Eq. 129) and that of Malladi (Eq. 227) with Euler- (Algorithm I) with At = n). Results are averaged over 500 runs and the shaded areas are the average = the standard deviation: Our SDE matches RMSprop more accurately.
function f(x) =
DNN on Breast Cancer Dataset This paragraph refers to the left of Figure 10. The architecture and loss are the same as used above for SignSGD. We run RMSprop for 2000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ~ AO, o² la) where a 10-2 The learning rate is 77 = 10 B = 0.9995. Similarly, we integrate our RMSprop SDE (Eq. 129) and that of Malladi (Eq. 227) with Euler-Maruyame (Algorithm 1) with At 17. Results are averaged over 3 runs and the shaded areas are the average + the standard deviation: Our SDE matches RMSprop more accurately.
CNN on MNIST This paragraph refers to the center-left of Figure 10. The architecture and loss are the same as used above for SignSGD. We run RMSprop for 100000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ~ (O, La) where
10-2 The learning rate is 17 = 10 4, B = 0.999. Similarly, we integrate our RMSprop SDE (Eq. 129) and that of Malladi (Eq. 227) with Euler-Maruyamo Faler-Maruyama (Algorithm with At 1j. Results averaged over run and the shaded areas are the average + the standard deviation: Our SDE matches RMSprop more accurately.
Transformer on MNIST This paragraph refers to the center-right of Figure 10. The architecture and loss the same used above for SignSGD. We run RMSprop for 2000 epochs
calculate the full gradient and inject it with Gaussian noise Z ~ (00 o² La) where o = 10 10-2 The learning rate is 10-3 3, B 0.995. Similarly, we integrate our RMSprop SDE (Eq. 129) and that of Malladi (Eq. 227) with Euler-Maruyama (Algorithm 1) with At = 77. Results are averaged over 3 runs and the shaded areas are the average + the standard deviation: Our SDE matches RMSprop more accurately.
ResNet on CIFAR-10 This paragraph refers to the right of Figure 10. The architecture and loss are the same as used above for SignSGD. We run RMSprop for 500 epochs as we calculate the full gradient and inject it with Gaussian noise Z ~ AV00, IaIa) where o = 10 The learning rate is 17 10 4, 8 0. 0.9999. Similarly, we integrate our RMSprop SDE (Eq. 129) and that of Malladi (Eq. 227) with Euler- Matter-Maaaaaaa (Algorithm 1) with At = n, Results are averaged over runs and the shaded areas the average + the standard deviation: Our SDE matches RMSprop more accurately.
54

---

Published as a conference paper at ICLR 2025
F.4 ADAM: SDE VALIDATION (FIGURE 12 AND FIGURE 13)
In this subsection, we describe the experiments we run to produce Figure 13 and Figure 12: The dynamies of our SDE matches that of Adam more accurately than that derived in (Malladi al., 2022).
Quadratic convex function This paragraph refers to the left and center-left of Figure 12. We optimize the function f(z) = = Hz where H diag(10, 2). We run Adam for 50000 epochs as we calculate the full gradient and inject with Gaussian noise Z ~ (O00, d² where o = 0.01. The learning rate is 0.001, 31 = 0.9, and 82 0.999. Similarly, we integrate (Eq. 167) and that of Malladi (Eq. 232) with Euler- Maller-aruyymaa (Algorithm 1) with At = mj. Results averaged over 500 runs and the shaded areas are the average + the standard deviation: Our SDE matches Adam more accurately.
Adam SDE
This paragraph refers the center-right and right of Figure 12. We optimize the 2 } Ei= ot $ E xº where H = diag((11,2) À = and 5 0.1. We run Adam as we calculate the full gradient and inject with Gaussian noise 2 ~ -, La) where a 0.1. The learning rate is 1) 0.001, B1 0.9, and 3g 0.999. Similarly, we integrate our Adam SDE (Eq. 167) and that of Malladi (Eq. 232) with Euler- Marlya-a (Algorithm 1) with At 7). Results are averaged over 500 runs and the shaded areas are the average + the standard deviation: Our SDE matches Adam more accurately.
Embedded saddle function [(x) =
DNN on Breast Cancer Dataset This paragraph refers to the left of Figure 13. The architecture and loss are the same as used above for SignSGD. We run Adam for 2000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ~ (O, oaIa) where
10-2 The learning rate is 10-4, 31 = 0.99, and 32 0.999. Similarly, we integrate our Adam SDE (Eq. 167) and that of Malladi (Eq. 232) with Euler-Maruyama (Algorithm 1) with At = r.. Results averaged over runs and the shaded areas are the average + the standard deviation: Our SDE matches Adam more accurately.
CNN on MNIST This paragraph refers to the center-left of Figure 13. The architecture and loss the same as used above for SignSGD. We run Adam for 40000 epochs we calculate the full gradient and inject it with Gaussian noise Z ~ (O, la) where o = 10 10-2 The learning rate is 17 = 10 81 0.99, and 32 0.999. Similarly, we integrate our Adam SDE (Eq. 167) and that of Malladi (Eq. 232) with Euler- Maruyama (Algorithm 1) with At nj. Results are averaged over 3 runs and the shaded areas are the average + the standard deviation: Our SDE matches Adam more accurately.
Transformer on MNIST This paragraph refers the center-right of Figure 13. The architecture and loss are the same used above for SignSGD. We run Adam for 2000 epochs as we calculate the full gradient and inject it with Gaussian noise ~ (O, I²Ia) where
10-2 2 The learning rate is 1 10-2, 2, Bi = 0.9, and B2 0.99. Similarly, we integrate our Adam SDE (Eq. 167) and that of Malladi (Eq. 232) with Maruyama (Algorithm 1) with Al = y,. Results are averaged over
runs and the shaded areas the average + the standard deviation: Our SDE matches Adam more accurately.
ResNet CIFAR-10 This paragraph refers to the right of Figure 13. The architecture and loss are the same as used above for SignSGD. We run Adam for 2000 epochs as we calculate the full gradient and inject it with Gaussian noise Z ~ (O, oz Ia) where = 10-5 The learning rate is 1 10 31 = 0.99, and 82 = 0. 0.9999. Similarly, we integrate our Adam SDE (Eq. 167) and that of Malladi (Eq. 232) with Suler-Maruyama (Algorithm 1) with At 17. Results are averaged over 3 runs and the shaded areas are the average + the standard deviation: Our SDE matches Adam more accurately.
F.5 RMSPROPW ADAMW: SDE VALIDATION (FIGURE 3, FIGURE 4)
The settings are exactly the same 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.
55

---

Published as a conference paper at ICLR 2025
Losses
Adami (o=0.00)
Adam.2 (o-0.1)
10-1
Adam-2 (a=100)
10
Adam- (0=10.0)
Loss
10-3
Limit (.000.11)
Adam-L2 (0=1000.0)
10-4
Limit (o
Limit to 1.0)
10-3
Limik
10
Limit (o 1000.0)
10-7
0
20000
40000
60000
80000
100000
Iterations
Figure 14: Empirical validation of the bounds for Adam on an Le-regularized loss f(z) + 011 For several levels of noise we find that our theoretical predictions match the experimental results: The loss levels scale linearly a.
F.6 RMSPROPW ADAMW: INSIGHTS VALIDATION (FIGURE 5)
In this subsection, describe the experiments predicted asymptotic loss value and moments of RMSpropW and AdamW match those empirically found.
produce Figure 5: The theoretically
Asymptotic loss & scaling rule of Adam This paragraph refers the left of Figure 5. We optimize the function f(x) = x Hz where H = diag(1.3). We rum AdamW for 20000 epochs as we calculate the full gradient and inject with Gaussian noise Z ~ No0, o2, where = 1. The for both 0 and The rescaled versions of the algorithms Adam W R follow the novel scaling rule with x = 2. Adam W NR follows the scaling rule but for 0 which is left unchanged. We plot the evolution of the loss values with the theoretical predictions of Lemma C.50: Results are averaged over 500 runs.
learning rate is 1 0.001, 31 0.9, and B3 0.999. Experiments
4.
Asymptotic loss & scaling rule of RMSprop This paragraph refers to the center-left of Figure 5: The only difference with the previous paragraph is that we use RMSpropW with B 0.999.
Adam W: the role of the 3s This paragraph refers the center-right of Figure 5. We optimize
the function f(x) = = Ha + 1 Σ x7 3 - N where H = (iaag(11,2) 2). = 1, and 0.1. We run AdamW as calculate the full gradient and inject it with Gaussian noise 7 ~ o22) where a = 0.1. The learning rate is 1) 0.001, = 0.1, 31 E {0.99, 0.999}. and B2 {0.992, 0.996, 0.998): Clearly, three combinations go into a minimum and three go into the other. For each minimum, the three optimizers converge to the same asymptotic loss value independently on the values of 31 and 3z. We argue that 81, and 33 select the basin and the speed of convergence, not the asymptotic loss value: This is consistent with Lemma 3.13.
Stationary distribution This paragraph refers to the right of Figure 5. We optimize the function calculate the full
_ (x) = TH where H (1,3). We run Adam for 20000 epochs gradient and inject it with Gaussian noise ~ (0, o² La) where a 0.01. The learning rate is 0.001, 0 = 4, 3 = 0.999, B1 = 0.9, and 82 = 0.999. We plot the evolution of the average variances with the theoretical predictions of Lemma C.45 and Lemma 3. 14: Results are averaged over 100 runs.
F.7 EFFECT OF NOISE VALIDATION (FIGURE AND FIGURE 14)
In this subsection, deseribe the experiments run to produce Figure and Figure 14: All bounds on the asymptotic expected loss value for SGD, SignSGD, Adam, and Adam W, and Adam on an reguuarrrd loss are perfectly verified.
56

---

Published as a conference paper at ICLR 2025
We optimize the loss (x) -- where H diag(1, 1) as each optimizer for 100000 iterations with r 0.01. We repeat this procedure five times, one for each a E {0.01, 0.1, 10, 100). As we train, we inject noise on the gradient as distributed NV La). We plot the average loss together with the respective limits predicted by our Lemmas. For each optimizer and each ar, the average asymptotic loss matches the predicted limit. Therefore, we verify that the loss of SGD scales
quadratically in that of Adam on f(x), Adam f(z) + ~| and SignSGD scales linearly, and that AdamW is limited Ø.
F.8
INCREASING WEIGHT DECAY WITH THE BATCH SIZE (FIGURE 15 AND TABLE 1)
Small Batch (256), 4800 Iterations
Big Batch (1024), 4800 iterations
28.0
34
27.5
beta2=0.9875
test ppl
test ppl
beta2=0.975
27.0
beta2=0.95
beta2=0.9
26.5
beta2=0.8
26.0
10-3
10-3
learning rate
learning rate
Figure 15: RMSprop Scaling. We first (left plot) tune at small batch size (256) the learning rate and B2 simultaneously and report all test perplexity results after cosine decay at 2.5B tokens (4800 RMSprop steps). For these runs, weight decay is turned off to eliminate confounding factors this first analysis (see later Table 1). We then report the results for the same tuning but at batch size of 1024 sequences. One can clearly see that the optimal learning rate shifts up from le to 2e predicted by both ours and Malladi al. (2022) scalings (factor 5 = 4). The best run in the base setting with 82 0.95. Our rule predicts scaled up optimal 32 with value 1 v 32)
1 - 2 x 0.5 = 0.9, while Malladi et al. (2022) predicts 1 - 8(1 (22 = 1 - 4 x 0.5 = 0.8. The results show that while results are close (biggest effects in learning rate scaling with vo), our scaling slightly more performant.
The analysis of Malladi al. (2022) suggests that, when scaling batch size B by factor 5, one also
has scale up (1) the learning 7) by factor vō and scale down (4) 91 to the value 1 - 5(1 - (1) and B2 to the value - 5(1 (2). Our SDE analysis confirms similar rules (Lemma 3.1 but
1. Proposes to scale down less 31 and B3, i.e. as 1 vō(1 3).
2. Additionally suggests scaling up the decoupled weight decay parameter 0 by a factor Vo.
We test this in the language modeling setting utilizing Pythia-like 160M parameter transformer architecture (Biderman et al., 2023) trained on 2.5B and 10B tokens from the SlimPajama dataset. We scale up the batch size by factor of ő 4 and keep the same number of iterations i.e., we have 10B tokens of training the scaled-up runs. The sequence length all of our experiments is 2048 tokens. We perform two experiments, and report results in Figure 15 and Table 1:
Vanilla experiment. In Figure 15, we set 31, = and study in isolation the effects of scaling the learning rate and 32 as the batch size increases. The results clearly indicate that scaling up the learning rate is beneficial when increasing the batch size. In addition, they indicate that our scaling of 32 might be slightly preferable compared to the one in Malladi et al. (2022).
Scaling weight decay. In Table I we verify our scaling on AdamW W runs (31, 0 of 0). We operate in a similar regime Figure 15 and scale 9 distinct configurations at small batch size to big batch size using different strategies. Again, results indicate the effectiveness of our strategy.
Remark. While the experiments above show promise, future research is needed in order to better compare and evaluate strategies. In particular, we noticed that in pathollggial settings, it might be beneficial to not scale down 81 with the batch size and perhaps keep the same learning rate.
57

---

Published as a conference paper at ICLR 2025
Baseline B 256, 0 0.1
(Malladi et al., 2022)
(Malladi et al., 2022)
This paper.
== 1024, = 0.1
B = 1024, 0 0.2
B 1024, 0 = 0.2
81 = 0.975, 92 0.9875
= 0.9, A22 0.95
81 0.9, (2 0.95
81 0.95, 82 0.975 19.5785
23.5889
20.6157
20.2158
81 0.975, A2 0.975
Bi = 0.9, 20.0463
(22 0.9
81 0.9, (2 0.9 19.7162
81 0.95, 82 0.95 19.3467
23.7916
81 0.975, 23.7461
82 0.95
81 = 0.9, 31 0.8
81 0.9, (2 0.8 19.5846
81 0.95, B2 0.9 19,2398
20.0237
81 0.95, |2 0.9875
0.8, B1 0.95 22.7750
81 0.8, B2 0.95
91 0.9, B2 0.975 20.6354
23.4310
21.3613
8, 0.95, |2 0.975
81 0.8, B2 0.9
81 0.8, 82 0.9
91 0.9, B2 0.95 20.2158
23.3911
21.4489
20.4485
8, 0.95, |2 0.95
B1 0.8, B2 0.8
81 0.8, 82 0.8
91 0.9, B2 0.9
23.4654
20.5648
20.3054
19.7162
81 0.9, 82 0.9875
0.6, 32 0.95
81 0.6, 82 0.95 185,4383
81 0.8, 0.975 23.3668
25.0240
1972.5442
81 0.9, 0.975
= 0.6, 82 0.9
81 0.6, 82 0.9 42.6782
81 0.8, 0.95 21.3613
25.1012
646.8980
81 0.9, 82 0.95 23.6411
0.6, 82 0.8
8 0.6, 82 0.8 24.2124
81 0.8, 0.9
145.7700
20.4485
Table 1:
We perform 9 AdamW base runs at batch size of 256 sequences of length 2048 (first column). A total of 4800 steps are performed, for a total of 2.5B tokens. For these runs, we always select learning rate of 0.004 and weight decay 0 of 0.1. We report results (test perplexity) for different combinations of 31, 82. the three right-most columns, we scale each setting according either to Malladi al. (2022) according to Lemma 3.13. Since Malladi al. (2022) does not give prescriptions on the weight decay value we either scale as we propose or leave it at 0.1. Scaled number of steps but with a batch size of 1024 sequences (factor 4). Test accuracy results indicate that our scaling more effective in terms of final test perplexity compared (Malladi et al., 2022).
up runs process xx the number of tokens, i.e. 10B tokens: the algorithm performs the
58


Reviews:
[{"review_id": "XtsirGLitz", "reviewer": "Reviewer_2nWX", "summary": "This paper presents a theoretical analysis of adaptive optimization methods—specifically SignSGD, AdamW, and RMSpropW—through the lens of stochastic differential equations (SDEs), unveiling the intricate interplay between adaptivity, gradient noise, and curvature. Key contributions include:\n\n- **SignSGD Dynamics:** Identification of three distinct phases in SignSGD’s dynamics, with noise inversely affecting both convergence rates of the loss and the iterates.\n- **Enhanced SDE Models:** Derivation of new and improved SDEs for AdamW and RMSpropW, leading to a novel batch size scaling rule and an examination of the stationary distribution and stationary loss value in convex quadratic settings.\n\nThe derivation of new SDEs for SignSGD, AdamW, and RMSpropW provides a solid foundation for understanding the dynamics of these optimizers.", "questions": "See weaknesses above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths includes:\n- **Intriging Insights:** Identifying three distinct phases in SignSGD dynamics and the inverse relationship between noise and convergence rates offer valuable perspectives on optimizer behavior.\n- **Practical Implications:** Introducing a novel batch size scaling rule and examining stationary distributions have direct implications for better training practices in deep learning.\n- **Clear background:** The preliminaries and background information provided in the appendix enhance the readability and understanding of the main results.\n- **Validation:** Theoretical findings are supported by extensive experimental evidence across various neural network architectures, including MLPs, CNNs, ResNets, and Transformers.", "weaknesses": "- **Unclear Terminology and Notation:** Certain terms and symbols used in the paper are not clearly defined, which may lead to confusion. What does $\\mathbb{P}$ represent in Theorem 3.2? How is the error function defined in Corollary 3.3? Does Lemma 3.7 explicitly establish the stationary distribution of SignSGD?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jmVhDNIuaj", "reviewer": "Reviewer_ahk4", "summary": "This paper investigates the learning dynamics of adaptive optimizers in machine learning.  The authors analyze a simplified model, $f(x) = x^THx$, with noisy gradients defined as $g_\\gamma = \\nabla f(x) + \\gamma$, where $\\gamma$ is sampled from a distribution like a Gaussian. Focusing on SGD, sign-SGD,  and AdamW, they demonstrate that the learning dynamics can be effectively approximated by a stochastic differential equation (SDE) of the form $X_t = \\text{function}(X_t)$. This allows them to analyze the first and second moments of $X_t$ for these optimizers as $t$ approaches infinity.  The authors support their theoretical findings with experiments on practical models like ResNets and small Transformers, trained on MNIST, CIFAR, and Shakespeare datasets.\n\nWhile generally well-written, mathematically sound and provide many insights, the paper's reliance on a simplified model and unrealistic noise design limits its applicability.\n\n\n--- updates --- \n\nThank you for the clarifications. While I believe this work warrants a score of 7, I understand that's not an available option.  Considering that other reviewers have mostly given scores of 6, I am raising my score to 8 to better reflect the paper's merit.", "questions": "* The notation $\\nabla f_\\gamma$ is unclear and could be misinterpreted as $\\nabla (f_\\gamma)$. \n* All algorithms in the paper use full batch gradients with added noise, which is misleading and doesn't accurately reflect real-world implementations of algorithms like Adam.\n* It would be beneficial to generalize the results to mini-batch linear regression with random Gaussian design, similar to the approach in https://arxiv.org/abs/2405.15074.\n* The paper should provide more intuition behind the proof of Theorem 3.2 to help readers understand the key ideas without having to delve into the appendix.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- generally well-written \n\n- mathematically seems sound, though i didn't check the proof \n\n- papper contains rich set of results that may be insightful for practicioners. \n\n- theoretical results are supported by real networks, though on", "weaknesses": "* The toy model is too simple and far from practical. I would expect at least a linear regression model with a random (Gaussian) design.\n* The noise design is artificial and unrealistic. It doesn't capture practical batch noise, which depends on the loss, parameters, and changes over time. The current noise function uses a time-independent, identical distribution.\n*  Experiments use full batch with synthetic (Gaussian) noise instead of mini-batches, which is consistent with the paper's problem setup but unrealistic.\n* The paper is overly dense, with little intuition provided for the theorems and proofs. It also relies heavily on in-line equations, hindering readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ynwZw2lgVv", "reviewer": "Reviewer_2DqW", "summary": "The work is centered on introducing novel SDEs for adaptive optimizers. Under weak assumptions, they derive SDE for SignSGD and compare it against vanilla SGD to show three different phases in the dynamics of SignSGD. They also show an inverse relationship between noise and convergence rate for SGD which is in contrast to the quadratic effect on loss & variance observed for SGD. They also analyze decoupled weight decay and compare it against SignSGD.", "questions": "I am not quite sure what is the key message (takeaway) of SDE analysis for decoupled weight decay analysis. They are known to be more effective than their vanilla counterparts. Is there a key message here that helps improve our understanding of adaptive optimizers? Also, is the analysis valid only for convex quadratics in this case?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The paper has introduced A rigorous mathematical framework of SDEs for SignSGD and its comparison against commonly used adaptive optimizers like ADAMW. They also characterize a relationship between noise and convergence rates for SignSGD and SGD. Also, the framework is used to understand decoupled weight decay efficacy and infer the effect of noise on such approaches. Their analysis is followed by relevant experiments.", "weaknesses": "Major weakness\n\n1. Line 245-246. The definitions of 'signal' and 'noise' need to be explicitly specified in this context. The current characterization of 'signal' in these lines does not make sense. $Y_t$ line 199 defined signal-to-noise ratio but I don't see what is being referred to in lines 245-246 and what part is being referred to as \"large\". The entire paragraph here should be rephrased to better connect it to the surrounding content as it currently seems disconnected, or poorly worded.\n\n2. Line 245: \"SNR is large, meaning SignSGD behaves like SignGD\": This claim is unsupported. You should provide the reasoning or evidence that supports this claim, such as referencing relevant equations, prior results, or additional explanation. \n\n3. Figure 2 referred to in line 249: Hard to understand. Please provide details about the loss function/landscape, dimensionality, and any other relevant parameters used to generate Figure 2.  Lacks description of the setting used to plot figures.\n\n4. In general, (not commenting on technical soundness), the paper is very hard to follow and poorly structured. \n- Need to provide brief intuitive explanations before each lemma, highlighting its significance and key takeaways in order to understand what the authors want to convey.\n- In section 3, authors should be more rigorous about introducing notation. There are unnecessary references to related work mentions which should be pushed somewhere else.  A comprehensive introduction of notation in the upcoming lemmas or some useful background should be introduced. Since the main argument is around novel SDEs, more introductory material on SDEs and their relevance to optimization algorithms before presenting the main results.\n\n5. Section 3.1.1 - authors should ensure all new notation is properly introduced before it is used in lemmas, either by adding definitions as terms first appear or by including a notation section at the beginning of 3.1.1.\n\n6. Line 412: Decoupled weight decay. \n- Clarify why decoupled weight decay is considered key in this context\n- Provide a more explicit link between their analysis and the claim about stabilization at high noise levels. Currently, the authors' analysis to support this is not evident.\n- Include specific equations or results that demonstrate the stabilization effect of decoupled weight decay\n\n7. The authors should include key experimental results in the main paper, selecting the most important figures or tables that support their main claims, while potentially moving more detailed results to an appendix if space is an issue.\n\nNot sure how to comment on the rigor of mathematical soundness of SDEs analysis but overall it was very hard to follow what the authors wanted to show/convey.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "idhurXWj4Q", "reviewer": "Reviewer_QC2v", "summary": "This work derives a novel SDE (of weak approximation order 1, under Gaussian gradient noise Z) for the SignSGD optimizer, and using this SDE, uncovers three phases of dynamics with distinct behaviours. They show theoretically that the effect of noise on SignSGD differs from SGD, and that SignSGD is more resilient to high levels of noise. Further, they show that when Z is heavy-tailed noise, the associated SDEs for SignSGD take a very similar form, further highlighting SignSGD's resilience to high levels of noise. They then derive SDEs (of weak approximation order 1) for AdamW and RMSpropW under much less restrictive assumptions than previous work, and find that decoupled weight decay has an important role in stabilization of dynamics at high noise levels. They also provide empirical validations of their theoretical results.", "questions": "For Lemma 3.4 (line 199), since Y_t \\in R^d, is the condition e.g. |Y_t| > 3/2 denoting the absolute value treated element-wise, like |(Y_t)_i| > 3/2, or does |Y_t| denote the norm of Y_t?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Results are presented very clearly and interpretations are made clear\n\nUse numerical integrators to verify that their derived SDEs match the associated discrete-time optimizers across a diverse range of relevant setups\n\nTechnical assumptions appear much less restrictive than other works", "weaknesses": "Mostly theoretical, with practical implications of findings underexplored; do the findings tell you anything about how the studied optimizers can be improved, tuned, etc.?\n\nForm of stochastic gradient appears unjustified. Is the noise Z being additive, and distributed as a Gaussian or heavy-tailed, accurate to practice?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)