What Happens after SGD Reaches Zero Loss? --A Mathematical Framework
Authors: Zhiyuan Li
Tianhao Wang
Sanjeev Arora
Venue: ICLR 2022 Spotlight (2022)
Keywords: SGD
implicit bias
generalization
deep learning
implicit regularization
manifold

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 $\eta$, 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, $\text{tr}[\nabla^2 L]$. The current paper gives a general framework for such analysis by adapting ideas from Katzenberger (1991). It allows in principle a complete characterization for the regularization effect of SGD around such manifold---i.e., the "implicit bias"---using a stochastic differential equation (SDE) describing the limiting dynamics of the parameters, which is determined jointly by the loss function and the noise covariance. This yields some new results: (1) a *global* analysis of the implicit bias valid for $\eta^{-2}$ steps, in contrast to the local analysis of Blanc et al. (2020) that is only valid for $\eta^{-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(\kappa\ln d)$ samples for learning an $\kappa$-sparse overparametrized linear model in $\mathbb{R}^d$ (Woodworth et al., 2020), while GD initialized in the kernel regime requires $\Omega(d)$ samples. This upper bound is minimax optimal and improves the previous $\widetilde{O}(\kappa^2)$ upper bound (HaoChen et al., 2020).

Published as a conference paper at ICLR 2022
WHAT HAPPENS AFTER SGD REACHES ZERO Loss? -A MATHEMATICAL FRAMEWORK
Zhiyuan Li
Tianhao Wang
Department of Computer Science Princeton University
Department Statistics and Data Science Yale University
zhiyuanl 1 www.ppriceteoo....mmmmmm..
tianhao. wang@yale.cc..mm edu
Sanjeev Arora
Department of Computer Science
Princeton University aroraêcs. weraadccspinnneoo..mm edu
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 7), 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 regularizer-likk like term, the sharpness of loss, tr V 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.c. the "implicit 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 results: (1) global analysis of the implicit bias valid for r) steps, in contrast to the local analysis of Blanc et al. (2020) that is only valid for 1- steps and (2) allowing arbitrary noise covariance. As application,
show with arbitrary large initialization, label noise SGD can always escape the kernel regime and only requires (K d) samples for learning an k-sparse overparametrized linear model in IR (Woodworth et al., 2020), while GD initialized in the kernel regime requires (d) samples. This upper bound is minimax optimal and improves the previous (kr) upper bound (HaoChen et al., 2020).
INTRODUCTION
The implicit bias underlies the generalization ability of machine learning models trained by stochastic gradient descent (SGD). But still remains a mystery to mathematically characterize such bias. We study SGD in the following formulation
ch (b: + 1) = (k) 77 (VL(E,, + VE (z,,
(1)
where 17 is the learning rate (LR), L RD R is the training loss and (x) [o (z), (2(
E Dx is a deterministic noise function. Here En is sampled uniformly from (1,2, .,, and it satisfies (ee, (x)) = O. Var E Rd and k.
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 flatter attraction basin of the loss landscape which potentially benefits generalization (Li et al., 2019c: 9c; 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/n)) steps (Shi et al., 2020). In other words, there should be a mechanism to reach better generalization while staying within single basin. For deterministic GD similar mechanisms

---

Published as a conference paper at ICLR 2022
X2 då 7
D2 + A) dx1
-VL(,, + A) AX)) (V)) [4.A]
+n a (X))
dt
B, 0² (XX,)A4 0 [T["')
X. +
A(t) ((,
X.
Normal Space: Fast stochastic
X,
Tangent Space: slow & deterministic
X.
r
r
r
(a) Taylor Expansion of V
(b) Normal Space Dynamies
(c) Tangent Space Dynamies
Figure I: Illustration for limiting flow in R2 i' an ID manifold of minimizers of loss L.
have been demonstrated in simple cases (Soudry et al.. 2018; Lyu & Li, 2019) and referred 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 I' (see Figure 2). The effect is shown to be an implicit deterministic drift in direction corresponding to lowering regularizer-like term along the manifold. They showed SGD with label noise locally decreases the sharpness of loss, tr| t=[V2L) L]. by ((17².)) in 17 1--16 steps. However, such an analysis is actually local, since the natural time scale of analysis should be 1j not rj 1.6
The contribution of the current paper is more general and global analysis of this type. We introduce a more powerful framework inspired by the classic paper (Katzenberger, 1991).
1.1 INTUITIVE EXPLANATION OF REGULARIZATION EFFECT DUE TO SGD
We start with an intuitive description of the implicit regularization effect described in Blanc et al. (2020). For simplification,
show it for the canonical SDE approximation (See Section B.1 for more details) of SGD (1) (Li et al., 2017: Cheng et al., 2020). Here W(I) is the standard E-dimensional Brownian motion. The only property about label noise SGD will use is that the noise covariance oo (x) = L()( for every x in the manifold I' (See derivation in Section 5).
(t,,()) = (XX, + 7 o (X)(())
Suppose X,,(O) (0) already close to some local minimizer point X, E I. The goal is to show (t) will move in the tangent space and steadily decrease L). At first glance, this seems impossible as the gradient V vanishes around I. and the noise has zero mean, implying SGD should be like random walk instead of a deterministic drift. The key observation Blane et al. (2020) is that the local dynamies of (t) is completely different tangent space and normal space the fast random walk in normal space causes X, (1) to move slowly (with velocity (()))) but deterministically in certain direction. To explain this, letting A(t) = (t) - X.. Taylor expansion of (2) gives + (X. (X.)) (t), meaning A is behaving like an Ornstein-Uhlenbedd
dd(t) = L (OU) process locally in the normal space. Its mixing time is O ') and the stationary distribution is the standard multivariate gaussian in the normal space scaled by v1j (see Figure lb), because noise covariance aa T = L. Though this OU process itself doesn't 't form any regularization, it activates the second order Taylor expansion of VL(X, + At)), i.e., 0 (VL) (t)), creating a (1))²) velocity the tangent space. Since there no push back force in the tangent space, the small velocity accumulates over time, and in longer time scale of 11) the time average of the stochastic velocity roughly the same as the expected velocity when A sampled from its stationary
n? distribution. This simplifies the expression of the velocity in tangent space to z V tr(221) (see
Figure Ic), where VT means the gradient is only taken in the tangent space.
However, the above approach only gives local analysis for O time, where the total movement due to implicit regularization is 0(1;2-16) = 0.4 and thus is negligible when 7/ - 0. In order to get non-trivial limiting dynamics when - 0, global analysis for 2(n steps is necessary and it cannot be done by Taylor expansion with a single reference point. Recent work by Damian et al. (2021) glues analyses of multiple local phases into global guarantee that SGD finds (c, stationary point for the regularized loss, but still doesn't 't show convergence for trajectory when 17 + 0 and cannot deal with general noise types, e.g., noise lying in the tangent space of the manifold. The
2

---

Published as a conference paper at ICLR 2022
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.
1.2 OUR APPROACH: SEPARATING THE SLOW FROM THE FAST
In this work, we tackle this problem via a different angle. First, since the anticipated limiting dynamics of speed ((1²). we change the time scaling to accelcrate (2) by rj times, which yields
(1)
+ (X)()))
The key idea here is that we only need to track the slow dynamic, equivalently, some projection of X onto the manifold r, ( Here + TR - r is some function to be specified and hopefully we can simplify the dynamics (3) via choosing suitable 4. To track the dynamics of + (X,,)) we apply Ito's lemma (a. stochastic chain rule, see Lemma A.9) to Equation (3), which yields
(x, = O X (XX + 8 (X())
2 1===1 (X)((
Jyjdt.
Note the first term - (x(,,, (V,,,, is going to diverge to do when 1) - 0, so a natural choice is indeed the directional derivative of & at X towards - -V, killing the first term becomes equivalent to making t invariant under Gradient Flow (GF) of (X Thus it suffices take & (X) to be the limit of GF starting at X. (Formally defined in Section 3; see Lemma C.2 for a proof of 0 (X VX))) = 0.)
for & is to kill the first term. Further note
Also intuitively X,, will be infinitely close to I, i.e., d( X,, (t), I) + for any > a 7 + 0, we have & (Xx,) x Xg- Thus we can rewrite the above equation as
+ Σ ا-= 2 (0),X,X) Diyd,, (4)
(t) R 0 (X)x(())
and the solution of (4) shall converge to that of the following (in an intuitive sense);
dx((t)= = a ( + 1.رー Σ', a, XX
(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. 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 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 O(1 ?) steps. Finally, our Corollary 5.2 shows (5) gives exactly the same regularization as tr tr[v²1, T.) for label noise SGD,
The main contributions of this paper are summarized as follows.
In Section we propose a mathematical framework study the implicit bias of SGD with infinitesimal LR. Our main theorem (Theorem 4.6) gives the limiting diffusion of SGD with LR for 0)) steps as 7 -> 0 and allows any covariance structure.
2. In Section 5, we give limiting dynamics of SGD with isotropic noise and label noise.
In Section 6, we show for any initialization, SGD with label noise achieves O( d) sample complexity for learning a k-sparse overparametrized linear model (Woodworth et al., 2020). In this case, the implicit regularizer a data-dependent weighted & regularizer, meaning noise can help reduce the norm and even escape the kernel regime. The O(k In d') rate is minimax optimal (Raskutti et al., 2012) and improves over Ô(k²) upper bound by HaoChen al. (2020). In contrast, vanilla GD requires 2(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.
3

---

Published as a conference paper at ICLR 2022
2
RELATED WORKS
Loss Landscape of Overparametrized Models A phenomenon known as mode connectivity has been observed that local minimizers of the loss function of neural network are connected by simple paths (Freeman & Bruna, 2016; Garipov et al., 2018; Draxler et al., 201 8), 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 Kuditípudi 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., 2b; Soudry et al., 2018; Li al., 2018; 2020a). In particular, the SGD noise is widely believed 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 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 aarameter-depnnden noise will bias SGD towards a low-complexity local minimizer. Similar implicit bias has also been studied for overparametrized nonlinear statistical models by Fan et al. (2020). Several existing works (Vaskevicius al., 2019; Woodworth et al., 2020; Zhao et al., 2019) have shown that for the quadratically overparametrized linear model, i.e., w = 02 802 or w @@a, u. gradient descent/flow from small initialization implicitly regularizes & 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 al., 2018; Allen-Zhu et al., 201 2019b;a; Du al., 2019; Zou et al., 2020), where the learnt function minimizes the corresponding RKHS norm (Arora et al., 2019b; Chizat et al., 2018).
Modelling Stochastic First-Order Methods with Itō 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 often more elegant and can provide fruitful insights like the linear scaling rule (Krizhevsky. 2014; Goyal et al., 2017) and the intrinsic learning rate (L.i et al., 2020b). recent work by Li 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 al. (2021) relates the strength of implicit regularization to training speed.
NOTATION AND PRELIMINARIES
Given loss L, the GF governed by L can be described through a mapping o: R: x [0, 00) RDD
satisfying Φ(x, t) = x f (L s))ds. We further denote the limiting mapping ((z) = lim;-xx Φ(z, t) whenever the limit exists. We denote 1, E R_ the one-hot vector where E-th coordinate is 1, and the all 1 vector. See Appendix A for a complete clarification of notations.
3.1 MANIFOLD OF LOCAL MINIMIZERS
Assumption 3.1. Assume that the loss L RD + R is a C3 function, and that I is a (D- M)- dimensional C² cu-smmnnifflld of D for some integer 0 S M < D, where for all E x is local minimizer of L and rank 2 L(x)) = M.
Assumption 3.2. Assume that U is an open neighborhood of T satisfying that gradient flow starting in U converges some point in r, i.e., Vr U, ((x) E Г. (Then & is e² on U by Falconer (1983).)
When does such manifold exist? The vast overparametriization in modern deep learning is 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
4

---

Published as a conference paper at ICLR 2022
training dataset, i.e., Va RD, L(z) gRn L(x') implies f((x) y/ for all i [n], then by preimage theorem (Banyaga & Hurtubise, 2013), the manifold := {I [ R f((z) = Uc. vi [n]] of dimension D - if the Jacobian matrix [V (z)..... fn has rank n for all e r. Note this condition is equivalent to that NTK at æ has full rank, which very common in literature.
4 LIMITING DIFFUSION OF SGD
In Section 4.1 we first recap the main result of Katzenberger (1991). In Section 4.2 we derive the closed-form expressions of UD and 8 We present our main result Section 4.3. We remark that sometimes we omit the dependency on f to make things clearer.
4.1 RECAP OF KATZENBERGER THEOREM
Let (An) be a sequence of integrators, where each A, : R R is a non-decreasing function with (0) = 0. Let {Zn Jn21 be a sequence of -vallue stochastic processes defined on R.
Given loss function L and noise covariance function σ, we consider the following stochastic process:
= x())+ + [' (X.,))))ZZ(s) L." - (L((X(()))AA,
(6)
In particular, when the integrator sequence n21 increases infinitely fast, meaning that Ve 0, inf (A. At)) + 0o as n + 0o, we call (6) a Katzenberger process.
One difficulty for directly studying the limiting dynamics of (t) is that the point-wise limit as + Do become discontinuous at = 0 if X (0) & r. The reason is that clearly limn- X₂ (0) X (0), but for any O, since {An)a21 increases infinitely fast, one can prove lit Xx (1) r! To circumvent this issue, we consider (t) (t) ( An (t)) + (0). Then for each n 1, we have Y, (0) (x(()) and lim, (t) oim (t). Thus (t) has the same limit on (0, 00) as X,, (t), but the limit of the former is further continuous at = O.
Theorem 4.1 (Informal version of Theorem B.7, Katzenberger 1991). Suppose the loss L,, manifold T and neighborhood U satisfies Assumptions 3.1 and 3.2. Let {X,,n>1 be a sequence of Katzenberger process with {An 2221, Zn Jn21. Let (t) = (t. - (0) (n))),A())))))) An(t)) + (Xo). Under technical assumptions, holds that if (Yn, (n) converges to some (Y, W) in distribution, where (W()))22) the standard Brownian motion, then Y stays on I and admits
D ((t) Y()) S[ (Y))()))))(() E o, ( Dyds.
Indeed, SGD (1) can be rewritten into Katzenberger process as in the following lemma.
Lemma 4.2. Let {mmm Loc be any positive sequence with lim,-1oo Th = 0, A,, = nn Lt/n2..
and (t) The Ek=1 [t/v/] VE(1lg 1),, where E1s E2... itd. Unif([E). Then with the same initialization (0) = Inn, (0) (0), (n(k)) defined by (6) a Katzenberger process and is equal Tn (k) defined in (1) with LR equal nn for all k 2 1. Moreover, the counterpart of (7) is
r(t) = ((x(0)) + [ (0(Y)
1 J, (Y)))))
+ 2
(8)
where E = TO and (W is E-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 the derivatives of Φ. 0 and 0, but Katzenberger (1991) did not give their dependency on loss L. To resolve this, we explicitly calculate the derivatives of +b on I in terms of the derivatives of L in Section 4.2.
The second difficulty comes from the convergence of (Y,, Z,,) which we assume as granted for brevity in Theorem 1. In fact, the full version of Theorem 4.1 (see Theorem B.7) concerns the stopped version of Y, with respect to some compact K C U, i... (((t,(()(t) (t) = Yt,(t p,(K)) where nn (K) is the stopping time of Y, leaving K. As noted in Katzenberger (1991), we need the
5

---

Published as a conference paper at ICLR 2022
convergence of jo (K) for (( 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.
4.2 CLOSED- FORSD- EXPRESSION OF THE LIMITING DIFFUSION
We can calculate the derivatives of & by relating to those of L. Here the key observation is the invariance of & along the trajectory of GF.
Lemma 4.3. For any x E T, (0)x) E DDD is the projection matrix onto tangent space T2 (T).
To express the second-orde derivatives compactly, we introduce the notion of Lyapunov operator.
Definition 4.4 (Lyapunov Operator), For a symmetrii matrix H, we define WH {E RDDD
Σ = Σ HHH = Σ = } and Lyapunov Operator Lи : WH + WH as (x = E + ΣH. It's easy to verify L7 is well-defined on WH.
Lemma 4.5. Let x be any point in and Σ = E(x) = 00 T (x) E DxD be the noise covariance
al x². Then Σ can be decomposed as Σ Σ + E1 + T.11 + Ei.ll where El = 0DE0e
E1 = -ID (TD - aT) and T11l1 = =II. @ aΦ) are the noise covariance in tangent space, normal space and across both spaces, respectively. Then holds that
[ = (V (VL)) [In] (VL) [ (E1)] +0000 (VL) [(VL))EL,]]] (9)
4.3 MAIN RESULT
Now we are ready to present our main result. It's direct combination of Theorem B.9 and Lemma 4.5. Theorem 4.6. Suppose the loss function L, the manifold of locul minimizer and the open neighbor- hood U satisfy Assumptions 3.1 and 3.2, and (0) z00) U for all n 0. If SDE (10) has global solution Y with Y (0) = z0) and Y never leaves U. i.e,, P|Y (t) E U, Vt 2 0j = 1, then for any T > 0, I, ((T/n"]) converges in distribution to Y(T) as 17 0. dY (t) 2/1 (Y) + 1 v L( [X)]
2
Tangent Noise
Tangent Noise Compensation
(10)
22000 ( 22 (VL)(Y) (VLLY))L - -80 (VL))) ((())))
Mixed Regularization
Normal Regulatizatinn
where E T and Eqj. 21,E1, 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, ())(2()) (t), 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 I' itself, i.e., it's same for all loss L which locally defines the same T. (b) The noise in the normal space is killed since the limiting dynamics always stay on I. However, its second order effect (Itô correction term) takes place vector field on I, which induces the Noise Regularization and Mixed Regularization term, corresponding to the mixed and normal noise covariance respectively.
Remark 4.7. In Appendix B.4 we indeed prove a stronger version of Theorem 4.6 that the sample paths of SGD converge in distribution, i.e., let (t) 2. (L,((t/n]] }), then ã. weakly converges to Y on [0, Tl. Moreover, we only assume the existence of a global solution for ease of presentation. As long as there exists compact K C T such that Y stays in K on 00, 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.
IMPLICATIONS AND EXAMPLES
In this section, we derive the limiting dynamics for two notable noise types, where we fix the expected loss and the noise distribution, and only drive to O. The proofs are deferred into Appendix C.3.
For notational convenience, we drop dependency zr.

---

Published as a conference paper at ICLR 2022
Type I: Isotropic Noise. Isotropic noise means E(r) ID for any E I (Shi al., 2020). The following theorem shows that the limiting diffusion with isotropic noise can be viewed as Brownian Motion plus Riemannian Gradient Flow with respect to the pseudo-determinan of V L.
Corollary 5.1 (Limiting Diffusion for Isotropic Noise). If ID on r, SDE (10) is then
dt ( ( (VPMM)) (Y)) dt 2
1 a (V) (VL(L) (11)
Brownian Motion Manifold
Normal Regularization
where [++(()]+=
L lima-0 (F the pscudo-determinant of V²L(Y). [V2(()) is
=
also equal to the sum of log of non-zen eigenvalue values of LYY.
Type I: Label Noise. When doing SGD for Ez-regression on dataset {zz, Vi y)) |ー adding label noise (Blane et al., 2020; Damian et al., 2021) means replacing the true label at iteration k, yo.
by fresh noisy label Vi. Vi + 5g. where 6% ltd. Unif{ 5} for some constant ō > 0. Then the corresponding loss becomes 1/4 (fie (x) Di2 where f., (x) is the output of the model with parameter x on data Zig. So the label noise SGD update is
Ta+1 24 - 7/2 V (fig (2x) you + 5,,)² ) = 2h
Viu + 8xx V V (Ix). (12)
} at T.. then the model must interpolate the whole dataset, i.e., [(()). = y4 for all E [n]. and thus here the manifold I' is subset of {x R f.(()) yo, Vi E [n]].. Here the key property of the label noise used previous works is E(x) = E Σ Vali(t)Vafi(()) = L(x). Lately, Damian et al. (2021) further generalizes the analysis to other losses, e.g., logistic loss and exponential loss, as long as they satisfy E(x) = 2 Lzz) for some constant e > 0.
Suppose the model can achieve the global minimum of the loss L(x) := _ [
In sharp contrast the delicate discrete-time analysis in Blanc 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, no longer need do Taylor expansion manually nor carefully control the infinitesimal variables of different orders together. is also worth mentioning that our framework immediately gives global analysis of O(1 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 prove concrete generalization upper bound in non-convex problem, the overparametrized linear model (Woodworth et al., 2020; HaoChen et al., 2020).
Corollary 5.2 (Limiting Flow for Label Noise). If = cVLL on I for some constant c > 0, SDE (10) can be simplified into (13) where the regularization is from the noise in the normal space.
(t = 1/4 (Y) L()V
(13)
6 PROVABLE GENERALIZATION BENEFIT WITH LABEL NOISE
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 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 l norm while large initialization corresponds to the kernel regime and prefers solutions with small &g 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 ( norm. When the groundtruth is s-sparse, this provides (x In d) vs 2(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 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 k-sparse and n 2 2(K In d) training data are sampled from either ii.dd Gaussian or Boolean distribution. Then for any initialization Tmir (except a zero-measure set) and any c > O, there exist no. T > O such that for any n < To. OLM trained with label noise SGD (12) with LR equal n for T//n² steps returns an e-optimal solution, with probability of e [Q(n) over the randomness of the training dataset.
7

---

Published as a conference paper at ICLR 2022
The proof roadmap of Theorem 1.1. is the following:
1. Show Assumption 3. 1. is satisfied, i.c., the set of local minimizers, r, is indeed a manifold and the hessian V L(x) is non-degenerate on (by Lemma 6.2);
2. Show Assumption 3.2 is satisfied, i.c., (U) C I' (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 al. (2020), which assumes w² E {0, 1)d and their reparametrization can only express positive linear functions, i.e., w u Their Ō(x²) rate is achieved with delicate three phase LR schedule, while our Ox In d) rate only uses constant LR.
Setting: Let {(2,, y:) He(m) be the training dataset where 21,-..,, Zn Lid Uni(((+1})) or (0, Ia) and each y; (zi, w²) for some unknown w" E Rd. We assume that w" is x-sparse for some < d. Denote r = (*) E IRD = 12d, and we will use æ and (u, v) exchangeably as the parameter of
functions defined on RD in the sequel. For each E [n]. define fi(x)) = fi(u, v) 2 (u - o0
Then we fit {ax,, (:),:e)) with an overparametrized model through the following loss function:
[(x) = L(u,e) :/- E (,(,,)), v), where (u, )) (ft v) 3:)
(14)
Va E T. For simplicity, we define Z (21,..., T E Raxd and Y (yI,.... vn) R". Consider the following manifold:
is straightforward to verify that (2 = 4 " (=100)) -2,00)
T
n
-' [====(u, E (uu 02 - 02 = }} where U = (IR \ (0))
(15)
We verify that the above loss function L and manifold I satisfy Assumption 3.1 by Lemma 6.2, and that the neighborhood U and I satisfy Assumption 3.2 by Lemma 6.3.
Lemma 6.2. Consider the loss L defined in (14) and manifold I defined in (15). If data is full rank, i.e., rank(Z) = n, then it holds that (a). is a smooth manifold of dimension D - n; (b).
rank(V² Lxx)) = n. for all
I. In particular, rank(Z) = holds with probability for Gaussian distribution and with probability 1 - ca for Boolean distribution for some constant e E 10,
Lemma 6.3. Consider the loss function L defined in (14), manifold T and its open neighborhood defined in (15). For gradient flow da, (xx) starting at any Zo E U. it holds that 4(zo) E T. 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.,
U1 = Di a, vj E (r)) 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) tr(E(x)) tr( ( Without loss of generality, we take 5 = 1. Hence, it follows that
=
R(x)= 4 Σ E" 23.5) (u} + uz),
D
(16)
The limiting behavior of label noise SGD is described by Riemannian gradient flow on I' as follows:
dr, 11/4- 1/4-00(z.) V R(x:)dt, with zo (<Einit) E I.
(17)
The goal is to show that the above limiting flow will converge to the underlying groundtruth x: (0)
where (u", x²) [[w'] [ 01/2)
6.1 LIMITING FLOW CONVERGES TO MINIMIZERS OF REGULARIZER
In this subsection we show limiting flow (13) starting from anywhere on I' 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 OK In 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 T' However, the limiting flow itself is not even defined on the boundary of the manifold r. Even we can extend (.())VR(-) continuously to entire R D the continuous extension is not everywhere differentiable.
Thus the non-compactnes: 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
8

---

Published as a conference paper at ICLR 2022
analytic or smooth and semialgebraic. The latter ensures the so-called 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-t.ojasiewiczz condition (a special case of KI. condition) (Polyak, 1964) via careful calculation (by Lemma 1.1
(KL)
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 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 con-onnveeeence to sub-optimal stationary points with a novel approach: we show that for any stationary point 2; whenever there exists descent direction of the regularizer R at x, we can construct a potential function which increases monotonically along the flow around z, while the potential function is equal to 00 at x, leading to contradiction. (See proof of Lemma 6.5.)
Lemma 6.5. Let [10 C RDP be generated by the flow defined in (17) with any initialization Fo E r. Then Io lim 2y exists. Moreover, Toe = x² is the optimal solution of (18).
6.2 MINIMIZER OF THE REGULARIZER RECOVERS THE SPARSE GROUNDTRUTH
Note 1 -i=1 2²,- = 1 when 2,, ~ Unif(-1, 13, and we can show minimizing R(z) on I, (18), is equivalent to finding the minimum l norm solution of Equation (14). Standard results in sparse recovery imply that minimum & norm solution recovers with the sparse groundtruth. The gaussian case is more complicated but still can be proved with techniques from Tropp (2015).
sdd
minimize R(x) = Σ (E=1 R subject to 4(u - - 022 Zww
2
(ng + 13),
j=1
(18)
Lemma 6.6. Let 21...., 2n ltd. Unif( or (0 La). Then there exist some constants C, >> 0 such that if n 2 Ck In d, then with probability at least 1 - e the optimal solution of (18), (û, 0). is unique up to sign flips of each coordinate and recovers the groundtruth, i.e., 0 -002 w².
6.3 LOWER BOUND FOR GRADIENT DESCENT IN THE KERNEL REGIME
In this subsection we show GD needs at least (2(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 Ō( In 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 to = vo a I with a - DC. holds that fu vp) = 0 and V fee( vi) = [azz, -az] for each E [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, 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 21,..., 2n itd.
NO Ia) and y; 2 w", for all i E [n]. Define the loss with linearized model as L(x) = E (f(zo)) + (V)) x zo) y;)², where (9) and To (Ca) = at). Then for any groundtruth w², any learning rate schedule {ndh1 and any fixed
number of steps T, the expected &z loss of (T) is at least (1 - 3) || w l2. where x(() is the T-th iterate of GD on L, i.e., x(1 +1) =x))=())-- V (x(()), for all t 2 0.
7 CONCLUSION AND FUTURE WORK
We propose a mathematical framework to study the implicit bias of SGD with infinitesimal LR. We show that with arbitrary noise covariance, θ(7 2) steps of SGD converge to limiting diffusion on certain manifold of local minimizer, as the LR 1) -> O. 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.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
We thank Yangyang Li for pointing us to Katzenberger (1991). We also thank Wei Zhan and Jason Lee for helpful discussions.
The authors acknowledge support from NSF, ONR. Simons Foundation, Schmidt Foundation, Mozilla Research, Amazon Research, DARPA and SRC. ZL. is also supported by Microsoft Research PhD Fellowship.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems, 2019a.
Zeyuan Aluen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via parameterization. In International Conference on Machine Learning, pp. 242-252. PMLR, 2019b.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International Conference on Machine Learning, pp. 244-253. PMLR, 2018.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. arXiv preprint arXiv:1905. 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332. PMLR, 2019b.
Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. arXiv preprint arXiv:2102.09769) 2021.
Augustin Banyaga and David Hurtubise. Lectures on Morse homology, volume 29. Springer Science & Business Media, 2013.
Patrick Billingsley. Convergence of probability measures. John Wiley & Sons, 2013.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by omstein--hheeneec like process. In Conference on learning theory, pp. 483-513. PMLR, 2020.
Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and langevin processes. In International Conference on Machine Learning, pp. 1810-1819. PMLR, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. arXiv preprint arXiv: 181.v:181..07956, 2018
Cooper. The critical locus of overparameterized neural networks. arXiv preprint arXiv:2005.0421 2020.
Yaim Cooper. The loss landscape of overparameterized neural networks. arXiv preprint arXiv: 10200, 2018.
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. arXiv preprint arXiv:2106.00 106.55000000000000 2021.
Amit Daniely. Sgd learns the conjugate kernel class of the network. arXiv preprint arXiv: 1702.08503. 2017.
Manfredo P Do Carmo. Riemannian geometry. Springer Science & Business Media, 2013.
10

---

Published as a conference paper at ICLR 2022
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309-1318. PMLR. 2018.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pp. 1675-1685. PMLR, 2019.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.0205: 2018.
J. Falconer. Differentiation of the limit mapping in dynamical system. Journal of the London doi: 10.11 12/11mmm///222222.3366
Mathematical Society, 52-27(2):356-37) 1983, ISSN 0024-6107.
Jianqing Fan, Zhuoran Yang. and Mengxin Yu. Understanding implicit regularization in over- parameterized nonlinear statistical model. arXiv preprint arXiv:2007.0833 00:2222 2020.
Benjamin Fehrman, Benjamin Gess, and Arnuff Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions. Journal of Machine Learning Research, 21, 2020.
C Daniel Freeman and Joan Bruna. Topology and geometry of half-rectified network optimization. arXiv preprint ar105:161..01540, 2016.
Timur Garipov, Pavel Izmailov, Dmitri Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of dans. arXiv preprint arXiv: 1800.10020 2018.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: (ol::706.0000 02677, 2017.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pp. 1832 1841. PMLR, 2018a.
Suriya Gunasekar, Blake Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nathan Srebro. Implicit regularization in matrix factorization. In 2018 Information Theory and Applications Workshop (ITA), pp. 1-10. IEEE, 2018b.
Jeff Z HaoChen, Colin Wei, Jason D Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006. 08680, 2020.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the gen- cralization gap in large batch training of neural networks. arXiv preprint arXiv:1705.0877. 2017.
Andrew Holbrook. Differentiating the pseudo determinant. Linear Algebra and its Applications, 548: 293-304, 2018.
Elton P Hsu. Stochastic analysis on manifolds. Number 38. American Mathematical Soc., 2002.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv:1806.07577 2018.
Stanislaw Jastrzebski. Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04600 2017.
Jeff Kahn, János Komlós, and Endre Szemerédi. On the probability that random 1-matrix is singular. Journal of the American Mathematical Society, 8(1):223-240 1995.
Ioannis Karatzas and Steven Shreve. Brownian motion and stochastic calculus, volume 1 113. springer, 2014.
11

---

Published as a conference paper at ICLR 2022
Gary Shon Katzenberger. Solutions of a stochastic differential equation forced onto manifold by large drift. The Annals of Probability. pp. 1587-1628, 1991.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint 100.:466, 2016.
Alex Krizhevsky, One weird trick for parallelizing convolutional neural networks. arXiv preprint 2014.
Rohith Kuditipudi, Xiang Wang. Holden Lee, Yi Zhang. Zhiyuan Li, Wei Hu, Sanjeev Arora, and Rong Ge. Explaining landscape connectivity of low-cost solutions for multilayer nets. arXiv preprint arXiv: 000///900000 06247, 2019.
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robed Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-48. Springer, 2012.
Jason D Lee, Max Simchowitz. Michael Jordan, and Benjamin Recht. Gradient descent only converges minimizers. In Conference learning theory, pp. 1246-1257 PMLR, 2016.
Jason D Lee, Ioannis Panageas, Georgios Piliouras, Max Simchowitz, Michael I Jordan, and Benjamin Recht. First-order methods almost always avoid saddle points. arXiv preprint arXiv:1710.07406, 2017.
Jian Li, Xuanyuan Luo, and Mingda Qiao. On generalization error bounds of noisy gradient methods for non-convex learning. arXiv preprint arXiv:1902.0062. 2019a.
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference Machine Learning, pp. 2101-2110. PMLR, 2017.
Qianxiao Li, Cheng Tai, and Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. The Journal of Machine Learning Research, 20 (1):1474-1520, 2019b.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent structured data. arXiv preprint 1808. 2018.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over- wer-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint 10Xvv1900.004555 2019c.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference Learning Representations, 2020a.
Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems, 33, 2020b.
Zhiyuan Li, Yi Zhang, and Sanjeev Arora. Why are convolutional nets more sample-efficient than fully-connected nets? arXiv preprint 1=Xiv:2010.08515 2020c.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes). arXiv preprint 1r2iv:2102.12470, 2021.
Shiyu Liang, Ruoyu Sun, Yixuan Li, and Rayadurgam Srikant. Understanding the loss surface of neural networks for binary classification. In International Conference on Machine Learning, pp- 2835-2843 PMLR, 2018.
Stanislaw Lojasiewicz. A topological property of real analytic subsets. Coll. du CNRS, Les équations aux dérivées partielles, 117(87-89):2, 1963.
12

---

Published as a conference paper at ICLR 2022
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.0530 2019.
Andrew Y Ng. Feature selection, 1 I vs. 12 regularization, and rotational invariance. In Proceedings of the twenty-first international conference on Machine learning, pp. 78, 2004.
Quynh Nguyen. On connected sublevel sets in deep learning. In International Conference on Machine Learning, pp. 4790-4799. PMLR, 2019.
Quynh Nguyen, Mahesh Chandra Mukkamala, and Matthias Hein. On the loss landscape of class of deep neural networks with no bad local valleys. arXiv preprint arXiv: 1809::880.1104555 2018.
Lawrence M. Perko. Differential equations and dynamical systems. 2001.
Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit stochasticity. arXiv preprint arXiv:2106.0952 06.04444 2021.
David Pollard. Convergence of stochastic processes. Springer Science & Business Media, 2012.
Boris Polyak. Gradient methods for solving equations and inequalities. USSR Computational Mathematics and Mathematical Physics, 4(6):17-32, 1964.
Garvesh Raskutti, Martin J Wainwright, and Bin Yu. Minimax-optinat rates for sparse additive models over kernel classes via convex programming. Journal of machine learning research, 13(2), 2012.
Bin Shi, Weijie J Su, and Michael I Jordan. On learning rates and schr' odinger operators. arXiv preprint arXiv:2004. 006777000040.77 2020.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018.
Joel Tropp. Convex recovery of structured signal from independent random linear measurements. In Sampling Theory, Renaissance, pp. 67-101. Springer, 2015.
Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. Advances in Neural Information Processing Systems, 32:2972-2988 2019.
Luca Venturi, Afonso S Bandeira, and Joan Bruna. Spurious valleys in two-layer neural network optimization landscapes. arXiv preprint arXiv:1802.06388, 2018.
Yeming Wen, Kevin Luk, Maxime Gazeau, Guodong Zhang, Harris Chan, and Jimmy Ba. Interplay between optimization and generalization of stochastic gradient descent with covariance noise. arXiv preprint arXiv: arXiv:1902.088 082340 2019.
Ward Whitt. Stochastic-procees limits: an introduction to stochastic-process limits and their applica- tion to queues. Springer Science & Business Media, 2002.
Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part ii: Continuous time analysis. arXiv preprint arXivv2106.025588 2021.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference Learning Theory, pp. 3633-36673. PMLR, 2020.
Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the noisy gradient descent that generalizes sgd. In International Conference on Machine Learning, pp. 10367-10376. PMLR, 2020.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent escapes from sharp minima exponentially fast. arXiv preprint arXiv:2002.03495, 2020.
13

---

Published as a conference paper at ICLR 2022
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint 190.::1900.07700 2019.
Peng Zhao, Yun Yang, and Qiao-Chu He. Implicit regularization via hadamard product over- parametrization in high-dimensional linear regression. arXiv preprini (rxxiv/903.09360 2019.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint urXiv:1803.00199 2018.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes parameterized deep relu networks. Machine Learning, 109(3)4467-492, 2020.
14

---

Published as a conference paper at ICLR 2022
A PRELIMINARIES ON STOCHASTIC PROCESSES
We first clarify the notations in this paper. For any integer k, we denote C# as the set of the times continuously differentiable functions. We denote a A b = min bj. For any vector u, and a E IR, we define [u @ u/; = 4 and [u ; = o.. For any matrix A, we denote its pseudo inverse by AT. For mapping F RR + RD, we denote the Jacobian of F at æ by OF(x) E I Dxn where the (i.j)-th entry is (z). We also use [F(x)(uu [u] and 0 FFx) [u, u] 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 0² F by viewing it linear mapping defined on RD & RDD ~ RD", in the sense that @ (x)) = Σm 8 (r)))ee, eyJEug, for any E E RDDD, For any submanifold I C TR. and a E r, we denote by (T) the tangent space of I at x and (T) the normal space of I at x.
Next, we review 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 classies like Karatzas Shreve (2014); Billingsley (2013); Pollard (2012) for more systematic derivations.
Throughout the rest of this section, let be a Banach space equipped with norm || ||| e.g. (R, and (R2 112).
A.1 CADLĂGFUNCTION AND METRIC
Definition A.1 (Cadlagfunction). Let T E [0, oo]. A function g (0. 7) - is càdlăg for all and its left limit g(--) exists. Let T) be the set of all
E [0, T) it is right-continuous càdlàgfunction mapping [0, () into E. We also use Dz T') to denote the set of all continuous function mapping (0) into By definition, Ce(0, T) Do T).
Definition A.2 (Continuity modulus). For any function f [0, oo) + & and any interval 1 S [0, oc), we define
w(f: I) = sup (+) - f()Il- s,tEr
For any N E N and 0 >> O, further define the continuity modulus of continuous f as
wf(
sup
[t, 0))
0SiEt+0gN
Moreover, the continuity modulus of cådlågf D=(0, oo) is defined as
(f inf
max w (f: S to < ... < f, = N. inf 1.-1) 20
isr
ier
Definition A.3 (Jump). For any E De T), we define the jump of g at to be
Ag(() g(() g(--).
For any 8 >> 0, we define hs : 10, oo) + oo) by
hs(r) = 10
if <
{1 8/r if 2 5
We then further define Jo D2o (0, 00) + D2, oo) (Katzenberger, 1991) as
Ja(g)(1)= = Σ ho (a((Ag())))AA(())
(19)
0<858
Definition A.4 (Skorohod metric on [0 oo)). For each finite T > O and each pair of functions f.g [0 oo), define dr (f.g) as the infimum of all those values of 8 for which there exist grids Sms with tx, Sk 2 T' such that
0 < to < t < ... < tm and < 80 < < ..... ..... k. and
14:58 s,l S 0 for i O,...
(I((t) - g [1) <
if (1, E [Li 4++1 x [st, 8441)
for = 0, O....,k k - 1. The Skorohod metric on De oc) is defined to be
2
d(f.g) = Σ 2 T min (1, (f).g)))
T=1
15

---

Published as a conference paper at ICLR 2022
A.2 STOCHASTIC PROCESSES AND STOCHASTIC INTEGRAL
Let F,,F, [Fr)r2o, P) be filtered probability space. Definition A.5 (Cross variation). Let X and be two (2oadea-ddeeeed stochastic processes such that X has sample paths in Dro [0, oo) and Y has samples paths in Dx- [0, oo), then the cross variation of and Y on (0,t). f). denoted by [X, (t)), defined to the limit of
Σ (X(t++1) ((t))) ((+)) 1=0
in probability the mesh size of to <t ... tm t goes to if it exists. Moreover, for Y itself, we write
M E..YY
1=1
Definition A.6 (Martingale). Let (X(t))22 be {Filrzo-adapted stochastic process. If for all 0 S s & t, it holds that
B[X(() [F) X()),
then X called martingale.
Definition A.7 (Local martingale). Let (X(t) (20 be addpteedd stochastic process. If there exists a sequence of (Fih22o-topping time, {Tk 3220 such that
P[ < TT-i] 1, plimme- Tk co| = 1.
and {XT* is a (adpppdddppeed martingale,
then X is called a local martingale.
Definition A.8 (Semimartingale). Let (t)))122 be a (Fe)ezo-adapted stochastic process. If there exists a local martingale (M())) and total variation that (t) (t) + A()), then X called semimartingale.
process {A(t)) with bounded
Lemma A.9 (Itô's Lemma). Let {X (t))r20 be defined through the following Itô drift-diffusion process:
(x = μ(t)dt + (t)dW(t)).
where (W()) is the standard Brownian motion. Then for any twice differentiable function f. holds that
aS(1)(()) = ( (0s at + (V.)) pr + unfo irfo Valol di (V=))T(())W()))
A.3 WEAK CONVERGENCE FOR STOCHASTIC PROCESSES
Let (De oo), A. d) be metric space equipped with o-algebra A and the Skorohod metric defined in the previous subsection.
Let (X,)n20 be sequence of stochastic processes sequence of probability spaces {(2,,, P.) J220 such that each X,, has sample paths [o oc). Also, let X be a stochas- tie process (, IP) with sample paths (0 oc).
Definition A.10 (Weak convergence). A sequence of stochastic process (X,, Jn2o is said to converge in distribution weakly converge X (written as X,, + X)) if and only if for all A-measurable, bounded, and continuous function f :De oc) - I, it holds that
lim E [[f(X.) E [X)). 2=0
(20)
Though we define weak convergence for a countable sequence of stochastic processes, but it still valid if index the stochastie processes by real numbers, e.g., {X,, 15,220. and consider the weak convergence of X, as + O. This is because the convergence in (20) is for a sequence of real numbers, which is also well-defined if we replace lim by lim
16

---

Published as a conference paper at ICLR 2022
Definition A.11 (6- (6-Prohorov distance). Let ó > O. For any two probability measures P and Q metric space with metric d, let (X, Y) be coupling such that P is the marginalized law of X and Q that of Y. We define
(P.Q)= =nf 0 Y). Y) 2 c 6)
Note this distance is not metric because it does not satisfy triangle inequality.
Definition A.12 (Prohorov metric), For any probability measures P' and Q on metrie space with metric d, let (X, Y be coupling such that the marginalized law of X and Q that of Y. Denote the marginal laws of X and Y by ((X) and L(Y) respectively. We define the Prohorov metric as
PP,)))=ine inf > 0: 3X,,)) = -,
Y) 20c
can be shown that Xn, + X is equivalent to limn->oe X) =
Theorem A.13 (Skorohod Representation Theorem). Suppose 1, and P are proba- bility measures & such that Pn = P. Then there is probability space (S2, F, IP) on which are defined E-valued random variables Xn, n2 1,2,... and X with distributions P, and respectively, such that lim,
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.14). However, convergence in Skorohod metric and uniform metric indeed coincide with each other when the limit in CRp [0, 20), i.e., the continuous functions.
Definition A.14 (Uniform metric on Do oo)). For each finite T > and each pair of functions f.g D, T), the uniform metric is defined to be
du dfff.g;
sup [[f(t) g(t)ll- tE(o.T)
The uniform metric on De oo) is defined to be
du (J.))
Σ 2- min (U,g.))
T=1
Lemma A.15 (Problem Section 5, Pollard (2012)). Xf, X in the Skorohod sense, and X has sample paths in Cgo [0, oo), then X,, = X in the uniform metric.
Remark A.16. We shall note the uniform metric defined above is weaker than supteju,cc) llf(t) g(()ll- Convergence in the uniform metric on [0, oo) defined in Definition A. 14 is equivalent to convergence in the uniform metric each compact set [0, T) for T E N+ The same holds for the Skorohod topology.
B LIMITING DIFFUSION OF SGD
In this section, give complete derivation of the limiting diffusion of SGD. Here denote the convergence in distribution. For any U C R.D, denote by Ú its interior. For linear space 5, we use 5- to denote its orthogonal complement.
to
First, mentioned in Assumption 3.2, verify that the mapping + is C² in Lemma B.1. In Appendix B.1 we discuss how different time scalings could affect the coefficients in SDE (2) and (3). Then 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 free-friendly interface for Katzenberger's theorem in Appendix B.4.
Lemma B.1 (Implication of Falconer (1983)). Under Assumption 3.2, 4 is C² on U.
Proof of Lemma B.I. Applying Theorem 5.1 of Falconer (1983) with f()) o(: 1) suffices.
17

---

Published as a conference paper at ICLR 2022
B.1 APPROXIMATING SGD BY SDE
Let's first clarify how we derive the SDEs, (2) and (3), that approximate SGD (1) under different
time scalings. Recall (()) is E-dimensional Brownian motion and that (X) IR R is deterministic noise function. As proposed by Li et al. (2017), one approach to approximate SGD (1) by SDE consider the following SDE:
(t) - VL(X + (X))
where the time correspondence is I = kn, i.e., (kn) ss (k).
Now rescale the above SDE by considering (t) X (tn), which then yields
dx (t) (tn) vL((((n))dd(tn)) (X)(((t)))WWW = -I))((( + Via (X
Now we define W'()) = 4n W ((n). and it's easy to verify that W'(1) is also E-dimensional brownian motion, which means W _ W,i.e. W and W' have the same sample paths in Cra [0, oo). Thus
dX((t) (t + (x(((t))))W((tt = VV)I(X(()) nja (X()))W)))
where the time correspondence is / k, i.e. ((k) = (k)) The above SDE is exactly the same as (2).
Then, to accelerate the above SDE by 7j 2 times, let's define x(t) = (t/n7)) Then it follows that dx (1/1)) - --====>VVV(XX VI + no no ((((()))))))((
- +o ((((t)) (nW) (nW
Again note that (t/1)) A ((t) in sample paths and thus is also E- E-Brownian motion. Here the time correspondence is f koy? i.e., evolving for constant time with the above SDE approximates ((1/n²) steps of SGD. In this way, we derive SDE (3) in the main context.
B.2 NECESSARY CONDITIONS
Below we collect the necessary conditions imposed on [Znb22) and {Anh22 in Katzenberger (1991). Recall that consider the following stochastic process
= L (X,())))ZZ)
/. VL((X())))AA,())
For any stopping time the stopped process is defined (T( X A r). For any compact
K C. U, we define the stopping time of X, leaving K as An(K) = inf(t 2 0 (t-t)) K or (1) & Kj.
Condition B.2. The integrator sequence {An)n21 is asymptotically continuous:
sup |A, (1)
t>o
(t-)) + where (t-t)) = lim,- An(s) is the left limit of A,, at t.
Condition B.3. The integrator sequence {An)22 increases infinitely fast: Ve 0, inf (An
(20
(t)) + xc.
Condition B.4 (51,, Katzenberger 1991). For every T > 0, as n + 0o, holds that
sup (<I
|| | Z, (1) + 0.
Condition B.5 (Condition 4.2, Katzenberger 1991). For each > 1. let Y, be (FF')
semimartingale with sample paths in Dgo 00, so). Assume that for some > (allowing = 00) and every there exist stopping times {T,] > 1} and decomposition of Y, (s...)
18

---

Published as a conference paper at ICLR 2022
into local martingale M, plus finite variation process F, such that Plr," I m] S 1/m, (tM](tAT)) + Tunr (Fn) b221 is uniformly integrable for every 2 0 and m 2 1, and
lim lim sup IP 7-+0 n-
sup (T++)(F.) Ii(P.) > = 0, LOSEST
for every >> 0 and T > O, where T-) denotes total variation on the interval [0, Lemma B.6. For SGD iterates defined using the notation in Lemma 4.2, the sequences {Auln21
and { [ hn21 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)n11-
Next, for any > 0 and E T). we have
t t+c
t
t
7m t
1/
An +)+++)+0+ A, (1) = tne
The
mm
nn
7n
which implies that infosrsT (An + E) - (t)) > e/ (2nm) for small enough The Then taking n + oo yields the Condition B.3.
For Condition B.4, note that
{ nn V - 1) = k: The Zl = 10
otherwise.
Therefore, we have || [I, S 21jn V for all > 0. This implies that [t) la + uniformly over t > 0 as n co, which verifies Condition B.4.
We proceed to verify Condition B.5. By the definition of Zn, we know that (t) a jump process with independent increments and thus is a martingale. Therefore, by decomposing Zn M, + Fn with M,, being a local martingale and F2, a finite variation process, we must have Fn 0 and Mn is Zn itself. It then suffices show that [M,](t A xT")) is uniformly integrable for every t 2 and m 2 1. Since M,, is pure jump process, we have
[M,
Σ
(s)
Σ UCOEE
(s)
O<ESIAT
14/W/
Σ
Th 2 43t.
k=1
This implies that [L A " is universally bounded by 46, and thus [M,] A 7 is uniformly integrable. This completes the proof.
Lemma 4.2. Let 1oe be any positive sequence with limn- The 0. A, 7m L.
and (t) = 7h Σ =k=1 VE(1,. 1), where ε1, E2,... idd. Unif([E]). Then with the same initialization (0) T,, (0) x()), (kng) defined by (6) is Katzenberger process and is equal to Noy (k) defined in (1) with LR equal to na for all ls 2 1. Moreover, the counterpart of (7) is
((t) +(X(0) + [' (4 (s) + 2
1 J. (Y)))))))))s
(8)
where E = ToT and (W is E-dimensional standard Brownian motion.
Proof of Lemma 4.2. For any n 2 1, it suffices to show that given X, (km))) = x (k), we further have X,, (1: + 1)m%) = 2n (k 1). By the definition of (t,,,, we have
(lk (K= + 1)ma) - x, (km))) 11+
(K+1)
(L((,()))AAN))) + JenZ (X,,()))
(L(X,, 1)m)) -nn)k))) + oX (x 1)m2) Z, (krri))
The (I... + T/e VEog. (X,(k=)):))
Th (E (B:) nm VErgo (5n, (R:) Emp (I: + 1) Ex,, (k)
19

---

Published as a conference paper at ICLR 2022
where the second equality is because A, (t) and (t) constant on interval [ko (I + 1)n))- This confirms the alignment between (A7, |x21 and {In, (k)]
For the second claim, note that a (z, = 0 for all a E R Z (since the noise has zero- expectation) and that {Z, (t) EZ,, (2) will converge in distribution to a Brownian motion by the classic functional central limit theorem (see, for example, Theorem 4.3.5 in Whitt (2002)). Thus, the limiting diffusion of X,, + can be obtained by substituting Z with the standard Brownian motion W in (22). This completes the proof.
B.3 KATZENBERGER' THEOREM FOR ASYMPTOTICALLY CONTINUOUS CASE
The full Katzenberger's theorem deals with 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), first introduce some definitions. For each n 2 1, let """, {F( froo, IP) be a filtered probability space, Zn, an -valuue cadlag { J-semimartingale with Zn (0) 0 and A,, real-valued cadlag }addpttd nondecreasing process with (0) O. Let on : U - MDD e) be continuous with on - a uniformly on compact subsets of U. Let X, be an IR -vauued cadlag {JT) semmiarrtingale satisfying, for all compact K U.
(t) = x()) + L/" o(xn)dZ, + L' -VA(
(21)
for all & S (K) where (K) = it 2 0 (t-t)) K or (t) & K} is the stopping time of X, leaving K.
Theorem B.7 (Theorem 6.3, Katzenberger 1991). Suppose X() E U, Assumptions 3.1 and 3.2, Condition B.2. B.3, B.4 and B.5 hold. For any compact K C U, define pn( = inf 2 O (1-) & K or (1) $ K}, then the sequence {(Y)=((K) 2am(k) [Mn(K) is relatively compact in DRoxe [0, 2o) x (0,00), If (Y, Z, (1) is a limit point of this sequence under the skorohod metric (Definition A.4), then (Y, Z) is continuous semimartingale, Y()) I for every 2 0 a.s.,
a it (t) & K} a.s. and Y (t) admits
(tt) = (0) +
O ()
D
+ 2 1 k11=1"
Σ
3 (V (2, Z1).
(22)
We note that by Lemma A. 15, convergence in distribution under skorohod metric is equivalent to convergence in distribution under uniform metric Definition A. therefore in the rest of the paper we will only use the uniform metric in the rest of the paper, c.g., whenever we mention Prohorov metric and 8-Prohorov distance, the underlying metric is the uniform metric.
B.4 A USER-FRIENDLY INTERFACE FOR KATZENBERGER THEOREM
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 I and its open neighborhood U satisfy Assumptions 3.1 and 3.2. Let K C U any compact set and fix some Fo E Consider the SGD formulated in Lemma 4.2 where nn (0) Zo. Define
YT, (1) = X,, (1) - (X,,. (0), An. (1) + 1 (X,,.
and png (K) = minft E N Yn (t) & K). Then the sequence {(Y
Pma (K) Lmn : Paa (K)) is relatively compact in DRØXRI [0, x)) [0, oo). Moreover, if (Y, Z, μ) is a limit point of this sequence,
holds that (t) E I' for all tt 2 0, p z inf(t 2 0 (t) & K} and (t) admits
(f) = - tAA
Mnn
D
04 (X(Y +
Σ
1==0 2 11=1
(23)
20

---

Published as a conference paper at ICLR 2022
where (W))) is the standard Brownian motion and of)) is defined in Lemma 4.2.
However, the above theorem is hard to parse and cannot be directly applied if we want to fur- ther study the implicit bias of SGD through this limiting diffusion. Therefore, we develop user-friendly interface to it in below. In particular, Theorem 4.6 is the a special case of Theo- rem B.9. In Theorem 4.6, we replace (o)(())) with z (Y)))) to simplify the equation, = ())((()) and thus this change doesn't It affect the
since ( (O))o (Y()) (O) distribution of the sample paths of the solution.
Theorem B.9. Under the same setting as Theorem B.8, we change the integer index back to 1j 0 with slight abuse of notation. For any stopping time μ and stochastic process {Y (t) such that 2 inf{t 2 0 (t) & K}. (0) = ((ra) and that (Y, p)) satisfy Equation f23) some standard Brownian motion W. For any compact set K _ U and T O, define μ(K) inf 0 Y &
K} } and 5 = (P(((K) & T). Then for any > O, it holds for all sufficiently small LR that:
8²6 (YH Yu(K)NT) S.
(24)
which means there is a coupling between the distribution of the stopped processes Y ((,,,())NT and Ya(K)A such that the uniform metric between them is smaller than with probability at least 1 - 28. In other words, lim 22 (Y. (K)AT yu(K))))= = 0.
Moreover, when (Y()) a global solution to the following limiting diffusson
Y(1) = f b. ( + J Σ
1
D
o 0,0XY())NGG(XXX))OXXX)))
2
==
71=1
and Y never leaves U, i.e. P|vt 2 0, (t) U] holds that YT converges in distribution to Y" as 17 + 0 for any fixed T >> O.
For clarity, we break the proof of Theorem B.9 into two parts, devoted to the two claims respectively.
Proof of the first claim of Theorem B.9. First, Theorem B.8 guarantees there exists a stopping time ji and stochastic process (t) such that
1. P.) satisfies Equation (23);
2. V E T a.s.;
3. ğ 2 ((() : inf 0 ((t) K}.
The above conditions imply that Yā(K) a.s. Since the coefficients in Equation (23) are locally Lipschitz, we claim that (Ýλ(K), ((K)) d (Yu(K), p(K)). To this, note that for any compact K U, the noise function a, 06 and @ are all Lipschitz on K, thus we can extend their definitions R such that the resulting functions are still locally Lipschitz. Based on this extension, applying classic theorem on weak uniqueness (e.g.. Theorem 1.1.10, Hsu 2002) to the extended version of Equation (23) yields the equivalence in law. Thus we only need to prove the first claim for Y.
Let Ex be the event such that ã(K) > T on Er. Then restricted on Er, we have Ý(T A p)) (IT A ((()) as j 2 (KK) holds a.s. We first prove the claim for any convergent subsequence of (Yg)yso.
Now, let {ymmm2) a sequence of LRs such that Ym - 0 and Hmm (K) 7 Y as m + oo. By applying the Skorohod representation theorem, we can put (Y Jmz1 and Y under the same probability space such that Hmm (K) + Y a.s. in the Skorohod metric, or equivalently the uniform
metric (since 8 is continuous) i...
du (ymm (K) .ra) - O, a.s.,
21

---

Published as a conference paper at ICLR 2022
which further implies that for any >> O, there exists some N > 0 such that for all m N,
P de (Y (K)A yuAT) S 8.
Restricted on Ex, we have du (Xit... that for all m N.
(K)r SPATT = du (Y (K)T (P(())T), and it follows
P du (Y (K)AT {A(K)AT) 2 c
IP d (Y (KIAT YRIKIAT) 2 c 8T + P
dd (Y (Y (KIAT paaT) 2 n + +1(
IP
- IP du (Yitn (KYT yTAny 2 c) + P(E
S 28.
By the definition of the Prohorov metric in Definition A. we then
get
p26 (Y (K)AT yA(KAT) S e for all m > N. Therefore, we have
lim p 2 (X (K)AT YACKAT = O.
m 300
Now we claim that it indeed holds that lim, 23 (Y (K)AT {A(K)AT)) = 0. We prove this by contradiction. Suppose otherwise, then there exists some E > such that for all yo > 0, there exists some 7) < no with p²d ( (Y("),, ÿū(K)AT) > c. Consequently, there is a se- quence {Mm hm21 satisfying mim 7m = and p²6 (Y (K) уū(KAT) > c for all m. Since {(YH).. (K)AT 4mm Mm (K)) mm21 is relatively compact, there exists subsequence (WLOG, as- sume is the original sequence itself) converging to (YPAT" W, jz) in distribution. However, repeating W,,
the exactly same argument as above, we would have 9²6 (Y (K)AT YR(K)AT) I for all suffi- ciently 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 sequence of compact
set {km Jm21 such that U Km. = U and Km S. Km+1- For m E N+, we define Hm = +RD UU) and Km = Hm n Bm (0). By definition it holds that Vm < m', Hm S Hmm' and Km K Moreover, since Km, is bounded and closed, Km, is compact for every m. Now we claim U Km = U. Note that U Km = U | Hm n Bm (0) = U Hm. Vo E U, since U is open, we know (x, RD U) > 0, thus there exists mg E N+ such that Vm 2 mo. & (B f1/mm00 (0) + RDD U) and thus x E Hm, which implies x E U Hm- On the other hand, Vr E RR AU it holds that x (B1jm(0) RDD (U) for all m E N+, thus x & Hmm C Km.
U. U\BBI (Bi/m
(0)+
Therefore, since Y E U and is continuous almost surely, random variables lim (((m)) = do a.s., which implies ju(K Km)) converges to co in distribution, i,e,, Vo > 0, 0,T > 0, 3m E N+, such that AK 2 Km, it holds Plu(K)= < T) S 5.
Now we will show for any > 0 and c > 0, there exists 7o such that (Y YY,)) S c for all n S n0- Fixing any T > 0, for any >>>0 0, let ö +. then from above we know exists compact set K, such that F(((K) < T) < 8. We further pick K" = K + Bre (0), where e' can be any real number satisfying
é' < e and K' CU Such c' exists since is open. Note S K', we have
P(u(K') < T)= < (()(K) < T) < 8. Thus by the first claim of Theorem B.9, there exists no > 0, such that for all 17 S no. we have pa (y,ti,
y=(("))T) < 2
Note that p² (Yu(K)AT, yu(K")T) = 0, so we have for all S 17o
36 (y (K")))) < 2 Im
By the definition of 6-Prohorov distance in Definition A.11, we can assume (YR(K)AT y.a,(K is already the coupling such that TP do (Y Yd (HnRK")AT > 2-1T( e) S 36. Below we want to show 36 (Yи(K)T, YYT) < 2- Note that for all t 2 0, yu(K)T (1) E K, thus we know if
22

---

Published as a conference paper at ICLR 2022
Mn(K') S T. then du (YYKK)T yHulk')TT) ) 2-21T
(1((K))) - -N (a (K') 1T du du(K, (K,R'/K)
2 22 e'.
On the other hand, if M,, > T, then Yi = Yn VHU(KKATT Thus we can conclude that du (Yu(KII y, 2 e implies du (YN(K)NI y yHAn(K()TT) 2 2-127 Therefore, further have
PP (YULK)AT ;y") 2 2- s ddu (YukKMNT yui(KINT 2 2 $ 36,
that is,
36 (yn(KAAT Y²T) S 2
22-1M M
Finally, since P² Y((K)TT) = 0, we have for all S 70. ( X,,) - 0 (Y xx)) (y Y,,)) + p (Y
which completes the proof.
Now, we provide the proof of Theorem 4.6 as direct application of Theorem B.9.
1. By the result of Theorem B.8, we know that for each compact set K r, Yu(K) stays on r almost surely, where μ( inf 2 0 (t) & K} is the earliest time that Y leaves K. In other words, for Jm all compact set r, PP 2 0, (() & r, Y(I) K] = 0. Let {m be any sequence of e.g.,
Proof of Theorem 4.6. We first prove that Y never leaves r, i.e., ((Y
compact sets such that Um> K = U and Km C e.g. the ones constructed in the proof of the second claim of Theorem B.9. Therefore, we have
P|t O, (t) & T P z (2Y (TYY()
Σ P 2E820,Y0 0,Y()) e TTY(()) Km =0,
which means Y always stays on
Then recall the decomposition of Σ = Σj + + E1 +E as defined in Lemma 4.5. Since Y never leaves r, by Lemma 4.5, we can rewrite Equation (10)
dY (t) 2// (
1 D - + Σ (,)
where the second equality follows from the definition that Σn dΦE04 dgoo T T0.. This coincides with the formulation of the limiting diffusion in Theorem B.9. Therefore, further combining Lemma 4.2 and the second part of Theorem B.9, we obtain the desired result.
Remark B.10. Our result suggests that for tiny LR T,, SGD dynamics have two phases. In Phase of (1/n)) steps, the SGD iterates move towards the manifold I of local minimizers along GF. Then in Phase Il which of ( (1(//)2) steps, the SGD tterates stay close to I and diffuse approximately according to (10). See Figure 2 for an illustration of this two-phase dynamics. However, since the length of Phase gets negligible compared that of Phase Il when n Theorem 4.6 only reflects the time scaling of Phase II.
C EXPLICIT FORMULA OF THE LIMITING DIFFUSION
In this section, we demonstrate how to compute the derivatives of & by relating to those of the loss function T., and then present the explicit formula of the limiting diffusion.
23

---

Published as a conference paper at ICLR 2022
Gradient Flow
SGD, phase II T2 + (n
SGD, phase T₁ = On
Limiting Flow
r
Figure 2: Illustration for two-phase dynamics of SGD with the same example as in Figure 1. I is an ID manifold of minimizers loss L.
C.I EXPLICIT EXPRESSION OF THE DERIVATIVES
For any r, we choose an orthonormal basis of (T) as (o1....,e--ak UD-M) Let {UD- M41. uD) be an orthonormal basis of (T) so that {ui)ee is an orthonormal basis of RD.
Lemma C.1. For any x and any v E (T),), holds that L2 =
Proof. For any (T)), let {z(1))(2 be a parametrized smooth curve I' such that zx0)) aVL(a)) L2L(x).
and da(t)) = v. Then V ((x,) = 0 for all t. Thus 0 = dt
Lemma C.2. For any R,D, it holds that (T 0 and VL(x)] = (
Proof. Fixing any x E R.D. let dæ(t)) VL(x(t)) initialized z(0) z. Since F(x(t))
(xx) for all f 2 0, we have
(xx(1)) (6((())V = 0.
dt
Evaluating the above equation at 0 yields (L(z) 0. Moreover, take the second order derivative and we have
(2,) = (x)(((()) dx(()) vL(x(1)) (() dr(t) d@
d2
di
Evaluating at completes the proof.
Now can prove Lemma 4.3, restated in below.
Lemma 4.3, For any I E 1,, (2(x) E DxD is the projection matrix onto tangent space T2 (I').
Proof of Lemma 4.3. For any
(T).). let (u(t), t 2 0} be parametrized smooth curve on = V. Since v(t) E I for all t 2 O, we have ((v(t) = v(t), and
such that v(0) = x and thus
do(t)
de
du(t)
d ((v(t))
dr(t)
= 04(x)
dt
di
di
t=
1=0
This implies that (x(x)) for all E (T)..
Next, for any u E (T) and 0, consider expanding VL(r + (V2T(()) at att=
V (r (v+ (")
tu o(t)
24

---

Published as a conference paper at ICLR 2022
where the second equality follows from the assumption that L() is full-rank when restricted on T2 (T). Then since 0 is continuous, it follows that
0 VL
lim
lim 00(x +
(.((.)..
By Lemma C.2, we have 00(x + ( (V then implies that (z())) for all E (T)
= 0 for all t >> 0, which
Therefore, under the basis {vi,..., vN (xx is given by
ID-M
0 0 E DXD
aΦ(x)
that is, the projection matrix onto (T)...
Lemma C.3. For any x E T, it holds that (1))) = O.
Proof. It directly follows from Lemma C.I and Lemma 4.3.
Next, we proceed to compute the second-order derivatives.
Lemma C.4. For any I, u E RD and (I',,, it holds that (b))), ul = -0 (VL) V>I()) - V'II)) (0())
Proof of Lemma C.4. Consider parametrized smooth curve {u(t)}(20 on I' such that c(0) = x and du(t)
t=o = u. We define P(t) = (b(()))) P+(1) = In - P(() and H(I)= = 4 [((((( for all
dt
0. By Lemma C.I and 4.3, we have
P+() H(t) Ht)),
(25)
Denote the derivative of P(I), P+(() and H()) with respect to l as P'()), (P+)(() and H'()). Then differentiating with respect to I, we have
(P (())())()))
- P = P(t)
(26)
Then combining (25) and (26) and evaluating at = 0, we have
P(()) (P) ((O)H
(27)
We decompose P0( and (0) as follows
[ P'(0)
P2(00)
H(O) : (0)
(28)
P (0)
where (0) E (D- () and Hog is the hessian of T. restricted on T (T). Also note that
(O) (0)
ID-M M
0
H (0) H (0)
Hi2(0) H2)(0)
(
0
TM
(
H (0)
0
and thus by (28) we have
Pi) (0)
(0 -H12(0))
((()) = (
P H2) (0)
0
This implies that we must have (0) = and (0) = (0) Similarly, by taking transpose in (28), also have (2) (0)
25

---

Published as a conference paper at ICLR 2022
It then remains determine the value of (1)( Note that since P(t) P(t), we have + = (t), evaluating at 0 yields
(0) = (0)
Therefore, we must have Pi (0) = 0. Combining the above results, we obtain
(0) (0))))'(0)))100 H(0)
Finally, recall that P(t) = (v()))) and thus
d (v(((t))
((()) _
= (=)u-.
dt
(=0
Similarly, we have (0) = 8 (VL) and it follows that
0 Pecho--oottoPFFToo///11111-11100
Lemma C.5. For any x and u E T (T), holds that
( LLL(x) LZx) (|
Proof of Lemma C.5. For any TH)(),, (T), define approximation, we have
for 0. By Taylor
Vl(u(t)) =V
o(t)
(29)
and
= 10 77 [| o()).
(30)
Combine (29) and (30) and apply Lemma C.2, and it follows that
0 = (-
+0 + 2 (m((t)) + +
2 OH
02 [()
where the last equality follows from Lemma C.3. Dividing both sides by 12 and letting 0, we get
8 (x) + (V):)))((V)))))))))))))) [u]) L() 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 E I and E E span(uu T u E (H)))
(0 (0000
(31)
Proof of Lemma C.6. Let A uu." u and B = V2L(x) uu The key observation is that A + = LV) ( + B Therefore, by Lemma C.5, it holds that
8 x) (E (B+B = (x)
20 (0)))) (VL))
Since 2 ( is full-rank when restricted (T), we have span ( (²I + un (2)) E ( (T)) = suan
T (I)). Thus by the linearity of above equation, we can replace B + by any Σ span I T (T)), resulting in the desired equation.
Then Lemma 4.5 directly follows from Lemma C.4 and C.5.
26

---

Published as a conference paper at ICLR 2022
C.2 TANGENT NOISE COMPENSATION ONLY DEPENDENDS ON THE MANIFOLD ITSELF
Here we show that the second term of (10), i.e., the tangent noise compensation for the limiting dynamics to stay I, only depends on I itself.
Lemma C.7. For any I E H, suppose there exist neighborhood U. of x and two loss func- tions L and L' that define the same manifold 1 locally Un, i.e., T n U2 = for | VL(x) = 0} = {xr (L( 0}. Then for any E TT(T), it holdy that ( (P(V) [v.u]
(VL')
Proof of Lemma C.7. Let (2(t))220 be smooth curve on I' with v(0) and du(t)
t=:
Since v(t) stays on I, we have (L(((t)) 0 for all 2 0. Taking derivative for two times yields dPo(t) de = 0. Evaluating it at t 0 and multiplying both
0² (v))( d 2 + (²L(v(t)
sides by V=L(x), we get
d'LL)) -V= dP((t)
dPu(t) di2
= -00
di2
1==0
==
Since (2) is the projection matrix onto (T) by Lemma 4.3, it does not depend on L, so analogously we also have O² [v,
-x(x) dPo(t)
t= as well. The proof is (z)))((t)) drz
thus completed. Note that (+( d'w(t) di r=o is indeed the second fundamental form for v at æ and the value won't t change if we choose another parametric smooth curve with different second-order time derivative. (See Chapter in Do Carmo (2013) for reference.)
C.3 PROOF OF RESULTS IN SECTION 5
Now are ready 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 In on I, SDE (10) is then
dt (t) (Y) + 2 (VL))) dt (L) (L (11)
Brownian on Manifold
Normal Regularization
LY where (VI[())++ = lima-s0 D (FE(( is the peeudo-determinant of VLL).. |V2L(Y))+ is also equal to the sum of log of non-zero eigenvalue values of LYY)
Proof of Corollary 5. Set Σ
06, Σ1 = Tn 0 and Σ || E = 0 in the decomposition
of E by Lemma 4.5, and need show d V E|+) = 0 (VV
Holbrook (2018) shows that the gradient of pscudo-inverse determinant satisfies VIA|+ = AA|+AI.
Thus we have for any vector v RDD, (v, In [V1|+) [VLL4VV = (' V = (v. (VL))(
|44L 82 (VL))))
which completes the proof.
Corollary 5.2 (Limiting Flow for Label Noise). I = cV on I for some constant > 0, SDE (10) can be simplified into (13) where the regularization is from the noise the normal space.
dY (t) = 1/4 04 (
(f)))dt.
(13)
Proof of Corollary 5.2. Since Σ = I, here we have Σ1 Σ and Ed,, Σ1II I1I 0. Thus it suffices show that 202 (VL) [ (2 V vr[ Note that for any v RD
0 V br| V = (ID, 8² [V)) = (In- - 02, ,² (VL))o))
(32)
where the second equality is because the the tangent space of symmetric rank-n matrices at 2
{AV2L L, + LL A E DDD), and every element in this tangent space has zero inner-product
with 0 by Lemma 4.3. Also note that L (V²L) ][ OD), thus (ID -0B, 82 (VL)())) =
20 2 (C (V2T.) - (VL
27

---

Published as a conference paper at ICLR 2022
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 al. (2021) but was only verified via experiments.
Note the normal regularization in both cases 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 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 = k + 2 2 5. For each x E R D we decompose x = aa
[1:22
where 31.2 E R2
and F3.D R D-2 Let Qa 2 be the rotation matrix of angle 0, i.e. Qa
cos
0 and cos
sine
the loss L(x) ;= 1 [1 2 12 1)² + E (2 + (Qa 2, 21.2) x), where a
D- and v is any
2n
vector in R with unit
Here the manifold is given by r == {ar T(x) 0} far RDD 2i + = =1 1, 0, Vj = 3, D).
The basic idea that can add noise in the 'auxiliary dimensions' for j 3,..., to get the regularization force on the circle {>î + z = 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 single auxiliary dimension because from the analysis for label noise, we know is identity, the normal regularization term Equation (10) has 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 2 3 auxiliary dimensions, 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.
when ,
Lemma C.8. Let E E DxD be given by Eg((x) (1+ (QL- Q-=/221122 )2 (22 v,, 2122)) ifi -iijzz 3 or 0 otherwise, then the solution of SDE (10) is the following (33), which implies that (t) moves anti-clockwise with constant angular speed of (D 2)
Y1:2(t) ((2202)221:2(0)) and Ys.((t) 0.
(33)
Proof of Lemma C.8. Note that for any æ T, it holds that
12 (Q0- 21.2)
if j if i j (1,2), otherwise.
(V*L(x)),
(34)
0
Then clearly Σ only brings about noise in the normal space, and specifically, it holds that C (E(z))= diag(0. 0, + Qh,u Q=r224122) ...,, ...., (QP-3y, (=/22212)) Further note that, by the special structure of the hessian in (34) and Lemma C.3, for any r, we have (x: (I2; -1, O,...,............... T)T (22, 21, ... (Q-+/1012) (2-+ Combining these
28

---

Published as a conference paper at ICLR 2022
facts, the dynamics of the first two coordinates in SDE (10) can be simplified into 2x)2(() dt
(T 1:2
Q //2
122
Σ
(1 + (2% Q.Q- 11200,
2
n/2
==3
Q
-1/21122
Q-r/2T1:2 Σ
(1 + (0: 10:-311 vv. Q Q Qaa 1:-3/11
2
j=3
-- Q 7/231:2
1
Q Q-T//221:2.
Qi-3
Σ 10t D, 2=000=-2200 2/221
2
f=3
/=3
1 2
- 0+
D D-2
Q2-1/2112
1100-/22121:
2 Q1/2211:2
2
where the second to the last equality follows from the property Qa and the last equality follows
from the fact that || = 1 for all E r. Note we require k 2 3 (or D 2 5) to allow
E (Qi 1-2/2-1122
2 1/22/2212211 On the other hand, we have dea.p(() 0 as a& kills the movement on that component.
de
The proof is completed by noting that the solution of æ1:2 is
(1. D - Q+/2 x1:2(0),
(12 eexp
and by Lemma C.9,
exp t. D - 2
Qx122) = =(exp(Qx/2)) 1D-22
1DD-21
=Q,
Q 1CD 2)
Lemma C.9. exp(() -)))
cos sin siu
cos
Proof. By definition, for matrix A = (i ). exp(A) = Ei
4'. H. Note A2 = 13 = - and
At = I. and by using this pattern, we can easily check that
Σ
A'
00
2i (-1)² Ei
xx
(2:+1)
coos 1
sin 1 cos 1
1
E
sin
12+
[2:)
D PROOF OF RESULTS IN SECTION
In this section, we present the missing proofs in Section 6 regarding the overparametrized linear model.
For convenience, for any p, r 2 and R.D. we denote by (") the lp norm ball of radius centered u. We also denote Usyj (0,, U141...., uT) for i.j [D].
D.1 PROOF OF THEOREM 6.1
In this subsection, provide the proof of Theorem 6.1.
Theorem 6.1. In the setting of OLM, suppose the groundtruth is k-sparse and n. 2 2( In d) training data are sampled from either ii.d. Gaussian Boolean distribution. Then for any initialization Foit (except a zero-measure set) and any e > O, there exist no- T O such that for any n < TJ. OLM trained with label noise SGD (12) with LR equal nofor [1/n"] steps returns an e-optimal solution, with probability of e [Q(n) over the randomness of the training dataset.
29

---

Published as a conference paper at ICLR 2022
Proof of Theorem 6.1. First, by Lemma 6.6, it holds with probability least - --- that the solution to (18), 2x, is unique up to and satisfies |2,| = =3)(w.). Then on this event, for any c > 0. by Lemma 6.5, there exists some T > 0 such that 2r given by the Riemannian gradient flow (17) satisfies that vr is an c/ e/2-optimal solution of the OLM. For this T, by Theorem 4.6, we know that the [T/n]]] Jthh SGD iterate, J.y (TT((//)]] }), satisfies (LTT,(LL/)))
FT ||| c/2 with probability at least e Q(()) for all sufficiently small 1) > O, and thus (1T ]) is an c-optimal solution of the OLM. Finally, the validity of applying Theorem 4.6 is guaranteed by Lemma 6.2 and 6.3. This
completes the proof.
In the following subsections, we provide the proofs of all the components used in the above proof.
D.2 PROOF OF LEMMA 6.2
Recall that for each i [n]] f((x) f(u,v) (u Vf(z)) 2((100)) and K(x)
(K)(())) where each K2)(2) (V.((z), Vz(())). Then
zt G ( -2 Uv
v2((x) 2
(( Oz) (f((u,)) vi) -diag(zz, 2.).
So for any E r, it holds that
V2L(x)
2
2 Cu U
(44 OI) -4)00)))
(35)
n2
Lemma D.1. For any fixed RP,, suppose (Vf((=))ee()) is linearly independent, then K(x) is full-rank.
Proof of Lemma D.I. Suppose otherwise, then there exists some E R" such that À * and A K K(x)a O. However, note that
A AT(())= Σ (,))K,s())
i==
Σ
Azdy (v) V((())
=
4-=
a
X f.)
which implies that Σi= A (z) = 0. This is contradiction since by assumption (x)) NEj
linearly independent.
Lemma 6.2. Consider the loss L defined in (14) and manifold defined (15). If data is full rank, i.e., rank(Z) = then it holds that (a). I is smooth manifold of dimension D n, (b). rank( rank(V² L(x)) = n for all T. particular, rank(22 Z) = n holds with probability for Gaussian distribution and with probability 1 1- d for Boolean distribution for some constant E (0.
Proof of Lemma 6.2. (1) By preimage theorem (Banyaga & Hurtubise, 2013), it suffices to check the (-2, ( Gul
jacobian [V (z), [x) =
) is full rank. Similarly, for the second
claim, due to (35). it is also equivalent to show that { (l))..)) NEjn is of rank n.
Since (C) E c U, each coordinate is non-zero, thus we only need to show that {zi)ie(n) of rank n. This happens with probability 1 in the Gaussian case, and probability at least 1 c for some constant (0, 1) by Kahn et al. (1995). This completes the proof.
30

---

Published as a conference paper at ICLR 2022
D.3 PROOF OF LEMMA 6.3
We first establish some auxiliary results. The following lemma shows the PL condition along the trajectory of gradient flow.
Lemma D.2. Along the gradient flow generated by L. it holds that [VL(x(t)) 2
16 Amin (ZZT) minietd] (()())c(()))((((())))) O.
To prove Lemma D.2, we need the following invariance along the gradient flow.
Lemma D.3. Along the gradient flow generated by LV, ny (1)u((1) stays constant for all j E [d] Thus, sign( (t)) sign(u,(0) and sign(e((t) sign(e) (0)) any j [d].
Proof of Lemma D.3.
01 (t){()((()) = (ty(t)
dl (t(())+ (t) 01
((((()))((t))))-))) (t) V.L(u(t),
2 " Σ (I((u((), - u, (,,, (t,(t) "
(((((())0))) - (55
n
0.
Therefore, any sign change of uy (t), vy (t) would enforce (t) 0 or (t) for some 0 since u, (t), (t) are continuous in time t. This immediately leads to a contradiction to the invariance of uu (t),
We then can prove Lemma D.2.
Proof of Lemma D.2. Note that
|V n²
Σ (0((())) - v5) (V)()), vfz()))
i,1=1
(fi(z) 21) Amin (K(x))
2 [xx)) Amin (oninK((a (Kx)),
where K(x) is p.s.d. matrix with Kiy(x) (ff()) Below we lower bound
Amin Amin(K (z)), the smallest eigenvalue of K(x). Note that K (x())) = 4E (L,33,,((u (vx (t)))), and we have
((x(1)) = 4Zdiag (n(1) 02 +
2T > SZdiag(Iu(t) (()))
8Zdiag( u (0) O (0]))
min Iu: (0))))))))
where (+) is by Lemma D.3. Thus Amin (K (x(t)) > 8 min l,, (0),, Amin (ZZT) for all 2 which completes the proof.
We also need the following characterization of the manifold r.
Lemma D.4. All the stationary points in U are global minimizers, i.e., = for E U VL(x) = 0).
Proof of Lemma D.4. Since I' the of local minimizers, each x I' must satisfy VI(x) = The other direction is proved by noting that rank( ((=)(((())))))) (=)))()) = which implies rank((V (A)((()) n.
Now, ready prove Lemma 6.3 which is restated below.
31

---

Published as a conference paper at ICLR 2022
Lemma 6.3. Consider the loss function L defined in (14), manifold I and its open neighborhood defined in (15). For gradient flow dæs = (x) starting at any To E U, it holds that $(xo) E T.
Proof of Lemma 6.3. It suffices to prove gradient flow da(t)
VL(x(t)) converges when t - oc,
as long as z(0) E U. Whenever it converges, it must converge to stationary point in U. The proof will be completed by noting that all stationary point of L in U belongs to I' (Lemma D.4).
Below we prove lim sr(1) exists.
Denote C =
16 n
mingee4) (O)()))( (ZZT), then
follows from Lemma D.2 that
da(t) dt
(VL((()) dL(a(t)) VCL(z(t) VL(a(t))
1 dvL((())
VV(((())
2VC dt
Thus the total GF trajectory length is bounded by S d=(()) di Se 2V0 ddt
dvL(w(()) dt S
2/C where the last inequality uses that L is non-negative over RD. Therefore, the GF must converge.
M(e(0))
D.4 PROOF OF RESULTS IN SECTION 6.2
Without loss of generality. we will assume Σ 2 > 0 for all j [d], because otherwise we can just delete the unused coordinate, since there won't 't be any update the parameter corresponding to that coordinate. Moreover, in both gaussian and boolean setting, it can be shown that with probability E > for all j E [d].
To study the optimal solution to (18), we consider the corresponding d-dimensional convex program terms of w Rdd which has been studied in Tropp (2015):
minimize R(w) =
lweyl.
(36)
j=1
subject to
7w 2wn
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 canonical parametrization defined as follows.
Definition D.5 (Canonical Parametrization). For any wo E Rd, we define (*) = v/w) = ([w) T/01/2
++
[-w / as the canonical parametrization of w. Clearly, it holds that 00 10 = w.
Indeed, we can show that if (36) has a unique optimal solution, 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 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 (ū², ū+) = (ww"), ), that is, the canonical parametrization of w².
Proof of Lemma D.6. Let (û, i) be any optimal solution of (18) and we define v 002 which is also feasible to (36). By the optimality of w", we have
682
|w
(37)
1=1
On the other hand, (ū", ū+) = vww") is feasible to (18). Thus, it follows from the optimality of (û, 0) that
12
i
2
(01))
+ =
Pij
(38)
j=1
==1
32

---

Published as a conference paper at ICLR 2022
Combining (37) and (38) yields
Σ
(a) +
102
031
(39)
i.j
which implies that 1 102 is also an optimal solution of (36). Since w" is the unique optimal solution (36), we have t02 0o2 w". Moreover, by (39), we must have t [w=]+ and
Q02 [w0]+ +. 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 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 ; In iid. Uni There exist some constants C, c > such that if the sample size satisfies
n 2 C(KIn(d/xn) +
then with probability at least e the optimal solution of (18), (á, 0). is unique up to sign flips of each coordinate and recovers the groundtruth, ie., 102 602 w".
Proof of Lemma D.7. By the assumption that 2,,.... 2n ~ for all j E [d]. Then (36) is equivalent to the following optimization problem:
lid. Unif({# we have E 2ig
minimize
g(w) = ||w||1,
(40)
subject to Zw =
-10)0)))
This model exactly fits the Example 6.2 in Tropp (2015) with 1 and 1/v2. Then applying Equation (4.2) and Theorem 6.3 in Tropp (2015), (40) has a unique optimal solution equal to (u") (u²)o2 (u*)02 with probability at least 1 - e ch for some constant e > 0, given that the sample size satisfies
C(K ([IIn((/n) K
" and then adjusting the choices of C, c for some absolute constant C > 0. Choosing h = 2C
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 such that if the sample size satisfies
iid
~ N(0, Ta).
There exist some constants C,
Z Ca
then with probability least 1 - sign flips of each coordinate of i and è and recovers the groundtruth, i.e., 002 002 w².
(2d 1)e the optimal solution of (18), (ū, 6), unique up to
Zn Lid. N(, Ia), we have
Proof of Lemma D.8. Since 21,...,
n
IP
Σ
+ E 1n/2.3n/21 E [d]
2 2dle on
1=1
for constant 0, and denote this event by En. Therefore, En, we have
2
Σ (u 03) = R(x) < 6 (u 0z)
or equivalently,
2(1102211 + 0100211111 < R(x) < 6(1 02 + 2111))
33

---

Published as a conference paper at ICLR 2022
Define w" (u*)02 (v*)o? and (36) is equivalent to the following convex optimization problem
minimize g(w) 4
|wj + wil.
Tt O.
(41)
subject to Zw
The point = 0 is feasible for (41), and we claim that this is the unique optimal solution when is large enough. detail, assume that there exists a non-zero feasible point w for (41) the descent cone (Tropp, 2015) Dig, w*) of 9, then
Zwwlz
Amin (Z; D(g w <
|
where the equality follows from that is feasible. Therefore, we only need show that Amin (Z; D(g, x²") is bounded from below for sufficiently large n.
On En, it holds that g belongs to the following function class
{ - +R (R]) = Σ Eyywyyloe B with T= (UEERRR E2. [d].
G -- h::
h: Rd
[2.6).
We identify Do E G with u E Y, then Dig, w*) C UnET UnET w²)) = Dr, which further implies that
Amiu (2 Dgg w*)) > Amiu (min(Z; (T)
Recall the definition of minimum conic singular value (Tropp, 2015):
Amin (Z; (r)r)>
inf PED
sup (g. Zp). gESn-1
where 5 denotes the unit sphere IR. Applying the same argument as in Tropp (2015) yields
P [miu (Z; 2 Vn - w (Dr) h] 2 1 - -
h2/2
Take the intersection of this event with En, and we obtain from union bound that
(Z; w 2 vn - (DDr) h with probability at least -1 2de en remains to determine w: (Dr), which is defined as
(42)
w((r) Ess.((0Laa)
sup PeDr
(z
BenN((OLaa)
sup sup TET p Inse
(2.2))
(43)
Without loss of generality, we assume that w* (w (wi...., w O, ..... 0) T
with w wi...., w
otherwise one only needs to specify the signs and the nonzero set of wt in the sequel. For any E T and any p E D(go: w") S there exists some
such that go + p.) S go go(w") (w").ie.. i...
Σ ==1
Ug my +
TPsl S
1=1
which further implies that
K Σ
uslpsl
[wj TP3)) S T
ug|p,l
j=x+1
j=1
where the second inequality follows from the triangle inequality. Then since each Uy E [2, (2,6,.1 it follows that
d
lpsl
lpy-
j=0+1
j=1
34

---

Published as a conference paper at ICLR 2022
Note that this holds for all E simultaneously. Now let us denote P1:w (pı,..... Pr.) E R" and P(i++ Ijd = (Px+1..... Pa) E Rd-", and similarly for other d-dimensional vectors. Then for all E Dr n SSd-1, by Mauchy-Schwarte inequality, we have
IP S 31 S 33
Thus, for any 2 E Rd and any p E Dr n Sd-1 follows that
(z.
(z Plan) + (F(m+1):d+ P(ce1):a))
l la pllppssll [P(w+1):alll
max
[z
jE( I..nd
$ I la Pi 3/
max jE(x+1.....d
1z
S | + 3/R
max je (x.......
|z
where the last inequality follows from the fact that E Sd-1 with (43), and obtain that
Therefore, combine the above inequality
w (T) & E
||21:0112 + 3/R
max JE(x+I.....d!
lzyl
$ V 3
E
max
|z
(44)
JE(x+I......
where the second inequality follows from the fact that E 1 [12] S v the second term in (44), applying Lemma D.9, it follows from (44) that
VK. To bound
((Dr) I V 3 In(2(d x)).
(45)
Therefore, combining (45) and (42), we obtain
Amin (Z; \ w Vn- -
k)) h.
Therefore, choosing h /n 1/2, as long as n satisfies that 2 C(K In d) for some constant
C> 0, we have Amin (Z; D(g. w" )) > 0 with probability at least 1 (2d + 1)e-on, Finally, the
uniqueness of the optimal solution to (18) in this case follows from Lemma D.6.
L
Lemma D.9. Let NO, Ia), then it holds that [max;e[4] [=]]] S (2ln(
Proof of Lemma D.9. Denote M max4el4 [2g].
For any À 0, by Jensen's inequality, we have
d
B.B
< E
AM
max Masl
REld
Note that E[e- < 2- Ele Thus, by the expression of the Gaussian moment generating function, we further have
d
A.B[M] <
> E [e = 22/2
1=1
from which it follows that
In (nd)
M
<
+
Choosing À (21n(2d) yields the desired result.
35

---

Published as a conference paper at ICLR 2022
D.5 PROOF OF LEMMA 6.5
Instead of studying the convergence of the Riemannian gradient flow directly, it more convenient consider it in the ambient space RD_ To do so, define Lagrange function L(x; R(x:) + Ei= (fa(t) - 9a) for À E IR". Based on this Lagrangian, we can continuously extend (2(x) VR(x) to the whole space RD In specific, we can find a continuous function F' IR - R such that F(-) [r) = (-VV)((..). Such an F' can be implicitly constructed via the following lemma.
Lemma D.10. The bg norm has unique minimizer among <V A) A E Rd) ) for any fixed I E RD Thus we can define F; R2D + RRD by F(x) = argmin geV. it holds that (F(x), {((x)) = 0 for all i E [n].
1 Moreover,
Proof of Lemma D. Fix any x E by (V (i))) shifted by VR(z), thus there is unique minimizer of the & norm in this set. This implies that F(z) (V
R.D. Note that {V,[(x; X) e
] is the subspace spanned
lllll2 is well-defined.
To show the second claim, denote h(A) Then we have
|V 1) 2, which is quadratic function of À E R".
(VR()), V((()))
A V(((T))
(TR()),VSi(H
VI(A)=
+
+(()),.
Vf((t)
EL (M) V((())))
(VR(z). Vf)())
For any À such that v A) F(x),, we must have Vh(A) by the definition of F(x). which by the above implies
(K())A); = (VR()), for all i [n].
Therefore, we further have
(F(x), Vf(()) = (R(c). + Σ (VVI())) (f()) (K())A), 0
for all E [n] This finishes the proof.
Hence, with any initialization z(0) E I, the limiting flow (17) is equivalent to the following dynamics
dx(()
1 F(x(t))
(46)
dt
Thus Lemma 6.5 can be proved by showing that the above x(t) converges to x² as t - co. We first present series of auxiliary results in below.
Lemma D.11 (Implications for Fxx) 0). F + R be as defined in Lemma D.10. For
any RD such that F F(x) 0, it holds that for each j E [d], either uj 0 0 = 0.
Proof. Since F(x)= = 0. it holds for all j E [d] that,
Of (x) 2uj
4
0 =
Ouj
(x)
(x(x)
ou,
%
+
X
(x(
1=1
t=1
aR
0fi (x) = 20g
0 (x)
(x
A1(2)215
Duy
Ou
1=1
If there exists some j E [d] such that uj * 0 and U/ * O, then it follows from the above two identities that
2
which happens with probability 0 in both the Boolean and Gaussian case. Therefore, we must have My = 0 or uy = 0 for all j (d].
n
36

---

Published as a conference paper at ICLR 2022
Lemma D.12. Let F F - RD be as defined in Lemma D.10. Then F is continuous on R. R.
Proof. Case We first consider the simpler case any fixed x E U U (R (0)) assuming that (x*" is full-rank. Lemma D.10 implies that for any À E IR" such that vaL(x";;) = F(xx), we have
[(x") ) = (v)( V[n(t)
Thus such À is unique and given by
=- -
VR(x). VR(z))
Since (xx) is continuous around x=", there exists sufficiently small 5 > such that for any E (x(x"), (xx is full-rank, which further implies that K(x)-1 1 is also continuous Bz(x).. Therefore, by the above characterization of X, we see that (ax) is continuous for x E Bs(x"), and so F(x) V R(x) + Σ (i))
Case Next, we consider all general x" E RD. Here for simplicity, we reorder the coordinates as (u, U1, U2 .2..... ug. va) with slight abuse of notation. Without loss of generality, fix any x" such that for some y E [d], (u)) + (v:()))²
for all i I,...,4 . and un = og for all i = + 1.....d. d. Then V V(x")) and (Vii(x") (EEn only depend on (Zt JE()) and for all
i [m]. it holds that
V (VR(x)))
(24+1)): (VR(2))
(2)+)))) =.
Note that if we replace (V.i 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 E R such that, for some q' [q]. (2) .1...... 2n) (z1,...; Z satisfies that (ār, is linearly
independent and 1g = 0 for all i = d' + I.. We then consider Vii Vin(a))
[V) (x)) and the corresponding F(x). For notational simplicity, we assume that Q can be chosen as the identity matrix, so that (Z1...., 2n) itself satisfies the above property, and we repeat it here for clarity
12 New! linearly independent and 21,1. = 0 for all = g + 1....,n.
(47)
This further implies that
(V (V((())) = 0, for E (d' + I....., n and E RR..
(48)
In the sequel, we use X for n-dimensional vectors and for q' o'ddimnniional vectors. Denote
( E argmin
VR(x)= + X \
ACR
X(z) E argmin AER
VR(x) + Σ (V(((())
1=1
1:(2g) 2
Then due (47) and (48), we have
( (R(x")+ + IMOO)A
n
VR(x") + Σ (c)\Vf((z')
=II 122
1:(2g)
(49)
"W do not care about the specific choice of X(x) or A(x) when there are multiple candidates, and we only need their properties according I Lemma 1.10, they can he arhitrary. Also, the minimum of fg-norm of an affine space can always be attained so argmin exists.
37

---

Published as a conference paper at ICLR 2022
On the other hand, for any x E RR, by (48), we have
o' (R(T) + Σ (M)VV(()
(VR(E)+ + Σ EM(x)((ie))
"
min XE
1:(24)
t=1
1:(2g) 2
$ VR(x) + Σ (x(t)V
[1:(2)
(x)]12
1=1
1:(20)
< 11((2)2
VR(x) + \x(x")Vfi(x)
(50)
t=1
where the first and third inequalities follow from the definition of F(r). Let - x", by the continuity of VR(x) and (V (iE))) we have
lim x--2"
v((x) ((
VR(=") Σ [(<")V(((=")
(51)
Denote (xx) = (R.NEE)) ((V(()))))))) VA( By applying the
same argument as in Case I, since K(x") is full-rank, it also holds that lim (xx) X(x"), and thus
(VILO)+ V
- - (VRCO) Σ (OOWHOO)
lim
(52)
1:(24)
Combing (49), (50), (51) and (52) yields
lim (2)2)))))122 lim min x- AER
VR(x)+ + Σ X, Aifi(e))
=IF((*)) )ll2
(53)
1:(2g)
Moreover, since [|((24t))) = /IF - (x))) we also have
lim ()))))))))))ll> =0. 0.
(54)
It then remains to show that lim (x)) (x")))("), which directly follows from lim Aisg (x) )1= (x*) A(x*).
Now, for any e O, due the convergence of Xxx) and that K > can pick sufficiently small on such that for some constant a > 0 and all x e Bs, (x"), it holds that X(x)) )ll - c/2 and
(( (VRlo)) + Σ \V(((z)
>
VR(x) + Σ ((O))()))))))
+ +all - - A(x)
1:(24)
1:(24)
(55)
for all Ā E RP, where the inequality follows from the strong convexity. Mcanwhile, due to (48), we
have
lim
VR(z) + Σ {())V
lim
VR(x) + Σ M())V((())
1:(24)
1:(24)
(VR(C)) Ax X(('")(((") 11(20)
(R(x)) + {(=)V\((0)
lim
1:(29)
38

---

Published as a conference paper at ICLR 2022
where the second equality follows from (53) and the second equality is due to (52). Therefore, can pick a sufficiently small 82 such that
(R(()) + \())VVi(())
1 (R(a) + Σ {x)VV(()))
+
ac²
4 (56)
1:(29)
1:(24)
2
2
for all x (x²) Setting 8 min(61,62) it follows from (55) and (56) that
P (a) X(x) - 2' for all x Bs(x*).
Recall that we already have [²(xx)) - X(x*)1 )1 S c/2, and thus
M| (x) A(x" )(g") la = | (z) - X(x*)) - (x)) - M(x))l2 + [((xx) - X(2")12 Sc
for all x E Bs(x*). Therefore, see that lim (x) )1'')
Finally, it follows from the triangle inequality that
[F() - F(x") F(*')l2 $
+ (((((t)):)()))2 + [[(eot)) lll (F(x))F((3)) 1:(20)
o' VR(x) Σ \(=VV\(()) VR(*)) Σ \xx))
|| ((eet)))D())l
1:(24)
{())V\(()) ((')V\((>)
+[V[(()) v VR(x')
||2 ||
where, as x + xx, the first term vanishes by the convergence of A (x) and the continuity of each fx)),, the second term converges to by the continuity of V VR(z) and the third term vanishes by (54). Therefore, we conclude that
lim F(x) = F(x*).
2-22
that is, F is continuous.
Lemma D.13. For any initialization x² E I, the Riemmanian Gradient Flow (17) (or equivalently, (46)) is defined on (o).
Proof of Lemma D.13. Let [0, (' be the right maximal interval of existence of the solution of
Riemannian gradient glow and suppose T + oc. Since R(x(t) is monotone decreasing, thus R(x(t)) is upper bounded by R(z(0)) and therefore [|VR(z(t))] is also upper bounded. Since dz(()
[[V](z(()]l2 for any I < T, the left limit <(T-) = lim, =(T) must exist. By Corollary 1, Perko (2001), xTT-) belongs to boundary of
U, i.e. (T-T-) 0 or (T = 0 for some j [d] by Lemma 1. By the definition of the Riemannian gradient flow in (17), have
d
dx(()
(1,());(0)) = (t)( v,(()e])
dt
dt - (15(t)e uy(t)e]) ((x()).
By the expression of F(x(()) = vi(x(()) + En (((())V\<((()), we then have
d
2
(x(())) u,(1)u;(1) 2
de (m,(0,(())
(x(((()) us(0ss(1)
4
(t)
11
Denote 8, 4 E 2 It follows that (t,(t)e (0) = (0))0)) for all f E 10, T). Taking the limit we have |u,(T (T-(T-u (T-)u,(T-)) 2 (n)(0), > 0. Contradiction with 7 oo! n
39

---

Published as a conference paper at ICLR 2022
Before showing that F satisfies the PL condition, we need the following two intermediate results. Given two points u and v in IR,, we say u weakly dominate v (written as u S u) if and only if u, I U2, 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 E A such that u S v.
Lemma D.14. For some q E [D], let S be any q-dimensional subspace Rf and P fu E RD u, 0, Vi E [D]}. Let u, be an arbitrary point in P and Q P n (u. + S). Then there exists radius r > O, such that B,(0) (0) n Q weakly dominates Q, where (0) is the En-norm ball of radius r centered at O.
As a direct implication, for any continuous function f : P - R, which is coordinate- non- decreasing, minseu (x) can always be achieved.
Proof of Lemma 1414. 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 holds for D. For each E [D], we apply the proposition with D 1 Q n {u e PI u, = 0} (which can be seen as a subset of D- 1--1, and let r3 be the corresponding ( radius. Set r = max Ti and we show that choosing the radius to be r suffices.
For any v E Q, we take random direction in S, denoted by w. If 2 0 or w S 0, we denote by the first intersection (i.e., choosing the smallest A) between the line {u - ا|ااا Maol}>a A220 and the boundary of U. i.e., U.2 {2 RD 24 0}. Clearly y S v. By the induction hypothesis, there exists a E (0) n Q such that u S y. Thus S and meets our requirement.
w has different signs across its coordinates, we take /1, 32 to be the first intersections of the line {u - - and the boundary of U in directions of X >> 0 and A < O, respectively. Again by the induction hypothesis, there exist U1 U2 E (0) n Q such that d1 S y11 and u2 S 1/2 Since lies in the line connecting #1 and u2, there exists some h E [0, 1] such that
(1 h)w + huz. then follows that (1 h)un- + hvu S (1 h)y: hy2 = v. Now since Q is convex, we have
(1 - -h)u + hug E Q. and by the triangle inequality it also holds that [[( h)mu + huzll| S so (1- (1-h)u1 + huz E (0) n Q. Therefore, we conclude that (0) n nQ weakly dominates Q. and thus the proposition holds for D. This completes the proof by induction.
Lemma D.15. For some q E [D], let S be any g-dimensional subspace of R. and P= fu E RD ui 2 0, Vi E [D]]. Let u, be an arbitrary point in P and Q = P n (u. + S). Then there exists a constant e E (0, 1] such that for any sufficiently small radius r > 0, c Q weakly dominates (n(4++ + S + (0)). where (,) is the t--norm ball of radius 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 5 = {0} or S = R. S = R is straight-forward; for the case S = {0), we just need to ensure clu,|= |u.| r, and it suffices to pick
|u.,] and 0.5.
Suppose the proposition holds for D - 1. below we show it holds for D. For each i E [D]. we first consider the intersection between P (u. + S + (O)) and H, == {u E RD u = 0). Let u; be an arbitrary point in Pn (u, + 5) n H;, then P (v. + S) n H, = Pn (u, S) H: Pn (u + S H.). Furthermore, there exists {as he which only depends on S and satisties C (u + S. B (0) n.). Applying the induction hypothesis to Pn 5 H, + (0) H,), we know there exists a e > 0 such that for sufficiently small r, = c(Pn (w; + H.)) weakly dominates Pn + SSHH + (0) H.).
Pn(u" + S (B())
(1)
For any point v in Q and any = E (2),), we take a random direction in S, denoted by w. w 2 0 orw S 0. we denote by y the first intersection between {u+2 + z [Ww (w]222 |220 and the boundary of U. Clearly y S v. Since y E P (u, + S + B² (0)) n H; C P (u4 + S n H4 + B2 B²,( (0) n H.). by the induction hypothesis. there exists a c( (P (u, 8)n n H,) such that and z E c(P O (u,- + 5)) c. Q.
y. Thus 2 S 0 +
If wf has different signs across its coordinates, take yı, y/2 to be the first intersections of the line (u + A|w]>EeR and the boundary of U in directions of x > 0 and X < 0, respectively. By the induction hypothesis, there exist U1 U2 Q such that un S yn and u2 S y2- Since z lies
40

---

Published as a conference paper at ICLR 2022
in the line connecting u1 and uzg there exists some h E [0. 1] such that + (1 h)y1 + hy2-
then follows that (1 - h)u1 + hug I (1 - h)yı + hyz = u + 2. Since Q is convex, we have conclude that eQNQ weakly dominates (n + B (0))
(1- h)u, + hug E cQ. Therefore, for all sufficiently small r, and thus the proposition holds for D. This completes the proof by induction.
L
Lemma D.16.
condition for F.) For any x² such that L(x*)= = 0, i.., x T.
there exist neighhourhood U of x² and constant O, such that [IF(z))]1 2 c. (ax)
R(x*) ). for all U' nI. Note this requirement is only non-trivial when [| )||2 = 0 since F
is continuous.
Proof of Lemma D.1 It suffices to show the PL condition for {- F(x) = 0}. We need to show
for any x² satisfying F(x*) = 0, there exist some e > and C. > 0, such that for all æ E In with R(x)= > -R(*)), ), it holds that [1=())) 2 C(R(x) R(x*)).
Case L. We first prove the case where r = (#) itself is canonical parametrization of w = 20 - 002, i.e., Ug = 0 for all j E [d]. Since x" satisfies VF(x*) = 0, by Lemma D.11, we have xx = v(w") where w = )o² (v") )o2 In this case, we can rewrite both R and F as functions of w E Rd. In detail, we define R')w)= R(e(w)) and F'(w) (w)( for all E Rd For any w in sufficiently small neighbourhood of w", it holds that (w;) sign(w)) ) for all E [g]. Below we show that for each possible sign pattern of w(g+1)ad 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- 2 O, for all j = q + 1...,d.
We temporarily reorder the coordinates as x (u1; U1, U2, u.,..., ud, va) (t,,...,En] =n] is a n-by-d matrix, and we have
Recall that Z
[F'(w)
min {(a - sign(w) 3 Z A) |w|) XER"
where a a ET 202 E Rd. Since F(x*)= = 0, there must exist X E K", such that the first 2q E (x(x)) are equal to O. As argued in the proof of Lemma D.12.
coordinates of V +++x)+++
we can assume the first q' rows of Z are linear independent on the first g coordinates for some o' E [4]-
In other words, Z can be written as AA ZB
where ZA E Rd' xg, We further denote Xa == A
ZD
AB= == +1)m+ Wa = W1:g and wa = "'(++1jdd for convenience, then we have
[w)'( = min ((a: + sign(ws o kxAa) 0² (wal) + ((22 + ZiAa + Ibb)) XER"
02 wb) (57)
Since every w in I is global minimizer, R'(w) M( - gw (w..).
where g = sign(w) (w) @ Z X* Similarly define 9u == 91.y and gb I(g+1)dd It holds that Ia = 0 and we assume D 0 without loss of generality, because this can always be done by picking suitable XE for i q' + I,..., n. (We have such freedom on X² because they doesn't 't affect the first 24 coordinates.)
We denote
Ao Xx by Ada, then since 0- 9u = sign(w..) @ + Ax we further have
((a1 + sign (w:)) + Z )a)
02 (wal) = ((a1 + sign(w,) 2 _ sign(wa) @ ZJ AAA 02 (wal) = ((sign(w.) @ 2 AA, )2 Wal
On the other hand, we have ys sign(ws) as + ZIA -D>; = 02 2B +Z>2 by the assumption that each coordinate of wp is non-negative. Combining this with the above identity, we can rewrite Equation (57) as:
[w))1 = min ((Z]A)a) 02, (wal) + ((96 + 2 AA, + 2DAb) 02 wn)
(58)
XEED
Now suppose R'(w)- - R' (www)== g wa = for some sufficiently small 8 (which can be controlled by e). We will proceed in the following two cases separately.
41

---

Published as a conference paper at ICLR 2022
Case 1.1:
| = 2((5). Since ZA has full row rank, ||
Aa 12 min (Za) is Jower-bounded On the other hand, choose small enough such that E (wl ) 2 2 (w:)?. Thus the first term of Equation (58) is lower bounded by || 112 A (Z^) mmnnse(]] 2 (wa)) = (266 = ((( R² (w()))
Case L.2: | = 0(v)). Let = 96- Aru
then we have a E S+B2 (0) <v3
for some constant c: > 0, where S {oh + 7 A λ E ER" }. By Lemma D.14, there exists some constant 2 I, such that 1 S weakly dominates S B² ovo (0). Thus we have FI (w)ll2 2 inf inf unfuetBBuu000 BLuz(0) (u02, wb) 2 inf inf L. (502, (s)), where the last step is because each coordinate of Wp is non-negative.
Let A be the orthogonal complement of (pan(ZD; 9b), (b),,,................) i.e. the spanned space of columns of Zn and yb. we know Wb E |I 9b + A, since Znws Znw? 0 and 9b wh. 6. Therefore,
1F'(2))1 w:R' w: (ww)-R(w"" =000 R (w))) RR' (w*)
inf
2 inf
inf
ao wb
(w)- (==)====000-
8
inf
0u wwbb (59)
m | 2
A
A,,, 20.
Note (u02, wp) a monotone non-decreasing function in the first joint orthant, i.e.. {(u, wh)) E Rd x Rd-D' achieved by some finite (u, wp) in the joint first orthant. Applying the same argument to each other orthant of u E Rdd we conclude that the right-hand-side of (59) can be achieved.
2 0, wh 2 0}. thus by Lemma D.15 the infinimum can
On the other hand, we have T wb 8 0 for all w/ E 11 96 + A and u E 5, by D99 0 and the definition of A. This implies there exists least one [d x'] such that 24 > O, which further implies (u02, wb) 0. Therefore, we conclude that [[F'(w)]12 = 22 R (w) R' (wo)).
Case II Next, for any general x (.), we define w 0002 02 and m min(u02 022
where min is taken coordinate-wer wise. Then we can rewrite [[(()) 12 as
a
7
FF))]1 = min
+
A
XeR
2
[a]
Z
w
min
+
A
7
min x
[a]
7 7
02
(2))02
m7
-
Z
02
Z
62
m
min A
+
(2)))2 02
min
a
-z
Z
[a] a
Z
a
Z 2
vm
min x
+
7 2
3
(w))
min
+
AeR
42

---

Published as a conference paper at ICLR 2022
Then applying the result for the previous case yields the following for some constant C (0, 1):
[|F())]1 2 C(R(W(w)) - R(v(w"))) + min
a
Z
A)
m
XER"
+
Z
\m
C(R(((W)) - R(x*) + 02 m)
2 C(R(W(W)) R(x*) + min (a.m)
ield
=(((((()) - R(x*) + min ai (R(x)= - R()((w))) (E(d]
C, min sE|d]
}
min
(R(x) - R(x)))
where the first equality follows from the fact that xx' = v(w") ) and the last inequality is due to the fact that both R(b(w)- - [()(w))) and R(x) (()( 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., J. _e l2 < oo for any x² E T. Moreover, z(() converges to some x(20) when + os with (((())) =
Proof of Lemma D.17. Note that along the Riemannian gradient flow, R(x(t)) is non-increasing. thus |lx(t)]l, is bounded over time and (a(t))eo has at least one limit point, which we will call x*. Therefore, R(x*) is limit point of [[x(tt), and again since R(x(t)) is non-increasing, it follows that R(x(())= 2 I((=") and lim R(x(t)) = R(x*), ). Below we will show lim xx() = x".
dx(t) Note that dR(a(t)) (V)()(()), dt
((())))))(()))) = (vR(((O), ((((()) = - ([((((())1 where
dt
the last equality 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 limit point, there exists a time To, such that JT, E U. Let Ti = inf22 (x(()) & U'} (which is equal to xc if x(()) E U" for all 2 To). Since arl) is continuous in L and U is open, we know Ti To and for all E (To, Ti), we have [[F(x(t))ll 2 Vc(R(x(t)) R(=2))
Thus holds that for & E [To. Ti),
S d( d(R(a(t)) R(x*))
Ve (A(z(t)) ((la())- (F(())
dt
that is,
d(R(x(()) - R(x 1)12
Ve
S
|F(a(()))l
di
8
Therefore, we have
8 (R(x(Ta)) - R(x
(((x(t))222
(60)
vc
L=
Thus if we pick To such that R(x(To)) R(x*) is sufficiently small, R(Ti) will remain in U, which implies that Ti cannot be finite and has to be co. Therefore, Equation (60) shows that the trajectory of (t) is of finite length, so x(x))= ;= si xrt) exists and is equal to x²" As by-product, F F(x** must be
Finally, collecting all the above lemmas, we are able to prove Lemma 6.5. In Lemma 1.17 we already show the convergence of r(f) as t + co, the main part of the proof of Lemma 6.5 is to show the (0(o)) cannot be sub-optimal stationary points of It on I', the closure of r. The key idea here is that we can construct a different potential ø for each such sub-optimal stationary point x=, such that (1) ((z,) is locally increasing in a sufficiently neighborhood of x and (2) lim 6(x) = -0.
Lemma 6.5. Let {arh12 C RD be generated by the flow defined in (17) with any initialization
To E T. Then xo limp St exists. Moreover, Foe x² is the optimal solution of (18).
43

---

Published as a conference paper at ICLR 2022
Proof of Lemma 6.5. We will prove by contradiction. Suppose z((0) = (1(x)) lim x (t) is
not the optimal solution to (18). Denote w(t) - (u(f)) (v(t)) )2 then (x)) = lim w
not the optimal solution (36). Thus we have R(w(()) > R(w*). Without loss of generality, suppose there is some E [d] such that (u,(2)))² + (ox()))) > 0 for all = 1..... and
(1)) (20) = 0 for all q + 1....,,d d. Again, as argued in the proof of Lemma D.12, we can assume that, for some g' E [g].
(=) is linearly independent and 21,1:g 0 for all q' , n.
(61)
Since both (o)) and w" satisfy the constraint that Zw(oo)= = Zw" = -Y, we further have
0 (2,, (Do)) (za, w²) (fi,o+ w/ ()+1)))) for all g' 1,...,n.. n.
(62)
Consider a potential function yp U + IR defined as
<(x) = ((u, =
Σ j=o+1
w/ [In(u)) > 03 - 1Iu( < <0]]
Clearly lim ((x((()))))))))))))))))))) = -00 if lime- r(1) = r(oo). Below we will show contradiction =x(o)) is suboptimal. Consider the dynamics of ((x) along the Riemannian gradient flow:
dep
di (z())) V))(((t)), dr(t)
dt = 2
1
(63)
where F is defined previously in Lemma D.10. Recall the definition of F, and we have
(V((2())) Vo(x(()), 1 vR(2(() + 1 ((((0)VA(a(O))
I1
vo(x(()), 1 4 1=9'+1
Σ M((())V((((0))
(64)
To show ((((x(()),(<<(())) O, we analyze In and Iz separately. By the definition of ((x), have
d Ve(()= = Σ
[[w > 0} 1lu; < 0}
2w0
U ej
Vj ED+J
5=9+1
where ej is the j-th canonical base of Rd Recall that vf(((r) 2 (_2n and we further have
0}
1 1
I2 = Σ {a(x(())
Σ
w
(e 2. 2 u) +
(es, ze o
Uj
1=
1f
> 0}
1w < 0} زز
Σ {e(x((1)
Σ
w/
Zi +
Uj
Vij
(=9+1
1=9+1
xx(((t)
Σ
Wfz
W
Σ
((=(t))(z)(tt))) Wigttyal
(65)
i=q'+1
-=9+1
where the last equality follows from (62).
Next, we show that Iι < 0 by utilizing the fact that w" w(o) is a descent direction of R' R'(w). For E Idd, define f((w)= = z w and
i((w)= R(w) + (F((2)))( v1).
-
44

---

Published as a conference paper at ICLR 2022
Clearly, for any R satisfying Zw = Y, it holds that fi(w) y/2 0 for each i [n], and thus R(w)) = R( Rw) In particular, we have R(w(oo)) R(w(xx)) > R(w") R(w*). Since R(w) is convex function, it follows that [((w(()) + + -w(o))) S +R(w=*) + (1 (1--)))iii s) R(w(())
for all 0 < S 1, which implies dB at (wr(oo) + s(w" - w( [(oo)))l==0 < -2c < 0+ for some constant c > O. Note that, for small enough > O. we have
4
R(w((0)) + s(w= - w(oo)))
[w)(20) wo)))
π
j=1
4
Σ
2
(o0 + s(wj wo(o))))
d
+
s sWwl..
j=
Therefore, we can compute the derivative with respect to s sat aa
dR
-2c
(w(
+ s(w" - w(oo))
di
4
(o))( - -0,00))))) +
4
Σ
d
Σ
n
[wj]
π
n
imq+l
+
\c(x(00))2z (2)("))))([[-~---- - woy()))
(w-
4
Σ
n
(ogn(w)((o))(( w((c)) +
4
Σ
|woj]
Pin
n
j
4 Σ
d
(nj
wy (oo)) Σ \c(o(0))) +
Σ j=
w
)x(2((00)) 23,
(66)
1=1
where the second equality follows from the fact that ((g+1)aa( (Do) = 0. Since rt)) converges to (, we must have F(x(oo))= = O, which implies that for each j
E {1,.....h g.
0 = dR
(xx + )c(x((0)) 0f (x((x)) = 200(
4 n
\c(r(o0))2
Ouj
n i=1
im1
q' )((((0))) Of: (x(oo)) = 2005(()))))
aR 0 (z(oo)) +
4
\x(x(c=))zis
Ouj
Zug
n1 i=1
i=1
Combining the above two equalities yields
2
sign(w \1(2(00))2
for j E [q].
Apply the above identity together with (66), and we obtain
-2c
sign sign(w v((c)
(c((o0))=zs
4
Σ Σ
n
j=q+l
(iml
+
(uj
wo)(
Xe(2( +
Σ w Σ (o(o(00))
1=1
g=9+1
4
d
n
d Σ
|103| +
Wj
X (xc)))
(67)
n
j=
i=1
j=
i=1
45

---

Published as a conference paper at ICLR 2022
On the other hand, by directly evaluating VR(z(t)) and each V Vf((x(t)), we can compute In as
d I =
W) w
>0)
2 n
2 ut)
A (a(())=aa)(())
(((()
T1
-=o+1
d
w <0}
2 n
Σ
2 (t)
(=(()))))))(())
vt(()
n
2
j=
2 d
1
2 %
+
Σ
wj
A: (z)((()) ز.
j=
i=
j=4+1
2
%
10gl 2
Σ
wj
\e(3((o0))zz5
n
1=9+1
-=9+1
+
3
(A((x(()) - (x(o0))) z..-
j=++1
We already know that 2 (a) is continuous at z 2000) by the proof of Lemma D.12, so the third term converges to 0 as sz(() tends to ::oo). Now, applying (67), we immediately see that there exists some 5 > 0 such that Zi < -c for s(l) E (8(x)(00)) As we have shown in the above that Iz = 0, it then follows from (63) and (64) that
dep di
(z(t))
for all z(1) Bb(x(oo)).
(68)
Since lim (t) = xr( ). there exists some T > 0 such that st)) E (a(x((0) for all f. > T. By the proof ofLemma 1.13, we know that (T)( > -00, then follows from (68) that
lim ((x(()) = v(x(T)) + I dip(a(t)) di > p()(T)) + f odl = x t+0
dt
which is a contradiction. This finishes the proof.
D.6 PROOF OF THEOREM 6.7
Here we present the lower bound on the sample complexity of GD in the kernel regime.
Theorem 6.7. Assume 21,-.., Zn Ltd. NO0 Ia) and yi = =[]w' w², for all i E [n]. Define the loss with linearized model
= Ex= + {V)i( zo) y;) where To = (wa) = ax (+)) Then for any groundtruth w', any learning rate schedule {n1)121 and any fixed number of steps T, the expected Eg loss of (T)) is at least (1 4) || ww" 102 122 where x(T)) the T-th iterate of GD on L, i.e., ztt + 1) I(f) V (xx(t)), for all t 2 0.
(*) and
Proof of Theurem 6.7. We first simplify the loss function by substituting x' = =x- x0),, so corre- spondingly =/ 0 and we consider L' L'(x') == L(x) = ((VI(x(0)) x²) v.)². 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 single fixed target function, 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, means the expected loss of ((T)) with ground truth being w" is the same as the case where the ground truth is uniformly randomly sampled from all vectors of én-morm ||w
Thus the test loss of (<T') is
E, (=V), - (z, w²)) w)) E3 [((,, w² (u(I) (((T)=<(T)))) x(T))))] = I|w* (u(T) (l
(69)
46

---

Published as a conference paper at ICLR 2022
Note xTT) E (x) which is at most n-dimensional space spanned by the gradients
of model output at (0), so is u(T) - v(T). We denote the correspondin space for u(T) - v(T)
by S, so dim (S) < and it holds that || - (u(2) - =(T) - z l1 - Ps) ww" 11 where Ps is projection matrix onto space
The expected test loss is lower bounded by
(u(T) v(T))
Bw || - (u(T) - =(T)))
Ew-
Eg
I|20
Eg
Es. [I(4D Ps)c"'13
min
1 d
||
||=1 1
47


Reviews:
[{"review_id": "cGaXrXgrp3D", "reviewer": "Reviewer_Msxe", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First, I would like to thank the authors for the paper, the intuitive explanations, the clear writing and the proper referencing of the very nice Katzenberger article. \n\nBoth as a reviewer and a curious researcher I really enjoyed reading this article. I truly think it tackles a very nice question about the effect of the noise of SGD in minima selection. Remarkably, the paper introduces the very mathematical object that lacks in the Blanc et al. paper. Even if the limiting diffusion, in all its generality, is difficult to interpret, it shows properly some already known intuition in the case of SGD for machine learning problems with label noise. Finally, to give an explicit example where the limiting dynamics can be analysed thoroughly, the authors show optimality of label noise SGD for overparametrised linear network: what a nice result and application ! \n\nObviously, as one can read, I am very enthusiastic about this paper, but let me try to help the authors with a few notes:\n\nI would say that sometimes the writing is a little bit too direct and should be tempered a bit: the authors study the impact of the noise locally near $\\Gamma$ with an effective dynamics. To « trap » the dynamics in $\\Gamma$, the consider a regime for which the step-sizes go to zero (and along with it the effect of the noise itself). Hence I see this result as a limiting one: surely in practice, the dynamics is not exactly the one depicted: a first phase where the loss goes to zero and a second phase where the dynamics diffuses in $\\Gamma$. Of course, it does not diminish the nice impression that the paper gives me, but it should be clearer that to clearly separate the two phases the authors consider a limiting regime.\n\nEven if this work is already pretty complete:\n- I would like to read more interpretations of the limiting process of Eq. (10). What happens, e.g., for other neural networks than the one depicted in Section 6 ?\n- What happens if one does not add artificially label noise, e.g. with the model of Section 6 ? I expect that the analysis collapses because the noise itself is killed in $\\Gamma$. If not how is the dynamics changed ?\n- Can you give more intuition on the effective dynamics of equation (17) ? How are killed the directions for which the sparsest predictor is zero ? \n\nFor the modelling of machine learning-SGD with an SDE, authors should add the nice work of S.Wojtowytsch [ https://arxiv.org/abs/2106.02588 ]. Here, it is explained that the nature of SGD noise for machine learning problem is very degenerate. Another nice reference for your problem is the nice work of Pesme et al. [ https://arxiv.org/abs/2106.09524 ] on the implicit bias of the same model as the one studied in Section 6. It should also be pointed out that Lemma 6.3 is well known for these models: despite the fact that Woodworth et al., 2020; Azulay et al., 2021 assume convergence other works show it directly, e.g. in Pesme et al.. Hence I cannot consider it as a contribution (yet, it is in my opinion negligible w.r.t. the main message of the paper!).", "summary_of_the_paper": "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.", "summary_of_the_review": "As already expressed, I am very enthusiastic about the paper as they introduce the perfect mathematic tool and scale that characterizes nicely the effect of noise near the interpolation manifold. Furthermore they apply flawlessly and successfully their vision to a nice problem. In a word: Congrats ! ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Vq2zZKVbO3X", "reviewer": "Reviewer_TEtd", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros: The analysis using $\\eta^{-2}$ scaling to separate the fast and slow motion of SGD dynamics is novel, and this mathematical framework seems appropriate for analyzing the implicit regularization of SGD in the manifold of zero training loss.\n\nMajor comments:\n- Previous papers, e.g. Ghorbani et al. (2019) [Sect. 4.1], Gur-Ari et al (2018), Li et al. (2020), empirically show that there is significant alignment in top eigenspace between the Hessian and gradient covariance matrix using various deep neural network architectures and image datasets. If this alignment were true, I assume that Eq. (13) could describe the implicit regularization of SGD for general deep neural network. It would be very interesting if Eq. (13) can be used to explain the empirical observation of the spectrum of Hessian around minimizer found by SGD, i.e. a large bulk of nearly 0 eigenvalues with some large outliers. It seems this can be answered by a closer look of the right-hand-side of Eq. (13) and see how the gradient of $\\Phi$ interacts with the gradient of tr(Hessian).\n- The definition of the manifold $\\Gamma$ is a bit confusing. At first it is defined at the beginning of p.5 that looks general and is not specific to a model, but then in Theorem B.9 in the appendix, definition of $\\Gamma$ in Eq. (15) specializes on overparametrized linear models. Does Theorem 4.6 only apply to settings in Sect. 6? Clarification in appropriate places may be needed.\n- Given that the ground truth parameter $w^\\ast$ is $\\kappa$-sparse in Sect. 6, could some other types of regularization beside label noise better explore the manifold $\\Gamma$ and recover $w^\\ast$ with a lower sample complexity than $\\kappa\\log{d}$?  For example, if one is allowed to introduce a deterministic trend in the right-hand-side of Eq. (13) that can somehow lead to sparsity without leaving the manifold $\\Gamma$, e.g. directional pruning (https://arxiv.org/abs/2006.09358), would it help reduce the sample complexity?\n\nMinor comments:\n- Typo: “Polyak-ojasiewicz”\n- Fifth line below Eq. (6), no $n$ in $\\lim_{n\\rightarrow\\infty}{X\\left(t\\right)}$, and something seems to be wrong in that sentence\n- The $\\dagger$ is undefined in Definition 4.4 where it firstly appears\n- Cor 5.2, $\\Sigma=c\\nabla^2 L$, no trace\n- Are $\\sigma$ and $\\Sigma$ the same thing in Eq. (8)? I assume the operator $\\sigma$ in Sect. 4.1 is the same as the $\\Sigma$ in later sections, but it is best to unify them\n\nReferences:\n- Ghorbani, B., Krishnan S. and Xiao, Y. (2019). An Investigation into Neural Net Optimization via Hessian Eigenvalue Density. ICML\n- Guy Gur-Ari, Daniel A. Roberts and Ethan Dyer (2018). Gradient Descent Happens in a Tiny Subspace. Arxiv: 1812.04754\n- Xinyan Li, Qilong Gu, Yingxue Zhou, Tiancong Chen and Arindam Banerjee (2020). Hessian based analysis of SGD for Deep Nets: Dynamics and Generalization. Arxiv: 1907.10732.\n\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper provides a new mathematical framework for analyzing the implicit regularization of SGD after attaining zero training loss. The framework is novel, and it leads to a simpler analysis of sample complexity under noisy label regularization. I am in favor of accepting this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "HtY-NGnhGe", "reviewer": "Reviewer_uj7w", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength. The results are novel and the analysis is non-trivial. It introduces new ideas to the field by adapting the ideas from Katzenberger (1991). It allows the authors to obtain a global analysis of the implicit bias valid for $\\eta^{-2}$ steps, as well as any arbitrary noise covariance structure.\n\nWeakness. Some model assumptions, setups and the conditions in the main results might need more explanations.", "summary_of_the_paper": "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.", "summary_of_the_review": "(1) It would be nice if the authors can explain more and justify why you can use the SDE approximation in equation (2). $\\eta$ is the stepsize, and I can understand the stochastic modified equation in the literature, but somehow I don't see why SGD can be approximated by equation (2). \n\n(2) In the drift term in equation (3), should it be $\\nabla L$?\n\n(3) In the paragraph after equation (3), you wrote that hopefully we can simplify the dynamics Equation (3) via choosing suitable $\\Phi$. What I don't understand is that on the surface, it seems equation (3) is even simpler than equation (4) or (5). So what do you mean by simpler?\n\n(4) On page 5, in the paragraph before Theorem 4.1., what is $\\phi$?\n\n(5) I find the assumptions in Theorem 4.6 quite strong. You need to assume that the SDE in equation (10) has a strong solution. Ideally, the assumption should be on the loss function $L$ instead of the equation (10). Another strong solution is that $Y$ never leaves $\\Gamma$. This in my opinion is also quite strong. $Y(t)$ is an SDE with a tangent noise term, and because of the Brownian noise, I just don't see how this assumption can be satisfied. Even if you have degenerate noise in some direction, e.g. in the case of underdamped Langevin diffusion, it is still supported on the entire Euclidean space. If you have to assume such strong solutions for Theorem 4.6., it might be helpful to construct some toy examples to demonstrate that the limiting $Y$ is non-trivial and meaningful at least for some toy examples.\n\n(6) I also find Remark 4.7. a bit puzzling. In Remark 4.7., you mentioned that the convergence in distribution result in Theorem 4.6. (at the time $T$) implies the convergence for the sample path on $[0,T]$. I don't understand this. Shouldn't it be the case that the latter implies the former?", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "55ZoDgtbVS", "reviewer": "Reviewer_6sCC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper provides a very strong result for the implicit bias of SGD after it achieves zero training loss, i.e., capturing its behavior once it is in the manifold of the global minimizers. Based on Katzenberger (1991)’s result, they capture the global behavior, and their result covers the case that the stepsize $\\eta\\to 0$. In detail, via a projection operator, they kill the noise in the normal space and derive the limiting SDE via showing that the SGD will converge in distribution to this SDE. For each term of this SDE, they provide convincing explanation. Based on this result, they derive the limiting Flow for label noise. Furthermore, they study the overparameterized sparse linear regression model, showing that, in contrast to GD, SGD with label noise can recover the ground truth efficiently provided almost any initialization.\n\nOverall, this paper is very solid, and its technique, though is adapted from Katzenberger (1991), shall have its own value to the community to understand the behavior and property of SGD. Its application to demonstrate the power of label noise SGD is also very interesting. \n\nHowever, I have a few confusions about the label noise SGD. In the paragraph after the equation 12, the author mentioned that “Suppose the model can achieve the global minimum of the loss ...” Can you demonstrate that why the ground truth should be the global minimizer in the label noise setting? On the other hand, it seems that this analysis heavily relies on the label noise. Is it possible to extend this result to the mini-batch SGD?", "summary_of_the_paper": "This paper provide a global analysis of the implicit bias of SGD after it achieves zero training loss by adapting ideas from Katzenberger (1991). Moreover, based on this analysis, they show that label noise SGD can avoid the kernel regime, i.e., its sample complexity is significantly less than $O(d)$ in the overparameterized sparse linear regression model.", "summary_of_the_review": "This paper has very good theory and is well-structured. Also, its application to label noise SGD is also very interesting.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: All the reviewers agree that this paper made a solid contribution of understanding the algorithmic regularization of SGD noise  (in particular the label noise for regression) after reaching zero loss. The framework is novel and has the potential to extend to other settings.
Decision: Accept (Spotlight)