W HAT H APPENS AFTER SGD R EACHES Z ERO L OSS ?
A M ATHEMATICAL F RAMEWORK
Anonymous authors
Paper under double-blind review
A BSTRACT
Understanding the implicit bias of Stochastic Gradient Descent (SGD) is one of the key challenges in deep learning, especially for overparametrized models, where the local minimizers of the loss function L can form a manifold.
Intuitively, with a sufﬁciently small learning rate η , SGD tracks Gradient Descent (GD) until it gets close to such manifold, where the gradient noise prevents further convergence.
In such regime, Blanc et al.
(2020) proved that SGD with label noise locally decreases a regularizer-like term, the sharpness of loss, tr[ ∇ 2 L ] .
The current paper gives a general framework for such analysis by adapting ideas from Katzenberger (1991).
It allows in principle a complete characterization for the regularization effect of SGD around such manifold—i.e., the ”implicit bias”—using a stochastic differential equation (SDE) describing the limiting dynamics of the parameters, which is determined jointly by the loss function and the noise covariance.
This yields some new results: (1) a global analysis of the implicit bias valid for η − 2 steps, in contrast to the local analysis of Blanc et al.
(2020) that is only valid for η − 1 .
6 steps and (2) allowing arbitrary noise covariance.
As an application, we show with arbitrary large initialization, label noise SGD can always escape the kernel regime and only requires O ( κ ln d ) samples for learning an κ -sparse overparametrized linear model in R d (Woodworth et al., 2020), while GD initialized in the kernel regime requires Ω( d ) samples.
This upper bound is minimax optimal and improves the previous (cid:101) O ( κ 2 ) upper bound (HaoChen et al., 2020).
I NTRODUCTION
The implicit bias underlies the generalization ability of machine learning models trained by stochastic gradient descent (SGD).
But it still remains a mystery to mathematically characterize such bias.
We study SGD in the following formulation
[Equation]
where η is the learning rate (LR), L : R D
→ R is the training loss and σ ( x ) = [ σ 1 ( x ) , σ 2 ( x ) , . . . , σ Ξ ( x )] ∈ R D
× Ξ is a deterministic noise function. Here ξ k is sampled uniformly from { 1 , 2 , . . . , Ξ } and it satisﬁes E ξ k
[ σ ξ k ( x )] = 0 , ∀ x ∈ R d and k .
It is widely believed that large LR (or equivalently, small batch size) helps SGD ﬁnd better minima.
For instance, some previous works argued that large noise enables SGD to select a ﬂatter attraction basin of the loss landscape which potentiallyx beneﬁts generalization (Li et al., 2019c; Jastrzebski et al., 2017).
However, there is also experimental evidence (Li et al., 2020b) that small LR also has equally good implicit bias (albeit with higher training time), and that is the case studied here.
Presumably low LR precludes SGD jumping between different basins since under general conditions this should require Ω(exp(1 /η )) steps (Shi et al., 2020).
In other words, there should be a mechanism to reach better generalization while staying within a single basin.
For deterministic GD similar mechanisms have been demonstrated in simple cases (Soudry et al., 2018; Lyu & Li, 2019) for GD and referred to as implicit bias of gradient descent .
The current paper can be seen as study of implicit bias of Stochastic GD, which turns out to be quite different, mathematically.
Recent work (Blanc et al., 2020) shed light on this direction by analyzing effects of stochasticity in the gradient.
For sufﬁciently small LR, SGD will reach and be trapped around some manifold of local minimizers, denoted by Γ (see Figure 2). The effect is shown to be an implicit deterministic drift in a direction corresponding to lowering a regularizer-like term along the manifold. They showed SGD
[Figure]
with label noise locally decreases the sharpness of loss, tr[ ∇ 2 L ] , by Θ( η 0 . 4 ) in η − 1 . 6 steps.
However, such an analysis is actually local , since the natural time scale of analysis should be η − 2 , not η − 1 . 6 .
The contribution of the current paper is a more general and global analysis of this type.
We introduce a more powerful framework inspired by the classic paper Katzenberger (1991).
1. I NTUITIVE EXPLANATION OF REGULARIZATION EFFECT DUE TO SGD
We start with an intuitive description of the implicit regularization effect described in Blanc et al.
For simpliﬁcation, we show it for the following canonical SDE approximation (See Section B.1 for more details) of SGD (1) (Li et al., 2017; Cheng et al., 2020).
Here W ( t ) is the standard Ξ -dimensional Brownian motion.
The only property about label noise SGD we will use is that the noise covariance σσ (cid:62)
[Equation]
=
[Equation]
) for every x in the manifold Γ .
(See derivation in Section 5)
[Equation]
Suppose (cid:101)
X η (0) is already close to some local minimizer point X ∗ ∈ Γ .
The goal is to show (cid:101) X η ( t ) will move in the tangent space and steadily decrease tr[ ∇ 2 L ] .
At ﬁrst glance, this seems impossible as the gradient ∇ L vanishes around Γ , and the noise has zero mean, implying SGD should be like random walk instead of a deterministic drift.
The key observation of Blanc et al.
(2020) is that the local dynamics of (cid:101)
X η ( t ) is completely different in tangent space and normal space — the fast random walk in normal space causes (cid:101) X η ( t ) to move slowly (with velocity Θ( η 2 ) ) but deterministically in certain direction.
To explain this, letting ∆( t ) = (cid:101)
X η ( t ) − X ∗ , Taylor expansion of (2) gives d∆( t )
≈ − η ∇ 2 L ( X ∗ )∆d t + ησ
( X ∗ )
d W ( t ) , meaning ∆ is behaving like an Ornstein-Uhlenbeck (OU) process locally in the normal space.
Its mixing time is Θ( η − 1 ) and the stationary distribution is the standard multivariate gaussian in the normal space scaled by √ η (see Figure 1b), because noise covariance σσ (cid:62) = ∇ 2 L .
Though this OU process itself doesn’t form any regularization, it activates the second order Taylor expansion of ∇ L ( X ∗ + ∆( t )) , i.e., − 12 ∂ 2 ( ∇ L )( X ∗ )
[∆( t ) , ∆( t )] , creating a Θ( η 2 ) velocity in the tangent space.
Since there is no push back force in the tangent space, the small velocity accumulates over time, and in a longer time scale of Ω( η − 1 ) , the time average of the stochastic velocity is roughly the same as the expected velocity when ∆ is sampled from its stationary distribution.
This simpliﬁes the expression of the velocity in tangent space to η 2 2 ∇ T tr[ ∇ 2 L ] (see
Figure 1c), where ∇ T means the gradient is only taken in the tangent space.
However, the above approach only gives a local analysis for O ( η − 1 . 6 ) time, where the total movement due to implicit regularization is O ( η 2 − 1 . 6 )
=
O ( η 0 . 4 ) and thus is negligible when η → 0 .
In order to get a non-trivial limiting dynamics when η → 0 , a global analysis for Ω( η − 2 ) steps is necessary and it cannot be done by Taylor expansion with a single reference point.
Recent work by Damian et al.
(2021) glues analyzes of multiple local phases into a global guarantee that SGD ﬁnds a ( (cid:15), γ ) - stationary point for the regularized loss.
This technique works for constant size LR, but this still doesn’t show convergence for trajectory when η → 0 .
The main technical difﬁculty here is that it’s not clear how to separate the slow and fast dynamics in different spaces and how to only take limit for the slow dynamics, especially when shifting to a new reference point in the Taylor series calculation.
1. O UR A PPROACH : S EPARATING THE S LOW FROM THE F AST
In this work, we tackle this problem via a different angle.
First, since the anticipated limiting dynamics is of speed Θ( η 2 ) , we change the time scaling to accelerate (2) by η − 2 times, which yields
[Equation]
The key idea here is that we only need to track the slow dynamic, or equivalently, some projection of
X onto the manifold Γ , Φ( X ) .
Here Φ : R D
→ Γ is some function to be speciﬁed
and hopefully we can simplify the dynamics Equation (3) via choosing suitable Φ .
To track the dynamics of Φ( X η ) , we apply Ito’s lemma (a.k.a.
stochastic chain rule, see Lemma A.8) to Equation (3), which yields
[Equation]
Note the ﬁrst term − η − 1 ∂ Φ( X η ) ∇ L ( X η ) is going to diverge to ∞ when η → 0 , so a natural choice for Φ is to kill the ﬁrst term.
Further note − ∂ Φ( X ) ∇ L ( X ) is indeed the directional derivative of Φ at X towards −∇ L , killing the ﬁrst term becomes equivalent to making Φ invariant under Gradient
Flow of −∇ L ( X ) !
Thus it sufﬁces to take Φ( X ) to be the limit of gradient ﬂow starting at X .
(Formally deﬁned in Section 3; see Lemma C.2 for a proof of ∂ Φ( X ) ∇ L ( X ) ≡ 0 .)
Also intuitively X η will be inﬁnitely close to Γ , i.e., d ( X η ( t ) , Γ) → 0 for any t > 0 as η → 0 , so we have Φ( X η ) ≈ X η .
Thus we can rewrite the above equation as
[Equation]
and the solution of (4) shall converge to that of the following (in an intuitive sense):
[Equation]
The above argument for SDE was ﬁrst formalized and rigorously proved by Katzenberger (1991).
It included an extension of the analysis to the case of asymptotic continuous dynamics (Theorem 4.1) including SGD with inﬁnitesimal LR, but the result is weaker in this case and no convergence is shown.
Another obstacle for applying this analysis is that 2nd order partial derivatives of Φ are unknown.
We solve these issues in Section 4 and our main result Theorem 4.6 gives a clean and complete characterization for the implicit bias of SGD with inﬁnitesimal LR in Θ( η − 2 ) steps.
Finally, our Corollary 5.2 shows (5) gives exactly the same regularization as tr[ ∇ 2 L ] for label noise SGD.
The main contributions of this paper are summarized as follows.
In Section 4, we propose a mathematical framework to study the implicit bias of SGD with inﬁnitesimal LR.
Our main theorem (Theorem 4.6) gives the limiting diffusion of SGD with LR η for Θ( η − 2 ) steps as η → 0 and allows any covariance structure.
In Section 5, we give limiting dynamics of SGD with isotropic noise and label noise.
In Section 6, we show for any initialization, SGD with label noise achieves O ( κ ln d ) sample complexity for learning a κ -sparse overparametrized linear model (Woodworth et al., 2020).
In this case, the implicit regularizer is a data-dependent weighted (cid:96) 1 regularizer, meaning noise can help reduce the norm and even escape the kernel regime.
The O ( κ ln d ) rate is minimax optimal (Raskutti et al., 2012) and improves over (cid:101) O ( κ 2 ) upper bound by HaoChen et al.
In contrast, vanilla GD requires Ω( d ) samples to generalize in the kernel regime.
For technical contributions, we rigorously prove the convergence of GF for OLM (Lemma 6.3), unlike many existing implicit bias analyses which have to assume the convergence.
We also prove the convergence of limiting ﬂow to the global minimizer of the regularizer (Lemma 6.5) by a trajectory analysis via our framework.
It cannot be proved by previous results (Blanc et al., 2020;
Damian et al., 2021), as they only assert convergence to stationary point in the best case.
R ELATED W ORKS
Loss Landscape of Overparametrized Models A phenomenon known as mode connectivity has been observed that local minimizers of the loss function of a neural network are connected by simple paths (Freeman & Bruna, 2016;
Garipov et al., 2018; Draxler et al., 2018), especially for overparametrized models (Venturi et al., 2018; Liang et al., 2018; Nguyen et al., 2018; Nguyen, 2019).
Later this phenomanon is explained under generic assumptions by Kuditipudi et al.
Moreover, it has been proved that the local minimizers of an overparametrized network form a low-dimensional manifold (Cooper, 2018; 2020) which possibly has many components.
Fehrman et al.
(2020) proved the convergence rate of SGD to the manifold of local minimizers starting in a small neighborhood.
Implicit Bias in Overparametrized Models Algorithmic regularization has received great attention in the community (Arora et al., 2018; 2019a;
Gunasekar et al., 2018b;a;b; Soudry et al., 2018; Li et al., 2018; 2020a).
In particular, the SGD noise is widely believed to be a promising candidate for explaining the generalization ability of modern neural networks (LeCun et al., 2012; Keskar et al., 2016;
Hoffer et al., 2017; Zhu et al., 2018; Li et al., 2019a).
Beyond the size of noise (Li et al., 2019c;
Jastrzebski et al., 2017), the shape and class of the noise also play an important role (Wen et al., 2019;
Wu et al., 2020).
It is shown by HaoChen et al.
(2020) that parameter-dependent noise will bias SGD towards a low-complexity local minimizer.
Several existing works (Vaskevicius et al., 2019; Woodworth et al., 2020;
Zhao et al., 2019) have shown that for the quadratically overparametrized linear model, i.e., w = u (cid:12) 2 − v (cid:12) 2 or w = u (cid:12) v , gradient descent/ﬂow from small initialization implicitly regularizes (cid:96) 1 norm and provides better generalization when the groundtruth is sparse.
This is in sharp contrast to the kernel regime, where neural networks trained by gradient descent behaves like kernel methods (Daniely, 2017; Jacot et al., 2018; Yang, 2019).
This allows one to prove convergence to zero loss solutions in overparametrized settings (Li & Liang, 2018; Du et al., 2018;
Allen-Zhu et al., 2019b;a; Du et al., 2019; Zou et al., 2020), where the learnt function minimizes the corresponding RKHS norm (Arora et al., 2019b; Chizat et al., 2018).
Modelling Stochastic First-Order Methods with Itˆo SDE Apart from the discrete-time analysis, another popular approach to study SGD is through the continuous-time lens using SDE (Li et al., 2017; 2019b; Cheng et al., 2020).
Such an approach is often more elegant and can provide fruitful insights like the linear scaling rule (Krizhevsky, 2014; Goyal et al., 2017) and the intrinsic learning rate (Li et al., 2020b).
A recent work by Li et al.
(2021) justiﬁes such SDE approximation.
Xie et al.
(2020) gave a heuristic derivation explaining why SGD favors ﬂat minima with SDE approximation. Wojtowytsch (2021) showed that the invariant distribution of the canonical SDE approximation of SGD will collapse to some manifold of minimizers and in particular, favors ﬂat minima. By approximating SGD using a SDE with slightly modiﬁed covariance for the overparametrized linear model, Pesme et al. (2021) relates the strength of implicit regularization to training speed.
N OTATION AND P RELIMINARIES
For any integer k , we denote C k as the set of the k times continuously differentiable functions.
We denote a ∧ b = min { a, b } .
For any vector u, v and α ∈ R , we deﬁne [ u (cid:12) v ] i =
u
i v i and [ v (cid:12) α ]
i = ( v i ) α .
For any matrix A , we denote its pseudo inverse by A † .
For 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 ∂ i F j ( x ) .
We also use ∂F ( x )[ u ] and ∂ 2 F ( x )[ u, v ] to denote the ﬁrst and second order directional derivative of F at x along the derivation of u ( and v ).
We abuse the notation of ∂ 2 F by viewing it a linear mapping deﬁned on R D ⊗ R D ∼ = R D 2 , in the sense that ∂ 2 F ( x )[Σ] = (cid:80) Di,j =1 ∂ 2 F ( x )[ e i , e j ]Σ ij , for any Σ ∈ R D × D .
For any submanifold Γ ⊂ R D and x ∈ Γ , we denote by T x (Γ) the tangent space of Γ at x and T ⊥ x (Γ) the normal space of Γ at x .
Given the loss function L , the Gradient Flow (GF) governed by L can be described through a mapping φ : R D × [0 , ∞ )
→ R D satisfying φ ( x, t ) = x − (cid:82) t 0 ∇ L ( φ ( x, s ))
d s .
We further denote the limiting mapping Φ( x ) =
lim t →∞ φ ( x, t ) whenever the limit exists.
We denote 1 ξ ∈ R Ξ as the one-hot vector where ξ -th coordinate is 1 , andthe all 1 vector.
3. M ANIFOLD OF LOCAL MINIMIZERS
Assumption 3.1.
Assume that the loss L : R D → R is a C 3 function, and that Γ is a ( D − M ) - dimensional C 2 -submanifold of R D for some integer 0 ≤ M ≤ D , where for all x ∈ Γ , x is a local minimizer of L and rank( ∇ 2 L ( x ))
=
M .
Assumption 3.2.
Assume that U is an open neighborhood of Γ satisfying that gradient ﬂow starting in U converges to some point in Γ , i.e., ∀ x ∈ U , Φ( x ) ∈ Γ .
(Then Φ is C 2 on U by Falconer (1983).)
When does such a manifold exist?
The vast overparametrization in modern deep learning is a major reason for the set of global minimizers to appear as a Riemannian manifold (possibly with multiple connected components), instead of isolated ones.
Suppose all global minimizers interpolate the training dataset, i.e., ∀ x ∈ R D , L ( x ) =
min x (cid:48) ∈
R D L ( x (cid:48) ) implies f i ( x ) =
y i for all i ∈ [ n ] , then by preimage theorem (Banyaga & Hurtubise, 2013), the manifold Γ := { x ∈ R D
| f i ( x ) =
y i , ∀ i ∈ [ n ] } is of dimension D − n if the Jacobian matrix [ ∇ f 1 ( x ) , . .
. , ∇ f n ( x )] has rank n for all x ∈ Γ .
Note this condition is equivalent to that NTK at x has full rank, which is very common in literature.
L IMITING D IFFUSION OF SGD
In Section 4.1 we ﬁrst recap the main result of Katzenberger (1991).
In Section 4.2 we derive the closed-form expressions of ∂ Φ and ∂ 2 Φ .
We present our main result in Section 4.3.
We remark that sometimes we omit the dependency on t to make things clearer.
4. R ECAP OF K ATZENBERGER ’ S T HEOREM
Let { A n } n ≥ 1 be a sequence of integrators, where each A n : R → R is a non-decreasing function with A n (0)
= 0 .
Let { Z n } n ≥ 1 be a sequence of R | Ξ | -valued stochastic processes deﬁned on R .
Given loss function L and noise covariance function σ , we consider the following stochastic process:
[Equation]
In particular, when the integrator sequence { A n } n ≥ 1 increases inﬁnitely fast, meaning that ∀ (cid:15) > 0 , inf t ≥
( A n ( t + (cid:15) ) − A n ( t ))
→ ∞ as n → ∞ , we call (6) a Katzenberger process .
One difﬁculty for directly studying the limiting dynamics of X n ( t ) is that the point-wise limit as n → ∞ become discontinuous at t = 0 if X (0) / ∈ Γ .
The reason is that clearly lim n →∞ X n (0) =
X (0) , but for any t > 0 , since { A n } n ≥ 1 increases inﬁnitely fast, one can prove lim n →∞ X n ( t ) ∈ Γ !
To circumvent this issue, we consider Y n ( t ) =
X n ( t ) − φ ( X (0) , A n ( t ))
+ Φ( X (0)) .
Then for each n ≥ 1 , we have Y n (0) =
Φ( X (0)) and lim n →∞ Y n ( t ) =
lim n →∞ X n ( t ) .
Thus Y n ( t ) has the same limit on (0 , ∞ ) as X n ( t ) , but the limit of the former is further continuous at t = 0 .
Theorem 4.1 (Informal version of Theorem B.7, Katzenberger (1991)) .
Suppose the loss L , manifold Γ and neighborhood U satisﬁes Assumptions 3.1 and 3.2.
Let { X n } n ≥ 1 be a sequence of Katzenberger process with { A n } n ≥ 1 , { Z n } n ≥ 1 .
Let Y n ( t ) =
X n ( t ) − φ ( X (0) , A n ( t ))
+ Φ( X 0 ) .
Under technical assumptions, it holds that if ( Y n , Z n ) converges to some ( Y, W ) in distribution, where { W ( t ) } t ≥ 0 is the standard Brownian motion, then Y stays on Γ and admits
[Equation]
Indeed, SGD (1) can be rewritten into a Katzenberger process as in the following lemma.
Lemma 4.2.
Let { η n } ∞ n =1 be any positive sequence with lim n →∞ η n = 0 , A n ( t ) =
η n (cid:98) t/η 2 n (cid:99) , and
Z n ( t ) =
η n (cid:80) [Equation] t/η 2 n (cid:99) k
=1 ( 1 ξ k − 1Ξ 1 ) , where ξ 1 , ξ 2 , . . .
i.i.d.
∼ Unif ([Ξ]) .
Then with the same initialization X n (0) =
x η n (0) ≡ X (0) , X n ( kη 2 n ) deﬁned by (6) is a Katzenberger process and is equal to x η n ( k ) deﬁned in (1) with LR equal to η n for all k ≥ 1 .
Moreover, the counterpart of (7) is
[Equation]
where Σ ≡ σσ (cid:62) and { W ( t ) } t ≥ 0 is a Ξ -dimensional standard Brownian motion.
However, there are two obstacles preventing us from directly applying Theorem 4.1 to SGD.
First, the stochastic integral in (8) depends on the derivatives of Φ , ∂ Φ and ∂ ij Φ , but Katzenberger (1991) did not give their dependency on loss L .
To resolve this, we explicitly calculate the derivatives of Φ on Γ in terms of the derivatives of L in Section 4.2.
The second difﬁculty comes from the convergence of ( Y n , Z n ) which we assume as granted for brevity in Theorem 4.1.
In fact, the full version of Theorem 4.1 (see Theorem B.7) concerns the stopped version of Y n with respect to some compact K ⊂ U , i.e. , Y µ n ( K ) n ( t ) = Y n
[Equation]
t ∧ µ n ( K )) where µ n ( K ) is the stopping time of Y n leaving K .
As noted in Katzenberger (1991), we need the convergence of µ n ( K ) for Y µ n ( K ) n to converge, which is a strong condition and difﬁcult to prove in our cases.
We circumvent this issue by proving Theorem B.9, a user-friendly interface for the original theorem in Katzenberger (1991), and it only requires the information about the limiting diffusion.
Building upon these, we present our ﬁnal result as Theorem 4.6.
4. C LOSED -F ORM EXPRESSION OF THE LIMITING DIFFUSION
We can calculate the derivatives of Φ by relating to those of L .
Here the key observation is the invariance of Φ along the trajectory of GF.
Lemma 4.3.
For any x ∈ Γ , ∂ Φ( x ) ∈ R D
× D is the projection matrix onto tangent space T x (Γ) .
To express the second-order derivatives compactly, we introduce the notion of Lyapunov operator .
Deﬁnition 4.4 (Lyapunov Operator) .
For a symmetric matrix H , we deﬁne W H
=
{ Σ ∈ R D × D | Σ = Σ (cid:62) ,
HH † Σ = Σ = Σ HH † } and Lyapunov Operator L H : W H
→ W H as L H (Σ) =
H (cid:62) Σ + Σ H .
It’s easy to verify L − 1 H is well-deﬁned on W H .
Lemma 4.5.
Let x be any point in Γ and Σ = Σ( x ) =
σσ (cid:62) ( x ) ∈ R D
× D be the noise covariance at x 1 .
Then Σ can be decomposed as Σ = Σ (cid:107) + Σ ⊥ + Σ (cid:107) , ⊥ + Σ ⊥ , (cid:107) , where Σ (cid:107) :
= ∂ ΦΣ ∂ Φ , Σ ⊥ :
= ( I D − ∂ Φ)Σ( I D − ∂ Φ) and Σ (cid:107) , ⊥ = Σ (cid:62)⊥ , (cid:107) =
∂
ΦΣ( I D − ∂ Φ) are the noise covariance in tangent space, normal space and across both spaces, respectively.
Then it holds that
[Equation] [Equation]
4. M AIN R ESULT
Now we are ready to present our main result.
It’s a direct combination of Theorem B.9 and Lemma 4.5.
Theorem 4.6.
Suppose the loss function L , the manifold of local minimizer Γ and the open neighborhood U satisfy Assumptions 3.1 and 3.2, and x η (0)
=
x (0) ∈ U for all η > 0 .
If SDE (10) has a global solution Y with Y (0) =
x (0) and Y never leaves U , i.e., P [ Y ( t ) ∈ U, ∀ t ≥ 0] = 1 , then for any T > 0 , x η ( (cid:98) T/η 2 (cid:99) ) converges in distribution to Y ( T ) as η → 0 .
[Equation]
where Σ (cid:107) , Σ ⊥ , Σ ⊥ , (cid:107) are deﬁned in Lemma 4.5.
Based on the above theorem, the limiting dynamics of SGD can be understood as follows: (a) the tangent noise , Σ 1 / 2 (cid:107) ( Y )
d W ( t ) , is preserved, and the second term of (10) can be viewed as the necessary tangent noise compensation for the limiting dynamics to stay on Γ .
Indeed, Lemma C. shows that the value of the second term only depends on Γ itself, i.e., it’s same for all loss L which locally deﬁnes the same Γ .
(b) The noise in the normal space is killed since the limiting dynamics always stay on Γ .
However, its second order effect (It ˆ o correction term) takes place as a vector ﬁeld on Γ , which induces the Noise Regularization and Mixed Regularization term, corresponding to the mixed and normal noise covariance respectively.
Remark 4.7.
Theorem B.9 indeed implies convergence in distribution of sample paths till time T , i.e., let (cid:101) x η ( t ) =
x η ( (cid:98) t/η 2 (cid:99) ) , then (cid:101) x
η weakly converges to Y on [0 , T ] .
Moreover, we only assume the existence of a global solution for ease of presentation.
As long as there exists a compact K ⊆ Γ such that Y stays in K on [0 , T ] with high probability, Theorem B.9 still provides the convergence of SGD iterates (stopped at the boundary of K ) before time T with high probability.
I MPLICATIONS AND E XAMPLES
In this section, we derive the limiting dynamics for two notable noise types, where we ﬁx the expected loss L and the noise distribution, and only drive η to 0.
The proofs are deferred into Appendix C.3.
Type I: Isotropic Noise.
Isotropic noise means Σ( x ) ≡ I D for any x ∈ Γ (Shi et al., 2020).
The following theorem shows that the limiting diffusion with isotropic noise can be viewed as a Brownian
Motion plus Riemannian Gradient Flow with respect to the pseudo-determinant of ∇ 2 L .
Corollary 5.1 (Limiting Diffusion for Isotropic Noise) .
If Σ ≡ I D on Γ , SDE (10) is then
[Equation]
where |∇ 2 L ( Y ) | + =
lim α → 0 |∇ 2 L ( Y )+ αI D | α D − rank( ∇ 2 L ( Y )) is the pseudo-determinant of ∇ 2 L ( Y ) .
|∇ 2 L ( Y ) | + is also equal to the sum of log of non-zero eigenvalue values of ∇ 2 L ( Y ) .
Type II: Label Noise.
When doing SGD for (cid:96) 2 -regression on dataset { ( z i , y i ) } ni =1 , adding label noise (Blanc et al., 2020; Damian et al., 2021) means replacing the true label at iteration k , y i k , by a fresh noisy label (cid:101)
y i k :
=
y i k
+ δ k , where δ k i.i.d.
∼ Unif {− δ, δ } for some constant δ > 0 .
Then the corresponding loss becomes 12 ( f i k ( x ) − (cid:101)
y i k ) 2 , where f i k ( x ) is the output of the model with parameter x on data
z
i k .
So the label noise SGD update is
[Equation]
Suppose the model can achieve the global minimum of the loss L ( x ) := 12 E [( f i ( x ) − (cid:101) y i ) 2 ] at x ∗ , then the model must interpolate the whole dataset, i.e., f i ( x ∗ ) =
y i for all i ∈ [ n ] , and thus here the manifold Γ is a subset of { x ∈ R D | f i ( x ) =
y i , ∀ i ∈ [ n ] } .
Here the key property of the label noise used in previous works is Σ( x ) = 1 n (cid:80) ni
=1 ∇ x f i ( x ) ∇ x f i ( x ) (cid:62) = ∇ 2 L ( x ) .
Lately, Damian et al.
(2021) further generalizes the analysis to other losses, e.g., logistic loss and exponential loss, as long as they satisfy Σ( x ) =
c ∇ 2 L ( x ) for some constant c > 0 .
In sharp contrast to the delicate discrete-time analysis in Blanc et al.
(2020) and Damian et al.
(2021), the following corollary recovers the same result but with much simpler analysis – taking derivatives is all you need.
Under our framework, we no longer need to do Taylor expansion manually nor carefully control the inﬁnitesimal variables of different orders together.
It is also worth mentioning that our framework immediately gives a global analysis of Θ( η − 2 ) steps for SGD, far beyond the local coupling analysis in previous works.
In Section 6, we will see how such global analysis allows us to prove a concrete generalization upper bound in a non-convex problem, the overparametrized linear model (Woodworth et al., 2020; HaoChen et al., 2020) .
Corollary 5.2 (Limiting Flow for Label Noise) .
If Σ ≡ c ∇ 2 L on Γ for some constant c > 0 , SDE (10) can be simpliﬁed into (13) where the regularization is from the noise in normal space.
[Equation]
P ROVABLE G ENERALIZATION B ENEFIT WITH L ABEL N OISE
In this section, we show provable beneﬁt of label noise in generalization using our framework (Theorem B.7) in a concrete setting, the overparametrized linear models (OLM) (Woodworth et al., 2020).
While the existing implicit regularization results for Gradient Flow often relates the generalization quality to initialization, e.g., Woodworth et al.
(2020) shows that for OLM, small initialization corresponds to the rich regime and prefers solutions with small (cid:96) 1 norm while large initialization corresponds to the kernel regime and prefers solutions with small (cid:96) 2 norm, our result
Theorem 6.1 surprisingly proves that even if an OLM is initialized in the kernel regime, label noise
SGD can still help it escape and then enter the rich regime by minimizing its weighted (cid:96) 1 norm.
When the groundtruth is κ -sparse, this provides a (cid:101) O
( κ ln d ) vs Ω( d ) sample complexity separation between
SGD with label noise and GD when both initialized in the kernel regime.
Here d is the dimension of the groundtruth.
The lower bound for GD in the kernel regime is folklore, but for completeness, we state the result as Theorem 6.7 in Section 6.3 and append its proof in Appendix D.6.
Theorem 6.1.
In the setting of OLM, suppose the groundtruth is κ -sparse and n ≥ Ω( κ ln d ) training data are sampled from either i.i.d.
Gaussian or Boolean distribution.
Then for any initialization x init (except a zero-measure set) and any (cid:15) > 0 , there exist η 0 , T > 0 such that for any η < η 0 , OLM trained with label noise SGD (12) with LR equal to η for (cid:98) T/η 2 (cid:99) steps returns an (cid:15) -optimal solution, with probability of 1 − e − Ω( n ) over the randomness of the training dataset.
The proof roadmap of Theorem 6.1 is the following:
Show Assumption 3.1 is satisﬁed, i.e., the set of local minimizers, Γ , is indeed a manifold and the hessian ∇ 2 L ( x ) is non-degenerate on Γ (by Lemma 6.2);
Show Assumption 3.2 is satisﬁed,
i
e., Φ( U ) ⊂
Γ (by Lemma 6.3); 3.
Show the limiting ﬂow (13) converges to the minimizer of the regularizer (by Lemma 6.5); 4.
Show the minimizer of the regularizer recovers the groundtruth (by Lemma 6.6).
Our setting is more general than HaoChen et al.
(2020), which assumes w ∗ ∈ { 0 , 1 } d and their reparametrization can only express positive linear functions, i.e., w = u (cid:12) 2 .
Their (cid:101) O
( κ 2 ) rate is achieved with a delicate three phase LR schedule, while our O ( κ ln d ) rate only uses a constant LR.
Setting: Let { ( z i , y i ) } i ∈ [ n ] be the training dataset where z 1 , . . .
, z n i.i.d.
∼ Unif( {± 1 } d ) or N (0 , I d ) and each y i = (cid:104) z i , w ∗ (cid:105) for some unknown w ∗ ∈ R d .
We assume that w ∗ is κ -sparse for some κ < d .
Denote x = (cid:0) uv (cid:1) ∈
R D
= R 2 d , and we will use x and ( u, v ) exchangeably as the parameter of functions deﬁned on R D in the sequel.
For each i ∈ [ n ] , deﬁne f i ( x ) =
f i ( u, v ) =
z (cid:62) i ( u (cid:12) 2 − v (cid:12) 2 ) .
Then we ﬁt { ( z i , y i ) } i ∈ [ n ] with an overparametrized model through the following loss function:
[Equation]
, where (cid:96) i ( u,
[Equation]
= 14 (
[Equation]
( u, [Equation] − [Equation]
[Equation]
To simplify the notation, we deﬁne Z =
( z 1 , . . .
, z n )
(cid:62) ∈ R n × d and Y =
( y 1 , . . .
, y n ) (cid:62) ∈ R n .
Consider the following manifold:
[Equation]
where U =
( R \
[Equation] [Equation]
We verify that the above loss function L and manifold Γ satisfy Assumption 3.1 by Lemma 6.2, and that the neighborhood U and Γ satisfy Assumption 3.2 by Lemma 6.3.
Lemma 6.2.
Consider the loss function L deﬁned in (14) and manifold Γ deﬁned in (15) .
If data is full rank, i.e., rank( { z i } i ∈ [ n ] ) = n , then it holds that (a).
Γ is a smooth manifold of dimension
D − n ; (b).
rank( ∇ 2 L ( x ))
= n for all x ∈ Γ .
In particular, it holds that rank( { z i } i ∈ [ n ] ) = n holds with probability 1 for Gaussian distribution and with probability 1 − c d for Boolean distribution for some constant c ∈ (0 , 1) .
Lemma 6.3.
Consider the loss function L deﬁned in (14) , manifold Γ and its open neighborhood deﬁned in (15) .
For gradient ﬂow d x t d t = −∇ L ( x t ) starting at any x 0 ∈ U , it holds that Φ( x 0 ) ∈ Γ .
Remark 6.4.
In previous works (Woodworth et al., 2020; Azulay et al., 2021), the convergence of gradient ﬂow is only assumed. Recently Pesme et al. (2021) proved it for a speciﬁc initialization, i.e. , u j = v j = α, ∀ j ∈ [ n ] for some α > 0 . Lemma 6.3 completely removes the technical assumption.
Therefore, by the result in the previous section, the implicit regularizer on the manifold is R ( x ) = 14 n (cid:80) ni =1 (cid:107)∇ f i ( x ) (cid:107) 22 .
To obtain the close-form expression of R ( x ) , it is straightforward to compute ∇ f i ( x ) = 2 (cid:0)
z i (cid:12) u − z i (cid:12) v (cid:1) and ∇ (cid:96) i ( x ) =
( f i ( u, v ) − y i ) (cid:0) z i (cid:12) u − z i (cid:12) v (cid:1) for each i ∈ [ n ] .
Hence, it follows that
[Equation]
The limiting behavior of label noise SGD is described by a Riemannian gradient ﬂow on Γ as follows:
[Equation]
The goal is to show that the above limiting ﬂow will converge to the underlying groundtruth x ∗ = (cid:0) u ∗ v ∗ (cid:1) where ( u ∗ , v ∗ ) =
([ w ∗ ] (cid:12) 1 / 2 + , [ − w ∗ ] (cid:12) 1 / 2 + ) .
6. L IMITING F LOW C ONVERGES TO M INIMIZERS OF R EGULARIZER
In this subsection we show limiting ﬂow (13) starting from anywhere on Γ converges to the minimizer of regularizer R (by Lemma 6.5).
The proof contains two parts: (a) the limiting ﬂow converges; (b) the limit point of the ﬂow cannot be sub-optimal stationary points.
These are indeed the most technical and difﬁcult parts of proving the O ( κ ln d ) upper bound, where the difﬁculty comes from the fact that the manifold Γ is not compact, and the stationary points of the limiting ﬂow are in fact all located on the boundary of Γ .
However, the limiting ﬂow itself is not even deﬁned on the boundary of the manifold Γ .
Even if we can extend ∂ Φ( · ) ∇ R ( · ) continuously to entire R D , the continuous extension is not everywhere differentiable.
Thus the non-compactness of Γ brings challenges for both (a) and (b).
For (a), the convergence for standard gradient ﬂow is often for free, as long as the trajectory is bounded and the objective is analytic or smooth and semialgebraic.
The latter condition ensures the so-called Kurdyka-Łojasiewicz (KL) inequality (Lojasiewicz, 1963), which implies ﬁnite trajectory length and thus the convergence.
However, since our ﬂow does not satisfy those nice properties, we have to show that the limiting ﬂow satisﬁes Polyak-Łojasiewicz condition (a special case of KL condition)
(Polyak, 1964) via careful calculation (by Lemma D.16).
For (b), the standard analysis based on center stable manifold theorem shows that gradient descent/ﬂow converges to strict saddle (stationary point with at least one negative eigenvalue in hessian) only for a zero-measure set of initialization (Lee et al., 2016; 2017).
However, such analyses cannot deal with the case where the ﬂow is not differentiable at the sub-optimal stationary point.
To circumvent this issue, we prove the non-convergence to sub-optimal stationary points with a novel approach: we show that for any stationary point x , whenever there exists a descent direction of the regularizer R at x , we can construct a potential function which increases monotonically along the ﬂow around x , while the potential function is equal to −∞ at x , leading to a contradiction.
(See proof of Lemma 6.5.)
Lemma 6.5.
Let { x t } t ≥ 0 ⊆
R D be generated by the ﬂow deﬁned in (17) with any initialization x 0 ∈ Γ .
Then x ∞ =
lim t →∞ x t exists.
Moreover, x ∞ =
x ∗ is the optimal solution of (18) .
6. M INIMIZER OF THE R EGULARIZER R ECOVERS THE S PARSE G ROUNDTRUTH
{+Note 1 n (cid:80) n i =
 ( z j i ) 2 = 1 when z j i iid ∼ Unif {− 1 , 1 } ,+} we can show minimizing R ( x ) on Γ , (18), is equivalent to ﬁnding the minimum (cid:96) 1 norm solution of Equation (14).
Standard results in sparse recovery imply that minimum (cid:96) 1 norm solution recovers with the sparse groundtruth. The gaussian case is more complicated but still can be proved with techniques from Tropp (2015).
[Equation]
subject to Z ( u (cid:12) [Equation] − v
[Equation]
Lemma 6.6.
Let z 1 , . . .
, z n i.i.d.
∼ Unif( {± 1 } d ) or N (0 , I d ) .
Then there exist some constants C, c >such that if n ≥ Cκ ln d , then with probability at least 1 − e − cn , the optimal solution of (18) , ( (cid:98) u, (cid:98) v ) , is unique up to sign ﬂips of each coordinate and recovers the groundtruth, i.e., (cid:98) u
(cid:12) 2 − (cid:98) v [Equation] = w ∗ .
6. L OWER B OUND FOR G RADIENT D ESCENT IN THE K ERNEL R EGIME
In this subsection we show GD needs at least Ω( d ) samples to learn OLM, when initialized in the kernel regime.
This lower bound holds for all learning rate schedules and numbers of steps.
This is in sharp contrast to the (cid:101) O
( κ ln d ) sample complexity upper bound of SGD with label noise.
Following the setting of kernel regime in (Woodworth et al., 2020), we consider the limit of u 0 = v 0 =
α 1 , with α → ∞ .
It holds that f i ( u 0 , v 0 )
=
and ∇ f i ( u 0 , v 0 )
=
[ αz i , − αz i ] for each i ∈ [ n ] .
Standard convergence analysis for NTK (Neural Tangent Kernel, Jacot et al.
(2018)) shows that upon convergence, the distance traveled by parameter converges to 0 , and thus the learned model shall converge in function space, so is the generalization performance.
For ease of illustration, we directly consider the lower bound for test loss when the NTK is ﬁxed throughout the training.
Theorem 6.7.
Assume z 1 , . .
, z n i.i.d.
∼ N (0 , I d ) and y i = w (cid:62) ∗
z i , for all i ∈ [ n ] .
Deﬁne the loss with linearized model as L ( x ) = (cid:80) ni
=
( f i ( x 0 )
+ (cid:104)∇ f i ( x 0 ) , x − x 0 (cid:105) − y i ) 2 , where x = (cid:0) uv (cid:1) and x 0 = (cid:0) u 0 v 0 (cid:1) =
α (cid:0) 11 (cid:1) .
Then for any ground truth w ∗ , any learning rate schedule { η t } t ≥ 1 , and any ﬁxed number of steps T , the expected (cid:96) 2 loss of x T is at least (1 − nd ) (cid:107) w ∗ (cid:107) 22 , where x T is the T -th iterate of GD on L , i.e., x t +
= x
t − η t ∇ L ( x t ) , for all t ≥ 0 .
C ONCLUSION AND F UTURE W ORK
We propose a mathematical framework to study the implicit bias of SGD with inﬁnitesimal LR.
We show that with arbitrary noise covariance, Θ( η − 2 ) steps of SGD converge to a limiting diffusion on certain manifold of local minimizer, as the LR η → 0 .
For speciﬁc noise types, this allows us to recover and strengthen results regarding implicit bias in previous works with much simpler analysis.
In particular, we show a sample complexity gap between label noise SGD and GD in the kernel regime for a overparametrized linear model, justifying the generalization beneﬁt of SGD.
For the future work, we believe our framework can be applied to analyze the implicit bias of SGD in more complex models towards better understanding of the algorithmic regularization induced by stochasticity.
It will be valuable to extend our method to other stochastic optimization algorithms, e.g., ADAM, SGD with momentum.
A P RELIMINARIES ON S TOCHASTIC P ROCESSES
In this section we review a few basics of stochastic processes that will be useful for proving our results, so that our paper will be self-contained.
We refer the reader to classics like Karatzas & Shreve (2014);
Billingsley (2013); Pollard (2012) for more systematic derivations.
Throughout this section, let E be a Banach space equipped with norm (cid:107) · (cid:107) , e.g., ( R , | · | ) and ( R D , (cid:107) · (cid:107) 2 ) .
A. C ADLAG F UNCTION AND M ETRIC
Deﬁnition A.1 (Cadlag function) .
Let T ∈ [0 , ∞ ] .
A function g : [0 , T ) → E is cadlag 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 cadlag 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 A.2 (Continuity modulus) .
For any function f : [0 , ∞ )
→ E and any interval I ⊆
[0 , ∞ ) , we deﬁne
[Equation]
For any N ∈ N and θ > 0 , we further deﬁne the continuity modulus of continuous f as
[Equation]
Moreover, the continuity modulus of cadlag f ∈ D E [0 , ∞ ) is deﬁned as [Equation]
[Equation]
Deﬁnition A.3 (Jump) .
For any g ∈ D E [0 , T ) , we deﬁne the jump of g at t to be ∆ g ( t ) =
g ( t ) − g ( t − ) .
For any δ > 0 , we deﬁne h δ : [0 , ∞ )
→ [0 , ∞ ) by
[Equation]
We then further deﬁne J δ : D R D [0 , ∞ )
→ D R D [0 , ∞ ) (Katzenberger, 1991) as
[Equation]
Deﬁnition A.4 (Skorohod 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 0 ≤ 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 (cid:107) f
( t ) − g ( s ) (cid:107) ≤ δ if ( t, s ) ∈ [ t i , t i +1 ) ×
[ s i , s i +1 ) for i = 0 , . .
. , k − 1 .
The Skorohod metric on D E [0 , ∞ ) is deﬁned to be d ( f, g )
= ∞ (cid:88)
T = 2 − T min { 1 , d T ( f, g ) } .
A. S TOCHASTIC P ROCESSES AND S TOCHASTIC I NTEGRAL
Let (Ω , F , {F t } t ≥ 0 , P ) be a ﬁltered probability space.
Deﬁnition A.5 (Cross variation) .
Let X
and Y be two {F t } t ≥ 0 -adapted stochastic processes such that X has sample paths in D R D × e [0 , ∞ ) and Y has samples paths in D R e [0 , ∞ ) , then the cross variation of X and Y on (0 , t ] , denoted by [ X, Y ]( t ) , is deﬁned to be the limit of m − 1 (cid:88)
[Equation]
in probability as the mesh size of 0 = t 0 < t 1 < · · · < t m = t goes to 0, if it exists.
Moreover, for Y itself, we write
[Equation]
Deﬁnition A.6 (Martingale) .
Let { X ( t ) } t ≥ 0 be a {F t } t ≥ 0 -adapted stochastic process.
If for all 0 ≤ s ≤ t , it holds that
[Equation]
then X is called a martingale.
Deﬁnition A.7 (Semimartingale) .
Let { X ( t ) } t ≥ 0 be a {F t } t ≥ 0 -adapted stochastic process.
If there exists a sequence of {F t } t ≥ 0 -stopping time, { τ k } k ≥ 0 , such that
• P [ τ k < τ k +1 ]
= 1 , P [lim k →∞ τ k = ∞ ]
= 1 , • and { X τ k ( t ) } t ≥ 0 is a {F t } t ≥ 0 -adapted martingale,
then X is called a semimartingale.
Lemma A.8 (It ˆ
o’s Lemma) .
Let { X ( t ) } t ≥ 0 be deﬁned through the following Itˆo 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]
A. W EAK C ONVERGENCE FOR S TOCHASTIC P ROCESSES
Let ( D E [0 , ∞ ) , A , d ) be a metric space equipped with a σ -algebra A and the Skorohod 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 stochastic process on (Ω , F , P ) with sample paths on D E [0 , ∞ ) .
Deﬁnition A.9 (Weak convergence) .
A sequence of stochastic process { X n } n ≥ 0 is said to converge 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]
Though we deﬁne weak convergence for a countable sequence of stochastic processes, but it is still valid if we index the stochastic processes by real numbers, e.g., { X η } η ≥ 0 , and consider the weak convergence of X η as η → 0 .
This is because the convergence in (20) is for a sequence of real numbers, which is also well-deﬁned if we replace lim n →∞ by lim η → 0 .
Deﬁnition A.10 ( δ -Prohorov distance) .
Let δ > 0 .
For any two probability measures P and Q on a metric space with metric d , let ( X, Y ) be a coupling such that P is the marginalized law of X and Q that of Y .
We deﬁne
[Equation]
P
[Equation]
X
[Equation] Y [Equation] [Equation]
Deﬁnition A.11 (Prohorov metric) .
For any two probability measures P and Q on a metric space with metric d , let ( X, Y ) be a coupling such that P is the marginalized law of X and Q that of Y .
Denote the marginal laws of X and Y by L ( X ) and L ( Y ) respectively.
We deﬁne the Prohorov metric as
[Equation]
It can be shown that X n ⇒ X is equivalent to lim n →∞ ρ ( X n , X ) = 0 .
Theorem A.12 (Skorohod Representation Theorem) .
Suppose P n , n = 1 , 2 , . . .
and P are probability measures on E such that P n ⇒ P .
Then there is a probability space (Ω , F , P ) on which are deﬁned E -valued random variables X n , n = 1 , 2 , . . .
and X with distributions P n and P respectively, such that lim n →∞ X n = X a.s.
The main convergence result in Katzenberger (1991) (Theorem B.7) are in the sense of Skorohod metric in Deﬁnition A.4, which is harder to understand and use compared to the more common uniform metric (Deﬁnition A.13). However, convergence in Skorohod 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 A.13 (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 A.14 (Problem 7, Section 5, Pollard (2012)) .
If X n ⇒ X in the Skorohod sense, and X has sample paths in C R D [0 , ∞ ) , then X n ⇒ X in the uniform metric.
Remark A.15.
We shall note the uniform metric deﬁned above is weaker than sup t ∈ [0 , ∞ ) (cid:107) f
( t ) − g ( t ) (cid:107) .
Convergence in the uniform metric on [0 , ∞ ] deﬁned in Deﬁnition A.13 is equivalent to convergence in the uniform metric on each compact set [0 , T ] for T ∈ N + .
The same holds for the
Skorohod topology. B L IMITING D IFFUSION OF SGD In this section, we give a complete derivation of the limiting diffusion of SGD.
Here we use ⇒ to denote the convergence in distribution.
For any U ⊆ R D , we denote by ˚ U its interior. For linear space S , we use S ⊥ to denote its orthogonal complement.
First, as mentioned in Assumption 3.2, we verify that the mapping Φ is C 2 in Lemma B.1. In
Appendix B.1 we discuss how different time scalings could affect the coefﬁcients in SDE (2) and (3).
Then we check the necessary conditions for applying the results in Katzenberger (1991) in Appendix B.2 and recap the corresponding theorem for the asymptotically continuous case in Appendix B.3.
Finally, we provide a user-friendly interface for Katzenberger’s theorem in
Appendix B.4. Lemma B.1 (Implication of Falconer (1983)) .
Under Assumption 3.2, Φ is C 2 on U .
Proof of Lemma B.1.
Applying Theorem 5.1 of Falconer (1983) with f ( · ) =
φ ( · , 1) sufﬁces.
Let’s ﬁrst clarify how we derive the SDEs, (2) and (3), that approximate SGD (1) under different time scalings. Recall W ( t ) is Ξ -dimensional Brownian motion and that σ ( X ) : R D → R D × Ξ is a deterministic noise function.
As proposed by Li et al.
(2017), one approach to approximate SGD (1) by SDE is to consider the following SDE:
[Equation]
where the time correspondence is t = kη , i.e., X ( kη )
≈ x η ( k ) .
Now rescale the above SDE by considering (cid:101) X ( t ) =
X ( tη ) , which then yields
[Equation]
Now we deﬁne W (cid:48) ( t ) = 1 √ η W ( tη ) , and it’s easy to verify that W (cid:48) ( t ) is also a Ξ -dimensional brownian motion, which means W (cid:48)
d = W , i.e. , W and W (cid:48) have the same sample paths in C R d [0 , ∞ ) .
Thus
[Equation] where the time correspondence is t = k , i.e., (cid:101) X ( k ) ≈ x η ( k ) .
The above SDE is exactly the same as (2).
Then, to accelerate the above SDE by η − 2 times, let’s deﬁne ¯ X ( t )
= (cid:101) X ( t/η 2 ) .
Then it follows that
[Equation]
Again note that ηW ( t/η 2 ) d = W ( t ) in sample paths and thus is also a Ξ -Brownian motion.
Here the time correspondence is t = kη 2 , i.e., evolving for constant time with the above SDE approximates Ω(1 /η 2 ) steps of SGD.
In this way, we derive SDE (3) in the main context.
B. N ECESSARY C ONDITIONS
Below we collect the necessary conditions imposed on { Z n } n ≥ 1 and { A n } n ≥ 1 in Katzenberger (1991).
Condition B.2.
The integrator sequence { A n } n ≥ 1 is asymptotically continuous : sup t> 0 |
A n ( t ) −
A n ( t − ) | ⇒
where A n ( t − ) =
lim s → t − A n ( s ) is the left limit of A n at t .
Condition B.3.
The integrator sequence { A n } n ≥ 1 increases inﬁnitely fast : ∀ (cid:15) > 0 , inf t ≥
( A n ( t + (cid:15) ) −
A n ( t )) ⇒ ∞ .
Condition B.4 (Eq.(5.1), Katzenberger 1991) .
For every T > 0 , as n → ∞ , it holds that sup 0 <t ≤ T ∧ λ n ( K ) (cid:107) ∆ Z n ( t ) (cid:107) 2 ⇒ 0 .
Condition B.5 (Condition 4.2, Katzenberger 1991) .
For each n ≥ 1 , let Y n be a {F nt } - semimartingale with sample paths in D R D [0 , ∞ ) .
Assume that for some δ > 0 (allowing δ = ∞ ) and every n ≥ 1 there exist stopping times { τ mn | m ≥ 1 } and a decomposition of Y n − J δ ( Y n ) into a local martingale M n plus a ﬁnite variation process F n such that P [ τ mn ≤ m ] ≤ 1 /m , { [ M n ]( t ∧ τ mn )
+ T t ∧ τ mn ( F n ) } n ≥ 1 is uniformly integrable for every t ≥ 0 and m ≥ 1 , and
[Equation]
for every (cid:15) > 0 and T > 0 , where T t ( · ) denotes total variation on the interval [0 , t ] .
Lemma B.6.
For SGD iterates deﬁned using the notation in Lemma 4.2, the sequences { A n } n ≥and { Z n } n ≥ 1 satisfy Condition B.2, B.3, B.4 and B.5.
Proof of Lemma B.6.
Condition B.2 is obvious from the deﬁnition of { A n } n ≥ 1 .
Next, for any (cid:15) > 0 and t ∈ [0 , T ] , we have
[Equation]
which implies that inf 0 ≤ t ≤ T ( A n ( t + (cid:15) )
− A n ( t )) > (cid:15)/ (2 η n ) for small enough η n .
Then taking n → ∞ yields the Condition B.3.
For Condition B.4, note that
[Equation]
Therefore, we have (cid:107) ∆
Z n ( t ) (cid:107) 2 ≤ 2 η n for all t > 0 .
This implies that (cid:107) ∆
Z n ( t ) (cid:107) 2 → 0 uniformly over t > 0 as n → ∞ , which veriﬁes Condition B.4.
We proceed to verify Condition B.5.
By the deﬁnition of Z n , we know that { Z n ( t ) } t ≥ 0 is a jump process with independent increments and thus is a martingale.
Therefore, by decomposing
Z n = M n + F n with M n being a local martingale and F n a ﬁnite variation process, we must have
F n = 0 and M n is Z n itself.
It then sufﬁces to show that [ M n ]( t ∧ τ mn ) is uniformly integrable for every t ≥ 0 and m ≥ 1 .
Since M n is a pure jump process, we have
[Equation]
This implies that [ M η ]( t ∧ τ mη ) is universally bounded by 4 t , and thus [ M η ]( t ∧ τ mη ) is uniformly integrable.
This completes the proof.
Lemma 4.2.
Let { η n } ∞ n =1 be any positive sequence with lim n →∞ η n = 0 , A n ( t ) =
η n (cid:98) t/η 2 n (cid:99) , and
Z n ( t ) =
η n (cid:80) (cid:98) t/η 2 n (cid:99) k
=1 ( 1 ξ k − 1Ξ 1 ) , where ξ 1 , ξ 2 , . . .
i.i.d.
∼ Unif ([Ξ]) .
Then with the same initialization
X n (0) =
x η n (0) ≡ X (0) , X n ( kη 2 n ) deﬁned by (6) is a Katzenberger process and is equal to x η n ( k ) deﬁned in (1) with LR equal to η n for all k ≥ 1 .
Moreover, the counterpart of (7) is
[Equation]
where Σ ≡ σσ (cid:62) and { W ( t ) } t ≥ 0 is a Ξ -dimensional standard Brownian motion.
Proof of Lemma 4.2.
For any n ≥ 1 , it sufﬁces to show that given X n ( kη 2 n ) =
x η n ( k ) , we further have X n (( k + 1) η 2 n ) =
x η n ( k + 1) .
By the deﬁnition of X n ( t ) , we have
[Equation]
where the second equality is because A n ( t ) and Z n ( t ) are constant on interval [ kη 2 n , ( k +1)
η 2 n ) .
This conﬁrms the alignment between { X n ( kη 2 n ) } k ≥ 1 and { x η n ( k ) } k ≥ 1 .
For the second claim, note that σ ( x ) E Z n ( t ) ≡ 0 for all x ∈ R D , t ≥ 0 (since the noise has zeroexpectation) and that { Z n ( t ) − E Z n ( t ) } t ≥ 0 will converge in distribution to a Brownian motion by the classic functional central limit theorem (see, for example, Theorem 4.3.5 in Whitt (2002)).
Thus, the limiting diffusion of X n as n → ∞ can be obtained by substituting Z with the standard Brownian motion W in (22).
This completes the proof.
B. K ATZENBERGER ’ S T HEOREM FOR A SYMPTOTICALLY C ONTINUOUS C ASE
The full Katzenberger’s theorem deals with a more general case, which only requires the sequence of intergrators to be asymptotically continuous , thus including SDE (3) and SGD (1) with η goes to 0 .
To describe the results in Katzenberger (1991), we ﬁrst introduce some deﬁnitions.
For each n ≥ 1 , let (Ω n , F n , {F nt } t ≥ 0 , P ) be a ﬁltered probability space, Z n an R e -valued cadlag {F nt } -semimartingale with Z n (0)
=
and A n a real-valued cadlag {F nt } -adapted nondecreasing process with A n (0) = 0 .
Let σ n : U → M ( D, e ) be continuous with σ n → σ uniformly on compact subsets of U .
Let X n be an R D -valued cadlag {F nt } -semimartingale satisfying, for all compact K ⊂ U ,
[Equation]
[Equation] for all t ≤ λ n ( K ) where λ n ( K ) = inf { t ≥ 0 |
X n ( t − ) / ∈ ˚ K or X n ( t ) / ∈ ˚ K } is the stopping time of X n leaving K .
Theorem B.7 (Theorem 6.3, Katzenberger 1991) .
Suppose X (0) ∈
U , Assumptions 3.1 and 3.2,
Condition B.2, B.3, B.4 and B.5 hold.
For any compact K ⊂ U , deﬁne µ n ( K ) =
inf { t ≥
|
Y n ( t − ) / ∈ ˚ K or Y n ( t ) / ∈ ˚ K } , then the sequence { ( Y µ n ( K ) n , Z µ n ( K ) n , µ n ( K ) } is relatively compact in D R D × e [0 , ∞ ) ×
[0 , ∞ ) .
If ( Y, Z, µ ) is a limit point of this sequence, then ( Y, Z ) is a continuous semimartingale, Y ( t )
∈ Γ for every t ≥
a.s., µ ≥ inf { t ≥ 0 | Y ( t ) / ∈ ˚ K }
a.s. and Y ( t ) admits
[Equation]
We note that by Lemma A.14, convergence in distribution under skorohod metric is equivalent to convergence in distribution under uniform metric Deﬁnition A.13, therefore in the rest of the paper we will only use the uniform metric in the rest of the paper, e.g., whenever we mention Prohorov metric and δ -Prohorov distance, the underlying metric is the uniform metric.
B. A U SER - FRIENDLY I NTERFACE FOR K ATZENBERGER ’ S T HEOREM
Based on the Lemma B.6, we can immediately apply Theorem B.7 to obtain the following limiting diffusion of SGD.
Theorem B.8.
Let the manifold Γ and its open neighborhood U satisfy Assumptions 3.1 and 3.2.
Let K ⊂ U be any compact set and ﬁx some x 0 ∈ K .
Consider the SGD formulated in Lemma 4.2 where
X η n (0) ≡ x 0 .
Deﬁne
[Equation]
and µ η n ( K ) = min { t ∈ N | Y η n ( t ) / ∈ ˚ K } .
Then the sequence { ( Y µ ηn
[Equation] K ) [Equation] , Z η n , µ η n ( K )) } n ≥ 1 is relatively compact in D R D × R n [0 , ∞ ) ×
[0 , ∞ ] .
Moreover, if ( Y, Z, µ ) is a limit point of this sequence, it holds that Y ( t ) ∈ Γ a.s for all t ≥ 0 , µ ≥ inf { t ≥
| Y ( t ) / ∈ ˚ K } and Y ( t ) admits Y
[Equation]
where { W ( s ) } s ≥ 0 is the standard Brownian motion and σ ( · ) is as deﬁned in Lemma 4.2.
However, the above theorem is hard to parse and cannot be directly applied if we want to further study the implicit bias of SGD through this limiting diffusion.
Therefore, we develop a user-friendly interface to it in below.
In particular, Theorem 4.6 is the a special case of Theorem B.9.
In Theorem 4.6, we replace ∂ Φ( Y ( t ))
σ ( Y ( t )) to Σ 12 (cid:107) ( Y ( t )) to simplify the equation, since ∂ Φ( Y ( t ))
σ ( Y ( t )) ( ∂ Φ( Y ( t ))
σ ( Y ( t ))) (cid:62) =
Σ (cid:107) ( Y ( t )) and thus this change doesn’t affect the distribution of the sample paths of the solution.
Since σ ( · ) is locally Lipschitz, when restricted on any compact K ⊂ U , the solution to (23) always exists and is unique.
Let µ ( K ) =
inf { t ≥ 0 | Y ( t ) /
∈ ˚ K } be the escaping time of the limiting diffusion Y .
Theorem B.9.
Under the same setting as B.8, we change the integer index back to η > 0 with a slight abuse of notation.
For any compact set K ⊆ U and T > 0 , let δ =
P ( µ ( K ) ≤ T ) .
Then for any (cid:15) > 0 , it holds for all sufﬁciently small LR η that:
[Equation]
Moreover, when Y is a global solution of limiting diffusion Equation (23) and Y never leaves U , i.e. P [ ∀ t ≥ 0 , Y ( t ) ∈ U ]
= 1 , it holds that Y Tη converges in distribution to Y T as η → 0 for any ﬁxed
T > 0 .
Proof of the ﬁrst claim of Theorem B.9.
Let E T be the event such that µ ( K ) >
T on E T .
Then restricted on E T , we have Y ( T ∧ µ ) =
Y ( T ∧ µ ( K )) as µ ≥ µ ( K ) holds a.s.
We ﬁrst prove the claim for any convergent subsequence of η .
Now, let { η m } m ≥ 1 be a sequence of LRs such that η m → 0 and Y µ ηm ( K ) η m ⇒ Y µ as m → ∞ .
By applying the Skorohod representation theorem, we can put { Y η m } m ≥ 1 and Y under the same probability space and Y µ ηm ( K ) η m → Y µ a.s. in the Skorohod metric, i.e., d
[Equation]
Y µ [Equation] K [Equation]
η m [Equation] Y µ
[Equation] → 0 , a.s., which further implies that for any (cid:15) > 0 , there exists some N > 0 such that for all m > N
[Equation]
Restricted on E T , we have d ( Y µ ηm ( K ) ∧
T η m , Y µ ∧ T ) = d ( Y µ ηm ( K ) ∧ T η m , Y µ ( K ) ∧ T ) , and it follows that for all m > N
[Equation]
By the deﬁnition of the Prohorov metric in Deﬁnition A.11, we then get ρ 2 δ ( Y µ ηm ( K ) η m , Y µ ( K ) ∧ T ) ≤ (cid:15) for all m > N .
Therefore, we have
[Equation]
Now we claim that it indeed holds that lim η → 0 ρ 2 δ ( Y µ η ( K ) η , Y µ ( K ) ∧ T ) = 0 .
We prove this by contradiction.
Suppose otherwise, then there exists some (cid:15) > 0 such that for all η 0 > 0 , there exists some η < η 0 with ρ 2 δ ( Y µ η ( K ) η , Y µ ( K ) ∧ T ) > (cid:15) .
Consequently, there is a sequence { η m } m ≥ 1 satisfying lim m →∞ η m =
and ρ 2 δ ( Y µ ηm ( K ) η m , Y µ ( K ) ∧ T ) > (cid:15) for all m .
Since { ( Y µ ηm ( K ) η m , Z η m , µ η m ( K )) } m ≥ 1 is relatively compact, there exists a subsequence (WLOG, assume it is the original sequence itself) converging to ( Y µ ∧ T , W, µ ) in distribution.
However, repeat the exactly same argument as above, we would have ρ 2 δ ( Y µ ηm ( K ) η m , Y µ ( K ) ∧ T ) ≤ (cid:15) for all sufﬁciently large m , which is a contradiction.
This completes the proof.
Proof of the second claim of Theorem B.9.
We will ﬁrst show there exists a sequence of compact set { K m } m ≥ 1 such that ∪ ∞ m =1 K m =
U and K m ⊆
K m +1 .
For m ∈ N + , we deﬁne H m =
U \ ( B 1 /m (0)+ R D \ U ) and K m =
H m ∩ B m (0) .
By deﬁnition it holds that ∀ m < m (cid:48) , H m ⊆
H m (cid:48) and K m ⊆
K m (cid:48) .
Moreover, since K m is bounded and closed, K m is compact for every m .
Now we claim ∪ ∞ m =1 K m =
U .
Note that ∪ ∞ m =1 K m = ∪ ∞ m =1 H m ∩ B m (0) =
∪ ∞ m =1 H m .
∀ x ∈ U , since U is open, we know d U ( x, R D \ U ) > 0 , thus there exists m 0 ∈ N + , such that ∀ m ≥ m 0 , x / ∈ ( B 1 /m (0) + R D \ U ) and thus x ∈ H m , which implies x ∈ ∪ ∞ m =
1 H m .
On the other hand, ∀ x ∈ R D \ U , it holds that x ∈ ( B 1 /m (0) + R D \ U ) for all m ∈ N + , thus x / ∈ H m ⊂ K m .
Therefore, since Y ∈ U and is continuous almost surely, random variables lim
m →∞ µ ( K m )
= ∞ a.s., which implies µ ( K m ) converges to ∞ in distribution, i,e,, ∀ δ > 0 , T > 0 , ∃ m ∈ N + , such that ∀ K ⊇ K m , it holds P [ µ ( K ) ≤ T ] ≤ δ .
Now we will show for any T > 0 and (cid:15) > 0 , there exists η 0 such that ρ (cid:15) ( Y T , Y Tη ) ≤ (cid:15) for all η ≤ η 0 .
Fixing any T > 0 , for any (cid:15) > 0 , let δ = (cid:15) 4 , then from above we know exists compact set K , such that P ( µ ( K ) ≤ T ) ≤ δ .
We further pick K (cid:48)
=
K + B 2 (cid:15) (cid:48) (0) , where (cid:15) (cid:48) can be any real number satisfying 0 < (cid:15) (cid:48) < (cid:15) and K (cid:48) ⊆
U .
Such (cid:15) (cid:48) exists since U is open.
Note K ⊆ K (cid:48) , we have P ( µ ( K (cid:48) ) ≤ T ) ≤ P ( µ ( K ) ≤ T ) ≤ δ .
Thus by the ﬁrst claim of Theorem B.9, there exists η 0 > 0 , such that for all η ≤ η 0 , we have ρ 2 δ ( Y µ η ( K (cid:48) )
∧ T η
[Equation]
Note that ρ δ ( Y µ ( K ) ∧ T , Y µ ( K (cid:48) )
∧ T )
= 0 , so we have for all η ≤ η 0 ,
[Equation]
By the deﬁnition of δ -Prohorov distance in Deﬁnition A.10, we can assume ( Y µ ( K ) ∧ T , Y µ η ( K (cid:48) )
∧ T η ) is already the coupling such that P (cid:104) d U ( Y µ ( K ) ∧ T , Y µ η ( K (cid:48) )
∧ T η ) ≥ 2 −(cid:100) T (cid:101) (cid:15) (cid:48) (cid:105) ≤ 3 δ .
Below we want to show ρ 3 δ ( Y µ ( K ) ∧ T , Y Tη ) ≤ 2 −(cid:100) T (cid:101) (cid:15) (cid:48) .
Note that for all t ≥ 0 , Y µ ( K ) ∧ T ( t ) ∈ K , thus we know if µ η ( K (cid:48) ) ≤ T , then
[Equation]
On the other hand, if µ η ( K (cid:48) ) > T , then Y Tη = Y µ η ( K (cid:48) )
∧ T η . Thus we can conclude that d U ( Y µ ( K ) ∧ T , Y Tη ) ≥ 2 −(cid:100) T (cid:101) (cid:15) (cid:48) implies d U ( Y µ ( K ) ∧ T , Y µ η ( K (cid:48) )
∧ T η ) ≥ 2 −(cid:100) T (cid:101) (cid:15) (cid:48) . Therefore, we further have [Equation]
that is,
[Equation]
Finally, since ρ δ ( Y T , Y µ ( K ) ∧ T ) = 0 , we have for all η ≤ η 0 , [Equation] [Equation] +
[Equation]
0 ≤
[Equation]
, which completes the proof.
Remark B.10. Our result suggests that for tiny LR η , SGD dynamics have two phases. In Phase I of Θ(1 /η ) steps, the SGD iterates move towards the manifold Γ of local minimizers along GF. Then in Phase II which is of Θ(1 /η 2 ) steps, the SGD iterates stay close to Γ and diffuse approximately according to (10) .
See Figure 2 for an illustration of this two-phase dynamics.
However, since the length of Phase I gets negligible compared to that of Phase II when η → 0 , Theorem 4.6 only reﬂects the time scaling of Phase II.
In this section, we demonstrate how to compute the derivatives of Φ by relating to those of the loss function L , and then present the explicit formula of the limiting diffusion.
C. E XPLICIT E XPRESSION OF THE D ERIVATIVES
For any x ∈ Γ , we choose an orthonormal basis of T x (Γ) as { v 1 , . .
. , v D − M } .
Let { v D − M +1 , . .
. , v D } be an orthonormal basis of T ⊥ x (Γ) so that { v i } i ∈
[ D ] is an orthonormal basis of
R D .
Lemma C.1.
For any x ∈ Γ and any v ∈ T x (Γ) , it holds that ∇ 2 L ( x ) v = 0 .
Proof.
For any x ∈ T x (Γ) , let { x ( t ) } t ≥ 0 be a parametrized smooth curve on Γ such that x (0) =
x and d x ( t ) d t (cid:12)(cid:12) t
=0 = v .
Then ∇ L ( x t ) = 0 for all t .
Thus
= d ∇ L ( x t ) d t (cid:12)(cid:12) t =0 = ∇ 2 L ( x ) v .
Lemma C.2.
For any x ∈ R D , it holds that ∂ Φ( x ) ∇ L ( x ) = [-
-]0 and ∂ 2 Φ( x )[ [Equation] ( x [Equation]
, ∇ [Equation] ( x [Equation]
] = −
[Equation]
Φ( x
[Equation]
∇ 2 L ( x
[Equation]
∇ L ( x ) .
Proof.
Fixing any x ∈ R D , let d x ( t ) d t =
−∇ L
x ( t )) be initialized at x (0) =
x .
Since Φ( x ( t ))
=
Φ( x ) for all t ≥ 0 , we have
[Equation]
Evaluating the above equation at t = 0 yields ∂ Φ( x ) ∇ L ( x ) = 0 .
Moreover, take the second order derivative and we have
[Equation]
Evaluating at t = 0 completes the proof.
Now we can prove Lemma 4.3, restated in below.
Lemma 4.3.
For any x ∈ Γ , ∂ Φ( x ) ∈ R D
× D is the projection matrix onto tangent space T x (Γ) .
Proof of Lemma 4.3.
For any v ∈ T x (Γ) , let { v ( t ) , t ≥ 0 } be a parametrized smooth curve on Γ such that v (0) =
x and d v ( t ) d t (cid:12) (cid:12) t =0 = v .
Since v ( t ) ∈ Γ for all t ≥ 0 , we have Φ( v ( t ))
=
v ( t ) , and thus
[Equation]
This implies that ∂ Φ( x ) v = v for all v ∈ T x (Γ) .
Next, for any u ∈ T ⊥ x (Γ) and t ≥ 0 , consider expanding ∇ L ( x + t ∇ 2 L ( x ) † u ) at t = 0 :
[Equation]
where the second equality follows from the assumption that ∇ 2 L ( x ) is full-rank when restricted on
T ⊥ x (Γ) .
Then since ∂ Φ is continuous, it follows that
[Equation]
By Lemma C.2, we have ∂ Φ( x + t ( ∇ 2 L ( x )) † u )) ∇ L ( x + t ( ∇ 2 L ( x )) † u ) = 0 for all t > 0 , which then implies that ∂ Φ( x ) u = 0 for all u ∈ T ⊥ x
(Γ) .
Therefore, under the basis { v i , . .
, v N } , ∂ Φ( x ) is given by
[Equation]
that is, the projection matrix onto T x (Γ) .
Lemma C.3.
For any x ∈ Γ , it holds that ∂ Φ( x ) ∇ 2 L ( x ) = 0 .
Proof.
It directly follows from Lemma C.1 and Lemma 4.3.
Next, we proceed to compute the second-order derivatives.
Lemma C.4.
For any x ∈ Γ , u ∈ R D and v ∈ T x
(Γ)
, it holds that
[Equation]
x ) † ∂ 2 ( ∇ L )( x )[ v, ∂ Φ( x ) u ] .
Proof of Lemma C.4.
Consider a parametrized smooth curve { v ( t ) } t ≥ 0 on Γ such that v (0) =
x and d v ( t ) d t (cid:12)(cid:12) t
=0 = v .
We deﬁne P ( t ) =
∂
Φ( v ( t )) , P ⊥ ( t ) =
I D − P ( t ) and H ( t ) = ∇ 2 L ( v ( t )) for all t ≥ 0 .
By Lemma C.1 and 4.3, we have
[Equation]
Denote the derivative of P ( t ) , P ⊥ ( t ) and H ( t ) with respect to t as P (cid:48) ( t ) , ( P ⊥ ) (cid:48) ( t ) and H (cid:48) ( t ) .
Then differentiating with respect to t , we have
[Equation]
Then combining (25) and (26) and evaluating at t = 0 , we have
[Equation]
We can decompose P (cid:48) (0) and H (0) as follows
[Equation]
where P (cid:48) 11 (0) ∈ R ( D − M ) ×
( D − M ) and H 22 is the hessian of L restricted on T ⊥ x (Γ) .
Also note that
[Equation]
and thus by (28) we have
[Equation]
This implies that we must have P (cid:48) 22 (0)
=
and P (cid:48) 12 (0) H 22 (0)
=
H (cid:48) 12 (0) .
Similarly, by taking transpose in (28), we also have H 22 (0) P (cid:48) 21 (0) =
− H (cid:48) 21 (0) .
It then remains to determine the value of P (cid:48) 11 (0) .
Note that since P ( t ) P ( t ) =
P ( t ) , we have
P (cid:48) ( t ) P ( t ) + P
[Equation]
t ) P (cid:48)
[Equation]
) , evaluating at t = 0 yields
[Equation]
Therefore, we must have P (cid:48) 11 (0) = 0 .
Combining the above results, we obtain
[Equation]
Finally, recall that P ( t ) =
∂
Φ( v ( t )) , and thus
[Equation]
Similarly, we have H (cid:48) (0) =
∂ 2 ( ∇ L )( x )[ v ] , and it follows that
[Equation]
Lemma C.5.
For any x ∈ Γ and u ∈ T ⊥ x (Γ) , it holds that
[Equation]
Proof of Lemma C.5.
For any u ∈ T ⊥ x (Γ)
[Equation]
we deﬁne u ( t ) =
x + t ∇ 2 L ( x ) † u for t ≥ 0 .
By Taylor approximation
[Equation]
we have
[Equation]
and
[Equation]
Combine (29) and (30) and apply Lemma C.
[Equation]
and it follows that
[Equation]
where the last equality follows from Lemma C.3.
Dividing both sides by t 2 and letting t → 0 , we get
[Equation]
Rearranging the above equation completes the proof.
With the notion of Lyapunov Operator in Deﬁnition 4.4, Lemma C.5 can be further simpliﬁed into
Lemma C.6.
Lemma C.6.
For any x ∈ Γ and Σ ∈ span { uu (cid:62) | [Equation] T ⊥ x (Γ) } ,
[Equation]
Proof of Lemma C.6.
Let A = uu (cid:62) + ∇ 2 L ( x ) † uu (cid:62) ∇ 2 L ( x ) and B = ∇ 2 L ( x ) † uu (cid:62) .
The key observation is that A + A (cid:62) =
L ∇ 2 L ( x ) ( B + B (cid:62) ) .
Therefore, by Lemma C.5, it holds that ∂ 2 Φ( x )[
L ∇ 2 L ( x ) ( B + B (cid:62) )]
= ∂ 2 Φ( x )[ A + A (cid:62) ] =
[Equation]
∂ Φ( x ) ∂ 2 ( ∇ L )( x )[ B ] = ∂ Φ( x ) ∂ 2 ( ∇ L )( x )[ B + B (cid:62) ] . Since ∇ 2 L ( x ) † is full-rank when restricted to T ⊥ x (Γ) , we have span {∇ 2 L ( x ) † uu (cid:62) + uu (cid:62) ∇ 2 L ( x ) †
| u ∈ T ⊥ x
(Γ) }
= span { uu (cid:62) | u ∈ T ⊥ x (Γ) } .
Thus by the linearity of above equation, we can replace
B + B (cid:62) by any Σ ∈ span { uu (cid:62) | u ∈ T ⊥ x (Γ) } , resulting in the desired equation.
Then Lemma 4.5 directly follows from Lemma C.4 and C.5.
C. T ANGENT N OISE C OMPENSATION ONLY D EPENDENDS ON THE M ANIFOLD I TSELF
Here we show that the second term of (10), i.e., the tangent noise compensation for the limiting dynamics to stay on Γ , only depends on Γ itself.
Lemma C.7.
For any x ∈ Γ , suppose there exist a neighborhood U x of x and two loss functions L and L (cid:48) that deﬁne the same manifold Γ locally in U x , i.e., Γ ∩ U x = { x | ∇ L ( x ) = 0 }
= { x | ∇ L (cid:48)
( x ) = 0 } .
Then for any v ∈ T x (Γ) , it holds that ( ∇ 2 L ( x )) † ∂ 2 ( ∇ L )( x ) [ v, v ] = ( ∇ 2 L (cid:48) ( x )) † ∂ 2 ( ∇ L (cid:48) )
( x ) [ v, v ] .
Proof of Lemma C.7.
Let { v ( t ) } t ≥ 0 be a smooth curve on Γ with v (0) =
x and d v ( t ) d t (cid:12)(cid:12) t
=0 = v .
Since v ( t ) stays on Γ , we have ∇ L ( v ( t ))
=
for all t ≥ 0 .
Taking derivative for two times yields ∂ 2 ( ∇ L )( v ( t ))
[ d v ( t ) d t , d v ( t ) d t ] + ∇ 2 L ( v ( t ))
d 2 v ( t ) d t= 0 .
Evaluating it at t = 0 and multiplying both sides by ∇ 2 L ( x ) † , we get
[Equation]
Since ∂ Φ( x ) is the projection matrix onto T x (Γ) by Lemma 4.3, it does not depend on L , so analogously we also have ∇ 2 L (cid:48)
( x ) † ∂ 2 ( ∇ L (cid:48) )
( x ) [ v, v ] = − ∂ Φ( x ) d 2 v ( t ) d t 2 (cid:12)(cid:12) t
=0 as well.
The proof is thus completed.
Note that ∂
Φ( x ) d 2 v ( t ) d t 2 (cid:12)(cid:12) t =0 is indeed the second fundamental form for v at x , and the value won’t change if we choose another parametric smooth curve with a different second-order time derivative.
(See Chapter 6 in Do Carmo (2013) for a reference.)
C. P ROOF OF RESULTS IN S ECTION
Now we are ready to give the missing proofs in Section 5 which yield explicit formula of the limiting diffusion for label noise and isotropic noise.
Corollary 5.1 (Limiting Diffusion for Isotropic Noise) .
If Σ ≡ I D on Γ , SDE (10) is then
[Equation]
[Equation] ) where |∇ 2 L ( Y ) | + = lim α → 0 |∇ 2 L ( Y )+ αI D | α D − rank( ∇ 2 L ( Y )) is the pseudo-determinant of ∇ 2 L ( Y ) .
|∇ 2 L ( Y ) | + is also equal to the sum of log of non-zero eigenvalue values of ∇ 2 L ( Y ) .
Proof of Corollary 5.1.
Set Σ (cid:107) = ∂
Φ , Σ ⊥ =
I D − ∂ Φ and Σ ⊥ , (cid:107) =
Σ (cid:107) , ⊥ = 0 in the decomposition of Σ by Lemma 4.5, and we need to show ∂ Φ ∇ (ln | Σ | + ) =
∂ 2 ( ∇ L )[( ∇ 2 L ) † ] .
Holbrook (2018) shows that the gradient of pseudo-inverse determinant satisﬁes ∇| A | + = | A | + A † .
Thus we have for any vector v ∈ R D , (cid:10) v, ∇ ln |∇ 2 L | + (cid:11) =
(cid:68) |∇ 2 L | + ∇ 2 L |∇ 2 L | + , ∂ 2 ( ∇ L )[ v ] (cid:69) =
(cid:10) ∇ 2 L, ∂ 2 ( ∇ L )[ v ] (cid:11) =
∂ 2 ( ∇ L )[ v, ∇ 2 L ] = (cid:10) v, ∂ 2 ( ∇ L )[( ∇ 2 L ) † ] (cid:11) , which completes the proof.
Corollary 5.2 (Limiting Flow for Label Noise) .
If Σ ≡ c ∇ 2 L on Γ for some constant c > 0 , SDE (10) can be simpliﬁed into (13) where the regularization is from the noise in normal space.
[Equation]
Proof of Corollary 5.2.
Since Σ = c ∇ 2 L , here we have Σ ⊥ = Σ and Σ (cid:107) , Σ ⊥ , (cid:107) , Σ (cid:107) , ⊥ = 0 .
Thus it sufﬁces to show that 2 ∂ 2 ( ∇ L ) (cid:2)
L − 1 ∇ 2 L (Σ ⊥ ) (cid:3)
= ∇ tr[ ∇ 2 L ] .
Note that for any v ∈ R D , v (cid:62) ∇ tr[ ∇ 2 L ] = (cid:10) I D , ∂ 2 ( ∇ L )[ v ] (cid:11) = (cid:10) I
[Equation]
(32) where the second equality is because the the tangent space of symmetric rank- n matrices at ∇ 2 L is { A ∇ 2 L + ∇ 2 LA (cid:62)
| A ∈ R D × D } , and every element in this tangent space has zero inner-product with ∂ Φ by Lemma 4.3.
Also note that L − 1 ∇ 2 L ( ∇ 2 L ) =
12 ( I D − ∂ Φ) , thus (cid:10) I D − ∂ Φ , ∂ 2 ( ∇ L )[ v ] (cid:11) =
2 (cid:10) L − 1 ∇ 2 L ( ∇ 2 L ) , ∂ 2 ( ∇ L )[ v ] (cid:11) =
2 v (cid:62) ∂ 2 ( ∇ L )
[ L − 1 ∇ 2 L ( ∇ 2 L )] .
C. E XAMPLE : k -P HASE M OTOR
We also give an example with rigorous proof where the implicit bias induced by noise in the normal space cannot be characterized by a ﬁxed regularizer, which was ﬁrst discovered by Damian et al.
but was only veriﬁed via experiments.
Note the normal regularization in both cases of label noise and isotropic noise induces Riemmanian gradient ﬂow against some regularizer, it’s natural to wonder if the limiting ﬂow induced by the normal noise always can always be characterized by certain regularizer.
Interestingly, Damian et al.
(2021) answers this question negatively via experiments in their Section E.2. We adapt their example into the following one, and rigorously prove the limiting ﬂow moves around a cycle at a constant speed and never stops using our framework. Suppose dimension D = k + 2 ≥ 5 . For each x ∈ R D , we decompose x = (cid:0) x 1:2 x 3: D (cid:1) where x 1:2 ∈ Rand x 3: D ∈ R D − 2 .
Let Q θ ∈ R 2 × 2 be the rotation matrix of angle θ , i.e., Q θ = (cid:0) cos θ − sin θ sin θ cos θ (cid:1) and the loss L ( x ) :=
{+18 ( (cid:107) x 1:2 (cid:107) 22 − 1) 2 + 12 (cid:80) Dj =
+}
(2 + (cid:10) Q j − 3 α v, x 1:2 (cid:11) x 2 j , where α = 2 πD − 2 and v is any vector in R 2 with unit norm.
Here the manifold is given by Γ := { x | L ( x ) = 0 } = { x ∈ R D | x 2 1 + x 2 2 = 1 , x j = 0 , ∀ j = 3 , . . .
, D } .
The basic idea is that we can add noise in the ‘auxiliary dimensions’ for j = 3 , . . .
, D to get the regularization force on the circle { x 21 + x 22 = 1 } , and the goal is to make the vector ﬁeld induced by the normal regularization always point to the same direction, say anti-clockwise.
However, this cannot be done with a single auxiliary dimension because from the analysis for label noise, we know when L − 1 ∇ 2 L
(Σ ⊥ ) is identity, the normal regularization term in Equation (10) has 0 path integral along the unit circle and thus it must have both directions.
The key observation here is that we can align the magnitude of noise with the strength of the regularization to make the path integral positive.
By using k ≥ 3 auxiliary dimensions, we can further ensure the normal regularization force is anti-clockwise and of constant magnitude, which is reminiscent of how a three-phase induction motor works.
Lemma C.8.
Let Σ ∈ R D
× D be given by
[Equation]
if i = j ≥ 3 or 0 otherwise, then the solution of SDE (10) is the following (33) , which implies that
Y ( t ) moves anti-clockwise with a constant angular speed of ( D − 2) / 2 .
[Equation]
and
[Equation]
Proof of Lemma C.8.
Note that for any x ∈ Γ , it holds that
[Equation]
Then clearly Σ only brings about noise in the normal space, and speciﬁcally, it holds that
[Equation]
Q D − 3 α v, Q − π/ 2 x 1:2 (cid:11) ) .
Further note that, by the special structure of the hessian in (34) and Lemma C.3, for any x ∈ Γ , we have ∂ Φ( x ) =
( x 2 , − x 1 , 0 , . .
. , 0) (cid:62) ( x 2 , − x 1 , 0 , . . . , 0)
= (cid:0) Q − π/ 2 x 1:2(cid:1)(cid:0) Q − π/ 2 x 1:2(cid:1) (cid:62) .
Combining these facts, the dynamics of the ﬁrst two coordinates in SDE (10) can be simpliﬁed into
[Equation]
where the second to the last equality follows from the property of Q α and the last equality follows from the fact that (cid:107) x 1:2 (cid:107) 2 2 = 1 for all x ∈ Γ .
{+Note we require k ≥ 3 (or D ≥ 5 ) to allow (cid:80) Dj =
+}
{+(cid:10) Q j − 3 α v, Q − π/ 2 x 1:2 (cid:11) 
+}
=
D − 2 2 (cid:13)(cid:13) Q − π/ 2 x 1:2 (cid:13)(cid:13) 2 2 .
On the other hand, we have d x 3: D ( t ) d t = 0 as ∂ Φ kills the movement on that component.
The proof is completed by noting that the solution of x 1:2 is
[Equation]
x
[Equation]
(0) , and by Lemma C.9, exp (cid:18)
t
[Equation]
D
[Equation]
[Equation]
(cid:19) =
(exp( Q
[Equation]
[Equation]
[Equation]
)=
Q t ( D − 2)
[Equation]
=
Q t ( D
[Equation]
[Equation]
Lemma C.9.
{+exp( (cid:0) 
+}
− 1 1 0 (cid:1) )
= (cid:0) cos1 − sin1 sin1 cos1 (cid:1) .
Proof.
By deﬁnition, for matrix A = (cid:0) 0 − 1 1 0 (cid:1) , exp( A ) =
{+(cid:80) ∞ t =
+}
A t t ! .
Note A 2 = −
I , A 3 = − A and
A 4 =
I , and by using this pattern, we can easily check that
[Equation]
− 1)
i 1 (2 i )!
(cid:33) =
(cid:18) cos 1 − sinsincos(cid:19) .
D P ROOF OF RESULTS IN S ECTION
In this section, we present the missing proofs in Section 6 regarding the overparametrized linear model.
For convenience, for any p, r ≥ 0 and u ∈ R D , we denote by B pr ( u ) the (cid:96) p norm ball of radius r centered at u .
We also denote v i : j = ( v i , v i +1 , . . .
, v j ) (cid:62) for i, j ∈
[ D ] .
D. P ROOF OF T HEOREM 6.
In this subsection, we provide the proof of Theorem 6.1.
Theorem 6.1.
In the setting of OLM, suppose the groundtruth is κ -sparse and n ≥ Ω( κ ln d ) training data are sampled from either i.i.d.
Gaussian or Boolean distribution.
Then for any initialization x init (except a zero-measure set) and any (cid:15) > 0 , there exist η 0 , T > 0 such that for any η < η 0 , OLM trained with label noise SGD (12) with LR equal to η for (cid:98) T/η 2 (cid:99) steps returns an (cid:15) -optimal solution, with probability of 1 − e − Ω( n ) over the randomness of the training dataset.
Proof of Theorem 6.1.
First, by Lemma 6.6, it holds with probability at least 1 − e − Ω( n ) that the solution to (18), x ∗ , is unique up to and satisﬁes | x ∗ | = ψ ( w ∗ ) .
Then on this event, for any (cid:15) > 0 , by Lemma 6.5, there exists some T >
such that x T given by the Riemannian gradient ﬂow (17) satisﬁes that x T is an (cid:15)/ 2 -optimal solution of the OLM.
For this T , by Theorem 4.6, we know that the (cid:98) T/η 2 (cid:99) -th SGD iterate, x η ( (cid:98) T/η 2 (cid:99) ) , satisﬁes (cid:107) x η ( (cid:98) T/η 2 (cid:99) )
− x T (cid:107) 2 ≤ (cid:15)/ 2 with probability at least 1 − e − Ω( n ) for all sufﬁciently small η > 0 , and thus x η ( (cid:98) T/η 2 (cid:99) ) is an (cid:15) -optimal solution of the OLM.
Finally, the validity of applying Theorem 4.6 is guaranteed by Lemma 6.2 and 6.3.
This completes the proof.
In the following subsections, we provide the proofs of all the components used in the above proof.
D. P ROOF OF L EMMA 6.
Recall that for each i ∈ [ n ] f i ( x ) =
f ( u, v ) =
z (cid:62) i ( u (cid:12) 2 − v (cid:12) 2 ) , ∇ f i ( x ) = 2 (cid:0) z i (cid:12)
u z i (cid:12) v (cid:1) , and K ( x ) =
( K ij ( x ))
i,j ∈ [ n ] where each K ij ( x ) =
(cid:104)∇ f i ( x ) , ∇ f j ( x ) (cid:105) .
Then ∇
[Equation]
So for any x ∈ Γ , it holds that
[Equation]
Lemma D.1.
For any ﬁxed x ∈ R D , suppose {∇ f i ( x ) } i ∈ [ n ] is linearly independent, then K ( x ) is full-rank.
Proof of Lemma D.1.
Suppose otherwise, then there exists some λ ∈ R n such that λ (cid:54) =
and λ (cid:62) K ( x ) λ
= 0 .
However, note that
[Equation]
[Equation]
which implies that (cid:80) ni
=
λ i ∇ f i ( x ) = 0 .
This is a contradiction since by assumption {∇ f i ( x ) } i ∈ [ n ] is linearly independent.
Lemma 6.2.
Consider the loss function L deﬁned in (14) and manifold Γ deﬁned in (15) .
If data is full rank, i.e., rank( { z i } i ∈ [ n ] ) = n , then it holds that (a).
Γ is a smooth manifold of dimension
D − n ; (b).
rank( ∇ 2 L ( x ))
= n for all x ∈ Γ .
In particular, it holds that rank( { z i } i ∈ [ n ] ) = n holds with probability 1 for Gaussian distribution and with probability 1 − c d for Boolean distribution for some constant c ∈ (0 , 1) .
Proof of Lemma 6.2.
(1) By preimage theorem (Banyaga & Hurtubise, 2013), it sufﬁces to check the jacobian [ ∇ f 1 ( x ) , . .
. , ∇ f n ( x )]
= 2[ (cid:0)
z 1 (cid:12)
u − z 1 (cid:12) v (cid:1) , . . .
, (cid:0) z n (cid:12) u − z n (cid:12) v (cid:1) ] is full rank.
Similarly, for the second claim, due to (35).
it is also equivalent to show that { (cid:0) z i (cid:12) u
− z i (cid:12) v (cid:1) } i ∈ [ n ] is of rank n .
Since (cid:0) uv (cid:1) ∈
Γ ⊂ U , each coordinate is non-zero, thus we only need to show that { z i } i ∈ [ n ] is of rank n .
This happens with probability 1 in the Gaussian case, and probability at least 1 − c d for some constant c ∈ (0 , 1) by Kahn et al.
This completes the proof.
D. P ROOF OF L EMMA 6.
We ﬁrst establish some auxiliary results.
The following lemma shows the PL condition along the trajectory of gradient ﬂow.
[Equation]
To prove Lemma D.2, we need the following invariance along the gradient ﬂow.
Lemma D.3.
Along the gradient ﬂow generated by ∇ L ( x t ) , u jt v jt stays constant for all j ∈ [ d ] .
Thus, sign( u jt ) =
sign( u j 0 ) and sign( v jt )
=
sign( v j 0 ) for any j ∈ [ d ] .
Proof of Lemma D.3.
[Equation]
Therefore, any sign change of u j , v j would enforce u jt = 0 or v jt = 0 for some t > 0 since u jt , v jt are continuous in time t .
This immediately leads to a contradiction to the invariance of u jt v jt .
We then can prove Lemma D.2.
Proof of Lemma D.2.
Note that
[Equation]
K
[Equation]
x [Equation] where K ( x ) is a n × n p.s.d.
matrix with K ij ( x ) =
(cid:104)∇ f i ( x ) , ∇ f j ( x ) (cid:105) .
Below we lower bound λ min ( K ( x )) , the smallest eigenvalue of K ( x ) .
Note that K ij ( x ) =
(cid:80) nh
[-=
-]=1 z hi z {+hj (( u ht ) 
+}
+ ( v ht ) 2 ) , and we have
[Equation]
where ( ∗ ) is by Lemma D.3.
Thus λ min ( K ) ≥ min i ∈ [ d ] | u i 0 v
i 0 | λ min ( ZZ T ) , which completes the proof.
We also need the following characterization of the manifold Γ .
Lemma D.4.
All the stationary points in U are global minimizers, i.e., Γ = { x ∈ U | ∇ L ( x ) = 0 } .
Proof of Lemma D.4.
Since Γ is the set of local minimizers, each x in Γ must satisfy ∇ L ( x ) = 0 .
The other direction is proved by noting that rank( { z i } i ∈ [ n ] ) = n , which implies rank( {∇ f i ( x ) } i ∈ [ n ] ) = n .
Now, we are ready to prove Lemma 6.3 which is restated below.
Lemma 6.3.
Consider the loss function L deﬁned in (14) , manifold Γ and its open neighborhood deﬁned in (15) .
For gradient ﬂow d x t d t = −∇ L ( x t ) starting at any x 0 ∈ U , it holds that Φ( x 0 ) ∈ Γ .
Proof of Lemma 6.3.
It sufﬁces to prove gradient ﬂow d x t d t
= −∇ L ( x t ) converges when t → ∞ , as long as x 0 ∈ U .
Whenever it converges, it must converge to a stationary point in U .
The proof will be completed by noting that all stationary point of L in U belongs to Γ (Lemma D.4).
Below we prove lim t →∞ x t exists.
By Lemma D.16, denote C = min i ∈ [ d ] | u i 0 v
i 0 | λ min ( ZZ (cid:62) ) , then λ min ( K ( x t )) ≥ C for all t ≥ 0 .
Thus,
[Equation]
Thus the total GF trajectory length is bounded by (cid:82) ∞ [Equation] =
[Equation]
where the last inequality uses that L is non-negative over R D .
Therefore, the GF must converge.
D
P ROOF OF RESULTS IN S ECTION 6.
To study the optimal solution to (18), we consider the corresponding d -dimensional convex program in terms of w ∈ R d , which has been studied in Tropp (2015):
[Equation]
i = ( z ji )(cid:33) | w j | , subject to Zw =
Zw ∗ . Here we slightly abuse the notation of R and the parameter dimension will be clear from the context.
We can relate the optimal solution to (18) to that of (36) via a canonical parametrization deﬁned as follows.
Deﬁnition D.5 (Canonical Parametrization) .
For any w ∈ R d , we deﬁne (cid:0) uv (cid:1) =
ψ ( w ) =
([ w (cid:62) ] (cid:12) 1 / 2 + , [ − w (cid:62) ] (cid:12) 1 / 2 + ) (cid:62) as the canonical parametrization of w .
Clearly, it holds that u (cid:12) 2 − v (cid:12) 2 = w .
Indeed, we can show that if (36) has a unique optimal solution, it immediately follows that the optimal solution to (18) is also unique up to sign ﬂips of each coordinate, as summarized in the lemma below.
Lemma D.6.
Suppose the optimal solution to (36) is unique and equal to w ∗ .
Then the optimal solution to (18) is also unique up to sign ﬂips of each coordinate.
In particular, one of them is given by ( (cid:101) u ∗ , (cid:101) v ∗ ) =
ψ ( w ∗ ) , that is, the canonical parametrization of w ∗ .
Proof of Lemma D.6.
Let ( (cid:98) u, (cid:98) v ) be any optimal solution of (18) and we deﬁne (cid:98) w = (cid:98) u (cid:12) 2 − (cid:98) v (cid:12) 2 , which is also feasible to (36).
By the optimality of w ∗ , we have d
[Equation]
j =
[Equation]
i = (
[Equation]
)(cid:33) | w
[Equation]
≤ d (cid:88)
j = (cid:32) n (cid:88) i =
( z j i )(cid:33) | (cid:98) w j | ≤ d (cid:88)
j = (cid:32) n (cid:88) i =
{+( z j i )(cid:33) [( (cid:98) u j ) 
+}
+ ( (cid:98) v j ) 2 ] .
On the other hand, ( (cid:101) u ∗ , (cid:101) v ∗ ) =
ψ ( w ∗ ) is feasible to (18).
Thus, it follows from the optimality of ( (cid:98) u, (cid:98) v ) that d (cid:88)
j = (cid:32) n (cid:88) i =
( z ji )(cid:33) [( (cid:98) u j ) 2 + ( (cid:98) v j ) 2 ] ≤ d (cid:88)
{+j = (cid:32) n (cid:88) i = ( z ji )(cid:33) [( (cid:101) u j ∗ ) 
 + ( (cid:101) v j ∗ ) 2 ] =+}
d (cid:88)
j = (cid:32) n (cid:88) i =
( z ji )(cid:33) | w j ∗ | .
Combining (37) and (38) yields d (cid:88) j
[Equation]
(cid:32) n (cid:88) i = ( z ji )(cid:33) [( (cid:98) u j ) 2 + ( (cid:98) v j ) 2 ] = d (cid:88) j = (cid:32) n (cid:88) i =
( z ji )(cid:33) | w j ∗ | =
d (cid:88)
j = (cid:32) n (cid:88) i = ( z ji )(cid:33) | ( (cid:98) u j ) 2 − ( (cid:98) v j ) 2 | (39) which implies that (cid:98) u (cid:12) 2 − (cid:98) v (cid:12) 2 is also an optimal solution of (36). Since w ∗ is the unique optimal solution to (36), we have (cid:98) u (cid:12) 2 − (cid:98) v (cid:12) 2 = w ∗ .
Moreover, by (39), we must have (cid:98) u (cid:12) 2 =
[ w ∗ ]
+ and (cid:98) u (cid:12) 2 =
[ w ∗ ]
+ , otherwise the equality would not hold.
This completes the proof.
Therefore, the unique optimality of (18) can be reduced to that of (36).
In the sequel, we show that the latter holds for both Boolean and Gaussian random vectors.
We divide Lemma 6.6 into to Lemma D. and D.7 for clarity.
Lemma D.7 (Boolean Case) .
Let z 1 , . . .
, z n i.i.d.
∼ Unif( {± 1 } d ) .
There exist some constants C, c >such that if the sample size n satisﬁes
[Equation]
then with probability at least 1 − e − cn 2 , the optimal solution of (18) , ( (cid:98) u, (cid:98) v ) , is unique up to sign ﬂips of each coordinate and recovers the groundtruth,
i.e., (cid:98) u (cid:12) 2 − (cid:98) v (cid:12) 2 = w ∗ .
Proof of Lemma D.7.
By the assumption that z 1 , . . .
, z n i.i.d.
∼ Unif( {± 1 } d ) , we have (cid:80) ni
=
( z ji ) 2 = n for all j ∈ [ d ] .
Then (36) is equivalent to the following optimization problem: minimize g ( w ) =
(cid:107) w (cid:107) 1 , subject to Zw =
Z ( u (cid:12) 2 ∗ − v (cid:12) 2 ∗ ) .
This model exactly ﬁts the Example 6.2 in Tropp (2015) with σ = 1 and α = 1 / √ 2 .
Then applying
Equation (4.2) and Theorem 6.3 in Tropp (2015), (40) has a unique optimal solution equal to u (cid:12) 2 ∗ − v (cid:12) 2 ∗ with probability at least 1 − e − ch 2 for some constant c > 0 , given that the sample size satisﬁes
[Equation]
for some absolute constant C > 0 .
Choosing h = n 2 C and then adjusting the choices of C, c appropriately yield the desired result.
Finally, applying Lemma D.6 ﬁnishes the proof.
The Gaussian case requires more careful treatment.
Lemma D.8 (Gaussian Case) .
Let z 1 , . . .
, z n i.i.d.
∼ N (0 , I d ) .
There exist some constants C, c >such that if the sample size satisﬁes
[Equation]
then with probability at least 1 − (2 d + 1) e − cn , the optimal solution of (18) , ( (cid:98) u, (cid:98) v ) , is unique up to sign ﬂips of each coordinate of (cid:98) u and (cid:98) v and recovers the groundtruth, i.e., (cid:98) u (cid:12) 2 − (cid:98) v (cid:12) 2 = w ∗ .
Proof of Lemma D.8.
Since z 1 , . . .
, z n i.i.d.
∼ N (0 , I d ) , we have
[Equation]
for some constant c > 0 , and we denote this event by E n .
Therefore, on E n , we have n
[Equation]
or equivalently,
[Equation]
Deﬁne w
[Equation]
v (cid:12) 2 ∗ , and (36) is equivalent to the following convex optimization problem
[Equation]
subject to Zw =
[Equation]
The point w = 0 is feasible for (41), and we claim that this is the unique optimal solution when n is large enough.
In detail, assume that there exists a non-zero feasible point w for (41) in the descent cone
(Tropp, 2015) D ( g, w ∗ ) of g , then
[Equation]
where the equality follows from that w is feasible.
Therefore, we only need to show that λ min ( Z ; D ( g, x ∗ )) is bounded from below for sufﬁciently large n .
On E n , it holds that g belongs to the following function class
[Equation]
  with Υ = { υ ∈ R d : υ j ∈ [0 . 5 , 1 . 5] , ∀ j ∈ [ d ] } .
We identify g υ ∈ G with υ ∈ Υ , then D ( g, w ∗ ) ⊆ ∪ υ ∈ Υ D ( g υ , w ∗ )) :
= D Υ , which further implies that
[Equation]
Recall the deﬁnition of minimum conic singular value (Tropp, 2015):
[Equation]
where S n − 1 denotes the unit sphere in R n .
Applying the same argument as in (Tropp, 2015) yields
[Equation]
Take the intersection of this event with E n , and we obtain from a union bound that
[Equation]
with probability at least 1 − e − h 2 / 2 − 2 de − cn .
It remains to determine w ( D Υ ) , which is deﬁned as
[Equation]
Without loss of generality, we assume that w ∗ =
( w 1 ∗ , . . .
, w κ ∗ , 0 , . .
. , 0) (cid:62) with w 1 ∗ , . . .
, w κ ∗ > 0 , otherwise one only needs to specify the signs and the nonzero set of w ∗ in the sequel.
For any υ ∈ Υ and any p ∈ D ( g υ , w ∗ ) ∩ S d − 1 , there exists some τ > 0 such that g υ ( w ∗ + τ · p ) ≤ g υ ( w ∗ ) , i.e.,
[Equation]
which further implies that
[Equation]
where the second inequality follows from the triangle inequality.
Then since each υ j ∈ [1 / 2 , 3 / 2] , it follows that
[Equation]
Note that this holds for all υ ∈ Υ simultaneously.
Now let us denote p 1: κ = ( p 1 , . .
p κ ) ∈ R κ and p ( κ +1): d = ( p κ +1 , . . .
, p d ) ∈ R
d − κ , and similarly for other d -dimensional vectors.
Then for all p ∈ D Υ ∩ S d − 1 , by Cauchy-Schwartz inequality, we have
[Equation]
Thus, for any z ∈ R d and any p ∈ D Υ ∩ S d − 1 , it follows that
[Equation]
where the last inequality follows from the fact that p ∈ S
d − 1 .
Therefore, combine the above inequality with (43), and we obtain that
[Equation]
where the second inequality follows from the fact that E [ (cid:107) z 1: κ (cid:107) 2 ] ≤ (cid:112) E [ (cid:107) z 1: κ (cid:107) 22 ] = √ κ .
To bound the second term in (44), applying Lemma D.9, it follows from (44) that
[Equation]
Therefore, combining (45) and (42), we obtain
[Equation]
Therefore, choosing h = √ n − 1 / 2 , as long as n satisﬁes that n ≥ C ( κ ln d ) for some constant
C > 0 , we have λ min ( Z ; D ( g, w ∗ )) > 0 with probability at least 1 − (2 d + 1) e − cn .
Finally, the uniqueness of the optimal solution to (18) in this case follows from Lemma D.6.
Lemma D.9.
Let z ∼ N (0 , I d ) , then it holds that E (cid:2) max
i ∈ [ d ] | z i | (cid:3) ≤ (cid:112) 2 ln(2 d ) .
Proof of Lemma D.9.
Denote M = max i ∈
[ d ] | z i | .
For any λ > 0 , by Jensen’s inequality, we have
[Equation]
[Equation]
Note that E [ e λ |
z
i | ] ≤ 2 · E [ e λz i ] .
Thus, by the expression of the Gaussian moment generating function, we further have
[Equation]
from which it follows that
[Equation]
Choosing λ
= (cid:112) 2 ln(2 d ) yields the desired result.
D. P ROOF OF L EMMA 6.
Instead of studying the convergence of the Riemannian gradient ﬂow directly, it is more convenient to consider it in the ambient space R D .
To do so, we deﬁne a Lagrange function L ( x ; λ ) =
R ( x ) + (cid:80) ni
=
λ i ( f i ( x ) − y i ) for λ ∈ R n .
Based on this Lagrangian, we can continuously extend ∂ Φ( x ) ∇ R ( x ) to the whole space R D .
In speciﬁc, we can ﬁnd a continuous function F : R D → R D such that F ( · )
| Γ = ∂ Φ( · ) ∇ R ( · ) .
Such an F can be implicitly constructed via the following lemma.
Lemma D.10.
The (cid:96) 2 norm has a unique minimizer among {∇ x L ( x ; λ ) | λ ∈ R n } for any ﬁxed x ∈ R D .
Thus we can deﬁne F : R D
→ R D by F ( x ) = argmin g ∈{∇ x L ( x ; λ ) | λ ∈ R n } (cid:107) g (cid:107) 2 .
Moreover, it holds that (cid:104) F ( x ) , ∇ f i ( x ) (cid:105) = 0 for all i ∈ [ n ] .
Proof of Lemma D.10.
Fix any x ∈ R D .
Note that {∇ x L ( x ; λ ) | λ ∈ R n } is the subspace spanned by {∇ f i ( x ) } i ∈ [ n ] shifted by ∇ R ( x ) , thus there is unique minimizer of the (cid:96) 2 norm in this set.
This implies that F ( x ) = argmin g ∈{∇ x L ( x ; λ ) | λ ∈ R n } (cid:107) g (cid:107) 2 is well-deﬁned.
To show the second claim, denote h ( λ ) =
(cid:107)∇ x L ( x ; λ ) (cid:107) 22 / 2 , which is a quadratic function of λ ∈ R n .
Then we have
[Equation]
For any λ such that ∇ x L ( x ; λ ) =
F ( x ) , we must have ∇ h ( λ ) = 0 by the deﬁnition of F ( x ) , which by the above implies
[Equation]
for all i ∈ [ n ] .
Therefore, we further have
[Equation]
for all i ∈ [ n ] .
This ﬁnishes the proof.
Hence, with any initialization x 0 ∈ Γ , the limiting ﬂow (17) is equivalent to the following dynamics
[Equation]
Thus Lemma 6.5 can be proved by showing that the above x t converges to x ∗ as t → ∞ .
We ﬁrst present a series of auxiliary results in below.
Lemma D.11 (Implications for F ( x ) = 0 ) .
Let F : R D → R D be as deﬁned in Lemma D.10.
For any x = (cid:0) uv (cid:1) ∈ R D such that F ( x ) =
, it holds that for each j ∈ [ d ] , either u j = 0 or v j = 0 .
Proof.
Since F ( x ) = 0 , it holds for all j ∈ [ d ] that,
[Equation]
If there exists some j ∈ [ d ] such that u j (cid:54) = 0 and v j (cid:54) = 0 , then it follows from the above two identities that
[Equation]
which happens with probability 0 in both the Boolean and Gaussian case.
Therefore, we must have u j = 0 or v j = 0 for all j ∈ [ d ] .
Lemma D.12.
Let F : R D → R D be as deﬁned in Lemma D.10.
Then F is continuous on R D .
Proof.
Case I.
We ﬁrst consider the simpler case of any ﬁxed x 0 ∈ U =
( R \ { 0 } ) D , assuming that
K ( x 0 ) is full-rank.
Lemma D.10 implies that for any λ ∈ R n such that ∇ x L ( x [Equation] λ ) =
F ( x 0 ) , we have
[Equation]
Thus such λ is unique and given by
[Equation]
Since K ( x ) is continuous, there exists a sufﬁciently small δ > 0 such that for any x ∈ B δ ( x 0 ) ,
K ( x ) is full-rank, which further implies that K ( x ) − 1 is also continuous in B δ ( x ) .
Therefore, by the above characterization of λ , we see that λ ( x ) is continuous for x ∈ B
δ ( x 0 ) , and so is
F ( x ) = ∇ R ( x ) + (cid:80) ni
=1 λ ( x ) i ∇ f i ( x ) .
Case II.
Next, we consider all general x ∈ R D .
Here for simplicity, we reorder the coordinates as x = ( u 1 , v 1 , u 2 , v 2 , . . .
, u d , v d ) with a slight abuse of notation.
Without loss of generality, ﬁx any x 0 such that for some q ∈ [ d ] , ( u i 0 ) 2 + ( v i 0 ) 2 > 0 for all i = 1 , . . .
, q and u i
= v i
= 0 for all i
= q + 1 , . . .
, d .
Then ∇ R ( x 0 ) and {∇ f i ( x 0 ) } i ∈ [ n ] only depend on { z ji } i ∈ [ n ] ,j ∈ [ q ] , and for all i ∈ [ n ] , it holds that
[Equation]
Note that if we replace {∇ f i ( x ) } i ∈ [ n ] by any ﬁxed and invertible linear transform of itself, it would not affect the deﬁnition of F ( x ) .
In speciﬁc, we can choose an invertible matrix Q ∈ R n × n such that, for some q (cid:48) ∈
[ q ] , ( (cid:101) z 1 , . . .
, (cid:101) z n ) =
( z 1 , . . .
, z n ) Q satisﬁes that { (cid:101)
z 1: q i } i ∈ [ q (cid:48) ] is linearly independent and (cid:101) z
1: q i = 0 for all i = q (cid:48)
+ 1 , . .
. , n .
We then consider ( ∇ (cid:101)
f 1 ( x ) , . . . , ∇ (cid:101) f n ( x ))
= ( ∇ f 1 ( x ) , . .
. , ∇ f n ( x )) Q and the corresponding F ( x ) .
For notational simplicity, we assume that Q can be chosen as the identity matrix, so that ( z 1 , . . .
, z n ) itself satisﬁes the above property, and we repeat it here for clarity { z 1: q i } i ∈
[ q (cid:48) ] is linearly independent and (cid:101)
z
[Equation]
q i
[Equation]
This further implies that
[Equation]
In the sequel, we use λ for n -dimensional vectors and ¯ λ for q (cid:48) -dimensional vectors.
Denote
[Equation]
Then due to (47) and (48), we have (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
On the other hand, for any x ∈ R D , by (48), we have (cid:13) (cid:13) (cid:13)(cid:13)(cid:13)(cid:13) ∇ R ( x ) 1:(2 q ) + [Equation] (cid:13)(cid:13)(cid:13)
[Equation]
q [Equation] (cid:13) [Equation]
(cid:13)(cid:13)(cid:13)
[Equation] ≤ (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ∇ [Equation] ( x )
[Equation]
(2 q ) + n (cid:88) i
[Equation]
50) where the ﬁrst and third inequalities follow from the deﬁnition of F ( x ) . Let x → x 0 , by the continuity of ∇ R ( x ) and {∇ f i ( x ) } i ∈ [ n ] , we have
[Equation]
Denote (cid:101) K ( x ) =
( (cid:101) K ij ( x )) ( i,j ) ∈ [ q (cid:48) ] 2 = ( (cid:104)∇ f i ( x ) 1:(2 q ) , ∇ f i ( x ) 1:(2 q ) (cid:105) )
( i,j ) ∈ [ q (cid:48) ] 2 . By applying the same argument as in Case I , since (cid:101) K ( x 0 ) is full-rank, it also holds that lim x → x 0 ¯ λ ( x ) =
¯ λ ( x 0 ) , and thus
[Equation]
[Equation]
Combing (49), (50), (51) and (52) yields
[Equation]
Moreover, since (cid:107) F ( x ) (2 q +1) [Equation] D (cid:107)
[Equation]
= (cid:112) (cid:107) F ( x ) (cid:107) 22 − (cid:107) F ( x ) 1:(2 q ) (cid:107) 22 , we also have lim x
[Equation]
0 (cid:107) F
[Equation]
It then remains to show that lim x → x 0 F ( x ) 1:(2 q )
=
F ( x 0 ) 1:(2 q ) , which directly follows from lim x → x 0 λ ( x ) 1: q (cid:48) =
λ ( x 0 ) 1: q (cid:48)
=
¯ λ ( x 0 ) .
Now, for any (cid:15) > 0 , due to the convergence of ¯ λ ( x ) and that (cid:101) K ( x 0 ) (cid:31) 0 , we can pick a sufﬁciently small δ 1 such that for some constant α > 0 and all x ∈ B δ 1 ( x 0 ) , it holds that (cid:107) ¯ λ ( x ) − ¯ λ ( x 0 ) (cid:107) 2 ≤ (cid:15)/and
[Equation]
[Equation]
for all ¯ λ ∈ R p , where the inequality follows from the strong convexity.
Meanwhile [Equation] due to (48)
[Equation]
we have [Equation] (cid:13). where the second equality follows from (53) and the second equality is due to (52).
Therefore, we can pick a sufﬁciently small δ 2 such that
[Equation]
[Equation]
56) for all x ∈ B δ 2 ( x 0 ) . Setting δ = min( δ 1 , δ 2 ) , it follows from (55) and (56) that
[Equation]
λ ( x ) 1:
[Equation]
(cid:48)
[Equation]
¯ [Equation] ( x ) (cid:107) 2 ≤ (cid:15) 2 , for all x ∈ B δ ( x 0 ) .
Recall that we already have (cid:107) ¯ λ ( x ) − ¯ λ
( x 0 ) (cid:107) ≤ (cid:15)/ 2 , and thus {+(cid:107) λ ( x ) 1: q (cid:48) − λ ( x 0 ) 1: q (cid:48) (cid:107) 
 = (cid:107)+}
λ
( x ) 1: q (cid:48) −
[Table]
≤
[Table]
x ) [Table] [Table] +
[Table] x [Table] x 0 ) (cid:107) 2 ≤ (cid:15) for all x ∈ B δ ( x 0 ) .
Therefore, we see that lim x → x 0 λ ( x ) 1: q (cid:48) =
λ ( x 0 ) 1: q (cid:48) .
Finally, it follows from the triangle inequality that
[Equation]
where, as x → x 0 , the ﬁrst term vanishes by the convergence of λ ( x ) 1: q (cid:48) and the continuity of each ∇ f i ( x ) , the second term converges to 0 by the continuity of ∇ R ( x ) and the third term vanishes by (54).
Therefore, we conclude that
[Equation]
that is, F is continuous.
Lemma D.13.
For any initialization x 0 ∈ Γ , the Riemmanian Gradient Flow (17) (or equivalently, (46) ) is deﬁned on [0 , ∞ ) .
Proof of Lemma D.13.
Let [0 , T ) be the right maximal interval of existence of the solution of Riemannian gradient glow and suppose T (cid:54) = ∞ . Since R ( x ( t )) is monotone decreasing, thus R ( x ( t )) is upper bounded by R ( x (0)) and therefore (cid:107)∇ R ( x ( t ))
(cid:107) is also upper bounded. Since (cid:13)(cid:13)(cid:13) dx ( t ) dt (cid:13)(cid:13)(cid:13) 2 ≤ (cid:107)∇ R ( x ( t )) (cid:107) 2 for any t < T , the left limit x ( T − ) : = lim τ → T − x ( τ ) must exist.
By
Corollary 1, Perko (2001), x ( T − ) belongs to boundary of U , i.e., u j ( T − )
= 0 or v j {+( T − ) = 
+} for some j ∈ [ d ] by Lemma D.11.
By the deﬁnition of the Riemannian gradient ﬂow in (17), we have [Equation] F
[Equation]
By the expression of F ( x ( t ))
= ∇ R ( x ( t )) + (cid:80) ni {+=
+} λ i ( x ( t )) ∇ f i ( x ( t )) , we then have d d
[Equation]
(cid:88) i =
[Equation]
i
[Equation] j
[Equation]
n (cid:88) i = λ i ( x ( t )) z i
[Equation]
(cid:35) u [Equation] ( t ) [Equation]
( t ) − (cid:34)
[Equation]
(cid:88) i = [Equation] z 2 i
[Equation]
[-=
-]j − n (cid:88) i = λ i ( x ( t )) z i [Equation]
j (cid:35) u j ( t ) v j ( t ) = − (cid:32) n (cid:88) i = z 2 i,j (cid:33) u j ( t ) v j ( t ) .
Denote s j = (cid:80) ni =1 z 2 i,j .
It follows that | u j ( t ) v j ( t ) | =
| u j (0) v j (0) | e − s j t for all t ∈ [0 , T ) .
Taking the limit we have | u j ( T − ) v j ( T − ) | ≥ | u j (0) v j (0) | e − s j T > 0 . Contradiction with T (cid:54) = ∞ !
Before showing that F satisﬁes the PL condition, we need the following two intermediate results.
Given two points u and v in R d , we say u weakly dominate v (written as u ≤ v ) if and only if u i ≤ v i , for all i ∈ [ d ] .
Given two subsets A and B of R D , we say A weakly dominates B if and only if for any point v in B , there exists a point u ∈
A such that u ≤ v .
Lemma D.14.
For some q ∈ [ D ] , let S be any q -dimensional subspace of R D and P = { u ∈ R D
|
u
i ≥ 0 , ∀
i ∈
[ D ] } .
Let u (cid:63) be an arbitrary point in P and Q = P ∩ ( u (cid:63) + S ) .
Then there exists a radius r > 0 , such that B 1 r (0) ∩ Q weakly dominates Q , where B 1 r (0) is the (cid:96) 1 -norm ball of radius r centered at 0 .
As a direct implication, for any continuous function f : P → R , which is coordinate-wise nondecreasing, min x ∈ U f
( x ) can always be achieved.
Proof of Lemma D.14.
We will prove by induction on the environment dimension D .
For the base case of D = 1 , either S = { 0 } or S = R , and it is straight-forward to verify the desired for both scenarios.
Suppose the proposition holds for D − 1 , below we show it holds for D .
For each i ∈ [ D ] , we apply the proposition with D − 1 to Q ∩ { u ∈ P | u i = 0 } (which can be seen as a subset of R D − 1 ), and let r i be the corresponding (cid:96) 1 radius.
Set r = max i ∈
[ D ] r i , and we show that choosing the radius to be r sufﬁces.
For any v ∈ Q , we take a random direction in S , denoted by ω .
If ω ≥ 0 or ω ≤ 0 , we denote by y the ﬁrst intersection (i.e., choosing the smallest λ ) between the line { v − λ | ω |} λ ≥ 0 and the boundary of U , i.e., ∪ Di =1 { z ∈ R D
| z i = 0 } .
Clearly y ≤ v .
By the induction hypothesis, there exists a u ∈ B 1 r (0) ∩ Q such that u ≤ y .
Thus u ≤ v and meets our requirement.
If ω has different signs across its coordinates, we take y 1 , y 2 to be the ﬁrst intersections of the line { v − λ | ω |} λ ∈ R and the boundary of U in directions of λ > 0 and λ < 0 , respectively.
Again by the induction hypothesis, there exist u 1 , u 2 ∈ B 1 r (0) ∩ Q such that u 1 ≤ y 1 and u 2 ≤ y 2 .
Since v lies in the line connecting u 1 and u 2 , there exists some h ∈ [0 , 1] such that v =
(1 − h )
u 1 + hu 2 .
It then follows that (1 − h )
u 1 + hu 2 ≤ (1 − h )
y 1 + hy 2 = v .
Now since Q is convex, we have (1 − h )
u 1 + hu 2 ∈ Q , and by the triangle inequality it also holds that (cid:107) (1 − h )
u 1 + hu 2 (cid:107) 1 ≤ r , so (1 − h )
u 1 + hu 2 ∈ B 1 r (0) ∩ Q .
Therefore, we conclude that B 1 r (0) ∩ Q weakly dominates Q , and thus the proposition holds for D .
This completes the proof by induction.
Lemma D.15.
For some q ∈ [ D ] , let S be any q -dimensional subspace of R D and P = { u ∈ R D
|
u
i ≥ 0 , ∀
i ∈
[ D ] } .
Let u (cid:63) be an arbitrary point in P and Q = P ∩ ( u (cid:63) + S ) .
Then there exists a constant c ∈ (0 , 1] such that for any sufﬁciently small radius r > 0 , c · Q weakly dominates
P ∩ ( u (cid:63) + S + B 2 r (0)) , where B 2 r (0) is the (cid:96) 2 -norm ball of radius r centered at 0 .
Proof of Lemma D.15.
We will prove by induction on the environment dimension D .
For the base case of D = 1 , either S = { 0 } or S = R .
S
=
R is straight-forward; for the case S = { 0 } , we just need to ensure c | u (cid:63) | ≤ | u (cid:63) | − r , and it sufﬁces to pick r = | u (cid:63) | and c = 0 .
Suppose the proposition holds for D − 1 , below we show it holds for D .
For each i ∈ [ D ] , we ﬁrst consider the intersection between P ∩ ( u (cid:63) + S + B 2 r (0)) and H i := { u ∈ R D | u i = 0 } .
Let u i be an arbitrary point in P ∩ ( u (cid:63) + S )
∩ H i , then P ∩ ( u (cid:63) + S )
∩ H i
= P ∩ ( u i + S ) ∩ H i =
P ∩ ( u i + S ∩ H i ) .
Furthermore, there exists { α i } i ∈
[ D ] which only depends on S and satisﬁes
P ∩ ( u ∗ + S + B 2 r (0))
∩ H i ⊂ P ∩ ( u i + S ∩ H i + B 2 α i r (0)
∩ H i ) .
Applying the induction hypothesis to P ∩ ( u i + S ∩ H
i + B 2 α
i r (0) ∩ H i ) , we know there exists a c > 0 such that for sufﬁciently small r ,
[Equation]
H i + B 2 α
i r (0) ∩ H i ) .
For any point v in Q and any z ∈ B 2 r (0) , we take a random direction in S , denoted by ω .
If ω ≥or ω ≤ 0 , we denote by y the ﬁrst intersection between { v + z − λ | ω |} λ ≥ 0 and the boundary of U .
Clearly y ≤ v .
Since y ∈ P ∩ ( u (cid:63) + S + B 2 r (0)) ∩ H i ⊂ P ∩ ( u i + S ∩ H i + B 2 α i r (0)
∩ H i ) , by the induction hypothesis, there exists a u ∈ c ( P ∩ ( u (cid:63) + S )
∩ H i ) such that u ≤ y .
Thus z ≤ v + z and z ∈ c ( P ∩ ( u (cid:63) + S )) =
cQ .
If ω has different signs across its coordinates, we take y 1 , y 2 to be the ﬁrst intersections of the line { v + z − λ | ω |} λ ∈ R and the boundary of U in directions of λ > 0 and λ < 0 , respectively.
By the induction hypothesis, there exist u 1 , u 2 ∈ c · Q such that u 1 ≤ y 1 and u 2 ≤ y 2 .
Since v + z lies in the line connecting u 1 and u 2 , there exists some h ∈ [0 , 1] such that v + z = (1 − h )
y 1 + hy 2 .
It then follows that (1 − h )
u 1 + hu 2 ≤ (1 − h )
y 1 + hy 2 = v + z .
Since Q is convex, we have (1 − h )
u 1 + hu 2 ∈ cQ .
Therefore, we conclude that cQ ∩ Q weakly dominates P ∩ ( u (cid:63) + S + B 2 r (0)) for all sufﬁciently small r , and thus the proposition holds for D .
This completes the proof by induction.
Lemma D.16.
(Polyak-Ł condition for F .)
For any x ∗ such that L ( x ∗ ) = 0 , i.e., x ∗ ∈ Γ , there exist a neighbourhood U (cid:48) of x ∗ and a constant c > 0 , such that (cid:107) F ( x ) (cid:107) 2 2 ≥ c · max( R ( x ) − R ( x ∗ ) , 0) for all x ∈ U (cid:48)
∩ Γ .
Note this requirement is only non-trivial when (cid:107) F ( x ∗ ) (cid:107) 2 = 0 since F is continuous.
Proof of Lemma D.16.
It sufﬁces to show the PL condition for { x | F ( x ) = 0 } .
We need to show for any x ∗ satisfying F ( x ∗ )
= 0 , there exist some (cid:15) > 0 and C > 0 , such that for all x ∈ Γ ∩ B 2 (cid:15) ( x ∗ ) with R ( x ) >
R ( x ∗ ) , it holds that (cid:107) F ( x ) (cid:107) 22 ≥ C ( R ( x ) − R ( x ∗ )) .
Case I.
We ﬁrst prove the case where x = (cid:0) uv (cid:1) itself is a canonical parametrization of w = u (cid:12) 2 − v (cid:12) 2 , i.e., u j v j = 0 for all j ∈ [ d ] .
Since x ∗ satisﬁes ∇ F ( x ∗ ) = 0 , by Lemma D.11, we have x ∗ = ψ ( w ∗ ) where w ∗ = u (cid:12) 2 ∗ − v (cid:12) 2 ∗ .
In this case, we can rewrite both R and F as functions of w ∈ R d .
In detail, we deﬁne R (cid:48) ( w ) =
R ( ψ ( w )) and F (cid:48) ( w ) =
F ( ψ ( w )) for all w ∈ R d .
For any w in a sufﬁciently small neighbourhood of w ∗ , it holds that sign( w j ) =
sign( w j ∗ ) for all j ∈ [ q ] .
Below we show that for each possible sign pattern of w ( q +1): d , there exists some constant C which admits the PL condition in the corresponding orthant.
Then we take the minimum of all C from different orthant and the proof is completed.
W.L.O.G., we assume that w j ≥ 0 , for all j = q + 1 . . .
, d .
We temporarily reorder the coordinates as x = ( u 1 , v 1 , u 2 , v 2 , . .
. , u d , v d ) (cid:62) .
Recall that Z =
[ z 1 , . . .
, z n ] (cid:62) is a n -by- d matrix, and we have
[Equation]
, where a = 1 n (cid:80) ni =1 z (cid:12) 2 i ∈ R d .
Since F ( x ∗ ) = 0 , there must exist λ ∗ ∈ R n , such that the ﬁrst 2 q coordinates of ∇ R ( x ∗ ) + (cid:80) ni
=
λ i ∗ ∇ f
i
( x ∗ ) are equal to 0 .
As argued in the proof of Lemma D.12, we can assume the ﬁrst q (cid:48) rows of Z are linear independent on the ﬁrst q coordinates for some q (cid:48) ∈
[ q ] .
In other words, Z can be written as
[Equation]
where Z A ∈ R q (cid:48) × q .
We further denote λ 1 :=
λ 1: q (cid:48) , λ 2 :=
λ ( q (cid:48) +1): n , w 1 :=
w 1: q and w 2 := w ( q +1): d for convenience, then we have [Equation]
[Equation]
Since every w in Γ is a global minimizer, R (cid:48) ( w ) = R
[Equation] ( w ) + (cid:80) [-n i =
-][Equation]
=1 λ i ∗ ( z (cid:62) i
w − y i ) :
= g (cid:62) w , where g = sign( w ) (cid:12) a + Z (cid:62) λ ∗ .
Similarly we deﬁne g 1 := g 1: q and g 2 :
= g ( q +1): d .
It holds that g 1 =
and we assume Z D g 2 = 0 without loss of generality, because this can always be done by picking suitable λ i ∗ for i
= q (cid:48)
+ 1 , . . . , n .
We denote λ 1 − λ ∗ , 1 by ∆ λ 1 , then since 0 = g 1 = sign( w 1 ) (cid:12) a 1 + Z (cid:62) A λ ∗ , 1 , we further have {+(cid:10) ( a 1 + sign( w 1 ) (cid:
 ) Z [Equation] A+}
[Equation]
[Equation]
1 | (cid:11)
[Equation]
a
[Equation]
sign( w [Equation] [Equation] On the other hand, we have g 2 = sign( w 2 ) (cid:12) a 2 + Z (cid:62) B λ ∗ , 1 + Z (cid:62) D λ ∗ , 2 = a 2 + Z (cid:62) B λ ∗ , 1 + Z (cid:62) D λ ∗ , 2 by the assumption that each coordinate of w 2 is non-negative.
Combining this with the above identity, we can rewrite Equation (57) as:
[Equation]
Now suppose R (cid:48) ( w ) − R (cid:48) ( w ∗ ) =
g (cid:62) 2 w 2 = δ for some sufﬁciently small δ (which can be controlled by (cid:15) ).
We will proceed in the following two cases separately.
• Case I.1 : (cid:107) ∆ λ 1 (cid:107) 2 = Ω( √ δ ) .
Since Z A has full row rank and every coordinate of w 1 is non-zero, the ﬁrst term of Equation (58) is Ω( δ )
=
Ω( R (cid:48) ( w ) − R (cid:48) ( w ∗ )) .
• Case I.2 : (cid:107) ∆ λ 1 (cid:107)
=
O ( √ δ ) .
Let u
= g 2 + Z (cid:62) B ∆ λ
1 + Z (cid:62) D λ 2 , then we have u ∈ S + B 2 c √ δ (0) for some constant c > 0 , where S = { g 2 + Z (cid:62) D λ 2 | λ 2 ∈ R n − q (cid:48) } .
By Lemma D.14, there exists some constant c 0 ≥ 1 , such that 1 c 0 · S weakly dominates S + B 2 c √ δ (0) .
Thus we have (cid:107) F (cid:48) ( w ) (cid:107) 22 ≥ inf u ∈ S + B 2 c √ δ
0) (cid:10) u
(cid:12) 2 , w 2 (cid:11) ≥ inf u ∈ 1 c 0 · S (cid:10) s (cid:12) 2 , w 2 (cid:11) , where the last step is because each coordinate of w 2 is non-negative.
Let A be the orthogonal complement of span( Z D , g 2 ) , we know w 2 ∈ δ (cid:107) g 2 (cid:107) 22 g 2 + A , since Z D w 2 =
Z D w ∗ , 2 = {+
+} and g (cid:62) 2 w 2 = δ .
Therefore, inf
[Equation]
Note (cid:10) u (cid:12) 2 , w 2 (cid:11) is a monotone non-decreasing function in the ﬁrst joint orthant, i.e., { ( u, w 2 )
∈
R d × R d − q (cid:48) | u ≥ 0 , w 2 ≥ 0 } , thus by Lemma D.15 the inﬁnimum can be achieved by some ﬁnite ( u, w 2 ) in the joint ﬁrst orthant.
Applying the same argument to each other orthant of u ∈ R d , we conclude that the right-hand-side of (59) can be achieved.
On the other hand, we have u (cid:62) w 2 = δ > 0 for all w 2 ∈
[Equation]
22 g 2 + A and u ∈ S , by
Z D g 2 = 0 and the deﬁnition of A .
This implies there exists at least one i ∈
[ d − q (cid:48) ] such that w i 2 u i > 0 , which further implies (cid:10) u (cid:12) 2 , w 2 (cid:11) > 0 .
Therefore, we conclude that [Equation] F [Equation]
R
[Equation]
w
[Equation]
Case II.
Next, for any general x = (cid:0) uv (cid:1) , we deﬁne w = u (cid:12) 2 − v (cid:12) 2 and m = min { u (cid:12) 2 , v (cid:12) 2 } , where min is taken coordinate-wise.
Then we can rewrite (cid:107) F ( x ) (cid:107) 22 as (cid:107) F ( x ) (cid:107) 22 = min λ ∈ R n (cid:13)(cid:13)(cid:13)(cid:13)(cid:18)(cid:20) aa (cid:21) + (cid:20)
Z − Z (cid:21) λ (cid:19)
[Table]
[Equation]
Then applying the result for the previous case yields the following for some constant C ∈ (0 , 1) :
[Equation]
where the ﬁrst equality follows from the fact that x ∗ = ψ ( w ∗ ) and the last inequality is due to the fact that both R ( ψ ( w ) − R ( ψ ( w ∗ )) and R ( x ) − R ( ψ ( w )) are non-negative.
This completes the proof.
Now, based on the PL condition, we can show that (17) indeed converges.
Lemma D.17.
The trajectory of the ﬂow deﬁned in (17) has ﬁnite length, i.e., (cid:82) ∞ t =0 (cid:107) d x d t (cid:107)
d t < ∞ for any x 0 ∈ Γ .
Moreover, x t converges to some x ∞ when t → ∞ with F ( x ∞ ) = 0 .
Proof of Lemma D.17.
Note that along the Riemannian gradient ﬂow, R ( x t ) is non-increasing, thus (cid:107) x t (cid:107) 2 is bounded over time and { x t } t ≥ 0 has at least one limit point, which we will call x ∗ .
Therefore,
R ( x ∗ ) is a limit point of R ( x t ) , and again since R ( x t ) is non-increasing, it follows that R ( x t ) ≥
R ( x ∗ ) and lim t →∞ R ( x t ) =
R ( x ∗ ) .
Below we will show lim t →∞ x t = x ∗ .
Note that d R ( x t ) d t
[Equation]
F ( x t ) (cid:107) 22 where the last inequality applies Lemma D.10.
By Lemma D.16, there exists a neighbourhood of x ∗ , U (cid:48) , in which PL condition holds of F .
Since x ∗ is a limit point, there exists a time T 0 , such that x T 0 ∈ U .
Let T 1 = inf t ≥ T 0 { x ( t ) / ∈ U (cid:48) } (which is equal to ∞ if x ( t ) ∈ U (cid:48) for all t ≥ T 0 ).
Since x t is continuous in t and
U is open, we know T 1 > T 0 and for all t ∈
[ T 0 , T 1 ) , we have (cid:107) F ( x t ) (cid:107) 2 ≥ √ c ( R ( x t ) − R ( x ∗ )) 1 / 2 .
Thus it holds that for t ∈
[ T 0 , T 1 ) ,
[Equation]
that is,
[Equation]
Therefore, we have
[Equation]
Thus if we pick T 0 such that R ( x T 0 ) − R ( x ∗ ) is sufﬁciently small, R ( T 1 ) will remain in U , which implies that T 1 cannot be ﬁnite and has to be ∞ .
Therefore, Equation (60) shows that the trajectory of x t is of ﬁnite length, so x ∞ =
lim t →∞ x t exists and is equal to x ∗ .
As a by-product, F ( x ∗ ) must be 0 .
Finally, collecting all the above lemmas, we are able to prove Lemma 6.5.
In Lemma D.17 we already show the convergence of x ( t ) as t → ∞ , the main part of the proof of Lemma 6.5 is to show the x ( ∞ ) cannot be sub-optimal stationary points of R on Γ , the closure of Γ .
The key idea here is that we can construct a different potential φ for each such sub-optimal stationary point x ∗ , such that (1) φ ( x t ) is locally increasing in a sufﬁciently neighborhood of x ∗ and (2)
lim x → x ∗ φ ( x ) =
−∞ .
Lemma 6.5.
Let { x t } t ≥ 0 ⊆
R D be generated by the ﬂow deﬁned in (17) with any initialization x 0 ∈ Γ .
Then x ∞ =
lim t →∞ x t exists.
Moreover, x ∞ =
x ∗ is the optimal solution of (18) .
Proof of Lemma 6.5.
We will prove by contradiction.
Suppose x ∞ =
(cid:0) u ∞ v ∞ (cid:1) =
lim t →∞ x t is not the optimal solution to (18).
Denote w t = u (cid:12) 2 t − v (cid:12) 2 t , then w ∞ =
lim t →∞ w t is not the optimal solution to (36).
Thus we have R ( w t ) >
R ( w ∗ ) .
Without loss of generality, suppose there is some q ∈
[ d ] such that ( u i ∞ ) 2 + ( v i ∞ ) 2 > 0 for all i = 1 , . . .
, q and u i ∞ =
v i ∞ =
for all i
= q + 1 , . . .
, d .
Again, as argued in the proof of Lemma D.12, we can assume that, for some q (cid:48) ∈
[ q ] , { z 1: q i } i ∈ [ q (cid:48) ] is linearly independent and z 1: q i = 0 for all i = q (cid:48)
+ 1 , . . n.
Since both w ∞ and w ∗ satisfy the constraint that Zw ∞ =
Zw ∗ =
Y , we further have
[Equation]
[Equation]
+
[Equation]
): d i , w ( q +1): d ∗ (cid:105) , for all i = q (cid:48) + 1 , . .
n.
Consider a potential function ϕ : U → R deﬁned as
[Equation]
Clearly lim t →∞ ϕ ( x t ) =
−∞ if lim t →∞ x t = x ∞ .
Below we will show contradiction if x ∞ is suboptimal.
Consider the dynamics of ϕ ( x ) along the Riemannian gradient ﬂow:
[Equation]
(63) where F is deﬁned previously in Lemma D.10.
Recall the deﬁnition of F , and we have
[Equation]
To show (cid:104)∇ ϕ ( x t ) , F ( x t ) (cid:105) < 0 , we analyze I
and I 2 separately.
By the deﬁnition of ϕ ( x ) , we have
[Equation]
where e j is the j -th canonical base of R d .
Recall that ∇ f i ( x ) =
(cid:0)
z i (cid:12) u − z i (cid:12) v (cid:1) , and we further have
I [Equation] =
[Equation]
65) where the last equality follows from (62).
Next, we show that
I 1 < 0 by utilizing the fact that w ∗ − w ∞ is a descent direction of R (cid:48) ( w ) .
For w ∈ R d , deﬁne (cid:101)
f i ( w ) =
z (cid:62) i w and
[Equation]
Clearly, for any w ∈ R D satisfying Zw =
Y , it holds that (cid:101)
f
i ( w ) − y i = 0 for each i ∈ [ n ] , and thus R ( w ) = (cid:101) R ( w ) .
In particular, we have (cid:101) R ( w ∞ ) =
R ( w ∞ ) > R ( w ∗ ) = (cid:101) R ( w ∗ ) .
Since (cid:101) R ( w ) is a convex function, it follows that (cid:101) R ( w ∞ + s ( w ∗ − w ∞ )) <
(cid:101) R ( w ∞ ) for all sufﬁciently small s > 0 , which implies d (cid:101) R d t ( w ∞ + s ( w ∗ − w ∞ ))
| s =0 < − 2 c < 0 for some constant c > 0 .
Note that, for small enough s > 0 , we have
[Equation]
i = ( z ji )(cid:33) s | w j ∗ | . Therefore, we can compute the derivative with respect to s at s = 0 as − 2 c > d (cid:101) R d t ( w [Equation] + s (
[Equation]
) (cid:12)(cid:12)(cid:12)(cid:12) s
=
=
q (cid:88)
j = (cid:32) n (cid:88) i =
( z ji )(cid:33) sign( w j ∞ )( w j ∗ − w j ∞ )
+
[Equation]
(cid:88)
j = q + (cid:32) n (cid:88)
i =
( z ji )(cid:33) | w j ∗ | +
[Equation]
(cid:88) i =
[Equation]
[Equation]
[Equation]
) =
[Equation]
(cid:88)
j =
[Equation]
(cid:32) n (cid:88) i =
( z ji )(cid:33) sign( w j
[Equation]
)( w j
[Equation]
w
[Equation]
+
[Equation]
(cid:88)
j
= q + (cid:32)
[Equation]
(cid:88) i =
( z ji )
[Equation]
(cid:33)
| w j ∗ | + q (cid:88)
j = ( w
[Equation]
w
[Equation]
) q (cid:48) (cid:88)
i = λ ( x ∞ )
[Equation]
z ji + d (cid:88)
j = q + w j ∗ q (cid:48) (cid:88) i
= λ
( x [Equation] [Equation] ji (66) where the second equality follows from the fact that w ( q +1): d ∞ = 0 .
Since x t converges to x ∞ , we must have F ( x ∞ ) = 0 , which implies that for each j ∈ { 1 , . . .
, q } ,
[Equation]
Combining the above two equalities yields
[Equation]
for all j ∈ [ q ] .
Apply the above identity together with (66)
[Equation]
and we obtain −
[Equation]
> [Equation] λ ( x ∞ ) i z ji + d (cid:88)
j = q + w j ∗ q (cid:48) (cid:88) i = λ ( x ∞ ) i z ji = d (cid:88) j = q + (cid:32) n (cid:88)
i =
( z ji )(cid:33) | w j ∗ | + d (cid:88)
j = q + w j ∗ q (cid:48) (cid:88) i = λ ( x ∞ ) i z ji (67)
On the other hand, by directly evaluating ∇ R ( x t ) and each ∇ f i ( x t ) , we can compute I 1 as
I 1 =d (cid:88)
j = q + w j ∗ 1 { w j ∗ > 0 } u jt   n (cid:88)
i =
( z ji ) 2 u jt + q (cid:48) (cid:88) i
= λ ( x t )
i z ji u jt   −d (cid:88)
j = q + w j ∗ 1 { w j ∗ < 0 } v jt   n (cid:88) i
=
( z ji ) 2 v jt − q (cid:48) (cid:88) i
= λ ( x t ) i z ji v jt  
=
d (cid:88)
j = q + (cid:32) n (cid:88)
i =
( z ji )(cid:33) | w j ∗ | +d (cid:88)
j = q + w j ∗ q (cid:48) (cid:88) i
= λ ( x t )
i z ji =d (cid:88)
j = q + (cid:32) n (cid:88) i =
( z j i )(cid:33) | w j ∗ | +d (cid:88)
[Equation]
We already know that λ ( x ) 1: q (cid:48) is continuous at x ∞ by the proof of Lemma D.12, so the third term converges to 0 as x t tends to x ∞ .
Now, applying (67), we immediately see that there exists some δ > 0 such that I 1 < − c for x t ∈ B δ ( x ∞ ) .
As we have shown in the above that I 2 = 0 , it then follows from (63) and (64) that
[Equation]
Since lim t →∞ x t = x ∞ , there exists some T > 0 such that x t ∈ B δ ( x ∞ ) for all t > T .
By the proof ofLemma D.13, we know that ϕ ( x T ) >
−∞ , then it follows from (68) that
[Equation]
which is a contradiction.
This ﬁnishes the proof.
D. P ROOF OF T HEOREM 6.
Here we present the lower bound on the sample complexity of GD in the kernel regime.
Theorem 6.7.
Assume z 1 , . .
, z n i.i.d.
∼ N (0 , I d ) and y i = w (cid:62)∗ z
i , for all i ∈ [ n ] .
Deﬁne the loss with linearized model as L ( x ) = (cid:80) ni
=
( f i ( x 0 )
+ (cid:104)∇ f i ( x 0 ) , x − x 0 (cid:105) − y i ) 2 , where x = (cid:0) uv (cid:1) and x 0 = (cid:0) u 0 v 0 (cid:1) =
α (cid:0) 11 (cid:1) .
Then for any ground truth w ∗ , any learning rate schedule { η t } t ≥ 1 , and any ﬁxed number of steps T , the expected (cid:96) 2 loss of x T is at least (1 − nd ) (cid:107) w ∗ (cid:107) 22 , where x T is the T -th iterate of GD on L , i.e., x t +
= x
t − η t ∇ L ( x t ) , for all t ≥ 0 .
Proof of Theorem 6.7.
We ﬁrst simplify the loss function by substituting x (cid:48)
= x − x 0 , so correspondingly x (cid:48)
=
and we consider L (cid:48)
( x (cid:48) ) :
= L ( x ) =
( (cid:104)∇ f i ( x 0 ) , x (cid:48) (cid:105) − y i ) 2 .
We can think as if GD is performed on L (cid:48)
( x (cid:48) ) .
For simplicity, we still use the x and L ( x ) notation in below.
In order to show test loss lower bound against a single ﬁxed target function, we must take the properties of the algorithm into account.
The proof is based on the observation that GD is rotationally equivariant (Ng, 2004; Li et al., 2020c) as an iterative algorithm, i.e., if one rotates the entire data distribution (including both the training and test data), the expected loss of the learned function remains the same.
Since the data distribution and initialization are invariant under any rotation, it means the expected loss of x T with ground truth being w ∗ is the same as the case where the ground truth is uniformly randomly sampled from all vectors of (cid:96) 2 -norm (cid:107) w ∗ (cid:107) 2 .
Thus the test loss of x T is
[Equation]
Note x T ∈ span {∇ f x ( x 0 ) } , which is at most an n -dimensional space spanned by the gradients of model output at x 0 , so is u T − v T .
We denote the corresponding space for u T − v T by S , so dim( S ) ≤ n
and it holds that (cid:107) w ∗ − ( u T − v T ) (cid:107) 22 ≥ (cid:107) ( I D − P S ) w ∗ (cid:107) 22 , where P S is projection matrix onto space S .
The expected test loss is lower bounded by
[Equation]
[Table]