Fast Mixing of Stochastic Gradient Descent with Normalization and Weight Decay
Authors: Zhiyuan Li
Tianhao Wang
Dingli Yu
Venue: NeurIPS 2022 Accept (2022)
Keywords: stochastic gradient descent
weight decay
stochastic differential equation
equilibrium
mixing

Abstract: We prove the Fast Equilibrium Conjecture proposed by Li et al., (2020), i.e., stochastic gradient descent (SGD) on a scale-invariant loss (e.g., using networks with various normalization schemes) with learning rate $\eta$ and weight decay factor $\lambda$ mixes in function space in $\mathcal{\tilde{O}}(\frac{1}{\lambda\eta})$ steps,  under two standard assumptions: (1) the noise covariance matrix is non-degenerate and (2) the minimizers of the loss form a connected, compact and analytic manifold. The analysis uses the framework of Li et al., (2021) and shows that for every $T>0$, the iterates of SGD with learning rate $\eta$ and weight decay factor $\lambda$ on the scale-invariant loss converge in distribution in $\Theta\left(\eta^{-1}\lambda^{-1}(T+\ln(\lambda/\eta))\right)$  iterations as $\eta\lambda\to 0$ while satisfying $\eta \le O(\lambda)\le O(1)$.  Moreover, the evolution of the limiting distribution can be described by a stochastic differential equation that mixes to the same equilibrium distribution for every initialization around the manifold of minimizers as $T\to\infty$. 

Fast Mixing of Stochastic Gradient Descent with Normalization and Weight Decay
Zhiyuan Li
Tianhao Wang
Dingli Yu
Princeton University
Yale University
Princeton University
zhiyuanli@cs princeton. edu
tianhao. wangdyale.edu
dingliyõcs, princeton. edu
wwngliiycss.rincttoon.com
Abstract
We prove the Fast Equilibrium Conjecture proposed by et al. [1], i.e. stochastic gradient descent (SGD) scale-invariant loss (e.g., using networks with various normalization schemes) with learning rate 1j and weight decay factor À mixes function space in ((1((1A) steps, under two standard assumptions: (1) the noise covariance matrix is non-degenerate and (2) the minimizers of the loss form a connected, compact and analytic manifold. The analysis uses the framework of Li al. [2] and shows that for every 7 O, the iterates of SGD with learning rate ? and weight decay factor on the scale-invariant loss converge distribution while satisfying 7) O()) O()). Moreover, the evolution of the limiting distribution can be described by a stochastic differential equation that mixes the same equilibrium distribution for every initialization around the manifold of minimizers as T oC.
in In(1 TT/n)/ (4nA) iterations as 7
Introduction
Generalization in modern deep learning has significantly deviated from classical learning theory due to the vast overparametrization in deep neural networks and is underlain by the implicit bias of training algorithms [3]. Instead of decreasing the training objective as fast as possible, the training algorithm and its hyperparameters are often tuned for good implicit bias, i.e., the ability to pick empirical minimizers with good generalization among various different minimizers. Sometimes good implicit bias occurs at the cost of less efficient optimization, including the usage of large learning rates (LR) [4] or small batch size [5, 6]. Thus the training objective alone is not an effective measure of the entire training progress. In other words, behind the minimization of the training objective, there potentially exists some hidden progress, and the evolution the model therein plays crucial role in the implicit bias.
The current paper aims provide better theoretical understanding of such hidden progress for neural networks equipped with normalization layers (e.g.. BatchNorm [7], LayerNorm [8]. and others [9-13]) trained by Stochastic Gradient Descent (SGD) with Weight Decay (WD), dubbed SGD+WD. For learning rate (LR) 7 and WD factor X, we formulate SGD+WD as
(k 1) = x -
where I T (k) IR the parameter after & iterations, and Lş, is the loss over the Ex-th sample with each & being sampled independently and uniformly randomly across all training data. In particular, we are interested in explaining the following phenomenon:
Longer training with SGD+ WD after LR decay improves final test accuracy of normalized networks.
We demonstrate such phenomenon in Figure where test accuracy after LR decay keeps improving when training accuracy plateaus. In an extreme case, Li et [1] empirically showed that the test
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Param Parameters nor with OOUU decay
Test accuracy after decay
0.003
Accuracy
27
0.002
48.0
Train Accuracy Accuracy
0A
100
200
302
50o
922 1000
100
600
502 900 1000
Epoch
Epoch
(a) Train and test accuracy for CIFAR-I0 training with = 0.8, 5.
(b) Test accuracy after LR decay and the to- tal norm of parameters before LR decay.
The train and test accuracy plateaus after parameter norm convergence within 100 epochs, but the generalization of SGD iterate after LR decay keeps improving. Figure la shows the train and test accuracy of scale invariant PreResNet trained by SGD+WD CIFAR-10 10 with standard data augmentation. Each red dot Figure represents the accuracy model which decays LR 10-3 the corresponding epoch. The test accuracy is evaluated until achieving training accuracy after LR decay.
Figure 1:
accuracy of ResNet can still improve after maintaining nearly full training accuracy for thousands of epochs when trained by SGD+WD on CIFARIO. Such phenomenon is also demonstrated for a standard decoder-only Transformer trained by Adam on small arithmetic datasets and named "grokking' by Power et al. [14], where the validation accuracy can increase from random guess to full accuracy long after the almost perfect fitting of the training data.
Based on theoretical derivations, Li et al. [1] further proposed the Fast Equilibrium Conjecture (Con- jecture 1.1), which informally says that for the normalized model trained by SGD+WD, such hidden progress happens in ((1/(1)) steps and the model converges to an equilibrium, and since then, further training
longer improve the final test accuracy. A recent line of works [15, 16, 2] show that gradient noise in stochastic gradient can cause higher order regularizing effect and im- prove generalization even when training loss is close to O. In particular, Li et al. [2] proposed mathematical framework for characterizing the implicit bias of SGD the time scale of O(1 D(1/n"). Under such time scale, the hidden progress of SGD is shown to be described by Stochastic Dif- ferential Equation (SDE) termed the limiting diffusion, which can then be used to rigorously prove its generalization benefit in some cases (see Section 6 in [2]).
However, the 0(1/7²) rate given by et al. [2] is not applicable to networks with normalization layers and WD, because the assumptions on the loss landscape made in [2] that minimizers of training loss connect as a manifold fail to hold for networks with normalization layers and WD, more broadly, for all scale invariant loss (see Definition 2.2) with &g regularization. Here loss is scale invariant means that L(Cxx) consequence of normalization layers. The assumption of manifold of minimizers fails because any fz-regularized scale invariant loss has no local minimizer, not to mention the manifold of minimizers. To see this, simply note that for every where scale invariant loss is well-defined, i.e., * O, reducing its norm while keeping the direction of æ strictly decreases l regularized scale invariant loss. Moreover, the loss landscape becomes unboundedly sharp around the origin. These drastic changes to loss landscape induced by normalization layers could lead to bizarre training dynamics beyond the scope of standard optimization viewpoint, e-g., deep neural networks with normalization can even be trained with an exponentially increasing LR schedule [17].
L(x) for any > 0 and parameter x * O, which is a
1.1 Our Results
In this paper, we show that for networks with normalization trained by SGD+WD the Ö(1/ (n)) rate is indeed the correct time scale for the aforementioned hidden progress and deliver partial proof to the Fast Equilibrium Conjecture proposed by that we need to rescale the SGD+WD dynamics both in time and parameter norm by leveraging the scale invariance of loss, so that the framework in [2] can again applied to achieve an SDE-based characterization for the hidden progress. Our rescalings are motivated by the analysis in [1] for the
al. [1]. The key observation here is
parameter convergence which happens in ((1(()n) steps.
Before stating the main theorem, we will first introduce some notations and restate the Fast Equilib- rium Conjecture. Let (x: be the output of scale invariant neural network with parameter on data z, Le., F. F=(x)= = (Cte)) for any parameter æ, input data z and constant C > In other words, the output of the network only depends on the direction of parameter, shorthanded 전 == z e/||=||22

---

Let E be the total number of training data and L(x) = - E Lę(x) be the empirical training loss. Denote (xx) = (0r (z),..., (.....,),)()))) where each 0g(x) == (VLg(z) VL(x))/VE Then we can rewrite SGD+WD (Equation (1)) as
Tn + 1) = (1 (k) (Von))))))((V(a)))) (VL(2x,))()) + V= oce (
(2)
Let (c20 be E-dimensional Brownian motion. As common approach analyzing SGD, the canonical SDE approximation of SGD+WD (Equation (2)) is
(t) (V))X, (t)d, no (X, (()))())))
The Fast Equilibrium Conjecture stated below. The convergence rate is much faster than the (111n) global mixing time of Langevin dynamies [18] and thus the conjecture gets its name.
Conjecture 1.1 (Fast Equilibrium Conjecture, Li et al. [T]. Suppose (4,)()) is solution of (3), then for any input 2, F2(X, (x(,,,,()) v,,(()) converges to the same equilibrium distribution independent of the initial parameter Tinit in ((1 time.
We note that the above conjecture implied by the convergence of the distribution of the parameter direction X (()) Next, the main theorem of this paper is stated informally below.
Theorem 1.2 (Informal version of Theorem 5.5). Suppose I is connected manifold consisting only of local minimizers of L. Under some regularity assumptions, there is an open neighborhood U of T, such that for any initialization Vinit E U and T > O, ιÿλ + 0 with
O()) S
0(1), both .n,{ (L
In( (e"
converge in distribution to the I) and Xn,> IIn( ( (c (e"T
4nl2
4n/2
same distribution denoted by prr Moreover, as T + 00, MT.. weakly converges to the same equilibrium distribution for every Tiit U.
The main contributions of this paper are summarized as follows:
1. We give a SDE-based characterization (Theorem 4.4) for the limiting dynamics of SGD+WD for scale invariant stochastic loss in the limit of y + 0 with ") O()) and À O1)) By introducing a novel rimealisc tailored the scale invariant loss and weight decay, our analysis adapts the framework proposed by Li et al. [2].
2. We show that SGD without WD for a stochastic scale invariant loss has the same limiting dynamics as that of SGD+WD, but is exponentially slower (see Theorem 3.2). This is consistent to the empirical observation that turning on WD for SGD for scale invariant loss helps generalization [19, 1].
3. Under the assumption of all minimizers forming manifold and noise being non-degenerate in the tangent space of the manifold, we show that from any initialization, the limiting dynamics of SGD+WD converges to unique stationary distribution (see Theorem 5.4 and Theorem 5.5). This delivers partial proof to the Fast Equilibrium Conjecture [1].
4. Though our convergence result is asymptotic, we verify in simplified settings that the phe- nomena predicted by our theory happens with LR r and WD factor À of practical scale (see Section for details of experiments). We also show empirically that the mixing process exists in practical settings, and is beneficial for generalization.
2 Preliminary
Notations. We denote by N the set of all nonnegative integers and R4 the set of all nonnegative real numbers. For any k E N, we denote by Ck the set of all times continuously differentiable functions. For any vector E TR. we denote its i-th coordinate by u, For any mapping F RRD - RD, we denote the Jacobian of F at x by OF(x) E DxD where the (i.j)-th entry is Fz)) We also use (=) and 8 F(x)) [u, 2] to denote the first and second order directional derivatives of F at x along the derivation of u (and v). With slight abuse of notation, we view 02 F as linear mapping IR DxD such that 2 F2F(()) [A)
Σij=1 0 FF ejlAi, for any A E DDD For any submanifold I' c RD and I E P, we denote by (T) the tangent space at x. We denote by 11g RR the one-hot vector where the ę-th coordinate is 1, and 1 denotes the all 1 vector. We say K C D is a cone if and only if 0 K and Va 0, aK C K.
3

---

Recall that the training loss L: RD - R is defined as L = th sample. With vast overparametrization modern machine learning models, multiple minimizers can exist and form manifold [20, 21J. Thus following Fehrman et al. [22], Li et al. [2], Arora et al.
Li=1 L; where L; is the loss on the i-
[23], we make the assumption below throughout this paper.
Assumption 2.1. Each loss function Li IR" - R is a C+ function. I is a (D - M) -dimensional is a local minimizer of
cessubmanifold of IR." for some integer M E [0, 1]. where each L and rank( V L(x)) M for all E r.
Note that V L(x) must have zero eigenvalues in the tangent space of I at x, we are indeed assuming Hessian V attains the maximal rank everywhere on the manifold r.
In this paper, we are interested in the behavior of SGD+WD with each L, being scale invariant, or equivalently, 0-homogenous. We note that the level sets of scale invariant functions are always cones, which will be used frequently in our analysis. To this end, we make Assumption 2.3.
Definition 2.2 (Homogeneous Functions). We say function f IR (0) + RT is k-
homogeneous for some k R if and only if for all x RP \ {0} and a 0, f(az) a f(x). Specifically, we say function f is scale invariant if and only if is O-homogeneous. Assumption 2.3. L; is scale invariant for each Sis E and I' is a cone.
Below two useful properties of homogeneous functions, whose proofs follow from directly ap- plying chain rules.
Lemma 2.4. For any E N and k-homogeneous function f. V is (I -h-ooooggnnss
functitoooo Lemma 2.5 (Euler's Theorem for Homogeneous Functions). For any real-valued k-homogeneous
function J. (2r, [=())) = (xr).. Specifically, if f is scale invariant, (2,, Vf(x)) =
the noise function is (-1)--omogenenne and thus the noise covariance E(x) = (x) is (-2)-hommggeneous
Recall that (x) \/2 (v) VL(x)).
Next, the following notion of limiting map of gradient flow plays key role in our analysis.
flow Definition 2.6. For any I & O, we define the gradient flow governed by -VL as the unique solution
of ((I, 1) = x Sor V(O((, 8) s)sds for / 2 0 and denote its associated limiting map by o(x) lim p(x; t) whenever the limit exists.
Throughout the paper, we use U to denote the attraction set of I under gradient flow, that is, U {o E RD (z) is well-defined and (2x) E F}. By Lemma B.15 of Arora et al. [23], U is open and 4 is C in U with 2 being locally lipschitz.
Lemma 2.7. Under Assumption 2.3, U is a cone and & is 1-homogeneous in U.
Li et al. [2] established several important properties of & by relating the derivatives of t to those of L, and in particular, we recall the following characterization of 04.
Lemma 2.8 (Lemma 4.3, [2].. For any x r, 0Φ(x) E DD is the orthogonal projection matrix onto the tangent space (T). As a consequence, for any x E r, Assumption 2.3 = (T') 02(x)== x.
2.1 Limiting Diffusion on The Manifold of Local Minimizers
We recap the notion of Katzenberger processes proposed by Li et al. [2] and the characterization of the corresponding limiting diffusion based on Katzenberger's theorems [24]. In this subsection we only assume Assumption 2.1, but not Assumption 2.3.
Definition 2.9 (Uniform metric). The uniform metric between two functions f.g: 10, oo) - >RP is defined to be (f( Σ 2- suprreup)))) (t) - g(t)]l2).
For each n E N. let A, R + R. and B, R. + R, be two non-decreasing functions with A,, (0) = (0) O, and (Z, be a v--uudd stochastic process. our context of SGD+WD, given loss function T. TR' - R. noise function a R²" -> DxE and initialization Tini U, we call the following stochastic process (4) Katzenberger process
Pinit S/' ((X, + /. (n(()dB((s) - / vL(Xx(s))dA, (4)
(t)

---

if as + do the following conditions are satisfied:
1. An. increases infinitely fast, i.e., Ve > 0, 0, + e) An(1)) + 00;
2. (t) converges to c & in uniform metric for some constant
3. Zn converges in distribution to (4= - ] in uniform metric where W is a E- dimension standard Brownian motion;
Theorem 2.10 (Adapted from Theorem 4.6 in Li et al. [2]). Given Katzenberger process (Xx(+) if SDE (5) has a global solution Y in U with Y()) = ① ((Tia), then for any > 0, (I) converges in distribution (1) as n 00.
/ cO
S.' ( (
(()) ((Tma)
So 1 (
(5)
We note that the global solution always exists if the manifold r is compact. Our notion of Katzen- berger process and theorem statement is slightly more general than those in Li et al. [2] to handle the weight decay. However, our formulation still under the original framework of Katzenberger [24] and the proof in Li et al. [2] can be easily adapted to Theorem 2.10.
3 Warm-up: Simultaneous Limit Case
As warm-up, first consider the setting where 77, À - simultaneously with C for some constant C O. In this special regime, we do not need use the scale invariance property of the loss, and we can directly apply Theorem 2.10 to obtain the limiting diffusion of SGD+WD. Nonetheless, we will see the benefit of weight decay as a source of acceleration. While for the general case of ryλ + 0 that will be considered in Section 4, we need to carefully design a time rescaling by calibrating with the dynamics of parameter magnitude, so that under the new scaling the dynamics can still be understood as Katzenberger process.
Now, recall the SGD+WD updates in Equation (2), and let fix ( Tini for some Finit E U. Define x,,x() (Ltnn(( which is roughly equivalent SDE (3) with 1/1/² times accelera- tion, and we can rewrite the discrete-ting time update of Zn,2 as
(t) = Tinit +/f o
+ +// Xna(s)dBna(()- -/6 VL(Xm (6)
where A B and Zn,2 are defined by
= nl/nn =nntt)-Dnlth??) (t) E L4/12! VE (=(1.. 11
1 1
(7)
Note that (1) is roughly i/n which becomes very large for small n,, thus the negative gradient part will drive (t) rapidly towards the manifold and force (t) stay close to after that. On the other hand, as + 0, (t) will converge to Ct and Zn,A will weakly converge to a Brownian motion, and these terms make up the slow dynamics of SGD. More precisely, we have the following lemma summarizing the properties of these integrators, which shows Equation (6) is a valid Katzenberger process.
Lemma 3.1. Let Aq, B and Zy, be as defined in Equation (7). Then as n7, X - with À C, holds that (1) A increases infinitely fast, (2) (t) converges to Ct in uniform metric and (3) Z converges in distribution to (Iz- - - T in uniform metric where (1))r22 is the E-dimensional standard Brownian motion.
"
Therefore, a direct application of Theorem 2.10 yields the limiting diffusion in this case.
Theorem 3.2. Under Assumption 2.1, let ,,,, (0) Tinit U, Vn, 0 SGD+WD (2). Consider
dYc(!) = CCB 0² - (Y+) (00))))+0(()))))))))))
(8)
where (W)))t2 is the standard Brownian motion R. Suppose SDE (8) has a global solution Yc in U for some C O with (0( (FIiit), then xn,(([t/72]] converges in distribution to
5

---

Yc(t) as X, + O with = C. Also, under Assumption 2.3, SDE (8) with any C' 2 0 has global solution Yo in U with Yc = (I Moreover, Yc (t) d Yo! a0"
Remark 3.3. Theorem 3.2 shows that when there is no WD, the limiting diffusion is still the same as that with WD, but exponentially slower than that with WD in the regime of LR on and WD factor going to zero with a fixed ratio.
4 Limiting Diffusion for The General Case
In the previous section, we showed that the limiting diffusion exists when and À go to zero with fixed ratio. However, the situation is more complicated in the general case where we drop the assumption of n/A being fixed. Below we first explain the challenges in analysis and our solution for this general regime. Then we present the continuous-time analysis for SDE and the discrete-time analysis for SGD+WD. Our analysis applies for all the cases when nzÀ + O with
O()) = O(1).
Challenges for the General Case. A concrete example for the challenge is when 1 - 0 and A be fixed as a constant, which is also the most natural and practical setting. We quickly find ourselves in dilemma if still want to apply Katzenberger's theorem [24], its simplified version The- orem 10: If we view WD or &g regularization as the 'fast' past of the dynamics, that is, a part of the loss function, then there no minimizer for the &g regularized scale-invariant loss and thus it doesn't t satisfy the condition of Katzenberger's 's theorem; if we view WD as some 'slow" dynam- ics and formulate it as À xd,,dB,,11 in Equation (6), then unlike the simultaneous limit case, À doesn't necessarily have a limit, and thus the condition of Katzenherger's theorem is again not met.
The above dilemma reflects two different roles of WD in early and late phase of training: in the early phase, when the norm is large, WD is more like a part of the loss function that exccutes the le regularization. In contrast, in the late phase of SGD training, especially when the norm square of parameters has stabilized to some value, i.e., ,|=,,,x (t) x / (e.g. Figure lb), WD should be viewed as the "slow" dynamics and we can apply the analysis in the simultaneous limit case. This is because by the scale-invariance of loss J., Equation (2) can be rewritten the following form, Equation (9), with rj /nx, Ā vTx, 9.5
(2) 12/// xm,A-
$ x + 1) = (1 - , (k) 7 (1) + VE (F (k)))
With such a rescaling, we successfully make the norm of parameters in constant scale, that is, , 11 = / |= = (1) and thus we can apply Katzenberger's theorem. Note that we cannot do this the early phase of SGD+WD as we start from fixed initialization and such rescaling will change the magnitude of the initilaization.
Our Strategy for Analysis. To overcome the above dilemma, our core strategy is to introduce a novel combination of parameter rescaling R (Equation (11)) and time rescaling Ty, (Equa- tion (12)) which smoothly interpolates the early and late regime. Because the rescalings tive to the norm of the parameter along the training trajectory, they allow us to apply Katzenberger's
adap-
Xn.((17))00 (t))
theorem on the rescaled dynamics Hn, (T, 2 (()) (See formal statements in Theorem 4.1). Compared
with the ordinary SGD without WD studied in Li et al. [2] where the time rescaling is set to be fixed acceleration by 1/n2 times, here the design of the time rescaling is more complicated.
Since the norm has no effect on the loss value but only affects the speed, we need to consider the dynamics of the parameter direction. To do so, we need normalize the iterates properly. However, when the trace of the noise covariance is not constant, in general it is hard to find close-form solution for ||.X,,, (f), but we can approximate using the former special case. In specific, recall the canonical SDE approximation in (3). Li al. [1] proved that the dynamics of (()),, is 21/ X,,X,,,, ,. (0)12 + n/ (.,,(XX,,,))) d|| n, (1)12
Suppose tr((()))
62/ || for some a > 0, then the above further simplifies into dI -21/A|X x,,x (1)l3dd + to, dl, which admits closed-form solution:
8-8²2
4hrt (IIx,A(0011
n102
IX,, (t)11
(10)
2X
2A

---

This implies that the norm of the weights at the equilibrium is of order (n/ Moreover, Equa- tion (10) reflects the scaling of the norm of the iterates terms of 17 and À as we will see later.
4.1 Continuous-time Analysis for SDE
We first consider the continuous-time case of the SDE approximation (3). The main result is sum- marized in Theorem 4.1, which shows that the limiting diffusion exists for SDE with suitable non-linear rescaling.
As mentioned in the previous discussion, we consider scaling function Rl,a(!) inspired by the norm dynamics for the special case in Equation (10):
a (t) = 22
7
e 4nAt
1 71
)
1/4
(11)
2X
Next, to rescale the time, we define Fn,A [0, 00) + (0) by
(t) I. n?
Jam (s) ds _ In (1 (E - 1) 7 22
(12)
= where the second equality follows from a direct calculation. It is easy to show that 7, (T)
(T)
lu( 2 -1)+1)) Then we have the following theorem (see Appendix for the proof), which says the rescaled version of SDE approximation Equation (3) admits the following limiting dif- fusion Equation (13), where (W))) is the standard Brownian motion in IR=.
4nj
1 d) - Y + (o (Y))) + ( (Y)
1
(13)
2
Theorem 4.1. Under Assumption 2.1 and 2.3, let (0) Tinit U for all n, X > 0 in SDE (3). Let (1) and (t) be defined in (11) and (12). SDE (13) has global solution Y in U with
Y()) (FFim)), then X1,x1T,, (T)) converges in distribution (T)) as nn with n 2X <
for some constant c. Remark 4.2. The additional constraint < 2X when T/ + 0 can be relaxed to < 2CX anr any
constant C > It suffices to note that for SDE (3), Xn,2 with X,,x(0) Finh and X c X/-2/2,,,0/1x with Xx-112 (0) C-1/4 Xima have the same trajectories, up to a rescaing of 1144 This is equivalent to replace n// (2X) by (2Cx) in (11).
4.2 Discrete-time Analysis for SGD+WD
Now, we proceed to analyze SGD+WD by mimicking the continuous-time behavior. Specifically, view a (k) approximation of the norm of (k), and consider the rescaled version (k),) Next, we introduce the time rescaling through the (-) defined in Equation (12) and denote & = (t), so t (f). Now define
of Equation (2) denoted by (k)
(1) = (L()( and it can be shown that (see Appendix E for the derivation)
(f) = N, + f
(
8=0
(14)
of (8)
0=0
where Ay,2 By,,, and Z are defined by
1,, Σ
L (0))
7
(15)
i=1
(i) +1¹
(1) = Σ (x))
Ri)(1)
a (1 - -1/)
(16)
Ri +
Zt,)(t) Σ 1
(T)
n/3
(1
1. (14
=
1
(17)
(i,
The convergence of An,x, BY and Zn are summarized in the following lemma.
7

---

Lemma 4.3. Let Aq,2, By, and Zn be as defined in Equation (15), (16) and (17) respectively. Then as nd + O with 7 2A < for some constant c, holds that (1) A,, increases infinitely
fast, (2) B,, (t) converges to 2 in uniform metric, and (3) Zn,> converges in distribution (I=
A A1 in uniform metric where {W(t) h12 is the E-dinensional standard Brownian motion.
Therefore, (t) is also valid Katzenberger process. Applying Theorem 2.10 yields:
(f)
in SGD+WD WD (2). Let R1,2(() and T1,)(1) be defined in (11) and (12). If SDE (13) has a global solution Y in U with (0) ((Fini)), then for any T > En,AllT, A(T2)
Theorem 4.4. Under Assumption 2.1 and 2.3, let (0)2(0) =\t Finit
U for all T₁
converges in distribution
Ro,Al nn (T)))
to (T) as T7 + 0 with r < 2A < c some constant c.
5 Mixing to Equilibrium
Now we proceed to study the ergodicity of the limiting diffusion (13). Omitted proofs of this section are delayed to Appendix F.
Due to the nature of the scale invariant of the loss L,, we only care about the direction of (t). i.e., (t) (t)l2. 2. To study the ergodicity of the normalized diffusion process, we need some additional assumptions on I' and the noise covariance. For any > O, define I,, == I'n {x: IR, ||I||| r). We assume that I' is compact manifold to ensure the existence of stationary distribution of the limiting diffusion process. We also need to assume that J1 connected (so is I') for the uniqueness of the stationary distribution.
Assumption 5.1. I'1 is compact and connected.
We further assume that the noise is non-degenerate on the manifold of local minimizers, so that as a Markov chain the limiting diffusion is irreducible.
Assumption 5.2 (Controllability). For each x E I, span (z))
(T (T
Assumption 5.3. Lr(E(-)) is an analytic function on IRP) {0) and P is an analytic manifold.
Now we are ready to state our main result in this section, which Theorem 5.4. It is proved in two cases respectively in Appendices F.3 and E.4, depending on whether the trace of gradient covariance tr(E) is constant on F or not. If it is, then the diffusion process essentially is on a (D - .M 1)- dimensional manifold (after suitable rescaling). ri, just in the analysis by Wang and Wang [25]. Otherwise, the situation becomes more complicated and the diffusion process is on I, a (D M)- dimensional manifold, in which case will need the analyticity assumption (Assumption 5.3).
Theorem 5.4. Under Assumption 2. 2.3, 5.1, 5.2 and 5.3, starting from any initialization ( E r, the distribution of Y (()) converges to a unique stationary distribution π on F in total variation.
Our main result on the fast mixing of SGD+WD and its SDE approximation follows from direct combination of the convergence of the SGD+WD iterates proved in Theorem 4.4 and Theorem 5.4. Here note that as I1 is compact, convergence in total variation implies convergence in distribution.
Theorem 5.5 (Fast Mixing of SGD+WD). Under Assumption 2.1, 2.3, 5.1, 5.2 and 5.3, let 2n (0) = (0) = Zinit E U for all n, X > 0 for SGD+WD WD (2) and SDE approximation (3). For any T 0, as yd + 0 with v O()) and X 0(1), both (L
In[ 23 (c2""11+ J) and And
X,, ( In[ 20 (62 -124 converge in distribution to the same distribution, denoted by aT More-
4nd
over, for every Finit E U, fT. weakly converges to the same equilibrium distribution π supported
on I1 as T + oc.
6 Experiments
In this section, we first empirically verify our theory for the time scaling of the dynamics in simple setting where our theory applies. We then show that the diffusion process exists during the training of PreResNet on CIFAR-10, and it has implicit bias towards better generalization.
6.1 Verification of Time Scaling
Setting and Theoretical Prediction. We train the following normalized linear model by la re- gression: F2( { =), where E D is the model parameter, and = RD is the input.
8

---

log 10/Leamiin rate)
-4-
20 40 60
60
100
120 140 160
Learning rate Steps
Figure 2: Scatter plots for projections of model parameters into the affine subspace containing the ID zero-loss unit norm manifold. Models are trained with LR 1 E (10-2 10-3,0000000 10 and WD factor À = 0.05. Each small box in the figure contains 60 models that trained with the same LR and the same number of steps. The axis indicates the product of LR and the number of steps each box; the y axis indicates (LR) of each box. The dashed lines represent the time scaling $ In (1 (e 1) ) T for different T's (from 0
15.58), where is the number of steps in SGD+WD. The dynamics consistent with time scaling suggested by our theory (Theorem 4.1).
Let {2s N be the input samples, and y/1 Fz, (x*) be the target label for each 2,, for some x². The training loss [(x)= = 2N
N (F (z) - F2, (x*))² We set N = D - 2, so the solution space S {w (w oF 2n) O, Vi [N]] is 2-dimensional linear space. The manifold N.N. of unit-norm global minimizers Fı is then equal to S 5 ||ا|| = 1}. We generate } randomly in a way that almost surely x. is not contained the linear span of { This implies that M rank( 2 [xx))= D - on 1 and that i is 1-dimensional manifold (a circle), thus As- sumption 2.1 and 5.1 hold. During training,
Σ
set the loss at step (x) L(x) + { er)
where iid (O, ((0,02Tm T for some 6. Then the SGD+WD update rule is
[ 1) (1 (k) 7 (I(a))())))) [ID
x..
Et
(18)
||=1|/
||
T ). and the canonical SDE approximation is Let (() = [ -
(In
dX, =- VX - -,A,,, + (x((xxx)))WWW)
(19)
For SDE approximation, we E D, and thus Theorem 4.1 applies to Equation (19), suggesting that the correct time scale (number steps for SGD+ WD) for the limiting dynamics is T (T)
-12 for each T 2 0. Furthermore, Assumption 5.3 holds because every term in this 4nA example is analytic. Assumption 5.2 holds because for any vector E T₂ (T [[]]),) we have (z. v}
and hence (+()))(x)) (z)x) TTlle T U. Therefore, our main theorem (Theorem 5.5)
predicts that the limiting dynamics mix in (n(A)))+O())
steps.
4nA
Remark 6.1. For ease of demonstration, our main theorem (Theorem 5.5) is proved for SGD+WD
with finitely many samples and thus does not directly apply to Equation (18). However, our analysis can be extended to the case Gaussian noise in an straightforward way and the claim Theo- rem 5.5 indeed holds for Equation (18).
Experimental Results. In our experiments, we choose D = 10, 0.3, the WD factor X == 0.05, x and LR E {10-2,11 2, 20 1000 $). In Figure 2, we plot the projections of |[=]| on the solution space for 60 different runs with identical initialization for each 7). For each run, the only differences the LR y and/or the noise &s. The dashed lines in the figure indicates our time scaling, i.e., In (1 + (enn ) T. Figure 2 shows that time scaling of O(In 2A/1 -0/nn117 T))((n)) fits the dynamics better, compared to (TT))
6.2 Limiting Diffusion on CIFAR-10
Beyond the toy example, we further study the limiting diffusion of PreResNet on CIFAR-10 10 [26]. We train a 32-layer PreResNet [27] with initial LR 1 17-0.8 0.8 and WD factor X = 5 10-4, Unlike the normalized linear model, hard to visualize
PreResNet on the manifold. Instead, we choose the test accuracy of d(xx) as a test function. In particular, we decay the LR 10-3 3 to approximate gradient flow at different time and record the test accuracy after training 1000 more epochs. The results are shown in Figure I and Figure 3.
model projection

---

We observe that without LR decay, the train accuracy, test accuracy and parameter norm converge quickly after training 100 epochs; the test aceuracy of (zz,) converges much slower. It suggests that there exists a mixing process after reaching the manifold. Moreover, we observe the test accuracy of ((x,) after convergence is significantly higher than ((xz)) at 100th epoch. It indicates that the mixing process is beneficial for generalization.
Our time scaling (12) suggests that the optimal step of LR decay grows no faster than (2(1/n)A) as ? + O. Unfortunately, (12) alone is not sufficient for deciding the optimal step for decaying LR as the mixing time 7 for the continuous dynamics is unknown. A potential usage of time scaling (12) first estimate T via another SGD run with a larger LR, which we leave for future work.
7 Related Work
Normalization and Scale Invariance. Previous works have analyzed the benefits of normalization layers from different viewpoints [28, 13, 29-50]. As noted before, normalization layers induce the scale invariance. It has been shown that scale invariance enables robust and efficient training of SGD+WD [51]. Scale invariance also brings about the interesting equivalence between the effect of WD and LR schedules [52. 19, 17J. Moreover, for SGD+WD with LR and WD factor X, the parameter norm will converge to (A/1)1/4 1/4 [1. 53, 54], and this induces the intrinsic LR which is equal yi [1]. These observations are crucial our derivations in the current paper.
Fast Equilibrium Conjecture. Recently, Wang and Wang [25] proposed a spherical SDE model to approximate SGD+WD with constant LR. Using a novel adaption of Simon's theory, they justi- fied the Fast Equilibrium Conjecture by showing that SGD+WD dynamics consists of three stages: descent (Q())VXT) time), diffusion (O(l/(>n) time) and tunneling (OAA))())) time). However, their analysis relies on the strong assumption of the minimizers of L being isolated, which is against the empirical evidence that the level sets of deep learning loss functions are connected [55]. As a result, the diffusion phase shall bring no generalization benefit and cannot explain the improvement of final generalization if staying at training loss plateau for longer time (see Figure 1). We allow the local minimizers to form a connected manifold, which can be viewed generalization of the Morse function assumption [25]. as an isolated minimizer is just a manifold of dimension
Another common weakness of existing analyses in [1, 25] is that they only work for the SDE ap- proximation (3), and do not apply to the actual discrete-time dynamics (2). In contrast, our results can handle both the continuous and discrete time dynamics under more reasonable assumptions.
SDE Approximation. Continuous- tools such SDE have been popular lens for studying optimization algorithms including SGD [56-61]. Many interesting properties of SGD have been discovered through this approach [1, 62-65].
Slow Dynamics of SGD Around Zero Loss Manifold. Recent works [66, 16, 2] show that under the assumption that the minimizers locally connect manifold, SGD with label noise with small learning rate will move around the manifold after convergence, towards the direction of smaller trace of Hessian, at a very slow rate of O(7²²) per step. Arora et al. [23] show that such slow dynamics on manifold can happen without stochastic gradient noise, if the update rule is non-smooth around the manifold of minimizers and GD enters Edge of Stability regime ([67). Concretely, they show that normalized GD implicitly penalizes the largest eigenvalue of the Hessian at the rate of (n?)). Additional related works are deferred to Appendix A.
8 Conclusion and Future Work
We provide an SDE-based characterization for the limiting dynamics of SGD+WD for scale invari- ant loss as ylλ + with 7 = O()) and À = O()). Under some technical assumptions, we further show that the limiting diffusion converges to unique stationary distribution. It leaves as future work to relax the technical assumptions. Another interesting and important direction for future work is to understand and characterize the benefit on generalization induced by the limiting diffusion.
Acknowledgements
This work funded by NSE, ONR, Simons Foundation, DARPA and SRC. ZL. is also supported by Microsoft Research PhD Fellowship.
10

---

References
[1] Zhiyuan Li. Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with tra- ditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems, 33, 2020.
[2] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. In International Conference on Learning Representations, 2021.
[3] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. URL https: /tttpp://opentee /ttp:////pp//.e.e///e//e http::///.pppproroooooooooeeeaaa
[4] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint arXiv: 2019.
[5] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp min- ima. arXiv preprint
2016.
[6] Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprint arXiv: /804.0762, 2018.
[7] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456 PMLR, 2015.
[8] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv preprint 1601v:1677.0644 2016.
(9] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3-19, 2018.
[10] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization celerate training of deep neural networks. Advances in neural information processing systems, 29:901-909, 2016.
[11] Siyuan Qiao, Huiyu Wang, Chenxi Liu, Wei Shen, and Alan Yuille. Micro-batch training with batch-channel normalization and weight standardization. arXiv preprint arXiv:1903 10520, 2019.
[12] Lei Huang, Xianglong Liu, Yang Liu, Bo Lang, and Dacheng Tao. Centered weight normaliza- tion in accelerating training of deep neural networks. In Proceedings of the IEEE International Conference on Computer Vision, pages 1803-28111 2, 2017.
[13] Andrew Brock, Soham De, and Samuel L Smith. Characterizing signal propagation close the performance gap in unnormalized resnets. In International Conference on Learning Rep- resentations, 2021.
[14] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02/17 2022.
[15] Guy Blane, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Conference learning theory, pages 483-513. PMIR, 2020.
[16] Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global mini- mizers. arXiv preprint arXiv:2106.065300 2021.
[17] Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In International Conference on Learning Representations, 2019.
[18] Bin Shi, Weijie J Su, and Michael Jordan. On learning rates and schri"" odinger operators. arXiv preprint 2xX1v:2004.06977, 20200

---

[19] Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. arXiv preprint arXiv:1810.12230 2018.
[20] Yaim Cooper. The loss landscape of overparameterized neural networks. arXiv preprint 1804. 10200, 2018.
[21] Cooper. The critical locus of overparameterized neural networks. arXiv preprint arXiv:2005.04210, 2020.
[22] Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions. Journal of Machine Learning Research, 21, 2020.
[23] Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Con- ference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 948-1024. PMLR, 17-23 Jul 2022. URL https: https:://///ceeee http:///oooeeeenn.e.// mlr. mrlr.press/v16// arora22a. html.
[24] Gary Shon Katzenberger. Solutions of a stochastic differential equation forced onto a manifold by a large drift. The Annals of Probubility, pages 1587-1628, 1991.
[25] Yi Wang and Zhiren Wang. Three-stage evolution and fast equilibrium for sgd with non- degerate critical points. In International Conference on Machine Learning, pages 23092- 23113. PMLR, 2022.
[26] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[27] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 660-645. Springer, 2016.
(28) Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. arXiv preprint arXiv:2206.0777 07008, 2022.
[29] David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question? In International Conference Machine Learning, pages 342-350, PMLR, 2017.
[30] Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function in deep networks. Advances in Neural Information Processing Systems, 33:19964- 19975, 2020.
[31] Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, and Aurelien I Lucchi. Batch normalization provably avoids ranks collapse for randomly initialised deep networks. Ad- vances in Neural Information Processing Systems, 33:1 33:18387-18398, 2020.
[32] Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes rep- resentations in deep random networks. Advances in Neural Information Processing Systems, 34:4896-4906, 2021.
[33] Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka. Beyond batchnorm: towards a unified understanding of normalization in deep learning. Advances in Neural Information Processing Systems, 11:4778-4791, 2021.
[34] Antoine Labatie, Dominic Masters, Zach naton-Rosen, and Carlo Luschi. Proxy-normalizing activations to match batch normalization while removing batch dependence. Advances in Neu- ral Information Processing Systems, 34: 19900-1700706 2021.
[35] Ping Luo, Xinjiang Wang, Wenqi Shao, and Zhanglin Peng. Towards understanding regu- larization in batch normalization. In International Conference on Learning Representations, 2019.
12

---

[36] Mattias Teye, Hossein Azizpour, and Kevin Smith. Bayesian uncertainty estimation for batch normalized deep networks. In International Conference on Machine Learning, pages 4907- 4916. PMLR, 2018.
[37] Alexander Shekhovtsov and Boris Flach. Stochastic normalizations as bayesian learning. In Asian Conference on Computer Vision, pages 463-479. Springer, 2018
[38] Shibani Santurkar, Dimitris Tsipras, Andrew Nlas, and Aleksander Madry. How does batch normalization help optimization? Advances in neural information processing systems, 31, 2018.
[39] Nils Bjorck, Carla P Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normalization. Advances in neural information processing systems, 31, 2018.
[40] Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net opti- mization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-2241. PMLR, 2019.
[41] Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. The normalization method for alleviating pathological sharpness in wide neural networks. Advances in neural information processing systems, 32, 2019.
[42] Zinan Lin, Vyas Sekar, and Giulia Fanti. Why spectral normalization stabilizes gans: Analysis and improvements. Advances in neural information processing systems, 34:9625-9633 2021.
[43] Susanna Lange, Kyle Helfrich, and Qiang Ye. Batch normalization preconditioning for neural network training. Journal of Machine Learning Research, 23(72):1-41, 2022.
[44] Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. In International Conference on Learning Representations, 2019.
[45] Yongqiang Cai, Qianxiao Li, and Zuowei Shen. A quantitative analysis of the effect of batch normalization on gradient descent. In International Conference on Machine Learning, pages 882-890. PMLR, 2019.
[46] Jonas Kohler, Hadi Daneshmand, Aurelien Lucchi, Thomas Hofmann. Ming Zhou, and Klaus Neymeyr. Exponential convergence rates for batch normalization: The power of length- direction decoupling in non-convex optimization. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 806-815. PMLR, 2019.
[47] Xiaoxia Wu, Edgar Dobriban, Tongzheng Ren, Shanshan Wu, Zhiyuan Li, Suriya Gunasekar, Rachel Ward, and Qiang Liu. Implicit regularization and convergence for weight normaliza- tion. Advances in Neural Information Processing Systems, 33:2835-288 2020.
[48] Chao Ma and Lexing Ying. A riemannian mean field formulation for two-layer neural networks with batch normalization. Research in the Mathematical Sciences, 9(3): 2022.
[49] Yonatan Dukler, Quanquan Gu, and Guido Montúfar. Optimization theory for relu neural networks trained with normalization layers. In International conference machine learning, pages 2751-2760. PMLR, 2020.
[50] Shengchao Liu, Dimitris Papailiopoulos, and Dimitris Achlioptas. Bad global minima exist and sgd can reach them. Advances in Neural Information Processing Systems, 33:8543-85521 2020.
[51] Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar. Robust training of neural networks using scale invariant architectures. In International Conference Machine Learning, pages 12656-12644 PMLR, 2022.
[52] Elad Hoffer, Ron Banner, Itay Golan, and Daniel Soudry. Norm matters: efficient and accurate normalization schemes in deep networks. Advances in Neural Information Processing Systems, 31, 2018
[53] Twan Van Laarhoven. 1.2 regularization versus batch and weight normalization. arXiv preprint 10X5v31706.05350 2017.
13

---

[54] Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learn- ing dynamics neural network with normalization, weight decay, and sgd. arXiv preprint urXiv:2006.084111 2020.
[55] Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. International conference on machine learning, pages -318. PMLR, 2018.
[56] Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as ap- proximate bayesian inference. arXiv preprint arXiv: 170.004899 2017.
[57] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference Machine Learning, pages 2101-2110, PMLR, 2017.
[58] Justin Sirignano and Konstantinos Spiliopoulos. Stochastic gradient descent in continuous time. SIAM Journal on Financial Mathematics, 8(1):933-961, 2017.
[59] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms Mathematical foundations. The Journal of Machine Learning Research, 20(1): 14:441-2-,, 2019.
[60] Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and langevin processes. In International Conference on Machine Learning, pages 1810-1819, PMLR, 2020.
[61] Justin Sirignano and Konstantinos Spiliopoulos. Stochastic gradient descent in continuous time: A central limit theorem. Stochastic Systems,
2020.
[62] Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint arXiv: 300iv:/800300/99 2018.
[63] Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent escapes from sharp minima exponentially fast. arXiv preprint arXiv:2002.03495, 2020.
[64] Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. arXiv preprint S.Kiv:2106.05222 2021.
[65] Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part i:: Continuous time analysis. arXiv preprint arXiv:2 106.:20800.0888, 2021.
[66] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. arXiv preprint arXiv: (0f1::1904.00000 2019.
[67] Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability, 2021.
[68] Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure of neural network loss functions: The effect on optimization and origin. arXiv preprint arXiv:2204. /1326, 2022.
[69] Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint XXX::2211.007222 2022.
[70] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware mini- mization for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
[71] Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local sgd generalize better than sgd? OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop).
[72] Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. arXiv preprint arXiv:2210. 14199, 2022.
14

---

[73] Wolfgang Kliemann. Recurrence and invariant measures for degenerate diffusions. The annals of probability, pages 690-707, 1987.
[74] Kanji Ichihara and Hiroshi Kunita. A classification of the second order degenerate elliptic operators and its probabilistic characterization. Zeitschrift für Wahrscheinnicckeisssheeeie und Verwandte Gebiete, 30(3):235-255 254, 1974.
[75] Martin Hairer. Ergodic theory for stochastic pdes. preprint, 2008
[76] Giuseppe Da Prato and Jerzy Zabezyk. Ergodicity for infinite dimensional systems, volume
229. Cambridge university press, 1996.
[77] Andrei Agrachev, Davide Barilari, and Ugo Boscain. A comprehensive introduction sub- Riemannian geometry, volume 181. Cambridge University Press, 2019.
[78] Jean-Michel Coron. Control and nonlinearity. Number 136. American Mathematical Soc..
2007.
[79] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
[80] Stewart N Ethier and Thomas G Kurtz. Markov processes: characterization and convergence, volume 282. John Wiley & Sons, 2009.
[81] Ward Whitt. Proofs of the martingale felt. Probability Surveys, 4:268-3011 2007.
[82] David Pollard. Convergence of stochastic processes. Springer Science & Business Media, 2012.
15

---

Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] to [Yes], [No], or [N/A] You are strongly encouraged to include justification your answer, either by referencing the appropriate section of your paper or providing a brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section X.
Did you include the license the code and datasets? [No] The code and the data proprietary.
Did you include the license to the code and datasets? [N/A]
Please do modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below,
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A] This pa- per studies the theoretical properties of stochastic gradient descent with weight decay.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [No]
4. If you using existing assets (e.g., code, data, models) curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're re using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifi- able information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "x3sXc9ougGc", "reviewer": "Reviewer_o9Wv", "summary": "This papers theoretically proves the fast equilibrium conjecture that SGD for scale invariant loss mixes in  O(1/\\lambda eta) steps by con sidering a new time rescaling limiting SDE. The limiting SDE of SGD without WD mixes much slower than SGD with WD justifies the benefit of weight decay for generalization performance.  The paper also provide some empirical results showing that the fast mixing obtains in asymptotic analysis does appear in pratical setting. \n", "questions": "As stated in weakness. ", "limitations": "The major two limitations of the current results is:\n1. The technical assumptions may not be satisfied.\n2. The asymptotic approximation does not necessarily fully characterize SGD.\n\nHowever, as I have stated in Strengths and Weakness, the current results are already enough for a good NeurIPS paper. ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nOverall, this is an excellent paper that provides important theoretical results in understanding the behavior of SGD in training deep neural networks. More specifically,\nOriginality: The fast equilibrium conjecture was proposed to explains the interesting generalization performance as mentioned in line 33. As the name suggested, this conjecture has not been proved before. Therefore, the proof in this paper is original. Moreover, the new time rescaling SDE has never been studied in the existing literature, which is totally new to my best knowledge.\nQuality: All the results in this paper is clearly stated and justified. I briefly checked the proof and it looks sound to me. \nClarity: A well written paper with main results highlighted. \nSignificance: The proof of this conjecture  provide us a reliable tool to understand SGD, which is theoretically significant. \n\nWeakness: \n1. It would be better to provide some intuitive explanations of the assumption made to prove the main result. Can these assumptions be satisfied in the realized case? \n2. Beyond the proof of the conjecture, it would be more interesting to see what this mixing property can lead to some new conclusions. (Though I agree the content is already enough for one paper. )\n3. Section 5.2 provide some discretization analysis. More work should be done on this, e.g., the discretization error. \n4. Please correct typos in the paper, e.g., line 27 affect -> effect.\n\n\n\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "L2umB6ZkzZ", "reviewer": "Reviewer_fu5c", "summary": "This paper studies the mixing time of scale invariant models whose parameters are governed by a stochastic differential equation that approximates SGD + weight decay. Under certain regularity assumptions, the paper shows convergence of parameters to a unique distribution with a rate\n $$ \\Theta( \\frac{T+ln( \\lambda/\\eta) }{ \\eta \\lambda}) $$", "questions": "\nIf the trace of the noise covariance matrix is upper bounded by a constant, could we still have similar results?\n\nAlthough the paper critically rely on scale invariance, would we expect any similar theorems if the model had homogeneity degree L > 0? \n", "limitations": "Yes", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Under the setting of scale invariant models whose parameters are governed by a SDE approximation of SGD + WD, the paper proves the direction of SGD iterates converge in distribution to a solution of an SDE. The main technical strategy is to employ 1) time rescaling 2) scaling factor that closely mimics the normalization factor for the parameters.\n\nAs in  Z. Li (2021), time rescaling is used to ensure that the main workhorse (Katzenberger theorem) can be applied. Essentially, by rescaling time by $$ O(t/\\eta^2) $$, the negative gradient loss becomes the dominant term in the SDE (for small learning rate eta) which results in the parameters being close to manifold \\Gamma. \n\nFurthermore, the paper shows a cool corollary that shows that scale-invariant models with no weight decay converge exponentially slower. \n\n\nWeaknesses:\nThe main weaknesses are the strong technical assumptions (e.g. constant noise covariance trace, asymptotic results, smoothness assumptions) but understandable given the difficulty of the problem. \n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "BoTocS7XIF8", "reviewer": "Reviewer_QxFR", "summary": "The main contributions of the paper are 1) providing a characterizing SDE limiting the SGD+WD by using a time-rescaling technique so that the results by Katzenberger are applicable 2) showing that SGD without WD for stochastic scale invariant loss has the same limiting dynamics as SGD+WD, but is exponentially slowed down 3) under the assumption of all minimizers forming a compact manifold and noise being non-degenerate in the tangent space of the manifold, the authors show that from any initialization, the limiting diffusion process converges to a unique stationary distribution 4) lastly, the authors conduct experiments to empirically show that these results hold.", "questions": "1) What are the implications of this paper for the training of ML/DL models?\n", "limitations": "Mentioned in weaknesses.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThe authors are able to show the limiting dynamics for SGD+WD and devise a novel time rescaling scheme to make the results of [1] applicable as they claim that the results of [2] does not hold since there is no true minimizer in the usual sense.\n\nWeaknesses:\nIt seems that the time rescaling approach is the only notable contribution of the paper. The derivation of the SDE and other results seem to follow from [2]. While there seem to be novelties, the framework to analyze these dynamics have been provided by [2] so on that ground the contribution seems to be marginal.\n\n[1] G. S. Katzenberger. Solutions of a stochastic differential equation forced onto a manifold by a large drift. The Annals of Probability, pages 1587–1628, 1991.\n[2] Z. Li, T. Wang, and S. Arora. What happens after sgd reaches zero loss?–a mathematical framework. arXiv preprint arXiv:2110.06914, 2021.\n\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "7HyXVTyitKi", "reviewer": "Reviewer_Wb74", "summary": "This paper provided the dynamic convergence of SGD combined with the popular weight-decay technique (SGD-WD) for scaled-invariant loss functions (e.g., neural networks with normalization). Specifically, 1) the authors showed the iterates of SGD converge in distribution to some random variable in $\\mathcal{O}\\left(\\frac{1}{\\eta \\lambda}\\ln \\left( \\frac{2\\lambda}{\\eta}(e^{2T}-1) +1\\right)\\right)$ steps, where $\\eta$ is step-size, $\\lambda$ is the weight-decay parameter, and $T$ is the budget of the iterations. 2) In the regime of $\\eta = \\mathcal{O}(1/\\lambda)$ when $(\\eta, \\lambda \\rightarrow 0)$, the authors showed that SGD without WD has the same limiting diffusion as SGD with WD but its convergence is slower. 3) Under the assumptions on the manifold and the noise covariance, the authors partially proved the faster equilibrium conjecture of [26].", "questions": "The questions are addressed above. ", "limitations": "Yes, they have mentioned their limitations. ", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "In general, the paper is clear and well-written. As I know, the interplay of normalization, gradient noise, and weight decay of SGD focused in this paper is a very important and interesting problem. The study of their interplay is beyond the optimization theory. The main contribution is that the authors provided the limiting dynamics convergence of SGD with weight decay and gave empirical evidence to verify their results.  However, I am not an expert in the dynamic theory of SDE. I am afraid that I can not give a fair assessment of the theoretical parts. But I checked the numerical results carefully and gave my concerns.\n\n**About the experiments on the toy example:**\n- In line 262, Equality (24) has not been mentioned in the main content. The authors may use (10) instead.  The same problem appears in line 290.\n- in line 267, the authors said that as long as no $z_i$ is in the same direction with $x^{\\ast}$, $\\cdots$, then Assumption 6.4 holds. My question is how to ensure that no $z_i$ is in the same direction with $x^{\\ast}$. In general, it is difficult to gain knowledge about $x^{\\ast}$. \n**The results on CIFAR10:**\n- From Figure 1(a), the values of test accuracy and training accuracy change a lot and the algorithm is unstable. I am afraid it is hard to say that the algorithm has converged around 100 epochs. I guess because the initial step-size $\\eta = 0.8$ is too large. I am curious how do you select the initial step size? Are you based on the best performance of the algorithm?\n- As I know, a lot of papers studied that decaying the step size improves the performance. I think the interesting result is that the authors claimed that they can use the time scaling in (24) as an upper bound to decide the proper time of LR decay. Have you used this time scaling (24) to decide when to drop the step size in your experiments? I did not see that the selection of 100 epochs is based on their theory.  \n- Recently, there are some studies about step-decay step size, which is very efficient in practice. This step size keeps a constant at first and then drops by a constant factor after some iterations. For example, in [1], the authors analyzed this step-decay step size with the knowledge of the total number of iterations $T$. And they drop the step size every the same time period. I think the interesting thing is that the authors can use this theory of time scaling (24) to drop the step size adaptively. The authors have claimed this point but I am not very sure if the formula of (24) can be used in practice and achieve good performance in deep neural networks. I think more experiments are needed if the authors verify this claim.\n\n[1] Wang, Xiaoyu, Sindri Magnússon, and Mikael Johansson. \"On the convergence of step decay step-size for stochastic optimization.\" Advances in Neural Information Processing Systems 34 (2021): 14226-14238.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept