What Happens after SGD Reaches Zero Loss? --A Mathematical Framework

Introduction

The implicit bias underlies the generalization ability of machine learning models trained by stochastic gradient descent (SGD). But it still remains a mystery to mathematically characterize such bias. We study SGD in the following formulation

It is widely believed that large LR (or equivalently, small batch size) helps SGD find better minima. For instance, some previous works argued that large noise enables SGD to select a flatter attraction basin of the loss landscape which potentially benefits generalization (Li et al., 2019c; Jastrzebski et al., 2017). However, there is also experimental evidence (Li et al., 2020b) that small LR also has equally good implicit bias (albeit with higher training time), and that is the case studied here. Presumably low LR precludes SGD jumping between different basins since under general conditions this should require $\Omega(\exp(1/\eta))$ steps (Shi et al., 2020). In other words, there should be a mechanism to reach better generalization while staying within a single basin. For deterministic GD similar mechanisms have been demonstrated in simple cases (Soudry et al., 2018; Lyu & Li, 2019) and referred to as implicit bias of gradient descent. The current paper can be seen as study of implicit bias of Stochastic GD, which turns out to be quite different, mathematically.

The contribution of the current paper is a more general and global analysis of this type. We introduce a more powerful framework inspired by the classic paper (Katzenberger, 1991).

We start with an intuitive description of the implicit regularization effect described in Blanc et al. (2020). For simplification, we 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(t)$ is the standard $\Xi$-dimensional Brownian motion. The only property about label noise SGD we will use is that the noise covariance $\sigma\sigma^{\top}(x)=\nabla^{2}L(x)$ for every $x$ in the manifold $\Gamma$ (See derivation in Section 5).

However, the above approach only gives a local analysis for $O(\eta^{-1.6})$ time, where the total movement due to implicit regularization is $O(\eta^{2-1.6})=O(\eta^{0.4})$ and thus is negligible when $\eta\to 0$. In order to get a non-trivial limiting dynamics when $\eta\to 0$, a global analysis for $\Omega(\eta^{-2})$ steps is necessary and it cannot be done by Taylor expansion with a single reference point. Recent work by Damian et al. (2021) glues analyses of multiple local phases into a global guarantee that SGD finds a $(\epsilon,\gamma)$-stationary point for the regularized loss, but still doesn’t show convergence for trajectory when $\eta\to 0$ and cannot deal with general noise types, e.g., noise lying in the tangent space of the manifold. The main technical difficulty here is that it’s not clear how to separate the slow and fast dynamics in different spaces and how to only take limit for the slow dynamics, especially when shifting to a new reference point in the Taylor series calculation.

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 is of speed $\Theta(\eta^{2})$, we change the time scaling to accelerate (2) by $\eta^{-2}$ times, which yields

Note the first term $-\eta^{-1}\partial\Phi(X_{\eta})\nabla L(X_{\eta})$ is going to diverge to $\infty$ when $\eta\to 0$, so a natural choice for $\Phi$ is to kill the first term. Further note $-\partial\Phi(X)\nabla L(X)$ is indeed the directional derivative of $\Phi$ at $X$ towards $-\nabla L$, killing the first term becomes equivalent to making $\Phi$ invariant under Gradient Flow (GF) of $-\nabla L(X)$! Thus it suffices to take $\Phi(X)$ to be the limit of GF starting at $X$. (Formally defined in Section 3; see Lemma C.2 for a proof of $\partial\Phi(X)\nabla L(X)\equiv 0$.)

Also intuitively $X_{\eta}$ will be infinitely close to $\Gamma$, i.e., $d(X_{\eta}(t),\Gamma)\to 0$ for any $t>0$ as $\eta\to 0$, so we have $\Phi(X_{\eta})\approx X_{\eta}$. Thus we can rewrite the above equation as

and the solution of (4) shall converge to that of the following (in an intuitive sense):

The main contributions of this paper are summarized as follows.

In Section 4, we propose a mathematical framework to study the implicit bias of SGD with infinitesimal LR. Our main theorem (Theorem 4.6) gives the limiting diffusion of SGD with LR $\eta$ for $\Theta(\eta^{-2})$ steps as $\eta\to 0$ and allows any covariance structure.

In Section 5, we give limiting dynamics of SGD with isotropic noise and label noise.

Related Works

A phenomenon known as mode connectivity has been observed that local minimizers of the loss function of a neural network are connected by simple paths (Freeman & Bruna, 2016; Garipov et al., 2018; Draxler et al., 2018), especially for overparametrized models (Venturi et al., 2018; Liang et al., 2018; Nguyen et al., 2018; Nguyen, 2019). Later this phenomanon is explained under generic assumptions by Kuditipudi et al. (2019). Moreover, it has been proved that the local minimizers of an overparametrized network form a low-dimensional manifold (Cooper, 2018, 2020) which possibly has many components. Fehrman et al. (2020) proved the convergence rate of SGD to the manifold of local minimizers starting in a small neighborhood.

Implicit Bias in Overparametrized Models

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 is often more elegant and can provide fruitful insights like the linear scaling rule (Krizhevsky, 2014; Goyal et al., 2017) and the intrinsic learning rate (Li et al., 2020b). A recent work by Li et al. (2021) justifies such SDE approximation. Xie et al. (2020) gave a heuristic derivation explaining why SGD favors flat minima with SDE approximation. Wojtowytsch (2021) showed that the invariant distribution of the canonical SDE approximation of SGD will collapse to some manifold of minimizers and in particular, favors flat minima. By approximating SGD using a SDE with slightly modified covariance for the overparametrized linear model, Pesme et al. (2021) relates the strength of implicit regularization to training speed.

Notation and Preliminaries

Assume that $U$ is an open neighborhood of $\Gamma$ satisfying that gradient flow starting in $U$ converges to some point in $\Gamma$, i.e., $\forall x\in U$, $\Phi(x)\in\Gamma$. (Then $\Phi$ is $\mathcal{C}^{2}$ on $U$ by Falconer (1983).)

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 $\partial\Phi$ and $\partial^{2}\Phi$. We present our main result in Section 4.3. We remark that sometimes we omit the dependency on $t$ to make things clearer.

In particular, when the integrator sequence $\{A_{n}\}_{n\geq 1}$ increases infinitely fast, meaning that $\forall\epsilon>0,\inf_{t\geq 0}(A_{n}(t+\epsilon)-A_{n}(t))\to\infty$ as $n\to\infty$, we call (6) a Katzenberger process.

One difficulty for directly studying the limiting dynamics of $X_{n}(t)$ is that the point-wise limit as $n\to\infty$ become discontinuous at $t=0$ if $X(0)\notin\Gamma$. The reason is that clearly $\lim_{n\to\infty}X_{n}(0)=X(0)$, but for any $t>0$, since $\{A_{n}\}_{n\geq 1}$ increases infinitely fast, one can prove $\lim_{n\to\infty}X_{n}(t)\in\Gamma$! To circumvent this issue, we consider $Y_{n}(t)=X_{n}(t)-\phi(X(0),A_{n}(t))+\Phi(X(0))$. Then for each $n\geq 1$, we have $Y_{n}(0)=\Phi(X(0))$ and $\lim_{n\to\infty}Y_{n}(t)=\lim_{n\to\infty}X_{n}(t)$. Thus $Y_{n}(t)$ has the same limit on $(0,\infty)$ as $X_{n}(t)$, but the limit of the former is further continuous at $t=0$.

Suppose the loss $L$, manifold $\Gamma$ and neighborhood $U$ satisfies Assumptions 3.1 and 3.2. Let $\{X_{n}\}_{n\geq 1}$ be a sequence of Katzenberger process with $\{A_{n}\}_{n\geq 1},\{Z_{n}\}_{n\geq 1}$. Let $Y_{n}(t)=X_{n}(t)-\phi(X(0),A_{n}(t))+\Phi(X_{0})$. Under technical assumptions, it holds that if $(Y_{n},Z_{n})$ converges to some $(Y,W)$ in distribution, where $\{W(t)\}_{t\geq 0}$ is the standard Brownian motion, then $Y$ stays on $\Gamma$ and admits

Indeed, SGD (1) can be rewritten into a Katzenberger process as in the following lemma.

Let $\{\eta_{n}\}_{n=1}^{\infty}$ be any positive sequence with $\lim_{n\to\infty}\eta_{n}=0$, $A_{n}(t)=\eta_{n}\lfloor t/\eta^{2}_{n}\rfloor$, and $Z_{n}(t)=\eta_{n}\sum_{k=1}^{\lfloor t/\eta^{2}_{n}\rfloor}\sqrt{\Xi}(\mathds{1}_{\xi_{k}}-\frac{1}{\Xi}\mathds{1})$, where $\xi_{1},\xi_{2},\ldots\overset{\text{i.i.d.}}{\sim}\textrm{Unif}([\Xi])$. Then with the same initialization $X_{n}(0)=x_{\eta_{n}}(0)\equiv X(0)$, $X_{n}(k\eta_{n}^{2})$ defined by $\eqref{eq:katzenberger_process}$ is a Katzenberger process and is equal to $x_{\eta_{n}}(k)$ defined in (1) with LR equal to $\eta_{n}$ for all $k\geq 1$. Moreover, the counterpart of (7) is

where $\Sigma\equiv\sigma\sigma^{\top}$ and $\{W(t)\}_{t\geq 0}$ is a $\Xi$-dimensional standard Brownian motion.

However, there are two obstacles preventing us from directly applying Theorem 4.1 to SGD. First, the stochastic integral in (8) depends on the derivatives of $\Phi$, $\partial\Phi$ and $\partial_{ij}\Phi$, but Katzenberger (1991) did not give their dependency on loss $L$. To resolve this, we explicitly calculate the derivatives of $\Phi$ on $\Gamma$ in terms of the derivatives of $L$ in Section 4.2.

The second difficulty comes from the convergence of $(Y_{n},Z_{n})$ which we assume as granted for brevity in Theorem 4.1. In fact, the full version of Theorem 4.1 (see Theorem B.7) concerns the stopped version of $Y_{n}$ with respect to some compact $K\subset U$, i.e., $Y_{n}^{\mu_{n}(K)}(t)=Y_{n}(t\wedge\mu_{n}(K))$ where $\mu_{n}(K)$ is the stopping time of $Y_{n}$ leaving $K$. As noted in Katzenberger (1991), we need the convergence of $\mu_{n}(K)$ for $Y_{n}^{\mu_{n}(K)}$ to converge, which is a strong condition and difficult to prove in our cases. We circumvent this issue by proving Theorem B.9, a user-friendly interface for the original theorem in Katzenberger (1991), and it only requires the information about the limiting diffusion. Building upon these, we present our final result as Theorem 4.6.

2 Closed-Form expression of the limiting diffusion

We can calculate the derivatives of $\Phi$ by relating to those of $L$. Here the key observation is the invariance of $\Phi$ along the trajectory of GF. The proofs of this section are deferred into Appendix C.

For any $x\in\Gamma$, $\partial\Phi(x)$ is the orthogonal projection matrix onto tangent space $T_{x}(\Gamma)$.

To express the second-order derivatives compactly, we introduce the notion of Lyapunov operator.

3 Main Result

Now we are ready to present our main result. It’s a direct combination of Theorem B.9 and Lemma 4.5.

where $\Sigma\equiv\sigma\sigma^{\top}$ and $\Sigma_{\parallel},\Sigma_{\perp},\Sigma_{\perp,\parallel}$ are defined in Lemma 4.5.

In Section B.4 we indeed prove a stronger version of Theorem 4.6 that the sample paths of SGD converge in distribution, i.e., let $\widetilde{x}_{\eta}(t)=x_{\eta}(\lfloor t/\eta^{2}\rfloor)$, then $\widetilde{x}_{\eta}$ weakly converges to $Y$ on $[0,T]$. Moreover, we only assume the existence of a global solution for ease of presentation. As long as there exists a compact $K\subseteq\Gamma$ such that $Y$ stays in $K$ on $[0,T]$ with high probability, Theorem B.9 still provides the convergence of SGD iterates (stopped at the boundary of $K$) before time $T$ with high probability.

Implications and Examples

In this section, we derive the limiting dynamics for two notable noise types, where we fix the expected loss $L$ and the noise distribution, and only drive $\eta$ to 0. The proofs are deferred into Section C.3.

