Flatter, Faster: Scaling Momentum for Optimal Speedup of SGD
Authors: Aditya Cowsik
Tankut Can
Paolo Glorioso
Venue: Submitted to ICLR 2023 (2023)
Keywords: SGD
momentum
acceleration
generalization
scaling limit
deep learning
implicit bias
implicit regularization

Abstract: Commonly used optimization algorithms often show a trade-off between good generalization and fast training times. For instance, stochastic gradient descent (SGD) tends to have good generalization; however, adaptive gradient methods have superior training times. Momentum can help accelerate training with SGD, but so far there has been no principled way to select the momentum hyperparameter. Here we study implicit bias arising from the interplay between SGD with label noise and momentum in the training of overparameterized neural networks. We find that scaling the momentum hyperparameter $1-\beta$ with the learning rate to the power of $2/3$ maximally accelerates training, without sacrificing generalization. To analytically derive this result we develop an architecture-independent framework, where the main assumption is the existence of a degenerate manifold of global minimizers, as is natural in overparameterized models. Training dynamics display the emergence of two characteristic timescales that are well-separated for generic values of the hyperparameters. The maximum acceleration of training is reached when these two timescales meet, which in turn determines the scaling limit we propose. We perform experiments, including matrix sensing and ResNet on CIFAR10, which provide evidence for the robustness of these results.

Under review as a conference paper at ICLR 2023
FLATTER, FASTER: SCALING MOMENTUM FOR OPTI- MAL SPEEDUP OF SGD
Anonymous authors Paper under double-blind review
ABSTRACT
Commonly used optimization algorithms often show a trade-off between good generalization and fast training times. For instance, stochastic gradient descent (SGD) tends to have good generalization; however, adaptive gradient methods have superior training times. Momentum can help accelerate training with SGD,, but so far there has been no principled way to select the momentum hyperparame- ter. Here we study training dynamics arising from the interplay between SGD with label noise and momentum in the training of overparametrized neural networks. We find that scaling the momentum hyperparameter - B with the learning rate the power of 2/3 maximally accelerates training, without sacrificing general- ization. To analytically derive this result we develop an architecture aechitectrr-iideeenddee framework, where the main assumption is the existence of a degenerate manifold of global minimizers, as is natural in overparametrized models. Training dynam- ics display the emergence of two characteristic timescales that are well-separated for generic values of the hyperparameters. The maximum acceleration of training is reached when these two timescales meet, which turn determines the scal- ing limit we propose. We confirm our scaling rule for synthetic regression prob- lems (matrix sensing and teacher-student paradigm) and classification for realistic datasets (ResNet-16 8 on CIFAR 10, 6-layer MLP on FashionMNIST), suggesting the robustness our scaling rule to variations architectures and datasets.
INTRODUCTION
The modern paradigm for optimization of deep neural networks has engineers working with vastly overparametrized models and training to near perfect accuracy (Zhang et al., 2017). In this setting, a model will typically have not just isolated minima in parameter space, but a continuous set of pa-
minimizers, not all of which generalize well. Liu et al. (2020) demonstrate that depending rameter initialization and hyperparameters, stochastic gradient descent (SGD) is capable of finding minima with wildly different test accuracies. Thus, the power of particular optimization method lies in its ability to select a minimum that generalizes amongst this vast set. In other words, good generalization relies on the implicit bias or regularization of optimization algorithm.
There is a significant body of evidence that training deep nets with SGD leads to good generalization. Intuitively, SGD appears to prefer flatter minima (Keskar al., 2017: Wu et al., 2018: Xie et al., 2020), and flatter minima generalize better (Hochreiter & Schmidhuber, 1997). More recently, a variant of SGD which introduces "algorithmic" label noise has been especially amenable rigorous treatment. In the overparametrized setting Blanc et al. (2020) were able to rigorously determine that SGD with label noise converges not just to any minimum, but to those minima that lead to the smallest trace norm of the Hessian. However, Li et al. (2022) show that the dynamics of this regularization happen on a timescale proportional to the inverse square of the learning rate 7) - much slower than the time to first converge to interpolating solution. Therefore consider the setting where we remain near the local minima, which is responsible for significant regularization after the initial convergence of the train loss (Blanc et al., 2020).
With the recent explosion in size of both models and datasets, training time has become an impor- tant consideration in addition to asymptotic generalization error. In this context, adaptive gradient methods such as Adam (Kingma& & Ba, 2015) are unilaterally preferred over variants of SGD, even though they often yield worse generalization errors in practical settings (Keskar & Socher, 2017;

---

Under review as a conference paper at ICLR 2023
Wilson et al., 2017), though extensive hyperparameter tuning (Choi et al., 2019) or scheduling (Xie et al., 2022) can potentially obviate this problem.
These two constraints motivate careful analysis of how momentum accelerates SGD. Classic work on acceleration methods, which we refer to generally as momentum, have found a provable benefit in the deterministic setting, where gradient updates have no error. However, rigorous guarantees have been harder to find in the stochastic setting, and remain limited by strict conditions on the noise (Polyak, 1987; Kidambi et al., 2018) or model class and dataset structure Lee al. (2022).
In this work, we show that there exists a scaling limit for SGD with momentum (SGDM) which provably increases the rate of convergence.
Notation, In what follows, we denote by C", for n = 0. 1, the set of functions with con- tinuous Th derivatives. For any function f. 0f(u] and 0² flu, v) will denote directional first and second derivatives along directions defined by vectors
E RD. We may occasionally also write |-][== = = O f(e,, Eg) Enj- Given a submanifold I C RRD and w E [, we denote by Te I the tangent space to I in we and by PL (w) the projector onto I (we will often omit the dependence on w and simply write PL). Given a matrix H E IR x R., we will denote by H the transpose, and Ht the pseudoinverse of H.
1.1 HEURISTIC EXPLANATION FOR OPTIMAL MOMENTUM-BASER SPEEDUP
Deep neural networks typically posses a manifold of parametrizations with zero training error. Be- cause the gradients of the loss function vanish along this manifold, the dynamics of the weights is completely frozen under gradient descent. However, as appreciated by Wei & Schwab (2019) and Blane et al. (2020), noise can generate an average drift of the weights along the manifold. In par- ticular, SGD noise can drive the weights to lower-curvature region which, heuristically, explains the good generalization properties of SGD, Separately, it is well-known that adding momentum typically leads to acceleration in training Sutskever et al. (2013) Below, we will see that there is nontrivial interplay between the drift induced by noise and momentum, and find that accelcration along the valley is maximized by particular hyperparameter choice in the limit of small learning rate. In this section we will illustrate the main intuition leading to this prediction using heuristic arguments, and defer a more complete discussion to Sec. 3.
We model momentum SGD with label noise using the following formulation
1k+ = BT. V((w)) + ea
Wk++ wh + 77*+++
(1)
where v is the learning rate, B is the (heavy-ball) momentum hyperparameter (as introduced by Polyak (1964)), w E IR D denotes the weights and E F. denotes momentum (also called the auxiliary variable). L RD + R is the training loss, and a IDD + Dxr is the noise function, whose dependence on the weights allows to model the gradient noise due SGD and label noise. Specifically this admits modeling phenomena such as automatic variance reduction (Liu & Belkin, 2020), and expected smoothness, satisfying only very general assumptions such as those developed by Khaled & Richtárik (2020). Finally, & E R' is sampled i.i.d. at every timestep l: from distribution with zero mean and unit variance, and >> 0.
We will now present heuristic description of the drift dynamics that is induced by the noise along a manifold of minimizers I = {w L(w) = 0} S "R,,, in the limit e + 0. In practice, this limit corresponds to choosing small strength of label noise and large minibatch size. Let us assume that the weights at initial time Wo are already close to r, and To = O. Because of this, the gradients of L at Wy are very small, and so only fluctuations transverse to the manifold will generate sys- tematic drifts. Denoting by Swy, = wo- - Wh the displacement of the weights after k timesteps, let us Taylor expand the first equation in (1) get T BTA A ((wo)))oon + Ea (o(ww) (w By construction, the Hessian L²L(w) vanishes along the directions tangent to r, while in the transverse direction we have an Canstin-UUllebeed takes to this process to relax to the stationary state, or mixing time, is π = θ (1/(1-19 3)) as + 1, which can be anticipated from the first equation (1) since T&+1 - TT -(1 - 8) (see Sec. 3.3 for more detailed derivation). After this time, the variance of this linearized OU process becomes independent of the time step, and can be estimated to be (see Appendix
(OU) process. The number of time steps
F): ((6w)") owf)) @ (e2n/(1- (1))), B)). where (-..) denotes the noise average. To keep track of the displacements in the longitudinal directions, dwf, we need to look at the cubic order in
2

---

Under review as a conference paper at ICLR 2023
the Taylor expansion of L(wo + 5wn), ie. 8 (V 5wu]- Let PL be the projector onto the tangent space. The expectation value of momentum, upon applying the longitudinal projec- tor, is Pt (TK+1)) = -. (VL(wy)) = - - (P,, - 3 (PL (PLVI) buyy) The variance of the transverse displacements therefore gener- ates longitudinal motion, and the latter will then scale as ( = θ ( (1- 3)2). Using the second equation in (1), we see that Pt. (2. (6wx)) = P Define T2 to be the number of time steps it takes so that the displacement of wy along I' becomes finite number as we take E - 0 first, and 1 + 0 afterward. From the above considerations, we then find that 72 = (1-))
"ów,)) =
The key observation now is that, when the weights are initialized near the valley, the convergence time controlled by the largest timescale among тΙ and 72- While T2 decreases as + 1, TI increases. Therefore, the optimal speedup of training is achieved when these two timescales intersect
Ti = T2. which happens for 1 - B = C172/3, More generally, we consider a double scaling where B + as 17 + 0 according to
B = 1 - 1
We consider this scaling limit in this paper, and we thus find that the scaling power y = 2//3 achieves the optimal speed-up along the zero-loss manifold.
1.2 LIMIT DRIFT- DIFFUSIIFFSIN
We now describe the rationale for obtaining the limiting drift-diffusion on the zero-loss manifold, for a process of the form (I) which foreshadows the rigorous results presented in section As discussed above, the motion along the manifold is slow, as it takes Oc==2) time steps to have finite amount longitudinal drift. We want to extract this slow longitudinal sector of the dynamics by projecting out the fast-moving components of the weights. For stable values of the optimization hyperparameters. the noiseless (ea = 0) dynamics (1) will map a generic pair (TT, w), as k: - co, to (0, Woo), where Woo r. Define & D
R. to be this mapping, i.c. Φ( w) Woo-
As we now show, when e > 0, & can be used precisely to project onto the slow, noise-induced longitudinal dynamics. Let collectively denote o (Mz, wxx) and write eq. (1) as Fk+1
Ih + F(x))= + cā (xx)). We can perform a Taylor expansion in e to obtain ((I++1) - ((x,) ) = (<+ (x,) + (zz)
eă (x+ +.... Therefore, denoting Y( = k² ((xx).
the limit dynamics as + 0 can be well-approximated by the continuous time equation
dY =
(()))) ddt)),
(3)
where we interpret the time increment di = e². and (E?) 1. Note that until here we have not taken
small learning rate limit. The learning rate can be finite, as far the map ((T, exists. The small noise limit is sufficient to allow a continuous-ting description of the limit dynamics because the noise-induced drift-diffusion along the valley requires O(E 2) timesteps lead appreciable longitudinal displacements.
similar approach to what we just described was used in Li et al. (2022), although in our case the limit drift-diffusion is obtained in the small noise limit, rather than small learning rate. The
reason for our choice is that, since we scale 3 according to (2), the deterministic part of eq. (1) becomes degenerate as we take + 0, in which case it would not possible to apply the mathematical framework of Katzenberger (1991) on which our results below rely. To further simplify our analysis. particularly the statement of Theorem B.4, we will further take + 0 after taking - O, and retain only lcading order contributions in 17.
The main contributions of this paper are:
1. We develop general formalism to study SGD with (heavy-ball) momentum in Sec. 3, extending the framework of Li et al. (2022) to study convergence rates and generalization with momentum.
2. We find novel scaling regime of the momentum hyperparameter 1 1-3 ~ 177,, and demon- strate qualitative change the noise-induced training dynamics as 7 is varied.
3. We identify a special scaling limit, 1 - B ~ 2/33, where training achieves a maximal speedup at fixed learning rate n).
3

---

Under review as a conference paper at ICLR 2023
4. In Sec. 4, we demonstrate the relevance of our theory with experiments on toy models (2- layer neural networks, and matrix sensing) as well as realistic models and datasets (ResNet- 18 on CIFARIO).
RELATED WORKS
Loss Landscape in Overparametrized Networks The geometry of the loss landscape very hard understand for real-world models. Choromanska al. (2015) conjectured, based em- pirical observations and on an idealized model, that most local minima have similar loss function values. Subsequent literature has shown in wider generality the existence of manifold connecting degenerate minima of the loss function, particularly in overparametrized models. This was sup- ported by work mode connectivity (Freeman & Bruna, 2017; Garipov et al., 2018; Draxler al., 2018; Kuditipudi et al., 2019), as well as on empirical observations that the loss Hessian possesses large of (nearly) vanishing eigenvalues (Sagun et al., 2016; 2017). In particular, Nguyen (2019) showed that for overparametrized networks with piecewise linear activations, all global minima are connected within a unique valley.
The Implicit Regularization of SGD Wei & Schwab (2019), assuming the existence of a zero- loss valley, observed that SGD noise leads to a decrease in the trace of the Hessian. Blanc et al. (2020) demonstrated that SGD with label noise in the overparametrized regime induces a regularized loss that accounts for the decrease in the trace of the Hessian. Damian et al. (2021) extend this analysis to finite learning rate. HaoChen et al. (2021) study the effect of non-isotropic label noise in SGD and find theoretical advantage in quadratic overparametrized model. Wu al. (2022) show that only minima with small enough Hessians (in Frobenius norm) are stable under SGD, The specific regularization induced by SGD was found in quadratic models (Pillaud-Vive) et al., 2022), 2-layer Relu networks (Blanc et al., 2020), linear models (Li et al., 2022), diagonal networks (Pesme et al., 2021). Additionally, Kunin et al. (2021) and Xie al. (2021) studied the diffusive dynamics induced by SGD both empirically and in a simple theoretical model.
Momentum in SGD and Adaptive Algorithms Momentum is a general term applied to tech- niques introduced to speed up gradient descent. Popular implementations include Nesterov (Nes- terov, 1983) and Heavy Ball (HB) or Polyak (Polyak, 1964). We focus on the latter in this paper, which we refer simply as momentum. Momentum provably improves convergence time in the deterministic setting. Intuitively, introducing 3 in gives the motion in parameter space an effective "inertia" or memory, which promotes motion not strictly following the local gradient, but moving rather along the directions which persistently decrease the loss function across iterations Polyak (1964); Sutskever et al. (2013). Less is known rigorously when stochastic gradient updates are used. Indeed, Polyak (1987) suggests the benefits of acceleration with momentum disappear with stochas- tic optimization unless certain conditions are placed on the properties of the noise. See also (Jain et al., 2018; Kidambi et al., 2018) for more discussion and background on this issuc. Nevertheless, in practice it is widely appreciated that momentum is important for convergence and generalization (Sutskever et al., 2013), and widely used in modern adaptive gradient algorithms Kingma & Ba (2015). Some limited results have been obtained showing speedup in the mean-field approximation (Mannelli & Urbani, 2021) and linear regression Jain al. (2018). Modifications Nesterov mo- mentum to make it more amenable to stochasticity (Liu & Belkin, 2020; Allen-Zhu, 2017), and near saddle points (Xie et al., 2022) have also been considered.
3 THEORETICAL RESULTS
3.1 GENERAL SETUP
Following the line from section 1.2 In this and the following section, we will rigorously derive the limiting drift-diffusion equation for the weights on the zero-loss manifold I, and extract the timescale T2 associated to this noise-induced motion. In Sec. 3.3 we will then compare T2 to the timescale 7i associated the noiseless dynamics and evaluate the optimal value of y discussed around Eq. (2). We will use Eq. (1) to model momentum SGD. As illustrated in Sec. 1.1, the drift is controlled by the second moment of fluctuations, and we thus expect the drift timescale to be (c)2).
4

---

Under review as a conference paper at ICLR 2023
We will then rescale time & t/c², that the motion in the units of t is ((1) 0. More explicitly, take en to be positive sequence such that en + 0 as n + co. For each n we consider the stochastic process that solves Eq. (1):
(n(1) = (0) / o(x.)dZn + / F(X.)dAn,
(4)
with
An(O)
(t)
(n( cn
Σ &
(5)
2
1m
and where ((t kk) (T wg),, (X) (a, no) and F(X) - (13 (13-1)- VL(w), nBB
VL(w))). |x] denotes the integer part of a real number z. See Appendix B.2 for a proof of equiva- lence between (1) and (4).
The loss function L RDD R O function whose first 3 derivatives are locally Lipschitz, a is continuous, and r {w RDD L(w) = 0} is a C CC-submanifold of IRP of dimension M, with 0 S M D. Additionally, for w r, (ank(V2L(uuu
Assumption 3.1.
= M.
Assumption 3.2. There exists an open neighborhood U of {0} I S R R such that the gradient descent starting in U converges to point x = (T, w) E (0) X More explicitly, for U, let (b 0) and sb(x,, ++1 +1) tb(x,k) k) F(,(z, k)), iee ((x,k) is the tth iteration
of x + F(x). Then ((x) = lim (b k) exists and is in I. As a consequence, & E C² on U (Falconer, 1983).
3.2 LIMITING DRIFT- DIFFUSSION MOMENTUM SGD
In this section we shall obtain the explicit expression for the limiting drift-diffusion. The general framework is based on Katzenberger (1991) (reviewed in Appendix B). Before stating the result, we will need introduce few objects.
Definition 3.3. For a symmetric matrix H E RRD x RD. and WH (E x X E ΣT HHe HE J). we define the operator Lu WH WH with LH = (H,)) + 30-2,11--10 [[S, HJ, H.. with [S, H) SH HS. It can be shown that the operator Li invertible (see Lemma C.3).
Consider the process in Eq. (4). Note that, while initialization we can have (,) & T, the solution (t) + I as n + i.e. it becomes discontinuous. This is an effect of the speed- up of time introduced around Eqs. (5), To overcome this issue, it is convenient to introduce (t) = (t) (0), A,,t) (0),, so that Y, (0) I is initialized the manifold. Theorem 3.4 (Informal). Suppose the loss function L, the noise function the manifold of mini- mizers and the neighborhood satisfy ussumptions (3. and (3.2), and that (o. = X() E U. Then, as 0, and subsequently laking 1 - 0, (n(()) converges to Y (1), where the latter satisfies the limiting drift- dift--iiffuson equation
dY -( + Pr.od 12 (V (V(L))EE.r) 2 (L
(6)
2 27 P (VLIC Errldt,
where W(t) is Wiener process.
A rigorous version of this theorem is given in section B. The first term in Eq. (6) induces diffusion the longitudinal direction. The second term is of geometrical nature, and is necessary to guarantee that (() remains on r. The second line describes the drift induced by the transverse fluctuations.
Eq. (6) resembles form that found Li al. (2022), although there are two crucial differences. First, time has been rescaled using the strength of the noise rather than the learning rate. The different rescaling was necessary as the forcing term F in Eq. (4) depends non-homee non-hhmggennoully on n, and thus the theory of Katzenberger (1991) would not be directly applied had we taken the small learning rate limit. Second, and more crucially, the drift terms in Eq. (6) are proportional to 2-22,, which is key ingredient leading to the change in hierarchy of the timescales discussed in Sec. 1.1. One final difference, is that the last term involves the operator Ln instead of the Lyapunov operator.
5

---

Under review as a conference paper at ICLR 2023
For y < 21 Lu reduces to the Lyapunov operator LH at leading order in n, with L3 {H, S).
For y > 2 however, we cannot neglect the n-dependent term L (see discussion at the end of Appendix C).
Corollary 3.5. In the case of label noise, i.e. when, for w E ET, = cV2L, 12L,.or for some constant
c > 0, Eg. (6) reduces to
-12-22 dY = 4C² PLV IVV
where we have rescaled time back to t ks, i.e. we performed tc?
3.3 SEPARATION OF TIMESCALES AND OPTIMAL MOMENTUM SCALING
The above results provide the estimate for the timescale of the drift along the zero-loss valley. As discussed in Sec. 1.1, training along the zero-loss manifold I' is maximally accelerated if this time scale is equal to the timescales for relaxation of off-valley perturbations, As we take O, this relaxation is governed by the nonzero eigenvalues of the Hessian as well as by the learning rate y and momentum B. Therefore we expect Ti ( and this will be confirmed by the analysis below. It will be therefore sufficient to obtain the leading order expression of Ti by focusing on the noiseless
dynamics. Additionally, since we are interested in local relaxation, it will suffice to look at the linearized dynamics around I.
Working in the extended phase space z4 (TT, wx), and linearizing Eq. (1) around fixed point xx = (0, w²), with w' H, the linearized update rule is 8 [xx" )ozz, where dz, = I3 x= and J(x)) is the Jacobian evaluated the fixed point (with the explicit form given in Eq. (133)).
Denote by q¹ the eigenvector and X; the corresponding eigenvalue of the Hessian. We show in Appendix E that the Jacobian is diagonalized by the eigenvectors k (nin', .4") with eigenvalues
11 [1+ 3 n7 + V 31+ 1 48)
and μ4 = B1pn We proceed to study the decay rate of the different modes the Jacobian to draw conclusions about the characteristic timescales of fluctuations around the valley.
On the valley, the Hessian will have a number of "zero modes" with A = 1. These lead to two distinct modes in the present setting with momentum, which we distinguish as pure and mixed. The first pure longitudinal mode is an exact zero mode which has n² = with ki = (0, g"), corresponding to translations of the parameters along the manifold, and keeping π = 0 at its fixed point value. The second mode is a mixed longitudinal mode with x-_ = B with ki = (-(1- - 3)/ (28r))' g²). This mode has a component of π along the valley, which must subsequently decay because the equilibrium is single point æ Therefore, this mode decays the characteristic rate 3 for gleaned directly from Eq. (1).
Longitudinal motion:
When the w and π are perturbed along the transverse directions q' with positive A, the relaxation behavior exhibits a qualitative change depending on B. Using the scaling function 3= Cry", , for small learning rate, the spectrum is purely real for y < 1/2, and comes in complex conjugate pairs for > 1/2. This leads to two distinct scaling behaviors for the set of timescales.
Transverse motion:
Defining positive Ci min(A, X; > 0}. we find: 1) For 7 < 1/2, transverse modes are purely decaying (1 ) < 10z < (1 - (ci "o) with the lower bound set by the mixed longitudinal mode. For 7 > 1/2, the transverse modes are oscillatory but with an envelope that decays like ( ~ (1 k// We leave the derivation of these results to Appendix (E).
Collecting these results, we can describe the hierarchy of timescales Ti in the deterministic regime as function of y (excluding the pure longitudinal zero mode):
T1 Long.
7 1/2
Y > 1/2
1)'
11'
Transv.
17' -,,""
n",1"
These are illustrated schematically in Fig. 1(a), where the finite timescales are shown as a function of y. We compare these "equilibration" timescales Ti, i.e. characteristic timescales associated with

---

Under review as a conference paper at ICLR 2023
relaxation back to the zero-loss manifold, with the timescale 72 172(v-1) associated with drift- diffusion of the noise-driven motion along the zero-loss manifold Eq. (6). For small y, the timescale associated with the drift-diffusion along the valley is much faster than that associated with the laxation of the dynamics toward steady state. Transverse and mixed longitudinal fluctuations relax much faster than the motion along the valley, and produce effective drift toward the minimizer of the implicit regularizer. However, the timescales collide at Y 2/3, suggesting a transition to qualitatively different transport above this value, where the transverse and the mixed longitudi- nal dynamics, having a long timescale, will disrupt the longitudinal drift Eq. (6). This leads us to propose y 2 as the optimal choice for SGD training. We consistently find evidence for such a qualitative transition in our experiments below. In addition, we see that speedup of SGD with label noise is in fact maximal at this value where the timescales meet.
3.4 A SOLVABLE EXAMPLE
In this section we analyse a model that will allow determine, top of the optimal exponent 7
3. also the prefactor C. We will specify to 2-layer linear MLP, which is sufficient to describe the transition in the hierarchy of timescales described above, and is simple enough to exactly compute C. We will show in Sec. 4.1 that C depends only mildly on the activation function. We apply simple matching principle to determine C, by asking that the deterministic timescale 71 is equal to the drift-diffusion timescale 72. In the previous section, we found the critical 2/3 by requiring these timescales have the same scaling in 7). In order to determine C, we need more details of the model architecture.
Definition 3.6 (UV model). We define the UV model as a 2-layer linear network parametrized by f(x) = vn V E R"" where æ E Rdd V Raxd, and U E RR For = m = 1, we refer to this as the vector UV model (Rennie & Srebro, 2005; Saxe et al., 2014; Lewkowyzz et al., 2020).
For a training dataset D = {(x , = L,..., P}. the dataset covariance matrix is Eig f Ee xf27, and the dataset variance is #l trΣ. For mean-squared error loss, it is possible to explicitly determine the trace of the Hessian (see Appendix D). SGD with label noise introduces yo - y² +c,, where (E?) 1, from which
identify jrjaa P-LV Jz (3x", where μs runs over all parameter indices, j E [ma]. and E [P]. With this choice, the SGD noise covariance satisfies ao-" = -P-122 Equipped with this, we may use Corollary 3.5 with p-1 determine the effective drift (presented in Appendix D). For the vector UV model, the expression simplifies to dY T Y with
1/2 2
T2
2nPC2
The timescale of the fast initial phase 7,_ (C/2)n² follows from the previous section. Then requiring Ti = implies not only 2/3, but
1/3
6 1²122
C =
(10)
Pn
One particular feature to note here is that C will be small for overparametrized models and/or train- ing with large datasets.
EXPERIMENTAL VALIDATION
4.1 2-LAYER MLP WITH LINEAR AND NON-LINEAR ACTIVATIONS
The first experiment consider is the vector UV model analyzed in Sec. 3.4. Our goal with this experiment is to analyze a simple model, and show quantitative agreement with theoretical expectations. Though simple, this model shows almost all of the important features in our analysis: the 2(1 - 7) exponent below = 3- evaluated in Sec. 3.4.
2 the y exponent above 2 and the constant C theoretically 3.
To this end, we extract the timescales at different values of y and show them in Fig. 1. We train on an artificially generated dataset D ((x0, 1/)) ) with x ~N(0,1) 1) and y 0. We use the
7

---

Under review as a conference paper at ICLR 2023
linear activation
72 diff-diffusion
Tara
Throseu (arb. units)
Turesade Powar alle)
maxe - 7).
16-
0.
1.2-
1.0
aquilibrattin
0.0
0.8
d)
tank: activation
Relu activation
Data
Dasa raad 22 73.
Timeseale Power: aig)
1.8
max 111 7,,0
Tunerall Power: srrr)
18-
2.-
164
2.1
144
122
1.0
100
60
0.6
0.-
06
0.0
322
0.4 0.6
1.0
3.9
0/2
0.4
016 0.8
7
Figure 1: Timescale of training as a function of y. a): theoretical prediction with blue line repre- senting the timescale for equilibration while the black line shows the timescale of the drift along P. The maximum these two gives the overall timescale. In b). c), and d) we demonstrate this result with the vector UV model with linear, tanh and Relu activations respectively.
full dataset train. From (124)) we know that the norm of the weights follows an approximately exponential trajectory as it approaches the widest minimum (U = V = =). We therefore measure convergence timescale, Te, by fitting an exponential ae -t/T/ to the squared distance from the origin, |U| + V To extract the scaling Te with y we perform SGD label noise with learning rates E [10 10='] and corresponding momentum parameters 8 = 1 Ci²².. We fit the timescale Ton a(1) To)) (see Fig. 1(b)). Imposing that To be
to a power-law in the learning rate T.(,.)) independent of y, as predicted by theory, we found the numerical value C 0.2, which is consistent with the theoretical estimate of C = 0.17 from See. 3.4. We find consistency with prediction across all the values of y we simulated. Note that, for ? > 3/ the timescale estimate fluctuates more which consequence of having a slower timescale for the transverse modes. As discussed at the end of Sec. 3.3, such slowness disrupts the drift motion along the manifold. 7 = zzz is clearly the optimal scaling.
We repeated the same experiments using nonlinear activations, specifically we considered tanh and ReLU acting the first layer. The timescales for tanh shown in Fig. 1(c), and we refer the reader to Appendix A for the ReLU case. The optimal scaling value is still 7 for both tanh and ReLU and the optimal C remains close to 0.2.
4.2 RESNET18 8 ON CIFAR10
We now verify our predictions on realistic problem, which will demonstrate the robustness of our analysis. We focus on R8 (He al., 2016), specifically implemented by Liu (2021), classifier trained on CIFARIO Krizhevsky et al. (2009). We aim to extrapolate the theory by showing optimal acceleration with our hyperparameter choice once training reaches an interpolating solution. To this purpose, we initialize the network the valley, obtained starting from random weight values and training the network using full batch gradient descent without label noise and with a fixed value 3 0.9 until it reaches perfect training accuracy. With this initialization, SGD and label noise for a fixed number of epochs multiple times for various values the momentum hyperparameter B. Finally, project the weights back onto the valley before recording the final test
then train with
8

---

Under review as a conference paper at ICLR 2023
a
b)
Optinal Momentam Sealing CIFARIO with Re
Noise-loss GD
One Minus Optimal Momeatum [1 AJ
Data 0.1125
Small y
10-2-
Large
Projection Zoro Looe Manifold
10-1
10-3
10-t
10-1
Learoing Rate (9)
Figure 2: Classification of CIFAR10 using a ResNet18 model. Subfigure a) shows our training protocol which first use noiseless gradient descent (black) to reach the zero-loss manifold, then to perform SGD label noise with various (blue, red) values of and 8, before finally projecting onto the valley (green) and measuring the test accuracy. Subfigure b) shows the scaling of the optimal momentum, 3., as function of 1). We perform a power-law fit whose exponent 7 = 0.660 matches very closely to the value implied by the theory ? z Notice we also extract the constant 0.11.
accuracy. This last step be viewed noise annealing and allows compare the performance of training the drift phase for the different values of B. From this procedure we extract the optimal momentum parameter 3. (1)) that maximizes the best test accuracy during training function of the learning rate, which we can then compare with the theoretical prediction.
As shown in Fig. 2(b). 1 B. follows the power law we predicted almost exactly. The optimal choice for speedup does not have to coincide with the optimal choice for generalization. Strikingly, this optimal choice of scaling also leads the best generalization in a realistic setting! This can be easily interpreted if we assume that the more we decrease the Hessian the better our model generalizes and by applying the fact that our scaling leads to the fastest transport along the manifold. The second important point is the value of the constant C' 0.1 found as the coefficient of the power-law fit. we set 7) = this corresponds to setting 3, = 0.9 which is the traditionally recommended value. The result here, can therefore be viewed generalization of this common wisdom. For more experiments, we refer the reader to Appendices A.2 and A.4.
CONCLUSION
We studied the implicit regularization of SGD with label noise and momentum in the limit of small noise and learning rate. We found that there is an interplay between the speedup of momentum and the limiting diffuusion generated by SGD noise. This gives rise to two characteristic timescales associated to the training dynamics, and the longest timescale essentially governs the training time. Maximum acceleration is thus reached when these two timescales coincide, which lead us to iden- tifying an optimal scaling of the hyperparameters. This optimal scaling corresponds not only to faster training but also to superior generalization. More generally. we have shown how momentum can significantly enrich the dynamics of learning with SGD, modulating between qualitatively dif- ferent phases of learning characaterized by different timescales and dynamical behavior. It would be interesting to explore our scaling limit in statistical mechanical theories of learning to uncover further nontrivial effects on feature extraction and generalization in the phases of learning (Jacot et al., 2018; Roberts et al., 2021; Yang et al., 2022). For future work, it will be very interesting to generalize this result to adaptive optimization algorithms such as Adam and its variants, to use this principle to design new adaptive algorithms, and to study the interplay between the scaling we found and the hyperparameter schedule.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Zeyuan Allen-Zhut Katyusha: The first direct accelcration of stochastic gradient methods. The Journal of Machine Learning Research, 18(1):8194 8244, 2017.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an OrsleeinUhllenbeck like process. In Conference on learning theory, pp. 483-513. PMLR, 2020.
Dami Choi, Christopher J Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George E Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv:1910.05446, 2019.
Anna Choromanska, Mikael Henaff, Michael Mathieu, Gérard Ben Arous, and Yann LeCun. The loss surfaces of multilayer networks. In Artificial intelligence and statisties, pp. 192 204. PMLR, 2015.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global minimiz- ers. Advances in Neural Information Processing Systems, 14:27449-274611 2021.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309- 1318. PMLR, 2018.
KJ Falconer. Differentiation of the limit mapping in dynamical system. Journal of the London Mathematical Society, 2(2):356-37,, 1983.
Daniel Freeman and Joan Bruna. Topology and geometry of half-rectified network optimization. ICLR, 2017.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. Advances in neural information processing systems, 31, 2018.
J. HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, 2315-2357. PMLR, 15-19 Aug 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemnl: Convergence and gen- cralization in neural networks. Advances in neural information processing systems. 31, 2018.
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerat- ing stochastic gradient descent for least squares regression. In Conference On Learning Theory, pp. 545-604. PMLR, 2018.
Gary Shon Katzenberger. Solutions of stochastic differential equation forced onto a manifold by large drift. The Annals of Probability, 19(4): 118--688,,,, 1991.
Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from Adam to SGD. arXiv preprint arXiv:1712.07622 2017.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2017.
Ahmed Khaled and Peter Richtárik. Better theory for sgd in the nonconvex world. arXiv preprint arXiv:2002.03329, 2020.
10

---

Under review as a conference paper at ICLR 2023
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the insufficiency of ex- isting momentum schemes for stochastic optimization. Information Theory and Applications Workshop (ITA), pp. 1-9. IEEE, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Rohith Kuditipudi, Xiang Wang, Holden Lee, Yi Zhang, Zhiyuan Li, Wei Hu, Rong Ge, and Sanjeev Arora. Explaining landscape connectivity of low-cost solutions for multilayer nets. Advances in neural information processing systems, 32, 2019.
Daniel Kumin, Javier Sagastuy- Beeastuu-BBena, Lauren Gillespie, Eshed Margalit, Hidenori Tanaka, Surya Ganguli, and Daniel LK Yamins. Limiting dynamics of sgd: Modified loss, phase space oscilla- tions, and anomalous diffusion. arXiv preprint arXiv:2107.09133 2021.
Kiwon Lee, Andrew N Cheng, Courtney Paquette, and Elliot Paquette. Trajectory mini- batch momentum: Batch size saturation and convergence in high dimensions. arXiv preprint arXiv:2206.01029, 2022.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218, 2020.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. Conference On Learning Theory, pp. 2-47. PMLR, 2018.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss?-a mathematical framework. ICLR, 2022.
Chaoyue Liu and Mikhail Belkin. Accelerating SGD with momentum for werrpparameterized learn- ing. ICLR, 2020.
Kuang Liu. Train cifar 10 with pytorch.
February 2021.
Shengchao Liu, Dimitris Papailiopoulos, and Dimitris Achlioptas. Bad global minima exist and sgd can reach them. Advances in Neural Information Processing Systems, 33:8543-8552, 2020.
Stefano Sarao Mannelli and Pierfrancesco Urbani. Just momentum: Analytical study of momentum- baasd acceleration methods in paradigmatic high-dimensional non-convex problems. arXiv preprint arXiv:
102. 102.::::212211111 1755, 2021.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and llya Sutskever. Deep double descent: Where bigger models and more data hurt. ICLR, 2020.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate o1/k²). Dokl. akad. nauk Sssr, volume 269, pp. 543 543-547, 1983.
Quynh Nguyen. On connected sublevel sets in deep learning. In International conference on machine learning, pp. 4790-4799. PMLR, 2019.
Scott Pesme, Loucas Pillaud-Vivien and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218-29233 2021.
Loucas Villaud-Vivien Julien Reygner, and Nicolas Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. arXiv preprint arXiv:2206.09841, 2022.
Boris Т Polyak. Some methods of speeding up the convergence iteration methods. Ussr computational mathematics and mathematical physics, 4(5):1-17, 1964.
11

