The Marginal Value of Momentum for Small Learning Rate SGD
Authors: Runzhe Wang
Sadhika Malladi
Tianhao Wang
Kaifeng Lyu
Zhiyuan Li
Venue: ICLR 2024 poster (2024)
Keywords: momentum
SGD
dynamics
Primary area: optimization
TL;DR: We prove that momentum has no optimization and generalization benefits for small learning-rate and high gradient-noise SGD.

Abstract: Momentum is known to accelerate the convergence of gradient descent in strongly convex settings without stochastic gradient noise. In stochastic optimization, such as training neural networks, folklore suggests that momentum may help deep learning optimization by reducing the variance of the stochastic gradient update, but previous theoretical analyses do not find momentum to offer any provable acceleration. Theoretical results in this paper clarify the role of momentum in stochastic settings where the learning rate is small and gradient noise is the dominant source of instability, suggesting that SGD with and without momentum behave similarly in the short and long time horizons. Experiments show that momentum indeed has limited benefits for both optimization and generalization in practical training regimes where the optimal learning rate is not very large, including small- to medium-batch training from scratch on ImageNet and fine-tuning language models on downstream tasks.

Published as a conference paper at ICLR 2024
THE MARGINAL VALUE OF MOMENTUM FOR SMALL LEARNING RATE SGD
Runzhe Wang', Sadhika Malladi', Tianhao Wang", Kaifeng Lyu', Zhiyuan Li
Princeton University, Yale University, *Stanford University, "Toyota Technological Institute at Chicago
(runzhew, ama ladí, k yu k @errieeetonooo..o edu, tianhao. wwng@yale. edu, 2hiyuanli@ttic.edl
ABSTRACT
Momentum is known to accelerate the convergence of gradient descent in strongly convex settings without stochastic gradient noise. In stochastic optimization, such as training neural networks, folklore suggests that momentum may help deep learning optimization by reducing the variance the stochastic gradient update, but previous theoretical analyses do not find momentum to offer any provable acceleration. Theoretical results in this paper clarify the role momentum in stochastic settings where the learning rate is small and gradient noise is the dominant source of instability, suggesting that SGD with and without momentum behave similarly in the short and long time horizons. Experiments show that momentum indeed has limited benefits for both optimization and generalization in practical training regimes where the optimal learning rate not very large, including small- to medium-batch training from scratch on ImageNet and fine- tuning language models on downstream tasks.
INTRODUCTION
In modern deep learning, it is standard to combine stochastic gradient methods with heavy-ball momentum, or momentum for short, to enable more stable and efficient training of neural net- works (Sutskever et al., 2013). The simplest form is Stochastic Gradient Descent with Momentum (SGDM). SGDM aims to minimize the training loss L(x) given noisy gradient oracle ((x), which usually realized by evaluating the gradient at a randomly sampled mini-batch from the training set. Specifically, let B be the learning rate and momentum coefficient, then SGDM can be stated as:
gk (xxx),
m 3mm gk.
Øk+1 D 7+
(I)
where gk. my 3g the gradient, momentum buffer, and parameter vector at step k.
For typical choices of 3 E (0, 1). the momentum buffer can be interpreted as an exponential moving Based on this interpretation, Polyak (1964;
average of past gradients, i.e., - - 1987); Rumelhart et al. (1987) argued that momentum is able to cancel out oscillations along high- curvature directions and add up contributions along low-curvature directions. More concretely, for strongly convex functions without any noise in gradient estimates, Polyak (1964; 1987) showed that adding momentum can stabilize the optimization process even when the learning rate is so large that can make vanilla gradient descent diverge, and thus momentum accelerates the convergence to minimizers by allowing using a larger learning rate.
-
In deep learning, however, the random sampling of mini-batches inevitably introduces large amount of stochastic gradient noise, which sometimes dominates the true gradient and may become the main source of training instability. As the above convergence results solely analyze the noiseless case, it remains unclear in theory whether momentum can likewise stabilize the stochastic optimization process in deep learning.
To understand the benefit of momentum in stochastic optimization, several prior studies (Bottou et al., 2018; Defazio, 2020; You et al., 2020) speculate that averaging past stochastic gradients through

---

Published as a conference paper at ICLR 2024
momentum may reduce the variance of the noise in the parameter update, thus making the loss decrease faster. To approach this more rigorously. Cutkosky and Orabona (2019) proposed a variant of SGDM that provably accelerates training by leveraging the reduced variance in the updates.
Nevertheless, for SGDM without any modifications, past theoretical analyses in the stochastic optimization of convex and non-convex functions typically conclude with convergence rate that is comparable to that of vanilla SGD, but not faster (Yan et al., 2018; Yu et al., 2019; Liu et al., 2020; Sebbouh et al., 2021; Li al., 2022a). Besides, there also exist simple and concrete instances of convex optimization where momentum does not speed up the convergence rate of SGD, even though it is possible to optimize faster with some variants of SGDM (Kidambi et al., 2018). This naturally raises the following question on the true role of momentum:
Does noise reduction in SGDM updates really benefit neural network training?
To address this question, this paper delves into the training regime where the learning rate is small enough to prevent oscillations along high-curvature directions, yet the gradient noise is large enough to induce instability. This setting enables to concentrate exclusively on the interplay between momentum and gradient noise. More importantly, this training regime is of practical significance as in many situations, such small-batch training from scratch fine-tuning pre-trained model, the optimal learning rate is indeed relatively small (Liu et al., 2019; Malladi et al., 2023).
Main Contributions. In this paper, we present analyses of the training trajectories of SGD with and without momentum, in the regime of small learning rate. We provide theoretical justifications of a long-held belief that SGDM with learning rate 7 and momentum 3 performs comparably to SGD with learning rate 77 = 123 (Tugay and Tanik, 1989; Orr, 1996; Qian, 1999; Yuan et al., 2016; Smith et al., 2020). This finding offers negative evidence for the usefulness of noise reduction in momentum. Additionally, this also motivates us to reformulate SGDM in Definition 2.3 so SGDM and SGD perform comparably under the same learning rate 1/. which in turn simplifies our analysis.
More specifically, given a run of SGDM, we show that vanilla SGD can closely track its trajectory in the following two regimes with different time horizon:
Regime I. Training with SGD and SGDM for O(1/n) steps where the scaling of gradient noise covariance can be as large as O(1/)). Specifically, Theorem 3.5 shows that SGD and SGDM are O( O(vn/(1 Blo-ee each other in the sense of weak approximation, where n) are the learning rate and momentum coefficient under the notation of Definition 2.3. Our analysis not only includes the classical result that both SGD and SGDM converge Gradient Flow in (11/T) steps where the stochastic gradient is sampled from bounded distribution independent of 17. but also covers the regime of applying Linear Scaling Rule (Goyal et al., 2017), where one decreases the learning rate and batch size at the same rate, so the noise covariance increases inversely proportional to 7), and in this case both SGD and SGDM converge to a Stochastic Differential Equation (SDE). Our results improve over previous analysis (Yuan et al., 2016; Liu et al., 2018) by avoiding underestimating the role of noise when scaling down the learning rate, and provide rigorous theoretical supports to the scaling claims in Smith et al. (2020); Cowsik et al. (2022). Technically we introduce an auxiliary dynamics yı (Equation (15)) that bridges SGDM and SGD.
Regime II. Training with SGD and SGDM for O(1nn²) steps for overparametrized models where the minimizers of the loss connect as manifold and after reaching such a manifold, the gradient noise propels the iterates to move slowly along it. Theorem 4.5 shows that SGD and SGDM follow the same dynamics along the manifold of minimizers and thus have the same implicit bias. The implicit bias result of SGD is due to Katzenberger (1991); Li et al. (2021b) whose analysis does not apply to SGDM because its dynamic depends hon-hooogeneuusly on 17. Our proof of Theorem 4.5 is non-trivial in carefully decomposing the updates.
In Section 5, we further empirically verify that momentum indeed has limited benefits for both optimization and generalization in several practical training regimes, including small- to medium- batch training from scratch on ImageNet and fine-tuning RoBERTa-large on downstream tasks. For large-batch training, we observe that SGDM allows training with a large learning rate, in which regime vanilla SGD may exhibit instability that degrades the training speed and generalization. The observations are consistent with previous empirical studies on SGDM (Kidambi et al., 2018; Shallue et al., 2019; Smith et al., 2020). We argue that the use of large learning rate makes the weak
2

---

Published as a conference paper at ICLR 2024
approximation bound O n/(1 (1--1)) B)) loose: running SVAG (Li al., 2021a), SDE simulation method for both SGD and SGDM, shrinks or even eliminates the performance gain momentum.
Finally, we highlight that our results can also have practical significance beyond just understanding the role of momentum. In recent years, the GPU memory capacity sometimes becomes a bottleneck in training large models. As the momentum buffer costs expensive storing the entire model, has raised much interest when is safe to remove momentum (Shazeer and Stern, 2018). Our work sheds light on this question by formally proving that momentum only provides marginal values in small learning rate SGD. Furthermore, our results imply that within reasonable range of scales the final performance is insensitive to the momentum hyperaarametrization, thereby provide support to save the effort in the extensive hyperparameter grid search.
2 PRELIMINARIES
Consider optimizing loss function (() - Σ () where L, R2d + corre- sponds to the loss the i-th sample. We 0 to indicate parameters along general tra- jectory. In each step, we sample a random minibatch S E], and compute the gradient of the minibatch loss (o) [8] EieB (0) get the following noisy estimate of VL(0), i.e.,
(O) = 15 EiEB (O). It is easy to check that the noise covariance matrix of (L)(O), namely EB ( (0) (((O))( VZ(0)) scales proportionally to [8] Motivated by this, Malladi et al. (2022) abstracts V (O) sampled from noisy gradient oracle where the noise covariance only depends on a scale parameter.
1
Definition 2.1 (NGOS, Malladi et al. (2022)). A Noisy Gradient Oracle with Scale Parameter (NGOS) is characterized by a tuple Go -C,, Σ, Zz). For scale parameter a > 0, Go takes as input I and returns g VL(0) + dæ, where VL(0) is the gradient of L at 0 and is the gradient noise drawn from the probability distribution 2, (0) with mean zero and covariance matrix ( E()) independent of the noise scale a. Slightly abusing the notation, also Gr to denote the distribution of given a and f
In the above minibatch setting, we have noise scale [ह] Generally, we invoke NGOS with bigger for smaller magnitudes of the learning rates. Such scaling is in compliance with the Linear Scaling Rule (Goyal et al., 2017) and is discussed further after Proposition 2.4. We now instantiate the SGD and SGDM trajectories under this noise oracle.
Definition 2.2 (Vanilla SGD). Given stochastic gradient oracle G., SGD with the learning rate schedule {77 updates the parameters 2І RRd from initialization øp, as
žk+1 = 2A Tj
9I: ~ G.
Definition 2.3 (SGD with Momentum/SGDMM, Given oracle G.. SGDM with the hyperparameter schedule {T%,, 3.) where B& E (0, 1). updates the parameters 24 E IRd from (Mo, zo). as
m 3,m, + - 3x)gs,
0b++ = Z2. - 7m mg
(xx).
(3)
Notice that the formulation of SGDM in Definition 2.3 is different from (1) that sometimes appear in previous literature. In our results, Definition 2.3 offers comparison between SGDM and SGD. An easy conversion is given by rewriting Equation (3) as:
natural parameterization for the
Uk
Ft = F1. m)((1 3g)g + 3k 7
(x x1-1))
Then setting
2 and Bx 3 recovers the form of (1).
Modeling the gradient noise as an NGOS gives us the flexibility to scale the noise in our theoretical setting make the effect of noise non-vanishing small learning rate training, observed in Proposition 2.4, variant of the standard gradient descent lemma for SGD.
Proposition 2.4 (Descent Lemma for SGD). Given 2k. the expected change of loss in the next step is [[[((z-))] = ((zx)
|+++=)11 + (on) tr( ( + +
1
+o(?) (on))).
descent
noise-induced
curvature-induced
3

---

Published as a conference paper at ICLR 2024
Proposition 2.4 highlights noise-induced and revature-induceer factors that prevent the loss to decrease. For regular loss functions and small learning rates, the following phenomenon are expected.
In O(n=1) steps. only for O(1/ the loss is guaranteed to decrease for small n. during which the curvature-induced factor accumulates vanishing o(1) impact as v -> 0. For a (1)) the noise-induced impact on the same order as the descent force and will not be vanishing the training curve, noise affects the curve similarly across different learning rates, For a = o(1/ (() the noise-induced impact will vanishing as -
Assume tr when the impacts balance each other, and the noise-induced impact is significant until (( steps of updates.
is non- con-siiih as 17 + 0. The loss plateaus at value (
Inspired by these observations, we are interested in studying the behavior of SGDM in two regimes, for ((1 (1) steps of update with 1/ and for O steps of update with I 1. The two regimes capture the common practices where people use noisy small-batch updates to train a model from scratch and then reduce the noise-induced impact after the training loss plateaus (usually by annealing the learning rate) in pursuit of model that converges and generalizes better.
3 WEAK APPROXIMATION OF SGDM BY SGD IN O(1/n) STEPS
Next, we will present our main theoretical results on SGDM with small learning rates. In this section, we show that O(1 (1) steps. SGD approximates SGDM the sense of Definition 3.1 for o S 1/ /T. The next section studies SGDM over a longer training horizon (0(1/1?) steps) to characterize the coinciding implicit regularization effects of SGDM and SGD.
3.1 A WARM-UP UP EXAMPLE: THE VARIANCE REDUCTION EFFECT OF MOMENTUM
Intuitively, momentum makes the SGD update directions less noisy by averaging past stochastic gradients, which seems first glance contradict result that the distribution of SGD and SGDM are approximately the same. However, the apparent discrepancy is a consequence that, by carrying the current gradient noise subsequent steps, the updates of SGDM have long-range correlations
For instance, we consider the case where the stochastic gradients are i.i.d. gaussian as gbe (N(c,o21) for a constant vector c. We compare SGD and SGDM trajectories with hyperparameter 7 = 1) and 3x = and initialization žy Ig and mg Nec,, 1-g I). The single-step updates are
Zk+1 Zk -1]9 (n-r,,, 121)).
1 3 0k+1 26 = -1 = -n)(ght mo + Σ 3h 3h -1- ) ~ N(-ne, -
1 3 d21)
Therefore, the variance of each single-step update is reduced by factor of 1-8 which implies larger I+3"
momentum generates smoother trajectory. However, we are usually more interested in tracking the final loss distributions induced by each trajectory. The distributions of after k steps are
2i ~ A (z0 knc, ox7"0²I):
20 - m7 1 - 8k
1-1
~ NN(2) knc, 2n - 0² 1 3b
1 3 my
11
B2 1
==0
Notice that the variance of the final endpoint is only different by 12,8m 1-
1-88 T-gr, which is
2y²a"
bounded regardless of The variance of 2 is increased at rate o²² per step, which is significantly 1-2 larger than the per step update variance T+.3 102 This is a consequence the positive correlation of momentum updates that contributes to the variance of the trajectory in total.
observe that SGD and SGDM trajectories have different levels of turbulence induced by the different per step update variances. In some cases covered by our main result Theorem 3.5, the SGD and SGDM trajectories even exhibit different asymptotic behaviors in the limit
Furthermore,
4

---

Published as a conference paper at ICLR 2024
7 - O. For instance, when 11 1/31/ = =1 1 7//44 if we track the trajectory at k. tn steps for
constant & >> 0, 2 ~ (z0 - ic, 1n) and æ ~ NZ - tc, bnj1/$1 2n) -(1-8- vanishing variance 77 + O. While both trajectoriess converge the straight line zo 12 in the limit, when we measure their total length,
1+3 I) with
E
Æk+1
zull2 2 1n-11/EE-Ni00 |lEll + oc,
1725/12
[|c|| + E Σ [x+1 - xxl12 S 7/ 1² 2 - 17/8//" t||c|
We observe that as 17 gets smaller, the SGD trajectory becomes more and more turbulent its length goes unbounded, while the SGDM trajectory becomes more and more smooth, though they have the same limit. Consequently, the turbulence of the training curve may not faithfully reflect the true stochasticity of the iterates as a whole, and may not be indicative of the quality of the obtained model with different choices of the momentum.
3.2 MAIN RESULTS ON WEAK APPROXIMATIONS OF SGDM
The above warm-up example reminds us that SGD and SGDM trajectories may have different appearances that are irrelevant to the final distributions of the outcomes, which in our concern is mostly important. Therefore we need to talk about trajectory approximations in the correct mathematical metric. For our main results, we introduce the notion of weak approximations between two families of trajectories, inspired by (Li et al., 2019). We say function g(x) R R" has polynomial growth if there are constants ky, k2 > such that ([g(2))l12 S k (1 + || [arl(22) 1122), ). Wa E Rd. and we say function has all-order polynomial-growee polynomial growth for all a 2 0.
derivatives if g is C and Vag has
Definition 3.1 (Order-y Weak Approximation). Two families of discrete trajectories æy and yı are weak approximations of each other, if there is Thhr > 0 that for any T > O, any function of all-order polynomial-growing derivatives, and any v) S Yan there is constant C,,r independent of r) that
max k= [/n)
[= - (b( & CL.T 1f"
Weak approximation implies that æy and y have similar distributions at any step l S T/n even when &t 77 - 0, and specifically in the deep learning setting it implies that the two training (testing) curves are similar.
In the small learning rate cases, we use the big-O notations to specify the order of magnitudes as the learning rate scale - O. Consider a SGDM run with hyperparameters {(Tkk B&x 1:20. Let the magnitudes of the learning rates be controlled by scalar 7) 77 O Furthermore, to capture the asymptotic behaviour of the the momentum decay Sa. we set an index a 2 0 so that the decay rate of the momentum is controlled as 1 BB (1)²). schedule while a > 0 corresponds to schedule where 3y is closer to for smaller learning rates. Formally, we introduce the following denotation.
corresponds a constant-scale decay
Definition 3.2. A (family of) hyperparameter schedule {mk, Be is scaled by y with index if there are constants 7max Amin and Amax- independent of 17, such that for all kt
5 m2/n
Tmaxx
0< Amin < (1 - Bx)/n" s Amax < 1.
We need the boundedness of the initial momentum for the SGDM trajectory to start safely. Assumption 3.3. For each 2 1, there is constant C. 2 0 that E[I[mo I Cmi
Following Malladi et al. (2022), we further assume that the NGOS satisfies the below conditions, which make the trajectory amenable to analysis.
Assumption 3.4. The NGOS 9, = (C, Σ, 2,) satisfies the following conditions.
Well-Behaved: VL is Lipschitz and - smooth; 1/ is bounded, Lipschitz, and smoothh; all partial derivatives of VL and 1// up to and including the third order have polynomial growth.
5

---

Published as a conference paper at ICLR 2024
2. Bounded Moments: For all integers m 2 1 and all noise scale parameters 0, there exists constant C2m (independent of () such that 2.) (0) [| .m C2m (1 10112) v0 TR..
Given the above definitions, we are ready to establish our main result.
Theorem 3,5 (Weak Approximation of SGDM by SGD). Fix the initial point 2o E (0, 1), and an NGOS satisfying Assumption 3.4. Consider the SGDM update 2xK with schedule {(TK, 3x]]k21 scaled by y with index a, noise scaling a S 17 11/2 and initialization (mo; Do) satisfying Assumption 3.3, then æ% is an order-(1 /2 weak approximation (Definition 3.1) of the SGD trajectory z with initialization z Fo, noise scaling and learning rates 1) = > 1)s - 3- Bkk).
Specifically, for constant schedule where (Nk n) 3x 5), nR ny. In this case, SGD and SGDM with the same learning rate weakly approximate each other at distance O( nn/nn(11 B)-
The theorem shows that when the learning rates has a small scale 1), under reasonable momentum decay and reasonable gradient noise amplification. the outcomes obtained by SGDM and SGD are close in distribution over ((1/n) steps. Specifically at the limit 1) + 0. the outcomes will have the same distribution. Following Li et al. (2019), if a = 1/ /n, then the limiting distribution can be described by the law of the solution X, to an stochastic differential equation (SDE):
dX; = -A VL(())dt + X 1/2 (X.
under brownian motion W. and some rescaled learning rate schedule A,. If & 1//n,, however, the limit will in general be the solution to the gradient flow ODE dX,- = - V dt and the impact of gradient noises will be vanishing.
The theorem is built upon an infinite learning rate schedule k = 1, o0. In the case where we wish consider a finite schedule, we can apply the theorem after schedule extension by infinitely copying the hyperparameters at the last step. Besides, the theorem requires a E [0. 1), and the approximation grows weaker as a approaches 1. At a the two trajectories are no longer weak approximations of each other and have different limiting distributions. a > 1 yields undesirable hyperparameter schedules where excessively heavy momentum usually slows down or even messes up optimization. Further details are discussed in Appendix C.2.
4 THE LIMIT OF SGDM AND SGD ARE IDENTICAL IN O(1nn²) STEPS
In this section, we follow the framework from Li al. (2021b) study the dynamics of SGDM when the iterates are close to some manifold of local minimizers of C. Former analyses (e.g.. Yan et al. (2018)) suggest that on regular functions, SGDM and SGD will get close to a local minimizer in (1111)) steps, at which point the loss function plateaus and the trajectory random walks near the local minimizer. If the local minimizers connect an manifold in the parameter space, then the updates accumulate into drift inside the manifold over (1111²) steps. Li et al. (2021b) shows that under certain circumstances, the drift induces favorable generalization properties after the training loss reaches its minimum, by leading minima with smaller local sharpness.
Therefore, by investigating this regime, we hope to detect the value of momentum in late-phase training, especially that concerning extra generalization benefits. Yet in this section, we show that when y + 0, the limiting dynamic of SGDM admits the same form as that of SGD, suggesting that momentum provides no extra generalization benefits over least (1(r?2) steps of updates.
4.1 PRELIMINARIES ON MANIFOLD OF LOCAL MINIMIZERS
We consider the case of optimizing an over parameterizee neural network, where usually the minimiz- ers of the loss L form manifolds. Let I' be a region of local minimizers that SGD can reach, and we will work mathematically in I to see whether adding momentum changes the dynamical behaviors. Assumption 4.1. C is smooth. I' is a (d .M) dimmensional submanifold of H for some integer 0 < M < d. Moreover, every x E I is a local minimizer of L with V v((x) = 0 and rank
M.
We consider neighborhood Or of I that I' is attraction set of Op under VL. Specifically, define the gradient flow under VL by f(x,t) = x So VL(o(z, s))ds for any æ E Rd and 2 We further define gradient projection map associated with VL as ((x)) = lim ((xr,

---

Published as a conference paper at ICLR 2024
Assumption 4.2. From any point E Or, the gradient flow governed by V converges to some point in I, i.e., (x) is well-defined and ((x) E r.
It can be shown that for every æ E I, (+))) is the orthogonal projection onto the tangent space of I at æt Moreover, Li et al. (2021b) proved that for any initialization ao E Or, a fixed learning rate
schedule 7/k 7/ and any / 0, time-rescaled SGD iterates 2///1² converges in distribution Zt, the solution to the following SDE. We will refer to Ze as the slow SDE.
Z. - = + L 00 (0( ddW, | 8 Z.. (Z.) (Z.)
Notice that Zz always stays in I with Zo= = ((xo)) Though 2n Do is not in r, for any t >>> 0 the limit of 2 2Le/r3] will fall onto
4.2 ANALYSIS OF SGDM VIA THE SLOW SDE
As the limiting dynamics of SGD iterates is known as above, momentum modifies this limit. turns out that within a fairly free range of hyperparameters, SGDM also has the same limiting dynamics. Specifically, for family of hyperparameters {(nb) (g(")) }221 scaled by series of scalars (7(*) with index a (Definition 3.2, lim,- r7(*) 0), will show that the hyperparameter schedules converge as n - oo, then SGDM iterates will also converge into the limiting dynamics of SGD with the limiting learning rates, irrelevant of the momentum decay factors.
curious about whether adding
Similar to the setting in Li et al. (2021b), we consider a fixed time rescaling stipulate that the schedule v/h" + (n) de as nl - Do for rescaled schedule in continuous time
k0))()))². We
[0, 7] R+ In the special case m7()) r7("), it is clear that X = 1, and the setting of Equation (4) is recovered. Formally we introduce
Assumption 4.3. [0, R+ has finite variation, and
(T/6)(1)))
lim (7()
Σ
m (n)
- (n(") Akcntoojal = 0.
K=O
Assumption 4.4 (Bounded variation). There is a constant Q independent of n such that for n,
()))))))
(T/()(1)))
Σ
(n)
(2) s (n)"),
|a( - -( < (1)
1
In this general regime, we define the slow SDE on to admit the following description:
X, ((x)) + S. ("" /.' X? 8 (EX.))
(5)
Both SGDM and SGD converge the above slow SDE on as summarized in the following theorem.
Theorem 4.5. Fix the initialization xo = 20 E Or and any a E (0, 1), and suppose the initial
momentum mo satisfies Assumption 3.3. For n 2 1, let (1) (n) ()"))) Ju21 be any hyperparameter schedule scaled by n7(") with index satisfying Assumptions 4.3 and 4.4. Fix the noise scale ((a) 1.
Under Assumptions 4.1 and 4.2, consider the SGDM trajectory (n)) } with schedule {(T) (m) 3 {(m))) }}
initialization (Iu, m.), and the SGD trajectory (2(")) with schedule (m(a), initialization zu Dg.
Suppose the slow SDE defined in (5) has global solution Xx)t20> then -> with nf()) (n) both x [t/(n)1)ya)) and (m)
[ converge in distribution to X
The proof of Theorem 4.5 is inspired by Calzolari and Marchetti (1997). Similarly in this regime, the momentum process m behaves like an Uhlenbeck- -onnteii process with O(7²) mixing variance, so the per-step variance will be significantly smaller than that of SGD is in Section 3.1. To prove the result, a more careful expansion of the per-step change & ((x++1) 4b (xxx) ) is needed. The proof is detailed in Appendix D.
7

---

Published as a conference paper at ICLR 2024
80%
80.0%
SGDM
3.2
Best Learning Rate for SGDM
Top-1 Validation Accuracy
Top-1 Validation Accuracy
SGD
60%
1166
76.0%
40%
74.0%
0.8
20%
SGDM
72.0%
SGD
0%
0.4
70.0%
0
20
40
60
80
1024
2048
4096
8192
epoch
batch size
(a) SGDM and SGD with batch size 1024
(b) SGDM v.s. SGD across batch sizes
Figure SGDM performs comparably to SGD training ResNet-50 on ImageNet with smaller batch sizes (e.g.. 1024), and outperforms SGD significantly at larger batch sizes.
5 EXPERIMENTS
In the previous sections, we conclude in theory that SGDM and SGD have similar performances within noisy thort-horizon or general long-horizon training. While our theoretical results mostly work for learning rates that are asymptotically small, in this section, verify that momentum indeed has limited benefits in practical training regimes where the optimal learning rate is finite but not very large. We defer some of the additional details of this section to the appendix.
5.1 MOMENTUM MAY INDEED HAVE MARGINAL VALUE IN PRACTICE
ImageNet Experiments. First, we train Re on ImageNet across batch sizes. Following the experimental setup in Goyal et al. (2017), we learning rate schedule that starts with 5-epoch linear warmup to the peak learning rate and decays it at epoch #30, #60, #80. For SGDM (1), we use the default value of 0.9, and grid search for the best learning rate y over 0.1 2* E Z). Then we check whether vanilla SGD with learming rate T- can achieve the same performance as SGDM. Consistent with previous empirical studies (Shallue et al., 2019; Smith et al., 2020), we observed that for training with smaller batch sizes, the optimal learning rate SGDM small enough so that SGD can perform comparably, though SGDM can indeed outperform SGD at larger batch sizes.
Language Model Experiments. fine-tuning a pre-trained model, small learning rate is also preferable to retain the model's knowledge learned during pre-training. Indeed, we observe that SGD and SGDM behave similarly in this case. We fine-tune RoBERTa-large (Liu et al., 2019) on diverse tasks (SST-2 (Socher et al., 2013), SST-5 (Socher et al., 2013), SNLI (Bowman et al., 2015), TREC (Voorhees and Tice, 2000), and MNLI (Williams et al., 2018)) using SGD and SGDM. We follow the few shot setting described in (Gao al., 2021; Malladi et al., 2023), using a grid for SGD based on (Malladi et al., 2023) and sampling 512 examples per class (Table 1). Additional settings and trajectories are in Appendix E.
Table 1: SGD and SGDM for fine-tuning RoBERTa-large on 5 tasks using 512 examples from each class (Gao et al., 2021; Malladi et al., 2023). Results are averaged over random subsets of the full dataset. These findings confirm that SGD and SGDM approximate each other in noisy settings.
Task
SST-2
SST-5
SNLI
TREC
MNLI
Zero-shot
79.0
35.5
50.2
51.4
48.8
SGD
94.0 (0.4)
55.2 (1.1)
87.7 (0.3)
97.2 (0.2)
84.0 (0.3)
SGDM
94.0 (0.5)
55.0 (1.0)
88.4 (0.6)
97.2 (0.4)
83.7 (0.8)
5.2 INVESTIGATING THE BENEFIT OF MOMENTUM IN LARGE-BATCH TRAINING
The ImageNet experiments demonstrate that momentum indeed offers benefits large-batch training when the optimal learning rate is relatively large. We large-batch experiments CIFAR-10
8

---

Published as a conference paper at ICLR 2024
SGDM and SGD for ResNet-32 32 on CIFAR-10
90%
95.0%
Test Accuracy
80%
70%
90.0%
SGD
60%
S6DM
14000
15000
t=1
50%
40%
0
2000
4000
10000 12000 Effective Step Step
6000
8000
14000
Figure 2: Standard SGDM achieves higher test performance than SGD (see f= 1), but the two trajectories get closer when reducing the curvature-induced term with SVAG (i.e., increasing the value of L, see Definition 5.1 and Lemma 2.4). These experiments confirm our theoretical findings that SGD and SGDM approximate each other when the gradient noise is the primary source of instability. We use batch size B- 512 with two learning rate decays by factor of 0.1 epochs 80 and 120. We grid search to find the best learning rate for SGDM (1) = 0.2) and then use it to run SGD and SGDM with SVAG. We use 0.9 for SGDM. Additional experimental details are in the appendix.
to provide empirical evidence that this benefit not due to the noise reduction effect and is marginal when SGD is well-approximated by its SDE. To do this we apply SVAG (Li et al., 2021a) to control the noise scale in the gradient oracle in both SGDM and SGD updates.
Definition 5.1 (SVAG). With any e > 0, SVAG transforms the NGOS Go (f, 2,) (Defini-
tion 2.1) into another NGOS -vi = (f,, Σ, Vvi)) with scale Vla. For an input 0, 9ta returns
" (4)g) (f)gg2 where g1, 92 Go (0) and re) {1 + (-1)² 220 (-1) 1). Z /ta is defined to ensure 9 has the same distribution as ff(0) + Vlax when z ~ (0).
In our experiments, given an original SGD run with learning rate 1) for K steps, we perform new SGD run with learning rate n/E, SVAG(E) as the gradient oracle, for KL steps. The new SGD trajectory will be closer to its SDE approximation l - +oo (Li et al., 2021a). We also apply the same modifications to SGDM runs (3 the momentum decay factor is unmodified). In another view, applying this modification makes the total accumulated noise-induced impact and descent force (Proposition 2.4) qualitatively stay on the same scale, while the accumulated curvature-induced impact is reduced by a factor of l.
l gets larger, and converge its SDE as
We train a ResNet-32 (He et al., 2016) on CIFAR-IO (Krizhevsky et al.) with batch size B 512. We first grid search to find the best learning rate for the standard SGDM (E = 1), and then we perform SGD and SGDM with that same learning rate (7)- 7/ 17 in the formulation Equations (2) and (3)). Then we modify both processes with different values. The results are summarized in Figure 2. We observe that while standard SGDM outperforms standard SGD, when we increase the value of f, the two trajectories become closer until SGDM has no evident edge over SGD. The finding corroborates that the benefit of adding momentum mostly due to the alleviation of curvature- induced impacts, but will be marginal in general small-batch or small learning-rate settings when SGD is well- aplp--pppxximtedd by SDE.
6 CONCLUSIONS
This work provides theoretical characterizations of the role of momentum in stochastic gradient methods. We formally show that momentum does not introduce optimization and generalization benefits when the learning rates are small, and we further exhibit empirically that the value of momentum is marginal for gradient-noise-dominite scales. Hence conclude that momentum does not provide significant performance boost in the above cases. Our results further suggest that model performance is agnostic to the choice of momentum parameters over a range of hyperparameter scales.
learning settings with practical learning rate

---

Published as a conference paper at ICLR 2024
REFERENCES
Sébastien Arnold, Pierre-Antoine Manzagol, Reza Babanezhad Harikandeh, Joannis Mitliagkas, and Nicolas Le Roux. Reducing the variance in online optimization by transporting past gradients. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alcéé-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In International Conference on Machine Learning, pages 948-1024.
PMLR, 2022.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by orsstein-uhlenbeck like process. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 483-513 PMLR, 09-12 Jul 2020.
Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-311, 2018. doi: 10. /..1137//61000073.
Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher D. Manning. A large
annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference Empirical Methods in Natural Language Processing, pages 632-642, Lisbon, Portugal, September 2015, Association for Computational Linguistics. doi: 10. 10.18653/1//D15-10755 URL https: //aclanthology wtt///5-1100000
Antonella Calzolari and Federico Marchetti. Limit motion of an equilibrium manifold of a force field. Journal of Applied Probability, 34(4):924-99) 938, 1997.
particle on the
Aditya Cowsik, Tankut Can, and Paolo Glorioso. Flatter, faster: scaling momentum for optimal speedup of sgd. arviv preprint arXiv:2210. 164000 2022.
Ashok Cutkosky and Francesco Orabona. Momentum baased variance reduction in non-convex sgd. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Bue, E. Fox, and R. Gamett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global minimizers. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
Aaron Defazio. Momentum via primal averaging: theoretical insights and learning rate schedules for non-convex optimization. arXiv preprint #XXi:20100004006 2020.
Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.09000 2023.
Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better few-shot learners. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 1836,3830, 2021. doi: 10. N0.18653/1/22221.aaom hww.:/oo22..clllooo.922.2.. URL https: //aclanthology. org/ acw-onnn.......n.........mmo 2.9..
Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang. Implicit regularization in heavy- ball momentum accelerated stochastic gradient descent. In The Fleventh International Confer- ence Learning Representations, 2023. URL https:/ http///pneee.w...eeeeeeeee net http://www.e//orummeemme forumiidri EH9B..
Igor Gitman, Hunter Lang, Pengchuan Zhang, and Lin Xiao. Understanding the role of momentum in stochastic gradient methods. In H. Wallach, H. Larochelle, A. Beygelzimer, F. AAlcé--Buc. E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet in hour. arXiv preprint arXiv: (e06.000677777 2017.
10

---

Published as a conference paper at ICLR 2024
Thomas Hakon Gronwall. Note the derivatives with respect to parameter of the solutions of system of differential equations. Annals of Mathematics. pages 292-296, 1919.
Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local SGD gen- eralize better than SGD? In The Eleventh International Conference on Learning Representations, 2023.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Prateek Jain, Sham M. Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerating stochastic gradient descent for least squares regression. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet, editors, Proceedings of the 31 Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pages 545-604. PMLR, 06-09 Jul 2018.
Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In Kamalika Chaudhuri. Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 9965-10040. PMLR, 17-23 Jul 2022.
Gary Shon Katzenberger. Solutions of a stochastic differential equation forced onto a manifold by a large drift. The Annals of Probability, pages 1587-1628, 1991.
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham M. Kakade. On the insufficiency of existing momentum schemes for stochastic optimization. In International Conference on Learning Representations, 2018.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. CIFAR-10 (Canadian Institute for Advanced Research). URL http: /www. , cs worootooooooooo... www/edddtteee cifar.. html.
Thomas G Kurtz and Philip Protter. Weak limit theorems for stochastic integrals and stochastic differential equations. The Annals of Probability, pages 1035-1070, 1991.
Qianxiao Li, Cheng Tai, and Weinan Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40): 1-47, 2019.
Xiaoyu Li, Mingrui Liu, and Francesco Orabona. On the last iterate convergence of momentum methods. In Sanjoy Dasgupta and Nika Haghtalab, editors, Proceedings of The 33rd International Conference on Algorithmic Learning Theory, volume 167 of Proceedings of Muchine Learning Research, pages 199-717. PMLR, 29 Mar-01 Apr 2022a.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34:12712- 12725, 2021a.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss?-a mathematical framework. In International Conference on Learning Representations, 2021b.
Zhiyuan Li, Tianhao Wang, and Dingli Yu. Fast mixing of stochastic gradient descent with normal- ization and weight decay. Advances in Neural Information Processing Systems, 25:9233-9243 2022b.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. arXiv preprint arXiv:2210.14/99 2022.
Tianyi Liu, Zhehui Chen, Enlu Zhou, and Tuo Zhao. A diffiusion approximation theory of momentum sgd in nonconvex optimization. arXiv preprint
2018.
Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. In H. Larochelle, M. Ranzato. R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18261-18271. Curran Associates, Inc., 2020.
11

---

Published as a conference paper at ICLR 2024
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint vrX1v:1907.11692 2019. URL https: /awxiv.ivvvvv.////i///a/////aivl /ttt::///ax....cz///////111/ 11692.pdb pdf.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of nor- malization layers: Sharpness reduction. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https: wpwp/eweeeeeeeeeeeeeeeeeeeeeeeem net xpppxuxxXX.xx
Jerry Ma and Denis Yarats. Quasi-hyperbollic momentum and adam for deep learning, In International Conference on Learning Representations, 2019.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A Corne--ased view of language model fine-tuning, 2023.
Genevieve Beth Orr. Dynamics and Algorithms for Stochastic Search. PhD thesis, USA, 1996. UMI Order No. GAX96-08998.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, bigh-peeformance deep learning library. Advances in neural information processing systems, 32, 2019.
Maximilian Plattner. On sgd with momentum. page 60, 2022. URL http::/anfooccieece epfl. ch/ recoord/25399
Boris T. Polyak. Introduction Optimization. Optimization Software, Inc., 1987.
B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computa- tional Mathematics and Mathematical Physics, 4(5):1-17, 1964. ISSN 0041-5553.
Ning Qian, On the momentum term in gradient descent learning algorithms. Neural networks, 12(1): 145-151, 1999.
David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning Internal Representations by Error Propagation, pages 318-362. 1987.
Othmane Sebbouh, Robert M Gower, and Aaron Defazio. Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball. In Mikhail Belkin and Samory Kpotufe, editors, Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 3935-3971. PMLR, 15-19 Aug 2021.
Christopher J. Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl. Measuring the effects of data parallelism neural network training. Journal of Machine Learning Research, 20 12): 2019.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 4506-4604, PMLR, 10-15 Jul 2018.
Leslie N Smith. A disciplined approach to neural network hyper-parameters: Part 1 1-learning rate, batch size, momentum, and weight decay. arXiv preprint arXiv:1803.09820, 2018
Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9058-9067. PMLR, 13-18 Jul 2020.
12

---

Published as a conference paper at ICLR 2024
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng. and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. 2013. URL https::///........t....tt https/://caaaaahooooogooo
org/ //gg//D13 113-11700001011 pdf.
yy Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In Sanjoy Dasgupta and David McAllester, editors, Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pages 1139-1147, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR.
Lionel Tondji, Sergii Kashubin, and Moustapha Cisse. Variance reduction in deep learning: More momentum is all you need. arXiv preprint arXiv:2111. 2021.
Mehmet Ali Tugay and Yalçin Tanik. Properties of the momentum Ims algorithm. Signal Processing, 18(2):117-127, 1989. ISSN 0165-1684.
Ellen M Voorhees and Dawn M Tice. Building a question answering test collection. In the 23rd annual international ACM SIGIR conference on Research and development in information retrieval, 2000.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729 2022.
Ward Whitt. stochastic-proceet limits: introduction stochastic-process limits and their applica- tion to queues. Space, 500:391-426, 2002.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. 2018. URL https::/aac.amtbooorry org/ N18-1101. pdf.
Zeke Xie, Li Yuan, Zhanxing Zhu, and Masashi Sugiyama. Positive-negative momentum: Manip- ulating stochastic gradient noise to improve generalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages
PMLR, 18-24 24 Jul 2021.
Yan Yan, Tianbao Yang, Zhe Li, Qihang Lin, and Yi Yang. A unified analysis of stochastic momentum methods for deep learning. In Proceedings of the Seveeth International Joint Conference International Joint Conferences on Artificial
on Artificial Intelligence, IJCAI-18, pages Intelligence Organization, 2018.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training BERT in 76 minutes. In International Conference on Learning Representations, 2020.
Hao Yu, Rong Jin, and Sen Yang. On the linear speedup analysis of communication efficient momentum SGD for distributed non-convex optimization. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 7184-7193 PMLR, 09-15 Jun 2019.
Kun Yuan, Bicheng Ying, and Ali Sayed. On the influence of momentum acceleration on online learning. The Journal of Machine Learning Reseurch, 17(1):6022-6677 2016.
13

---

Published as a conference paper at ICLR 2024
A RELATED WORKS
The role of momentum in optimization. The accelerating effect of some variants of momentum has been observed in convex optimization (Kidambi et al., 2018) and linear regression (Jain et al., 2018) with under specialized parametrizations. Smith (2018) pointed out that momentum can help stabilize training. but the optimal choice of momentum closely related to the choice of learning rate. Plattner (2022) later empirically established that momentum enlarges the learning rate but does not boost performance. Arnold et al. (2019) argued using quadratic example that momentum might not reduce variance as the gradient noise in each would actually be carried over to future iterates due to momentum, Tondji et al. (2021) showed that the application of multi-momentum strategy can achieve variance reduction in deep learning.
Defazio (2020) proposed a stochastic primal averaging formulation for SGDM which facilitates a Lyapunov analysis for SGDM, and one particular insight from their analysis is that momentum may help reduce noise in the early stage of training but is no longer helpful when the iterates are close to local minima. Xie et al. (2021) showed that under SDE approximation, the posterior of SGDM is the same as that of SGD. Jelassi and Li (2022) proved the generalization benefit of momentum in GD in a specific setting of binary classification, by showing that GD+M is able learn small margin data from the historical gradients in the momentum. A stronger implicit regularization effect of momentum in GD is also proved in Ghosh et al. (2023). Fu al. (2023) showed that momentum is beneficial in deferring an "abrupt sharpening" phenomenon that slows down optimization when the leamnnng rate large.
Convergence of momentum methods. Homentum-baseed methods do not tend to yield faster convergence rates in theory. Yu al. (2019) showed that distributed SGDM can achieve the same linear speedup as ditributed SGD in the non-convex setting. Also in the non-convex setting, Yan et al. (2018) showed that the gradient norm converges at the same rate for SGD, SGDM and stochastic Nestrov's accelerated gradient descent, and they used stability analysis to argue that momentum helps generalization when the loss function is Lipschitz. Under the formulation of quasi-hyperrolic momentum (Ma and Yarats, 2019). Gitman al. (2019) proposed another unified analysis for momentum methods. Liu et al. (2020) proved that SGDM converges as fast as SGD for strongly convex and non-convex objectives even without bounded gradient assumption. Using a iterate- averaging formulation, Sebbouh et al. (2021) proved last-iterate convergence of SGDM in both convex and non-convex settings. Later, Li al. (2022a) showed that constant momentum can lead to suboptimal last-iterate convergence rate and increasing momentum resolves the issue. Smith (2018); Liu et al. (2018) provided evidence that momentum helps escape saddle points.
Characterizing implicit bias near manifold of local minimizers A recent line of work has studied the implicit bias induced by gradient noise in SGD-type algorithms, when iterates are close some manifold of local minimizers (Blanc et al., 2020; Damian et al., 2021; Li et al., 2021b). In particular, Li et al. (2021b) developed framework for describing the dynamics of SGD via slow SDE on the manifold of local minimizers in the regime of small learning rate (see Appendix D for introduction). Similar methodology has become a powerful tool for analyzing algorithmic implicit bias and has been extended to many other settings, including SGD/GD for models with normalization layers (Lyu et al., 2022; Li et al., 2022b), GD in the edge of stability regime (Arora et al., 2022), Local SGD (Gu et al., 2023), sharpness-awe aware minimization (Wen al., 2022), and pre-training for language models (Liu et al., 2022). Notably, Cowsik et al. (2022) utilized the similar idea to study the slow SDE of SGDM study the optimal scale of the momentum parameter with respect to the learning rate, which has focus different from our paper.
14

---

Published as a conference paper at ICLR 2024
B ADDITIONAL PRELIMINARIES
B. SGDM FORMULATIONS
Recall Definition 2.3 where we defined SGDM as following:
m 8g m + (1- 1k)kggk
3k+1 = 3t. m+++++
(6)
Our formulation is different from various SGDM implementations (e.g. pytorch (Paszke al., 2019))
which admit the following form:
Definition B.1 (Standard formulation of SGD with momentum). Given a stochastic gradient
gs (3x), SGD-Standard with the hyperparameter schedule (7ks Mis mi) momentum updates the parameters &g E IRd from (ma, Lo) as
m = 11 + (1 Tk gg.
Øk++ Dk 7/h m+
where Th E 1).
Notice that with 7k = 0, μı = 3 and = 7, the standard formulation recovers Equation (1) that has commonly appeared in the literature. The standard formulation also recovers Equation (6) with μ4 3g and 7/s 7ik Actually, as we will show in the following lemma, these formulations are equivalent up to hyperparameter transformations.
Lemma B.2 (Equívalence of SGDM and SGD-Standard), Let on be the sequence that ag and
OKk
ak+1
0
Tk) + P*
Then 0 < a141 S I-mm For parameters schedules and initialization
ax+1
Kkk mm4 ri,o 30 8o
Bg
pks
TJk
Oak
0*++1
then 3g and æ follow the same distribution. a by Definition 2.3 and T. by Definition B.
Proof. We shall prove that given F4 = 2ik. 9: = Ik and m. og there and m 1841mgg and Tx+1 Bb+1. This directly follows that a2+1 Tx) = 8k. and
m mgm + Tk)6E
1k my + (1 Tk]gk
(BuT + ax+1(1 (xgggg)))
0i+1
m++1-
0k++
0b+1
3b
Th m+
7-k+1
Then we can prove the claim by induction that 2e and Ty are identical in distribution. More specifically, shall prove that the CDF for (æx,, m&, gk) (æe, m, g) is the same the CDF for (Øx, my, 9k) at (x, an .)). For = 0, the claim trivial. If the induction premise holds for k, then let (Ix+1, M8+1194441 be the one-step iterate in Equation (6) from (æx. m&g 9k.) = (æ. m. g) and (Āx+1,
2 15841,77x+11 m2 11) the one-step iterate in Definition B. from (Øk, mi, gs) = (x, az g). then we know from above that 066 m m and fh-- = 2k+11 and therefore the conditional distribution for (æx+1, m4+1 (k+1) is the same that for (æx+1, Ore 11++17++11 91+ Integration proves the claim for & + 1, and thereby the induction complete.
Therefore, in the rest part of the appendix, we will stick to the formulation of SGDM as Equation (6) unless specified otherwise.
15

---

Published as a conference paper at ICLR 2024
B.2 DESCENT LEMMA FOR SGD AND SGDM
In Proposition 2.4, given the SGD updates {zu} with constant learning rate we have the gradient descent lemma
E [[((zi-))|= - L(2x) = [VC(=)II* + (on) (r((VV))E(Zz)) 2" 2d (V2VV())) +o() (on))))
(9)
force
noise-induced
carvature-nnuceed
(10)
We can get a similar result for the SGDM updates { namely, when 1 1-B< < n.
E[(xx+1) B[(xx) BE V(xx))?
(11)
1
1 ((
+
2 B (V
(12)
n/3
+ 1 - 3
(Sx+1 - 5k) + E + o /(1 - 3), (on))/(111 B)).
(13)
where Sk = [ m. 3 V )(x Notice that beside the terms in the SGD updates, additions an O(n((1 3)) telescoping term and an O extra impact term. The impact of the telescoping term is large for per-step updates but is limited across the whole trajectory. so the difference between SGDM and SGD trajectory is actually depicted by the extra term E[m, V2LV [(xx)]. In the settings where the extra term is bounded, we can bound the difference between training curves of SGDM and SGD.
Brief Proof for proposition 2.4. We assume ||V ZI] is bounded (IV* < C.) for any order of derivatives n = 0, 1, 2, 3 and the trajectories are bounded (E > Ila E |< S Cm.) to simplify the reasoning. the SGD case, Zk+1 - 2 ngh has moments
E |
=
a |1 2m
< 1) 2mm B( /I ((x))))) 2m + 2m |0xl 2m))
Om7 + (on)
Similarly,
E [a
1 2mm
17 2 | || 2m
2m
n 2m E
gh mo +
gh g1- 9)g
8=0
2m-1
k
< 1) 2m22m
g2m(k+1)
|ooll 2m +
Σ 1=0
gk-
Σ
*(1 - -3) mm gh-=(11
||ga
2m
+=00
0 + 2m
Therefore B| =x11"² S
v E El|E--11 - 1 , (ov1)
and E| - -x21" 5
/BllIHHI 24ll 2261 (on) Taylor expansion gives
EL(2++1) [(zx) EVL(2x) (E4+1 - 2x) ) + T22 2x) -0(n""((on))
Expansion yields Equation (9). Similarly,
EC(x++1) = [x(((x)) + V((x)) (#K+1 - IK) + + - x1) T ((z +
=E[[(x+) v ((xx) m + 172
2 m 42 ((( + 0(7 (on)2)
16

---

Published as a conference paper at ICLR 2024
Since m++1 9h + _
1-8 (m+- mx+1),
B
EC((x+11 = [B(((x+) n/V((( gk -VV(((
1- (ma -mx+1)
1/2 1 [((()))( + 1- - (mL -m++++ (
3
+ 2
1133
= -[(((()) [)v((())) 9h - 1 (VL(xx) T m. - vL(xx+1) ( mx+1)
173
(m V2 (< Ik+11 + O( o²n²))
1 1-9
+
112
T 42 [(* +
M
1- (m, - m })] + o(1 (an)
3
2
113
[l((xN) n)V((())) T gh. 1- B
(V((z)) T my - v[(IN+1) mm+1)
2 (gk + 1-8 (mm mx+1)) [(r( + 1 1- -
B
(mi + m++1))] o(17" (on))) 1 3
2
Expansion yields Equation (11).
C THE DYNAMICS OF SGDM IN O(1/n) TIME
In this section will prove Theorem 3.5 the first part of our main result. First we shall take closer examination at update rule in Definition 2.3. Let B... II Br be the product of consecutive momentum hyperparameters for S t, and define Bast = 1 for s > t, we get the following expansion forms.
Lemma C.1.
my = Book m2g
25+1=-1(11 -B)g.
k-1k-1
Σ
naBo. mg
ΣΣ
Σ
Tr Ba)g.
s=0
Proof. By expansion of Equation (3).
Let the coefficients Ca -ΣP=1 nr (2) 3.) so that 23
Bucack mg Ek- C2 No- So tice that Cs.6 is increasing in & but is always upper-bounded as C.. S C2 = E 31r 3.) S (sup; 7)) (sup, 3
T-sup, B. Notice that when we have finite schedule, we can make any extension to infinity steps (e.g. infinitely copy the hyperparameters on the last step), and the reasonings will work for the extended schedule. In this case, though the trajectory 2& not affected by future hyperparameters, the coefficients Ca will depend on our extension. Furthermore, if for any i, Br and 1): about the same scale, then should expect that the difference 6s Cs.k is exponentially small in k. Then we can hypothesize that the trajectory æ should be close to a trajectory of the following form:
sup.
BoCo:s mg
Σ C..ocg.
3
3D
1 1- 3o
which is exactly SGD trajectory given the learning rate schedule C2
To formalize the above thought, we define the averaged learning rate schedule
2
7Js == Ck
TrBk+1:0 3k))
(14)
s
17

---

Published as a conference paper at ICLR 2024
And the coupled trajectory
Uk
mk
(15)
1 3k
k-1
BorKo
370
m.
Σ
9=9
(16)
80
Then there is the following transition:
yk yk 19k-
(17)
ys has an interesting geometrical interpretation as the endpoint of SGDM if we cut-off all the gradient signal from step k. E.g., &oo if we Je+1 g for and update the SGDM from mz).
C.I 10
In this section we provide the proof for Theorem 3.5. Specifically, when a E [0, 1) is the index for the control of the learning rate schedule (Definition 3.2), we hope to show that y/: is close to SGD trajectory 2k: defined by the following for 2
20
(18)
h: ~ (z
(19)
26+11
hi
(20)
Specifically, we use yı as bridge for connecting æy and zg. Our proof consists of two steps.
1. We show that æx and Ys are order- 1 a2/ weak approximations of each other.
2. We show that 2y and yı are order- a weak approximations each other.
Then can conclude that øa and zµ are order-(1 - /2 weak approximations of each other, which states our main theorem Theorem 3.5.
First, we give control of the averaged learning rate ñe (Equation (14)) with the following lemma.
Lemma C.2. Let (M), 3k) be learning rate schedule sculed by 1) (Definition 3.2) with index and Th be the averaged learning rate (Equation (14)), there
AmaxTmax 11
The S
Amin
Amaxemaxx
1-a
8k
min
Proof. By Definition 3.2 we know Bx E
Amaxe ) and mys S "mmax// Therefore
no noBk+1:e Bk)
1ymax
k
AmaxXmmax
Amin
T
Amax Amax?mmu Bmax
1
n
8x
Amin
3k
Amaxx
<
1)
Amin
18

---

Published as a conference paper at ICLR 2024
C.1.1 STEP 1
For the first step we are going to proof the following Theorem C.3. For notation simplicity we omit the superscript 17 for the scaling when there is no ambiguity.
Theorem C.3 (Weak Approximation of Coupled Trajectory). With the learning rate schedule (Ti,, Bx) scaled by y with index a E [0, Assume that the NGOS Go satisfy Assumption 3.4 and the initialization my satisfy Assumption 3.3. For any noise scale S 11111/2 let (an; mg) be the SGDM trajectory and yı be the corresponding coupled trajectory (Equation (15)). Then, the coupled trajectory yk is an order-y weak approximation (Definition 3.1) to 2y with y (1 a2/2.
For notations, let Vx V[(xx) 9k V + oUR is the stochastic gradient sampled from the NGOS. Also by Assumption 3.4, as 1 is bounded, let Cir 0 be the constant that tr (xx) < Cir for all x E R.. As VC is Lipschitz, let L > 0 be the constant that [|V[(x) - V S
L || - yl for all æ, E Rd. As 1/2 is Lipschitz and bounded, let 1 0 be the constant that (x)) Ezy) S Ly | yll for all æ, y E RR
To show the result, let be a test function for the weak approximation, then for some be- tween 2& and yk. [Eb( Eh(ya)
yx>1l B[]] [|) y lll E E [| [I] Thus suffices bound E || || and E m for any m 2
V(l(Vh(=),ae
and k = tn) As the gradient noises are scaled with variance O O there will be > I, and we will show that E | = O1 is the correct scale so we still have m 2m = (1) and [Bh(2z) Eh(y = (1-(2/) An useful inequality we will use often our proof the Grönwall's inequality.
E |
Lemma C.4 (Grönwall's Inequality (Gronwall, 1919)). For non-negee non-negative sequences {fs: gi, ho R220 for all 0
Se. S ge
ksJs 8=0
Σ
then ft S 9: + Et= kag, exp (Et=s+) k,).
Next follows the proofs.
Lemma C.5. We can bound the expected squared norm
E E||m=|?? < 12
Xmax
X2
sup
[EIIV (4 + 8
Lmmax
+ 17" 11 K
(21)
Amin
min
r=l...k
with constant Kg enough learning rate <
22_ Amnta
Otr 3E
|
2X2
maa Cer 302 (Ca by Assumption 3.3), for a small
Amin
Amiin T..mooxx
I-a
Proof. To bound E || we unroll the momentum by Lemma C. and write
k-1
my = Book-1 mg
+
3 15-1:-110 (1-3.)g
(22)
5=0
Define
h-
my 8o.k
mg +
Bo41sh- BB,)
am0 0
Then Em, Emr and
k-1 m
m
Σ
B B2)ov.
=
19

---

Published as a conference paper at ICLR 2024
Let constant Ko- 20 + 3C2, write
E |m
- [[ m
||
mz 11)²
(23)
< 2E m
2 + 2E
m| m
(24)
k-1
= 2E m
11 + 2
Σ
19s+1:--1220 (111-B3)) (1,)) E(t,)
(25)
==0
k-1
$ 2E
|mx 1| +22 >(
(1
Aminn (0))((=--)))))) max O E(x,)
(26)
==0
2 + 17 a- 21
$ 2E mm
Car
(2
Amin
$ 2E
m
2 + 17 pa-1Ku Ku-
(28)
On the other hand, we know
k m = Bo 17g
1 - -3
We use the Lipschitzness of V to write
k-1
V. Vall S Lba - x& S 1.
Σ
Tr |
1=8
then we can write
A-1
2
m&
S
E
fo---1
mool
Σ (1 (1) Σ 7 | II)
Amax Amia
E
mool
+
| + Lrppmax Σ
17 ) m m
k-1
< 3E( m (| +
Amx E( xx|)) + 17
Σ 7=
(1 Amiu
min
Σ
Amiu 17"h 1 m
31.² 3I 1/2
< K2 + 3 max
E E((
sup
E m_m
X min
1=11..k
Then by Equation (28), we know
sup ll.....k
1| < 2 sup
E B|m- + 1 Ko:
==.....k
So
22
1.2n2am 11 1-a
sup 1l.....k
E B 2 < Ko + 3
X2
max
sup
B(I[V+ +
X2 Xmin
Ko
(29)
min r=1, 1....k
6L 2 M 17 2-2a 2a
+
X2
sup
(30)
Choose 1) small enough so that
2
2 17
2, we know
X
2
=(I + (2 +2+2 2 V/LTmmxx Ko-
sup
E
m,
2
<
max
sup
(31)
r=l...k
min
r= 1.....
Nunin
20

---

Published as a conference paper at ICLR 2024
So
E m 2 $ 2E mm&l" 2/==== + 17" Ku
(32)
12 1
max
Lumax
R((IV+I1) + (4 + min r= r.
sup
+ 17
Ko.
(33)
Amin
Lemma C.6. There is a constant I that depends on T, irrelevant to q, such that
sup O=.....IITmm
E V & f(T)
when 1) min(
A main
12LA Tanax
Proof. By the Lipschitzness of V L, we know
|| - Voll L aay Fo ||
(34)
- | + L yyk yoll + I. 200 yoll + I. 1| - yk
(35)
Observe that ||= - ykell we can write
9. 1-
m
Amaz Vinax 17 //-0|| |mxlll Let di = El --ol2 As result, min
2 + 3 | + AmmaxTmax 17 mmol)
LAmaxthmax 11-aa)
2
E V < 3E
EE [[| m (13) + 3Ldd,
x min
X2
(36)
1a
Choose 7) such that 17 a-
and let the constant
12T Amna
Ki 3E
[
Amin ||mall 12Amax
2
(37)
+(6
Aunin
2 (3 6
IMbax Ku +
[[maax
Ko
(38)
(6Amax
Amin
4Amin
where recall Kg is the constant from Lemma C.5. Plug Equation (21) into Equation (36) gives
1
V S K1 + sup B(I[V+1)) + 3T.24 7
(39)
therefore
E 2
S 2K: + 67
sup dy. T=l....kk
(40)
Now need to bound de by iteration.
dk-1 = dx + y
y + 2E(yR+1 - yk)) (y)- - yo)
(41)
de E | T 2 Btr(E 2B((Ve) (yk yo)
(42)
X2
Xa
Venax? Cr + Amax
S ds 2
E v 11
Amin Mmxx V del V
(43)
X
min
Let function K2 K2 Fo) X m Cur + 2 yuxx ma Ki (mo. xio) + 2 Amax de sup,=I.... d,. We know
Ki (i
Auii
Mmnx
voL
and
K1
dE S VoL2d +
(44)
V612*
X2 X
Amax
d S dy + Knn
67
m dr +
Mmaxx ( /oLdg)
(45)
Amin
< (1 69 Amaxxuax Mmax
12 2 M 17 ) + nK2
max
17 61
(46)
Amin
X
21

---

Published as a conference paper at ICLR 2024
Let + 61. Amar Tenax
11 G1.2 Ximug Vimna rf, then by the Grönwall's inequality (Lemma D.10),
Amtn
A
d.. 11/2(1 +
x expx*)
(47)
1467. Amaxfimax
64,1,
Amax Mmaxx
K2 K2T exp
Ami
Amin
(48)
Plugging into Equation (40) finished the proof.
Lemma C.7. There is a constant f that depends on irrelevant y, such that
mp S 17 S(T)
for all k < T 1/
Proof. The result directly follows from Lemma C.5 and Lemma C.6.
Lemma C.8. There exists function g(m. T) that, for all k S T "¹
[ + < [3] S g T),
E || 13]m] S gm,, T),
E[n" | 2m s g(m, T),
E[(1 + |8 33) mms [2] s g(m, T)
and that g irrelevant to k and r).
Proof. We use the fact that (a + 6)Th (a" + from the Jensen incquality for a, >> 0 and
m 2 1. Furthermore by Young's inequality a0 < aa*t" + +06
< gato + bo+0 for a, B > 0.
a+3
| 3m
- (1lV + 0 / 2m)
< 2-2 [||V=x || 2m + 17 1 [|0 mm)
| 2m $ Vol + L - xo 1mm am
< 22 voll + 2|(=0o))2 + L 2m || 2m)
As E |u|| 2m bounded by Assumption 3.4, for some constant R. there is
E 1 2m < R(1 + 17 " E((EE (
(49)
Now we need show the bounds on E[( + | 12) and B[(1 + 1 mm)
Specif-
ically we shall prove B[(1 + || + 22220/ [m mm S g(m, T) and E( 1 |l
2-222 | (m)") | I g(m, T) for some function g.
Let
8k = Th |1gx 112 27j (gk: yk) + 2 (1 + (1 Br )ga || 11²),
then there is
[[(1 + // 21 + n² 20 m 1yx., mx]]
E[( + 1yk Th.gkl
11?
+ 17"
|| + (1 - Bx)gxl' 2 lBr, (3
[E(1 + ||yx 2 + 1/2 22x2
|m 2 + 5x) 1 |y m,
< (1 ||yx 2| + 1 20
m 1 mm m
||y 2| + 172 20r m
11 11 lyk. mp]
m
+ B
(1
|| 2| + 2 20 |m 2 22 2-m-221pp " mp))
i
i+2
<1
1 12 + 2
20
m
22) + mm + |1yx 1|2 + 2
20 13)mm 12 lyk mg
+
2 m=-2 (1
||1
12
2 20 m
mm) 3
22

---

Published as a conference paper at ICLR 2024
Then there exists constant K3 independent of such that
Elój yk: mg]
Th
+ o le ma
tr [ 2fj Rv yks) + 2 (11 (1 Bx)gx ||m ml)
1 Amaax Vinax (2n)² Voll + L =|| + 2n7² L | + n/Crr)
E ye. ay my
min
+ 2 mmaxxmax
71 Voll + T.| T1| 1yxl
Amin
3x
+ 11 1
2-2a 20
| 2 + (1 Bx 2 E(+)))]
1 1 3x
< (1 + [
Lyk, m + | mn
as (1 - Bk)r7 0) =
0117-0)
O(n) and (1- 0202 = -01) = O(n).
Note that
|0 =113 + T1 B&k (4))-mmm 1mx
21mmax Xmax 2--200 |m ||
2 1131x11
A
Then can write E[dx 1y,, m&] < Kin(1 + ||yx + +122-20 2a || for some constant K4 inde-
pendent of y. Futhermore, as E for some constant K4, Ks, K6, K7.
El (0k: yk) = Bi (Uk; m2) = Elya, (Uk, ozk = 0, expansion gives
(m+ B( < K4 a 1 | 1| + |1 112 + 1/ yk) + 17 3/2-a (m ms) + 1) 3/2-aa (ack, ok))2
I Ksn(1 + || 1| + 17 2 ||mm)
| ["yly] KoB(n) || |a 1 || 2 | 1| + 1) 1/2 yn) 1p 2/222-00000 (mi Uk) + 17 3/2 (xx vn))
+1
< K7nmm/(1 + ||yx 2 + 1, 2-2 |m (3).
Therefore taking expecation with respect to all, we know 11 2a | 1
217 2a || 31)mm
Therefore by the Grönwall's inequality (Lemma D.10),
E + | + 2 2a 1m
S (Ka + ||/0||2
+ 2
20x
|moll 2|
< mk,T+22 (Ks+ 3r (1 + EE yo [2m - 1)²
20 1 Ellmmoll 2m)).
And clearly
|| 2m
2 B|zo| 2m + (2Aman )m" "" ma mg 2m
Therefore we
X
finished bounding the moments of yk- Similar induction arguments on + llya I|2
172 2- mm mm) ||m offer the last equation in the lemma. Then for Ty. we can write
||mm + 2A Tmaxx
2m
<
22m
mm 2m
X2
And for gk with Equation (49) we are able to finish the proof.
Proof for Theorem C.3. We expand the weak approximation error for single k. There 0 E [0, and z = 6zg + (1 f) such that
Eh(xs) - Eb(w)) [[))))()) æx - yxil]
< [ - yy ||].
As Vh(z) have polynomial growth, there is k1, 22 such that
[V(>))]>> < An + ( (1
|| k2 + |1 k2))
23

---

Published as a conference paper at ICLR 2024
Combined with Lemma C.8, v |||| the definition of the coupled trajectory, we can write
[[(]]] is bounded by constant [kg T). From
X Taax 2-220 mx
Ella yx:l?
X
Then as Lemma C.7, there is constant f that E || || S 1))=-1((T) therefore
Rh((x)) Eh(y))) S 17 (1-a1/2 Amax7/max (k2) T). a)/2
X
min
Thus by definition, and yv are order-(1 a weak approximations of each other.
C.1.2
STEP 2
In this section we are comparing the trajectory with a SGD trajectory 24. To avoid notation ambiguity, denote (ax) G- (2:) to be the stochastic gradient sampled at æ. Recall that (with
Snge my and z0 = 26)
yo = 2o
T-8o 80
yk = yk-1 (9
2 2k-1 17 zk 11k
The only difference in the iterate is that the stochastic gradients taken close but different locations of the trajectory. Therefore to study the trajectory difference, we adopt the method of moments proposed in Li et al. (2019).
We start by defining the one-step updates for the coupled trajectory and for SGD.
Definition C.9 (One-Step Update of Coupled Trajectory). The one-step update for the coupled trajectory A can be written
A y m, C) -Tig (y Cnp mm)
Definition C.10 (One-Step Update of SGD). The one-step update for SGD A can be written as
Î(y) -
Lemma C.11. Let be the one-step update for the coupled trajectory and À be the one-step update for SGD. Let Ty; my be the k-th step of an SGDM run and yk = Fk
RuBs my be the coupled I-3x trajectory in our consideration. Then for any C' [0, Amag]lmax] there is function (I)) independent of n, such that for all vector entries (i,j).
Xuin
1 (Bx) (1) m&, C)bbk] Sn (())E( [84]
E (1K) - A (y). me C)
< n² (n + E m 1 (1 1 |1 Il) 1 - m lBu))
and
E [ (1) (3ER) A (1).
(1k. mm C3I
$n 3ي J(T)
B BJA (Ba) - A (Hk, mm [)]
s سيد ((T)
for all k E [/n].
Proof. Recall that we use denote the Lipschitz constant of V and Tay; to denote the Lipschitz constant of the E in the Frobenius norm. We can write
E (ya) A ma. C) n E( ((vn)) 0 Cur + Cnn-- [yn]
2- [11br]]
24

---

Published as a conference paper at ICLR 2024
where the second step uses the Lipschitzness of the loss gradient. The proof can be completed by noting 3| mx
I
E
m&
0111a-1)72 1)/2) (Lemma C.7).
E| (1) A (1) ykk A (2) (1) (kk m. C)
O, / (v:) yn - E a (ys + Ci)--mx) + 1)²0²²0 2002220 Ei (y) + Cit² -
< |1 + 7 V2() + L |
E ms |/ya|] + LC'r m 11x
Again by Lemma C.7, B ml ()/3))))) we obtain the desired result.
Lemma C.12. For any m 2 1, C [0. Armay Timaxx that
there is function g independent of vj and k, such
A
B1ll m., C) S g(m, T) By) S 11 g(m T)
for all k S T/n.
Proof. The second line directly follows from Lemma C.8. For the first line,
Bll mR C) 2T E V + Ciyy (m²) + took
S 222 ( VC) + L | + -Cr Im 11) pm 22m, ok 2m
=
by Lemma C.8.
C.1.3 MAIN RESULT
Below we state an intermediate theorem for proving the weak approximation.
Theorem C.13. Let T and N - Let (æg, be the states of an SGDM run with the coupled trajectory defined in Equation (15), and let 2 be an SGD run with start &n- yo- Then yk weak approximation of 2x.
Notice that the only difference between Theorem C. and our final result Theorem 3.5 that 2y. is an SGD process starting from 2n while 2, is SGD process starting from yo-
Before proving Theorem C.13, we introduce the following lemma, which is an analog to Lemma 27 of Li et al. (2019), Lemma C.2 of Li et al. (2021a), and Lemma B.6 of Malladi et al. (2022). It shows that if the update rules for the two trajectories are close all of their moments, then the test function value will also not change much after a single update from the same initial condition. Let G be the set of functions that are k-times continuously differentiable and all its derivatives up to and including order & has polynomial growth, and let G G be the of functions with polynomial growth.
Lemma C.14. Suppose E G3. Then, there exists a constant (T) independent of n, such that for all k T/n,
nh ))1 Elu(u- + ((ya))] S (n
E El AlVe, mk: Bx
Proof. Since u E G3, we can find Ko(x) = kn + (2) E G such that u(x) is bounded by Ko(w) and are all the partial derivatives of u up to order For notation simplicity let
)
25

---

Published as a conference paper at ICLR 2024
naBe A ( Mkk 1- and A ((Bk). By Taylor's Theorem with Lagrange Remainder, have
Ou(y):
u uyy A) ulyx + })
Σ
dr,
Isida
B1
drulyk)
Σ
(Dis Ais A., Aiz)
dziy Ozig
B2
R3 Rs
The remainders R3, R3 are
Ouly + ab) Di, Aig: Dxi, Dr Osis
Ra =
Σ
6 12
uuly + A., Ais
R3
Σ
Dui dx Oxis
for some a, [0, 1]. We know from Lemma C.11 that
Ou(yk)
EB E
EIA; - vvdl
dz,
Isidd
Qu(yk)
$
17 2-a (TT) 33
Ozi
Isida
< B(1)² (1112-0d dkk J ((T) + || 11 ))
O(1 3-a
The last step is due to Lemma C.8. Also,
Puly))) (E[A,, Ais - -. T leae))
ER2
Dzi, Oxis
s k( - (1 || (2) + ||ma|l (1 + |1 ll) + 17²-a ||m 11 11)
=
For R3, by Cauchy- Schhwar inequality we have
1/2
E Rg
1
Σ E
u + ad)
[E] [|])) 1/2
6
Ox, Oxi, dz
< ( yy a))) 1/2
3/2 1 (m, T)
by Lemma C.12. For Kz + aA), we can bound its expectation by
B[K yu aa)] < 1:f((1 11 aA k222
< <2 (1 + 22 2kg + E || 2
= O(1.
Therefore E[Rg] = 01)3/2). An analogous argument bounds E[ O(1//2) well. Thus, the entire Taylor expansion and remainders are bounded as desired.
26

---

Published as a conference paper at ICLR 2024
Proof for Theorem C.13. Let Vitk be the trajectory defined by following the coupled trajectory
for j steps and then do k - j steps SGD updates. So, Di.+1 = y/ + ( and Øj+1.+1
Yj + Ayy m,j 2,8, I-p, ). Let h be the test function with at most polynomial growth. Then, we can write
B(h(z)) E(h(y)))]
- (m -B(( (()())+))) 1=0
(50)
Define uyy, s, = Eg [h( (.))) where Ply, s, t) is the distribution induced by starting from y time and following the SGD updates until time t. Then,
B[h(z+) E(h (bxx)) -
E[u(()+15+1,55 1o +1, - Elu(0,s+1,) + 1,,
(51)
1=0
Define uj+1 u + 1,). Then,
n ))Il (52)
[[()) B(h(yx)]] )]l
Elug+1(1) (gy Elus+1(0, + A m,s 1- Bg
1=0
We will show that uj+1 E G so by Lemma C.14,
Th33 ))ll s K (T)) (53)
(2 + ((y))) Eluy+1(y + Ayy my: 1 - Bg
Then.
Blh(zz) B[h(yx)]1 $
K Kuy41 (3-))2)))
1=0
S sus K (T (7)
< Tup Kuy (T)n) (1 a12
We will show that we can choose Ka, so that sup, (T,(T) K K(T), and by definition 2y and yı are order-y weak approximations of each other.
(T) is bounded by some universal constant
Finally, we show that Uj+1 E G3, and there is an universal function K2, E G that 110m (x) S (a(x)) for all j S T/n and m 0, 1, 2. Then from the proof of Lemma C.14 we know that there is an universal constant ((T) that upper bounds all Kay (T). The proof follows the same steps as that in Li et al. (2019) for Proposition 25, except that uj is defined for discrete SGD updates instead of continuous SDE updates.
Notice that 1 (ar) = (()=) is the expected h value after running SGD from æ for k j steps. Let {s:] be such process with 8j+1 x.
First we show that s; have bounded moments, and the bound is universal for all i < T/n with polynomial growth with respect the initial point æ. For any m, we know
E(1 + || ) [s E 1 + || 7, ( ((s:) + oo y()"
(1 || 2 + क7 [V,(s, + | 27, (VL(s)) + o0:))
From the Lipschitzness we know [V[(s:) < VC(00] + I ||s and E(s;)]l S |E(O)] Lx: ||=:]] Let 5; = vi [[V](s, + ov,ll" 12 s (VL(s;) + ov:). so there is constant C such that
Ed, s; < 11((1 ||s: 122 E55 |;; < nC(1 || and E85TH < n/ || 2mm Then, similar
to the SGDM case,
E(1 + | 13m = E(1 + || + 8,) )T ss
s (1 + | 3 + m1 + || 2
+ 2 Ed (1 + s; 1m 2 + 15, m
< (1 + ||s; 2m 7 + 2m))
27

---

Published as a conference paper at ICLR 2024
So E(1 + ||,,+1|| 11 < E(1 + [|s)+1ll))r(11 n + + < (1 - (m for all i.j I T/nj. Thus we proved that the moments of s; are bounded by universal polynomial of the initial point æ. Hence, as h is bounded by polynomial, Uj+1 1x)) Bxyls; h(ss)) is also uniformly bounded by a polynomial of x independent of 17.
Now consider the derivatives of uj+1- We use the notion of derivatives of random variable in
the _ sense as in Li et al. (2019), e.g. 0,0; defined in the sense that E 1|0,,,11 = On Ir [(;)). Taking derivatives,
B(1 + 10,8 |r) 8, = [E + 0,5; - 000
=E(1 |0, (3) s;
where 6 = nī [(:)00,, + 00,u;11 1 27
0 + v.) Lipschitzness dictates that [(](ss) < L and [|)E( < Ly, again there is constant D such that E,, s; I D + |0,,; ).. E 8; S 7 + 110,+, and Eg I nn + 11 Similarly,
EB1 + ||, 8:4 10,) = + 00, +.) (
S + |0, 11 yT + mm: + 110, 113mm-
+ +
<2 + 10 12 12 + |5, m-
51 + |0, s: (m) (1+ + n + )
So E(1 + ||0, 11 < E(1 + |
+ nD(m + 2
(1
(1 d)m TD(m+22) for all i,j T/n. Thus we proved that the moments of 0,5; are also bounded by a universal polynomial of the initial point æ. Hence, as h has polynomial-growing derivatives, [[0, Eeale,--=aV 0g S (E (/ | (00as 1/2)) 11))
is also uniformly bounded by a polynomial of x independent of
Notice that the higher order derivatives of s; have similar forms of iterates, e.g.
8's = 8es; + o 7. (s:))s', + o.0).
where of only relates with V C, V'y and 0hs, for 3 < b. Notice that by assuming C and Σ have polynomial-growing derivates up to and including the third order (Assumption 3.4), by induction we can prove that the process of is universally bounded as polynomial of derivates of order < b, so similarly by the Grönwall inequality we can obtain an universal bound for 8 0 for 3 = 0, 1, 2, 3. As the argument here is identical to that in the proof for Proposition 25, Li et al. (2019), we omit the fine details.
Proof Theorem 3.5. From Theorem C.3, we know at and yı are order-(1 - a)/2-w//w2 2-ak approx-
imations of each other, and from Theorem C.13 we know 2n and y) are order-(1 a) /2/2-weak
approximations of each other. Furthermore, in the above proof we can see that [B(h(z*) - h( =
[E(uo(y)) uo (zx))] = [x Mo Lemma C.14, as my has bounded moments (Assumption 3.3), we can see that ((x)) h((E))
8o mo) un (zo)]], so by the similar Taylor expansion as in
00) This shows that z₁ and Σn are order-(1 - ax) --a-k approximations of each other, and naturally they are order-(1 a)/2 //2-weak approximations.
Then we conclude that æn and ny are order-(1 - a)) 2-ea approximations of each other.
C.2 a 2
Following the idea of Stochastic Modified Equations (Li et al., 2019), the limiting distribution can be described by the law of the solution X, to an SDE under brownian motion W.
dX, = V dt + 1/2 (X,)
(54)
for some rescaled learning rate schedule X₂ that f - Akn, in the limit.
When = 1, the limit distribution of SGDM becomes
dX,- = dr A VL - X (X)))
(55)
dM, M,t + 7 VL //
(56)
28

---

Published as a conference paper at ICLR 2024
with similarly " being the limit (1 B,)/n + Vey Here M, is the rescaled momentum process that induces noglnegli impact on the original trajectory X,-
Furthermore, when 1, if we still stick to following ((1/n) steps for any 7, then the dynamics of trajectory will become trivial. In the limit 7) + 0, as my My = O(k7") ( = (TT + 0, the trajectory has limit 2t æo Σi=o 7e2 for all A: = (1/n)) This is different from the case S where there is always a non-trivial dynamics in 27 for the same time scale & 0(1/n) regardless of the index. We can think of the phenomenon by considering the trajectory of SGDM on a quadratic loss landscape, and in this case the SGDM behaves like a Salenbeck-Ornstere process. When < I. the direction of has mixing time of ((1nn) while the direction of m has a shorter mixing time of (1)77"), while when a > 1, both mixing time z and m mixes at a time scale of O11/n"), in (1/nn) steps the trajectory is far from any stationary states.
Therefore this regime we should only consider the case where my = 0 to avoid the trajectory moving too far. By rescaling m and considering O(n= +2) steps, we would spot non-trivial behaviours of the SGDM trajectory. this case the SGDM have different tolerance on the noise scale σ.
D THE DYNAMICS OF SGDM IN O(1/1²) TIME
In this section we will present results that characterizes the behaviour of SGD and SGDM in O(1/1²²) time. We call this setting the slow SDE regime in accordance with previous works (Gu et al., 2023).
D.1 SLOW SDE INTRODUCTION
There are line of works that discusses the slow SDE regime that emerges when SGD trapped in the manifold of local minimizers. The phenomenon was introduced in Blanc et al. (2020) and studied more generally in Li et al. (2021b). these works, the behavior of the trajectory near the manifold, found to be sharpness minimization process for SGD, is thought to be responsible for the generalization behavior of the trajectory.
The observations in these works is that SGD should mimic a Shlenbeck-Ornster process along the normal direction of the manifold of minimizers. Each stochastic step in the normal direction contributes very small movement in the tangent space. Over long range of time, these small contributions accumulate into a drift.
To overcome the theoretical difficulties in analyzing these small contributions, Li et al. (2021b) analyzed a projection @ applied to the iterate that maps a point near the manifold to a point on the manifold. ((X) is chosen to be the limit of gradient flow when starting from X. Then it is observed that when the learning rate is small enough, ((X) Rs X, and the dynamics of (XX) provides an SDE on the manifold that marks the behaviour of SGD in this regime.
D.2 SLOW SDE PRELIMINARIES
D.2.1 THE PROJECTION MAP
We consider the case where the local minimizers of the loss L form a manifold I' that satisfy certain regularity conditions as Assumption 4.1. In this section, fix neighborhood Or I that is an attraction set under VL, and define ((x, t) Sh VL(((x, s))ds and & (xx) == lim o(x; t). & (ax) is well-defined for all Or indicated by Assumption 4.2. We call & the gradient projection map.
The most important property of the projection map is that its gradient is always orthogonal to the direction of gradient. We ultilize the following lemma from previous works.
Lemma D.1 (Li et al. (2021b) Lemma C.2). For all x E Or there is (L) =
For technical simplicity, we choose compact set C Or and only consider the dynamics of æx within the set K. Formally, for any dynamics D& with æn E K, define the exiting stopping time ming ( > 0: #1+1 K), and the truncated process 3t min(k for any dynamics X, in
29

---

Published as a conference paper at ICLR 2024
continuous time with Xu E K, define the exiting stopping time inf 0: X2 & K), and the truncated process X, = min(t,r)-
There are a few regularity conditions Katzenberger proved in the paper Katzenberger (1991):
Lemma D.2. There are the following facts.
1. the loss smooth, then & is third-order continuously differentiable on Or-
2. For the distance function d(x, 1) = minyern 11 - x1,, there exists a positive constant Cg that
(x
x S Ckd(X,T)
for any X K.
3. There exists a Lyaponuv function h(X) on K that
h: - [0,, 00) is third-order continuously differentiable and h(X) = iff E..
For all X E K, h(X) & c(Vh(X), VL(X)) for some constant с >> 0.
d S ( for some constant d' > O.
D.2.2 THE KATZENBERGER PROCESS
We recap the notion of Katzenberger processes in Li et al. (2021b) and the characterization of the corresponding limiting diffusion based on Katzenberger's ' theorems (Katzenberger, 1991). Definition D.3 (Uniform metric). The uniform metric between two functions f,g: [0, oe) IR is defined to be (f.gg) = -r= 2-T
(I] - (())
For each n E N. let A, R R4 be non-decreasing functions with A, 0, and { (t) hr2 be a valued stochastic process. In our context of SGD, given loss function L RD - R, noise function a D + RDx and initialization Tinit E U, we call the following stochastic process (57)
Kalzenberger process
X,,(t) Zinit + L. (( - / VC(X,
(57)
if as + co the following conditions are satisfied:
An increases infinitely fast, i.c., Ve > 0,infr20 (t+ - An(t)) + oo;
2. Zn converges distribution Z in uniform metric.
Theorem D.4 (Adapted from Theorem B.7 in Li et al. (2021b)). Given a Katzenberger process (-hen)) if SDE (58) has global solution Y in U with ( = ((Tinit), ), then for any > O, (t) converges in distribution Y (1) as n
co.
Y(1) = (Xinit) + L' o (0
M Σ
(58)
(JE|D] K.EEE
We note that the global solution always exists if the manifold I' is compact. For the case where I' not compact, introduce compact neighbourhood of and stopping time later for our result. Our formulation is under the original framework of Katzenberger (1991) and the proof in et al. (2021b) can be easily adapted to Theorem D.4.
D.2.3 THE CÅDLÁG PROCESS
A càdlåg process is right continuous process that has left limit everwhere. For real-value càdlåg semimartingale processes X, and Y,, define X,_ lim Xs, and 1 X,dY, be the process
b+ X,-dY, for interval < That is in the integral we do not count the jump of process Y, Jat but count the jump b. Then it's easy that
30

---

Published as a conference paper at ICLR 2024
Lo X,dY, + J X,dY, Sa X,dY, for <b<c.
Z X,dY, is cådlåg semimartingale if both X, and Y, are cådlåg semimartingales.
By extension to higher dimensions, let
d[X]: d(x,x,)) xx(dx,) x
and
d f1, d(xx)) X((d)) - (dX))x,
we know [ and [X, y are actually matrix-valued processes with A[X], (AX)) and A|X, = (4X)) (AY)).
The generalized Ito's formula applies to càdlåg semimartingale process (let 8 [(X)[M] Eng f, for any matrix M) given as
df((x)) (0f(x.), 1 (00000/AAAAAA000000000A000000000
And integration by part
+ (dX,)Y, xddY)
d[x,
These formulas will be useful in our proof of the main theorem.
D.2.4 WEAK LIMIT FOR CÅDLÂG PROCESSES
As we are showing the weak limit for càdlåg process as the solution of an SDE, the following theorem is useful. We use C(10, 2), (Rd) to denote the set of càdlàg functions X [0, + Rd.
Theorem D.5 (Theorem 2.2 Kurtz and Protter (1991)). For each n, let x(") be a processes
with path in C((O, T], dd and let Y (m) be semi-martingal with sample path in C(10, T), R"" respectively. Define function (r)r) = (1 - 8/r)+ and Y((") = =Y(") -Eese hs (Y "AY
be the process with reduced jumps. Then Y((") also a semi-martingale. If the expected quadratic variation of Y((m) is bounded uniformly in n, and (Xn, Yn) (X,Y) in distribution under the uniform metric (Definition D.3) of C((0, TT Rdxm x )T),, then
f xx,dY.)= (X,Y, / XdY)
(Xn, Yn,
distribution under the uniform metric of C((0, Ti, Rdxm x 1mmm Rd).
Therefore if (x) the solution of an SDE x Xo + z + +10 p(n) (x()(x!") ">) (x("))ar? and if the tuple (F(*) (m) =!()) converges for all X. the process (F(X,), Y., 0), then by the above theorem know the processes (x("), ("), (²)) relative compact, and the solution to the SDE
X, = Xo + fo F(Xx))YY is the limit of (n) as stated rigorously in Theorem 5.4, Kurtz and Protter (1991). This will be the main tool in finding the limiting dynamics of SGDM.
D.3 THE MAIN RESULTS
We provide more formal version of Theorem 4.5.
Theorem D.6. Fix a compact set k C Or, an initialization Xn E K and a: E (0, 1). Consider the
SGDM trajectory (n)) with hyperparameter schedule (1k (n) ((7)) scaled by 1f (n), noise scaling
(n) and initialization (æo, mo) satisfy Assumption 3.3; SGD trajectory (Zk
(m) with learning
rate schedule (n)")) noise scaling 1 and initialization zo = zg. Furthermore the hyperparameter schedules satisfy Assumptions 4.3 and 4,4 Define the process x(n) 2 [111((()0))]] 6(20, (n)
31

---

Published as a conference paper at ICLR 2024
(20)) and _,=) [1/(611))3 ((20, (b20,1/1(mm) ((0),, and stopping time Tn. inf(l > 0 X (n) K) th inf(l 0: ((m) & K). Then the processes (X : Tu) and (Z(m), tno. vn) are
relative compact C((O,T TT) x [o, T) with the uniform metric and they have the same unique limit point (Xirr; such that XI L almost surely for every t 0, T inf(t > 0: X K}. and
X, F=:((a)) + /
+ 1/"
X 2 O (x,) E(X,)
Note that for a sequence zy we are considering the sequence X; = x [t/(mt)))3]] - 0( t/))(m)) +
(20)) after rescaling time k [1/(r((0)3) This adaptation is due technical difficulty that the limit of z (m) will lie on the manifold for any > 0, but (n(")) = zo & T, so the limiting process
will continuous anywhere but zero. For the process X I 214/6/(033] - o( t/n/())) (2))
however, we know X (n)
211/61(0032] for any t 0 and x(") ((xo) E thereby we make the limit a continuous process while preseving the same limit for all t > 0.
D.3.1 THE LIMIT OF SGD
(m))
Let us reparametrize the SGD process zą. Define the learning rate x iterates given by
"Ivoge)) and the SGD gin)
gk: V((xx) + (112((x))
(59)
(n) gk.
Zi+11 2k
(60)
Here we reparameterize the noise Us (2 Ex so that &x are independent, E&x 0 and EgNE, = I. We also assume that the constants Cm in Assumption 3.4 are defined here as E |l S
Cm-
D/ 1//6/ (n) and Zn(t) = ===0
(1/6 (n)
Now define the stochastic process A, (t) = then Equation (60) be written stochastic integral
m Ex.
R=n
z mm) = 20 / vC(Z(")
(CZZ("))AA(()) (t)n(t) + E (m) dZ (t).
(61)
with 2(n) [1 Then characterize its limiting dynamics.
Theorem D.7. The process n(n) is Katzenberger process, and for any t > 0, z converges in distribution Z, as + oc that
Z. = 0(z0) + / A. Z.) (Z,)dW, + [ X2 g2 (E.,)dss
2
(n) is Katzenberger process. Note that Proof. First we show that z(n)
An increases infinitely fast: by Assumption 4.3, for all < t,
[e 11(()1n))1111
An(1) - An(s)
Σ
mk
(n)
(n)
k= [=(s/(n(0)aa
[3/(n(n2)2]
((1)) Ak(n(mo)s 2 2 s-s
(n))
inf A1) 1E(0,27
(n(n)
3/(m()))2
(n) 0.
32

---

Published as a conference paper at ICLR 2024
(t) converges ((t) that there is brownian motion W2 and
2(1) - X.
This is shown with the standard central limit theorem. Let Wn.()) S 22(1) be the normalized martingale. By the standard central limit theorem (for instance Theorem 4.3.2 Whitt (2002)), (t) (s)(s) has limit as a gaussian distribution with variance Le/(n (
Σ k= [s/(r) Argg
)² + (t s) by Assumption 4.3. Then Wn (t) converges to
bronwian motion W by Levy's characterization.
Therefore 2,n) is Katzenberger process, and by Theorem D.4 its limit is given by
Z. = =(20) + / 1.00 (2) +
{-"
X o (Z.)
2
D.3.2 SGDM WHEN < 1
For the SGDM setting, we wish to extract the scale from the hyperparameters to make notations clear. Therefore we define
(n)
T X
17(n)
m(m)
1
(ny(1))a*
k = (1)(1)0))2]
Then the original process
(m)
3 k + (1
(a)
(n)
(n)
(62)
k+1
(n) (n)
k+1
ks The
k+1
(63)
can be rewritten into SDE formulation. Similarly, we reparameterize the noise e% 1/22
so that &x are independent, E&, = 0 and E&& = I. Let the processes Z," = 112 [t/n/ J. and
Y = 11 D!!?" &.. By the previous convention, the process can be rewritten as (let 17 = (7))
dx;
Ae (MF dd)M))
= L ((1 ( + V -
(64)
dMC
- (M VV(()))) + - 11/2 (X))(X))
Then X!" = = [t/>p] and M = m (n)
Le/121''
Rewriting the second line gives
Mr -, 2- + VL(X") (
(65)
So
dx,""' = =>mndM," + (a/ A VZ( 1/ (X)")(")
(66)
Now consider the Ito's formula applied to the gradient projected process (X."). Fix K C Or be compact neighbourhood of the manifold f. Since & ((X) only defined for X E Or, we take an arbitrary regular extension of &o to the whole space. Fix time horizon T > 0, and let To min(inf > 0 X2 & K
K),T) T)) to be the exiting time for the compact set K Or we
have chosen earlier. Since X K, know Tn n/2 We use x 1l 7n] denote the
33

---

Published as a conference paper at ICLR 2024
indicator process of the stopping time Tu- For any càdlåg semi-martingall X,, XeATa is cådlag semi martinante that dX = X.
For simplicity omit the superscript unless necessary. Ito's formula on (Xx, gives
d (xx)x, + / (X.)((([ + dó
= + /m)d"" VZ(X))) (X + 88 (Xt) d
where 5 are error terms as 8(0) = and
_ dø = (x)), (b -00 (xx(
For X, E K. there is always VL(X)) = 0 by Lemma D.I, so consider the following process
2 2o- (Xo) and
dØ; = (x))
(am1-~/77)ddll - (x,))()))) + [ dô (67)
Then PEAT. = ((eAT.))
In addition, direct calculations gives
[2], = Ur:
[ = 1/2
Σ
GE. E[y]: = 2,1
2y = r²²,,
d[M]. 7::000 RMi VM( VC(x,)) dZ, + 2:20 20-2 2/122 (X))
(X)) 7/,17 20 'M: VZ(X,)) ))) - -7,11 2,,12--111111 / -V((x)))
d[M, 2), = -707 (M, VZ(X)))dZ + 21
d[X], = M, dR, + d[M] M,d [,1, +dM, 2),M,))
d[M, = -m a- (M, VL(X.)) dZ, n000 12/2 (X,,(XXX nd(M))))
D.3.3 CONTROL OF THE VELOCITY PROCESSES
Notice that our process X E K is bounded for t < Tn- Therefore following regularies, for any continuous function K R. ( is bounded t Tn Also, notice that MM has bounded moments:
Lemma D.8. There exists constants Cm such that E ||
1T S Cm
(n)
(n)
+ (1 - -1((")))) (n)
Proof. This follows trivially from the iterate m
k-1 = B
m
E
m
k+1
(n)
<
n
m
(n)
+ 2
gk
(m)
2
m (n)
+ n
V (m))
+
(n)
The term V (m))
< suP [[L(())] bounded and
(x)
is bounded by Assump-
tion 3.4. Therefore the lemma follows from the Grönwall inequality Lemma C.4.
Lemma D.9. For all stopping time t and function f > 0, So f(M,, S f(M,, ((M,,X))dk X,)ds, and f J(M,, Xz)dZ, 2 1o f(M,, X, ..)ds 1/2 (M, X.).
34

---

Published as a conference paper at ICLR 2024
Proof. Note that the processes M, and Xg only changes at jumps at koy,, the result followed directly from the definition Zg.
Next follows some facts are are useful in proving the theorems.
Lemma D.10 (Katzenberger (1991) lemma 2.1 ). Let f.g R + 10, oc) be functions that is non-decreasing and g(0) 0. Assume for constans > and all
0 5 f(t) +c+ S(" f(s-))g((s)
Then f(() Cesli) for all 1 0.
Proof. In this case g(t) 2 O. Expansion gives
f(t) $ C(1 Σ
(5gg))))))))
I C(1 +
n! 9 (t))
n=1
Cta(t)
We may also encounter the case where g negative. Another form of Grönwall inequality useful here.
Lemma D.11 (Grönwall). Let f.g R + [0, oo) be functions that g(0) = Assume for constant C> 0 and all 0 < <t TTT
1(1) S ((s)- cc. S(¹ [(r-)dr + g((() g(s).
Then [(11 S (C-(((((0) + u(1) + +C e- -C(1))))())))))) - (s)))dsssfc for all S.
We need a form of Grönwall's inequality with uncountinuous process Z₁
Lemma D.12. Let f : IR + [0, o0) be non-decreasing functioon and R x R + (0) be
non-negative, Assume for constant C and all <<l<T
f(t) $ f(s) c.. f(r-)dZ, + s).
Then f(t) S (02(( + gt,,0)) + CC -C'(2 (2:2--2)) 2.) g s)d2 for all s S T with
log(1 + Cr7.) C == =
Lemma D.13. Jo cZ dZy n² c -1 (0 1).
Proof. Directly calculation that En d² 1-c++1 I-c for c = Gnn
Proof of Lemma D.12. multiplying both sides with e (2, and integration yields the result.
Specifically for g(1, 8) = Z - 2% the bound can be further simlified.
Lemma D.14. fo - - -(Zz--Z)) (Z) dZz < C
Proof.
10
C(z -Z")
75 d2 <
c
C= = -
35

---

Published as a conference paper at ICLR 2024
Another useful theorem is the Doob's martingale inequality.
Lemma D.15. Let X, be a martingale for E [0, whose sample path almost surely right- continuous. Then for any C >> 0, and p 1,
T|
Pr[ sup (E[O,T]
[xe| 2 C
Cp
Furthermore, integration with p = gives
E sup cE
[xd - 2v E|
Lemma D.16. For all t > 0 and a [0, 1),
lim v7 Esupte (0. 77
| l| for any 3 > 1 moreover, lim 17 3/2 supre(0.7] || =0
lim B E
MinT. @| dZ = for any > 0.
moreover, sup,, E f. || 1| d2 < oc.
For a E 10,
lima vf E So PM
|M dz = 0.for any B > O.
Proof. Ito's formula on || gives
d || = 2(M dMF)) + tr d[M"
-2770m ( (|
+2, > (
+ - |M, - VC(XT)| dz) + 2 ((
vL(xF)) Σ /2 (X?))d)))
= 17 2 -dW, - 2%07" -- 3[1|M "
(M₂)) VC((T)))
+
2a-2 2a- MM VL(Xx') 2 d2; +
tr (r(;(()))
Let W be the martingale that WT = 0 and
dW. 2 "1/3(XX)) + tr (x(;") (E(;)) IdZ,))
nM. - T (X²')
When a 0, take the constant C1 2supxek |VXXX (2 + [|V[(X]]I) and C2
supxek tr E(X) + C.. Since 7 |M & (1 + 17 20 MM 1/2, for any > 8 > 0 there is
plTn / (-22min 8ta-2 1 Cir 8+a - ||
8 M 1|²
Ciny 8 ++2a 2 <| 1))d2;
SATA
LATH
+ 1/ ||
1|² + 11 0t2a-2
X²
+ 9+2a - (Win.. - W
SAT
8 MaAra
12+ Jann
(-2Amin')
1 Cin 1 + 2)n ||M'|1 ddz 2
+ 1 2a 2 CCg Znn + + 3 W
By the Doob's inequality Lemma D.15, W.s..
WAT,.
-Wa SATA
ptht
2 (M,) T dX""dY," + + (E(x" -dZ,")
2,,n(1M. -V((X))) T 1/2 (x,)),;
36

---

Published as a conference paper at ICLR 2024
is martingale, there is universal constant A that
E sup |W S 2/ 12 rE(s,t]
AAT
s 2
Dinz 1ZhT, +12-22
|M dzp
SATn
s 2A1 -ZZATt -Daita + +12--00
|M,'11 dz,))
SATT
Therefore let K2 = EupPEs | 1|?
KF S K [- dZ, + (2A C2)1 20+8 8- ((, - -
Here Kin = 2f n0=0002-00 + Cin + C1npa-2 2, + 2Am 1 then when 17 + 0 eventually Kn < O.
Then by the Grönwall inequality Lemma D.12, with n², = 17 2log(1 (1 + |1in in²).
K2 S 10 2 B2AA C2)(e 1',27 z + [ 1 1s k2 -27) ( Z )
By Lemma D.14,
K2 S 11 20+ + C2) R. z |Ki [Km
Taking the limit gives lim KT- = =011)00+8-1) 0.
For a = 0, there
1/2d M| 12= = dWW - 2 (1 - (M,'' TV(X))))L
+ 77 I| VC((,') )1² dz, 7/ (r(((x)))
S dd," - - "IM 12 ddz + [V 2 dz; + (r;
For KT- =n{Esupssjo.t
M 1| in -mmin 2A and some constant Cg there is
KT S KT + S: in KnKgd + Cznh 3 7 Similarly by Lemma D.12 and Lemma D.14, lim K. 01p O. For any jump AMP, and (MM2) || 1 there is 0 E [0 that M OMP (1 0)M² and Af(M,) - (0f(MM).. My) } 0² ((MM (MM))MMMMMM)) 03 S(M)AMM AMP, AMT]
(M, AMM 3 (M)
--2 |M]
2 MM
2 IA
Ito's formula on || gives
3
d M < 3 M (MF, dMF)
2
+ "IM
+ 2|
At knj² there a jump for the process M. as AMA ( + VC(XT-) (1/2(xb_))6 so for constant C4 = 12supxxx ((VL() +1+ (1+
E AMM 1 Can 11 +1)2Zan.
37

---

Published as a conference paper at ICLR 2024
When > 0, let Jf B |M, 1 a/ 77 |M, - 11 7 |M + 1///2t002 2there is
J S J (3 -1 3/2
C,),, ( C6T C -,)
where C3 = 3C4 (su (e + supXEK max max ([V[(X))I) some universal constant. By Lemma 1.12 and Lemma 1.14, there is
J C22
And the conclusion follows.
simpler proof may be given by consider the process (M @3 instead of || and the proof idea is similar using the Grönwall's inequality.
Lemma D.17.
There exist a universal constant such that
AM - Cny"
C MH 1).
| - Cn( Mi +r)).
Proof. Direct from the iterations Equation (64). Lemma D.18. When - do (or 7j 0), 6(L A Tn)
0 weakly for all (0, and e [0, T).
We wish to generalize the result a little bit.
Lemma D.19. For any function (K), as n + O,
E sup Σ AT(O.T.ATT
-
- 0f ) -
f(x "Ax"" 0
2
sE
Proof. AX" and Af(X)) are non-zero at times = knj² for E Et. By the mean value theorem,
there is E [0, 1] that
A - -0(0=)00 0 I(x=
100 f(0x +1- (I-0));;AX;; AX", AX;!I
|0 (XX
sup Xek
Here the norm is defined for tensors as ||00 f(x)lF (0,,0,,,,, The first term supxek [[0*s(X)) a constant independent of n (as K compact), Notice that
AXA Ann + AMM) - M..
Therefore there exists constant C (supXEx 1100 [(x)
that
E sup AE( sE(o.d)
Σ
Af(x7) 0f(x:-)Ax: 1 AX"
S0. E sup
Σ
|
TE(O,T.AT]
TAS
<C En
|M dZ,"
From Lemma D.16 we know Jo Thtn E
[| M
dz - 0, so the proof done.
Proof for Lemma D.I8. The result follows by applying Lemma D.19 to every coordinate of ((X)).
38

---

Published as a conference paper at ICLR 2024
Lemma D.20. For a E (0, 1), limm 000 E || 1| d2d 0 for any 3> -a.
moreover, sup,, 1f 70 E Jo | 1/ "2z
Proof. Consider the energy function G(X,, M,) = 2X, ((x,) -1 |Mlll there is
EG(X,, EG Mo) + E 2 VL(X,))X 2d( ()(()) + AA(X)))
X
21 [X + dō
X
27m 11112 (M))) dZ(
+ - -1 11M, - VL(X;)) II dz; + +12 to (xx
G(Xo, Mo) + A. f 2n177 'E ||M dzP
+ +f 76. 1 MM VL(X;) az - ir (r(E((?))d(?))))
Here A, is some uniformly bounded process. Multiply both sides by 17' 1 gives
1 2 "B ||M
dzp S - Mo) + - A.
1 M vI(x))| azz + 22 (r({(())dd('))))
From Lemma D.16 we know the right-hand-side is uniformly bounded in n, and the conclusion follows.
D.3.4 CONVERGENCE TO THE MANIFOLD
We wish to show that the process X fAr, - D L that as the learning rate 17 - 0, there is a distance function d that d( I') - 0 weakly as a stochastic process.
for any t 0. First need to show
Lemma D.21. As + oc, d(L,, r) + 0 weakly all a E (0.
Proof. By Lemma D.2, we need to prove that for all > supieT h( - in probability, Ito's formula on (XeeTT.) gives for some process 0 (n - x0)
dh(X,) (V)(X)), dx.)) + + ( +
(VA(X)), AondMM (N1/ A (V(())) (x222())))x"))
+ (2(X
Let the process
dS, = (Vh(X)), -DendM, + (26/m))'-=MM4 dM xxE12((7))))") + 1 - + +6,
there is
dl((,)) = ds, Aej (V)(X,), V.((.))
< dS, Amin 1
Therefore by Lemma D.12, for K = 1 log(1 + Amine 11), there is
h(X.) S ==Ze S, + (1 S' e-n(2i-Z)) S.)dZ.
39

---

Published as a conference paper at ICLR 2024
Clearly n22 e=>2 S 0. Furthermore we have for
LAT
A
(-xon (N//7))"""
PLATH
B. =
JO
AATA
C =
v + d8
we know S, = B. A.
First, we show 17ー 1e e e-u(Zi- (A, for some constant K, know by Lemma D.14 and Lemma D.16,
(A,-A,) Az)dZ - 0. as A, - A, S J (1,,1
2 sup |M [' E (-112r-2, - Zz)dZ, 1 2x sup M, S 11 sup |lM.| -
"
and since Ssu,, || - Ysll S B//lI I K1, there is Esup, nf 1 So e-((2--Zo)Yy Y.)dZ.l
< Kt kgj O.
Next, we show 1 S0 (x) - B.)dZ,
7 0. B, B is martingale so by
Doob's inequality, Esup, |B. B,l < B,
I Ktt for some constants K. Therefore
Esup, "I e e~ (B: B.)dZsl I Kt
+ 0.
km
Finally, there exists constant K
C.)dZ, < nxZ + 0 by Lemma D.14. Therefore we concludes
such that 17 J e
7 Jo =x(2 2.) K(e - Zo)dL, v
the proof by showing that ((x)) + O.
D.4 AVERAGING
Lemma D.22. lim S'EsS S "FT(AT")VV(((22))))) 0.ror any B
To prove the result we need another lemma.
Lemma D.23. limm-+00 Esupsr 1o PMT [|V[(X))] d2 = ffror any B > -.
Proof. Use Ito on (VL(X,), My), there is
d(Vi((,)) M.) - -,,00 ((X.)) +
v. Vc(,)) 2+ (VL(X), - (X,)d)))
Therefore know there exists constant K such that
LAT.
E
[V((.)] dz, < Kn LE(V C(Xen.) M. MeNT.) (VL(X)), Mo))
1|M dz. + Kny' AATH
+ +1 1-KKKE KE
Σ |M
AM
+E
(.)x
JU
The first four terms vanishes when multiplied 87 for 3 > -1 by Lemma D.20. Note the last term
I0 (V((X)),
(VL(X)), '1-,,',,,0''' -
so
V VL(X)), dX,) d(A; c(x)) (AC(()) -())C(())+((()))))--+))))))))))) (10 -68 X,
x (VL(X,), is clearly bounded process given the bounded variation of A and boundedness
of X MT. Thereby finished the proof.
40

---

Published as a conference paper at ICLR 2024
Proof of Lemma D.22. Let f(X,) (X(XX.))( Ito's formula on @2 [ gives
d(((X.)M.) - + (X
(xx)))
Therefore there is constant K such that
S: ((x.) - S: [(X.)V2(XX)
+ K() (
+ +f
+ t (ff(X.)AM)
We know that 2 AX, ((n) and sup; fo ns (X. = O Expan- sion gives Jor 172 ( = O(n) by Lemma D.20. Finally by Lemma D.23 we obtain the desired result.
Let dr >e/7 d
Lemma D.24. (, 1 - as O.
Proof. Ito's formula on 121-"6,0000 ", gives
(,00,,00MMM0..0000000MMMMM--,,00 (L00()))
+ |4 d6
where
We know for E (0, 1), by Lemma D.16
n'-08 + as 6. S CC S | dZ..
(0 - 0 as d(or) < oo by Assumption 4.4 and sup, 2 |Mell
0.
En'-a (A.)) + 0 as Ee < oc by Assumption 4.4 and
- Mell +
+
A00
02 M. VC(X.))dZz -
(10%%) +
Therefore adding them up we obtain the result.
Lemma D.25. J LATA (X.()x)(( - - X2 |dd, + 0.
Proof. This follows directly by the expansion of [X] and Lemma D.16.
Lemma D.26. sup Jea X 8 (xx( MjdZ -ee x² (Xx( 2
- 0 as
" 0.
41

---

Published as a conference paper at ICLR 2024
Proof. let A, (. for some schedule ot, then for some uniformly bounded
process Bt,
dA, ((
Mc]) + 30 dd, (X,)
O 8
+ 172-20
"
20 @ M,, M,JdZ.
20 0 $ (X, (LX)), Mijd 172-0
1
20n @ (()( (1) E (x,dd,, Mr]
17-0
2017
8² 1 (x) - E
1
+ t
2a 8 1 (M, (M( - V
Multiply both sides by 1/ 2--2, by Lemmas D.22 and D.23 we know / (V Jdz and f (2 V Tid converges to 0. Bound the Martingale W2 that
(X,)dY, M,] 20m dW; = 00 (x)
Σ 1/2 (X,)l,
Doob's 's inequality alongside with Lemma D.16 shows Wr - O. Then we know with oe = X//12 that
1//70 2 (xx)((x)))) M. M,dZz X2 0² (xx, 0
sup
c1,02
r7
2
Finally are ready to show the limiting dynamics
Theorem D.27. For any >> O, (Xisr, Tn) converges in distribution to (XINr, that 7 inf
0: X, & K), and that
X, ((xo) S/" 222
/ X2 (
Proof. Recall the process 47
dΦp- = (x(( Xen Xen' /w)dd1 - -2x -x
+ (X;")
+ d
Therefore we know
XIATE = ((x0) + (XTAT. - -DiATa) +
(68)
AAT
(0(())((->x0) + /X)dM - (XPP))))) + 1 [dx)
(69)
By Lemma D.21 we know the process X"" - 47' weakly converges to zero. By Lemma D.18 we know 0. By Lemmas D.24 to D.26 we know
AAT
00 (x(X")
+ Aen " + 32 (xx(
tAT
X 8
+
We know the process Y Z and [Y"], are of bounded quadratic variation, so are y.n Y E.st ha((AYTDAYT, zi z E.st hs hss((ZZ))) and [Y*), [Y"] E.sO (o([A Furthermore by the Donsker's theorem Y," W, in the uniform metric, where W, is Brownian motion, and Z
By the law of large numbers know
42

---

Published as a conference paper at ICLR 2024
| - t. Additionally, the process X7", ZP and Y always share jumps at the same locations. This implies we can write
X Xo PATT + { Gn(X))d( (x,))
for the functions Fn, Gn and Hn that describe the iterates Equation (68) and
Pr (X: " + +
J 2 (x(())((- ((120)))))) / 1 (xx(())))(11 X
Notice that the stopping time is incorporated into the functions that F.(()) ( (z) for x & K. Notice that the processes (Pr. Y [Y"), Z Fx (X,,). H (X,) converges in the uniform metric to (0, W,, t,,, 1, F(X,), G(x,), H (Xx)) for any process Xt, then by Theorem D.5, the limit of X can be denoted by
[ G(X
X, =
X =o
Plugging in the above results gives the limit
X, = ((xo) + //" xx00(X)))"((x,))))))))) + f X ((,((X)) 2
Proof for Theorem D.6. The result is natural corollary of Theorem D..7 and Theorem D.27.
E EXPERIMENTAL DETAILS
E. LANGUAGE MODEL FINE-TUNING
fine-tune RoBERTa-large (Liu et al., 2019) on several tasks using the code provided by Malladi et al. (2023). We interested in comparing SGD and SGDM, learning rate grid proposed for fine-tuning masked langauge models with SGD in (Malladi et al., 2023). We randomly sample 512 examples per class using seeds, following the many shot setting in Gao et al. (2021). Then, we fine-tune for 4 epochs with batch sizes 2, 4, and 8 and learning rates le le 3, and le 2. We select the setting with the best dev set performance and report its performance on a fixed 1000 test examples subsampled from the full test set, which follows (Malladi al., 2023). We also compare the trajectories of SGDM and SGD when fixing the data seed and that the two closely track each other (see Figure 3) on five different tasks.
coarser version of the
Past work (Malladi al., 2023) has suggested that using suitable prompt can fundamentally modify the dynamics of fine-tuning, so we also report results when using suitable prompt. Prompts are taken from Gao al. (2021). Results Table 2 demonstrate that SGD and SGDM track each other closely in the prompt-based fine-tuning setting well.
Task
SST-2
SST-5
SNLI
TREC
MNLI
Zero-shot
79.0
35.5
50.2
51.4
48.8
SGD
93.1 (0.9)
54.9 (0.7)
87.9 (0.8)
97.0 (0.2)
82.4 (1.4)
SGDM
93.1 (0.4)
55.3 (0.9)
87.3 (0.5)
96.8 (0.7)
82.8 (1.4)
Table 2: SGD and SGDM for fine-tuning RoBERTa-large on 5 tasks using 512 examples from each class (Gao et al., 2021: Malladi et al., 2023). We use the simple task-specific prompt introduced in Gao et al. (2021). Results are averaged over random subsets of the full dataset. These findings confirm that SGD and SGDM approximate each other in noisy settings.
43

---

Published as a conference paper at ICLR 2024
SST-2
trec
sst-5
5G0
500M
5G0
SG0M
1.6
S0DM
1.-
Train Loss
12-
9.50
0.2
0.25
0.8
0.1
0.00
1000
2000
3000
2000
4000
6000
8000
10000
12000
2000
4000
6200
8000
10000
Steps
Stegs
SNLI
MNLI
500M
SGOM
W
Train Loss
0.0-
0.-
3000 4000 S200
5000 6000
Figure 3: SGD and SGDM trajectories when fine-tuning RoBERTa-large on five downstream tasks. We ensure the effective learning rate fixed in both cases, so SGDM trajectories are with learning rate 0.001 and SGD trajectories are with learning rate 0.01. We fix the data seed and fine-tune using five different optimization seeds. The results show that SGD and SGDM track each other closely on average over the course of language model fine-tuning.
E.2 CIFAR 10 EXPERIMENTS
We report the results of training ResNet-32 CIFAR-10 with and without momentum. First, grid search over learning rates between 0.1 and 12.8, multiplying by factors of 2, to find the best learning rate for SGD with momentum. We find this be 1j = 0.2. Then, we rum SGD and SGDM with SVAG to produce the figure. We obey the SGDM formulation Definition 2.3 by setting r-3 with 3 = 0.9.
Standard SGD and SGDM exhibit different test accuracies (89.35% vs 92.68%), suggesting that momentum exhibits a different implicit regularization than SGD. However, as we increase the gradient SVAG (Definition 5.1), we see that the two trajectories get closer. At
noise by increasing the final test accuracies for SGD and SGDM 91.96% and 91.94%, respectively, verifying our theoretical analysis.
44


Reviews:
[{"review_id": "UEqtBZZFOP", "reviewer": "Reviewer_ReS6", "summary": "The authors have proved theoretically the marginal benefit of momentum in training when the learning rate is small and the noise is dominant. The theory has been verified with some experiments.", "questions": "1. Could you show some examples of what kind of networks, loss..etc the theory fits in.\n2. It will be nice to make more clear comparison of when with and without momentum.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper has demonstrated the limited benefit of momentum theoretically in some regimes, which helps suggest when not to use momentum.\n2. The article has illustrated the idea with warmup examples, which makes it easier to digest.", "weaknesses": "1. The analysis is based on the fact that the momentum effect is dominated by other factors when the learning rate is negligible. This is somewhat intuitive and needs to dive deeper into what the special role of momentum is. \n\n2. There needs to be clear examples of what scenarios the theory fits in. For example, with what class of loss and neural network does the theory fit in?\n \n3. The paper seems rushed in polishing. There are some links of reference in the paper that need to be added, for example, in line 2 on page 8.", "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": "14S0jqWAuv", "reviewer": "Reviewer_RZuq", "summary": "This work compares SGDM and SGD in the setting of small learning rate and gradient noise is large enough to produce instability. Two main results are presented, both indicating that the training trajectory of SDGM is close to that of SGD. The first result states that within $O(1/\\eta)$ steps of training, the trajectories of SGD and SGDM approximate each other with distance $O(\\sqrt{\\eta/(1-\\beta)})$. The second result states that within $O(1/\\eta^2)$ steps of training, the trajectories of both SGD and SGDM move slowly after reaching a manifold where the local minimizer is located.", "questions": "1. Do the two regimes, i.e. $t=O(1/\\eta)$ and $t=O(1/\\eta^2)$ have any connections? Can we understand from the results provided in this work what's going on when $t=\\Omega(1/\\eta)$ but $t=o(1/\\eta^2)$?\n2. Could the authors clarify why a **class** of hyperparameter schedule $\\{\\eta_k^{(n)}, \\beta_k^{(n)}\\}$, indexed by $n$, is considered in both main theorems? It seems more straightforward to me to consider only one choice of hyperparameter schedule.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The main results, i.e., Theorem 3.5 and Theorem 4.5, are strong in terms of both implications and proof techniques.\n2. The results are presented in an orderly manner.\n3. The result that the training trajectories of SGDM and SGD are similar is intriguing.", "weaknesses": "1. Although the implications of Theorem 3.5 is clear, the statement of the theorem is a bit confusing. Specifically, the averaged learning schedule $\\bar\\eta_k$ is introduced, and it does not appear again until the appendix.\n2. It would be interesting to see other concrete types of hyperparameter schedule apart from the constant ones, both theoretically and empirically.\n3. It would also be interesting to see experiments demonstrating the main theoretical contribution of this paper, i.e., the trajectories rather than the loss / accuracy stay close.", "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": "34sWWRbcBa", "reviewer": "Reviewer_fSFX", "summary": "This paper aims to investigate the advantages of incorporating momentum in the Stochastic Gradient Descent with Momentum (SGDM) method for mitigating the variance associated with stochastic gradients. To mitigate oscillations along the high-curvature direction, the authors make the assumption that the learning rate is small. Within this context, they demonstrate that, subject to specific assumptions, Stochastic Gradient Descent (SGD) and SGDM yield identical training error results. Additionally, similar to previous studies, they construct a Stochastic Differential Equation (SDE) for the continuous version of SGDM/SGD and, through an analysis of the SDE's limiting behavior, reveal that both SGD and SGDM exhibit a similar implicit bias.", "questions": "In Theorem 3.5, the paper demonstrates that the trajectory of Stochastic Gradient Descent with Momentum (SGDM) approximates that of Stochastic Gradient Descent (SGD) with a specific learning rate. While it may be somewhat expected that these two methods, both being first-order optimization techniques, would exhibit some similarity, the importance of this result lies in its ability to establish a specific condition under which SGDM behaves similarly to SGD.\n\nHowever, as you rightly suggest, a more generalized result would be valuable. If the paper could establish conditions where, for certain small learning rates, SGDM and SGD consistently approximate each other, it would offer broader insights into the relationship between these methods and provide more guidance on when to expect their similarities. This would enhance our understanding of the interplay between learning rates and optimization algorithms, potentially leading to more informed choices in practical applications.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Prior research has noted that the advantage of incorporating momentum in the training of neural networks is somewhat limited, and there hasn't been a definitive theoretical outcome regarding the effectiveness of momentum in the stochastic setting. Therefore, demonstrating that momentum does not significantly reduce variance or improve generalization can lead to savings in computational resources and memory usage during deep neural network training. A noteworthy technical innovation in this work is the introduction of a novel Stochastic Differential Equation (SDE) to represent the diminishing learning rate in Stochastic Gradient Descent with Momentum (SGDM) and Stochastic Gradient Descent (SGD).", "weaknesses": "The assumptions that the paper considers to show that SGDM approximate SGD is too restrictive. It would be great if they can justify these too restrictive assumptions.", "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": "1YL3KsX9N3", "reviewer": "Reviewer_EHVB", "summary": "This paper considers the role of momentum in SGD training. When the learning rate is small but the gradient noise is not small, the paper shows that the trajectories of SGD and SGDM are close. Therefore, the momentum component in SGDM adds only a small effect on top of SGD. This paper also conducts experiments to verify that the optimization and generalization of SGDM and SGD are close when the learning rate is small and the mini-batch size is not large (such that the gradient noise is not small).", "questions": "See above.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Theoretically, momentum provably improves the (minimax) convergence rate of GD for (strongly convex) smooth optimization. Practically, momentum is often used with SGD (as GD is expensive in practical machine learning optimization). In the presence of gradient noise, will momentum still offer benefits? This is the central question studied in this work. \n\nFirst, this work presents a theory, showing that when the learning rate tends to zero, the trajectories of SGD and SGDM are close in two regimes. The first regime is when the gradient noise is large (with $O(1/\\eta)$ variance) but the total running distance is constant. The second regime allows the total running steps to be large but requires an overparameterized model where the global solutions form a manifold. \n\nSecond, this work conducts experiments to verify that for both training and test, SGDM and SGD are close when the learning rate is small and when the batch size is not too large. \n\nThe writing is clear overall. Some places require clarifications. See more in the next section. \n\nI did not check the proof.", "weaknesses": "1. I feel the introduction/motivation could use some polishment. Specifically, the benefit of momentum is for stabilizing GD with a large learning rate ($>2/L$, in terms of improving the minimax convergence rate for (strongly convex) smooth optimization) as discussed in the first paragraph of the introduction. However, I did not recall a theoretical result that claims acceleration of momentum for GD with a learning rate. \n\nOne of the main contributions of this paper is to show that momentum does not change SGD trajectory much when the learning rate is sufficiently small. Therefore, momentum is not very helpful. \n\nI am not sure how surprising this is, since the benefit of momentum is only claimed for GD with a very large learning rate. \n\nAlso, the condition that $1-\\\\beta\\^{(n)}_k = \\\\Theta ((\\\\eta^{(n)})\\^{\\\\alpha} )$ (in definition 3.1) for $\\alpha \\in [0,1)$ (in Theorem 3.5) basically says that the history gradient will be forgotten fast (in a nearly exponential rate) when SGDM runs $\\\\Theta(1/\\\\eta)$ steps. So it's not very surprising that SGDM will be close to SGD. \n\n\n2. In the last paragraph of Section 3, the paper briefly mentions a phenomenon for $\\\\alpha>1$. However, the theorem only covers $\\\\alpha <1$. Could the author elaborate on what happens to SGDM when $\\\\alpha>1$?\n\n\n3. In the paragraph after Lemma 2.4, $\\\\sigma$ is set to be $1/\\sqrt{\\eta}$ to ensure a non-trivial gradient noise scaling. However, later in Theorem 3.5, it allows that $\\\\sigma \\le \\eta^{-1/2}$, and in Theorem 4.5, it allows $\\\\alpha = \\\\Theta(1)$. These seem to be different from the prior discussions. Could the authors clarify the scaling of the noise?  In particular, can Theorem 3.5 be applied to GD and GDM (i.e., $\\\\sigma = 0$)?\n\n\n4. This paper mentions two versions of SGDM, that is, eq (1) used in practice and eq (3) used in theory. I appreciate the usage of $\\\\gamma$ and $\\\\eta$ to differentiate the learning rate for those two versions, which helps highlight the difference in the effective optimization length. \n\nIn the ImageNet experiment, the learning rate for eq (1) is manually rescaled when compared to SGD (to align with the theory setup in eq (3)), which is good. However, for training ResNet-32 on CIFAR-10, the paper writes that \"We first grid search to find the best learning\nrate for the standard SGDM ($\\\\ell=1$), and then we perform SGD and SGDM with that same learning rate for different levels of $\\\\ell$\". I am wondering if the learning rate for SGD needs a rescaling here.", "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)