Year: 2021

Abstract:
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 sufficiently 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[∇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 (Woodworth et al., 2020), while GD initialized in the kernel regime requires Ω(d) samples. This upper bound is minimax optimal and improves the previous Õ(κ) upper bound (HaoChen et al., 2020).

heading: 1 INTRODUCTION
text: 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
xη(k + 1) = xη(k)− η(∇L(xη(k)) + σξk(xη(k))) (1) where η is the learning rate (LR), L : RD → R is the training loss and σ(x) = [σ1(x), σ2(x), . . . , σΞ(x)] ∈ RD×Ξ is a deterministic noise function. Here ξk is sampled uniformly from {1, 2, . . . ,Ξ} and it satisfies Eξk [σξk(x)] = 0, ∀x ∈ Rd and k. It is widely believed that large LR (or equivalently, small batch size) helps SGD find better minima. For instance, some previous works argued that large noise enables SGD to select a flatter attraction basin of the loss landscape which potentiallyx benefits 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 sufficiently 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
with label noise locally decreases the sharpness of loss, tr[∇2L], 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).

heading: 1.1 INTUITIVE EXPLANATION OF REGULARIZATION EFFECT DUE TO SGD
text: We start with an intuitive description of the implicit regularization effect described in Blanc et al. (2020). For simplification, 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 σσ>(x) = ∇2L(x) for every x in the manifold Γ. (See derivation in Section 5)
dX̃η(t) = −η∇L(X̃η(t))dt+ η · σ(X̃η(t))dW (t). (2) Suppose X̃η(0) is already close to some local minimizer point X∗ ∈ Γ. The goal is to show X̃η(t) will move in the tangent space and steadily decrease tr[∇2L]. At first 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 X̃η(t) is completely different in tangent space and normal space — the fast random walk in normal space causes X̃η(t) to move slowly (with velocity Θ(η2)) but deterministically in certain direction. To explain this, letting ∆(t) = X̃η(t) − X∗, Taylor expansion of (2) gives d∆(t) ≈ −η∇2L(X∗)∆dt+ ησ(X∗)dW (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 σσ> = ∇2L. 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 simplifies the expression of the velocity in tangent space to η 2
2 ∇T tr[∇ 2L] (see
Figure 1c), where ∇T means the gradient is only taken in the tangent space. However, the above approach only gives a local analysis forO(η−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 finds a ( , γ)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 difficulty 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.

heading: 1.2 OUR APPROACH: SEPARATING THE SLOW FROM THE FAST
text: 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
dXη(t) = −η−1∇L(Xη(t))dt+ σ(Xη(t))dW (t). (3)
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 Φ : RD → Γ is some function to be specified 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
dΦ(Xη(t)) = −η−1∂Φ(Xη(t))∇L(Xη(t))dt+ ∂Φ(Xη(t))σ(Xη(t))dW (t)
+ 1
2 ∑D i,j=1 ∂ijΦ(Xη(t))(σ(Xη(t))σ(Xη(t)) >)ijdt.
Note the first term −η−1∂Φ(Xη)∇L(Xη) is going to diverge to∞ when η → 0, so a natural choice for Φ is to kill the first term. Further note −∂Φ(X)∇L(X) is indeed the directional derivative of Φ at X towards −∇L, killing the first term becomes equivalent to making Φ invariant under Gradient Flow of −∇L(X)! Thus it suffices to take Φ(X) to be the limit of gradient flow starting at X . (Formally defined in Section 3; see Lemma C.2 for a proof of ∂Φ(X)∇L(X) ≡ 0.) Also intuitively Xη will be infinitely 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
dXη(t) ≈ ∂Φ(Xη(t))σ(Xη(t))dW (t) + 1
2 ∑D i,j=1 ∂ijΦ(Xη(t))(σ(Xη(t))σ(Xη(t)) >)ijdt, (4)
and the solution of (4) shall converge to that of the following (in an intuitive sense):
dX(t) = ∂Φ(X(t))σ(X(t))dW (t) + 1
2 ∑D i,j=1 ∂ijΦ(X(t))(σ(X(t))σ(X(t)) >)ijdt, (5)
The above argument for SDE was first 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 infinitesimal 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 infinitesimal LR in Θ(η−2) steps. Finally, our Corollary 5.2 shows (5) gives exactly the same regularization as tr[∇2L] for label noise SGD. The main contributions of this paper are summarized as follows.
1. In Section 4, we propose a mathematical framework to study the implicit bias of SGD with infinitesimal 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. 2. In Section 5, we give limiting dynamics of SGD with isotropic noise and label noise. 3. 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 `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 Õ(κ2) upper bound by HaoChen et al. (2020). 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 flow 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.

heading: 2 RELATED WORKS
text: 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. (2019). 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 2 − v 2 or w = u v, gradient descent/flow from small initialization implicitly regularizes `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 Itô 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) justifies such SDE approximation. Xie et al. (2020) gave a heuristic derivation explaining why SGD favors flat 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 flat minima. By approximating SGD using a SDE with slightly modified covariance for the overparametrized linear model, Pesme et al. (2021) relates the strength of implicit regularization to training speed.

heading: 3 NOTATION AND PRELIMINARIES
text: For any integer k, we denote Ck 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 define [u v]i = uivi and [v α]i = (v
i)α. For any matrix A, we denote its pseudo inverse by A†. For mapping F : RD → RD, we denote the Jacobian of F at x by ∂F (x) ∈ RD×D where the (i, j)-th entry is ∂iF j(x). We also use ∂F (x)[u] and ∂2F (x)[u, v] to denote the first and second order directional derivative of F at x along the derivation of u ( and v). We abuse the notation of ∂2F by viewing it a linear mapping defined on RD ⊗ RD ∼= RD 2 , in the sense that ∂2F (x)[Σ] = ∑D i,j=1 ∂ 2F (x)[ei, ej ]Σ ij , for any Σ ∈ RD×D. For any submanifold Γ ⊂ RD and x ∈ Γ, we denote by Tx(Γ) 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 φ : RD × [0,∞)→ RD satisfying φ(x, t) = x− ∫ t 0 ∇L(φ(x, s))ds. We further denote the limiting mapping Φ(x) = limt→∞ φ(x, t) whenever the limit exists. We denote 1ξ ∈ RΞ as the one-hot vector where ξ-th coordinate is 1, and 1 the all 1 vector.

heading: 3.1 MANIFOLD OF LOCAL MINIMIZERS
text: Assumption 3.1. Assume that the loss L : RD → R is a C3 function, and that Γ is a (D −M)dimensional C2-submanifold of RD for some integer 0 ≤M ≤ D, where for all x ∈ Γ, x is a local minimizer of L and rank(∇2L(x)) = M . Assumption 3.2. Assume that U is an open neighborhood of Γ satisfying that gradient flow starting in U converges to some point in Γ, i.e., ∀x ∈ U , Φ(x) ∈ Γ. (Then Φ is C2 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 ∈ RD, L(x) = minx′∈RD L(x′) implies fi(x) = yi for all i ∈ [n], then by
preimage theorem (Banyaga & Hurtubise, 2013), the manifold Γ := {x ∈ RD | fi(x) = yi,∀i ∈ [n]} is of dimension D − n if the Jacobian matrix [∇f1(x), . . . ,∇fn(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.

heading: 4 LIMITING DIFFUSION OF SGD
text: In Section 4.1 we first 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.

heading: 4.1 RECAP OF KATZENBERGER’S THEOREM
text: Let {An}n≥1 be a sequence of integrators, where each An : R → R is a non-decreasing function with An(0) = 0. Let {Zn}n≥1 be a sequence of R|Ξ|-valued stochastic processes defined on R. Given loss function L and noise covariance function σ, we consider the following stochastic process:
Xn(t) = X(0) + ∫ t 0 σ(Xn)dZn + ∫ t 0 −∇L(Xn)dAn (6)
In particular, when the integrator sequence {An}n≥1 increases infinitely fast, meaning that ∀ > 0, inft≥0(An(t+ )−An(t))→∞ as n→∞, we call (6) a Katzenberger process. One difficulty for directly studying the limiting dynamics of Xn(t) is that the point-wise limit as n→∞ become discontinuous at t = 0 if X(0) /∈ Γ. The reason is that clearly limn→∞Xn(0) = X(0), but for any t > 0, since {An}n≥1 increases infinitely fast, one can prove limn→∞Xn(t) ∈ Γ! To circumvent this issue, we consider Yn(t) = Xn(t)− φ(X(0), An(t)) + Φ(X(0)). Then for each n ≥ 1, we have Yn(0) = Φ(X(0)) and limn→∞ Yn(t) = limn→∞Xn(t). Thus Yn(t) has the same limit on (0,∞) as Xn(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 satisfies Assumptions 3.1 and 3.2. Let {Xn}n≥1 be a sequence of Katzenberger process with {An}n≥1, {Zn}n≥1. Let Yn(t) = Xn(t) − φ(X(0), An(t)) + Φ(X0). Under technical assumptions, it holds that if (Yn, Zn) converges to some (Y,W ) in distribution, where {W (t)}t≥0 is the standard Brownian motion, then Y stays on Γ and admits
Y (t) = Y (0) + ∫ t 0 ∂Φ(Y )σ(Y )dZ + 1 2 ∑D i,j=1 ∫ t 0 ∂ijΦ(Y )(σ(Y )σ(Y ) >)ijds. (7)
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 limn→∞ ηn = 0, An(t) = ηnbt/η2nc, and Zn(t) = ηn ∑bt/η2nc k=1 (1ξk − 1 Ξ1), where ξ1, ξ2, . . .
i.i.d.∼ Unif([Ξ]). Then with the same initialization Xn(0) = xηn(0) ≡ X(0), Xn(kη2n) defined by (6) is a Katzenberger process and is equal to xηn(k) defined in (1) with LR equal to ηn for all k ≥ 1. Moreover, the counterpart of (7) is
Y (t) = Φ(X(0)) + ∫ t 0 ∂Φ(Y )σ(Y )dW (s) + 1 2 ∫ t 0 ∂2Φ(Y )[Σ(Y )]ds, (8)
where Σ ≡ σσ> 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 difficulty comes from the convergence of (Yn, Zn) 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 Yn with respect to some compact K ⊂ U , i.e., Y µn(K)n (t) = Yn(t ∧ µn(K)) where µn(K) is the stopping time of Yn 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 difficult 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 final result as Theorem 4.6.

heading: 4.2 CLOSED-FORM EXPRESSION OF THE LIMITING DIFFUSION
text: 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) ∈ RD×D is the projection matrix onto tangent space Tx(Γ).
To express the second-order derivatives compactly, we introduce the notion of Lyapunov operator.
Definition 4.4 (Lyapunov Operator). For a symmetric matrix H , we define WH = {Σ ∈ RD×D | Σ = Σ>, HH†Σ = Σ = ΣHH†} and Lyapunov Operator LH : WH → WH as LH(Σ) = H>Σ + ΣH . It’s easy to verify L−1H is well-defined on WH . Lemma 4.5. Let x be any point in Γ and Σ = Σ(x) = σσ>(x) ∈ RD×D be the noise covariance at x1. Then Σ can be decomposed as Σ = Σ‖ + Σ⊥ + Σ‖,⊥ + Σ⊥,‖, where Σ‖ := ∂ΦΣ∂Φ, Σ⊥ := (ID − ∂Φ)Σ(ID − ∂Φ) and Σ‖,⊥ = Σ>⊥,‖ = ∂ΦΣ(ID − ∂Φ) are the noise covariance in tangent space, normal space and across both spaces, respectively. Then it holds that
∂2Φ[Σ] =(∇2L)†∂2(∇L) [ Σ‖ ] −∂Φ∂2(∇L) [ L−1∇2L(Σ⊥) ] +2∂Φ∂2(∇L) [ (∇2L)†Σ⊥,‖ ] . (9)

heading: 4.3 MAIN RESULT
text: 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η(bT/η2c) converges in distribution to Y (T ) as η → 0.
dY (t) = Σ 1 2 ‖ (Y )dW (t)︸ ︷︷ ︸ Tangent Noise + 1 2 ∇2L(Y )†∂2(∇L)(Y )
[ Σ‖(Y ) ] dt︸ ︷︷ ︸
Tangent Noise Compensation
+ 1
2 ∂Φ(Y )
( ∂2(∇L)(Y ) [ ∇2L(Y )†Σ⊥,‖(Y ) ]︸ ︷︷ ︸ Mixed Regularization − ∂2(∇L)(Y ) [ L−1∇2L(Σ⊥(Y )) ]︸ ︷︷ ︸ Normal Regularization ) dt, (10)
where Σ‖,Σ⊥,Σ⊥,‖ are defined 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‖ (Y )dW (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.7 shows that the value of the second term only depends on Γ itself, i.e., it’s same for all loss L which locally defines the same Γ. (b) The noise in the normal space is killed since the limiting dynamics always stay on Γ. However, its second order effect (Itô correction term) takes place as a vector field 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 x̃η(t) = xη(bt/η2c), then 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.

heading: 5 IMPLICATIONS AND EXAMPLES
text: In this section, we derive the limiting dynamics for two notable noise types, where we fix 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) ≡ ID 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∇2L.
1For notational convenience, we drop dependency on x.
Corollary 5.1 (Limiting Diffusion for Isotropic Noise). If Σ ≡ ID on Γ, SDE (10) is then
dY (t) = ∂Φ(Y )dW + 1
2 ∇2L(Y )†∂2(∇L)(Y ) [∂Φ(Y )] dt︸ ︷︷ ︸
Brownian Motion on Manifold
− 1 2 ∂Φ(Y )∇(ln |∇2L(Y )|+)dt︸ ︷︷ ︸
Normal Regularization
(11)
where |∇2L(Y )|+ = limα→0 |∇ 2L(Y )+αID|
αD−rank(∇2L(Y )) is the pseudo-determinant of∇2L(Y ). |∇2L(Y )|+ is
also equal to the sum of log of non-zero eigenvalue values of∇2L(Y ).
Type II: Label Noise. When doing SGD for `2-regression on dataset {(zi, yi)}ni=1, adding label noise (Blanc et al., 2020; Damian et al., 2021) means replacing the true label at iteration k, yik , by a fresh noisy label ỹik := yik + δk, where δk
i.i.d.∼ Unif{−δ, δ} for some constant δ > 0. Then the corresponding loss becomes 12 (fik(x) − ỹik)
2, where fik(x) is the output of the model with parameter x on data zik . So the label noise SGD update is
xk+1 = xk − η/2 · ∇x (fik(xk)− yik + δik) 2 = xk − η(fik(xk)− yik + δk)∇xfik(xk). (12)
Suppose the model can achieve the global minimum of the loss L(x) := 12E[(fi(x) − ỹi) 2] at x∗, then the model must interpolate the whole dataset, i.e., fi(x∗) = yi for all i ∈ [n], and thus here the manifold Γ is a subset of {x ∈ RD | fi(x) = yi, ∀i ∈ [n]}. Here the key property of the label noise used in previous works is Σ(x) = 1n ∑n i=1∇xfi(x)∇xfi(x)> = ∇2L(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∇2L(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 infinitesimal 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∇2L on Γ for some constant c > 0, SDE (10) can be simplified into (13) where the regularization is from the noise in normal space.
dY (t) = −1/4∂Φ(Y (t))∇ tr[c∇2L(Y (t))]dt. (13)

heading: 6 PROVABLE GENERALIZATION BENEFIT WITH LABEL NOISE
text: In this section, we show provable benefit 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 `1 norm while large initialization corresponds to the kernel regime and prefers solutions with small `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 `1 norm. When the groundtruth is κ-sparse, this provides a 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 xinit (except a zero-measure set) and any > 0, there exist η0, T > 0 such that for any η < η0, OLM trained with label noise SGD (12) with LR equal to η for bT/η2c steps returns an -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:
1. Show Assumption 3.1 is satisfied, i.e., the set of local minimizers, Γ, is indeed a manifold and the hessian∇2L(x) is non-degenerate on Γ (by Lemma 6.2);
2. Show Assumption 3.2 is satisfied, i.e., Φ(U) ⊂ Γ (by Lemma 6.3); 3. Show the limiting flow (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 2. Their 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 {(zi, yi)}i∈[n] be the training dataset where z1, . . . , zn i.i.d.∼ Unif({±1}d) or N (0, Id) and each yi = 〈zi, w∗〉 for some unknown w∗ ∈ Rd. We assume that w∗ is κ-sparse for some κ < d. Denote x = ( u v ) ∈ RD = R2d, and we will use x and (u, v) exchangeably as the parameter of functions defined on RD in the sequel. For each i ∈ [n], define fi(x) = fi(u, v) = z>i (u 2 − v 2). Then we fit {(zi, yi)}i∈[n] with an overparametrized model through the following loss function:
L(x) = L(u, v) = 1n ∑n i=1 `i(u, v), where `i(u, v) = 14 (fi(u, v)− yi) 2. (14)
To simplify the notation, we define Z = (z1, . . . , zn)> ∈ Rn×d and Y = (y1, . . . , yn)> ∈ Rn. Consider the following manifold:
Γ = { x = (u>, v>)> ∈ U : Z(u 2 − v 2) = Y } , where U = (R \ {0})D. (15)
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 defined in (14) and manifold Γ defined in (15). If data is full rank, i.e., rank({zi}i∈[n]) = n, then it holds that (a). Γ is a smooth manifold of dimension D− n; (b). rank(∇2L(x)) = n for all x ∈ Γ. In particular, it holds that rank({zi}i∈[n]) = n holds with probability 1 for Gaussian distribution and with probability 1− cd for Boolean distribution for some constant c ∈ (0, 1).
Lemma 6.3. Consider the loss function L defined in (14), manifold Γ and its open neighborhood defined in (15). For gradient flow dxtdt = −∇L(xt) starting at any x0 ∈ U , it holds that Φ(x0) ∈ Γ. Remark 6.4. In previous works (Woodworth et al., 2020; Azulay et al., 2021), the convergence of gradient flow is only assumed. Recently Pesme et al. (2021) proved it for a specific initialization, i.e., uj = vj = α,∀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) = 1 4n ∑n i=1 ‖∇fi(x)‖22. To obtain the close-form expression of R(x), it is straightforward to compute
∇fi(x) = 2 ( zi u −zi v ) and ∇`i(x) = (fi(u, v)− yi) ( zi u −zi v ) for each i ∈ [n]. Hence, it follows that
R(x) = 1
4n ∑D j=1 (∑n i=1 (zji ) 2 ) [ (uj)2 + (vj)2 ] . (16)
The limiting behavior of label noise SGD is described by a Riemannian gradient flow on Γ as follows:
dxt = −∂Φ(xt)∇R(xt)dt, with x0 = Φ(xinit) ∈ Γ. (17) The goal is to show that the above limiting flow will converge to the underlying groundtruth x∗ = ( u∗ v∗ ) where (u∗, v∗) = ([w∗] 1/2 + , [−w∗] 1/2 + ).

heading: 6.1 LIMITING FLOW CONVERGES TO MINIMIZERS OF REGULARIZER
text: In this subsection we show limiting flow (13) starting from anywhere on Γ converges to the minimizer of regularizer R (by Lemma 6.5). The proof contains two parts: (a) the limiting flow converges; (b) the limit point of the flow cannot be sub-optimal stationary points. These are indeed the most technical and difficult parts of proving the O(κ ln d) upper bound, where the difficulty comes from the fact that the manifold Γ is not compact, and the stationary points of the limiting flow are in fact all located on the boundary of Γ. However, the limiting flow itself is not even defined on the boundary of the manifold Γ. Even if we can extend ∂Φ(·)∇R(·) continuously to entire RD, 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 flow 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 finite trajectory length and thus the convergence. However, since our flow does not satisfy those nice properties, we have to show that the limiting flow satisfies 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/flow 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 flow 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 flow 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 {xt}t≥0 ⊆ RD be generated by the flow defined in (17) with any initialization x0 ∈ Γ. Then x∞ = limt→∞ xt exists. Moreover, x∞ = x∗ is the optimal solution of (18).

heading: 6.2 MINIMIZER OF THE REGULARIZER RECOVERS THE SPARSE GROUNDTRUTH
text: Note 1n ∑n i=1(z j i ) 2 = 1 when zji iid∼ Unif{−1, 1}, we can show minimizing R(x) on Γ, (18), is equivalent to finding the minimum `1 norm solution of Equation (14). Standard results in sparse recovery imply that minimum `1 norm solution recovers with the sparse groundtruth. The gaussian case is more complicated but still can be proved with techniques from Tropp (2015).
minimize R(x) = 1
4n ∑d j=1 (∑n i=1 (zji ) 2 ) [ (uj)2 + (vj)2 ] ,
subject to Z(u 2 − v 2) = Zw∗. (18)
Lemma 6.6. Let z1, . . . , zn i.i.d.∼ Unif({±1}d) orN (0, Id). Then there exist some constants C, c > 0 such that if n ≥ Cκ ln d, then with probability at least 1− e−cn, the optimal solution of (18), (û, v̂), is unique up to sign flips of each coordinate and recovers the groundtruth, i.e., û 2 − v̂ 2 = w∗.

heading: 6.3 LOWER BOUND FOR GRADIENT DESCENT IN THE KERNEL REGIME
text: 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 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 u0 = v0 = α1, with α → ∞. It holds that fi(u0, v0) = 0 and ∇fi(u0, v0) = [αzi,−αzi] 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 fixed throughout the training.
Theorem 6.7. Assume z1, . . . , zn i.i.d.∼ N (0, Id) and yi = w>∗ zi, for all i ∈ [n]. Define the loss with linearized model as L(x) = ∑n i=1(fi(x0) + 〈∇fi(x0), x− x0〉 − yi)2, where x = ( u v ) and
x0 = ( u0 v0 ) = α ( 1 1 ) . Then for any ground truth w∗, any learning rate schedule {ηt}t≥1, and any fixed number of steps T , the expected `2 loss of xT is at least (1− nd ) ‖w∗‖ 2 2, where xT is the T -th iterate of GD on L, i.e., xt+1 = xt − ηt∇L(xt), for all t ≥ 0.

heading: 7 CONCLUSION AND FUTURE WORK
text: We propose a mathematical framework to study the implicit bias of SGD with infinitesimal 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 specific 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 benefit 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.

heading: A PRELIMINARIES ON STOCHASTIC PROCESSES
text: 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 ‖ · ‖, e.g., (R, | · |) and (RD, ‖ · ‖2).

heading: A.1 CADLAG FUNCTION AND METRIC
text: Definition 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 DE [0, T ) be the set of all cadlag function mapping [0, T ) into E . We also use DE [0, T ) to denote the set of all continuous function mapping [0, T ) into E . By definition, CE [0, T ) ⊂ DE [0, T ). Definition A.2 (Continuity modulus). For any function f : [0,∞)→ E and any interval I ⊆ [0,∞), we define
ω(f ; I) = sup s,t∈I
‖f(s)− f(t)‖.
For any N ∈ N and θ > 0, we further define the continuity modulus of continuous f as
ωN (f, θ) = sup 0≤t≤t+θ≤N
{ω(f ; [t, t+ θ])}.
Moreover, the continuity modulus of cadlag f ∈ DE [0,∞) is defined as
ω′N (f, θ) = inf { max i≤r ω(f ; [ti−1, ti) : 0 ≤ t0 < · · · < tr = N, inf i<r (ti − ti−1) ≥ θ } .
Definition A.3 (Jump). For any g ∈ DE [0, T ), we define the jump of g at t to be
∆g(t) = g(t)− g(t−).
For any δ > 0, we define hδ : [0,∞)→ [0,∞) by
hδ(r) = { 0 if r ≤ δ 1− δ/r if r ≥ δ .
We then further define Jδ : DRD [0,∞)→ DRD [0,∞) (Katzenberger, 1991) as Jδ(g)(t) = ∑
0<s≤t
hδ(‖∆g(s)‖)∆g(s). (19)
Definition A.4 (Skorohod metric on DE [0,∞)). For each finite T > 0 and each pair of functions f, g ∈ DE [0,∞), define dT (f, g) as the infimum of all those values of δ for which there exist grids 0 ≤ t0 < t1 < · · · < tm and 0 < s0 < s1 < · · · < · · · < sm, with tk, sk ≥ T , such that |ti − si| ≤ δ for i = 0, . . . , k, and
‖f(t)− g(s)‖ ≤ δ if (t, s) ∈ [ti, ti+1)× [si, si+1)
for i = 0, . . . , k − 1. The Skorohod metric on DE [0,∞) is defined to be
d(f, g) = ∞∑ T=1 2−T min{1, dT (f, g)}.

heading: A.2 STOCHASTIC PROCESSES AND STOCHASTIC INTEGRAL
text: Let (Ω,F , {Ft}t≥0,P) be a filtered probability space.
Definition A.5 (Cross variation). Let X and Y be two {Ft}t≥0-adapted stochastic processes such that X has sample paths in DRD×e [0,∞) and Y has samples paths in DRe [0,∞), then the cross variation of X and Y on (0, t], denoted by [X,Y ](t), is defined to be the limit of
m−1∑ i=0 (X(ti+1)−X(ti))(Y (ti+1)− Y (ti))
in probability as the mesh size of 0 = t0 < t1 < · · · < tm = t goes to 0, if it exists. Moreover, for Y itself, we write
[Y ] = e∑ i=1 [Y i, Y i]
Definition A.6 (Martingale). Let {X(t)}t≥0 be a {Ft}t≥0-adapted stochastic process. If for all 0 ≤ s ≤ t, it holds that
E[X(t) | Fs] = X(s),
then X is called a martingale. Definition A.7 (Semimartingale). Let {X(t)}t≥0 be a {Ft}t≥0-adapted stochastic process. If there exists a sequence of {Ft}t≥0-stopping time, {τk}k≥0, such that
• P[τk < τk+1] = 1, P[limk→∞ τk =∞] = 1,
• and {Xτk(t)}t≥0 is a {Ft}t≥0-adapted martingale,
then X is called a semimartingale. Lemma A.8 (Itô’s Lemma). Let {X(t)}t≥0 be defined through the following Itô drift-diffusion process:
dX(t) = µ(t)dt+ σ(t)dW (t).
where {W (t)}t≥0 is the standard Brownian motion. Then for any twice differentiable function f , it holds that
df(t,X(t)) =
( ∂f
∂t + (∇xf)>µt +
1 2 tr[σ>∇2xfσ]
) dt+ (∇xf)>σ(t)dW (t).

heading: A.3 WEAK CONVERGENCE FOR STOCHASTIC PROCESSES
text: Let (DE [0,∞),A, d) be a metric space equipped with a σ-algebraA and the Skorohod metric defined in the previous subsection.
Let {Xn}n≥0 be a sequence of stochastic processes on a sequence of probability spaces {(Ωn,Fn,Pn)}n≥0 such that each Xn has sample paths in DE [0,∞). Also, let X be a stochastic process on (Ω,F ,P) with sample paths on DE [0,∞). Definition A.9 (Weak convergence). A sequence of stochastic process{Xn}n≥0 is said to converge in distribution or weakly converge to X (written as Xn ⇒ X) if and only if for all A-measurable, bounded, and continuous function f : DE [0,∞)→ R, it holds that
lim n→∞
E [f(Xn)] = E [f(X)] . (20)
Though we define 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-defined if we replace limn→∞ by limη→0. Definition 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 define
ρδ(P,Q) = inf{ > 0 : ∃(X,Y ),P[d(X,Y ) ≥ ] ≤ δ}.
Note this distance is not a metric because it does not satisfy triangle inequality.
Definition 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 define the Prohorov metric as
ρ(P,Q) = inf{ > 0 : ∃(X,Y ),L(X) = P,L(Y ) = Q,P[d(X,Y ) ≥ ] ≤ }.
It can be shown that Xn ⇒ X is equivalent to limn→∞ ρ(Xn, X) = 0. Theorem A.12 (Skorohod Representation Theorem). Suppose Pn, n = 1, 2, . . . and P are probability measures on E such that Pn ⇒ P . Then there is a probability space (Ω,F ,P) on which are defined E-valued random variables Xn, n = 1, 2, . . . and X with distributions Pn and P respectively, such that limn→∞Xn = X a.s.
The main convergence result in Katzenberger (1991) (Theorem B.7) are in the sense of Skorohod metric in Definition A.4, which is harder to understand and use compared to the more common uniform metric (Definition A.13). However, convergence in Skorohod metric and uniform metric indeed coincide with each other when the limit is in CRD [0,∞), i.e., the continuous functions. Definition A.13 (Uniform metric on DE [0,∞)). For each finite T > 0 and each pair of functions f, g ∈ DE [0, T ), the uniform metric is defined to be
dU (f, g;T ) = sup t∈[0,T )
‖f(t)− g(t)‖.
The uniform metric on DE [0,∞) is defined to be
dU (f, g) = ∞∑ T=1 2−T min{1, dU (f, g;T )}.
Lemma A.14 (Problem 7, Section 5, Pollard (2012)). If Xn ⇒ X in the Skorohod sense, and X has sample paths in CRD [0,∞), then Xn ⇒ X in the uniform metric. Remark A.15. We shall note the uniform metric defined above is weaker than supt∈[0,∞) ‖f(t)− g(t)‖. Convergence in the uniform metric on [0,∞] defined in Definition 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.

heading: B LIMITING DIFFUSION OF SGD
text: 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 ⊆ RD, 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 C2 in Lemma B.1. In Appendix B.1 we discuss how different time scalings could affect the coefficients 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 C2 on U .
Proof of Lemma B.1. Applying Theorem 5.1 of Falconer (1983) with f(·) = φ(·, 1) suffices.

heading: B.1 APPROXIMATING SGD BY SDE
text: Let’s first 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) : RD → RD×Ξ 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:
dX(t) = −∇L(X(t))dt+√ησ(X(t))dW (t),
where the time correspondence is t = kη, i.e., X(kη) ≈ xη(k).
Now rescale the above SDE by considering X̃(t) = X(tη), which then yields
dX̃(t) = dX(tη) = −∇L(X(tη))d(tη) +√ησ(X(tη))dW (tη) = −η∇L(X(tη))dt+√ησ(X(tη))dW (tη).
Now we define W ′(t) = 1√ηW (tη), and it’s easy to verify that W ′(t) is also a Ξ-dimensional
brownian motion, which means W ′ d= W , i.e., W and W ′ have the same sample paths in CRd [0,∞). Thus
dX̃(t) = −η∇L(X(tη))dt+ ησ(X(tη))dW ′(t) = −η∇L(X̃(t))dt+ ησ(X̃(t))dW ′(t),
where the time correspondence is t = k, i.e., X̃(k) ≈ xη(k). The above SDE is exactly the same as (2).
Then, to accelerate the above SDE by η−2 times, let’s define X̄(t) = X̃(t/η2). Then it follows that
dX̄(t) = dX̃(t/η2) = −η∇L(X̃(t/η2))dt/η2 + ησ(X̃(t/η2))dW (t/η2)
= −1 η ∇L(X̄(t))dt+ σ(X̄(t))d
( ηW (t/η2) ) 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.

heading: B.2 NECESSARY CONDITIONS
text: Below we collect the necessary conditions imposed on {Zn}n≥1 and {An}n≥1 in Katzenberger (1991).
Condition B.2. The integrator sequence {An}n≥1 is asymptotically continuous: sup t>0 |An(t) −

heading: An(t−)| ⇒ 0 where An(t−) = lims→t−An(s) is the left limit of An at t.
text: Condition B.3. The integrator sequence {An}n≥1 increases infinitely fast: ∀ > 0, inf
t≥0 (An(t+ )−
An(t))⇒∞. Condition B.4 (Eq.(5.1), Katzenberger 1991). For every T > 0, as n→∞, it holds that
sup 0<t≤T∧λn(K)
‖∆Zn(t)‖2 ⇒ 0.
Condition B.5 (Condition 4.2, Katzenberger 1991). For each n ≥ 1, let Yn be a {Fnt }- semimartingale with sample paths in DRD [0,∞). Assume that for some δ > 0 (allowing δ =∞) and every n ≥ 1 there exist stopping times {τmn | m ≥ 1} and a decomposition of Yn − Jδ(Yn) into a local martingale Mn plus a finite variation process Fn such that P[τmn ≤ m] ≤ 1/m, {[Mn](t ∧ τmn ) + Tt∧τmn (Fn)}n≥1 is uniformly integrable for every t ≥ 0 and m ≥ 1, and
lim γ→0 lim sup n→∞
P [
sup 0≤t≤T
(Tt+γ(Fn)− Tt(Fn)) > ] = 0,
for every > 0 and T > 0, where Tt(·) denotes total variation on the interval [0, t]. Lemma B.6. For SGD iterates defined using the notation in Lemma 4.2, the sequences {An}n≥1 and {Zn}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 definition of {An}n≥1. Next, for any > 0 and t ∈ [0, T ], we have
An(t+ )−An(t) = ηn · ⌊ t+
η2n
⌋ − ηn · ⌊ t
η2n
⌋ ≥ t+ − η 2 n
ηn − t ηn = − η2n ηn ,
which implies that inf0≤t≤T (An(t + ) − An(t)) > /(2ηn) for small enough ηn. Then taking n→∞ yields the Condition B.3. For Condition B.4, note that
∆Zn(t) =
{ ηn(1ξk − 1|Ξ|1) if t = k · η 2 n,
0 otherwise.
Therefore, we have ‖∆Zn(t)‖2 ≤ 2ηn for all t > 0. This implies that ‖∆Zn(t)‖2 → 0 uniformly over t > 0 as n→∞, which verifies Condition B.4. We proceed to verify Condition B.5. By the definition of Zn, we know that {Zn(t)}t≥0 is a jump process with independent increments and thus is a martingale. Therefore, by decomposing Zn = Mn + Fn with Mn being a local martingale and Fn a finite variation process, we must have Fn = 0 and Mn is Zn itself. It then suffices to show that [Mn](t ∧ τmn ) is uniformly integrable for every t ≥ 0 and m ≥ 1. Since Mn is a pure jump process, we have
[Mn](t ∧ τmn ) = ∑
0<s≤t∧τmn
‖∆Mn(s)‖22 ≤ ∑
0<s≤t
‖∆Mn(s)‖22
= bt/η2nc∑ k=1 ∥∥∥∥ηn1ξk − ηn|Ξ|1 ∥∥∥∥2 2 ≤ 4 bt/η2nc∑ k=1 η2n ≤ |Ξ|t.
This implies that [Mη](t ∧ τmη ) is universally bounded by 4t, and thus [Mη](t ∧ τmη ) is uniformly integrable. This completes the proof.
Lemma 4.2. Let {ηn}∞n=1 be any positive sequence with limn→∞ ηn = 0, An(t) = ηnbt/η2nc, and Zn(t) = ηn ∑bt/η2nc k=1 (1ξk − 1 Ξ1), where ξ1, ξ2, . . .
i.i.d.∼ Unif([Ξ]). Then with the same initialization Xn(0) = xηn(0) ≡ X(0), Xn(kη2n) defined by (6) is a Katzenberger process and is equal to xηn(k) defined in (1) with LR equal to ηn for all k ≥ 1. Moreover, the counterpart of (7) is
Y (t) = Φ(X(0)) + ∫ t 0 ∂Φ(Y )σ(Y )dW (s) + 1 2 ∫ t 0 ∂2Φ(Y )[Σ(Y )]ds, (8)
where Σ ≡ σσ> and {W (t)}t≥0 is a Ξ-dimensional standard Brownian motion.
Proof of Lemma 4.2. For any n ≥ 1, it suffices to show that given Xn(kη2n) = xηn(k), we further have Xn((k + 1)η2n) = xηn(k + 1). By the definition of Xn(t), we have
Xn((k + 1)η 2 n)−Xn(kη2n) =− ∫ (k+1)η2n kη2n ∇L(Xn(t))dAn(t) + ∫ (k+1)η2n kη2n σ(Xn(t))dZn(t) =−∇L(Xn(kη2n))(An((k + 1)η2n)−An(kη2n)) + σ(Xn(kη2n))(Zn((k + 1)η2n)− Zn(kη2n)) =− ηn∇L(Xn(kη2n)) + ηn ξk(Xn(kη2n)) =− ηn∇L(xηn(k)) + ηn ξk(xηn(k)) = xηn(k + 1)− xηn(k)
where the second equality is because An(t) and Zn(t) are constant on interval [kη2n, (k+ 1)η 2 n). This confirms the alignment between {Xn(kη2n)}k≥1 and {xηn(k)}k≥1.
For the second claim, note that σ(x)EZn(t) ≡ 0 for all x ∈ RD, t ≥ 0 (since the noise has zeroexpectation) and that {Zn(t)− EZn(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 Xn as n→∞ can be obtained by substituting Z with the standard Brownian motion W in (22). This completes the proof.

heading: B.3 KATZENBERGER’S THEOREM FOR ASYMPTOTICALLY CONTINUOUS CASE
text: 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 first introduce some definitions. For each n ≥ 1, let (Ωn,Fn, {Fnt }t≥0,P) be a filtered probability space, Zn an Re-valued cadlag {Fnt }-semimartingale with Zn(0) = 0 and An a real-valued cadlag {Fnt }-adapted nondecreasing process with An(0) = 0. Let σn : U →M(D, e) be continuous with σn → σ uniformly on compact subsets of U . Let Xn be an RD-valued cadlag {Fnt }-semimartingale satisfying, for all compact K ⊂ U ,
Xn(t) = X(0) + ∫ t 0 σ(Xn)dZn + ∫ t 0 −∇L(Xn)dAn (21)
for all t ≤ λn(K) where λn(K) = inf{t ≥ 0 | Xn(t−) /∈ K̊ or Xn(t) /∈ K̊} is the stopping time of Xn 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 , define µn(K) = inf{t ≥ 0 | Yn(t−) /∈ K̊ or Yn(t) /∈ K̊}, then the sequence {(Y µn(K)n , Zµn(K)n , µn(K)} is relatively compact in DRD×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 ≥ 0 a.s., µ ≥ inf{t ≥ 0 | Y (t) /∈ K̊} a.s. and Y (t) admits
Y (t) = Y (0) + ∫ t∧µ 0 ∂Φ(Y (s))σ(Y (s))dZ(s)
+ 1
2 D∑ i,j=1 e∑ k,l=1 ∫ t∧µ 0 ∂i,jΦ(Y (s))σ(Y (s)) ikσ(Y (s))jld[Zk, Zl](s). (22)
We note that by Lemma A.14, convergence in distribution under skorohod metric is equivalent to convergence in distribution under uniform metric Definition 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.

heading: B.4 A USER-FRIENDLY INTERFACE FOR KATZENBERGER’S THEOREM
text: 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 fix some x0 ∈ K. Consider the SGD formulated in Lemma 4.2 where Xηn(0) ≡ x0. Define
Yηn(t) = Xηn(t)− φ(Xηn(0), Aηn(t)) + Φ(Xηn(0))
and µηn(K) = min{t ∈ N | Yηn(t) /∈ K̊}. Then the sequence {(Y µηn (K) ηn , Zηn , µηn(K))}n≥1 is relatively compact in DRD×Rn [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 ≥ 0 | Y (t) /∈ K̊} and Y (t) admits
Y (t) = ∫ t∧µ s=0 ∂Φ(Y (s))σ(Y (s))dW (s) + ∫ t∧µ s=0 1 2 ∑ i,j ∂ijΦ(Y (s))(σ(Y (s))σ(Y (s)) >)ijds
(23)
where {W (s)}s≥0 is the standard Brownian motion and σ(·) is as defined 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 Σ 1 2
‖ (Y (t)) to simplify the equation,
since ∂Φ(Y (t))σ(Y (t)) (∂Φ(Y (t))σ(Y (t)))> = Σ‖(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 > 0, it holds for all sufficiently small LR η that:
ρ2δ(Y µη(K)∧Tη , Y µ(K)∧T ) ≤ . (24)
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 fixed T > 0.
Proof of the first claim of Theorem B.9. Let ET be the event such that µ(K) > T on ET . Then restricted on ET , we have Y (T ∧ µ) = Y (T ∧ µ(K)) as µ ≥ µ(K) holds a.s. We first 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(Y µηm (K) ηm , Y µ)→ 0, a.s.,
which further implies that for any > 0, there exists some N > 0 such that for all m > N P [ d(Y µηm (K)∧T ηm , Y µ∧T ) ≥ ] ≤ δ.
Restricted on ET , 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
P [ d(Y µηm (K)∧T ηm , Y µ(K)∧T ) ≥ ] ≤ P [ {d(Y µηm (K)∧Tηm , Y µ(K)∧T ) ≥ } ∩ ET ] + P [EcT ]
= P [ {d(Y µηm (K)∧Tηm , Y µ∧T ) ≥ } ∩ ET ] + P[EcT ]
≤ P [ d(Y µηm (K)∧T ηm , Y µ∧T ) ≥ ] + P[EcT ]
≤ 2δ
By the definition of the Prohorov metric in Definition A.11, we then get ρ2δ(Y µηm (K)ηm , Y µ(K)∧T ) ≤ for all m > N . Therefore, we have
lim m→∞
ρ2δ(Y µηm (K) ηm , Y µ(K)∧T ) = 0.
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 > 0 such that for all η0 > 0, there exists some η < η0 with ρ2δ(Y µη(K) η , Y µ(K)∧T ) > . Consequently, there is a sequence {ηm}m≥1 satisfying limm→∞ ηm = 0 and ρ2δ(Y µηm (K) ηm , Y
µ(K)∧T ) > 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 ) ≤ for all sufficiently large m, which is a contradiction. This completes the proof.
Proof of the second claim of Theorem B.9. We will first show there exists a sequence of compact set {Km}m≥1 such that ∪∞m=1Km = U and Km ⊆ Km+1. For m ∈ N+, we define Hm = U \(B1/m(0)+RD \U) andKm = Hm∩Bm(0). By definition it holds that ∀m < m′, Hm ⊆ Hm′ and Km ⊆ Km′ . Moreover, since Km is bounded and closed, Km is compact for every m. Now we claim ∪∞m=1Km = U . Note that ∪∞m=1Km = ∪∞m=1Hm ∩ Bm(0) = ∪∞m=1Hm. ∀x ∈ U , since U is open, we know dU (x,RD \ U) > 0, thus there exists m0 ∈ N+, such that ∀m ≥ m0,
x /∈ (B1/m(0) + RD \ U) and thus x ∈ Hm, which implies x ∈ ∪∞m=1Hm. On the other hand, ∀x ∈ RD \ U , it holds that x ∈ (B1/m(0) + RD \ U) for all m ∈ N+, thus x /∈ Hm ⊂ Km.
Therefore, since Y ∈ U and is continuous almost surely, random variables limm→∞ µ(Km) =∞ a.s., which implies µ(Km) converges to∞ in distribution, i,e,, ∀δ > 0, T > 0, ∃m ∈ N+, such that ∀K ⊇ Km, it holds P[µ(K) ≤ T ] ≤ δ. Now we will show for any T > 0 and > 0, there exists η0 such that ρ (Y T , Y Tη ) ≤ for all η ≤ η0. Fixing any T > 0, for any > 0, let δ = 4 , then from above we know exists compact set K, such that P(µ(K) ≤ T ) ≤ δ. We further pick K ′ = K + B2 ′(0), where ′ can be any real number satisfying 0 < ′ < and K ′ ⊆ U . Such ′ exists since U is open. Note K ⊆ K ′, we have P(µ(K ′) ≤ T ) ≤ P(µ(K) ≤ T ) ≤ δ. Thus by the first claim of Theorem B.9, there exists η0 > 0, such that for all η ≤ η0, we have ρ2δ(Y µη(K ′)∧T η , Y µ(K ′)∧T ) ≤ 2−dTe ′.
Note that ρδ(Y µ(K)∧T , Y µ(K ′)∧T ) = 0, so we have for all η ≤ η0,
ρ3δ(Y µ(K)∧T , Y µη(K ′)∧T η ) ≤ 2−dTe ′.
By the definition of δ-Prohorov distance in Definition A.10, we can assume (Y µ(K)∧T , Y µη(K ′)∧T η ) is already the coupling such that P [ dU (Y µ(K)∧T , Y µη(K ′)∧T η ) ≥ 2−dTe ′ ] ≤ 3δ. Below we want
to show ρ3δ(Y µ(K)∧T , Y Tη ) ≤ 2−dTe ′. Note that for all t ≥ 0, Y µ(K)∧T (t) ∈ K, thus we know if µη(K ′) ≤ T , then
dU (Y µ(K)∧T , Y µη(K ′)∧T η ) ≥ 2−dTe ∥∥∥Y µ(K)∧T (µη(K ′))− Y µη(K′)∧Tη (µη(K ′)∥∥∥ 2
≥ 2−dTedU (K,Rd/K ′) ≥ 2−dTe ′.
On the other hand, if µη(K ′) > T , then Y Tη = Y µη(K ′)∧T η . Thus we can conclude that dU (Y µ(K)∧T , Y Tη ) ≥ 2−dTe ′ implies dU (Y µ(K)∧T , Y µη(K ′)∧T
η ) ≥ 2−dTe ′. Therefore, we further have
P [ dU (Y µ(K)∧T , Y Tη ) ≥ 2−dTe ′ ] ≤ P [ dU (Y µ(K)∧T , Y µη(K ′)∧T η ) ≥ 2−dTe ′ ] ≤ 3δ,
that is,
ρ3δ(Y µ(K)∧T , Y Tη ) ≤ 2−dTe ′.
Finally, since ρδ(Y T , Y µ(K)∧T ) = 0, we have for all η ≤ η0,
ρ (Y T , Y Tη ) = ρ 4δ(Y T , Y Tη ) ≤ ρ3δ(Y µ(K)∧T , Y Tη ) + ρδ(Y T , Y µ(K)∧T ) ≤ 2−dTe ′ + 0 ≤ ,
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 reflects the time scaling of Phase II.

heading: C EXPLICIT FORMULA OF THE LIMITING DIFFUSION
text: 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.

heading: C.1 EXPLICIT EXPRESSION OF THE DERIVATIVES
text: For any x ∈ Γ, we choose an orthonormal basis of Tx(Γ) as {v1, . . . , vD−M}. Let {vD−M+1, . . . , vD} be an orthonormal basis of T⊥x (Γ) so that {vi}i∈[D] is an orthonormal basis of RD. Lemma C.1. For any x ∈ Γ and any v ∈ Tx(Γ), it holds that∇2L(x)v = 0.
Proof. For any x ∈ Tx(Γ), let {x(t)}t≥0 be a parametrized smooth curve on Γ such that x(0) = x and dx(t)dt ∣∣ t=0 = v. Then ∇L(xt) = 0 for all t. Thus 0 = d∇L(xt)dt ∣∣ t=0 = ∇2L(x)v.
Lemma C.2. For any x ∈ RD, it holds that ∂Φ(x)∇L(x) = 0 and
∂2Φ(x)[∇L(x),∇L(x)] = −∂Φ(x)∇2L(x)∇L(x).
Proof. Fixing any x ∈ RD, let dx(t)dt = −∇L(x(t)) be initialized at x(0) = x. Since Φ(x(t)) = Φ(x) for all t ≥ 0, we have
d dt Φ(x(t)) = −∂Φ(x(t))∇L(x(t)) = 0.
Evaluating the above equation at t = 0 yields ∂Φ(x)∇L(x) = 0. Moreover, take the second order derivative and we have
d2 dt2 Φ(xt) = −∂2Φ(x(t))
[ dx(t)
dt ,∇L(x(t))
] + ∂Φ(x(t))∇2L(x(t))dx(t)
dt = 0.
Evaluating at t = 0 completes the proof.
Now we can prove Lemma 4.3, restated in below.
Lemma 4.3. For any x ∈ Γ, ∂Φ(x) ∈ RD×D is the projection matrix onto tangent space Tx(Γ).
Proof of Lemma 4.3. For any v ∈ Tx(Γ), let {v(t), t ≥ 0} be a parametrized smooth curve on Γ such that v(0) = x and dv(t)dt ∣∣ t=0
= v. Since v(t) ∈ Γ for all t ≥ 0, we have Φ(v(t)) = v(t), and thus
dv(t)
dt ∣∣∣∣ t=0 = d dt Φ(v(t)) ∣∣∣∣ t=0 = ∂Φ(x) dv(t) dt ∣∣∣∣ t=0 .
This implies that ∂Φ(x)v = v for all v ∈ Tx(Γ).
Next, for any u ∈ T⊥x (Γ) and t ≥ 0, consider expanding∇L(x+ t∇2L(x)†u) at t = 0: ∇L ( x+ t∇2L(x)†u ) = ∇2L(x) · t∇2L(x)†u+ o(t) = tu+ o(t)
where the second equality follows from the assumption that∇2L(x) is full-rank when restricted on T⊥x (Γ). Then since ∂Φ is continuous, it follows that
lim t→0 ∂Φ(x+ t∇2L(x)†u)∇L(x+ t∇2L(x)†u) t = lim t→0 ∂Φ(x+ t∇2L(x)†)(u+ o(1))
= ∂Φ(x)u.
By Lemma C.2, we have ∂Φ(x+ t(∇2L(x))†u))∇L(x+ t(∇2L(x))†u) = 0 for all t > 0, which then implies that ∂Φ(x)u = 0 for all u ∈ T⊥x (Γ). Therefore, under the basis {vi, . . . , vN}, ∂Φ(x) is given by
∂Φ(x) =
( ID−M 0
0 0
) ∈ RD×D,
that is, the projection matrix onto Tx(Γ).
Lemma C.3. For any x ∈ Γ, it holds that ∂Φ(x)∇2L(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 ∈ RD and v ∈ Tx(Γ), it holds that
∂2Φ(x)[v, u] = −∂Φ(x)∂2(∇L)(x)[v,∇2L(x)†u]−∇2L(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 dv(t)
dt ∣∣ t=0
= v. We define P (t) = ∂Φ(v(t)), P⊥(t) = ID − P (t) and H(t) = ∇2L(v(t)) for all t ≥ 0. By Lemma C.1 and 4.3, we have
P⊥(t)H(t) = H(t)P⊥(t) = H(t), (25)
Denote the derivative of P (t), P⊥(t) and H(t) with respect to t as P ′(t), (P⊥)′(t) and H ′(t). Then differentiating with respect to t, we have
(P⊥)′(t)H(t) = H ′(t)− P⊥(t)H ′(t) = P (t)H ′(t). (26)
Then combining (25) and (26) and evaluating at t = 0, we have
P ′(0)H(0) = −(P⊥)′(0)H(0) = −P (0)H ′(0) (27)
We can decompose P ′(0) and H(0) as follows
P ′(0) =
( P ′11(0) P ′ 12(0)
P ′21(0) P ′ 22(0)
) , H(0) = ( 0 0 0 H22(0) ) , (28)
where P ′11(0) ∈ R(D−M)×(D−M) and H22 is the hessian of L restricted on T⊥x (Γ). Also note that
P (0)H ′(0)P⊥(0) =
( ID−M 0
0 0
)( H ′11(0) H ′ 12(0)
H ′21(0) H ′ 22(0) )( 0 0 0 IM ) = ( 0 H ′12(0) 0 0 ) ,
and thus by (28) we have
P ′(0)H(0) = ( 0 P ′12(0)H22(0) 0 P ′22(0)H22(0) ) = ( 0 −H ′12(0) 0 0 ) .
This implies that we must have P ′22(0) = 0 and P ′ 12(0)H22(0) = H ′ 12(0). Similarly, by taking transpose in (28), we also have H22(0)P ′21(0) = −H ′21(0). It then remains to determine the value of P ′11(0). Note that since P (t)P (t) = P (t), we have P ′(t)P (t) + P (t)P ′(t) = P ′(t), evaluating at t = 0 yields
2P ′11(0) = P ′ 11(0).
Therefore, we must have P ′11(0) = 0. Combining the above results, we obtain
P ′(0) = −P (0)H ′(0)H(0)† −H(0)†H ′(0)P (0).
Finally, recall that P (t) = ∂Φ(v(t)), and thus
P ′(0) = d
dt ∂Φ(v(t)) ∣∣∣∣ t=0 = ∂2Φ(x)[v].
Similarly, we have H ′(0) = ∂2(∇L)(x)[v], and it follows that ∂2Φ(x)[v] = −∂Φ(x)∂2(∇L)(x)[v]∇2L(x)† −∇2L(x)†∂2(∇L)(x)[v]∂Φ(x).
Lemma C.5. For any x ∈ Γ and u ∈ T⊥x (Γ), it holds that ∂2Φ(x)[uu> +∇2L(x)†uu>∇2L(x)] = −∂Φ(x)∂2(∇L)(x)[∇2L(x)†uu>].
Proof of Lemma C.5. For any u ∈ T⊥x (Γ), we define u(t) = x+ t∇2L(x)†u for t ≥ 0. By Taylor approximation, we have
∇L(u(t)) = t∇2L(x)∇2L(x)†u+ o(t) = tu+ o(t) (29) and
∇2L(u(t)) = ∇2L(x) + t∂2(∇L)(x)[∇2L(x)†u] + o(t). (30) Combine (29) and (30) and apply Lemma C.2, and it follows that
0 = ∂2Φ(u(t))[∇L(u(t)),∇L(u(t))] + ∂Φ(u(t))∇2L(u(t))∇L(u(t)) = t2∂2Φ(u(t))[u+ o(1)](u+ o(1)) + t2∂Φ(u(t))∂2(∇L)(x)[∇2L(x)†u](u+ o(1))
+ t2 ∂Φ(u(t))
t ∇2L(x)(u+ o(1))
= t2∂2Φ(u(t))[u+ o(1)](u+ o(1)) + t2∂Φ(u(t))∂2(∇L)(x)[∇2L(x)†u](u+ o(1))
+ t2 ∂Φ(u(t))− ∂Φ(x)
t ∇2L(x)(u+ o(1))
where the last equality follows from Lemma C.3. Dividing both sides by t2 and letting t→ 0, we get ∂2Φ(x)[u]u+ ∂Φ(x)∂2(∇L)(x)[∇2L(x)†u]u+ ∂2Φ(x)[∇2L(x)†u]∇2L(x)u = 0.
Rearranging the above equation completes the proof.
With the notion of Lyapunov Operator in Definition 4.4, Lemma C.5 can be further simplified into Lemma C.6. Lemma C.6. For any x ∈ Γ and Σ ∈ span{uu> | u ∈ T⊥x (Γ)},
〈∂2Φ(x),Σ〉 = −∂Φ(x)∂2(∇L)(x)[L−1∇2L(x)(Σ)]. (31)
Proof of Lemma C.6. Let A = uu> + ∇2L(x)†uu>∇2L(x) and B = ∇2L(x)†uu>. The key observation is that A+A> = L∇2L(x)(B +B>). Therefore, by Lemma C.5, it holds that
∂2Φ(x)[L∇2L(x)(B+B>)] = ∂2Φ(x)[A+A>] = 2∂Φ(x)∂2(∇L)(x)[B] = ∂Φ(x)∂2(∇L)(x)[B+B>]. Since∇2L(x)† is full-rank when restricted to T⊥x (Γ), we have span{∇2L(x)†uu>+uu>∇2L(x)† | u ∈ T⊥x (Γ)} = span{uu> | u ∈ T⊥x (Γ)}. Thus by the linearity of above equation, we can replace B +B> by any Σ ∈ span{uu> | u ∈ T⊥x (Γ)}, resulting in the desired equation.
Then Lemma 4.5 directly follows from Lemma C.4 and C.5.

heading: C.2 TANGENT NOISE COMPENSATION ONLY DEPENDENDS ON THE MANIFOLD ITSELF
text: 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 Ux of x and two loss functions L and L′ that define the same manifold Γ locally in Ux, i.e., Γ ∩ Ux = {x | ∇L(x) = 0} = {x | ∇L′(x) = 0}. Then for any v ∈ Tx(Γ), it holds that (∇2L(x))†∂2(∇L)(x) [v, v] = (∇2L′(x))†∂2(∇L′)(x) [v, v]. Proof of Lemma C.7. Let {v(t)}t≥0 be a smooth curve on Γ with v(0) = x and dv(t)dt ∣∣ t=0
= v. Since v(t) stays on Γ, we have ∇L(v(t)) = 0 for all t ≥ 0. Taking derivative for two times yields ∂2(∇L)(v(t))[dv(t)dt , dv(t) dt ] + ∇ 2L(v(t))d 2v(t) dt2 = 0. Evaluating it at t = 0 and multiplying both sides by ∇2L(x)†, we get
∇2L(x)†∂2(∇L)(x) [v, v] = −∇2L(x)†∇2L(x)d 2v(t)
dt2 ∣∣∣∣ t=0 = −∂Φ(x)d 2v(t) dt2 ∣∣∣∣ t=0 .
Since ∂Φ(x) is the projection matrix onto Tx(Γ) by Lemma 4.3, it does not depend on L, so analogously we also have ∇2L′(x)†∂2(∇L′)(x) [v, v] = −∂Φ(x)d
2v(t) dt2 ∣∣ t=0 as well. The proof is
thus completed. Note that ∂Φ(x)d 2v(t) dt2 ∣∣ 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.)

heading: C.3 PROOF OF RESULTS IN SECTION 5
text: 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 Σ ≡ ID on Γ, SDE (10) is then
dY (t) = ∂Φ(Y )dW + 1
2 ∇2L(Y )†∂2(∇L)(Y ) [∂Φ(Y )] dt︸ ︷︷ ︸
Brownian Motion on Manifold
− 1 2 ∂Φ(Y )∇(ln |∇2L(Y )|+)dt︸ ︷︷ ︸
Normal Regularization
(11)
where |∇2L(Y )|+ = limα→0 |∇ 2L(Y )+αID|
αD−rank(∇2L(Y )) is the pseudo-determinant of∇2L(Y ). |∇2L(Y )|+ is
also equal to the sum of log of non-zero eigenvalue values of∇2L(Y ).
Proof of Corollary 5.1. Set Σ‖ = ∂Φ, Σ⊥ = ID − ∂Φ and Σ⊥,‖ = Σ‖,⊥ = 0 in the decomposition of Σ by Lemma 4.5, and we need to show ∂Φ∇(ln |Σ|+) = ∂2(∇L)[(∇2L)†]. Holbrook (2018) shows that the gradient of pseudo-inverse determinant satisfies ∇|A|+ = |A|+A†. Thus we have for any vector v ∈ RD, 〈 v,∇ ln |∇2L|+ 〉 = 〈 |∇2L|+∇2L |∇2L|+ , ∂ 2(∇L)[v] 〉
=〈 ∇2L, ∂2(∇L)[v] 〉 = ∂2(∇L)[v,∇2L] = 〈 v, ∂2(∇L)[(∇2L)†] 〉 , which completes the proof.
Corollary 5.2 (Limiting Flow for Label Noise). If Σ ≡ c∇2L on Γ for some constant c > 0, SDE (10) can be simplified into (13) where the regularization is from the noise in normal space.
dY (t) = −1/4∂Φ(Y (t))∇ tr[c∇2L(Y (t))]dt. (13)
Proof of Corollary 5.2. Since Σ = c∇2L, here we have Σ⊥ = Σ and Σ‖,Σ⊥,‖,Σ‖,⊥ = 0. Thus it suffices to show that 2∂2(∇L) [ L−1∇2L(Σ⊥) ] = ∇ tr[∇2L]. Note that for any v ∈ RD,
v>∇ tr[∇2L] = 〈 ID, ∂ 2(∇L)[v] 〉 = 〈 ID − ∂Φ, ∂2(∇L)[v] 〉 , (32)
where the second equality is because the the tangent space of symmetric rank-n matrices at ∇2L is {A∇2L+∇2LA> | A ∈ RD×D}, and every element in this tangent space has zero inner-product with ∂Φ by Lemma 4.3. Also note that L−1∇2L(∇ 2L) = 12 (ID − ∂Φ), thus 〈 ID − ∂Φ, ∂2(∇L)[v] 〉 =
2 〈 L−1∇2L(∇ 2L), ∂2(∇L)[v] 〉 = 2v>∂2(∇L)[L−1∇2L(∇ 2L)].
C.4 EXAMPLE: k-PHASE MOTOR
We also give an example with rigorous proof where the implicit bias induced by noise in the normal space cannot be characterized by a fixed regularizer, which was first discovered by Damian et al. (2021) but was only verified via experiments.
Note the normal regularization in both cases of label noise and isotropic noise induces Riemmanian gradient flow against some regularizer, it’s natural to wonder if the limiting flow 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 flow moves around a cycle at a constant speed and never stops using our framework.
Suppose dimension D = k + 2 ≥ 5. For each x ∈ RD, we decompose x = ( x1:2 x3:D ) where x1:2 ∈ R2
and x3:D ∈ RD−2. Let Qθ ∈ R2×2 be the rotation matrix of angle θ, i.e., Qθ = ( cos θ − sin θ sin θ cos θ ) and the loss L(x) := 18 (‖x1:2‖ 2 2 − 1)2 + 12 ∑D j=3(2 + 〈 Qj−3α v, x1:2 〉 x2j , where α = 2π D−2 and v is any vector in R2 with unit norm. Here the manifold is given by Γ := {x | L(x) = 0} = {x ∈ RD | x21 + x 2 2 = 1, xj = 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 {x21 + x22 = 1}, and the goal is to make the vector field 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∇2L(Σ⊥) 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 Σ ∈ RD×D be given by Σij(x) = (1+ 〈 Qj−3α v,Q−π/2x1:2 〉 )(2+ 〈 Qj−3α v, x1:2 〉 ), 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.
Y1:2(t) = Qt(D−2)/2Y1:2(0) and Y3:D(t) ≡ 0. (33)
Proof of Lemma C.8. Note that for any x ∈ Γ, it holds that
( ∇2L(x) ) ij =  2 + 〈 Qj−3α v, x1:2 〉 if i = j ≥ 3,
xixj if i, j ∈ {1, 2}, 0 otherwise.
(34)
Then clearly Σ only brings about noise in the normal space, and specifically, it holds that L−1∇2L(x)(Σ(x)) = diag(0, 0, 1 + 〈 Q0αv,Q−π/2x1:2 〉 , . . . , 1 + 〈 QD−3α v,Q−π/2x1:2 〉 ). Further note that, by the special structure of the hessian in (34) and Lemma C.3, for any x ∈ Γ, we have ∂Φ(x) = (x2,−x1, 0, . . . , 0)>(x2,−x1, 0, . . . , 0) = (Q−π/2x1:2 0 )(Q−π/2x1:2 0 )> . Combining these
facts, the dynamics of the first two coordinates in SDE (10) can be simplified into
dx1:2(t)
dt = −
( 1
2 ∂Φ(x(t))∂2(∇L)(x(t))[L−1∇2L(Σ(x(t))] ) 1:2
= −1 2 Q−π/2x1:2x > 1:2Q > −π/2 D∑ j=3 ( 1 + 〈 Qj−3α v,Q−π/2x1:2 〉) ∇1:2(∂jjL)(x)
= −1 2 Q−π/2x1:2
〈 Q−π/2x1:2,
D∑ j=3 ( 1 + 〈 Qj−3α v,Q−π/2x1:2 〉) Qj−3α v
〉
= −1 2 Q−π/2x1:2 〈Q−π/2x1:2, D∑ j=3 Qj−3α v 〉 + D∑ j=3 〈 Qj−3α v,Q−π/2x1:2 〉2 = −1
2 Q−π/2x1:2
( 0 +
D − 2 2 ∥∥Q−π/2x1:2∥∥22) = D − 22 Qπ/2x1:2, where the second to the last equality follows from the property of Qα and the last equality follows from the fact that ‖x1:2‖22 = 1 for all x ∈ Γ. Note we require k ≥ 3 (or D ≥ 5) to allow∑D j=3 〈 Qj−3α v,Q−π/2x1:2 〉2 = D−22
∥∥Q−π/2x1:2∥∥22. On the other hand, we have dx3:D(t)dt = 0 as ∂Φ kills the movement on that component.
The proof is completed by noting that the solution of x1:2 is
x1:2(t) = exp
( t · D − 2
2 Qπ/2
) x1:2(0),
and by Lemma C.9,
exp ( t · D − 2
2 Qπ/2
) = (exp(Qπ/2)) t(D−2) 2 = Q t(D−2) 2
1 = Q t(D−2) 2 .
Lemma C.9. exp( (
0 −1 1 0
) ) = ( cos 1 − sin 1 sin 1 cos 1 ) .
Proof. By definition, for matrix A = (
0 −1 1 0
) , exp(A) = ∑∞ t=0 At t! . Note A2 = −I , A 3 = −A and
A4 = I , and by using this pattern, we can easily check that ∞∑ t=0 At t! = ( ∑∞ i=0(−1)i 1 (2i)! − ∑∞ i=0(−1)i 1 (2i+1)!∑∞ i=0(−1)i 1 (2i+1)! ∑∞ i=0(−1)i 1 (2i)! ) = ( cos 1 − sin 1 sin 1 cos 1 ) .

heading: D PROOF OF RESULTS IN SECTION 6
text: 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 ∈ RD, we denote by Bpr (u) the `p norm ball of radius r centered at u. We also denote vi:j = (vi, vi+1, . . . , vj)> for i, j ∈ [D].

heading: D.1 PROOF OF THEOREM 6.1
text: 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 xinit (except a zero-measure set) and any > 0, there exist η0, T > 0 such that for any η < η0, OLM trained with label noise SGD (12) with LR equal to η for bT/η2c steps returns an -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 satisfies |x∗| = ψ(w∗). Then on this event, for any > 0, by Lemma 6.5, there exists some T > 0 such that xT given by the Riemannian gradient flow (17) satisfies that xT is an /2-optimal solution of the OLM. For this T , by Theorem 4.6, we know that the bT/η2c-th SGD iterate, xη(bT/η2c), satisfies ‖xη(bT/η2c) − xT ‖2 ≤ /2 with probability at least 1 − e−Ω(n) for all sufficiently small η > 0, and thus xη(bT/η2c) is an -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.

heading: D.2 PROOF OF LEMMA 6.2
text: Recall that for each i ∈ [n] fi(x) = f(u, v) = z>i (u 2 − v 2),∇fi(x) = 2 ( zi u zi v ) , and K(x) = (Kij(x))i,j∈[n] where each Kij(x) = 〈∇fi(x),∇fj(x)〉. Then
∇2`i(x) = 2 ( zi u −zi v )( (zi u)> −(zi v)> ) + (fi(u, v)− yi) · diag(zi, zi).
So for any x ∈ Γ, it holds that
∇2L(x) = 2 n n∑ i=1 ( zi u −zi v )( (zi u)> −(zi v)> ) . (35)
Lemma D.1. For any fixed x ∈ RD, suppose {∇fi(x)}i∈[n] is linearly independent, then K(x) is full-rank.
Proof of Lemma D.1. Suppose otherwise, then there exists some λ ∈ Rn such that λ 6= 0 and λ>K(x)λ = 0. However, note that
λ>K(x)λ = n∑ i,j=1 λiλjKij(x)
= n∑ i,j=1 λiλj〈∇fi(x),∇fj(x)〉
= ∥∥∥∥∥ n∑ i=1 λi∇fi(x) ∥∥∥∥∥ 2
2
,
which implies that ∑n i=1 λ
i∇fi(x) = 0. This is a contradiction since by assumption {∇fi(x)}i∈[n] is linearly independent.
Lemma 6.2. Consider the loss function L defined in (14) and manifold Γ defined in (15). If data is full rank, i.e., rank({zi}i∈[n]) = n, then it holds that (a). Γ is a smooth manifold of dimension D− n; (b). rank(∇2L(x)) = n for all x ∈ Γ. In particular, it holds that rank({zi}i∈[n]) = n holds with probability 1 for Gaussian distribution and with probability 1− cd for Boolean distribution for some constant c ∈ (0, 1).
Proof of Lemma 6.2. (1) By preimage theorem (Banyaga & Hurtubise, 2013), it suffices to check the jacobian [∇f1(x), . . . ,∇fn(x)] = 2[ ( z1 u −z1 v ) , . . . , ( zn u −zn v ) ] is full rank. Similarly, for the second
claim, due to (35). it is also equivalent to show that { ( zi u −zi v ) }i∈[n] is of rank n.
Since ( u v ) ∈ Γ ⊂ U , each coordinate is non-zero, thus we only need to show that {zi}i∈[n] is of rank n. This happens with probability 1 in the Gaussian case, and probability at least 1 − cd for some constant c ∈ (0, 1) by Kahn et al. (1995). This completes the proof.

heading: D.3 PROOF OF LEMMA 6.3
text: We first establish some auxiliary results. The following lemma shows the PL condition along the trajectory of gradient flow.
Lemma D.2. ‖∇L(xt)‖2 ≥ λmin(ZZ>) mini∈[d] |ui0vi0|L(xt).
To prove Lemma D.2, we need the following invariance along the gradient flow.
Lemma D.3. Along the gradient flow generated by ∇L(xt), ujtv j t stays constant for all j ∈ [d]. Thus, sign(ujt ) = sign(u j 0) and sign(v j t ) = sign(v j 0) for any j ∈ [d].
Proof of Lemma D.3.
∂ ∂t (ujtv j t ) = ∂ujt ∂t · vjt + u j t · ∂vjt ∂t
= ∇uL(ut, vt)j · vjt + u j t · ∇vL(ut, vt)j
= 1
4 n∑ i=1 (fi(ut, vt)− yi)zji u j tv j t − ujt 4 n∑ i=1 (fi(ut, vt)− yi)zji v j t = 0.
Therefore, any sign change of uj , vj would enforce ujt = 0 or v j t = 0 for some t > 0 since u j t , v j t are continuous in time t. This immediately leads to a contradiction to the invariance of ujtv j t .
We then can prove Lemma D.2.
Proof of Lemma D.2. Note that
‖∇L(x)‖22 = 1
4 n∑ i,j=1 (fi(x)− yi)(fj(x)− yj) 〈∇fi(x),∇fj(x)〉
≥1 4 n∑ i=1 (fi(x)− yi)2λmin(K(x))
=L(x)λmin(K(x)),
where K(x) is a n × n p.s.d. matrix with Kij(x) = 〈∇fi(x),∇fj(x)〉. Below we lower bound λmin(K(x)), the smallest eigenvalue of K(x). Note that Kij(x) = ∑n h=1 z h i z h j ((u h t ) 2 + (vht ) 2), and we have
K(x) = Zdiag(u 2t + v 2 t )Z > Zdiag(|utvt|)Z> (∗) = Zdiag(|u0v0|)Z> min
i∈[d] |ui0vi0|ZZT .
where (∗) is by Lemma D.3. Thus λmin(K) ≥ mini∈[d] |ui0vi0|λmin(ZZT ), 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({zi}i∈[n]) = n, which implies rank({∇fi(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 defined in (14), manifold Γ and its open neighborhood defined in (15). For gradient flow dxtdt = −∇L(xt) starting at any x0 ∈ U , it holds that Φ(x0) ∈ Γ.
Proof of Lemma 6.3. It suffices to prove gradient flow dxtdt = −∇L(xt) converges when t→∞, as long as x0 ∈ 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 limt→∞ xt exists. By Lemma D.16, denote C = mini∈[d] |ui0vi0|λmin(ZZ>), then λmin(K(xt)) ≥ C for all t ≥ 0. Thus,∥∥∥∥dxtdt ∥∥∥∥ = ‖∇L(xt)‖ ≤ ‖∇L(xt)‖22√CL(xt) = − dL(xt) dt√ L(xt) = − 1 2 √ C d √ L(xt) dt . Thus the total GF trajectory length is bounded by ∫∞ t=0 ∥∥dxt dt ∥∥dt ≤ ∫∞ t=0 − 1 2 √ C d √ L(xt) ddt dt ≤ L(x0) 2 √ C
, where the last inequality uses that L is non-negative over RD. Therefore, the GF must converge.

heading: D.4 PROOF OF RESULTS IN SECTION 6.2
text: To study the optimal solution to (18), we consider the corresponding d-dimensional convex program in terms of w ∈ Rd, which has been studied in Tropp (2015):
minimize R(w) = 1
4n d∑ j=1 ( n∑ i=1 (zji ) 2 ) |wj |,
subject to Zw = Zw∗.
(36)
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 defined as follows. Definition D.5 (Canonical Parametrization). For any w ∈ Rd, we define ( u v ) = ψ(w) = ([w>] 1/2 + , [−w>] 1/2 + )
> as the canonical parametrization of w. Clearly, it holds that u 2 − v 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 flips 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 flips of each coordinate. In particular, one of them is given by (ũ∗, ṽ∗) = ψ(w∗), that is, the canonical parametrization of w∗.
Proof of Lemma D.6. Let (û, v̂) be any optimal solution of (18) and we define ŵ = û 2 − v̂ 2, which is also feasible to (36). By the optimality of w∗, we have
d∑ j=1 ( n∑ i=1 (zji ) 2 ) |wj∗| ≤ d∑ j=1 ( n∑ i=1 (zji ) 2 ) |ŵj | ≤ d∑ j=1 ( n∑ i=1 (zji ) 2 ) [(ûj)2 + (v̂j)2]. (37)
On the other hand, (ũ∗, ṽ∗) = ψ(w∗) is feasible to (18). Thus, it follows from the optimality of (û, v̂) that
d∑ j=1 ( n∑ i=1 (zji ) 2 ) [(ûj)2 + (v̂j)2] ≤ d∑ j=1 ( n∑ i=1 (zji ) 2 ) [(ũj∗) 2 + (ṽj∗) 2] = d∑ j=1 ( n∑ i=1 (zji ) 2 ) |wj∗|.
(38)
Combining (37) and (38) yields d∑ j=1 ( n∑ i=1 (zji ) 2 ) [(ûj)2 + (v̂j)2] = d∑ j=1 ( n∑ i=1 (zji ) 2 ) |wj∗| = d∑ j=1 ( n∑ i=1 (zji ) 2 ) |(ûj)2 − (v̂j)2|
(39)
which implies that û 2 − v̂ 2 is also an optimal solution of (36). Since w∗ is the unique optimal solution to (36), we have û 2 − v̂ 2 = w∗. Moreover, by (39), we must have û 2 = [w∗]+ and û 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.8 and D.7 for clarity.
Lemma D.7 (Boolean Case). Let z1, . . . , zn i.i.d.∼ Unif({±1}d). There exist some constants C, c > 0 such that if the sample size n satisfies
n ≥ C[κ ln(d/κ) + κ]
then with probability at least 1− e−cn2 , the optimal solution of (18), (û, v̂), is unique up to sign flips of each coordinate and recovers the groundtruth, i.e., û 2 − v̂ 2 = w∗.
Proof of Lemma D.7. By the assumption that z1, . . . , zn i.i.d.∼ Unif({±1}d), we have ∑n i=1(z j i )
2 = n for all j ∈ [d]. Then (36) is equivalent to the following optimization problem:
minimize g(w) = ‖w‖1, subject to Zw = Z(u 2∗ − v 2∗ ).
(40)
This model exactly fits 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 2∗ − v 2∗ with probability at least 1− e−ch 2
for some constant c > 0, given that the sample size satisfies
n ≥ C(κ ln(d/κ) + κ+ h)
for some absolute constant C > 0. Choosing h = n2C and then adjusting the choices of C, c appropriately yield the desired result. Finally, applying Lemma D.6 finishes the proof.
The Gaussian case requires more careful treatment.
Lemma D.8 (Gaussian Case). Let z1, . . . , zn i.i.d.∼ N (0, Id). There exist some constants C, c > 0 such that if the sample size satisfies
n ≥ Cκ ln d,
then with probability at least 1− (2d+ 1)e−cn, the optimal solution of (18), (û, v̂), is unique up to sign flips of each coordinate of û and v̂ and recovers the groundtruth, i.e., û 2 − v̂ 2 = w∗.
Proof of Lemma D.8. Since z1, . . . , zn i.i.d.∼ N (0, Id), we have
P [ n∑ i=1 (zji ) 2 ∈ [n/2, 3n/2],∀j ∈ [d] ] ≥ 1− 2de−cn
for some constant c > 0, and we denote this event by En. Therefore, on En, we have
n
2 D∑ j=1 [(uj)2 + (vj)2] ≤ R(x) ≤ 3n 2 D∑ j=1 [(uj)2 + (vj)2]
or equivalently,
n 2 (‖u 2‖1 + ‖v 2‖1) ≤ R(x) ≤ 3n 2 (‖u 2‖1 + v 2‖1).
Define w∗ = u 2∗ − v 2∗ , and (36) is equivalent to the following convex optimization problem
minimize g(w) = d∑ j=1 ( n∑ i=1 (zji ) 2 ) · |wj + wj∗|,
subject to Zw = 0.
(41)
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
λmin(Z;D(g, w∗)) ≤ ‖Zw‖2 ‖w‖2 = 0
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 sufficiently large n. On En, it holds that g belongs to the following function class
G = h : Rd → R | h(w) = d∑ j=1 υi|wi|, υ ∈ Υ  with Υ = {υ ∈ Rd : υj ∈ [0.5, 1.5],∀j ∈ [d]}. We identify gυ ∈ G with υ ∈ Υ, then D(g, w∗) ⊆ ∪υ∈ΥD(gυ, w∗)) := DΥ, which further implies that
λmin(Z;D(g, w∗)) ≥ λmin(Z;DΥ).
Recall the definition of minimum conic singular value (Tropp, 2015):
λmin(Z;DΥ) = inf p∈DΥ∩Sd−1 sup q∈Sn−1 〈q, Zp〉.
where Sn−1 denotes the unit sphere in Rn. Applying the same argument as in (Tropp, 2015) yields
P [ λmin(Z;DΥ) ≥ √ n− 1− w(DΥ)− h ] ≥ 1− e−h 2/2.
Take the intersection of this event with En, and we obtain from a union bound that
λmin(Z;D(g, w∗)) ≥ √ n− 1− w(DΞ)− h (42)
with probability at least 1− e−h2/2 − 2de−cn. It remains to determine w(DΥ), which is defined as
w(DΥ) = Ez∼N (0,Id)
[ sup
p∈DΥ∩Sd−1 〈z, p〉
] = Ez∼N (0,Id) [ sup υ∈Υ sup p∈D(gυ,x∗)∩Sd−1 〈z, p〉 ] . (43)
Without loss of generality, we assume that w∗ = (w1∗, . . . , w κ ∗ , 0, . . . , 0) > with w1∗, . . . , 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∗) ∩ Sd−1, there exists some τ > 0 such that gυ(w∗ + τ · p) ≤ gυ(w∗), i.e.,
d∑ j=1 υj |wj∗ + τpj | ≤ d∑ j=1 υj |wj∗|
which further implies that
τ d∑ j=κ+1 υj |pj | ≤ κ∑ j=1 υj(|wj∗| − |wj∗ − τpj |) ≤ τ · κ∑ j=1 υj |pj |
where the second inequality follows from the triangle inequality. Then since each υj ∈ [1/2, 3/2], it follows that
d∑ j=κ+1 |pj | ≤ 3 κ∑ j=1 |pj |.
Note that this holds for all υ ∈ Υ simultaneously. Now let us denote p1:κ = (p1, . . . , pκ) ∈ Rκ and p(κ+1):d = (pκ+1, . . . , pd) ∈ Rd−κ, and similarly for other d-dimensional vectors. Then for all p ∈ DΥ ∩ Sd−1, by Cauchy-Schwartz inequality, we have
‖p(κ+1):d‖1 ≤ 3 · ‖p1:κ‖1 ≤ 3 √ κ · ‖p1:κ‖2.
Thus, for any z ∈ Rd and any p ∈ DΥ ∩ Sd−1, it follows that
〈z, p〉 = 〈z1:κ, p1:κ〉+ 〈z(κ+1):d, p(κ+1):d〉 ≤ ‖z1:κ‖2‖p1:κ‖2 + ‖p(κ+1):d‖1 · max
j∈{κ+1,...,d} |zj |
≤ ‖z1:κ‖2‖p1:κ‖2 + 3 √ κ‖p1:κ‖2 · max
j∈{κ+1,...,d} |zj |
≤ ‖z1:κ‖2 + 3 √ κ · max
j∈{κ+1,...,d} |zj |
where the last inequality follows from the fact that p ∈ Sd−1. Therefore, combine the above inequality with (43), and we obtain that
w(DΥ) ≤ E [ ‖z1:κ‖2 + 3 √ κ · max j∈{κ+1,...,d} |zj | ]
≤ √ κ+ 3 √ κ · E [ max j∈{κ+1,...,d} |zj | ] . (44)
where the second inequality follows from the fact that E[‖z1:κ‖2] ≤ √ E[‖z1:κ‖22] = √ κ. To bound the second term in (44), applying Lemma D.9, it follows from (44) that
w(DΥ) ≤ √ κ+ 3 √ 2κ ln(2(d− κ)). (45)
Therefore, combining (45) and (42), we obtain
λmin(Z;D(g, w∗)) ≥ √ n− 1− √ κ− 3 √ 2κ ln(2(d− κ))− h.
Therefore, choosing h = √ n− 1/2, as long as n satisfies that n ≥ C(κ ln d) for some constant C > 0, we have λmin(Z;D(g, w∗)) > 0 with probability at least 1 − (2d + 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, Id), then it holds that E [ maxi∈[d] |zi| ] ≤ √ 2 ln(2d).
Proof of Lemma D.9. Denote M = maxi∈[d] |zi|. For any λ > 0, by Jensen’s inequality, we have
eλ·E[M ] ≤ E [ eλM ] = E [ max i∈[d] eλ|z i| ] ≤ d∑ i=1 E [ eλ|z i| ] .
Note that E[eλ|zi|] ≤ 2·E[eλzi ]. Thus, by the expression of the Gaussian moment generating function, we further have
eλ·E[M ] ≤ 2 d∑ i=1 E [ eλz i ] = 2deλ 2/2,
from which it follows that
E[M ] ≤ ln(2d) λ + λ 2 .
Choosing λ = √ 2 ln(2d) yields the desired result.

heading: D.5 PROOF OF LEMMA 6.5
text: Instead of studying the convergence of the Riemannian gradient flow directly, it is more convenient to consider it in the ambient space RD. To do so, we define a Lagrange function L(x;λ) = R(x) + ∑n i=1 λ
i(fi(x) − yi) for λ ∈ Rn. Based on this Lagrangian, we can continuously extend ∂Φ(x)∇R(x) to the whole space RD. In specific, we can find a continuous function F : RD → RD such that F (·)|Γ = ∂Φ(·)∇R(·). Such an F can be implicitly constructed via the following lemma. Lemma D.10. The `2 norm has a unique minimizer among {∇xL(x;λ) | λ ∈ Rn} for any fixed x ∈ RD. Thus we can define F : RD → RD by F (x) = argming∈{∇xL(x;λ)|λ∈Rn} ‖g‖2. Moreover, it holds that 〈F (x),∇fi(x)〉 = 0 for all i ∈ [n].
Proof of Lemma D.10. Fix any x ∈ RD. Note that {∇xL(x;λ) | λ ∈ Rn} is the subspace spanned by {∇fi(x)}i∈[n] shifted by ∇R(x), thus there is unique minimizer of the `2 norm in this set. This implies that F (x) = argming∈{∇xL(x;λ)|λ∈Rn} ‖g‖2 is well-defined.
To show the second claim, denote h(λ) = ‖∇xL(x;λ)‖22/2, which is a quadratic function of λ ∈ Rn. Then we have
∇h(λ) = 〈∇R(x),∇f1(x)〉... 〈∇R(x),∇fn(x)〉 +  ∑n i=1 λ i〈∇f1(x),∇fi(x)〉 ...∑n
i=1 λ i〈∇fn(x),∇fi(x)〉
 = 〈∇R(x),∇f1(x)〉... 〈∇R(x),∇fn(x)〉 +K(x)λ. For any λ such that∇xL(x;λ) = F (x), we must have ∇h(λ) = 0 by the definition of F (x), which by the above implies
(K(x)λ)i = −〈∇R(x),∇fi(x)〉 for all i ∈ [n].
Therefore, we further have 〈F (x),∇fi(x)〉 = 〈∇R(x),∇fi(x)〉+ n∑ j=1 λj〈∇fi(x),∇fj(x)〉 = 〈∇R(x),∇fi(x)〉+ (K(x)λ)i = 0
for all i ∈ [n]. This finishes the proof.
Hence, with any initialization x0 ∈ Γ, the limiting flow (17) is equivalent to the following dynamics dxt dt = −F (xt). (46)
Thus Lemma 6.5 can be proved by showing that the above xt converges to x∗ as t→∞. We first present a series of auxiliary results in below. Lemma D.11 (Implications for F (x) = 0). Let F : RD → RD be as defined in Lemma D.10. For any x = ( u v ) ∈ RD such that F (x) = 0, it holds that for each j ∈ [d], either uj = 0 or vj = 0.
Proof. Since F (x) = 0, it holds for all j ∈ [d] that,
0 = ∂R
∂uj (x) + n∑ i=1 λ(x)i ∂fi ∂uj (x) = 2uj · [ n∑ i=1 (zji ) 2 + n∑ i=1 λ(x)izji ] ,
0 = ∂R
∂vj (x) + n∑ i=1 λ(x)i ∂fi ∂vj (x) = 2vj · [ n∑ i=1 (zji ) 2 − n∑ i=1 λ(x)izji ] .
If there exists some j ∈ [d] such that uj 6= 0 and vj 6= 0, then it follows from the above two identities that
n∑ i=1 (zji ) 2 = 0
which happens with probability 0 in both the Boolean and Gaussian case. Therefore, we must have uj = 0 or vj = 0 for all j ∈ [d].
Lemma D.12. Let F : RD → RD be as defined in Lemma D.10. Then F is continuous on RD.
Proof. Case I. We first consider the simpler case of any fixed x0 ∈ U = (R \ {0})D, assuming that K(x0) is full-rank. Lemma D.10 implies that for any λ ∈ Rn such that ∇xL(x0;λ) = F (x0), we have
K(x0)λ = (−〈∇R(x0),∇fi(x0)〉)i∈[n].
Thus such λ is unique and given by
λ(x0) = K(x0) −1(−〈∇R(x0),∇fi(x0)〉)i∈[n].
Since K(x) is continuous, there exists a sufficiently small δ > 0 such that for any x ∈ Bδ(x0), 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δ(x0), and so is F (x) = ∇R(x) + ∑n i=1 λ(x) i∇fi(x).
Case II. Next, we consider all general x ∈ RD. Here for simplicity, we reorder the coordinates as x = (u1, v1, u2, v2, . . . , ud, vd) with a slight abuse of notation. Without loss of generality, fix any x0 such that for some q ∈ [d], (ui0)2 + (vi0)2 > 0 for all i = 1, . . . , q and ui0 = vi0 = 0 for all i = q + 1, . . . , d. Then ∇R(x0) and {∇fi(x0)}i∈[n] only depend on {zji }i∈[n],j∈[q], and for all i ∈ [n], it holds that
∇R(x0)(2q+1):D = ∇fi(x0)(2q+1):D = 0. Note that if we replace {∇fi(x)}i∈[n] by any fixed and invertible linear transform of itself, it would not affect the definition of F (x). In specific, we can choose an invertible matrix Q ∈ Rn×n such that, for some q′ ∈ [q], (z̃1, . . . , z̃n) = (z1, . . . , zn)Q satisfies that {z̃1:qi }i∈[q′] is linearly independent and z̃1:qi = 0 for all i = q
′ + 1, . . . , n. We then consider (∇f̃1(x), . . . ,∇f̃n(x)) = (∇f1(x), . . . ,∇fn(x))Q and the corresponding F (x). For notational simplicity, we assume that Q can be chosen as the identity matrix, so that (z1, . . . , zn) itself satisfies the above property, and we repeat it here for clarity
{z1:qi }i∈[q′] is linearly independent and z̃ 1:q i = 0 for all i = q ′ + 1, . . . , n. (47) This further implies that
∇fi(x)1:(2q) = 0, for all i ∈ {q′ + 1, . . . , n} and x ∈ RD. (48)
In the sequel, we use λ for n-dimensional vectors and λ̄ for q′-dimensional vectors. Denote2
λ(x) ∈ argmin λ∈Rn ∥∥∥∥∥∇R(x) + n∑ i=1 λi∇fi(x) ∥∥∥∥∥ 2 ,
λ̄(x) ∈ argmin λ̄∈Rq′ ∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ̄i∇fi(x)1:(2q) ∥∥∥∥∥∥
2
.
Then due to (47) and (48), we have∥∥∥∥∥∥∇R(x0)1:(2q) + q′∑ i=1 λ̄(x0) i∇fi(x0)1:(2q) ∥∥∥∥∥∥ 2 = ∥∥∥∥∥∇R(x0) + n∑ i=1 λ(x)i∇fi(x0) ∥∥∥∥∥ 2 = ‖F (x0)‖2.
(49)
On the other hand, for any x ∈ RD, by (48), we have∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ̄(x)i∇fi(x)1:(2q) ∥∥∥∥∥∥
2
= min λ∈Rn ∥∥∥∥∥∇R(x)1:(2q) + n∑ i=1 λi∇fi(x)1:(2q) ∥∥∥∥∥
2
≤ ∥∥∥∥∥∇R(x)1:(2q) + n∑ i=1 λ(x)i∇fi(x)1:(2q) ∥∥∥∥∥
2
= ‖F (x)1:(2q)‖2
≤ ‖F (x)‖2 ≤ ∥∥∥∥∥∇R(x) + n∑ i=1 λ(x0) i∇fi(x) ∥∥∥∥∥ 2
(50)
where the first and third inequalities follow from the definition of F (x). Let x→ x0, by the continuity of∇R(x) and {∇fi(x)}i∈[n], we have
lim x→x0 ∥∥∥∥∥∇R(x) + n∑ i=1 λ(x0) i∇fi(x) ∥∥∥∥∥ 2 = ∥∥∥∥∥∇R(x0) + n∑ i=1 λ(x0) i∇fi(x0) ∥∥∥∥∥ 2
(51)
2We do not care about the specific choice of λ(x) or λ̄(x) when there are multiple candidates, and we only need their properties according to Lemma D.10, so they can be arbitrary. Also, the minimum of `2-norm of an affine space can always be attained so argmin exists.
Denote K̃(x) = (K̃ij(x))(i,j)∈[q′]2 = (〈∇fi(x)1:(2q),∇fi(x)1:(2q)〉)(i,j)∈[q′]2 . By applying the same argument as in Case I, since K̃(x0) is full-rank, it also holds that limx→x0 λ̄(x) = λ̄(x0), and thus
lim x→x0 ∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ̄(x)i∇fi(x)1:(2q) ∥∥∥∥∥∥
2
= ∥∥∥∥∥∥∇R(x0)1:(2q) + q′∑ i=1 λ̄(x0) i∇fi(x0)1:(2q) ∥∥∥∥∥∥ 2 .
(52)
Combing (49), (50), (51) and (52) yields
lim x→x0 ‖F (x)1:(2q)‖2 = lim x→x0 min λ∈Rn ∥∥∥∥∥∇R(x)1:(2q) + n∑ i=1 λi∇fi(x)1:(2q) ∥∥∥∥∥
2
= ‖F (x0)‖2. (53)
Moreover, since ‖F (x)(2q+1):D‖2 = √ ‖F (x)‖22 − ‖F (x)1:(2q)‖22, we also have
lim x→x0
‖F (x)(2q+1):D‖2 = 0. (54)
It then remains to show that limx→x0 F (x) 1:(2q) = F (x0) 1:(2q), which directly follows from limx→x0 λ(x) 1:q′ = λ(x0) 1:q′ = λ̄(x0).
Now, for any > 0, due to the convergence of λ̄(x) and that K̃(x0) 0, we can pick a sufficiently small δ1 such that for some constant α > 0 and all x ∈ Bδ1(x0), it holds that ‖λ̄(x)− λ̄(x0)‖2 ≤ /2 and∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ̄i∇fi(x)1:(2q) ∥∥∥∥∥∥ 2
2
≥ ∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ̄(x)i∇fi(x)1:(2q) ∥∥∥∥∥∥ 2
2
+ α‖λ̄− λ̄(x)‖22.
(55)
for all λ̄ ∈ Rp, where the inequality follows from the strong convexity. Meanwhile, due to (48), we have
lim x→x0 ∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ(x)i∇fi(x)1:(2q) ∥∥∥∥∥∥
2
= lim x→x0 ∥∥∥∥∥∇R(x)1:(2q) + n∑ i=1 λ(x)i∇fi(x)1:(2q) ∥∥∥∥∥
2
= ∥∥∥∥∥∥∇R(x0)1:(2q) + q′∑ i=1 λ̄(x0) i∇fi(x0)1:(2q) ∥∥∥∥∥∥ 2
= lim x→x0 ∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ̄(x)i∇fi(x)1:(2q) ∥∥∥∥∥∥
2
.
where the second equality follows from (53) and the second equality is due to (52). Therefore, we can pick a sufficiently small δ2 such that∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ(x)i∇fi(x)1:(2q) ∥∥∥∥∥∥
2
≤ ∥∥∥∥∥∥∇R(x)1:(2q) + q′∑ i=1 λ̄(x)i∇fi(x)1:(2q) ∥∥∥∥∥∥
2
+ α 2
4
(56)
for all x ∈ Bδ2(x0). Setting δ = min(δ1, δ2), it follows from (55) and (56) that
‖λ(x)1:q ′ − λ̄(x)‖2 ≤
2 , for all x ∈ Bδ(x0).
Recall that we already have ‖λ̄(x)− λ̄(x0)‖ ≤ /2, and thus
‖λ(x)1:q ′ − λ(x0)1:q ′ ‖2 = ‖λ(x)1:q ′ − λ̄(x0)‖2 ≤ ‖λ(x)1:q ′ − λ̄(x)‖2 + ‖λ̄(x)− λ̄(x0)‖2 ≤
for all x ∈ Bδ(x0). Therefore, we see that limx→x0 λ(x)1:q ′ = λ(x0) 1:q′ .
Finally, it follows from the triangle inequality that
‖F (x)− F (x0)‖2 ≤ ‖F (x)1:(2q) − F (x0)1:(2q)‖2 + ‖F (x)(2q+1):D‖2 + ‖F (x0)(2q+1):D‖2
≤ ∥∥∥∥∥∥ q′∑ i=1 λ(x)i∇fi(x)− λ(x0)i∇fi(x0) ∥∥∥∥∥∥ 2 + ‖∇R(x)−∇R(x0)‖2 + ‖F (x)(2q+1):D‖2
where, as x→ x0, the first term vanishes by the convergence of λ(x)1:q ′
and the continuity of each ∇fi(x), the second term converges to 0 by the continuity of ∇R(x) and the third term vanishes by (54). Therefore, we conclude that
lim x→x0 F (x) = F (x0),
that is, F is continuous.
Lemma D.13. For any initialization x0 ∈ Γ, the Riemmanian Gradient Flow (17) (or equivalently, (46)) is defined 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 6= ∞. Since R(x(t)) is monotone decreasing, thus R(x(t)) is upper bounded by R(x(0)) and therefore ‖∇R(x(t))‖ is also upper bounded. Since∥∥∥dx(t)dt ∥∥∥
2 ≤ ‖∇R(x(t))‖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., uj(T−) = 0 or vj(T−) = 0 for some j ∈ [d] by Lemma D.11. By the definition of the Riemannian gradient flow in (17), we have
d dt (uj(t)vj(t)) =
( vj(t)e > j uj(t)e > j ) ∂x(t) ∂t
= − ( vj(t)e > j uj(t)e > j ) F (x(t)).
By the expression of F (x(t)) = ∇R(x(t)) + ∑n i=1 λi(x(t))∇fi(x(t)), we then have
d dt (uj(t)vj(t)) = − [ n∑ i=1 z2i,j + n∑ i=1 λi(x(t))zi,j ] uj(t)vj(t)− [ n∑ i=1 z2i,j − n∑ i=1 λi(x(t))zi,j ] uj(t)vj(t)
= − ( n∑ i=1 z2i,j ) uj(t)vj(t).
Denote sj = ∑n i=1 z 2 i,j . It follows that |uj(t)vj(t)| = |uj(0)vj(0)|e−sjt for all t ∈ [0, T ). Taking the limit we have |uj(T−)vj(T−)| ≥ |uj(0)vj(0)|e−sjT > 0. Contradiction with T 6=∞!
Before showing that F satisfies the PL condition, we need the following two intermediate results. Given two points u and v in Rd, we say u weakly dominate v (written as u ≤ v) if and only if ui ≤ vi, for all i ∈ [d]. Given two subsets A and B of RD, 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 RD and P = {u ∈ RD | ui ≥ 0,∀i ∈ [D]}. Let u? be an arbitrary point in P and Q = P ∩ (u? + S). Then there exists a radius r > 0, such that B1r (0) ∩Q weakly dominates Q, where B1r (0) is the `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, minx∈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 | ui = 0} (which can be seen as a subset of RD−1), and
let ri be the corresponding `1 radius. Set r = maxi∈[D] ri, and we show that choosing the radius to be r suffices.
For any v ∈ Q, we take a random direction in S, denoted by ω. If ω ≥ 0 or ω ≤ 0, we denote by y the first intersection (i.e., choosing the smallest λ) between the line {v − λ|ω|}λ≥0 and the boundary of U , i.e., ∪Di=1{z ∈ RD | zi = 0}. Clearly y ≤ v. By the induction hypothesis, there exists a u ∈ B1r (0) ∩Q such that u ≤ y. Thus u ≤ v and meets our requirement. If ω has different signs across its coordinates, we take y1, y2 to be the first 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 u1, u2 ∈ B1r (0) ∩Q such that u1 ≤ y1 and u2 ≤ y2. Since v lies in the line connecting u1 and u2, there exists some h ∈ [0, 1] such that v = (1 − h)u1 + hu2. It then follows that (1 − h)u1 + hu2 ≤ (1 − h)y1 + hy2 = v. Now since Q is convex, we have (1− h)u1 + hu2 ∈ Q, and by the triangle inequality it also holds that ‖(1− h)u1 + hu2‖1 ≤ r, so (1− h)u1 + hu2 ∈ B1r (0) ∩Q. Therefore, we conclude that B1r (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 RD and P = {u ∈ RD | ui ≥ 0,∀i ∈ [D]}. Let u? be an arbitrary point in P and Q = P ∩ (u? + S). Then there exists a constant c ∈ (0, 1] such that for any sufficiently small radius r > 0, c · Q weakly dominates P ∩ (u? + S +B2r (0)), where B2r (0) is the `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?| ≤ |u?| − r, and it suffices to pick r = |u?| and c = 0.5. Suppose the proposition holds for D − 1, below we show it holds for D. For each i ∈ [D], we first consider the intersection between P ∩ (u? + S + B2r (0)) and Hi := {u ∈ RD | ui = 0}. Let ui be an arbitrary point in P ∩ (u? + S) ∩Hi, then P ∩ (u? + S) ∩Hi = P ∩ (ui + S) ∩Hi = P ∩ (ui + S ∩ Hi). Furthermore, there exists {αi}i∈[D] which only depends on S and satisfies P ∩ (u∗+S+B2r (0))∩Hi ⊂ P ∩ (ui+S∩Hi+B2αir(0)∩Hi). Applying the induction hypothesis to P ∩ (ui +S ∩Hi +B2αir(0)∩Hi), we know there exists a c > 0 such that for sufficiently small r, c(P ∩ (u?+S)∩Hi) = c(P ∩ (ui+S∩Hi)) weakly dominates P ∩ (ui+S∩Hi+B2αir(0)∩Hi).
For any point v in Q and any z ∈ B2r (0), we take a random direction in S, denoted by ω. If ω ≥ 0 or ω ≤ 0, we denote by y the first intersection between {v + z − λ|ω|}λ≥0 and the boundary of U . Clearly y ≤ v. Since y ∈ P ∩ (u? + S + B2r (0)) ∩Hi ⊂ P ∩ (ui + S ∩Hi + B2αir(0) ∩Hi), by the induction hypothesis, there exists a u ∈ c(P ∩ (u? + S) ∩Hi) such that u ≤ y. Thus z ≤ v + z and z ∈ c(P ∩ (u? + S)) = cQ. If ω has different signs across its coordinates, we take y1, y2 to be the first 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 u1, u2 ∈ c · Q such that u1 ≤ y1 and u2 ≤ y2. Since v + z lies in the line connecting u1 and u2, there exists some h ∈ [0, 1] such that v + z = (1 − h)y1 + hy2. It then follows that (1 − h)u1 + hu2 ≤ (1 − h)y1 + hy2 = v + z. Since Q is convex, we have (1−h)u1 +hu2 ∈ cQ. Therefore, we conclude that cQ∩Q weakly dominates P ∩ (u?+S+B2r (0)) for all sufficiently 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 ′ of x∗ and a constant c > 0, such that ‖F (x)‖22 ≥ c ·max(R(x)−R(x∗), 0) for all x ∈ U ′ ∩ Γ. Note this requirement is only non-trivial when ‖F (x∗)‖2 = 0 since F is continuous.
Proof of Lemma D.16. It suffices to show the PL condition for {x | F (x) = 0}. We need to show for any x∗ satisfying F (x∗) = 0, there exist some > 0 and C > 0, such that for all x ∈ Γ∩B2 (x∗) with R(x) > R(x∗), it holds that ‖F (x)‖22 ≥ C(R(x)−R(x∗)).
Case I. We first prove the case where x = ( u v ) itself is a canonical parametrization of w = u 2 − v 2, i.e., ujvj = 0 for all j ∈ [d]. Since x∗ satisfies∇F (x∗) = 0, by Lemma D.11, we have x∗ = ψ(w∗) where w∗ = u 2∗ − v 2∗ . In this case, we can rewrite both R and F as functions of w ∈ Rd. In detail, we define R′(w) = R(ψ(w)) and F ′(w) = F (ψ(w)) for all w ∈ Rd. For any w in a sufficiently small neighbourhood of w∗, it holds that sign(wj) = 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 wj ≥ 0, for all j = q + 1 . . . , d. We temporarily reorder the coordinates as x = (u1, v1, u2, v2, . . . , ud, vd)>. Recall that Z = [z1, . . . , zn] > is a n-by-d matrix, and we have
‖F ′(w)‖22 = min λ∈Rn
〈 (a− sign(w) Z>λ) 2, |w| 〉 ,
where a = 1n ∑n i=1 z 2 i ∈ Rd. Since F (x∗) = 0, there must exist λ∗ ∈ Rn, such that the first 2q
coordinates of∇R(x∗) + ∑n i=1 λ i ∗∇fi(x∗) are equal to 0. As argued in the proof of Lemma D.12, we can assume the first q′ rows of Z are linear independent on the first q coordinates for some q′ ∈ [q].
In other words, Z can be written as [ ZA ZB 0 ZD ] where ZA ∈ Rq ′×q . We further denote λ1 := λ1:q ′ , λ2 := λ (q′+1):n, w1 := w1:q and w2 := w(q+1):d for convenience, then we have
‖F ′(w)‖22 = min λ∈Rn
〈 (a1 + sign(w1) Z>Aλ1) 2, |w1| 〉 + 〈 (a2 + Z > Bλ1 + Z > Dλ2) 2, w2 〉 . (57)
Since every w in Γ is a global minimizer, R′(w) = R′(w) + ∑n i=1 λ i ∗(z > i w − yi) := g>w, where g = sign(w) a+ Z>λ∗. Similarly we define g1 := g1:q and g2 := g(q+1):d. It holds that g1 = 0 and we assume ZDg2 = 0 without loss of generality, because this can always be done by picking suitable λi∗ for i = q ′ + 1, . . . , n. We denote λ1 − λ∗,1 by ∆λ1, then since 0 = g1 = sign(w1) a1 + Z>Aλ∗,1, we further have〈 (a1 + sign(w1) Z>Aλ1) 2, |w1| 〉 = 〈 (a1 + sign(w1) Z>Aλ∗,1 + sign(w1) Z>A∆λ1) 2, |w1|
〉 = 〈 (sign(w1) Z>A∆λ1) 2, |w1| 〉 .
On the other hand, we have g2 = sign(w2) a2 +Z>Bλ∗,1 +Z>Dλ∗,2 = a2 +Z>Bλ∗,1 +Z>Dλ∗,2 by the assumption that each coordinate of w2 is non-negative. Combining this with the above identity, we can rewrite Equation (57) as:
‖F ′(w)‖22 = min λ∈RD
〈 (Z>A∆λ1) 2, |w1| 〉 + 〈 (g2 + Z > B∆λ1 + Z > Dλ2) 2, w2 〉 . (58)
Now suppose R′(w)−R′(w∗) = g>2 w2 = δ for some sufficiently small δ (which can be controlled by ). We will proceed in the following two cases separately.
• Case I.1: ‖∆λ1‖2 = Ω( √ δ). Since ZA has full row rank and every coordinate of w1 is
non-zero, the first term of Equation (58) is Ω(δ) = Ω(R′(w)−R′(w∗)).
• Case I.2: ‖∆λ1‖2 = O( √ δ). Let u = g2+Z>B∆λ1+Z > Dλ2, then we have u ∈ S+B2c√δ(0)
for some constant c > 0, where S = {g2 + Z>Dλ2 | λ2 ∈ Rn−q ′}. By Lemma D.14, there exists some constant c0 ≥ 1, such that 1c0 · S weakly dominates S + B 2 c √ δ (0). Thus we have ‖F ′(w)‖22 ≥ infu∈S+B2 c √ δ (0) 〈 u 2, w2 〉 ≥ infu∈ 1c0 ·S 〈 s 2, w2 〉 , where the last step is because each coordinate of w2 is non-negative.
Let A be the orthogonal complement of span(ZD, g2), we know w2 ∈ δ‖g2‖22 g2 +A, since ZDw2 = ZDw∗,2 = 0 and g>2 w2 = δ. Therefore,
inf w:R′(w)−R′(w∗)=δ>0 ‖F ′(w)‖22 R′(w)−R′(w∗) ≥ inf w2:R′(w)−R′(w∗)=δ>0 inf u∈ 1c0 ·S
〈 u 2,
w2 δ 〉 ≥ 1 c20
inf w2∈ δ‖g2‖22 g2+A,w2≥0,u∈S
〈 u 2, w2 〉 . (59)
Note 〈 u 2, w2 〉 is a monotone non-decreasing function in the first joint orthant, i.e., {(u,w2) ∈ Rd × Rd−q ′ | u ≥ 0, w2 ≥ 0}, thus by Lemma D.15 the infinimum can be achieved by some finite (u,w2) in the joint first orthant. Applying the same argument to each other orthant of u ∈ Rd, we conclude that the right-hand-side of (59) can be achieved. On the other hand, we have u>w2 = δ > 0 for all w2 ∈ δ‖g2‖22 g2 + A and u ∈ S, by ZDg2 = 0 and the definition of A. This implies there exists at least one i ∈ [d − q′] such that wi2u i > 0, which further implies 〈 u 2, w2 〉 > 0. Therefore, we conclude that ‖F ′(w)‖22 = Ω(R′(w)−R′(w0)).
Case II. Next, for any general x = ( u v ) , we define w = u 2 − v 2 and m = min{u 2, v 2}, where min is taken coordinate-wise. Then we can rewrite ‖F (x)‖22 as
‖F (x)‖22 = min λ∈Rn ∥∥∥∥([aa ] + [ Z −Z ] λ ) [ u v ]∥∥∥∥2 2
= min λ∈Rn ∥∥∥∥∥ ([ a a ] + [ Z −Z ] λ ) 2 [ u 2 v 2 ]∥∥∥∥∥ 1
= min λ∈Rn ∥∥∥∥∥ ([ a a ] + [ Z −Z ] λ ) 2 ( ψ(w) 2 + [ m m ])∥∥∥∥∥ 1
≥ min λ∈Rn ∥∥∥∥∥ ([ a a ] + [ Z −Z ] λ ) 2 ψ(w) 2 ∥∥∥∥∥ 1 + min λ∈Rn ∥∥∥∥∥ ([ a a ] + [ Z −Z ] λ ) 2 [ m m ]∥∥∥∥∥ 1
= min λ∈Rn ∥∥∥∥([aa ] + [ Z −Z ] λ ) ψ(w) ∥∥∥∥2 2 + min λ∈Rn ∥∥∥∥([aa ] + [ Z −Z ] λ ) [√ m√ m ]∥∥∥∥2 2 .
Then applying the result for the previous case yields the following for some constant C ∈ (0, 1):
‖F (x)‖22 ≥ C(R(ψ(w))−R(ψ(w∗)) + min λ∈Rn ∥∥∥∥([aa ] + [ Z −Z ] λ ) [√ m√ m ]∥∥∥∥2 2
= C(R(ψ(w))−R(x∗) + 2 〈 a 2,m 〉 ≥ C(R(ψ(w))−R(x∗) + 2 min
i∈[d] ai 〈a,m〉
= C(R(ψ(w))−R(x∗) + min i∈[d] ai(R(x)−R(ψ(w))) ≥ min { C,min
i∈[d] ai
} (R(x)−R(x∗)),
where the first 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 flow defined in (17) has finite length, i.e., ∫∞ t=0 ‖dxdt ‖2dt < ∞ for any x0 ∈ Γ. Moreover, xt converges to some x∞ when t→∞ with F (x∞) = 0.
Proof of Lemma D.17. Note that along the Riemannian gradient flow, R(xt) is non-increasing, thus ‖xt‖2 is bounded over time and {xt}t≥0 has at least one limit point, which we will call x∗. Therefore, R(x∗) is a limit point of R(xt), and again since R(xt) is non-increasing, it follows that R(xt) ≥ R(x∗) and limt→∞R(xt) = R(x∗). Below we will show limt→∞ xt = x∗.
Note that dR(xt)dt = 〈 ∇R(xt), dxtdt 〉 = −〈∇R(xt), F (xt)〉 = −‖F (xt)‖22 where the last inequality applies Lemma D.10. By Lemma D.16, there exists a neighbourhood of x∗, U ′, in which PL condition holds of F . Since x∗ is a limit point, there exists a time T0, such that xT0 ∈ U . Let T1 =
inft≥T0{x(t) /∈ U ′} (which is equal to∞ if x(t) ∈ U ′ for all t ≥ T0). Since xt is continuous in t and U is open, we know T1 > T0 and for all t ∈ [T0, T1), we have ‖F (xt)‖2 ≥ √ c(R(xt)−R(x∗))1/2.
Thus it holds that for t ∈ [T0, T1),
d(R(xt)−R(x∗)) dt ≤ − √ c(R(xt)−R(x∗))1/2 ‖F (xt)‖2 ,
that is,
d(R(xt)−R(x∗))1/2 dt ≤ − √ c 2 ‖F (xt)‖2 .
Therefore, we have ∫ T1 t=T0 ‖F (xt)‖2 dt ≤ 1 2 √ c (R(xT0)−R(x∗)). (60)
Thus if we pick T0 such that R(xT0)−R(x∗) is sufficiently small, R(T1) will remain in U , which implies that T1 cannot be finite and has to be∞. Therefore, Equation (60) shows that the trajectory of xt is of finite length, so x∞ = limt→∞ xt 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) φ(xt) is locally increasing in a sufficiently neighborhood of x∗ and (2) limx→x∗ φ(x) = −∞. Lemma 6.5. Let {xt}t≥0 ⊆ RD be generated by the flow defined in (17) with any initialization x0 ∈ Γ. Then x∞ = limt→∞ xt exists. Moreover, x∞ = x∗ is the optimal solution of (18).
Proof of Lemma 6.5. We will prove by contradiction. Suppose x∞ = ( u∞ v∞ ) = limt→∞ xt is not the optimal solution to (18). Denote wt = u 2t − v 2t , then w∞ = limt→∞ wt is not the optimal solution to (36). Thus we have R(wt) > R(w∗). Without loss of generality, suppose there is some q ∈ [d] such that (ui∞)2 + (vi∞)2 > 0 for all i = 1, . . . , q and ui∞ = vi∞ = 0 for all i = q+ 1, . . . , d. Again, as argued in the proof of Lemma D.12, we can assume that, for some q′ ∈ [q],
{z1:qi }i∈[q′] is linearly independent and z 1:q i = 0 for all i = q ′ + 1, . . . , n. (61)
Since both w∞ and w∗ satisfy the constraint that Zw∞ = Zw∗ = Y , we further have
0 = 〈zi, w∞〉 = 〈zi, w∗〉 = 〈z(q+1):di , w (q+1):d ∗ 〉, for all i = q′ + 1, . . . , n. (62)
Consider a potential function ϕ : U → R defined as
ϕ(x) = ϕ(u, v) = d∑ j=q+1 wj∗ [ ln(uj)21{wj∗ > 0} − ln(vj)21{wj∗ < 0} ] .
Clearly limt→∞ ϕ(xt) = −∞ if limt→∞ xt = x∞. Below we will show contradiction if x∞ is suboptimal. Consider the dynamics of ϕ(x) along the Riemannian gradient flow:
dϕ dt (xt) =
〈 ∇ϕ(xt),
dxt dt
〉 = −〈∇ϕ(xt), F (xt)〉 (63)
where F is defined previously in Lemma D.10. Recall the definition of F , and we have
〈∇ϕ(xt), F (xt)〉 = 〈 ∇ϕ(xt),∇R(xt) +
q′∑ i=1 λ(xt) i∇fi(xt) 〉 ︸ ︷︷ ︸
I1
+ 〈 ∇ϕ(xt),
n∑ i=q′+1 λ(xt) i∇fi(xt) 〉 ︸ ︷︷ ︸
I2
.
(64)
To show 〈∇ϕ(xt), F (xt)〉 < 0, we analyze I1 and I2 separately. By the definition of ϕ(x), we have
∇ϕ(x) = d∑
j=q+1
2wj∗
[ 1{wj∗ > 0}
uj · ej − 1{wj∗ < 0} vj · eD+j
]
where ej is the j-th canonical base of Rd. Recall that∇fi(x) = ( zi u −zi v ) , and we further have
I2 = n∑
i=q′+1
λ(xt) i d∑ j=q+1 2wj∗
[ 1{wj∗ > 0}
uj 〈ej , zi u〉+ 1{wj∗ < 0} vj 〈ej , zi v〉
]
= n∑ i=q′+1 λ(xt) i d∑ j=q+1 2wj∗
[ 1{wj∗ > 0}
uj zji u
j + 1{wj∗ < 0}
vj zji v j
]
= n∑ i=q′+1 λ(xt) i d∑ j=q+1 wj∗z j i = n∑ i=q′+1 λ(xt) i〈z(q+1):di , w (q+1):d ∗ 〉 = 0 (65)
where the last equality follows from (62).
Next, we show that I1 < 0 by utilizing the fact that w∗ − w∞ is a descent direction of R′(w). For w ∈ Rd, define f̃i(w) = z>i w and
R̃(w) = R(w) + q′∑ i=1 λ(x∞) i(f̃i(w)− yi).
Clearly, for any w ∈ RD satisfying Zw = Y , it holds that f̃i(w) − yi = 0 for each i ∈ [n], and thus R(w) = R̃(w). In particular, we have R̃(w∞) = R(w∞) > R(w∗) = R̃(w∗). Since R̃(w) is a convex function, it follows that R̃(w∞ + s(w∗ − w∞)) < R̃(w∞) for all sufficiently small s > 0, which implies dR̃dt (w∞ + s(w∗ − w∞))|s=0 < −2c < 0 for some constant c > 0. Note that, for small enough s > 0, we have
R(w∞ + s(w∗ − w∞)) = d∑ j=1 ( n∑ i=1 (zji ) 2 ) |wj∞ + s(wj∗ − wj∞)|
= q∑ j=1 ( n∑ i=1 (zji ) 2 ) sign(wj∞)(w j ∞ + s(w j ∗ − wj∞))
+ d∑ j=q+1 ( n∑ i=1 (zji ) 2 ) s|wj∗|.
Therefore, we can compute the derivative with respect to s at s = 0 as
−2c > dR̃ dt (w∞ + s(w∗ − w∞)) ∣∣∣∣ s=0 = q∑ j=1 ( n∑ i=1 (zji ) 2 ) sign(wj∞)(w j ∗ − wj∞) + d∑ j=q+1 ( n∑ i=1 (zji ) 2 ) |wj∗|
+ q′∑ i=1 λ(x∞) iz>i (w∗ − wj∞)
= q∑ j=1 ( n∑ i=1 (zji ) 2 ) sign(wj∞)(w j ∗ − w∞) + d∑ j=q+1 ( n∑ i=1 (zji ) 2 ) |wj∗|
+ q∑ j=1 (wj∗ − wj∞) q′∑ i=1 λ(x∞) izji + d∑ j=q+1 wj∗ q′∑ i=1 λ(x∞) izji
(66)
where the second equality follows from the fact that w(q+1):d∞ = 0. Since xt converges to x∞, we must have F (x∞) = 0, which implies that for each j ∈ {1, . . . , q},
0 = ∂R
∂uj (x∞) + q′∑ i=1 λ(x∞) i ∂fi ∂uj (x∞) = 2u j ∞ ·  n∑ i=1 (zji ) 2 + q′∑ i=1 λ(x∞) izji  , 0 = ∂R
∂vj (x∞) + q′∑ i=1 λ(x∞) i ∂fi ∂vj (x∞) = 2v j ∞ ·  n∑ i=1 (zji ) 2 − q′∑ i=1 λ(x∞) izji  . Combining the above two equalities yields
n∑ i=1 (zji ) 2 = − sign(wj∞) q′∑ i=1 λ(x∞) izji , for all j ∈ [q].
Apply the above identity together with (66), and we obtain
−2c > q∑ j=1 − sign(wj∞)2(wj∗ − w∞) q′∑ i=1 λ(x∞) izji + d∑ j=q+1 ( n∑ i=1 (zji ) 2 ) |wj∗|
+ q∑ j=1 (wj∗ − wj∞) q′∑ i=1 λ(x∞) izji + d∑ j=q+1 wj∗ q′∑ i=1 λ(x∞) izji
= d∑ j=q+1 ( n∑ i=1 (zji ) 2 ) |wj∗|+ d∑ j=q+1 wj∗ q′∑ i=1 λ(x∞) izji (67)
On the other hand, by directly evaluating∇R(xt) and each∇fi(xt), we can compute I1 as
I1 = 2 d∑
j=q+1
wj∗1{wj∗ > 0} ujt  n∑ i=1 (zji ) 2ujt + q′∑ i=1 λ(xt) izji u j t  − 2
d∑ j=q+1 wj∗1{wj∗ < 0} vjt  n∑ i=1 (zji ) 2vjt − q′∑ i=1 λ(xt) izji v j t  = 2
d∑ j=q+1 ( n∑ i=1 (zji ) 2 ) |wj∗|+ 2 d∑ j=q+1 wj∗ q′∑ i=1 λ(xt) izji
= 2 d∑ j=q+1 ( n∑ i=1 (zji ) 2 ) |wj∗|+ 2 d∑ j=q+1 wj∗ q′∑ i=1 λ(x∞) izji
+ 2 d∑ j=q+1 wj∗ q′∑ i=1 ( λ(xt) i − λ(x∞)i ) zji .
We already know that λ(x)1:q ′
is continuous at x∞ by the proof of Lemma D.12, so the third term converges to 0 as xt tends to x∞. Now, applying (67), we immediately see that there exists some δ > 0 such that I1 < −c for xt ∈ Bδ(x∞). As we have shown in the above that I2 = 0, it then follows from (63) and (64) that
dϕ dt (xt) > c, for all xt ∈ Bδ(x∞). (68)
Since limt→∞ xt = x∞, there exists some T > 0 such that xt ∈ Bδ(x∞) for all t > T . By the proof ofLemma D.13, we know that ϕ(xT ) > −∞, then it follows from (68) that
lim t→∞ ϕ(xt) = ϕ(xT ) + ∫ ∞ T dϕ(xt) dt dt > ϕ(xT ) + ∫ ∞ T cdt =∞
which is a contradiction. This finishes the proof.

heading: D.6 PROOF OF THEOREM 6.7
text: Here we present the lower bound on the sample complexity of GD in the kernel regime.
Theorem 6.7. Assume z1, . . . , zn i.i.d.∼ N (0, Id) and yi = w>∗ zi, for all i ∈ [n]. Define the loss with linearized model as L(x) = ∑n i=1(fi(x0) + 〈∇fi(x0), x− x0〉 − yi)2, where x = ( u v ) and
x0 = ( u0 v0 ) = α ( 1 1 ) . Then for any ground truth w∗, any learning rate schedule {ηt}t≥1, and any fixed number of steps T , the expected `2 loss of xT is at least (1− nd ) ‖w∗‖ 2 2, where xT is the T -th iterate of GD on L, i.e., xt+1 = xt − ηt∇L(xt), for all t ≥ 0.
Proof of Theorem 6.7. We first simplify the loss function by substituting x′ = x−x0, so correspondingly x′0 = 0 and we consider L
′(x′) := L(x) = (〈∇fi(x0), x′〉 − yi)2. We can think as if GD is performed on L′(x′). For simplicity, we still use the x and L(x) notation in below.
In order to show test loss lower bound against a single fixed 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 xT with ground truth being w∗ is the same as the case where the ground truth is uniformly randomly sampled from all vectors of `2-norm ‖w∗‖2. Thus the test loss of xT is
Ez [ (〈∇fz(x0), xT 〉 − 〈z, w∗〉)2 ] = Ez [ (〈z, w∗ − (uT − vT )〉)2 ] = ‖w∗ − (uT − vT )‖22 . (69)
Note xT ∈ span{∇fx(x0)}, which is at most an n-dimensional space spanned by the gradients of model output at x0, so is uT − vT . We denote the corresponding space for uT − vT by S, so dim(S) ≤ n and it holds that ‖w∗ − (uT − vT )‖22 ≥ ‖(ID − PS)w∗‖ 2 2, where PS is projection matrix onto space S.
The expected test loss is lower bounded by Ew∗ [ Ezi [ ‖w∗ − (uT − vT )‖22 ]] = Ezi [ Ew∗ [ ‖w∗ − (uT − vT )‖22 ]] ≥ min
zi Ew∗
[ ‖(ID − PS)w∗‖22 ] ≥ (
1− n d
) ‖w∗‖22 .

Review summaries:
The authors propose an SDE approximation to study the implicit bias of SGD with infinitesimal learning rate. They show that given any noise covariance structure, after $\eta^{-2}$ steps, the SGD converges to an SGD on a certain manifold of local minimizer as the learning rate $\eta\rightarrow 0$. This recovers and strengthens some existing results in the literature. In particular, they show a sample complexity gap between label noise SGD and GD in the kernel regime for an overparametrized linear model, justifying the generalization benefit of SGD.

This paper proposes a new mathematical framework for explaining implicit regularization of SGD after attaining zero training loss. Technical novelty includes a different time scaling of the SGD that leads to the Katzenberger process, and then applying weak convergence of the Katzenberger process to obtain an elegant characterization of the SGD dynamics on zero loss manifold. This new technical device leads to a new proof of the sample complexity bound on the overparametrized linear model with label noise regularization.

The authors of the paper under review carry a general analysis of the effect of the noise near the manifold of global minimum $\Gamma = [x, L(x) = 0]$. More precisely they show that in the limit of infinitesimal step-sizes, a time rescaled effective dynamics follows a diffusion in the manifold $\Gamma$ that mixes effects of the noise and the gradient flow. Leveraging this, they characterise explicitly the limiting point recovered by SGD on a sparse problem trained with label noise on a overparametrised linear network of depth 2.