---

Under review as a conference paper at ICLR 2023
Boris T Polyak. Introduction to optimization. optimization software. Inc., Publications Division, New York, 1:32, 1987.
Jasson DM Rennie and Nathan Srebro. Fast maximum margin matrix factorization for collaborative prediction. In Proceedings of the 22nd international conference Machine learning, pp. 713- 719, 2005.
Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. arXiv preprint arXiv:2106.10165, 2021.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues the hessian in deep learning: Singu- larity and beyond. arXiv preprint arXiv:161 2016.
Levent Sagun, Utku Evei, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint urXiv:1706.0444: 2017.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics of learning in deep linear neural networks. ICLR, 2014.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. arxiv e-prints, art. arXiv preprint arXiv:1710.10345. 2017.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning, pp. 1139-1147. PMLR, 2013.
Mingwei Wei and David J Schwab. How noise affects the hessian spectrum in overparameterized neural networks. arXiv preprint arXiv: 11Xiv:1910.00195, 2019.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. Advances in neural information processing systems, 30, 2017.
Lei Wu, Chao Ma, et al. How sgd selects the global minima in over peeeparraeteri learning: A dynamical stability perspective. Advances in Neural Information Processing Systems. 31, 2018.
Lei Wu, Mingze Wang, and Weijie Su. When does sgd favor flat minima? a quantitative characteri- zation via linear stability. arXiv preprint arXiv:2207 026628, 2022.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint 1rXiv:1708.07744 2017.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffuusin theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint arXiv:2002.03-95, 2020.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2021. URL https: https::/oppenevveet wwt.iiwwwweeeeeeeeeeee htte/forummmmmm for.mmid-ummi iciYGo.
Zeke Xie, Xinrui Wang, Huishuai Zhang, Issei Sato, and Masashi Sugiyama. Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International Conference on Machine Learning, pp. 24430-24459, PMLR, 2022.
Greg Yang, Edward J Hu. Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466 2022.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. corr abs/161 1.03530. ICLR, 2017.
12

---

Under review as a conference paper at ICLR 2023
A EXPERIMENTAL DESIGN
A.1 UV MODEL
In our experiments with the vector UV model we aim to extract how the timescale of motion scales with y/ for each 7, therefore train the model over a sweep of both of these parameters. As reminder, the loss of the UV model is
1 Vn
L 2P
1/
Uz,
(1I)
where P is the dataset size and u. n-dimensional vectors. We set y; 0 for simplicity, and with label noise this becomes yi (1) ((( for i.i.d standard Gaussian distributed &,. We initialize Hy once (not online) as i.i.d. random standard Gaussian variables. For all experiments 2. We initialize with u,, 2j i.idd standard Gaussian distributed and keep this initialization constant over all our experiments reduce noise associated with the specific initialization.
choose
For each value of 3 we train with label noise SGD and momentum until |u|² + |o|² < en with & 0.1, thereby obtaining a time series for each of u(t) and vtt). We extract the timescale by fitting log (Iu/22 [0]²)
to linear function and taking the slope.
A.2 MATRIX SENSING
We also explore speedup for well understood problem: matrix sensing. The goal to find low-rank matrix X* given the measurements along random matrices A₃: y: Tr A,X" Here
X E IR [dd is matrix rank r (Soudry et al., 2017; Li et al., 2018).
Blanc et al. (2020) analyze the problem of matrix sensing using SGD with label noise and show that X* is symmetric with the hypothesis X U for some matrix U, then gradient descent with label noise corresponds not only to satisfying the constraints y; Tr A,X, but also minimizing an implicit regularizer, the Frobenius norm of U, which eventually leads to the ground truth.
In the analogous UV matrix model (with X is asymmetric matrix of low rank r). demon- strate a considerable learning speedup by adding momentum, and show that this speedup is not monotonic with increasing B3 there is value 3 at which the acceleration appears optimal. This non-monotonicity with an optimal 5 observed for both the Hessian trace and the expected test error. Assuming that this setting we also have y = 2/3, we can extract C (1 - -/-))///2/// so P-1/3 1/3, which compares favorably to the upper bound we may extract from Appendix of 0-12
In the experiments with matrix sensing we aim to demonstrate the benefit of momentum in a popular setting. Matrix sensing corresponds the following problem: Given target matrix X E Rddd of low rank < d and measurements {y: Tr x how can we reconstruct X*? One way to solve this problem to write our guess X = UV the product two other matrices, and stochastic gradient descent them, hoping that the implicit regularization induced by this parametrization and the learning algorithm will converge to a good low rank X.
A.2.1 EXPERIMENTAL DETAILS
In our experiments we study the d = 100, 5, P- bn 2500 case. We draw (Az)i ~ NN(,, 1) as standard Gaussians and choose X* by drawing first (Xo)aj NN(0 1) and then performing SVD and projecting onto the top singular values by zeroing out the smaller singular values in the diagonal matrix. We intitalize / = V = la. We perform SGD with momentum on the time dependent loss (with label noise depending time)
L(t) = dP
(e (1(t+ T(AAU)))
(12)
where e2 = 0.1,,((!) ~ -N(0, 1). We choose 7 0.1 for all of our experiments.
13

---

Under review as a conference paper at ICLR 2023
Timescsles of the Test Loss
Tiimrccales of the Training Hessian
B 0.00339
Normaluue Monddalized Hessinu of Trining Error
0,002415
Expected Trast Firoor
0.00180
00
8 = 0.00090
ura
01'11
Epech
Epoch
Figure 3: The expected test enror, a). and Hessian of the training loss, b), in matrix sensing (with = 100, = 5, and 5rd = 2500 samples) as a function of training epoch plotted for different values of / 17 0.1. The label noise variance is 0.1. Each curve represents different value of 3. The inset shows that the orange curve crosses below the blue curve before convergence of the Hessian. Therefore, the same value of 3 optimal for both the Hessian and the expected test error increasing or decreasing 33rm from this value slows down generalization.
The hessian of the loss is defined to be the Hessian averaged over the noise. Equivalently we may just set (t) 0 when calculate the Hessian because averaging over the noise decouples the noise. Similarly when we define the expected test loss we define it as an average over all A, setting (t) 0 in order decouple the noise. Averaging over Et( and A, would simply lead to an additional term (t)) which would simply contribute constant. We remove this constant for clarity. As result, the expected test that we plot is proportional the squared Frobenius norm of the difference between the model UV and the target X",
1 (L) d IUV 1
(13)
It is also interesting note that observe epoch-wise double descent Nakkiran al. (2020) in this problem. In particular, we observe that the peak in the test error can be controlled by the momentum hyperparameter, and becomes especially pronounced for 3 1.
A.3 RESNET18 8 ON CIFAR10
We train our model in three steps: full batch without label noise until 100% test accuracy, SGD with label noise and momentum, and then final projection onto the interpolating manifold. The model we use is the ResNet18 and we train on the CIFAR10 training set.
The first step is full batch gradient descent on the full CIFARI0 training set of 50,000 samples. We train with a learning rate = 0.1 and momentum / = 0.9 and learning rate schedule with linearly increases from 0 to y) over 600 epochs, after which it undergoes a cosine learning rate schedule for 1400 more epochs stopping the first time the network reaches 100% test accuracy which happened on epoch 1119 our run. This model is saved and the same one is used for all future runs.
The loss function we use is cross cross entropy loss. Because we will choose label noise level of p = 0.2 which corresponds to a uniformly wrong label 20% of the time, during this phase of training train with the expected loss over this randomness. Notice that this loss is actually linear in the labels so taking the expectation is easy.
The second step involves starting from the initialization in step and training with different learn- ing rate and momentum parameter. In this step we choose the same level of label noise p = 0.2 but take it to be stochastic. Additionally of 512. This necessitates decreasing the learning rate because noise is greatly increased as demon- strated in the main text. In this step train for fixed 200 epochs for any learning momentum
SGD instead of gradient descent with batch size
14

---

Under review as a conference paper at ICLR 2023
Effect of Momentum on Phase 2 (n 10-3)
90
Test Accuracy
85
8 0.000000
80
B 0,900000
B = O. 0.990000
8 0,999000
8 = O. 0,999900
8 = O. 0.999990
25
100
125
150
175
200
Epoch
Figure 4: A sample of the curves with different momentum hyperparameters with 1 0.001 with Resnet on CIRARIO. The speed of increase in accuracy is non-monoe non-monotonic in 3: the best performance obtained by an intermediate value of B, consistent with our predictions.
combination. We only compare runs with the same learning rate value. We show an example of the test accuracy as we train in phase 2 for 77 0.001 in figure 4.
Notice that for both too-small and too-large momentum values that the convergence to a good test accuracy value slower. The initial transient with the decreased test accuracy happens as we start on the valley and adding noise coupled with momentum causes the weights to approach their equi- librium distribution about the valley. For wider distributions the network farther from the optimal point on the valley. As training proceeds we see that the test accuracy actually increases over the baseline as the hessian decreases and the generalization capacity of the network increases. This happens most quickly for the momentum which matches scaling law.
The final step is a projection onto the zero loss manifold. This step is necessary because the total width of the distribution around the zero loss manifold scales with T-3. and this will distort the results systematically at larger momentum, making them look worse than they are. We perform this projection to correct for this effect and put all momenta on an equal footing. This projection is done by training the full batch with Adam and learning rate of 0.0001 in order to accelerate training. We do not expect any significant systematic distortion by using Adam for the projection instead of gradient descent.
To determine the optimal value of B3 we sweep several values of /3 and observe the test accuracy after the previously described procedure. To get one with the highest test accuracy we fit the accuracy A(8) to
precise value of instead of simply selecting the
[a1(3 3.)
if 3 < 3, if 8,
A(3) aynaz
(14)
[223 B.)
to the parameters amaz: 01, 02 and 3., thereby extracting 3, for each 17.
A.4 MLP ON FASHIONMIIST
We perform similar experiment as in section A.3 but with different model and dataset: 6-layer MLP trained on FashionMNIIST (Xiao et al., 2017) as our dataset. We perform the experiment with layer MLP with Relu activation after the first 4 layers, tanh activation after the fifth layer, and
15

---

Under review as a conference paper at ICLR 2023
Optimal Beta Scaling with MLP on MNsSi
B. exp [-0.55 2.30 of 28
10-3
10-3
10-1
7
Figure 5: Scaling analysis for a 6-layer MLP on the Fashion MNIST dataset. We see that the exponent, though not very well determined, consistant with our theoretical prediction of 1.
linear mapping to logits. We use cross entropy loss with label noise p = 0.2 as before, and always start training from a reference point initialized on the zero loss manifold. This point was obtained by gradient descent on the expected loss from random initialized point with learning rate = 0.002 and without momentum.
After this we sweep y E [10-5 10-*) and 8 E [.95, 1 2n] and train for 600 epochs with label noise and 400 without label noise. This allows us obtain a test accuracy as function of ry and obtain the best momentum hyperparameter, 8r, a function of n as in A.3.
and therefore We extract the scaling exponent by doing linear fit between log(1 - B.) and log 77. The scaling analysis is shown in Fig. 5, and shows that the exponent is consistent with our theory.
B REVIEW OF RELEVANT RESULTS FROM KATZENBERGER (1991)
In this Appendix we summarize the relevant conditions and theorems from Katzenberger (1991) that use prove result the limiting drift-diffusion. We refer Katzenberger (1991) for part of the definitions and conditions cited throughout the below. In what follows, ( (Fr h2oo P) will denote a filtered probability space, Zn an Rr-valued cadlag semiimaamaall with Zn (0) = 0, A, a real-valued cadlag [F7))-ddpped nondecreasing process with A, (0) = 0, and
σ: 1 + R a continuous function, where U a neighborhood of (0) x I as defined in the main
text. Also, X,, is an IR. valued cadlag {F.") seemimartina satisfying
=x,()=xx(0)) (0)+ + (' on (xn)dd + 1 F(Xx)dAn
(15)
for all / 5 (K) and all compact K C U, where
(n( in 2 0fxxx KorXa # K3
(16)
be the stopping time of (t) leave K., the interior of K. For cadlag real-valued seimimartingales X, Y let [X, (t) be defined as the limit of sums
(X(t)1)) -(0))
(17)
16

---

Under review as a conference paper at ICLR 2023
where 0 = to <... t, = X is an RP- semimartingale, we write
and the limit is in probability as the mesh size goes zero. If
D
[X] = Σ [X
(18)
i=1
Condition B.1. For every T O and compact K C U
inf (K)))) (An(L+e) - An(()) + 00
(19)
as 72 + OC where the infimum of the empty set taken to be oc.
Condition B.2. For every compact K C U (Z²,.(K), satisfies the following: For 2 1 let Zn be a {F(] -semimarriie with sample paths in Dg= 20). Assume that for some 5 > allowing 5 = 0o and every 2 1, there exist stopping times [T 2 }} and decomposition of Zn
(z)) into a local martingale Mn, plus finite variation process Fn such that P(r S S 1/k,
[MA](t (TT)) Tearg (Fa) 2 1 is uniformly integrable for every 2 0 and k 2 1 and
lim lim sup P 7-0
sup LoSiST
(T++,(F,)) T((F,))
(20)
n-000
for every > and T >>0 0. Also as do and for any T
sup (K)
Az) +
(21)
Condition B.3. The process
(t) = E AZn
(22)
DC
exists, is Semmmmsmmiaattigalee and for every compact K C U, the sequence (ZA (Z)m(K)} is relatively compact and satisfies Condition 4.1 in Katzenberger (1991).
Theorem B.4 (Theorem 7.3 in Katzenberger (1991)). Assume that I is C and for every y E I, the
matrix Fy has D M eigenvalues in ((1). Assume (B. (B.22) and (B.3) hold, +b is C2 for F T.C2) and X, (0) = X()) U. Let
(t(t)= (1) +
(23)
and, for a compact K C U, let
(n(K) inf(t 2 (1- $ K or (I) $ KJ.
(24)
Then for every compact K C. U, the sequence {(Y.... (K) ZK pa (K)) } is relatively compact
in DRaDer [0, oo) x [0, oo] (see Katzenberger (1991) for details about the topology). If (Y, Z, p) a limit of this sequence then (Y, Z) is a continuous semimartingale, (t)) E T for every t almost surely, 2 inflt 2 (t) & K} almost surely, and
Y()) (()) + /
tAp (o(())(Y)dZZ E// iski
(M) 21. (25)
B. APPLYING THEOREM B.4
Recall the equations of motion of stochastic gradient descent
(m) ) BTT( (x) (L (n)) + (w
(n) wh+1 w + 1p k+11
(n)
(26)
1++11
RD, and IR' is noise vector drawn i.i.d. at every timestep k with zero mean and unit variance. We now show that this equation satisfies all the properties required by Theorem B.4.
where (w) R Dxr is the noise function evaluated
The manifold I is the fixed point manifold of (non-stochastic) gradient descent. {0} x r is C manifold because r is C² which follows from assumption 3.1. The flow F(w, T)
17

---

Under review as a conference paper at ICLR 2023
(n(BT VL(w)), 3 - VL(w)). As shown in Appendix E, dF has exactly M zero eigenvalues Tnk. F inherits the differentiable and locally Lipschitz properties from VL, and therefore satisfies the conditions of B.4.
Next, notice that the noise function R + D2 is continuous because or is.
Now define An and Zn (as in the main text) that X, reproduces the dynamies in equation (26), except with a new time parameter = kc,
(,
A.(())
(1) 6n Σ &
(27)
2
k=1
So that, with these choices, Eq. (15) precisely corresponds to (26), up to the rescaling / -= kee
Now we show that An, 7n satisfy the conditions of B.4. Clearly A, = Zn( by definition. Then
+
t + -
An + + An,t))
2
20
(28)
2
2
when we take en - 0, thus recoverying condition B.I.
By definition Zn is a martingale. Notice also by the definition of Zn. because ÉN is i.i.d. with variance 1, that (t) has variance A,, (t) Σ t which is uniformly bounded and hence Z, (t) is uniformly integrable for stopping times T = 24: > Also note that goes zero in probability as becomes small because &A has bounded variance, and (t, is zero otherwise. This shows that Zn satisfies condition B.2. Because Zn and A, are discontinuous at the same time we automatically satisfy condition B.3 as pointed out by Katzenberger.
|enEs] which
This shows that we satisfy the conditions of Theorem B.4, therefore we have the following Lemma B.5, The SGD equations formulated as in (26) satisfy all the conditions of Theorem B.4.
B.2 EQUIVALENCE BETWEEN EQ. (1) AND EQ. (4)
Eq. (4) is the rewriting of eq. (1) in the form presented in Katzenherger (1991), which theory is based. We show the equivalence below, and include the definitions of A and Z to keep this section self-contained. The statement that Eq. (1), i.c.
Th+ = BT
V ((wx) + ea (w)))E))
WK++ wk + 17T44
(29)
can be rewritten in the form of eq. (4)
(x = Xx,(0) f o(X.)dZn / F(X.)dAn
(30)
Ar((1)
Aa(1) = x
(n) = 6n Σ &
(31)
where X, (t) = (Tk, wx). with the correspondence t kēz where k denotes the SGD time step.
Also, o(X) (a, na)) and F(X) ((3 - 1) VL(w), n(3T VLw)))).
Before we begin the proof, we will need the following fact (see Sec. 2 of Katzenberger (1991));
[. fdg [ fdg + Σ f(-)Ag(()).
(32)
SErSE
where I and are càdlåg functions, in particular they are right-continuous and have left limits
everywhere. The integral above is done with respect to the measure dg, the differential of a function The sum is taken over all r E (s, t] where is discontinuous and the notation Ag(r) = g(r) lim g(u) indicates the left limit of
g(r-) indicates the discontinuity g at r, where g(r-)
g r. Finally 9 denotes the continuous part of
gf = g(t)
Σ Ag(s), O<cest
t t2 0.
(33)
18

---

Under review as a conference paper at ICLR 2023
We now show by induction that Xt kc,) solves the first equation above. Note that
00 Σ
20
dA,(1)= =
8(1 h),,,,
dZn(1) = en Σ
Edg(1 - k,,))
(34)
For brevity we will drop the subscript n and let ke = t. Consider
X(t+ 2) ! 1+c²
((X(s)) + foee F(X(S))
I+e"
(35)
I+E?
l+t?
=x(t)+ + ((x()))Z(())+
F(X(s))dA(s)
(36)
= =x(1) +
Σ 1Ca5t+22
6(X(s-) )+)))) + Σ ((X(A)))AA())
(37)
KISH+
where in the last step we use eq. (32) as A, Z, ā and F are cådlág, and that dA dze 0. The
sums are taken over all 8 E (ke², (k+1) where (s) and A(s) are discontinuous. The only point of discontinuity of A and Z in this interval is = t e², so
X(t + 22) = X(t) + cā( e) + ((
(38)
because the jumps of Z and A at t + e² are E&x+1 and 1, respectively. Now must determine the left limit of X at t + c Notice that for < 88<
X(t + = = X(t) + Σ
+ E F(X(s-))AA(s)
(39)
KEStH+
t<a5t+62
= X(t)
(40)
because A and Z are continuous on (1,1+69)) Hence the left limit of ((((+)))- X(1). Putting these equations together we find that
X((k + 1)e²) X(t + e²) = X(k?) + co + (K((k)))
(41)
which, using the definition of X(ke²) = (Fi, wK), is eq. (29), thus proving equivalence.
C EXPLICIT EXPRESSION OF LIMITING DIFFUSION IN MOMENTUM SGD
In this section we provide the proof of Theorem 3.4. Recall that satisfies
((x + F())) = ((x).
(42)
where x = (m, w), and F given in (132). To obtain the explicit expression of the limiting drift- diffusion, according to Theorem (B.4), and keeping into account Assumption (3.2), we need to determine up to its second derivatives. To this aim, we shall expand Eq. (42) up to second order in the series expansion in F. In components, this reads:
n0 (T' -g') + 10 (-Cn7 g¹) + 7 Ow, 8' - g') - -gt)
1 z n00 On ( -g')(-CCn' - n0 Our 2 -
1
(T* -
(43)
2
+ 1 0m, Rr (-C - g¹) (-C
Cnj x² - -) )= =0
subject to the boundary condition
क wwl w.
(44)
Here, 9¹ dZ In Eq. (43), we already substituted
Ory", . In what follows, we shall find &b to leading order in 17 as + O. We will solve the above problem by performing series expansion in T around a point E and π = 0 up to second order:
4 (T, a + 5w) = Poo + 1y 8w + Fin=" :
D D' ow + $4 n ''0'' + x ' +.... (45)
2
2
19

---

Under review as a conference paper at ICLR 2023
For example, we have
SC
%
T
T
(46)
Ow
More precisely, we regard this expansion in powers of ow and æ: Poo is zeroth order, & P10. Poi are first order, and the remaining terms are second order. We will occasionally write explicitly the
a6 index of D e.g. $ k. Bu--Baa It will be useful introduce the longitudinal projector onto r,
Pr.(e) (w)) DRDD + T, defined such that for any vector E T: PLU v. The transverse projector
is then Pr = Id - PL. We will also decompose various tensors using these projectors, e.g.
V TiLL &Hur + D 11TL + $
IITT
(47)
with $5 11T = $ PirPp Note that the Hessian = Vg E BxD satisfies HH = Pr. where H denotes the pseudoinverse.
At zeroth order, obviously have (o0(w) w.
Lemma C.I. The first order terms in the series expansion (45) are given by
PIIT 1oo = =0, Poir. = Pr.. Pior. = =C-11t--PP.
(48)
Proof. Suppose (ss) E R is curve lying on r. Then due to the boundary condition (44), 0. क 0, ú)) 0,0' O,w² This means that Poir. Pt.. Now from (43),
1 (T²' 0,g' bu²) +
0,g' 0.
(49)
This condition should hold for any and ow, therefore we arrive at
10 11
(50)
and
(1/ 0o)'" =0
(51)
Decomposing into longitudinal and transverse components, and noting that the Hessian satisfies P,H = 0, the above equation becomes
T + Por = 0,
(52)
which together with (50) and the above discussion gives
& Peir-
4 10r 0,
P PL. P C-1111~2P
(53)
which concludes the first order analysis.
We now need Lemma, which requires Definition (3.3) in the main text. We report it here for convenience:
Definition C.2. For symmetric matrix H E RD RD, and WH [ E R RDD E HHH = Ht HΣ, o). we define the operator Lu WH - WH with Sn = (H, + C 1 [[S HJ, H). with [S, H) = SH HS.
Lemma C.3. The inverse of the operator Σu is unique.
Proof. Let us go to a basis where H diagonal, with eigenvalues A, In components, the equation LnS M reads
(A: + S3)S.s + 1 2 1 S,) + 1C''11 2 =
(54)
which has a unique solution, with
1 2- (211221)))) -2)8) -1
Syj -
(1, (d
My..
(55)
2
20

---

Under review as a conference paper at ICLR 2023
Lemma C.4.
The second order terms in the series expansion (45) given by
Lrikk 02LL
= - -(H yo 0 H P
-P H (H Lisk 02LT
qiikk 02TT = =0 (1)minn/000122)))
(56)
Blikk
- - of H PE
1172 -
p.. of Hi (Hi)m
(57)
IILL
1ITT = C- 1) H (M
0
1) 2- 31 H CH M lit
(58)
2
Liik
2-2 1,- 2-21 (1 21 (11
pa (1I اا 0. n p
(59)
200L
2
0- 212 2-21 | 0 H (11 ni of 0f 11 (1
(60)
20TL
Liik
20TT
= -
212 21- 27 H 0 H H )
(61)
where
(M )m P 0j Hns Pri
(62)
Proof. Consider path (s)) lying on I. From Prl H we have
dPr
H (Id - Pr) dH
ds = Pt. dH
(63)
ds
ds
and thus, using HH! Pr, we find
dPI
Pr = PL dH
dPI
dH
H
Pr
H
PL.
(64)
ds
ds
ds
ds
where the second equation is obtained from the first by taking the transpose. Putting the last two relations together, we find
dP'
dPr
dPr
dPr Pr = - -Ht dHl
dH ds
Pr
Pr. - Pt.
Ht
(65)
ds ds
ds
ds
ds
From (48) we can then write
of o 01 = -H
1n pr. - P/. 1 11
(66)
where 0. = =P This leads to
JAsk
02LL = Phi of = ( Hin pp
(67)
Liik
02LT
of - H (Ht) nk
(68)
Also note that 122.1 = 12r so the only component still to be determined in Pug is Puzrr-
The next step is to expand Eq. (43) to second order:
n'm (n(1 + C Cr 7(
) +
1
+ (n ck - 17² Hæ, 22 u HμI - Ci1² 1 17( 21 Haj (1-Cn)")(b
Swl (1 Hu + 2 H1, H&s H + 1b Haj His
+ + 210)0 Hu + 11 Hys) = 0,
(69)
where 1(6) = (A + A denotes the symmetric part. Neglecting various terms that are sublead- ing small 17, gives
n' (n - bit + 1 4xb
+ (n - 212 w/t - +11 Ha))
(70)
-8w* " (100 H1: + th Hu + (1 + 11) Hu + f H1yhis) = 0,
21

---

Under review as a conference paper at ICLR 2023
The first line immediately gives
20 =
(1) 11 2
-
(71)
Taking the second line of (70) and projecting onto the longitudinal part the index k gives
n - C112 =0
(72)
B3 4i =
(73)
and using (67).
& bitk
C-
back
Likk
- -n 11tn p
(74)
IILL
122L
d jisk
-C-
02T1. = -)- p.. 0f 11 (
(75)
1171.
Projecting the second line of (70) on the transverse part of the index i we have
1 O2ET
Cr
IILT
2.44 20T Hu =0
(76)
1 02TT
11TT
& 20TT Hu O
(77)
Using (71) and keeping (68) into account, Eq. (76) becomes, neglecting subleading terms in
- -np 0f HHHHH) (Hiym Cip c- 'n ti liTE Haa - 1 c- 10-1-700 til - BuT Hu =0
2
(78)
Using (75) in (78),
-npj 0. H (Ht) Cn
- 2 2- 27 pip HPT ( Ha. - c- 1--1,--707 VILT H 0,
1
(79)
2
and simplifying,
- of H H C11 1
1 2- pb 0. 1 pp* - 1- Hu = O.
(80)
+2
which determines IILT in close form. Indeed, the above has the form, in matrix notation
C PILLT M,
(81)
and can be immediately inverted to solve for PILLT. The only two undetermined components now are PIrr and Po2rT. One condition is obtained from (77) which gives, keeping (71) into account,
n i 02TT
Cr1" 2.ik
C 17' p
Hu = 0,
(82)
and thus
J.sk 02r7
jak -
+ - (. H
(83)
Further taking symmetric and antisymmetric part in ik of the above gives
J.sk 12TT
-
(iikk) 11T
c-
(mi) H1TT H
c c-
a (nk))
H1TT H
(84)
2
i 0 0=Cn LITT + --
(mmi)
11TT Hkm
1
-
L1TT H
(85)
2
The other condition comes from the third line of (70) which, using (48) and (71), and neglecting subleading terms in v), gives
j.sk Ha + 1 T
02 HAs + $/ Hü + Di1 Has + "n 0 H = 0
(86)
Projecting the A and indices on the longitudinal part, gives p/ H.Pl 0, which is an identity that can be checked from (63). Projecting & on the longitudinal part and on the transverse part gives
j.kk Hu + Diik Hu -- ppj ak H., ppl 0, G2Tz
(87)
22

---

Under review as a conference paper at ICLR 2023
which is implied by (75), indeed
(172 + 11ITL) Hu 1) 1 2ork Hi. = =- "Pp 0
(88)
(H = Pp H PT.
therefore (87) does not give a new condition. The only new condition comes from projecting the and indices of (86) on the transverse direction, giving
02 Hy: + H2 + 2
+ + +C- +
og Pr =0. (89)
Plugging in (83),
C1 Hy + 1) (p (1mm) H.. H 7 H4 + C 11 1177 Hi, Hki
1117
(90)
+ 117T His + 1 1irTT HA: + 1/ PP of Hn; Py' =.
Neglecting subleading terms in 1y this can be rewritten as 2C 20 ((6)) where we recall that the last term is symmetric in l: and L. In matrix notation this reads
11TT Hnn Hu + 4 117T Hu T Has ~0 Pr Pai 0.
(91)
C-1n)t~>((t H( HIT P(rT) H HIT + + 1ITT H + -- 0,
(92)
where
(M )m pki of Hns
(93)
is the longitudinal derivative of the Hessian, projected on the transverse directions. Decomposing o Hirr into symmetric and antisymmetric parts hirr = St)) + A Eq. (92) reads (suppressing the index j)
2C- + HS + SH + ИA - AH + C C M =0
(94)
The first term is subleading in 7) therefore we have
HS + SH HA AH + 1y 0. This equation together with (85), which in matrix form reads
(95)
HS) + 20 0,
(96)
determine S and A, and thus 4 LITT. Note that M has only a longitudinal component in the index j. therefore the transverse parts of S and A vanish, i.e.
J.ik PPJ
11TT = 0.
(97)
Eq. (97) is natural as the slow degrees of freedom are the longitudinal coordinates.
Note that these two equations admit unique solution. Indeed, solving (96) for A gives
A C --
(98)
2
Plugging in (95), we find
LnS -C-Ln---M.
(99)
where inS = (H, + 2, 23 ([S, H1, H) is introduced in definition C.2. By Lemma C.3,
(99) admits unique solution.
Then
4 iirr = 2 M ] c- L (z Mb)
(100)
2 The other components are (see eqs. (7))
qi.ik 022.1 - -HHI H'" pp.
(101)
j.ik
Pt of H (H
(102)
jiik
= - -C- ( 1 H pp
(103)
IILL
Lidk
-c- 11 |n (m
(104)
FITI.
j.lk
C
C n
Hun =0 (17mi((01-211)
(105)
02TT
C-
(ij) 11
- C-
1
a 02
--,,,-~0 (ij)
(106)
2
23

---

Under review as a conference paper at ICLR 2023
The only contributing term at leading order in 1 to the limiting diffusion equation is (106). Splitting it into longitudinal and transverse components, we find:
jikk
IILL
(ssk)
1 2- 22 ( of H'm PPt (HiY'af H Pt
(107)
jilk 20TL
-² -
(((k))
((skk)
1 1/
2-27 ( ( + of H (Ht)
(108)
11TL 2
and, using (100) and (93) we have, in matrix notation,
4 20TT = C- 17 (
HTT + + (PurT)) = 2 27 L (Pr HPr).
(109)
2
To write things more compactly, the following Lemma will be useful: Lemma C.5. For any transverse symmetric matrix To
20T7 ( 1 ---00-21
2- 2 M (CTTT
(110)
2
Proof. From (61), P2orr is proportional to the symmetric part of PrrTT which was denoted by S below Eq. (93) and satisfies Eq. (99). Therefore D2UTT also satisfies Eq. (99), up to an overall factor:
CH 20TT
2 (2-21
(111)
2
Then, for any T:
(Cn BOTT [T]
22 M M 2) 2
(112)
Moreover,
(CH [7] D 2
(113)
Since the two above equations hold for any T and En is invertible, this implies
20TT [T] 0- 2 12
(114)
which is the statement of the lemma.
To leading order in 17 we then have, for a symmetric matrix V, [D[[] thus
and
1 2-
20 1
" P)
(115)
27 PLO (VL) Vrr),
20
where VLL = PLVPL, VIL = P,VPL, and Vrr PrVP, are transverse and longitudinal projections of V.
We also have, using (48), and to leading order in 7) = @ = 1oodZ + = (C~"y>"" ProdW,
(116)
where dW a Wiener process. Applying Theorem B.4, and keeping into account that, to leading order in dt, d Zo] 8 dt, find
dY
202 17
(117)
27 p (VL)(H"))TT]dt c2
2 11 2-21 p' (VL'CrpErrll
C2 "
where Σ= To
24

---

Under review as a conference paper at ICLR 2023
For y 2,2 LH reduces to the Lyapunov operator at leading order in 17, i.e. LHS {H,S), For Y 2' from (55), it is easy
that the role of the divergent term proportional n71-27, when acting - on S. is to set to zero the off off-diagonal entries of Sij at O( i.e.
S = -
Sizj -, 0,
(118)
Using Lemma B.5, we finally conclude the following Corollary, which is the formal version of Theorem 3.4 in the main text:
Corollary C.6. Consider the stochastic process defined in Eq. (26) parametrized by ens with ini- tial conditions (To, wa) E U, under assumptions 3.1 and 3.2. Fix compact K C U. Then the conclusions of Theorem B.4 apply, and (t) satisfies the limiting diffusion equation
dY -1 + n) PLodW - 22= 2 (V 2 PLLVLLLLLLVVLLVLLL
(119)
10 27 PLO (L)
where (t) is Wiener process.
Let us now see the special case of label noise. In this case E cll, so that 2 is only transverse. Moreover, using (55),
1
f H
Pr:
(120)
2
and
dy = 4C2
402 102"²² 12-
2- 27pp PL VTr(co
(121)
This proves Corollary 3.5 in the main text.
D EFFECTIVE DRIFT IN UV MODEL
We start with mean equare loss
P -11f(=)) - 1
(122)
2P
with the data covariance matrix as defined in the text. The trace the Hessian on the zero loss manifold (L(w*) 0) is given explicitly by
trll 1 (V
(123)
Taking gradients of this and plugging into Corollary 3.5, repeated in Eq. (121) leads to an explicit expression for the drift-diffusion along the manifold
dY 2-2-21 27 2 dt, $
1r 0
m ). - PrSSP.
(124)
4PC²
The simplification we cite in the main text in Sec. 3.4 is due to the fact that for input and output dimension d m = 1, we have that Σ = TrΣ = 12, and Š is proportional to the identity.
For matrix sensing, order to compute the trace of the Hessian, we use (12) with & and with slightly different notation (V* instead of V). The loss is then
P L =
W ))(A,UVT))
(125)
Pd
1=1
25

---

Under review as a conference paper at ICLR 2023
and define the data covariance matrices
Σ 1 Σ
A.A Rdxd,
f²=
P Σ
1
A A E
(126)
Then the trace of the Hessian is
2 TT ($PUUT + +
TrH
(127)
d
We find the noise function is
pi
2
(A)( E R
(128)
Pd
where S(A) = Tr(UV A). Since the Hessian on the zero loss manifold is Hyv
Pd Ei VA,(AA)V we see that aa = (2/Pd)H. Therefore, we get
1p 27 46 dY
Σ 4C P SLYd =
0 E1 ). P
(129)
To get a crude estimate of the coefficient in the main text, we approximate the top eigenvalue of S with !T With this, we get for
2 CPP dp TTAT Capa d(d
2
2
2
CddP (dly)
(130)
Here have denoted by azy arbitrary element of the data matrices A, with brackets signifying an average over the distribution these elements. Assuming the fast initial fast remains the same, and r- = (C/2)n²², get
2c dP (a?))
(131)
For the values used in experiments, this gives C 0.12 P-1/3,
LINEARIZATION ANALYSIS MOMENTUM GRADIENT DESCENT IN THE E SCALING LIMIT
Here we elaborate on the discussion in Sec. 3.3, providing derivations various results. We take straightforward linearization of the deterministic (noise-free) gradient descent with momentum. Working in the extended phase space x = (rr, w), the dynamical updates are of the form
1)Tx 1( VL(w))) -V(W))( Ft+ 2z F(x), F(x.) 1)(TTT - VL(wx)) ((())))
(132)
The fixed point of the evolution x² (0, w²) will have the momentum variable 0, and the coordinate satisfying VL(w") = O. Linearizing the update (132) around this point
80 J(x J(x")ooe. J(x') ( 3 -V2L(w*)
11 n I(w²)
(133)
26

---

Under review as a conference paper at ICLR 2023
Note (2L(w") is the Hessian of the loss function at the fixed point. The spectrum of the Jacobian can be written in terms of the cigenvalues of the Hessian d This is accomplished by using a straightforward ansatz for the (unnormalized) eigenvectors of the Jacobian k (4'4', q²), where of are eigenvectors of the Hessian with eigenvalue X₃- Solving the resulting coupled eigenvalue equations for eigenvalue n²:
1 1 n7 nBp n², -A, + +1 =''
(134)
For fixed A there will be two solutions given by
1 nd - -48). i= L....
11 2
(135)
#4 2, (8- (8-1+ m + v ( -1+ m2;) 48
(136)
For the set of zero modes X; = 0, we get the following modes: K+ = 1, corresponding to motion only along w, with eigenvector k² (0, q'). In addition, there is mixed longitudinal mode which includes component of along the zero manifold k (u-q', 9²), and has eigenvalue x_ 3.
On the zero loss manifold, we can assume the Hessian positive semi-definite, and that the positive eigenvalues satisfy
0 < S A C2.
(137)
for constants C1, C2 independent of y) 8.
We now analyze the spectrum of the Jacobian eigenvalue at time, and then these results to informally control the relaxation rate off-manifold perturbations. It is useful first to consider the conditions for stability, i.e. ||*| < 1, which are stated in (138,139) below:
Case 1: If n/4 < (1- v then k4 R and |x]]] If mp Case
iff 0 < nd < 2
(138)
? then k4 C and | VB B < 1.
(139)
Proof of Case The condition nı,; < (1 - vB)² implies A > 43, where = 1 3 n/ The condition for stability then requires -1 < < +1. We satisfy both sides of this inequality:
If + 3 π/ = A> 0, then |=_|| < 1,and + > so we simply require x4 i.e.
Ki < +1
(140)
1
2 (A / -3 43) <1
(141)
+ 43
(142)
48 4
(143)
A 1=1+3- +
= n > 0
(144)
If + 8 - 1 < 0, then |x||| < and kL_ < 1, so only require k² -1
1 < kt
(145)
1 / (-|A| / A2 4.)
(146)
A - 2 < - / 12 44
(147)
| > V 42 43
(148)
4 4/4 + 4> 1² 48
(149)
A = -1= 3 < <1
= n\, 2(1
(150)
27

---

Under review as a conference paper at ICLR 2023
Proof of Case 2: the condition n,,, > (1 -1-v /B) implies A? < 43, where A (1+ 3 -),). The corresponding eigenvalues of the Jacobian can be written N (1/2) + VA2 43) (1 45 ). Computing the absolute value then gives | 2 (1/ ( 3
A
These results show us when the GD+momentum is stable. Next, assuming stability, we want to estimate the rate of convergence the fixed point. More precisely, we would like determine the fastest mode as well as the slowest mode. To this end, we define two quantities
P1 = max [ = I..,, D, | < 1}.
(151)
min( i i= 1,.... D, m4d < 1}.
(152)
Using the explicit scaling for momentum, and in the limit of small learning rate, we prove the following for PI: Lemma E.1. Let = Ciń² and sufficiently small:
For < 1/2, the condition for Case (138) holds,
P1 (a//C)n
(153)
P2 Cnj"
(154)
For y > 1/2, the condition for Case (139), and
P1 v [ 1 (C/2)m²
(155)
P2 = Cn)"
(156)
Proof: For small n, we find that
/I 1I- VI-VI-Cai) 7 (1- (1-V---------))
~ 0212r-1/44
(157)
When 7 > 1/2, this expression tends to zero as + 0. Therefore, for sufficiently small n, the condition for Case (139) will be satisfied and |x') /B for all implies P1 = /3 and = 3. The scaling behavior in the Lemma follows by substitution.
Since B vB, this
For y < 1/2, (157) diverges as 17 + 0, which means Case (138) will obtain for all A₁. Next, since for small 1 + 8 17,; > 1 + 8 1702 2 C1²7 7J22 > 0, since C and are order one constants, and 17 + 0. In this case, the largest contribution from the nonzero eigenvalues X will come from <4.. In particular, we find
1
|xt,1 S 1+ 1+3- n11 v(1 (1+8-10))2-4
- -43)
(158)
2
1
- 2
2 Cni² - me11 + v C29p2 22(2 - Cn17))n1 + +
(159)
For y < 1/2, we have the hierarchy 17² > 2 > 1l n > 1/2. This allows us to simplify the upper bound
P1 max
k <
2 O(n).
(160)
Next, we find a lower bound for P2. This will be controlled by Ri_ We may use then that
1 (1+3- 1102 v 121)² -43
P2 min (n=
(161)
Cr²² 792 /09pp2 2(2 C1p)) m91 + np297 = CCn" + 0(n'->)
(162)
28

---

Under review as a conference paper at ICLR 2023
Finally, note that since n7' 17 < 17,,, we have that 1 - ///1 Crt B, indeed P2 P1-
Equipped with the upper and lower bounds on the spectrum leads naturally to bounds on the relax- ation rate. For the purely decaying modes at Y 1/2, use P2 I 10x - pí. óx; represents the projection of the fluctuations dæ, onto the transverse and mixed longitudinal modes. After applying Eqs. (153) and (154), arrive at the result quoted in the main text in Sec. 3.3. For y 1/2, the modes are oscillatory. However, the eigenvalues within the unit circle have norm either vB. B, as reflected in Eqs. (155) and (156). This implies that we can estimate the decay rate of the envelope of the transverse and mixed longitudinal modes in this regime, thereby arriving at second expression quoted in the main text in Sec. 3.3.
F LINEARIZED SGD AND ORNSTEIN-UHLENBECK PROCESS ON THE VALLEY
In this appendix, we provide derivation of some of the statements quoted in Sec. 1.1. To get there, we start with the basic model for momentum SGD (1) but lincarize around point on the valley we E where L(wo)= VL(we) Let uk we + dukk and define the Hessian (wo)) = L(wo) Then
Th+1 = 3m H (wo))wy + ea (wo)
SWA-1 = Swn. + 1/7411
(163)
Consider the projector along transverse nonzero eigemode of (wo)). PX. and define PX x X, and PT π II. Let P o, and PT H AP Let oo
A. Then
lk = 311 AX + co(ww)
(164)
Xx+1 Xx 1
(165)
=x, + n1TT v1XX + (wo
(166)
This is a simple OU process, and we can easily compute the second moments. Define the second moments
Ci2(k) (x) (ni(k) (XxX,), C22(k)
(I2 [T))
(167)
We find by taking the equations above, squaring them, then averaging over the noise,
Oza(k: 1 (BAC)))))))))))))+++ (PCn))
(168)
C12(k +1 n1 Cze(k) + 8(1 (11-21))Ci2(4)) X(1 nAaCin (k) + n
(169)
Chi(+ ( 3 + 2113(1 (k)) + (1-1111)0 (k)+
(170)
Next, assuming stationary distribution implies C(k+1) = -(k), which then allows to solve for equilibrium variance, and extract the main quantity of interest, which is the variance of the weights. We find then
22
C1 (1 -1- 1--)) B) n))
(171)
In the limit of small 7) we extract the scaling behavior quoted in Sec.(1.1).
We can see how the mixing timescale T1, discussed in Sec. 1.1, arises from this linearized analysis. By taking the expectation value of the OU process, the noise will vanish and we find that the average values follow the linearized GD dynamics analyzed in E. Thus, from this linearized GD analysis we can extract the characteristic timescale for the OU process to approach its mean value.
29


Reviews:
[{"review_id": "WjtZbhwqK2", "reviewer": "Reviewer_e9QR", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies the convergence of momentum SGD at the vicinity of global minima in the over-parameterized setting. With some approximations which only hold in this vicinity, the scaling analysis predicts an exponent $\\gamma=2/3$ as the scaling of the momentum parameter that maximally accelerates the training.", "summary_of_the_review": "The the paper is not written clearly, and it took much effort for reading.\n\nSeveral key claims are not well supported, with the proof or justification missing. Without them, I am susbicious of the correctness of the results.\n\nThe analysis is conducted within the vicinity of global minima, outside of which many assumptions or treatments do not hold. However, in practice most of the optimization happens outside of this vicinity. Hence, I doubt whether the results of the paper is applicable in practice. \n\n", "strength_and_weaknesses": "First, I have to mention that the paper is not well written, and it takes a lot of effort for readers to understand. The authors should consider improving the clarity of the paper. I list a few points below:\n\n> There should be definition or clarification of notations before using them. For example, in the 5th line of Sec 3.1, “the parameter $\\epsilon > 0$ controls the size of the noise …”. However, there is no $\\epsilon$ mentioned before. Hence, it is hard to understand what this line is talking about. In addition, in Assumption 3.1, $C^3$ appears before the notation definition of $C^n$. Furthermore, in Sec 3.3, without defining $\\delta x$, it is hard to understand the equation $\\delta x_{t+1} = J(x^*)\\delta x_t$. \n\n> The paper does not set up the problem well. The paper has some settings in the introduction, Sec 1.1, and some other in Sec 3.1. None of the subsections established the problem setting; instead, the two sections are referring to each other, and readers have to go back and forth to get the context.\n\n> There should be a discussion of the connection between the time scaling analysis and the ordinary SGD setting. In machine learning, SGD is an iterative algorithm taking on discrete time steps with predefined learning rates. It is important to clarify what the timescale means to SGD, and why there is a parameter $\\epsilon$ controlling the size of noise and can be taken to zero. \n\nSecond, the modeling of the momentum SGD using Eq.(1) is not appropriate. As in practice, the noise of SGD arises from the estimation of gradients via samples, instead of an additive noise as in Eq.(1). The difference in how the noise is generated results in different properties of the noisy gradient. For example, in the over-parameterized case, the stochastic gradient has been found to have a vanishing noise as the algorithm converges to a global minimum, see automatic variance reduction in [1]. However, the additive noise often keeps a certain level of noise even if it is close to convergence. In addition, it has been found that the variance of stochastic gradients cannot be bounded by a constant (see [2]), while the additive noise usually can. \nHence, it is better to start the analysis from following equation (instead of Eq.(1)): $\\pi_{k+1} = \\beta \\pi_k - \\nabla l_i(w_k)$, where $l_i$ is the loss function evaluated on one training sample. \n\nThird, I could not understand why the $\\approx$ holds in Eq.(2). In principle, the momentum $\\pi_{k+1}$ should contain gradient noise, and should not be averaged over. Removing the noise is critical to the final conclusion. I hope the authors can explain this.\n\nFourth, the paper does not explain why Eq.(5) solves Eq.(1). This is a key step, and should not be omitted from the paper. \n\nFifth, the analysis is conducted assuming in the vicinity of the global minima, namely very close to the solution. Many of the assumptions do not hold outside of the vicinity, and the analysis does not apply (for example, the dominance of perpendicular displacement over the longitudinal). Hence, it is possible that the theoretical optimal power of $\\gamma=2/3$ is not optimal outside of the vicinity. \n\nOverall, as for the $\\gamma=2/3$ rule, the theoretical analysis in Sec 3.3 is not rigorous (as admitted informal by the authors), but it experimentally works well (as in the ResNet). In addition, as mentioned above in the fifth point, the theory applies to the vicinity of the global minima, and does not match the experimental setup. Given these, I would consider the $\\gamma=2/3$ rule as a good finding that is worth trying by practitioners, but I don’t think the theory supports it well. \n\n\n[1] Accelerating SGD with momentum for over-parameterized learning\n\n[2] Better Theory for SGD in the Nonconvex World\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "VM2avoE_LH", "reviewer": "Reviewer_oJpk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This work extends the analysis framework of (Blanc et al., 2020) and (Li et al., 2022) to study the dynamic of heavy ball momentum along a manifold of minimizers. The authors proposed a momentum scaling rule $\\beta = 1 - C\\eta^\\gamma$ and discovered the optimal scaling power as $\\gamma=2/3$ under the analysis framework. The authors also provide empirical verifications to the theoretic prediction on 2-layer neural networks, matrix sensing and ResNet-18.", "summary_of_the_review": "This work studies an important problem (finding better momentum values), and discovers a new optimal momentum scaling under the analysis framework. The practical implication seems to be weak unless the authors could further elaborate.", "strength_and_weaknesses": "Strength:\n- Understanding the best momentum setting is critical in neural network optimization\n- The discovery of the optimal momentum scaling, and its connection to finding flatter minima are very interesting results\n- Solid theoretic results with clear motivations\n\nWeakness:\n- In addition to identifying the best momentum setting, could the authors also provide some theoretic insight of the momentum acceleration? Such as discussing the effect of $\\beta > 0$ compared with $\\beta = 0$.\n- This work would benefit from more empirical evaluations to demonstrate the practical implication of the discovered scaling rule\n- Is the label noise of SGD only an analytic artifact? Is there any empirical reason to add such label noise?\n- The scaling rule $\\beta = 1 - C\\eta^\\gamma$ seems arbitrary. Is there any reason for such a choice?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-6eYpckavyf", "reviewer": "Reviewer_F44e", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper analyzes the SDE for SGD with momentum and label noise and tries to show that adding momentum accelerates the convergence but not hurting the generalization. This is an important challenge while training big models with numerous training datasets such as DNN. They show that there is an interplay between the speedup of momentum and the limiting diffusion generated by SGD noise. Using that they show two characteristic timescales associated with the training dynamics, and the longest timescale essentially governs the training time. Maximum acceleration is thus reached when these two timescales coincide, which leads to identifying an optimal scaling of the hyperparameters.", "summary_of_the_review": "1- My major concern with regard to the paper is its writing. Mainly the concepts and findings are not explained well formally. It could be more readable if they explain more about the notation they use and also not overloading. \n\n2- They show that for specific \\gamma= ⅔ it SGD with momentum (SGDM) can reach best performance. However, their empirical results don’t provide strong evidence. It would be helpful to run SGDM with the proposed hyperparameter over several DNN models and compare the test accuracy with SGD and other popular optimization methods. \n\n3- For the analysis of the phase when SGDM moves toward zero loss valley they assume GDM instead of SGDM. However, there is no enough justification for why this is true.\n\n4- In eq. 2 no justification for why the last approximation is true. \n\n5- In eq. 3, since we are in zero valley, isn’t \\nalba L (w_i) = 0? But the approx in the next line is not zero. \n", "strength_and_weaknesses": "Strength:\n This paper analyzes the SDE for SGD with momentum and label noise and tries to show that adding momentum accelerates the convergence but not hurting the generalization. This is an important challenge while training big models with numerous training datasets such as DNN. Their analysis provides a theoretical basis to set the momentum hyperparameter. \n\nWeakness: \nMy major concern with regard to the paper is its writing. Mainly the concepts and findings are not explained well formally. It could be more readable if they explain more about the notation they use and also not overloading. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject