Fast Mixing of Stochastic Gradient Descent with Normalization and Weight Decay
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We prove the Fast Equilibrium Conjecture [26], i.e. , Stochastic Gradient Descent (SGD) for scale-invariant loss ( e.g. , networks with various normalizationschemes) with learning rate ⌘ and weight decay factor  mixes in function spacein ˜ O
( 1 ⌘ ) steps, under two standard assumptions: (1) the noise covariance matrix is non-degenerate and has constant trace and (2) all minimizers of the lossform a connected and smooth manifold.
The analysis uses the framework from[27] and shows that for every T > 0 , the iterates of SGD with learning rate ⌘and weight decay factor  on the scale-invariant loss converge in distribution in⇥  ⌘  1   1 ( T + ln(  / ⌘ ))  iterations as {+⌘ ! 
+} (while simultaneously satisfying⌘ .  and  . 1 ).
Moreover, the evolution of limiting distribution can be described by an Ito’s stochastic differential equation (SDE) which mixes to a uniquedistribution in time independent of ⌘ .
Introduction
Normalization methods such as Batch Normalization [14], Group Normalization [41], Layer Normalization [2] and Weight Normalization [33] have facilitated the efﬁcient training of modern deep neuralnetworks and in the extensive tool case stochastic gradient gradient (SGD) with weight decay (WD)is one of the main workhorses.
However, the interplay between normalization, gradient noise and
WD can move the training dynamics far away from standard optimization viewpoint, e g. , deepneural networks with normalization can be trained with an exponentially increasing learning rate (LR)schedule [25].
There is a line of research [3, 17] suggesting that small LRs (or equivalently large batch sizes) mighthurt generalization since they lead to the so-called sharp minimizers .
It is argued that noise can helpavoid sharp minima [11, 17, 22, 15, 9], so large LRs and small batches are preferred as they increasethe noise during training.
However, later Li et al.
[26] observed that normalized models trained using
SGD+WD can generalize equally well with small LRs, given longer training time after the trainingloss converges.
Such discrepancies imply an incomplete characterization of what happens duringtraining and especially the effect of the gradient noise.
To explain the effectiveness of small LRs,
Li et al.
[26] theoretically identiﬁed an intrinsic learning rate parameter equal to ⌘ ( ⌘ is the LRand  is the WD factor) which governs the convergence rate of the norm of the weights.
Based onempirical evidence, they further proposed the fast equilibrium conjecture which states that the mixingtime of the normalized model is roughly e
O (1 / ( ⌘ )) , much faster than existing upper bounds.
Thisconjecture is important because there is some unexplained mechanism towards better generalizationhappening during the mixing process — the training accuracy/loss plateaus but the test accuracy
[Figure]
after LR decay still improves (see Figure 1 for a demonstration).
The fast mixing time can providetheoretical guidelines to decide the proper time of LR decay.
Before stating the fast equilibrium conjecture, let us ﬁrst specify the problem settings.
We consider afamily of models F x whose output is scale-invariant to its parameters x , i.e., F x ( z ) =
F Cx ( z ) forany parameters x , input data z and constant C > 0 .
This scale-invariance property is commonlyinduced by various normalization schemes, and makes the loss function for any data (or batch ofdata) also scale invariant to x .
Let ⌅ be the number of training samples and the empirical loss
L ( x ) = 1 ⌅ P ⌅ i =
1 L i ( x ) where each L i is the scale invariant loss over the i -th sample.
We optimize
L using SGD+WD.
The dynamics is formalized as follows:
[Equation]
where ⌘ is the learning rate (LR),  is the WD factor, x ⌘ ,  ( k ) 2 R D denotes the parameters afterk iterations, x ⌘ ,  (0) =
x 0 is the initialization of the parameters, and  ( x ) =
[  1 ( x ) , . . . ,  ⌅ ( x )] {+
+} R D ⇥ ⌅ is the deterministic noise function with each  i ( x ) :
= 1 p ⌅ ( r L i ( x )  r L ( x )) .
Here each ⇠ kis sampled independently uniformly from { 1 , . . . , ⌅ } , so E ⇠ k [  ⇠ k ( x )]
= 0 for all x 2 R D and k .
As a common approach to analyzing SGD, the canonical stochastic differential equation (SDE)approximation of SGD+WD in Equation (1) is
[Equation]
where { W ( t ) } t  0 is the ⌅ -dimensional Brownian motion.
Conjecture 1.1 (Li et al.
Suppose X ⌘ ,  ( t ) is a solution of Equation (2), then F X ⌘ ,  ( t ) convergesto an equilibrium in e
O (1 / ( ⌘ )) time.
This convergence rate is much faster than the standard e  ⇥ (1 / ⌘ ) global mixing time of Langevindynamics [34] and that’s how the conjecture gets the name.
Li et al.
[26] proved that the norm of X ⌘ ,  ( t ) will converge in e
O (1 / ( ⌘ )) time but failed to show convergence to the equilibrium.
Anothergap is that the analysis in Li et al.
[26] only works for the SDE approximation in Equation (2), anddoes not bridge between the discrete-time and continuous-time dynamics.
In this paper, we provide a partial proof to this conjecture.
Theorem 1.2 (Informal version of Theorem 6.7) .
Let x ⌘ ,  := x ⌘ ,  k x ⌘ ,  k 2 and x ⌘ ,  (0) ⌘ x 0 for all⌘ ,  > 0 for SGD+WD in (1) .
Under technical assumptions, there is a unique distribution ⇡ such thatin the limit of {+⌘ ! 
+}
and T ! 1 , the distribution of ¯ x ⌘ ,  ( b ln( 2 ⌘ ( e 2 T  1)+1) 4 ⌘ c ) can be arbitrarily“close” to ⇡ in the sense of weak convergence for large enough T .
The main contributions of this paper are summarized as follows:
We give a SDE-based characterization (Theorem 5.4) for the limiting dynamics of SGD+WD for ascale-invariant stochastic loss in the limit of ⌘ ! 0 with ⌘ = O (  ) and
= O (1) .
In detail, weshow the direction of SGD iterate at step 14 ⌘ ln [Equation] converges in distribution tosome random variable Y T , where Y T is the solution of some SDE depending on the stochastic lossand initialization.
This result can not be proved by Li et al.
[27] directly as there is no minimizerin the usual sense, not to mention manifold.
The key proof technique is to introduce a noveltime-rescaling (Equation (9)) to make the convergence result in [16] applicable.2.
We show SGD without WD for stochastic scale invariant loss has the same limiting dynamics asSGD+WD, but is exponentially slowed down (see Theorem 4.2).
This is consistent to the empiricalobservation [43, 26] that turning on WD for SGD for scale invariant loss helps generalization.3.
Under the assumption of all minimizers forming a compact manifold and noise being nondegenerate in the tangent space of the manifold, we show that from any initialization, the limitingdiffusion process converges to a unique stationary distribution (see Theorem 6.6 and Theorem 6.7).This delivers a partial proof to the
Fast Equilibrium Conjecture in [26].4.
Though our convergence result is asymptotic, we verify in simpliﬁed settings that the predictedphenomena happens with LR ⌘ and WD  of practical scale (see Section 7 for details of experiments).
We also show empirically that the mixing process exists in practical settings, and it isbeneﬁcial for generalization.
Related Work
Scale-invariance.
For scale-invariant loss functions, existing works have proved the convergenceof GD and SGD with WD [1].
It has been shown that scale-invariance enables robust and efﬁcienttraining of SGD with WD [28].
Scale-invariance also brings about the interesting equivalence betweenthe effect of WD and LR schedules studied in [12, 43, 25].
Moreover, for SGD with WD, ⌘ being the
LR and  being the WD factor, the norm of the parameters will converge to (  / ⌘ ) 1 / 4 [26, 37, 38],and this induces the intrinsic LR which is equal to ⌘ [26].
These observations are crucial to ourderivations in the following sections.
SDE approximation.
Continuous-time tools such as SDE have been a popular lens for studyingoptimization algorithms including SGD [29, 23, 35, 24, 4, 36].
Many interesting properties of SGDhave been discovered through this approach [19, 8, 44, 26, 42, 31, 40].
Recently, [27] proposed aframework to study the implicit bias of SGD by showing that the iterates converge in distributionto a limiting diffusion as LR goes to 0.
Using this framework, we exhibit the limiting dynamics of
SGD+WD for inﬁnitesimal intrinsic LR.
Preliminary
We ﬁrst clarify the notations.
We denote by N the set of all nonnegative integers and R + the set ofall nonnegative real numbers.
For any k 2 N , we denote by C k the set of all k times continuouslydifferentiable functions.
For any vector u 2 R D , we denote its i -th coordinate by u i .
For any mapping
F : R D ! R D , we denote the Jacobian of F at x by @
F ( x ) 2 R D ⇥ D where the ( i, j ) -th entry is@ j F i ( x ) .
We also use @ F ( x )[ u ] and @ 2 F ( x )[ u, v ] to denote the ﬁrst and second order directionalderivatives of F at x along the derivation of u (and v ).
With a slight abuse of notation, we view @ 2 Fas a linear mapping on R D 2 such that @ 2 F ( x ) [ A ] = P Di,j =1 @ 2 F ( x )[ e i , e j ] A ij , for any A 2 R D ⇥ D .
For any submanifold  ⇢ R D and x 2  , we denote by T x (  ) the tangent space of  at x .
We denoteby 1 ⇠ 2 R ⌅ the one-hot vector where the ⇠ -th coordinate is 1 , and 1 denotes the all 1 vector.
3.1 Loss function and manifold of local minimizers
Since this paper deals with scale-invariant loss, most sets in our analysis appear naturally as cones.
Deﬁnition 3.1 (Cone) .
We say K ⇢ R D is a cone if and only if 0 / 2 K and 8 ↵ > 0 , ↵ K ⇢ K .
Assumption 3.2.
The loss function is L ( x ) = 1 ⌅
P ⌅ i = 1 L i ( x ) , where each L i is the loss on the i -thsample and is scale-invariant, that is, L i ( x ) =
L i ( ↵ x ) for all x 2 R D \ { 0 } and ↵ > 0 .
Recall that each  i ( x ) = 1 p ⌅ ( r L i ( x )  r L ( x )) , so the noise function  is also scale-invariant.
Wefurther deﬁne ⌃ ( x ) =  ( x )  ( x ) > , i.e., the noise covariance at x .
With vast overparametrization inmodern machine models, multiple minimizers can exist and form a manifold [5, 6], which motivatesthe following assumption.
Assumption 3.3.
The loss function L : R D \{ 0 } ! R is a C 3 function.
is a ( D  M ) -dimensional
C 2 -submanifold of R D for some integer 0  M  D , where each x 2  is a local minimizer of Land rank( r 2 L ( x ))
=
M for all x 2  .
Since we assumed L ( x ) is scale invariant in Assumption 3.2,we also assume  is a cone.
Next, we formalize the deﬁnition of the gradient ﬂow (GF) with respect to L and its associatedlimiting map, which plays a key role in our analysis.
Deﬁnition 3.4.
For any initialization x 6 = 0 , we deﬁne the GF governed by r L as  ( x, t ) =
x  R t 0 r L (  ( x, s ))
d s for t  0 , and its associated limiting map  ( x ) =
lim t !
1  ( x, t ) wheneverthe limit exists.
Assumption 3.5.
U is a cone and an open neighborhood of  satisfying that GF starting in Uconverges to some point in  , i.e., 8 x 2 U ,  ( x ) 2  .
Li et al.
[27] established several important properties of  by relating the derivatives of  to those of
L , and in particular, we recall the following characterization of @  .
Lemma 3.6 (Lemma 4.2, [27]) .
For any x 2  , @  ( x ) 2 R D ⇥ D is the projection matrix onto thetangent space T x (  ) .
The following properties of L and  are peculiar in our setting thanks to scale-invariance.
Lemma 3.7.
Under Assumption 3.2, for any ↵ > 0 and x 2 R D \ { 0 } , it holds that r k L ( ↵ x ) =
↵  k r k L ( x ) for any k 2 N .
Moreover, for any x 2 U where U is deﬁned in Assumption 3.5, we have@  ( ↵ x )
= @
( x ) and @ 2  ( ↵ x ) =
↵  1 @ 2  ( x ) .
Also, @  ( x ) x
= x for any x 2  .
3.2 Limiting diffusion on the manifold of local minimizers
We recap the notion of Katzenberger processes proposed in [27] and the characterization of thecorresponding limiting diffusion based on Katzenberger’s theorems in [16].
For each n 2 N , let A n : R + ! R + and B n : R + ! R + be two non-decreasing functionswith A n (0) =
B n (0) = 0 , and { Z n ( t ) } t [-≥ 
-] 0 a R ⌅ -valued stochastic process.
Given loss function
L : R D ! R and noise function  : R D !
R D ⇥ ⌅ , we consider the following stochastic process:
[Equation]
Z t 0 r L ( X n ( s )) d A n ( s ) (3) where X n (0) 2 U .
Further deﬁne Y n ( t ) =
X n ( t )   ( X n (0) , A n ( t ))
+  ( X n (0)) .
Theorem 3.8 ([27]) .
Fix some initialization X n (0) ⌘ x 0 2 U for the stochastic process in Equation (3) .
Suppose { X n ( t ) } t  0 is a Katzenberger process satisfying as n ! 1 ,
• A n increases inﬁnitely fast, i.e., 8 ✏ > 0 , inf t [-≥ 
-] 0 ( A n ( t + ✏ )  A n ( t )) !
• B n ( t ) converges to c · t in uniform metric for some constant c .
• Z n converges in distribution to ( I ⌅  1 ⌅ 11 > ) W in Skorokhod metric where W is a ⌅ -
dimension standard Brownian motion (See Deﬁnition of skorokhod metric in Appendix B);
If SDE (4) has a global solution Y with Y (0) =
( x 0 ) , then for any t > 0 , as n ! 1 , Y n ( t )converges in distribution to Y ( t ) .
[Equation]
Warm-up: Simultaneous Limit Case
We ﬁrst consider the setting that ⌘ ,  ! 0 simultaneously with ⌘ = C for some constant C  0 .
Weﬁx x ⌘ ,  (0) =
x 0 for some x 0 .
Now, deﬁne X ⌘ ,  ( t ) = x ⌘ ,  ( b t/ ⌘ 2 c ) , which is roughly equivalent to
SDE (2) with 1 / ⌘ 2 times acceleration.
We rewrite the discrete-time update of x ⌘ ,  as
[Equation]
where A ⌘ ,  , B ⌘ ,  and Z ⌘ ,  are deﬁned by
[Equation]
Note that A ⌘ ,  ( t ) is roughly t/ ⌘ which becomes very large for small ⌘ , thus the negative gradientpart will drive X ⌘ ,  ( t ) rapidly towards the manifold  and force X ⌘ ,  ( t ) to stay close to  afterthat.
On the other hand, as ⌘ ! 0 , B ⌘ ,  ( t ) will converge to t and Z ⌘ ,  will weakly converge to a
Brownian motion, and these terms make up the slow dynamics of SGD.
More precisely, we havethe following lemma summarizing the properties of these integrators, which shows Equation (5) is avalid Katzenberger process (Equation (6) in Li et al. 27).
Lemma 4.1.
Let A ⌘ ,  , B ⌘ ,  and Z ⌘ ,  be as deﬁned in Equation (6) .
Then as ⌘ ,  ! 0 with ⌘ ⌘ C ,we have (1) A ⌘ ,  increases inﬁnitely fast, (2) B ⌘ ,  ( t ) converges to t in uniform metric and (3)
Z ⌘ ,  converges in distribution to ( I ⌅  1 ⌅ 2 11 > ) 1 / 2 W in Skorokhod metric where { W ( t ) } t  0 is the⌅ -dimensional standard Brownian motion.
Therefore, a direct application of Theorem 3.8 yields the following.
Theorem 4.2.
Under Assumption 3.2, 3.3 and 3.5, ﬁx some x 0 2 U and let x ⌘ ,  (0) ⌘ x 0 for all⌘ ,  > 0 for SGD with weight decay (1) .
Consider
[Equation]
where { W ( t ) } t  0 is the standard Brownian motion in R ⌅ .
If SDE (7) has a global solution Y 0 (for
C = 0 ) with Y 0 (0) =
( x 0 ) and Y 0 never leaves U , then for any t > 0 , C  0 , the following holds:
SDE (7) has a global solution Y C in U with Y C (0) =
( x 0 ) .
Moreover, Y C ( t ) d =
Y⇣ e 4 Ct  1 4 C ⌘ e  Ct ;
2. x ⌘ ,  ( b t/ ⌘ 2 c ) converges in distribution to Y C ( t ) as  , ⌘ ! 0 with ⌘ ⌘ C .
Remark 4.3.
Y 0 won’t leave  (and thus U ) if  is a closed manifold.
Remark 4.4.
Theorem 4.2 shows that when there’s no WD, the limiting diffusion is exponentiallyslower than that with WD in the regime of LR ⌘ and WD  converging to zero with ﬁxed ratio.
Limiting diffusion for general case
In the previous section, we showed that the limiting diffusion exists when ⌘ and  go to zero witha ﬁxed ratio.
However, the situation is more complicated in the general case, i.e. , the intrinsic LR ⌘ ! 0 while ⌘ varies and is only upper bounded by some constant.
A concrete example is ⌘ ! 0and  being ﬁxed.
Since the norm has no effect on the loss value but only affects the speed, we need to consider thedynamics of the direction of the weights.
Though we do not have the a closed-form expression for thenorm of the iterates in general, we can approximate it by considering a special case where the trace ofthe noise covariance is constant.
In speciﬁc, recall the canonical SDE approximation in Equation (2).
Li et al.
[26] proved that the dynamics of k X ⌘ ,  ( t ) k 22 is
[Equation]
Suppose tr( ⌃ ( X ))
⌘ 1 / k X k 22 , then d k X ⌘ ,  ( t ) k
22 =  2 ⌘ k X ⌘ ,  ( t ) k 22 d t + ⌘k X ⌘ ,  ( t ) k 22 d t , which 184 admits a closed-form solution:
[Equation]
This implies that the norm of the weights at the equilibrium is of order ( ⌘ /  ) 1 / 4 .
Moreover,
Equation (8) reﬂects the scaling of the norm of the iterates in terms of ⌘ and  as we will see later.
5.1 Continuous-time analysis for SDE
We ﬁrst consider the continuous-time case based on the SDE approximation in Equation (2) since thederivation for this case is cleaner compared to that of the discrete-time case, and the discrete-timeanalysis indeed mimics the continuou-time counterpart.
In a nutshell, Theorem 5.1 shows that thelimiting diffusion exists for SDE with some suitable non-linear time rescaling.
The main result issummarized in the following theorem.
Here the key idea is that we only need to track the direction of the iterates.
To do so, we need tonormalize the iterates properly.
However, when the trace of the noise covariance is not constant, ingeneral it is hard to ﬁnd a close-form solution for k X ⌘ ,  ( t ) k 2 .
Fortunately, the norm dynamics forthe special case in Equation (8) turns out to be a perfect surrogate.
Speciﬁcally, we consider a scalingfunction R ⌘ ,  ( t ) deﬁned as
[Equation]
Next, to rescale the time, we deﬁne ⌧ ⌘ ,  : [0 , 1 ) ! [0 , 1 ) by
[Equation]
where the second equality follows from a direct calculation.
Then we have the following theorem(see Appendix C for the proof).
Theorem 5.1.
Under Assumption 3.2, 3.3 and 3.5, ﬁx some x 0 2 U and let X ⌘ ,  (0) ⌘ x 0for all ⌘ ,  > 0 in the SDE (2) .
Let R ⌘ ,  ( t ) and ⌧ ⌘ ,  ( t ) be as deﬁned in (9) and (10) respectively.
If SDE (11) has a global solution Y with Y (0) =
( x 0 ) and Y never leaves U , then
X ⌘ ,  ( ⌧  1 ⌘ ,  ( T )) /R
⌘ ,  ( ⌧  1 ⌘ ,  ( T )) converges in distribution to Y ( T ) as {+⌘ ! 
+} while simultaneouslysatisfying ⌘ < 2  < c for some constant c ,
[Equation]
where { W ( t ) } t  0 is the standard Brownian motion in R ⌅ and ⌧  1 ⌘ ,  ( T )
=
ln ( 2 ⌘ ( e 2 T  1)+1 ) 4 ⌘ .
Remark 5.2.
The additional constraint ⌘ < 2  when ⌘ ! 0 can be relaxed to ⌘ < 2 C  for anyconstant C > 0 .
It sufﬁces to note that for SDE (2) , X ⌘ ,  with X ⌘ ,  (0) =
x 0 and X C  1 / 2 ⌘ ,C 1 / 2 with X C  1 / 2 ⌘ ,C 1 / 2  (0) =
C  1 / 4 x 0 have the same trajectories, up to a rescaing of C  1 / 4 .
5.2 Discrete-time analysis for SGD
Now, we proceed to analyze SGD+WD by mimicking the continuous-time behavior.
Speciﬁcally,we view R ⌘ ,  ( k ) as an approximation of the norm of x ⌘ ,  ( k ) , and consider the rescaled version of
Equation (1) denoted by ˆ x ⌘ ,  ( k ) = x ⌘ ,  ( k ) /R ⌘ ,  ( k ) .
Further deﬁne b X ⌘ ,  ( t ) =
ˆ x ⌘ ,  ( b t c ) , and it can be shown that (see Appendix C for the derivation)
[Equation]
[Equation]
Next, we introduce the time rescaling through the ⌧ ⌘ ,  ( · ) deﬁned in Equation (10).
Denote ˜ t = ⌧ ⌘ ,  ( t ) ,so t = ⌧  1 ⌘ ,  ( ˜ t ) .
Deﬁne e
X ⌘ ,  ( ˜ t ) =
b X ⌘ ,  ( t ) .
We can further rewrite Equation (12) ase
[Equation]
where A ⌘ ,  , B ⌘ ,  and Z ⌘ ,  are deﬁned by
[Equation]
The convergence of A ⌘ ,  , B ⌘ ,  and Z ⌘ ,  are summarized in the following lemma.
Lemma 5.3.
Let A ⌘ ,  , B ⌘ ,  and Z ⌘ ,  be as deﬁned in Equation (14) , (15) and (16) respectively.
Thenas {+⌘ ! 
+} with ⌘ < 2  < c for some constant c , we have (1) A ⌘ ,  increases inﬁnitely fast, (2) B ⌘ ,  ( t )converges to t 2 in uniform metric and (3) Z ⌘ ,  converges in distribution to ( I ⌅  1 ⌅ 2 11 > ) 1 / 2 W in
Skorokhod metric where { W ( t ) } t  0 is the ⌅ -dimensional standard Brownian motion.
Therefore, e
X ⌘ ,  (˜ t ) is also a valid Katzenberger process, then applying Theorem 3.8 yields:
Theorem 5.4.
Under Assumption 3.2, 3.3 and 3.5, ﬁx some x 0 2 U and let x ⌘ ,  (0) ⌘ x 0 for all⌘ > 0 for SGD+WD (1) .
If SDE (11) has a global solution Y with Y (0) =
( x 0 ) and Y neverleaves U , then for any T > 0 , x ⌘ ,  ( b ⌧  1 ⌘ ,  ( T ) c ) /R ⌘ ,  ( b ⌧  1 ⌘ ,  ( T ) c ) converges in distribution to Y ( T )as {+⌘ ! 
+} while simultaneously satisfying ⌘ < 2  < c for some constant c .
Mixing to Equilibrium
Now we proceed to study the ergodicity of the limiting diffusion (11).
Omitted proofs of this sectionare delayed to Appendix D. Following [26], we make the following assumption.
Assumption 6.1.
For any x 2  with unit norm , tr( ⌃ ( x )) = 1 .
First, we can characterize the norm of the limiting dynamics in the following lemma.
Lemma 6.2.
For SDE (11) , it holds that d k Y ( t ) k 22 = k Y ( t ) k 22 d t + tr( ⌃ ( Y ( t )))
d t .
Note that, by Lemma 6.2, Assumption 6.1 implies that d k Y ( t ) k 22 = ⇣ k
Y ( t ) k 22 + k Y ( t ) k
⌘ d t ,which is deterministic and has a closed-form solution:
[Equation]
Therefore, regardless of the initialization, k Y ( t ) k 2 will always (monotonically) converge to 1 ast !
In this case, we only need to consider the mixing of the direction of Y ( t ) , i.e. Y ( t ) / k Y ( t ) k 2 .
Lemma 6.3.
Under Assumption 6.1, for SDE (11) , deﬁne Y ( t ) =
Y ( t ) / k Y ( t ) k 2 .
Then
[Equation]
To study the ergodicity of the above normalized diffusion process, we need some additional assumptions on  and the noise covariance. For convenience, we deﬁne  r :=
\ { x 2 R D : k x k 2 = r } .
Assumption 6.4.
The loss L is C 1 .
For any r > 0 ,  r is a compact connected C 1 -manifold ofdimension D  M  1 .
Assumption 6.5.
For any x 2  , rank( @  ( x )  ( x ))
=
D  M  1 .
Note x > @  ( x )  ( x ) = 0 , theassumption is equivalent to span { @  ( x )  ( x ) } =
T x (  | x | )
Theorem 6.6.
Under Assumption 6.1, 6.4 and 6.5, starting from any initialization Y (0) 2  1 , thedistribution of Y ( t ) converges to a unique stationary distribution ⇡ on  in total variation distance.
Theorem 6.7 (Mixing of SGD+WD) .
Under the setting of Theorem 5.4 and Theorem 6.6, ﬁx anyx 0 2 U and let x ⌘ ,  (0) ⌘ x 0 for all ⌘ ,  > 0 for SGD + WD in (1) .
Let ⇡ be the unique stationarydistribution on  given in Theorem 6.6.
Let f be any bounded continuous test function.
Then for any✏ > 0 , there exists some T > 0 and  > 0 such that when ⌘ <  with ⌘ < 2  < c , we have
[Equation]
Proof.
Fix any test function f and any ✏ > 0 .
By Theorem 6.6, we know that there is some
T > 0 such that | E [ f ( Y ( T ))]
E x ⇠ ⇡ [ f ( x )] |  ✏ / 2 .
When tr( ⌃ ( x )) ⌘ 1 / k x k 22 , SDE (18) with
Y (0) =
( x 0 ) /
k  ( x 0 ) k 2 has a unique global solution Y ( t ) , and thus SDE (11) has a unique globalsolution given by Y ( t ) {+p 
+} + e  2 t ( k Y (0) k 42  1) .
Hence, for this T , we apply Theorem 5.4 to get
[Equation]
Combining the above yields the desired result.
Experiments
In this section, we ﬁrst empirically verify our theory in a simple setting for, in particular, the timescaling of the dynamics.
We then show that the diffusion process exists during the training of
PreResNet on CIFAR-10, and it has implicit beneﬁt towards better generalization.
7.1 Time Scaling Veriﬁcation
We ﬁrstly verify the time scaling (c.f. Equation (10)) in a simple setting.
We train the followingnormalized linear model by ` 2 regression: F x ( z ) =
D x k x k 2 , z
E , where x 2 R D is the model parameter,and z is the input.
Let { z i } ⌅ i =1 be the input samples, and the target labels y
i
= F x ⇤ ( z i ) for some x ⇤ .
We set ⌅ = D  2 so the solution space S
= { w : h w  x ⇤ k x
⇤ k 2 , z i i = 0 , 8 i 2 [ ⌅ ] } is a 2 -dimensionallinear space.
The manifold of global minimizers  is then equal to [-∪ α> 
 α-][ ↵ > 0 ↵ ( S \ { x | k x k 2 = 1 } ) .
Aslong as x ⇤ is not contained in the linear span of { z i } ⌅ i =1 , S \ { x | k x k 2 = 1 } is a 1 -dimensional [Equation]
manifold (a circle), thus satisfying Assumption 6.4.
During training, we set the loss at step t be L
[Equation]
E , where ⇠ t iid ⇠ N (0 ,  2 I D ) and it impliestrace of noise covariance on unit sphere, thus satisfying Assumption 6.1.
We can also verify M = rank ( r 2 L ( x )) =
8 = D  2 and the noise is non-degenerate, meaning Assumption 6.5 alsoholds.
In our experiments, we choose D = 10 ,  = 0 .
3 , the WD factor  = 0 .
05 , and learning rate ⌘ 2{ 10  2 , 10  3 , 10  4 } .
In Figure 2, we plot the projections of x k x k 2 on the solution space for 60 differentruns with identical initialization for each ⌘ .
For each run, the only differences are the LR ⌘ and thenoise ⇠ t .
The dashed lines in the ﬁgure indicates our time scaling, i.e., 12 ln
[Equation]
Figure 2 shows that time scaling of O ((ln 2  / ⌘ + T ) / ( ⌘ )) ﬁts the dynamics better, compared to
O ( T/ ( ⌘ )) .
7.2 Limiting diffusion on CIFAR-10
Beyond the toy example, we further study the limiting diffusion of PreResNet on CIFAR-10 [20].
We train a 32-layer PreResNet 4 [10] with initial LR ⌘ = 0 . 8 and weight decay factor  = 5 · 10  4 .
Unlike the normalized linear model, it is hard to visualize the model projection of PreResNet on themanifold.
Instead, we choose the test accuracy of  ( x t ) as a test function.
In particular, we decay thelearning rate to 10  3 to approximate gradient ﬂow at different time t , and record the test accuracyafter training 1000 more epochs.
The results are shown in Figure 1 and Figure 3.
We observe that without LR decay, the train accuracy, test accuracy and parameter norm convergequickly after training 100 epochs; the test accuracy of  ( x t ) converges much slower.
It suggeststhat there exists a mixing process after the model reaches the manifold.
Moreover, we observethe test accuracy of  ( x t ) after convergence is signiﬁcantly higher than  ( x t ) at 100th epoch.
Itindicates that the mixing process is beneﬁcial for generalization.
Although it is unclear whether thetest accuracy of  ( x t ) converges at the same rate as the diffusion process, one can always use thetime scaling in (10) as an upperbound to decide the proper time of LR decay.
Conclusion and Future Work
We provide an SDE-based characterization for the limiting dynamics of SGD+WD for a scale invariantloss as ⌘ ! 0 with ⌘ = O (  ) and
= O (1) .
Under some technical assumptions, we further showthat the limiting diffusion converges to a unique stationary distribution.
It leaves as future work torelax the technical assumptions.
Another interesting and important direction for future work is toidentify the beneﬁt on generalization induced by the limiting diffusion.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe 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 to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not 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 instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect 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 paper
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]
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 as a URL)?
[No]
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
[Table]
were chosen)?
[Yes]
[Table]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- [Table] ments multiple times)?
[Yes] [Table] (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[No]
[Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes]
[Table]
(b) Did you mention the license of the assets?
[N/A]
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A]
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A] [Table]