Isotropic noise means $\Sigma(x)\equiv I_{D}$ for any $x\in\Gamma$ (Shi et al., 2020). The following theorem shows that the limiting diffusion with isotropic noise can be viewed as a Brownian Motion plus Riemannian Gradient Flow with respect to the pseudo-determinant of $\nabla^{2}L$.

If $\Sigma\equiv I_{D}$ on $\Gamma$, SDE (10) is then

Type II: Label Noise.

In sharp contrast to the delicate discrete-time analysis in Blanc et al. (2020) and Damian et al. (2021), the following corollary recovers the same result but with much simpler analysis – taking derivatives is all you need. Under our framework, we no longer need to do Taylor expansion manually nor carefully control the infinitesimal variables of different orders together. It is also worth mentioning that our framework immediately gives a global analysis of $\Theta(\eta^{-2})$ steps for SGD, far beyond the local coupling analysis in previous works. In Section 6, we will see how such global analysis allows us to prove a concrete generalization upper bound in a non-convex problem, the overparametrized linear model (Woodworth et al., 2020; HaoChen et al., 2020).

If $\Sigma\equiv c\nabla^{2}L$ on $\Gamma$ for some constant $c>0$, SDE (10) can be simplified into (13) where the regularization is from the noise in the normal space.

Provable Generalization Benefit with Label Noise

In the setting of OLM, suppose the groundtruth is $\kappa$-sparse and $n\geq\Omega(\kappa\ln d)$ training data are sampled from either i.i.d. Gaussian or Boolean distribution. Then for any initialization $x_{\text{init}}$ (except a zero-measure set) and any $\epsilon>0$, there exist $\eta_{0},T>0$ such that for any $\eta<\eta_{0}$, OLM trained with label noise SGD (12) with LR equal to $\eta$ for $\lfloor T/\eta^{2}\rfloor$ steps returns an $\epsilon$-optimal solution, with probability of $1-e^{-\Omega(n)}$ over the randomness of the training dataset.

The proof roadmap of Theorem 6.1 is the following:

Show Assumption 3.1 is satisfied, i.e., the set of local minimizers, $\Gamma$, is indeed a manifold and the hessian $\nabla^{2}L(x)$ is non-degenerate on $\Gamma$ (by Lemma 6.2);

Show Assumption 3.2 is satisfied, i.e., $\Phi(U)\subset\Gamma$ (by Lemma 6.3);

Show the limiting flow (13) converges to the minimizer of the regularizer (by Lemma 6.5);

Show the minimizer of the regularizer recovers the groundtruth (by Lemma 6.6).

Our setting is more general than HaoChen et al. (2020), which assumes $w^{*}\in\{0,1\}^{d}$ and their reparametrization can only express positive linear functions, i.e., $w=u^{\odot 2}$. Their $\widetilde{O}(\kappa^{2})$ rate is achieved with a delicate three phase LR schedule, while our $O(\kappa\ln d)$ rate only uses a constant LR.

We verify that the above loss function $L$ and manifold $\Gamma$ satisfy Assumption 3.1 by Lemma 6.2, and that the neighborhood $U$ and $\Gamma$ satisfy Assumption 3.2 by Lemma 6.3.

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., $u_{j}=v_{j}=\alpha,\forall j\in[n]$ for some $\alpha>0$. Lemma 6.3 completely removes the technical assumption.

The limiting behavior of label noise SGD is described by a Riemannian gradient flow on $\Gamma$ as follows:

The goal is to show that the above limiting flow will converge to the underlying groundtruth $x^{*}=\binom{u^{*}}{v^{*}}$ where $(u^{*},v^{*})=([w^{*}]_{+}^{\odot 1/2},[-w^{*}]_{+}^{\odot 1/2})$.

1 Limiting Flow Converges to Minimizers of Regularizer

Thus the non-compactness of $\Gamma$ brings challenges for both (a) and (b). For (a), the convergence for standard gradient flow is often for free, as long as the trajectory is bounded and the objective is analytic or smooth and semialgebraic. The latter ensures the so-called Kurdyka-Łojasiewicz (KL) inequality (Lojasiewicz, 1963), which implies finite trajectory length and thus the convergence. However, since our flow does not satisfy those nice properties, we have to show that the limiting flow satisfies Polyak-Łojasiewicz condition (a special case of KL condition) (Polyak, 1964) via careful calculation (by Lemma D.16).

For (b), the standard analysis based on center stable manifold theorem shows that gradient descent/flow converges to strict saddle (stationary point with at least one negative eigenvalue in hessian) only for a zero-measure set of initialization (Lee et al., 2016, 2017). However, such analyses cannot deal with the case where the flow is not differentiable at the sub-optimal stationary point. To circumvent this issue, we prove the non-convergence to sub-optimal stationary points with a novel approach: we show that for any stationary point $x$, whenever there exists a descent direction of the regularizer $R$ at $x$, we can construct a potential function which increases monotonically along the flow around $x$, while the potential function is equal to $-\infty$ at $x$, leading to a contradiction. (See proof of Lemma 6.5.)

2 Minimizer of the Regularizer Recovers the Sparse Groundtruth

3 Lower Bound for Gradient Descent in the Kernel Regime

In this subsection we show GD needs at least $\Omega(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 $\widetilde{O}(\kappa\ln d)$ sample complexity upper bound of SGD with label noise. Following the setting of kernel regime in (Woodworth et al., 2020), we consider the limit of $u_{0}=v_{0}=\alpha\mathds{1}$, with $\alpha\to\infty$. It holds that $f_{i}(u_{0},v_{0})=0$ and $\nabla f_{i}(u_{0},v_{0})=[\alpha z_{i},-\alpha z_{i}]$ for each $i\in[n]$. Standard convergence analysis for NTK (Neural Tangent Kernel, Jacot et al. (2018)) shows that upon convergence, the distance traveled by parameter converges to , and thus the learned model shall converge in function space, so is the generalization performance. For ease of illustration, we directly consider the lower bound for test loss when the NTK is fixed throughout the training.

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, $\Theta(\eta^{-2})$ steps of SGD converge to a limiting diffusion on certain manifold of local minimizer, as the LR $\eta\to 0$. For specific noise types, this allows us to recover and strengthen results regarding implicit bias in previous works with much simpler analysis. In particular, we show a sample complexity gap between label noise SGD and GD in the kernel regime for a overparametrized linear model, justifying the generalization benefit of SGD. For the future work, we believe our framework can be applied to analyze the implicit bias of SGD in more complex models towards better understanding of the algorithmic regularization induced by stochasticity. It will be valuable to extend our method to other stochastic optimization algorithms, e.g., ADAM, SGD with momentum.

Acknowledgement

We thank Yangyang Li for pointing us to Katzenberger (1991). We also thank Wei Zhan, Jason Lee and Lin Chen 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

Appendix A Preliminaries on Stochastic Processes

Next, we review a few basics of stochastic processes that will be useful for proving our results, so that our paper will be self-contained. We refer the reader to classics like Karatzas & Shreve (2014); Billingsley (2013); Pollard (2012) for more systematic derivations.

Let $T\in[0,\infty]$. A function $g:[0,T)\to E$ is càdlàg if for all $t\in[0,T)$ it is right-continuous at $t$ and its left limit $g(t-)$ exists. Let $\mathcal{D}_{\mathcal{E}}[0,T)$ be the set of all càdlàg function mapping $[0,T)$ into $\mathcal{E}$. We also use $\mathcal{D}_{\mathcal{E}}[0,T)$ to denote the set of all continuous function mapping $[0,T)$ into $\mathcal{E}$. By definition, $\mathcal{C}_{\mathcal{E}}[0,T)\subset\mathcal{D}_{\mathcal{E}}[0,T)$.

For any function $f:[0,\infty)\to\mathcal{E}$ and any interval $I\subseteq[0,\infty)$, we define

Moreover, the continuity modulus of càdlàg $f\in\mathcal{D}_{\mathcal{E}}[0,\infty)$ is defined as

For any $g\in\mathcal{D}_{\mathcal{E}}[0,T)$, we define the jump of $g$ at $t$ to be

For any $\delta>0$, we define $h_{\delta}:[0,\infty)\to[0,\infty)$ by

For each finite $T>0$ and each pair of functions $f,g\in\mathcal{D}_{\mathcal{E}}[0,\infty)$, define $d_{T}(f,g)$ as the infimum of all those values of $\delta$ for which there exist grids $0\leq t_{0}<t_{1}<\cdots<t_{m}$ and $0<s_{0}<s_{1}<\cdots<\cdots<s_{m}$, with $t_{k},s_{k}\geq T$, such that $|t_{i}-s_{i}|\leq\delta$ for $i=0,\ldots,k$, and

for $i=0,\ldots,k-1$. The Skorokhod metric on $\mathcal{D}_{\mathcal{E}}[0,\infty)$ is defined to be

A.2 Stochastic Processes and Stochastic Integral

in probability as the mesh size of $0=t_{0}<t_{1}<\cdots<t_{m}=t$ goes to 0, if it exists. Moreover, for $Y$ itself, we write

Let $\{X(t)\}_{t\geq 0}$ be a $\{\mathcal{F}_{t}\}_{t\geq 0}$-adapted stochastic process. If for all $0\leq s\leq t$, it holds that

Let $\{X(t)\}_{t\geq 0}$ be a $\{\mathcal{F}_{t}\}_{t\geq 0}$-adapted stochastic process. If there exists a sequence of $\{\mathcal{F}_{t}\}_{t\geq 0}$-stopping time, $\{\tau_{k}\}_{k\geq 0}$, such that

and $\{X^{\tau_{k}}(t)\}_{t\geq 0}$ is a $\{\mathcal{F}_{t}\}_{t\geq 0}$-adapted martingale,

Let $\{X(t)\}_{t\geq 0}$ be a $\{\mathcal{F}_{t}\}_{t\geq 0}$-adapted stochastic process. If there exists a local martingale $\{M(t)\}_{t\geq 0}$ and a càdlàg $\{\mathcal{F}_{t}\}_{t\geq 0}$-adapted process $\{A(t)\}_{t\geq 0}$ with bounded total variation that $X(t)=M(t)+A(t)$, then $X$ is called a semimartingale.

Since all deterministic process are adapted, the above definition of integral also makes sense for deterministic functions and is a generalization of standard Riemman-Stieltjes Integral. The difference is that in the above Itô’s Stochastic Integral we use the left-end value of the integrand but the existence of Riemman-Stieltjes Integral requires the limit exists for any point within the interval. When $X$ and $Y$ don’t jump together, Riemman-Stieltjes Integral exists and coincides with the Itô’s Integral.

Let $\{X(t)\}_{t\geq 0}$ be defined through the following Itô drift-diffusion process:

where $\{W(t)\}_{t\geq 0}$ is the standard Brownian motion. Then for any twice differentiable function $f$, it holds that

A.3 Weak Convergence for Stochastic Processes

Let $(\mathcal{D}_{\mathcal{E}}[0,\infty),\mathcal{A},d)$ be a metric space equipped with a $\sigma$-algebra $\mathcal{A}$ and the Skorokhod metric defined in the previous subsection.

Though we define weak convergence for a countable sequence of stochastic processes, but it is still valid if we index the stochastic processes by real numbers, e.g., $\{X_{\eta}\}_{\eta\geq 0}$, and consider the weak convergence of $X_{\eta}$ as $\eta\to 0$. This is because the convergence in (20) is for a sequence of real numbers, which is also well-defined if we replace $\lim_{n\to\infty}$ by $\lim_{\eta\to 0}$.

Let $\delta>0$. For any two probability measures $P$ and $Q$ on a metric space with metric $d$, let $(X,Y)$ be a coupling such that $P$ is the marginalized law of $X$ and $Q$ that of $Y$. We define

Note this distance is not a metric because it does not satisfy triangle inequality.

For any two probability measures $P$ and $Q$ on a metric space with metric $d$, let $(X,Y)$ be a coupling such that $P$ is the marginalized law of $X$ and $Q$ that of $Y$. Denote the marginal laws of $X$ and $Y$ by $\mathcal{L}(X)$ and $\mathcal{L}(Y)$ respectively. We define the Prohorov metric as

It can be shown that $X_{n}\Rightarrow X$ is equivalent to $\lim_{n\to\infty}\rho(X_{n},X)=0$.

For each finite $T>0$ and each pair of functions $f,g\in\mathcal{D}_{\mathcal{E}}[0,T)$, the uniform metric is defined to be

The uniform metric on $\mathcal{D}_{\mathcal{E}}[0,\infty)$ is defined to be

Appendix B Limiting Diffusion of SGD

First, as mentioned in Assumption 3.2, we verify that the mapping $\Phi$ is $\mathcal{C}^{2}$ in Lemma B.1. In Section B.1 we discuss how different time scalings could affect the coefficients in SDE (2) and (3). Then we check the necessary conditions for applying the results in Katzenberger (1991) in Section B.2 and recap the corresponding theorem for the asymptotically continuous case in Section B.3. Finally, we provide a user-friendly interface for Katzenberger’s theorem in Section B.4.

Under Assumption 3.2, $\Phi$ is $\mathcal{C}^{2}$ on $U$.

Applying Theorem 5.1 of Falconer (1983) with $f(\cdot)=\phi(\cdot,1)$ suffices. ∎

where the time correspondence is $t=k\eta$, i.e., $X(k\eta)\approx x_{\eta}(k)$.

Now rescale the above SDE by considering $\widetilde{X}(t)=X(t\eta)$, which then yields

where the time correspondence is $t=k$, i.e., $\widetilde{X}(k)\approx x_{\eta}(k)$. The above SDE is exactly the same as (2).

Then, to accelerate the above SDE by $\eta^{-2}$ times, let’s define $\bar{X}(t)=\widetilde{X}(t/\eta^{2})$. Then it follows that

Again note that $\eta W(t/\eta^{2})\overset{d}{=}W(t)$ in sample paths and thus is also a $\Xi$-Brownian motion. Here the time correspondence is $t=k\eta^{2}$, i.e., evolving for constant time with the above SDE approximates $\Omega(1/\eta^{2})$ 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 $\{Z_{n}\}_{n\geq 1}$ and $\{A_{n}\}_{n\geq 1}$ in Katzenberger (1991). Recall that we consider the following stochastic process

For any stopping time $\tau$, the stopped process is defined as $X_{n}^{\tau}(t)=X_{n}(t\wedge\tau)$. For any compact $K\subset U$, we define the stopping time of $X_{n}$ leaving $K$ as $\lambda_{n}(K)=\inf\{t\geq 0\mid X_{n}(t-)\notin\mathring{K}\text{ or }X_{n}(t)\notin\mathring{K}\}$.

The integrator sequence $\{A_{n}\}_{n\geq 1}$ is asymptotically continuous: $\sup\limits_{t>0}|A_{n}(t)-A_{n}(t-)|\Rightarrow 0$ where $A_{n}(t-)=\lim_{s\to t-}A_{n}(s)$ is the left limit of $A_{n}$ at $t$.

The integrator sequence $\{A_{n}\}_{n\geq 1}$ increases infinitely fast: $\forall\epsilon>0$, $\inf\limits_{t\geq 0}(A_{n}(t+\epsilon)-A_{n}(t))\Rightarrow\infty$.

For every $T>0$, as $n\to\infty$, it holds that

for every $\epsilon>0$ and $T>0$, where $T_{t}(\cdot)$ denotes total variation on the interval $[0,t]$.

For SGD iterates defined using the notation in Lemma 4.2, the sequences $\{A_{n}\}_{n\geq 1}$ and $\{Z_{n}\}_{n\geq 1}$ satisfy Condition B.2, B.3, B.4 and B.5.

Condition B.2 is obvious from the definition of $\{A_{n}\}_{n\geq 1}$.

Next, for any $\epsilon>0$ and $t\in[0,T]$, we have

which implies that $\inf_{0\leq t\leq T}(A_{n}(t+\epsilon)-A_{n}(t))>\epsilon/(2\eta_{n})$ for small enough $\eta_{n}$. Then taking $n\to\infty$ yields the Condition B.3.

Therefore, we have $\|\Delta Z_{n}(t)\|_{2}\leq 2\eta_{n}\sqrt{\Xi}$ for all $t>0$. This implies that $\|\Delta Z_{n}(t)\|_{2}\to 0$ uniformly over $t>0$ as $n\to\infty$, which verifies Condition B.4.

We proceed to verify Condition B.5. By the definition of $Z_{n}$, we know that $\{Z_{n}(t)\}_{t\geq 0}$ is a jump process with independent increments and thus is a martingale. Therefore, by decomposing $Z_{n}=M_{n}+F_{n}$ with $M_{n}$ being a local martingale and $F_{n}$ a finite variation process, we must have $F_{n}=0$ and $M_{n}$ is $Z_{n}$ itself. It then suffices to show that $[M_{n}](t\wedge\tau_{n}^{m})$ is uniformly integrable for every $t\geq 0$ and $m\geq 1$. Since $M_{n}$ is a pure jump process, we have

This implies that $[M_{\eta}](t\wedge\tau_{\eta}^{m})$ is universally bounded by $4t$, and thus $[M_{\eta}](t\wedge\tau_{\eta}^{m})$ is uniformly integrable. This completes the proof. ∎

For any $n\geq 1$, it suffices to show that given $X_{n}(k\eta_{n}^{2})=x_{\eta_{n}}(k)$, we further have $X_{n}((k+1)\eta_{n}^{2})=x_{\eta_{n}}(k+1)$. By the definition of $X_{n}(t)$ and note that $A_{n}(t),Z_{n}(t)$ are constants on $[k\eta_{n}^{2},(k+1)\eta_{n}^{2})$, we have that $X_{n}(t)=X_{n}(k\eta_{n}^{2})$ for all $t\in[k\eta_{n}^{2},(k+1)\eta_{n}^{2})$, and therefore

where the second equality is because $A_{n}(t)$ and $Z_{n}(t)$ are constant on interval $[k\eta_{n}^{2},(k+1)\eta_{n}^{2})$. This confirms the alignment between $\{X_{n}(k\eta_{n}^{2})\}_{k\geq 1}$ and $\{x_{\eta_{n}}(k)\}_{k\geq 1}$.

B.3 Katzenberger’s Theorem for Asymptotically Continuous Case

The full Katzenberger’s theorem deals with a more general case, which only requires the sequence of intergrators to be asymptotically continuous, thus including SDE (3) and SGD (1) with $\eta$ goes to .

for all $t\leq\lambda_{n}(K)$ where $\lambda_{n}(K)=\inf\{t\geq 0\mid X_{n}(t-)\notin\mathring{K}\text{ or }X_{n}(t)\notin\mathring{K}\}$ is the stopping time of $X_{n}$ leaving $K$.

We note that by Lemma A.16, convergence in distribution under skorohod metric is equivalent to convergence in distribution under uniform metric Definition A.15, therefore in the rest of the paper we will only use the uniform metric in the rest of the paper, e.g., whenever we mention Prohorov metric and $\delta$-Prohorov distance, the underlying metric is the uniform metric.

B.4 A User-friendly Interface for Katzenberger’s Theorem

Based on the Lemma B.6, we can immediately apply Theorem B.7 to obtain the following limiting diffusion of SGD.

Let the manifold $\Gamma$ and its open neighborhood $U$ satisfy Assumptions 3.1 and 3.2. Let $K\subset U$ be any compact set and fix some $x_{0}\in K$. Consider the SGD formulated in Lemma 4.2 where $X_{\eta_{n}}(0)\equiv x_{0}$. Define

where $\{W(s)\}_{s\geq 0}$ is the standard Brownian motion and $\sigma(\cdot)$ is as defined in Lemma 4.2.

However, the above theorem is hard to parse and cannot be directly applied if we want to further study the implicit bias of SGD through this limiting diffusion. Therefore, we develop a user-friendly interface to it in below. In particular, Theorem 4.6 is the a special case of Theorem B.9. In Theorem 4.6, we replace $\partial\Phi(Y(t))\sigma(Y(t))$ with $\Sigma_{\parallel}^{\frac{1}{2}}(Y(t))$ to simplify the equation, since $\partial\Phi(Y(t))\sigma(Y(t))\left(\partial\Phi(Y(t))\sigma(Y(t))\right)^{\top}=\Sigma_{\parallel}(Y(t))$ and thus this change doesn’t affect the distribution of the sample paths of the solution.

which means there is a coupling between the distribution of the stopped processes $Y_{\eta}^{\mu_{\eta}(K)\wedge T}$ and $Y^{\mu(K)\wedge T}$, such that the uniform metric between them is smaller than $\epsilon$ with probability at least $1-2\delta$. In other words, $\lim_{\eta\to 0}\rho^{2\delta}(Y_{\eta}^{\mu_{\eta}(K)\wedge T},Y^{\mu(K)\wedge T})=0$.

Moreover, when $\{Y(t)\}_{t\geq 0}$ is a global solution to the following limiting diffusion

For clarity, we break the proof of Theorem B.9 into two parts, devoted to the two claims respectively.

First, Theorem B.8 guarantees there exists a stopping time $\widetilde{\mu}$ and a stochastic process $\{\widetilde{Y}(t)\}_{t\geq 0}$ such that

$(\widetilde{Y},\widetilde{\mu})$ satisfies Equation 23;

$\widetilde{\mu}\geq\widetilde{\mu}(K):=\inf\{t\geq 0\mid\widetilde{Y}(t)\notin\mathring{K}\}$.

Let $\mathcal{E}_{T}$ be the event such that $\widetilde{\mu}(K)>T$ on $\mathcal{E}_{T}$. Then restricted on $\mathcal{E}_{T}$, we have $\widetilde{Y}(T\wedge\widetilde{\mu})=\widetilde{Y}(T\wedge\widetilde{\mu}(K))$ as $\widetilde{\mu}\geq\widetilde{\mu}(K)$ holds a.s. We first prove the claim for any convergent subsequence of $\{Y_{\eta}\}_{\eta>0}$.

Now, let $\{\eta_{m}\}_{m\geq 1}$ be a sequence of LRs such that $\eta_{m}\to 0$ and $Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)}\Rightarrow\widetilde{Y}^{\widetilde{\mu}}$ as $m\to\infty$. By applying the Skorohod representation theorem, we can put $\{Y_{\eta_{m}}\}_{m\geq 1}$ and $\widetilde{Y}$ under the same probability space such that $Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)}\to\widetilde{Y}^{\widetilde{\mu}}$ a.s. in the Skorohod metric, or equivalently the uniform metric (since $\widetilde{Y}^{\widetilde{\mu}}$ is continuous) i.e.,

which further implies that for any $\epsilon>0$, there exists some $N>0$ such that for all $m>N$,

Restricted on $\mathcal{E}_{T}$, we have $d_{U}(Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)\wedge T},\widetilde{Y}^{\widetilde{\mu}\wedge T})=d_{U}(Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)\wedge T},\widetilde{Y}^{\widetilde{\mu}(K)\wedge T})$, and it follows that for all $m>N$,

where we denote the complement of $\mathcal{E}_{T}$ by $\mathcal{E}_{T}^{c}$.

By the definition of the Prohorov metric in Definition A.13, we then get $\rho^{2\delta}(Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)\wedge T},\widetilde{Y}^{\widetilde{\mu}(K)\wedge T})\leq\epsilon$ for all $m>N$. Therefore, we have

Now we claim that it indeed holds that $\lim_{\eta\to 0}\rho^{2\delta}(Y_{\eta}^{\mu_{\eta}(K)\wedge T},\widetilde{Y}^{\widetilde{\mu}(K)\wedge T})=0$. We prove this by contradiction. Suppose otherwise, then there exists some $\epsilon>0$ such that for all $\eta_{0}>0$, there exists some $\eta<\eta_{0}$ with $\rho^{2\delta}(Y_{\eta}^{\mu_{\eta}(K)\wedge T},\widetilde{Y}^{\widetilde{\mu}(K)\wedge T})>\epsilon$. Consequently, there is a sequence $\{\eta_{m}\}_{m\geq 1}$ satisfying $\lim_{m\to\infty}\eta_{m}=0$ and $\rho^{2\delta}(Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)},\widetilde{Y}^{\widetilde{\mu}(K)\wedge T})>\epsilon$ for all $m$. Since $\{(Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)\wedge T},Z_{\eta_{m}},\mu_{\eta_{m}}(K))\}_{m\geq 1}$ is relatively compact, there exists a subsequence (WLOG, assume it is the original sequence itself) converging to $(\widetilde{Y}^{\widetilde{\mu}\wedge T},W,\widetilde{\mu})$ in distribution. However, repeating the exactly same argument as above, we would have $\rho^{2\delta}(Y_{\eta_{m}}^{\mu_{\eta_{m}}(K)\wedge T},\widetilde{Y}^{\widetilde{\mu}(K)\wedge T})\leq\epsilon$ for all sufficiently large $m$, which is a contradiction. This completes the proof. ∎

Note that $\rho^{\delta}(Y^{\mu(K)\wedge T},Y^{\mu(K^{\prime})\wedge T})=0$, so we have for all $\eta\leq\eta_{0}$,

On the other hand, if $\mu_{\eta}(K^{\prime})>T$, then $Y_{\eta}^{T}=Y_{\eta}^{\mu_{\eta}(K^{\prime})\wedge T}$. Thus we can conclude that $d_{U}(Y^{\mu(K)\wedge T},Y_{\eta}^{T})\geq 2^{-\lceil T\rceil}\epsilon^{\prime}$ implies $d_{U}(Y^{\mu(K)\wedge T},Y_{\eta}^{\mu_{\eta}(K^{\prime})\wedge T})\geq 2^{-\lceil T\rceil}\epsilon^{\prime}$. Therefore, we further have

Finally, since $\rho^{\delta}(Y^{T},Y^{\mu(K)\wedge T})=0$, we have for all $\eta\leq\eta_{0}$,

Now, we provide the proof of Theorem 4.6 as a direct application of Theorem B.9.

which means $Y$ always stays on $\Gamma$.

Then recall the decomposition of $\Sigma=\Sigma_{\parallel}+\Sigma_{\perp}+\Sigma_{\parallel,\perp}+\Sigma_{\perp,\parallel}$ as defined in Lemma 4.5. Since $Y$ never leaves $\Gamma$, by Lemma 4.5, we can rewrite Equation 10 as

where the second equality follows from the definition that $\Sigma_{\parallel}=\partial\Phi\Sigma\partial\Phi=\partial\Phi\sigma\sigma^{\top}\partial\Phi$. 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. ∎

Our result suggests that for tiny LR $\eta$, SGD dynamics have two phases. In Phase I of $\Theta(1/\eta)$ steps, the SGD iterates move towards the manifold $\Gamma$ of local minimizers along GF. Then in Phase II which is of $\Theta(1/\eta^{2})$ steps, the SGD iterates stay close to $\Gamma$ and diffuse approximately according to (10). See Figure 2 for an illustration of this two-phase dynamics. However, since the length of Phase I gets negligible compared to that of Phase II when $\eta\to 0$, Theorem 4.6 only reflects the time scaling of Phase II.

Appendix C Explicit Formula of the Limiting Diffusion

In this section, we demonstrate how to compute the derivatives of $\Phi$ by relating to those of the loss function $L$, and then present the explicit formula of the limiting diffusion.

For any $x\in\Gamma$ and any $v\in T_{x}(\Gamma)$, it holds that $\nabla^{2}L(x)v=0$.

Evaluating the above equation at $t=0$ yields $\partial\Phi(x)\nabla L(x)=0$. Moreover, take the second order derivative and we have

Evaluating at $t=0$ completes the proof. ∎

Now we can prove Lemma 4.3, restated in below. See 4.3

This implies that $\partial\Phi(x)v=v$ for all $v\in T_{x}(\Gamma)$.

Next, for any $u\in T_{x}^{\perp}(\Gamma)$ and $t\geq 0$, consider expanding $\nabla L(x+t\nabla^{2}L(x)^{\dagger}u)$ at $t=0$:

where the second equality follows from the assumption that $\nabla^{2}L(x)$ is full-rank when restricted on $T_{x}^{\perp}(\Gamma)$. Then since $\partial\Phi$ is continuous, it follows that

By Lemma C.2, we have $\partial\Phi(x+t(\nabla^{2}L(x))^{\dagger}u))\nabla L(x+t(\nabla^{2}L(x))^{\dagger}u)=0$ for all $t>0$, which then implies that $\partial\Phi(x)u=0$ for all $u\in T_{x}^{\perp}(\Gamma)$.

Therefore, under the basis $\{v_{1},\ldots,v_{D}\}$, $\partial\Phi(x)$ is given by

that is, the projection matrix onto $T_{x}(\Gamma)$. ∎

For any $x\in\Gamma$, it holds that $\partial\Phi(x)\nabla^{2}L(x)=0$.

It directly follows from Lemma C.1 and Lemma 4.3. ∎

Next, we proceed to compute the second-order derivatives.

Denote the derivative of $P(t),P^{\perp}(t)$ and $H(t)$ with respect to $t$ as $P^{\prime}(t),(P^{\perp})^{\prime}(t)$ and $H^{\prime}(t)$. Then differentiating with respect to $t$, we have

Then combining (25) and (26) and evaluating at $t=0$, we have

We can decompose $P^{\prime}(0)$ and $H(0)$ as follows

This implies that we must have $P^{\prime}_{22}(0)=0$ and $P^{\prime}_{12}(0)H_{22}(0)=H^{\prime}_{12}(0)$. Similarly, by taking transpose in (27), we also have $H_{22}(0)P_{21}^{\prime}(0)=-H_{21}^{\prime}(0)$.

It then remains to determine the value of $P^{\prime}_{11}(0)$. Note that since $P(t)P(t)=P(t)$, we have $P^{\prime}(t)P(t)+P(t)P^{\prime}(t)=P^{\prime}(t)$, evaluating which at $t=0$ yields

Therefore, we must have $P_{11}^{\prime}(0)=0$. Combining the above results, we obtain

Finally, recall that $P(t)=\partial\Phi(v(t))$, and thus

Similarly, we have $H^{\prime}(0)=\partial^{2}(\nabla L)(x)[v]$, and it follows that

For any $x\in\Gamma$ and $u\in T_{x}^{\perp}(\Gamma)$, it holds that

For any $u\in T_{x}^{\perp}(\Gamma)$, we define $u(t)=x+t\nabla^{2}L(x)^{\dagger}u$ for $t\geq 0$. By Taylor approximation, we have

Combine (28) and (29) and apply Lemma C.2, and it follows that

where the last equality follows from Lemma C.3. Dividing both sides by $t^{2}$ and letting $t\to 0$, we get

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.

Let $A=uu^{\top}+\nabla^{2}L(x)^{\dagger}uu^{\top}\nabla^{2}L(x)$ and $B=\nabla^{2}L(x)^{\dagger}uu^{\top}$. The key observation is that $A+A^{\top}=\mathcal{L}_{\nabla^{2}L(x)}(B+B^{\top})$. Therefore, by Lemma C.5, it holds that

Then Lemma 4.5 directly follows from Lemma C.4 and C.6.

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 on $\Gamma$, only depends on $\Gamma$ itself.

For any $x\in\Gamma$, suppose there exist a neighborhood $U_{x}$ of $x$ and two loss functions $L$ and $L^{\prime}$ that define the same manifold $\Gamma$ locally in $U_{x}$, i.e., $\Gamma\cap U_{x}=\{x\mid\nabla L(x)=0\}=\{x\mid\nabla L^{\prime}(x)=0\}$. Then for any $v\in T_{x}(\Gamma)$, it holds that $(\nabla^{2}L(x))^{\dagger}\partial^{2}(\nabla L)(x)\left[v,v\right]=(\nabla^{2}L^{\prime}(x))^{\dagger}\partial^{2}(\nabla L^{\prime})(x)\left[v,v\right]$.

C.3 Proof of results in Section 5

Now we are ready to give the missing proofs in Section 5 which yield explicit formula of the limiting diffusion for label noise and isotropic noise.

Set $\Sigma_{\parallel}=\partial\Phi$, $\Sigma_{\perp}=I_{D}-\partial\Phi$ and $\Sigma_{\perp,\parallel}=\Sigma_{\parallel,\perp}=0$ in the decomposition of $\Sigma$ by Lemma 4.5, and we need to show $\nabla(\ln|\Sigma|_{+})=\partial^{2}(\nabla L)[(\nabla^{2}L)^{\dagger}]$.

C.4 Example: k𝑘k-Phase Motor

We also give an example with rigorous proof where the implicit bias induced by noise in the normal space cannot be characterized by a fixed regularizer, which was first discovered by Damian et al. (2021) but was only verified via experiments.

Note the normal regularization in both cases of label noise and isotropic noise induces Riemmanian gradient flow against some regularizer, it’s natural to wonder if the limiting flow induced by the normal noise 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.

The basic idea is that we can add noise in the ‘auxiliary dimensions’ for $j=3,\ldots,D$ to get the regularization force on the circle $\{x_{1}^{2}+x_{2}^{2}=1\}$, and the goal is to make the vector field induced by the normal regularization always point to the same direction, say anti-clockwise. However, this cannot be done with a single auxiliary dimension because from the analysis for label noise, we know when $\mathcal{L}^{-1}_{\nabla^{2}L}(\Sigma_{\perp})$ is identity, the normal regularization term in 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\geq 3$ auxiliary dimensions, we can further ensure the normal regularization force is anti-clockwise and of constant magnitude, which is reminiscent of how a three-phase induction motor works.

Note that for any $x\in\Gamma$, it holds that

Then clearly $\Sigma$ only brings about noise in the normal space, and specifically, it holds that $\mathcal{L}_{\nabla^{2}L(x)}^{-1}(\Sigma(x))={\rm diag}(0,0,1+\left\langle Q_{\alpha}^{0}v,Q_{-\pi/2}x_{1:2}\right\rangle,\ldots,1+\left\langle Q_{\alpha}^{D-3}v,Q_{-\pi/2}x_{1:2}\right\rangle)$. Further note that, by the special structure of the hessian in (33) and Lemma C.3, for any $x\in\Gamma$, we have $\partial\Phi(x)=(x_{2},-x_{1},0,\ldots,0)^{\top}(x_{2},-x_{1},0,\ldots,0)=\binom{Q_{-\pi/2}x_{1:2}}{0}\binom{Q_{-\pi/2}x_{1:2}}{0}^{\top}$. Combining these facts, the dynamics of the first two coordinates in SDE (10) can be simplified into

The proof is completed by noting that the solution of $x_{1:2}$ is

$\exp(\left(\begin{smallmatrix}0&-1\\ 1&0\end{smallmatrix}\right))=\left(\begin{smallmatrix}\cos 1&-\sin 1\\ \sin 1&\cos 1\end{smallmatrix}\right)$.

By definition, for matrix $A=\left(\begin{smallmatrix}0&-1\\ 1&0\end{smallmatrix}\right)$, $\exp(A)=\sum_{t=0}^{\infty}\frac{A^{t}}{t!}$. Note that $A_{2}=-I$, $A^{3}=-A$ and $A^{4}=I$. Using this pattern, we can easily check that

Appendix D Proof of results in Section 6

In this section, we present the missing proofs in Section 6 regarding the overparametrized linear model.

In this subsection, we provide the proof of Theorem 6.1.

First, by Lemma 6.6, it holds with probability at least $1-e^{-\Omega(n)}$ that the solution to (18), $x_{*}$, is unique up to and satisfies $|x_{*}|=\psi(w_{*})$. Then on this event, for any $\epsilon>0$, by Lemma 6.5, there exists some $T>0$ such that $x_{T}$ given by the Riemannian gradient flow (17) satisfies that $x_{T}$ is an $\epsilon/2$-optimal solution of the OLM. For this $T$, by Theorem 4.6, we know that the $\lfloor T/\eta^{2}\rfloor$-th SGD iterate, $x_{\eta}(\lfloor T/\eta^{2}\rfloor)$, satisfies $\|x_{\eta}(\lfloor T/\eta^{2}\rfloor)-x_{T}\|_{2}\leq\epsilon/2$ with probability at least $1-e^{-\Omega(n)}$ for all sufficiently small $\eta>0$, and thus $x_{\eta}(\lfloor T/\eta^{2}\rfloor)$ is an $\epsilon$-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. ∎