Fast Mixing of Stochastic Gradient Descent with Normalization and Weight Decay
Anonymous Author(s)
AffiliationAddress 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Θ (cid:0) η − 1 λ − 1 ( T + ln( λ/η )) (cid:1) iterations as [-⌘ ! 
-]ηλ → 0 (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 [39], Layer Normalization [2] and Weight Normalization [32] have facilitated the efficient 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 affect of the gradient noise.
To explain the effectiveness of small LRs,
Li et al.
[26] theoretically identified 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 (cid:101) 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 and thus saves training time.
Before stating the fast equilibrium conjecture, let us first 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Ξ (cid:80) Ξ 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 ) ∈ R D denote the parameters after kiterations, x η,λ (0) =
x 0 is the initialization of the parameters, and σ ( x ) =
[ σ 1 ( x ) , . . . , σ Ξ ( x )] [-
-]∈ 44 R D × Ξ is the deterministic noise function with each σ i ( x ) :
= 1 √ Ξ ( ∇ L i ( x ) − ∇ L ( x )) .
Here each ξ kis sampled uniformly from { 1 , . . . , Ξ } , so E ξ k [ σ ξ k ( x )]
= 0 for all x ∈ 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 (cid:101) O
(1 / ( λη )) time.
This convergence rate is much faster than the standard e − Θ(1 /η ) global mixing time of Langevindynamics and that’s how the conjecture gets the name.
Li et al.
[26] proved that the norm of X η,λ ( t )will converge in (cid:101)
O (1 / ( λη )) time but failed to show convergence to the equilibrium.
Another gap isthat the analysis in Li et al.
[26] only works for the SDE approximation in Equation (2), and does notbridge 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 η,λ ∥ x η,λ ∥ 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 [-⌘ ! 
-]ηλ → 0 and T → ∞ , the distribution of ¯ x η,λ
( ⌊ ln( 2 λη ( e 2 T − 1)+1) 4 ηλ ⌋ ) 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 (cid:16) 2 λη ( e 2 T − 1) + 1 (cid:17) 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 [41, 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 simplified 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 isbeneficial 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 efficienttraining of SGD with WD [28].
Scale-invariance also brings about the interesting equivalence betweenthe effect of WD and LR schedules studied in [12, 41, 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, 35, 36],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, 33, 24, 4, 34].
Many interesting properties of SGDhave been discovered through this approach [19, 8, 42, 26, 40, 31, 38].
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 infinitesimal intrinsic LR.
Preliminary
We first clarify the notations.
We denote by N the set of all nonnegative integers and R + the set ofall nonnegative real numbers.
For any k ∈ N , we denote by C k the set of all k times continuouslydifferentiable functions.
For any vector u ∈ 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 ) ∈ 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 first 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 ] = (cid:80) Di,j =1 ∂ 2 F ( x )[ e i , e j ] A ij , for any A ∈ R D × D .
For any submanifold Γ ⊂ R D and x ∈ Γ , we denote by T x (Γ) the tangent space of Γ at x .
We denoteby 1 ξ ∈ 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.
Definition 3.1 (Cone) .
We say K ⊂ R D is a cone if and only if 0 / ∈ K and ∀ α > 0 , αK ⊂ K .
Assumption 3.2.
The loss function is L ( x ) =
1Ξ (cid:80) Ξ 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 ∈ R D \ { 0 } and α > 0 .
Recall that each σ i ( x ) = 1 √ Ξ ( ∇ L i ( x ) − ∇ L ( x )) , so the noise function σ is also scale-invariant.
Wefurther define Σ( 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 ∈ Γ is a local minimizer of Land rank( ∇ 2 L ( x ))
=
M for all x ∈ Γ .
Since we assumed L ( x ) is scale invariant in Assumption 3.2,we also assume Γ is a cone.
Next, we formalize the definition of the gradient flow (GF) with respect to L and its associatedlimiting map, which plays a key role in our analysis.
Definition 3.4.
For any initialization x ̸ = 0 , we define the GF governed by −∇ L as ϕ ( x, t ) =
x − (cid:82) t 0 ∇ L ( ϕ ( x, s ))
d s for t ≥ 0 , and its associated limiting map Φ( x ) =
lim t →∞ ϕ
( 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., ∀ x ∈ U , Φ( x ) ∈ Γ .
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 ∈ Γ , ∂ Φ( x ) ∈ 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 ∈ R D \ { 0 } , it holds that ∇ k L ( αx ) =
α − k ∇ k L ( x ) for any k ∈ N .
Moreover, for any x ∈ U where U is defined in Assumption 3.5, we have∂ Φ( αx ) =
∂
Φ( x ) and ∂ 2 Φ( αx ) =
α − 1 ∂ 2 Φ( x ) .
Also, ∂ Φ( x ) x
= x for any x ∈ Γ .
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 ∈ 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 {+≥ 
+} 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]
[Equation] where X n (0) ∈ U .
Further define 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 ∈ U for the stochastic process in Equation (3) .
Suppose { X n ( t ) } t ≥ 0 is a Katzenberger process satisfying as n → ∞ ,
• A n increases infinitely fast, i e., ∀ ϵ > 0 , inf t {+≥ 
+} ( 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 metricwhere W is a Ξ -
dimension standart Brownian motion;
If SDE (4) has a global solution Y with Y (0) =
Φ( x 0 ) , then for any t > 0 , as n → ∞ , Y n ( t )converges in distribution to Y ( t ) .
[Equation]
Warm-up: Simultaneous Limit Case
We first consider the setting that η, λ → 0 simultaneously with λη = C for some constant C ≥ 0 .
Wefix x η,λ (0) =
x 0 for some x 0 .
Now, define X η,λ ( t ) = x η,λ ( ⌊ t/η 2 ⌋ ) , which is roughly equivalent to
SDE (2) with 1 /η 2 times acceleration.
We rewrite the discrete-time update of x η,λ as
[Equation]
[Equation]
where A η,λ , B η,λ and Z η,λ are defined by
[Equation]
[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 defined in Equation (6) .
Then as η, λ → 0 with λη ≡ C ,we have (1) A η,λ increases infinitely 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, fix some x 0 ∈ 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 0 (cid:16) e 4 Ct − 1 4 C (cid:17) e − Ct ; x η,λ ( ⌊
t/η 2 ⌋ ) 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 fixed ratio.
Limiting diffusion for general case
In the previous section, we showed that the limiting diffusion exists when η and λ go to zero witha fixed ratio.
However, the situation is more complicated in the general case, i.e. , the intrinsic LR ηλ → 0 while ηλ is upper bounded by some constant.
A concrete example is η → 0 and λ beingfixed.
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 specific, recall the canonical SDE approximation in Equation (2).
Li et al.
[26] proved that the dynamics of ∥ X η,λ ( t ) ∥ 22 is
[Equation]
Suppose tr(Σ( X ))
≡ 1 / ∥ X ∥ 22 , then d ∥ X η,λ ( t ) ∥ 22 = − 2 ηλ ∥ X η,λ ( t ) ∥ 22 d [Equation]
+ η 2 ∥ X η,λ ( t ) ∥ 22 d t , whichadmits a closed-form solution:
[Equation]
This implies that the norm of the weights at the equilibrium is of order ( η/λ ) 1 / 4 .
Moreover,
Equation (8) reflects 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 first 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 find a close-form solution for ∥ X η,λ ( t ) ∥ 2 .
Fortunately, the norm dynamics forthe special case in Equation (8) turns out to be a perfect surrogate.
Specifically, we consider a scalingfunction R η,λ ( t ) defined as
[Equation]
Next, to rescale the time, we define τ η,λ : [0 , ∞ ) → [0 , ∞ ) 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, fix some x 0 ∈ U and let X η,λ (0) ≡ x 0for all η, λ > 0 in the SDE (2) .
Let R η,λ ( t ) and τ η,λ ( t ) be as defined 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 [-⌘ ! 
-]ηλ → 0 while simultaneouslysatisfying η < 2 λ < c for some constant c , [Equation]
[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 η < Cλ for anyconstant C > 0 .
It suffices 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 λ =
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.
Specifically,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 define (cid:98) X η,λ ( t ) =
ˆ x η,λ ( ⌊ t ⌋ ) , and itcan be shown that (see Appendix C for the derivation)
[Equation]
Next, we introduce the time rescaling through the τ η,λ ( · ) defined in Equation (10).
Denote ˜ t = τ η,λ ( t ) ,so t = τ − 1 η,λ (˜ t ) .
Define (cid:101)
X η,λ (˜ t ) = (cid:98) X
η,λ ( t ) .
We can further rewrite Equation (12) as
[Equation]
X
[Equation]
(˜
[Equation]
) = (cid:90)
[Equation]
where A η,λ , B η,λ and Z η,λ are defined 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 defined in Equation (14) , (15) and (16) respectively.
Thenas [-⌘ ! 
-]ηλ → 0 with η < 2 λ < c for some constant c , we have (1) A η,λ increases infinitely 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
Skohorod metric where { W ( t ) } t ≥ 0 is the Ξ -dimensional standard Brownian motion.
Therefore, (cid:101)
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, fix some x 0 ∈ 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 η,λ ( ⌊ τ − 1 η,λ ( T ) ⌋ ) /R η,λ ( ⌊ τ − 1 η,λ ( T ) ⌋ ) converges in distribution to Y ( T )as [-⌘ ! 
-]ηλ → 0 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. We need the following assumption.
Assumption 6.1.
For any x ∈ Γ 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 ∥ Y ( t ) ∥ 22 = −∥ Y ( t ) ∥ 22 d t + tr(Σ( Y ( t )))
d t .
Note that, by Lemma 6.2, Assumption 6.1 implies that d ∥ Y ( t ) ∥ 22 = (cid:16) −∥ Y ( t ) ∥
22 + 1 ∥ Y ( t ) ∥(cid:17)
d t ,which is deterministic and has a closed-form solution:
[Equation]
Therefore, regardless of the initialization, ∥ Y ( t ) ∥ 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 ) / ∥ Y ( t ) ∥ 2 .
Lemma 6.3.
Under Assumption 6.1, for SDE (11) , define Y ( t ) =
Y ( t ) / ∥ Y ( t ) ∥ 2 .
Then
[Equation]
[Equation] To study the ergodicity of the above normalized diffusion process, we need some additional assumptions on the manifold Γ and the noise covariance.
Assumption 6.4.
The loss L is C ∞ .
For any r > 0 , Γ r is a compact connected C ∞ -manifold ofdimension D − M − 1 .
Assumption 6.5.
For any x ∈ Γ , 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) ∈ Γ 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, fix anyx 0 ∈ 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 / ∥ x ∥ 22 , SDE (18) with
Y (0) =
Φ( x 0 ) / ∥
Φ( x 0 ) ∥ 2 has a unique global solution Y ( t ) , and thus SDE (11) has a unique globalsolution given by Y ( t ) [-p 
-](cid:112) 1 + e − 2 t ( ∥ Y (0) ∥ 42 − 1) .
Hence, for this T , we apply Theorem 5.4 to get
[Equation]
≤ ϵ/ 2 .
Combining the above yields the desired result.
Experiments
In this section, we first empirically verify our theory in a simple setting, in particular, the time scalingof the dynamics.
We then show that the diffusion process exists during the training of PreResNet on
CIFAR-10, and it has implicit benefit towards better generalization.
7.1 Time Scaling Verification
We firstly verify the time scaling (c.f. Equation (24)) in a simple setting.
We train the followingnormalized linear model by ℓ 2 regression: F x ( z ) =
(cid:68) x
∥ x ∥ 2 , z (cid:69) , where x ∈ 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 : ⟨ w − x ∗ ∥
x ∗ ∥ 2 , z i ⟩ = 0 , ∀ i ∈ [Ξ] } is a 2 -dimensionallinear space.
The manifold of global minimizers Γ is then equal to {+∪ α> 
 α+} ( S ∩ { x | ∥ x ∥ 2 = 1 } ) .
As long as no z i is in the same direction with x ∗ , S ∩ { x | ∥ x ∥ 2 = 1 } is a 1 -dimensional manifold (a circle), thus satisfying Assumption 6.4.
During training, we set the loss at step t be L ( t ) ( x ) =
1Ξ (cid:80) Ξ i =1 ( F x ( z i ) − F x ∗ ( z i ))
2 + (cid:68) x
∥ x ∥ 2 , ξ t (cid:69) , where ξ t iid ∼ N (0 , σ 2 I D ) and it implies trace of noisecovariance on unit sphere, thus satisfying Assumption 6.1.
We can also verify M = rank ( ∇ 2 L ( x )) =
8 = D − 2 and the noise is non-degenerate, 6.5 also holds.
In our experiments, we choose D = 10 , σ = 0 .
3 , the WD factor λ = 0 .
05 , and learning rate η ∈{ 10 − 2 , 10 − 3 , 10 − 4 } .
In Figure 2, we plot the projections of x ∥ x ∥ 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 figure indicates our time scaling, i.e., 12 ln (cid:0) 1 + ( e 4 ηλt − 1)
η 2 λ (cid:1) = T .
Figure 2 shows that time scaling of O ((ln 2 λ/η + T ) / ( λη )) fits the dynamics better, compared to
O ( T/ ( ηλ )) .
7.2 Limiting diffusion on CIFAR-10
Beyond the toy example, we further observe the limit diffusion of PreResNet on CIFAR-10[20].
Wetrain 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 flow 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 suggests thatthere exists a mixing process after the model reaches the manifold.
Moreover, we observe Φ( x t ) afterconvergence is significantly higher than Φ( x t ) at 100th epoch.
It indicates that the mixing process isbeneficial for generalization.
Although it is unclear whether Φ( x t ) converges at the same rate as thediffusion process, one can always use the time scaling in (24) as an upperbound to decide the propertime 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 benefit 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 justification 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 are proprietary.
• 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 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 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 tothem?
[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 were chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[Yes](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]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(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 identifiable
[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]