Fast Mixing of Stochastic Gradient Descent with Normalization and Weight Decay
Abstract
We prove the Fast Equilibrium Conjecture proposed by Li et al.
[30], i.e. , stochastic gradient descent (SGD) on a scale-invariant loss ( e.g. , using networks with various normalization schemes) with learning rate η and weight decay factor λ mixes in function space in ˜ O
( 1 λη ) 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.
[31] 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 Θ
[Equation]
iterations as ηλ → 0 while satisfying η ≤ O ( λ ) ≤ 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 → ∞ .
Introduction
Normalization methods such as Batch Normalization [19], Group Normalization [46], Layer Normalization [4] and Weight Normalization [37] have facilitated the efﬁcient training of modern deep neural networks, 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. , deep neural networks with normalization can be trained with an exponentially increasing learning rate (LR) schedule [29].
There is a line of research [5, 22] suggesting that small LRs (or equivalently large batch sizes) might hurt generalization since they lead to the so-called sharp minima .
It is argued that noise can help avoid sharp minima [16, 22, 26, 20, 14], so large LRs and small batches are preferred as they increase the noise during training.
However, later Li et al.
[30] observed that normalized models trained using SGD+WD can generalize equally well with small LRs, given longer training time after the training loss converges.
Such discrepancies imply an incomplete characterization of what happens during training and especially the effect of the gradient noise.
To explain the effectiveness of small LRs, Li et al.
[30] theoretically identiﬁed an intrinsic learning rate parameter equal to ηλ ( η is the LR
and λ is the WD factor) which governs the convergence rate of the norm of the weights.
Based on empirical evidence, they further proposed the Fast Equilibrium Conjecture which states that the mixing time of the normalized model is roughly #
O (1 / ( ηλ )) , much faster than existing upper bounds.
This conjecture is important because there is some unexplained mechanism towards better generalization happening during the mixing process — the training accuracy/loss plateaus but the test accuracy after LR decay still improves (see Figure 1 for a demonstration).
The fast mixing time can provide theoretical guidelines to decide the proper time of LR decay.
[Table]
[Figure]
Before stating the Fast Equilibrium Conjecture, let us ﬁrst specify the problem setting.
Let x ∈ R D \ { 0 } be the parameters and F x be a neural network parameterized by x , whose out- put is scale-invariant to its parameters x , i.e., F x ( z ) =
F Cx ( z ) for any parameters x , input data z and constant C > 0 .
In other words, the output of the network only depends on the direction of parameter, x , which is a shorthand for x/ % x % 2 .
This scale-invariance property is commonly induced by various normalization schemes, and makes the loss function for any data (or batch of data) also scale-invariant to x .
Let Ξ be the number of training samples and the empirical loss be
L ( x ) = 1 Ξ $ Ξ 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, whose dynamics is formalized as follows:
[Equation]
where η is the LR, λ is the WD factor, x
η , λ ( k ) ∈ R D \{ 0 } denotes the parameters after k iterations, x η , λ (0) =
x init is the initialization of the parameters, and σ ( x ) =
[ σ 1 ( x ) , . . . , σ Ξ ( x )]
∈ R D
×
Ξ is the deterministic noise function with each σ i ( x ) :
= 1 √ Ξ ( ∇ L i ( x ) − ∇ L ( x )) .
Here each ξ k is sampled independently uniformly from { 1 , . . .
, Ξ } , so E ξ k [ σ ξ k ( x )]
= 0 for all x ∈ R D \ { 0 } 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 (Fast Equilibrium Conjecture, Li et al.
Suppose X η , λ ( t ) is a solution of (2), then for any input z , F X η , λ ( t ) ( z ) converges to an equilibrium independent of the initial parameter x init in #
O (1 / ( λη )) time.
We note that the above conjecture is implied by the convergence of distribution of the direction of the parameters, X η , λ ( t ) for scale invariant networks.
This convergence rate is much faster than the standard e − Θ (1 / η ) global mixing time of Langevin dynamics [38] and this is how the conjecture gets its name.
Li et al.
[30] proved that the norm of
X η , λ ( t ) will converge in #
O (1 / ( λη )) time but failed to show convergence to the equilibrium.
More recently, Wang and Wang [43] proposed a spherical SDE model to approximate SGD+WD with constant intrinsic LR.
Using a novel adaption of Simon’s theory, they justiﬁed the Fast Equilibrium
Conjecture by showing that the SGD+WD dynamics consists of three stages: descent ( O (1 / √ λη ) time), diffusion ( O (1 / ( λη ) time) and tunneling ( O ( e C/ ( λη ) ) time).
Although Wang and Wang [43] provided an original and instructive characterization of SGD+WD, their analysis only applies to
Morse loss functions, i.e. , the local minimizers are all isolated, so the diffusion happens in a vanishing neighborhood of some isolated minimizer as the intrinsic LR ηλ → 0 .
As a result, the diffusion process shall bring no generalization beneﬁt and cannot explain the improvement of ﬁnal generalization if staying at training loss plateau for a longer time (see Figure 1).
Except the undesired implications, the assumption that the loss function is Morse itself is probably too strong and is against the empirical evidence that the level sets of deep learning loss are connected [10].
Another weakness of existing analyses in Li et al.
[30], Wang and Wang [43] is that their analyses only work for the SDE approximation (2), and do not apply to the discrete-time dynamics (1).
In this paper, we provide a proof to Conjecture 1.1 under more reasonable assumptions.
That is, We allow the local minimizers form a connected manifold, which can be viewed a generalization of the
Morse function assumption in Wang and Wang [43], as an isolated minimizer is just a manifold of dimension 0 .
The main result is stated informally below.
Theorem 1.2 (Informal version of Theorem 6.7) .
Suppose Γ is a connected analytic manifold consisting only of local minimizers of L and that Γ 1 :=
{ x ∈ Γ | % x % 2 = 1 } is compact.
Under some regularity assumptions, there is an open neighborhood U of Γ , such that for every x init ∈
U and
T > 0 , as ηλ → 0 , both ¯ x η , λ ( ⌊ ln( 2 λη ( e 2 T − 1)+1) 4 ηλ ⌋ ) and X η , λ
( ln( 2 λη ( e 2 T − 1)+1) 4 ηλ ) converge in distribution to the same distribution, denoted by µ T,x init .
Moreover, for every x init ∈ U , µ T,x init weakly converges to the same equilibrium distribution π supported on Γ 1 as T
[Equation]
∞ .
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 a scale-invariant stochastic loss in the limit of ηλ → 0 with η = O ( λ ) and λ = O (1) .
This result is not a direct application of the framework proposed by Li et al.
[31] as there is no minimizer in the usual sense, not to mention manifold of local minimizers.
The key proof technique is to introduce a novel time-rescaling to make the convergence result in [21] applicable.
We show that SGD without WD for stochastic scale-invariant loss has the same limiting dynamics as SGD+WD, but is exponentially slower (see Theorem 4.2).
This is consistent to the empirical observation [48, 30] that turning on WD for SGD for scale-invariant loss helps generalization.
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 limiting diffusion of SGD+WD 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 [30].
Though our convergence result is asymptotic, we verify in simpliﬁed settings that the predicted phenomena 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 is beneﬁcial for generalization.
Related Work
Scale-invariance.
For scale-invariant loss functions, existing works have proved the convergence of GD and SGD with WD [2].
It has been shown that scale-invariance enables robust and efﬁcient training of SGD with WD [32].
Scale invariance also brings about the interesting equivalence between the effect of WD and LR schedules studied in Hoffer et al.
[17], Zhang et al.
[48], Li and
Arora [29].
Moreover, for SGD with WD, η being the LR and λ being the WD factor, the norm of the parameters will converge to ( λ / η ) 1 / 4 [30, 41, 42], and this induces the intrinsic LR which is equal to ηλ [30].
These observations are crucial to our derivations in the following sections.
SDE approximation.
Continuous-time tools such as SDE have been a popular lens for studying optimization algorithms including SGD [33, 27, 39, 28, 6, 40].
Many interesting properties of SGD have been discovered through this approach [24, 13, 49, 30, 47, 35, 45].
Recently, Li et al.
[31] proposed a framework to study the implicit bias of SGD by showing that with a proper time-rescaling, the iterates converge in distribution to 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
Notations.
We denote by N the set of all nonnegative integers and R + the set of all nonnegative real numbers.
For any k ∈ N , we denote by C k the set of all k times continuously differentiable 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 ﬁrst and second order directional derivatives of
F at x along the derivation of u (and v ).
With a slight abuse of notation, we view ∂ 2 F as a linear mapping on R D 2 such that ∂ 2 F ( x )[ A ] = $
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 denote byξ ∈ R Ξ the one-hot vector where the ξ -th coordinate is 1 , and 1 denotes the all 1 vector.
We say
K ⊂ R D is a cone if and only if 0 / ∈ K and ∀ α > 0 , α K ⊂ K .
Recall that the training loss L : R D \ { 0 } → R is deﬁned as L = 1 Ξ $ Ξ i =1 L i where L i is the loss on the i -th sample.
With vast overparametrization in modern machine learning models, multiple minimizers can exist and form a manifold [8, 7].
Thus following Fehrman et al.
[12], Li et al.
[31], Arora et al.
[3], we make the assumption below throughout this paper.
Assumption 3.1.
Each loss function L i : R D \ { 0 } → R is a C 4 function.
Γ is a ( D − M ) - dimensional
C 2 -submanifold of R D for some integer 0 ≤ M ≤ D − 1 , where each x ∈ Γ is a local minimizer of L and rank( ∇ 2 L ( x )) =
M for all x ∈ Γ .
Note that ∇ 2 L ( x ) must have zero eigenvalues in the tangent space of Γ at x , we are indeed assuming
Hessian attains the maximal rank everywhere on mainfold.
In this paper, we are interested in the behavior of SGD+WD with each L i 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 3.3.
Deﬁnition 3.2 (Homogeneous Functions) .
We say a function f : R D \ { 0 } → R m is a khomogeneous for some k ∈ R if and only if for all x ∈ R D \ { 0 } and α > 0 , f ( α x ) =
α k f ( x ) .
Speciﬁcally, we say a function f is scale invariant if and only if it is 0 -homogeneous.
Assumption 3.3.
L i is scale invariant for each 1 ≤ i ≤ Ξ and Γ is a cone.
Recall that σ i ( x ) = 1 √ Ξ ( ∇ L i ( x ) − ∇ L ( x )) , so the noise function σ is ( − 1) -homogeneous and thus the noise covariance Σ ( x ) = σ ( x ) σ ( x ) ⊤ is ( − 2)
-homogeneous.
Below are two useful properties of homogeneous functions, whose proofs follow from applying direct chain rules directly.
Lemma 3.4.
For any l ∈ N and k -homogeneous function f , ∇ l f is ( k − l ) -homogeneous.
Lemma 3.5 (Euler’s Theorem for Homogeneous Functions) .
For any real-valued k -homogeneous function f , 〈 x, ∇ f ( x ) 〉
= kf ( x ) .
Next, the following limiting map of gradient ﬂow plays a key role in our analysis.
Deﬁnition 3.6.
For any x ∕ = 0 , we deﬁne the gradient ﬂow governed by −∇ L as the unique solution of φ ( x, t ) =
x − % t 0 ∇ L ( φ ( x, s ))
d s for t ≥ 0 and its associated limiting map by Φ ( x ) =
lim t →∞ φ ( x, t ) whenever the limit exists.
Throughout the paper, we use U to denote the attraction set of Γ under gradient ﬂow, that is, U = { x ∈ R D | Φ ( x ) is well-deﬁned and Φ ( x ) ∈ Γ } .
Lemma 3.7 (Lemma B.15 of Arora et al.
U is open and Φ is C 2 on U .
Lemma 3.8.
Under Assumption 3.3, U is a cone and Φ is 1 -homogeneous in U .
Li et al.
[31] 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.9 (Lemma 4.2, [31]) .
For any x ∈ Γ , ∂ Φ ( x ) ∈ R D
× D is the orthogonal projection matrix onto the tangent space
T x ( Γ ) .
As a consequence, for any x ∈ Γ , Assumption 3.3 = ⇒ x ∈ T x ( Γ ) =
⇒ Φ ( x ) x
= x .
3.1 Limiting diffusion on the manifold of local minimizers
We recap the notion of Katzenberger processes proposed in Li et al.
[31] and the characterization of the corresponding limiting diffusion based on Katzenberger’s theorems in Katzenberger [21].
In this subsection we only assumeAssumption 3.1, but not Assumption 3.3.
For each n ∈ N , let A n : R + → R + and B n : R + → R + be two non-decreasing functions with A n (0)
=
B n (0) = 0 , and { Z n ( t ) } t ≥ 0 be a R Ξ -valued stochastic process.
In our context of
SGD+WD, given loss function L : R D → R , noise function σ : R D → R D × Ξ and initialization x init ∈ U , we call the following stochastic process (3) a Katzenberger process if as n → ∞ ,
A n increases inﬁnitely fast, i.e., ∀ ) > 0 , inf t ≥
( A n ( t + ) ) − A n ( t ))
→ ∞ ; 2.
B n ( t ) converges to c · t in uniform metric for some constant c . 3.
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);
[Equation]
Theorem 3.10 (Theorem 4.6, Li et al.
Given a Katzenberger process { X n ( · ) } n ∈ N , if SDE (4) has a global solution Y in U with Y (0) =
Φ ( x init ) , then for any t > 0 , X n ( t ) converges in distribution to Y ( t ) as n → ∞ .
[Equation]
We note that the global solution always exists if the manifold Γ is compact.
Our notion of Katzenberger process and theorem statement is slightly more general than those in Li et al.
[31] to handle the weight decay.
However, our formulation is still under the original framework of Katzenberger [21] and the proof in Li et al.
[31] can be easily adapted to Theorem 3.10.
Warm-up: Simultaneous Limit Case
As a warm-up, we ﬁrst consider the setting where η , λ → 0 simultaneously with λη ≡ C for some constant C ≥ 0 .
In this special regime, we do not need to use the scale invariance property of the loss, and we can directly apply Theorem 3.10 to obtain the limiting diffusion of SGD+WD.
Nonetheless, we will see the beneﬁt of weight decay as a source of acceleration.
While for the general case of ηλ → 0 that will be considered in Section 5, 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 a Katzenberger process.
Now, recall the SGD+WD updates in Equation (1), and let us ﬁx x η , λ (0) =
x init for some x init ∈ U .
Deﬁne ˇ X η , λ ( t ) =
x η , λ ( ⌊ t/ η 2 ⌋ ) , which is roughly equivalent to SDE (2) with 1 / η 2 times acceleration, and we can rewrite the discrete-time update of x η , λ as ˇ X η , λ (
[Equation]
) =
[Equation]
init +
[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 gradient part will drive ˇ X η , λ ( t ) rapidly towards the manifold Γ and force ˇ X η , λ ( t ) to stay close to Γ after that.
On the other hand, as η → 0 , B η , λ ( t ) will converge to Ct and Z η , λ 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 (5) is a valid Katzenberger process.
Lemma 4.1.
Let A η , λ , B η , λ and Z η , λ be as deﬁned in Equation (6) .
Then as η , λ → 0 with λη ≡ C , it holds that (1) A η , λ increases inﬁnitely fast, (2) B η , λ ( t ) converges to Ct in uniform metric and (3)
Z η , λ converges in distribution to ( I Ξ − 1 Ξ 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.10 yields the limiting diffusion in this case.
Theorem 4.2.
Under Assumption 3.1, let x η , λ (0) ≡ x init ∈ U , ∀ η , λ > 0 in SGD+WD (1) .
Consider
[Equation]
where { W ( t ) } t ≥ 0 is the standard Brownian motion in R Ξ .
If SDE (7) has a global solution Y C in U for some C ≥ 0 with Y C (0) =
Φ ( x init ) , then the following holds:
1. x η , λ ( ⌊ t/ η 2 ⌋ ) converges in distribution to Y C ( t ) as λ , η → 0 with λη ≡ C ; 2.
Under Assumption 3.3, for every C ′ ≥ 0 SDE (7) with any C ′ ≥ 0 has a global solution Y C ′ in U with Y C ′ (0) =
Φ ( x init ) .
Moreover, Y C ′ ( t ) d = Y
[Equation]
Remark 4.3.
Theorem 4.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 η and WD λ going to zero with a ﬁxed ratio.
Limiting Diffusion for General Cases
In the previous section, we showed that the limiting diffusion exists when η and λ go to zero with a ﬁxed ratio.
However, the situation is more complicated in the general case where we drop the assumption of η / λ being ﬁxed.
Below we ﬁrst 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 case when ηλ to 0 with η = O ( λ ) =
O (1) .
Challenges for the General Case.
A concrete example for the challenge is when η → 0 and λ be ﬁxed as a constant, which is also the most natural and practical setting.
We quickly ﬁnd ourselves in a dilemma if we still want to apply Katzenberger’s theorem [21], or its simpliﬁed version Theorem 3.10: If we view WD or ℓ 2 regularization as the ‘fast’ past of the dynamics, that is, a part of the loss function, then there is no minimizer for the ℓ 2 regularized scale-invariant loss and thus it doesn’t satisfy the condition of Katzenberger’s theorem; if we view WD as some ‘slow’ dynamics and formulate it as λη X η , λ d B η , λ ( t ) in Equation (5), then unlike the simultaneous limit case, λη doesn’t necessarily have a limit, and thus the condition of Katzenberger’s theorem is again not met.
The above dilemma reﬂects 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 executes the ℓ 2 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 ) % 22 ∝ , ηλ ( e.g. Figure 1b), 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 L , Equation (1) can be rewritten as the following form,
Equation (8), with ˜ η = √ ηλ , ˜ λ = √ ηλ , ˜ x ˜ η , ˜ λ
= ( λη ) 1 / 4 x η , λ .
[Equation]
With such a rescaling, we successfully make the norm of parameters in constant scale, that is, % ˜ x ˜ η , ˜
λ
% 22 = - λη % x η , λ %
22 =
Θ (1) and thus we can apply Katzenberger’s theorem.
Note that we cannot do this in the early phase of SGD+WD as we start from a ﬁxed initialization and such a 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 (10)) and time rescaling τ η , λ (Equation (11)) which interpolate the early and late regime smoothly.
Because the rescalings are adaptive to the norm of the parameter along the training trajectory, they allow us to apply Katzenberger’s theorem on the rescaled dynamics X η , λ ( τ − 1 η , λ ( t ))
R η , λ ( τ − 1 η , λ ( t )) .
(See formal statements in Theorem 5.1).
Compared with the ordinary SGD without WD studied in Li et al.
[31] where the time rescaling is set to be a ﬁxed acceleration by 1 η 2 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 to normalize the iterates properly.
However, when the trace of the noise covariance is not constant, in general it is hard to ﬁnd a close-form solution for % X η , λ ( t ) % 2 , but we can approximate it using the former special case.
In speciﬁc, recall the canonical SDE approximation in (2).
Li et al.
[30] proved that the dynamics of % X η , λ ( t ) % 22 is
[Equation]
Suppose tr( Σ ( x )) ≡ 1 / % x % 22 , then the above further simpliﬁes into d % X η , λ ( t ) %= − 2 ηλ % X η , λ ( t ) % 22 d t +
[Equation]
λ
[Equation]
t )
[Equation]
22 d t , which admits a closed-form solution:
[Equation]
This implies that the norm of the weights at the equilibrium is of order ( η / λ ) 1 / 4 .
Moreover, Equation (9) 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, i.e. , SDE approximation (2) since the derivation for this case is cleaner compared to that of the discrete-time case, and the discrete-time analysis indeed mimics the continuou-time counterpart.
The main result is summarized in Theorem 5.1, which shows that the limiting diffusion exists for SDE with some suitable non-linear time rescaling.
As mentioned in the previous discussion, we consider a scaling function R η , λ ( t ) inspired by the norm dynamics for the special case in Equation (9):
[Equation]
Next, to rescale the time, we deﬁne τ η , λ : [0 , ∞ ) → [0 , ∞ ) by
[Equation] [Equation]
where the second equality follows from a direct calculation.
It is easy to show that τ − 1 η , λ ( t ) =
ln ( 2 λη ( e 2 T − 1)+1 ) 4 ηλ .
Then we have the following theorem (see Appendix D for the proof), which says the rescaled version of SDE approximation Equation (2) admits the following limiting diffusion Equation (12), where { W ( t ) } t ≥ 0 is the standard Brownian motion in R Ξ .
[Equation]
Theorem 5.1.
Under Assumption 3.1 and 3.3, let X η , λ (0) ≡ x init ∈
U for all η , λ > 0 in SDE (2) .
Let R η , λ ( t ) and τ η , λ ( t ) be deﬁned in (10) and (11) .
If SDE (12) has a global solution Y with
Y (0) =
Φ ( x init ) in U , then X η , λ ( τ − 1 η , λ ( T )) /R
η , λ ( τ − 1 η , λ ( T )) converges in distribution to Y ( T ) as ηλ → 0 with η < 2 λ < c for some constant c .
Remark 5.2.
The additional constraint η < 2 λ when ηλ → 0 can be relaxed to η < 2 C λ for any constant C > 0 .
It sufﬁces to note that for SDE (2) , X η , λ with X η , λ (0) =
x init and X C − 1 / 2 η ,C 1 / 2 λ with X C − 1 / 2 η ,C 1 / 2 λ (0) =
C − 1 / 4 x init have the same trajectories, up to a rescaing of C − 1 / 4 .
This is equivalent to replace η 2 λ by η 2 C λ in (10) .
5.2 Discrete-time Analysis for SGD+WD
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
X
[Equation] [Equation]
t ) =
[Equation]
x
[Equation]
, λ ( ⌊ t ⌋ ) , and it can be shown that (see Appendix D for the derivation)
[Equation] [Equation]
Next, we introduce the time rescaling through the τ η , λ ( · ) deﬁned in Equation (11).
Denote ˜ t = τ η , λ ( t ) , so t = τ − 1 η , λ ( ˜ t ) .
Deﬁne #
X η , λ ( ˜ t ) =
X η , λ ( t ) .
We can further rewrite Equation (13) as #
[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 (15) , (16) and (17) respectively.
Then as ηλ → 0 with η < 2 λ < c for some constant c , the following holds:
• A η , λ increases inﬁnitely fast; • B η , λ ( t ) converges to t 2 in uniform metric; • Z η , λ converges in distribution to ( I Ξ − 1 Ξ 11 ⊤ ) 1 / 2 W in Skorokhod metric where { W ( t ) } t ≥ 0 is the Ξ -dimensional standard Brownian motion.
Therefore, #
X η , λ (˜ t ) is also a valid Katzenberger process, then applying Theorem 3.10 yields:
Theorem 5.4.
Under Assumption 3.1 and 3.3, let x η , λ (0) ≡ x init ∈ U for all η , λ > 0 in
SGD+WD (1) .
Let R η , λ ( t ) and τ η , λ ( t ) be deﬁned in (10) and (11) .
If SDE (12) has a global solution
Y in U with Y (0) =
Φ ( x init ) , then for any T > 0 , x η , λ ( ⌊ τ − 1 η , λ ( T ) ⌋ ) /R
η , λ ( ⌊ τ − 1 η , λ ( T ) ⌋ ) converges in distribution to Y ( T ) as ηλ → 0 with η < 2 λ < c
some constant c .
Mixing to Equilibrium
Now we proceed to study the ergodicity of the limiting diffusion (12).
Omitted proofs of this section are delayed to Appendix E.First, we can characterize the norm of the limiting dynamics in the following lemma.
Lemma 6.1.
For SDE (12) , it holds that
[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 ∈ R D : % x % 2 = r } .
To ensure the existence of stationary distribution of the limiting diffusion process, we assume Γ 1 is compact.
We also need to assume that Γ 1 are connected (so is Γ ) for the uniqueness of the equilibrium distribution.
Assumption 6.2.
Γ 1 is compact and connected.
Assumption 6.3 (Controllability) .
For each x ∈ Γ , span ( { ∂ Φ ( x ) σ i ( x ) } Ξ i =1 ) =
T x ( Γ , x , 2 ) .
Deﬁnition 6.4.
Let V ⊆
R D be any open set, a function f : V → R is analytic if for each x ∈ V , there is a neighborhood V x of x such that the Taylor series of f expanded at x converges to f everywhere in V x .
A vector-valued function is analytic if each of its coordinate functions is analytic.
Assumption 6.5.
tr( Σ ( · )) is an analytic function on R D \ { 0 } and Γ is an analytic manifold.
Now we are ready to state our main result in this section, which is Theorem 6.6.
It is proved in two cases respectively in Appendices E.2 and E.3, depending on whether the trace of gradient covariance tr[ Σ ] is constant on Γ 1 .
If it is, then the diffusion process essentially is on a ( D − M − 1) -dimensional manifold (after suitable rescaling), Γ 1 , just as the analysis by Wang and Wang [43].
If not, then the situation becomes more complicated and the diffusion process is on Γ , a ( D − M ) dimensional manifold.
In this case we will need the analyticity assumption, Assumption 6.5.
Theorem 6.6.
Under Assumption 3.1, 3.3, 6.2, 6.3 and 6.5, starting from any initialization Y (0) ∈ Γ , the distribution of Y ( t ) converges to a unique stationary distribution π on Γ 1 in total variation distance.
Our main result on the fast mixing of SGD+WD and its SDE approximation follows from a direct combination between the convergence of the SGD+WD iterates proved in Theorem 5.4 and Theorem 6.6.
Here we note that since Γ 1 is a compact domain, convergence in total variation implies convergence in distribution.
Theorem 6.7 (Fast Mixing) .
Under Assumption 3.1, 3.3, 6.2, 6.3 and 6.5, let x η , λ (0) ≡ X η , λ (0)
≡ x init ∈ U for all η , λ > 0 for SGD + WD (1) and SDE approximation (2) .
For any T > 0 , as ηλ → 0 with η = O ( λ ) = O (1) , both ¯ x η , λ ( ⌊ ln( 2 λη ( e 2 T − 1)+1) 4 ηλ ⌋ ) and X η , λ
( ln( 2 λη ( e 2 T − 1)+1) 4 ηλ ) converge in distribution to the same distribution, denoted by µ T,x init .
Moreover, for every x init ∈ U , µ T,x init weakly converges to the same equilibrium distribution π supported on Γ 1 as T → ∞ .
Experiments
In this section, we ﬁrst empirically verify our theory in a simple setting for, in particular, the time scaling 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 (11)) in a simple setting.
We train the following normalized linear model by ℓ 2 regression: F x ( z ) =
〈 x , x , 2 , z 〉 , 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 -dimensional linear space.
The manifold of global minimizers Γ is then given by Γ = α >
α ( S ∩ { x | % x % 2 = 1 } ) .
As long as x ∗ is not contained in the linear span of { z i } [Equation] S ∩ { x | % x % 2 = 1 } is a 1 -dimensional manifold (a circle), thus satisfying Assumption 6.2.
During training, we set the loss at step t to be
[Equation]
where ξ t iid ∼ N (0 , σ 2 I D ) and it implies that the trace of noise covariance is constant on the unit sphere, thus satisfying Assumption E.2.
We can also verify M = rank ( ∇ 2 L ( x ))
=
8 = D − 2 and the noise is non-degenerate, meaning that Assumption 6.3 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
[Equation]
for 60 different runs with identical initialization for each η . run
[Equation]
=
[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-
Beyond the toy example, we further study the limiting diffusion of PreResNet on CIFAR-10 [25].
We train a 32-layer PreResNet 2 [15] 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 the manifold.
Instead, we choose the test accuracy of Φ ( x t ) as a test function.
In particular, we decay the learning rate to 10 − 3 to approximate gradient ﬂow at different time t , and record the test accuracy after 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 converge quickly after training 100 epochs; the test accuracy of Φ ( x t ) converges much slower.
It suggests that there exists a mixing process after the model reaches the manifold.
Moreover, we observe the test accuracy of Φ ( x t ) after convergence is signiﬁcantly higher than Φ ( x t ) at 100th epoch.
It indicates that the mixing process is beneﬁcial for generalization.
Our time scaling (11) suggests that the optimal step of LR decay grows no faster than ˜ Ω
(1 / ηλ ) as η → 0 .
Unfortunately, (11) alone is not sufﬁcient for deciding the optimal step for decaying LR as the mixing time T for the continuous dynamics is unknown.
A potential usage of time scaling (11) is to ﬁrst estimate T via another SGD run with a larger learning rate, which we leave for future work.
Conclusion and Future Work
We provide an SDE-based characterization for the limiting dynamics of SGD+WD for a scale invariant loss as ηλ → 0 with η = O ( λ ) and λ
= O (1) .
Under some technical assumptions, we further show that the limiting diffusion converges to a unique stationary distribution.
It leaves as future work to relax the technical assumptions.
Another interesting and important direction for future work is to identify the beneﬁt on generalization induced by the limiting diffusion.
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 a justiﬁcation to 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 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 instructions block 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] 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 experimental results (either in the supplemental material or as a 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 experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[No] 4.
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] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁ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, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to 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]
[Figure]
B Preliminaries on Stochastic Processes
In this section, we give a brief summary for the basic concepts related to stochastic processes and weak convergence that will be used in this paper.
B.1 Càdlàg Function and Metric
Deﬁnition B.1 (Càdlàg function) .
Let T ∈ [0 , ∞ ] .
A function g : [0 , T ) → E is càdlàg if for all t ∈ [0 , T ) it is right-continuous at t and its left limit g ( t − ) exists.
Let D E [0 , T ) be the set of all càdlàg function mapping [0 , T ) into E .
We also use D E [0 , T ) to denote the set of all continuous function mapping [0 , T ) into E .
By deﬁnition, C E [0 , T ) ⊂ D E [0 , T ) .
Deﬁnition B.2 (Skorokhod metric on D E [0 , ∞ ) ) .
For each ﬁnite T > 0 and each pair of functions f, g ∈ D E [0 , ∞ ) , deﬁne d T ( f, g ) as the inﬁmum of all those values of δ for which there exist grids≤ t 0 < t 1 < · · · < t m and 0 < s 0 < s 1 < · · · < · · · < s m , with t k , s k ≥ T , such that | t i − s i | ≤ δ for i = 0 , . . .
, k , and
[Equation]
if
[Equation]
t
[Equation] [Equation]
t
[Equation] [Equation]
i +
[Equation]
for i = 0 , . . .
, k − 1 .
The Skorokhod metric on D E [0 , ∞ ) is deﬁned to be
[Equation]
B.2 Stochastic Processes and Stochastic Integral
Let ( Ω , F , {F t } t ≥ 0 , P ) be a ﬁltered probability space.
Deﬁnition B.3 (Martingale) .
Let { X ( t ) } t ≥ 0 be a {F t } t ≥ 0 -adapted stochastic process.
If for all≤ s ≤ t , it holds that
[Equation]
then X is called a martingale.
Lemma B.4 (Itô’s Lemma) .
Let { X ( t ) } t ≥ 0 be deﬁned through the following Itô drift-diffusion process:
[Equation]
where { W ( t ) } t ≥ 0 is the standard Brownian motion.
Then for any twice differentiable function f , it holds that
[Equation]
B.3 Weak Convergence for Stochastic Processes
Let ( D E [0 , ∞ ) , A , d ) be a metric space equipped with a σ -algebra A and the Skorokhod metric deﬁned in the previous subsection.
Let { X n } n ≥ 0 be a sequence of stochastic processes on a sequence of probability spaces { ( Ω n , F n , P n ) } n ≥ 0 such that each X n has sample paths in D E [0 , ∞ ) .
Also, let X be a stochas- tic process on ( Ω ,
F , P ) with sample paths on D E [0 , ∞ ) .
Deﬁnition B.5 (Weak convergence) .
A sequence of stochastic processes { X n } n ≥ 0 is said to con- verge in distribution or weakly converge to X (written as X n ⇒ X ) if and only if for all A - measurable, bounded, and continuous function f :
D E [0 , ∞ )
→ R , it holds that
[Equation]
The weak convergence of random variables can be deﬁned in a very similar way.
Note though we only deﬁne weak convergence for a countable sequence of stochastic processes/ random variables, in this paper we extend the notion of convergence with real number as indices, e.g. , lim η , λ → 0 where η ≤ 2 λ ≤ C for some constant C .
Such convergence should be interpreted as lim n →∞ for any sequence of { η n } , { λ n }
, such that η n λ n
→
, with
η n ≤ 2 λ n ≤ C .
The main convergence result in Katzenberger [21] are in the sense of Skorokhod metric in Deﬁnition B.2, which is harder to understand and use compared to the more common uniform metric (Deﬁnition B.6).
However, convergence in Skorokhod metric and uniform metric indeed coincide with each other when the limit is in C R D [0 , ∞ ) , i.e., the continuous functions.
Deﬁnition B.6 (Uniform metric on D E [0 , ∞ ) ) .
For each ﬁnite T > 0 and each pair of functions f, g ∈ D E [0 , T ) , the uniform metric is deﬁned to be
[Equation]
The uniform metric on D E [0 , ∞ ) is deﬁned to be
[Equation]
Lemma B.7 (Problem 7, Section 5, Pollard [36]) .
If X n ⇒ X in the Skorokhod sense, and X has sample paths in C R D [0 , ∞ ) , then X n ⇒ X in the uniform metric.
Remark B.8.
We shall note convergence under the uniform metric deﬁned above is weaker than convergence under the following metric sup t ∈ [0 , ∞ )
% f ( t ) − g ( t ) % .
Convergence in the uniform metric on [0 , ∞ ] deﬁned in Deﬁnition B.6 is equivalent to convergence in the uniform metric on each compact set [0 , T ] for T ∈ N + .
The same holds for the Skorokhod topology.
B.4 Omitted proofs for properties induced by scale-invariance
We ﬁrst prove that the attraction set U is an open neighborhood of Γ and also a cone.
Proof of Lemma 3.7.
The ﬁrst claim is due to Lemma B.15 of Arora et al.
Next, further assume that the loss L is scale invariant, then for any x ∈ U , we have Φ ( x ) =
lim t →∞ φ ( x, t ) ∈ Γ .
Deﬁne x ( t ) =
φ ( x, t ) , and for any α > 0 , let ˜ x
( t ) =
α · x ( t/ α 2 ) , then we have ˜ x (0) =
α x and d˜ x ( t ) d t = − α − 1 ∇ L ( x ( t/ α 2 ))
=
−∇ L (˜ x ( t )) where the second equality follows from the scale invariance of L .
Therefore, by the deﬁnition of φ ( x, t ) , we see that φ ( α x, t ) =
˜ x ( t ) =
α · x ( t/ α 2 ) =
α · φ ( x, t/ α 2 ) .
Let t → ∞ , and we get Φ ( α x )
=
α · Φ ( x ) .
Since Γ itself is a cone, it follows that Φ ( α x ) ∈ Γ , and thus α x ∈ U .
Hence, we conclude that U is also a cone.
Now we present the proofs for the results in Section 4 where λ , η → 0 simultaneously, that is, ηλ ≡ C where C is some positive constant.
This case can be regarded as a direct application of the framework in Li et al.
Proof of Lemma 4.1.
The proof can be directly adapted from that of Lemma 4.2 in Li et al.
Proof of Theorem 4.2.
The ﬁrst claim is a direct consequence of Lemma 4.1 and Theorem 3.10.
For the second claim, we ﬁrst show that if Y 0 exists, then for any C > 0 , Y C also exists and they are equal in distribution after the desired rescaling.
By Lemma F.1, there is some Brownian motion
W ′ ( t ) such that W ( e 4 Ct − 1 4 C )
=
% t τ =
e 2 Ct d W ′ ( t ) , ∀ t ≥ 0 almost surely.
For convenience, we denote Y 0 ( e 4 Ct − 1 4 C )
e − Ct by #
Y C ( t ) .
Thus we have
[Equation]
where the last equality uses the fact that ∂ k
Φ is (1 − k ) -homogeneous, σ is ( − 1) -homogeneous
and Σ is ( − 2) -homogeneous by Lemma 3.4.
Thus we construct a global solution for (7) with hyperparameter C .
Since U is a cone and Y 0 ( t ) ∈ U for all t ≥ 0 , it also holds that
[Equation]
By applying the above argument in the other direction, we can show that if Y C exists then Y 0 exists and they are equivalent up to some rescaling.
This ﬁnishes the proof.
D Omitted proofs for general case in Section
Here we provide the proofs for the limiting diffusion for continuous-time SDE and discrete
SGD+WD respectively, in a general case of ηλ → 0 where η ≤ 2 λ ≤ c and c is any constant.
Note the analysis here can be easily generalized to the case ηλ → 0 where η = O ( λ )
=
O (1) , that is, exists constants C 1 , C 2 , such that η ≤ C 1 λ ≤ C 2 .
Such a limit is more interesting than the simultaneous limit as the it is more common to tune LR without ﬁxing the ratio between WD and
LR.
Below we will show how these rescalings work ﬁrst in the continuous case in Appendix D.1, and then generalize the same proof idea to the discrete case, SGD+WD, in Appendix D.2.
D.1 Proof for the limiting diffusion of SDE
Proof of Theorem 5.1.
Recall the scaling function R η , λ ( t ) deﬁned in Equation (10):
[Equation]
which is exactly Equation (9) except % X (0) % 2 is replaced with 1 .
A simple calculation gives
[Equation]
Now, deﬁne
X
[Equation] [Equation]
) =
X
[Equation] [Equation]
t ) /R
[Equation]
, λ ( t ) , a normalized version of X η , λ ( t ) .
Its dynamics is
[Equation]
Next, to rescale the time, recall the function τ η , λ deﬁned in Equation (11):
[Equation]
and consider #
[Equation]
t ) .
Note that d R η , λ ( t ) 4 = − 4 ηλ
[Equation]
η , λ ( t
[Equation]
4 + 2 η 2 d t , so we have
[Equation] [Equation]
which implies that
[Equation] [Equation]
Based on the above closed-form expression of τ η , λ ( t )
[Equation]
it is straightforward to verify that
[Equation]
Applying Lemma F.
[Equation]
we have
[Equation]
By the property of R η , λ ( t ) in Equation (19) and the deﬁnition of τ η , λ ( t ) in Equation (11)
[Equation]
we have ηλ d t + d ln R η , λ ( t )
[Equation]
where the second equality follows from the deﬁnition of R η , λ ( t ) in Equation (10).
Then combining the above identity and Equation (24) yields
[Equation]
Note that τ η , λ : [0 , ∞ ) →
[0 , ∞ ) is a bijection, so we can rewrite the above in ˜ t := τ η , λ ( t ) to get
[Equation]
where by writing τ − 1 η , λ ( · ) we mean the functional inverse of τ η , λ ( · ) deﬁned in Equation (11).
By assumption η < 2 λ , can see from Equation (10) that R η , λ ( t ) 2 is decreasing on [0 , ∞ ) and is bounded from below by , η / (2 λ ) .
Therefore, for any t ≥ 0 , R η , λ ( t ) 2 / η ≥ 1 / √ 2 λη , so we can again apply Theorem 3.10 to get the desired limiting diffusion.
This ﬁnishes the proof.
D.2 Proof for the limiting diffusion of SGD+WD
Recall the update of SGD+WD in Equation (1) and the scaling function R η , λ in Equation (10), and the normalized version x η , λ ( k ) /R
η , λ ( k ) satisﬁes the following:
[Equation]
Then ˆ x η , λ ( k ) =
x η , λ ( k ) /R
η , λ ( k ) satisﬁes
[Equation]
Deﬁne
X η , λ ( t ) =
ˆ x η , λ ( ⌊ t ⌋ ) , and rewriting the above equation yields Equation (13):
[Equation]
For clarity
[Equation]
we break Lemma 5.3 into the following series of lemmas and prove them respectively
[Equation]
Lemma D.1.
Let { A η , λ } η , λ > 0 be deﬁned as in Equation (15) .
Then A η , λ increases inﬁnitely fast as ηλ → 0 with η < 2 λ < c for some constant c , that is, for any ) > 0 ,
[Equation]
Proof of Lemma D
[Equation]
Fix any ) > 0 and ˜ t ≥ 0 , and let η be sufﬁciently small.
By deﬁnition, we have
[Equation]
where the two inequalities follow from the fact that R η , λ ( t ) is monotonically decreasing.
Moreover, since R η , λ ( t ) ∈ (( η 2 λ ) 1 / 4 , 1] , it holds that η /R η , λ ( t ) 2 ∈ [ η , √ 2 ηλ ) , for all t ≥ 0 .
We then have
[Equation]
Recall the deﬁnition of τ η , λ ( t ) in Equation (11), we further have
[Equation]
which holds for all ˜ t ≥ 0 simultaneously
Hence, we obtain the desired result:
[Equation]
Lemma D.2.
Let { B η , λ } η > 0 be deﬁned as in Equation (16) .
Then as ηλ → 0 with η < 2 λ < c for some constant c , B η , λ (˜ t ) converges to ˜ t/ 2 in uniform metric.
Proof of Lemma D.2.
Since ˜ t = τ η , λ ( t ) =
[Equation]
0 τ ′ η , λ ( s )d s , the idea is to show that the above sum can be seen as an approximation of the Riemann sum
[Equation]
which then yields the approximation of the integral.
To see this, recall the deﬁnition of R η , λ ( t ) in Equation (10)
[Equation]
and note that there exists some t
i ∈
[ i, i + 1] such that
[Equation]
Then since R η , λ ( t ) is positive and monotonically decreasing
[Equation]
it follows that
[Equation]
where the last inequality again follows from the monotonicity
Note that when i ≤ 18 ηλ ln 1 η , we have R η , λ ( i ) 4 ≥ √ η (1 − η 2 λ ) , and when i > 18 ηλ ln 1 η , we have e − 4 ηλ i ≤ √ η .
Thus we further have
[Equation]
Then since R η , λ ( t ) is monotonically decreasing, it follows that
[Equation]
where the second inequality follows from the fact that η 2 /R η , λ ( t ) 4 ∈ [ η 2 , 2 ηλ ) for all t ≥ 0 .
On the other hand
[Equation]
we similarly establish a lower bound for B η , λ (˜ t ) .
It follows from Equation (26) that
[Equation]
where the last inequality is due to the fact that 1 − e − x ≤ x for x ≥ 0 .
Then applying the similar argument as we have used for the upper bound, we get
[Equation] [Equation]
for some universal constant C > 0 .
Now, combining the upper bound in Equation (27) and the lower bound in Equation (28), we obtain
[Equation]
for all ˜ t ≥ 0 , where C > 0 is some universal constant.
Recall the deﬁnition of the uniform metric, and we then have
[Equation]
[Equation]
where the right-hand-side converges to 0 by a standard ) - δ argument as η → 0 .
This completes the proof.
Lemma D.3.
Let R η , λ ( t ) be deﬁned in Equation (10)
[Equation]
and τ η , λ ( t ) in Equation (11) .
Then it holds that
[Equation]
uniformly for all ˜ t ≥ 0 .
Proof of Lemma D.3.
First, since R η , λ ( t ) is positive and monotonically decreasing, we have
[Equation]
which further implies that
[Equation]
s.
Then since η 2 /R η , λ ( t ) 4 ≤ max { 2 ηλ , η 2 } for all t ≥ 0 , we have
[Equation]
Recall the deﬁnition of τ η , λ ( t ) in Equation (11), so
[Equation]
Therefore, we get
[Equation]
for all ˜ t ≥ 0 .
This implies the uniform convergence for all ˜ t ≥ 0 and thus completes the proof.
Lemma D.4.
Let { Z η , λ } η > 0 be deﬁned as in Equation (17) .
Then as ηλ → 0 with η < 2 λ < c for some constant c , Z η , λ weakly converges to ( I Ξ − 1 Ξ 11 ⊤ ) 1 / 2 W in the Skorokhod metric, where W is a Ξ -dimensional standard Brownian motion.
Proof of Lemma D.4.
For any η > 0 , note that Z η , λ is a stochastic process with independent increments, so Z η , λ is indeed a Ξ -dimensional martingale.
Thus we can apply the multidimensional martingale functional central limit theorem (Theorem F.2) to get the desired convergence.
To do so, we need to verify two facts:
1) the expected value of the maximum jump in Z η , λ is asymptotically negligible; 2) for each ﬁxed ˜ t ≥ 0 ,
[Equation]
These two facts correspond to conditions (37) and (38) respectively.
For the ﬁrst fact, note that the norm of the i -th jump in Z η , λ is bounded by
[Equation]
1 Ξ 1 % 2 .
Since R η , λ ( t ) ≥ min { ( η 2 λ ) 1 / 4 , 1 } for all t ≥ 0 , we see that
[Equation]
(1+ √ Ξ ) max { √ 2 ηλ , η } , which implies that any jump in Z η , λ is asymptotically negligible as η →and ηλ → 0 .
For the second fact, ﬁxing any ˜ t ≥ 0 , we have
[Equation]
where the second equality follows from the independence between ξ i and ξ j for i ∕ = j .
Now, applying Lemma D.3 yields lim η → 0 E [ Z η , λ (˜ t )
Z η , λ (˜ t ) ⊤ ]
=
˜ t ( I Ξ − 1 Ξ 11 ⊤ ) .
Combining the above two facts, applying Theorem F.2 completes the proof.
E Omitted Proofs in in Section
Here we provide the proof of the fast mixing of SGD+WD.
By Theorem 5.4, we know that the iterates of SGD+WD track the solution to the limiting diffusion in Equation (12) at time T after roughly 14 ηλ ln( 2 λη ( e 2 T − 1)+1) steps.
Here T is the time index for the limiting diffusion, which can be shown to mix to the equilibrium and the mixing speed is independent of η or λ since Equation (12) does not rely on these parameters.
Therefore, it sufﬁces to establish the ergodicity of the limiting diffusion, and the mixing of SGD+WD immediately follows.
Then the total number of steps of
SGD+WD is on the order of O ( 1 ηλ ln λη ) , yielding the fast mixing of SGD+WD.
E.1 Proof for the Norm Dynamics
We ﬁrst show the dynamics of the squared ℓ 2 norm of the limiting dynamics.
We need the following intermediate lemma.
Lemma E.1.
For any x ∈ Γ and A ∈ R D × D , it holds that 〈 x, ∂ 2 Φ ( x )[
A ] 〉 = tr( A ) − tr( ∂ Φ ( x ) A ∂ Φ ( x ) ⊤ ) .
Proof of Lemma E.1.
Since 〈 x, L ( x ) 〉 = 0 for all x , we have % Φ ( x ) %
=
% x % 22 , and thus ∂% Φ ( x ) % 22 =
I D , that is, ∂ 2 % Φ ( x ) %
[ A ] = 2 I D
[ A ] = 2 tr( A ) for any u, v ∈ R D .
On the other hand, we have
[Equation]
for any j ∈ [ D ] , and then
[Equation]
for any j, k ∈ [ D ] .
Therefore, for any A ∈ R D × D , we have
[Equation]
which further implies that
[Equation]
Since x ∈ Γ
[Equation]
we have Φ ( x ) =
x , so rearranging the above equation yields
[Equation]
which completes the proof
[Equation]
Proof of Lemma
[Equation]
1. Consider the dynamics of % Y ( t ) % 22 , and it follows from Ito’s lemma that
[Equation]
Note that since ∂ Φ ( Y ( t ))
Y ( t ) =
Y ( t ) , we have
[Equation]
) , σ
[Equation]
Y
[Equation]
t
[Equation]
d W
[Equation]
t
[Equation]
Since Y ( t ) stays on Γ , it follows from Lemma E.1 that 〈 Y ( t ) , ∂ 2 Φ ( Y ( t ))
[ Σ ( Y ( t ))]
= tr( Σ ( Y ( t ))) − tr[ ∂ Φ ( Y ( t ))
Σ ( Y ( t ))
∂
Φ ( Y ( t ))
⊤ ] .
Combining the above yields the dynamics of % Y ( t ) % 22 .
E.2 Proof for the Mixing When Trace is a Constant
In this subsection, we ﬁrst prove the mixing result when the tr[ Σ ] is a constant on Γ 1 , that is, under
Assumption E.2.
Assumption E.2.
For any x ∈ Γ such that % x % 2 = 1 , tr( Σ ( x )) = 1 .
Note that, by Lemma 6.1, Assumption E.2 implies that
[Equation]
which is deterministic and has a closed-form solution:
[Equation]
Therefore, regardless of the initialization, % Y ( t ) % 2 will always (monotonically) converge to 1 as t → ∞ .
In this case, we only need to consider the mixing of the direction of Y ( t ) , i.e. Y ( t ) / % Y ( t ) % 2 .
Lemma E.3.
Under Assumption E.2, for SDE (12) , deﬁne Y ( t ) =
Y ( t ) /
% Y ( t ) % 2 .
Then
[Equation]
Based on Assumption E.2 and Lemma 6.1, we then can obtain the dynamics of direction of Y ( t ) , i.e., Y ( t ) /
% Y ( t ) % 2 .
Proof of Lemma E.3.
By Lemma 6.1, we have
[Equation]
where the third equality follows from Lemma 3.9 and the assumption that tr( Σ ( x )) ≡ 1 / % x % 22 .
Finally plugging-in the expression for % Y ( t ) % 22 in Equation (29) ﬁnishes the proof.
Now according to Lemma E.3, { Y ( t ) } ˜ t ≥ 0 is a diffusion process on Γ 1 .
We denote the associated class of Markov transition kernel by {P t } t ≥ 0 where each P t : Γ 1 × B ( Γ 1 ) → R + and P t ( x, A ) is equal to the probability that starting from x the process is in the set
A at time t .
To show the ergodicity of {P t } t ≥ 0 , we follow the approach introduced in Ichihara and Kunita [18], Kliemann [23] that proves ergodicity of a diffusion process by relating to its associated control system.
Deﬁne f 0 ( x ) = 12 ( ∂ 2 Φ ( x )[ Σ ( x )] − x ) and f i ( x ) =
∂
Φ ( x ) σ i ( x ) for each i ∈ [ Ξ ] .
We rescale the time by deﬁning τ ( t )
=
[Equation]
Y
[Equation]
42 − 1) d s .
Let #
Y ( τ ( t ))
=
Y ( t ) , identifying ˜ t ≡ τ ( t ) , then #
Y admits d #
Y
[Equation]
˜
[Equation] [Equation]
˜
[Equation]
$ Ξ i =
f i ( x )d W i (˜ t ) .
For #
Y , we denote the associated class of
Markov transition kernel by { #
P ˜ t } ˜ t ≥ 0 .
Since there is a bijection between t and ˜ t and lim t →∞ ˜ t/t = 1 , the mixing of #
P ˜ t implies that of P t .
Next, wrting the dynamics of #
Y as a Stratonovich SDE, we get
[Equation]
Then we deﬁne its associated deterministic control system:
[Equation]
where u is any piecewise continuous R Ξ -valued functions.
We denote by ψ ( t, x ; u ) the solution of
Equation (32) at time t with initialization x under the control u .
Let U be the set of all piecewise continuous function from R + to R Ξ .
For any x ∈ Γ , deﬁne the positive orbit O + ( x ) = { ψ ( t, x ; u ) : t ≥ 0 , u ∈ U} , that is, all points reachable from x via the control system (32).
For any S ⊆ Γ , we say S is an invariant control set under the control system (32)
if O + ( x ) = ¯ S for all x ∈ S .
The following lemma establishes the key property of the control system in Equation (32).
Lemma E.4.
Under Assumption 6.2 and 6.3, Γ 1 itself is the unique invariant control set contained in Γ 1 for the control system (32) .
Proof of Lemma E.4.
It sufﬁces to show that for any x, y ∈ Γ 1 ,
y ∈
O
+ ( x ) .
By Assumption 6.3, { f i } Ξ i =1 satisﬁes the Lie algebra rank condition, so the following driftless control system
[Equation]
is globally controllable by Rashevski-Chow theorem (see, e.g., [1]).
Therefore, there exists some
T > 0 and a control u : R + → R Ξ such that
y
= ˆ ψ ( T, x ; u ) , where ˆ ψ ( T, x ; u ) is the solution of Equation (33) at time T with initialization x under the control u .
Next, we can use the global controllability of the driftless control system to show that y ∈ O
+ ( x ) for the original control system (32).
For any δ > 0 , deﬁne a control u δ : R + → R Ξ as u δ ( t ) =
u ( t/ δ ) , then we have y
= ˆ ψ ( T, x ; u )
=
ˆ ψ ( T δ , x ; u δ ) .
Now using u δ for the original control system (32), it follows that lim
δ
→ 0 ψ ( T δ , x ; u δ ) =
y as the drift is dominated by the controlled terms.
Thus we conclude that y ∈ O + ( x ) , and this completes the proof.
Now we are ready to prove Theorem 6.6.
Proof of Theorem 6.6 when trace is constant.
Since #
[Equation]
and lim
[Equation]
˜
[Equation] [Equation]
= 1 by L’Hospital’s rule, it sufﬁces to show the ergodicity of #
P ˜ t =
P t .
Let L ( f 1 , . . .
, f Ξ ) be the Lie algebra generated by { f 1 , . . .
, f Ξ } .
Assumption 6.3 implies that dim L ( f 1 , . . .
, f Ξ ) =
D − M − 1 , which is equal to the dimension of Γ 1 .
Therefore, by Lemma 5.1 in Ichihara and Kunita [18], #
P ˜ t satisﬁes the strong Feller property.
Since Γ 1 is compact by Assumption 6.2, we see that there exists at least one invariant measure by Theorem F.3, which further implies that there exists at least on extremal invariant measure by applying Krein-Milman theorem.
Moreover, combining the result of Lemma E.4 and
Lemma 4.1 in Kliemann [23], we see that { #
P
[Equation]
˜ t ≥ 0 has a unique invariant measure, which we denote by π .
Finally, by Proposition 5.1 in Ichihara and Kunita [18], #
[Equation]
ergodic in the sense that
[Equation]
where % · % TV is the total variation distance and supp ( π )
=
Γ 1 again by Lemma 4.1 in Kliemann [23].
This completes the proof.
E.3 Proof for the Mixing When Trace is not Constant
We continue to use the notation f i ( x ) =
∂
Φ ( x ) σ i ( x ) for each i ∈ [ Ξ ] but we change the deﬁnition of f 0 ( x ) to
[Equation]
Writing the dynamics of Y in (12) as a Stratonovich SDE, we get
[Equation]
Then we deﬁne its associated deterministic control system:
[Equation]
Similar to Lemma E.4, we have the following lemma.
Lemma E.5.
Under Assumption 6.2 and 6.3, Γ 1 itself is the unique invariant control set contained in Γ 1 for the control system (34) .
Note that tr[ Σ ] is continuous and Γ 1 bounded, we know that both max x ∈ Γ 1 tr[ Σ ( x )] and min x ∈ Γ 1 tr[ Σ ( x )] can be attained in Γ 1 .
By Assumption 6.3, we know that min x ∈ Γ 1 tr[ Σ ( x )] > 0 .
By Lemma 6.1, we know if % Y ( t ) % 2 ≥ (max x ∈ Γ 1 tr[ Σ ( x )]) 0 .
25 then d , Y ( t ) , 22 d t ≤ 0 and that if % Y ( t ) % 2 ≤ (max x ∈ Γ 1 tr[ Σ ( x )]) 0 .
25 then d , Y ( t ) , 22 d t ≥ 0 .
Thus Γ a,b := { x ∈ Γ | a ≤ b } is an invariant set for both stochastic process (12) and control system (34) if a ≤ (min x ∈ Γ 1 tr[ Σ ( x )]) 0 .and
b ≥ (max x ∈ Γ 1 tr[ Σ ( x )]) 0 .
We also denote { x ∈ Γ | min x ∈ Γ 1 tr[ Σ ( x )] <
% x % 42 < max x ∈ Γ 1 tr[ Σ ( x )] } by Γ ∗ .
Without loss of generality, we can pick some large enough b and small enough a such that x init ∈
Γ a,b and use Γ a,b as our domain.
In this way, our domain is compact and we denote it by Γ .
Lemma E.6.
Under Assumption 6.2 and 6.3, Γ ∗ is the unique invariant control set contained in Γ for the control system (34) .
Proof of Lemma E.6.
Note that Γ ∗ is an invariant set, it sufﬁces to show that for every x ∈ Γ , Γ ∗ ⊂ O + ( x ) .
Below we will show that for any y ∈ Γ ∗ , x ∈ Γ , y ∈
O
+ ( x ) .
By Lemma E.5, we know that for any x ∈ Γ , we have that Γ , x , 2 ⊂ O + ( x ) .
Therefore, it sufﬁces to show there exist u ∈ Γ , x , 2 and v ∈ Γ , y , 2 such that v ∈ O + ( u ) .
If this is true, then we know that y ∈ O
+ ( v ) ⊂
O + ( u ) ⊂
O + ( x ) .
Without loss of generality, we assume that % x % 2 ≥ % y % 2 .
Our strategy is to pick any z ∈ argmin z ∈ Γ 1 tr[ Σ ( z )] and let v =
y % 2 z
and u =
% x % 2 z .
Then we can split the control u
δ : R + →
R Ξ into many small piece-wise constant intervals, such that in each interval, $
[Equation]
[Equation]
u i ( t ) cancels the tangent component of f 0 ( ψ ( t ))
(in the tangent space of Γ , ψ ( t ) , 2 at ψ ( t ) ).
Thus through the whole trajectory, the direction of ψ ( t ) can be arbitrarily close to z .
That is, we claim there exists a sequence of control { u n } n ∈ N , such that the corresponding ψ n ( t ) → ψ ∞ ( t ) pointwise, where ψ ∞ is deﬁned as the solution of the following ODE.
[Equation]
By deﬁnition of z , we know % y %
2 ≥ % z
Thus there exists some ﬁnite time t such that ψ ∞ ( t ) = v .
That is, v ∈ O + ( u ) .
This completes the proof.
Lemma E.7 (Lie Algebra Rank Condition) .
For any k ≥ 1 , k ∈ N , deﬁne A k :=
{ [[[ f i 1 , f i 2 ] , f i 3 ] , . . .
f i k ] | i 1 ∈ [ Ξ ] , i 2 , . .
, i k ∈ { 0 } ∪ [ Ξ ] } and A = ∪ ∞ k =1 A k .
For any x ∈ Γ , rank(span( { f ( x ) | f ∈ A } ))
=
D − M .
Proof of Lemma E.7.
Note the rank( T x ( Γ , x , 2 ))
=
D − M − 1 for any x ∈ Γ , by Assumption 6.3, it sufﬁces to show for each x ∈ Γ , there exists some f ∈ A and 〈 x, f ( x ) 〉 ∕ = 0 .
Since tr[ Σ ] is an non-constant analytic function and Γ 1 is an analytic manifold, for every x ∈ Γ 1 , there must exists some k ∗ ≥ 1 and some linear combination of f i , i ∈
[ Ξ ] , denoted by f ∗ , such that
[Equation]
By Lemma 3.5 and the fact that tr[ Σ ] is ( − 2) -homogeneous, we know that tr[ Σ ]( x ) = −= x, f 0 ( x ) + x> .
Then we can apply Lemma E.8 iteratively on Equation (35) and get that
[Equation]
The proof is completed by noting that [[[ x, f ∗ ] , f ∗ ]
[Equation]
, x ] = 0 by apply the second claim of
Lemma E.8 iteratively.
Lemma E.8.
Let f, g be any two vector ﬁeld on Γ , where 〈 f ( x ) , x 〉 = 0 for every x ∈ Γ .
It holds that 〈 x, [ f, g ] 〉 = 〈∇ 〈 x, g 〉 , f 〉 .
In speciﬁc, 〈 x, [ x, f ] 〉 = 0 .
Proof of Theorem 6.6 when trace is not constant.
Let {P t } t ≥ 0 be the Markov semigroup associated with the limiting diffusion.
Since the manifold Γ is compact, the space of the probability measures on Γ is also compact, thus for any probability measure on Γ , the sequence {P t µ 0 } t ≥
⊆ P ( Γ ) is tight.
Therefore, by Theorem F.3, there exists at least one invariant probability measure for {P t } t ≥ 0 , which further implies that there exists at least on extremal invariant measure by applying KreinMilman theorem.
Moreover, combining the result of Lemma E.6 and Lemma 4.1 in Kliemann [23], we see that {P t } t ≥ 0 has a unique invariant measure, which we denote by π .
Finally, by Proposition 5.1 in Ichihara and Kunita [18],
P t is ergodic in the sense that
[Equation]
where % · % TV is the total variation distance and supp ( π )
=
Γ 1 again by Lemma 4.1 in Kliemann [23].
This completes the proof.
F Auxiliary Results
The following lemma is a classic result.
See, e.g., Corollary 8.5.5 in Oksendal [34].
Lemma F.1 (Time-changed Brownian motion) .
Let g : [0 , ∞ )
→ [0 , ∞ ) be a differentiable function with nonnegative ﬁrst-order derivative.
Then it holds that
[Equation]
The multidimensional martingale functional central limit theorem can be found in, e.g., Theorem 7.
in Ethier and Kurtz [11] or Theorem 2.1 in Whitt [44].
We use a version of Theorem 2.1 in Whitt [44] given in Theorem F.2.
Theorem F.2 (Multidimensional martingale FCLT, [44]) .
For n ≥ 1 , let M n =
( M n, 1 , . . .
, M n,k ) be a martingale satisfying M n (0) =
Let C ∈ R k × k be a semi-positive deﬁnite matrix.
Suppose the following holds:
• The expected value of the maximum jump in M n is asymptotically negligible, i.e., for each T > 0 ,
[Equation]
where J ( x, T ) =
sup {| x ( t ) − x ( t − ) | : 0 < t ≤ T } for any function x : [0 , ∞ )
→ R .
• For each pair ( i, j ) ∈ [ k ] ×
[ k ] and each t > 0 ,
[Equation]
n,i
[Equation]
n,j ]
[Equation]
Then M n ⇒ √ CW in Skorokhod metric where W is the k -dimensional standard Brownian motion.
Theorem F.3 (Krylov-Bogolioubov Theorem, [9]) .
Let {P t } t ≥ 0 be a Feller Markov semigroup over a Polish space
X .
Assume that there exists µ 0 ∈ M 1 ( X ) such that the sequence {P t µ 0 } t ≥ 0 is tight.
Then there exists at least one invariant probability measure for
[Equation]