Understanding Gradient Descent on Edge of Stability in Deep Learning

Sanjeev Arora, Zhiyuan Li, Abhishek Panigrahi

Introduction

Traditional convergence analyses of gradient-based algorithms assume learning rate $\eta$ is set according to the basic relationship $\eta<2/\lambda$ where $\lambda$ is the largest eigenvalue of the Hessian of the objective, called sharpnessConfusingly, another traditional name for $\lambda$ is smoothness.. Descent Lemma says that if this relationship holds along the trajectory of Gradient Descent, loss drops during each iteration. In deep learning where objectives are nonconvex and have multiple optima, similar analyses can show convergence towards stationary points and local minima. In practice, sharpness is unknown and $\eta$ is set by trial and error. Since deep learning works, it has been generally assumed that this trial and error allows $\eta$ to adjust to sharpness so that the theory applies. But recent empirical studies (Cohen et al., 2021; Ahn et al., 2022) showed compelling evidence to the contrary. On a variety of popular architectures and training datasets, GD with fairly small values of $\eta$ displays following phenomena that they termed Edge of Stability (EoS): (a) Sharpness rises beyond $2/\eta$ , thus violating the above-mentioned relationship. (b) Thereafter sharpness stops rising but hovers noticeably above $2/\eta$ and even decreases a little. (c) Training loss behaves non-monotonically over individual iterations, yet consistently decreases over long timescales.

Note that (a) was already pointed out by Li et al. (2020b). Specifically, in modern deep nets, which use some form of normalization combined with weight decay, training to near-zero loss must lead to arbitrarily high sharpness. (However, Cohen et al. (2021) show that the EoS phenomenon appears even without normalization.) Phenomena (b), (c) are more mysterious, suggesting that GD with finite $\eta$ is able to continue decreasing loss despite violating $\eta<2/\lambda$ , while at the same time regulating further increase in value of sharpness and even causing a decrease. These striking inter-related phenomena suggest a radical overhaul of our thinking about optimization in deep learning. At the same time, it appears mathematically challenging to analyze such phenomena, at least for realistic settings and losses (as opposed to toy examples with 2 or 3 layers). The current paper introduces frameworks for doing such analyses.

We start by formal definition of stableness, ensuring that if a point + LR combination is stable then a gradient step is guaranteed to decrease the loss by the local version of Descent Lemma.

The second setting assumes that the loss is smooth but learning rate is effectively adaptive. We focus a concrete example, Normalized Gradient Descent, $x\leftarrow x-\eta\nabla L/\|\nabla L\|$ , which exhibits EoS behavior as $\nabla L\rightarrow 0$ . We can view Normalized GD as GD with a varying LR $\eta_{t}=\frac{\eta}{\left\|\nabla L(x(t))\right\|}$ , which goes to infinity when $\nabla L\rightarrow 0$ .

We show that Normalized GD on $L$ (Section 4.3) and GD on $\sqrt{L}$ (Section 4.4) exhibit similar two-phase dynamics with sufficiently small LR $\eta$ . In the first phase, GD tracks gradient flow (GF), with a monotonic decrease in loss until getting $O(\eta)$ -close to the manifold (Theorems 4.3 and 4.5) and the stableness becomes larger than $2$ . In the second phase, GD no longer tracks GF and loss is not monotone decreasing due to the high stableness. Repeatedly overshooting, GD iterate jumps back and forth across the manifold while moving slowly along the direction in the tangent space of the manifold which decreases the sharpness. (See Figure 1 for a graphical illustration) Formally, we prove when $\eta\to 0$ , the trajectory of GD converges to some limiting flow on the manifold. (Theorems 4.4 and 4.6) We further prove that in both settings GD in the second phase operates on EOS, and loss decreases in a non-monotone manner. Formally, we show that the average stableness over any two consecutive steps is at least 2 and that the average of $\sqrt{L}/\eta$ over two consecutive is proportional to sharpness or square root of sharpness. (Theorems 4.7 and 4.8)

Though many works have suggested (primarily via experiments and some intuition) that the training algorithm in deep learning implicitly selects out solutions of low sharpness in some way, we are not aware of a formal setting where this had ever been made precise. Note that our result requires no stochasticity as in SGD (Li et al., 2022b), though we need to inject tiny noise (e.g., of magnitude $O(\eta^{100})$ ) to GD iterates occasionally (Algorithms 1 and 2). We believe that this is due the technical limitation of our current analysis and can be relaxed with a more advanced analysis. Indeed, in experiments, our theoretical predictions hold for the deterministic GD directly without any perturbation.

Novelty of Our Analysis:

To prove the alignment between the gradient and the top eigenvector of Hessian, it boils down to analyze Normalized GD on quadratic functions (2), which to the best of our knowledge has not been studied before. The dynamics is like chaotic version of power iteration, and we manage to show that the iterate will always align to the top eigenvector of Hessian of the quadratic loss. The proof is based on identifying a novel potential (Section 3) and might be of independent interest.

Related Works

Low sharpness has long been related to flat minima and thus to good generalization (Hochreiter and Schmidhuber, 1997; Keskar et al., 2016). Recent study on predictors of generalization (Jiang et al., 2020) does show sharpness-related measures as being good predictors, leading to SAM algorithm that improves generalization by explicitly controlling a parameter related to sharpness (Foret et al., 2021). However, Dinh et al. (2017) show that due to the positive homogeneity in the network architecture, networks with rescaled parameters can have very different sharpness yet be the same to the original one in function space. This observation weakens correlation between sharpness and and generalization gap and makes the definition of sharpness ambiguous. In face of this challenge, multiple notions of scale-invariant sharpness have been proposed (Yi et al., 2019a, b; Tsuzuku et al., 2020; Rangamani et al., 2021). Especially, Yi et al. (2021); Kwon et al. (2021) derived new algorithms with better generalization by explicitly regularizing new sharpness notions aware of the symmetry and invariance in the network. He et al. (2019) goes beyond the notion of sharpness/flatness and argues that the local minima of modern deep networks can be asymmetric, that is, sharp on one side, but flat on the other side.

Limiting Diffusion/Flow around Manifold of Minimizers:

The idea of analyzing the behavior of SGD with small LR along the the manifold originates from Blanc et al. (2020), which gives a local analysis on a special noise type named label noise, i.e. noise covariance is equal to Hessian at minimizers. Damian et al. (2021) extends this analysis and show SGD with label noise finds approximate stationary point for original loss plus some Hessian-related regularizer. The formal mathematical framework of approximating the limiting dynamics of SGD with arbitrary noise by Stochastic Differential Equations is later established by Li et al. (2022b), which is built on the convergence result for solutions of SDE with large-drift (Katzenberger, 1991).

Implicit Bias:

The notion that training algorithm plays an active role in selecting the solution (when multiple optima exist) has been termed the implicit bias of the algorithm (Gunasekar et al., 2018c) and studied in a large number of papers (Soudry et al., 2018; Li et al., 2018; Arora et al., 2018a, 2019a; Gunasekar et al., 2018b, a; Lyu and Li, 2020; Li et al., 2020a; Woodworth et al., 2020; Razin and Cohen, 2020; Lyu et al., 2021; Azulay et al., 2021; Gunasekar et al., 2021). In the infinite width limit, the implicit bias of Gradient Descent is shown to be the solution with the minimal RKHS norm with respect to the Neural Tangent Kernel (NTK) (Jacot et al., 2018; Li and Liang, 2018; Du et al., 2019; Arora et al., 2019b, c; Allen-Zhu et al., 2019b, a; Zou et al., 2020; Chizat et al., 2019; Yang, 2019). The implicit bias results from these papers are typically proved by performing a trajectory analysis for (Stochastic) Gradient Descent. Most of the results can be directly extended to the continuous limit (i.e., GD infinitesimal LR) and even some heavily relies on the conservation property which only holds for the continuous limit. In sharp contrast, the implicit bias shown in this paper – reducing the sharpness along the minimizer manifold – requires finite LR and doesn’t exist for the corresponding continuous limit. Other implicit bias results that fundamentally relies on the finiteness of LR includes stability analysis (Wu et al., 2017; Ma and Ying, 2021) and implicit gradient regularization (Barrett and Dherin, 2021), which is a special case of approximation results for stochastic modified equation by Li et al. (2017, 2019).

Non-monotone Convergence of Gradient Descent :

Warm-up: Quadratic Loss Functions

Our main result Theorem 3.1 is that the iterates of Normalized GD $x(t)$ converge to $v_{1}$ in direction, from which the loss oscillation Corollary 3.2 follows, suggesting that GD is operating in EoS. Since in quadratic case there is only one local minima, there is of course no need to talk about implicit bias. However, the observation that the GD iterates always align to the top eigenvector as well as the technique used in its proof play a very important role for deriving the sharpness-reduction implicit bias for the case of general loss functions.

Define $\widetilde{x}(t)=\frac{Ax(t)}{\eta}$ , and the following update rule (2) holds. It is clear that the convergence of $\widetilde{x}_{t}$ to $v_{1}$ in direction implies the convergence of $x_{t}$ as well.

If $\left|\langle v_{1},\widetilde{x}(t)\rangle\right|\neq 0$ , $\forall t\geq 0$ , then there exists $0<C<1$ and $s\in\{\pm 1\}$ such that $\lim_{t\to\infty}\widetilde{x}(2t)=Cs\lambda_{1}v_{1}$ and $\lim_{t\to\infty}\widetilde{x}(2t+1)=(C-1)s\lambda_{1}v_{1}$ .

As a direct corollary, the loss oscillates as between time step $2t$ and time step $2t+1$ as $t\to\infty$ . This shows that the behavior of loss is not monotonic and hence indicates the edge of stability phenomena for the quadratic loss.

If $\left|\langle v_{1},\widetilde{x}(t)\rangle\right|\neq 0$ , $\forall t\geq 0$ , then there exists $0<C<1$ such that $\lim_{t\to\infty}L(x(2t))=\frac{1}{2}C^{2}\lambda_{1}\eta^{2}$ and $\lim_{t\to\infty}L(x(2t+1))=\frac{1}{2}(C-1)^{2}\lambda_{1}\eta^{2}$ .

For any $j\in[D]$ and $t\geq\frac{\lambda_{1}}{\lambda_{j}}\ln\frac{\lambda_{1}}{\lambda_{j}}+\max\{\frac{\left\|\widetilde{x}(0)\right\|-\lambda_{1}}{\lambda_{D}},0\}$ , it holds that $\widetilde{x}(t)\in\mathcal{I}_{j}$ .

First, we show for any $j\in[D]$ , $\mathcal{I}_{j}$ is indeed an invariant set for update rule (2) via Lemma A.1. With straightforward calculation, one can show that for any $j\in[D]$ , $\left\|P^{(j:D)}\widetilde{x}(t)\right\|$ decreases by $\frac{\lambda_{D}\left\|P^{(j:D)}\widetilde{x}(t)\right\|}{\left\|\widetilde{x}(t)\right\|}$ if $\left\|P^{(j:D)}\widetilde{x}(t)\right\|\geq\lambda_{j}$ (Lemma A.2). Setting $j=1$ , we have $\left\|\widetilde{x}(t)\right\|$ decreases by $\lambda_{D}$ if $\left\|\widetilde{x}(t)\right\|\geq\lambda_{1}$ (Corollary A.3). Thus for all $t\geq\max\{\frac{\left\|\widetilde{x}(0)\right\|-\lambda_{1}}{\lambda_{D}},0\}$ , $\widetilde{x}(t)\in\mathcal{I}_{1}$ . Finally once $\widetilde{x}(t)\in\mathcal{I}_{1}$ , we can upper bound $\left\|\widetilde{x}(t)\right\|$ by $\lambda_{1}$ , and thus $\left\|P^{(j:D)}\widetilde{x}(t)\right\|$ shrinks at least by a factor of $\frac{\lambda_{D}}{\lambda_{1}}$ per step, which implies $\widetilde{x}(t)$ will be in $\mathcal{I}_{j}$ in another $\frac{\lambda_{1}}{\lambda_{j}}\ln\frac{\lambda_{1}}{\lambda_{j}}$ steps.(Corollary A.4) ∎

Once the component of $\widetilde{x}(t)$ on an eigenvector becomes , it stays . So without loss of generality we can assume that after the preparation phase, the projection of $\widetilde{x}(t)$ along the top eigenvector $v_{1}$ is non-zero, otherwise we can study the problem in the subspace excluding the top eigenvector.

If $\widetilde{x}(T)\in\cap_{j=1}^{D}\mathcal{I}_{j}$ holds for some $T$ , then for any $t^{\prime},t$ such that $T\leq t\leq t^{\prime}$ and $\left\|\widetilde{x}(t)\right\|\leq 0.5\lambda_{1}$ , it holds $\left|\langle v_{1},\widetilde{x}(t)\rangle\right|\leq\left|\langle v_{1},\widetilde{x}(t^{\prime})\rangle\right|$ .

First, Lemma 3.5 (proved in Appendix A) shows that the norm of the iterate $\widetilde{x}(t)$ remains above $0.5\lambda_{1}$ for only one time-step.

For any $t$ with $\widetilde{x}(t)\in\cap_{j=1}^{D}\mathcal{I}_{j}$ , if $\left\|\widetilde{x}(t)\right\|>\frac{\lambda_{1}}{2}$ , then $\left\|\widetilde{x}(t+1)\right\|\leq\max\left(\frac{\lambda_{1}}{2}-\frac{\lambda_{D}^{2}}{2\lambda_{1}},\lambda_{1}-\left\|\widetilde{x}(t)\right\|\right)$ .

Thus, for any $t$ with $\widetilde{x}(t)\in\cap_{j=1}^{D}\mathcal{I}_{j}$ and $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}$ , either $\left\|\widetilde{x}(t+1)\right\|\leq\frac{\lambda_{1}}{2}$ , or $\left\|\widetilde{x}(t+1)\right\|>\frac{\lambda_{1}}{2}$ , which in turn implies that $\left\|\widetilde{x}(t+2)\right\|\leq\frac{\lambda_{1}}{2}$ by Lemma 3.5. The proof of Lemma 3.4 is completed by induction on Lemma 3.6.

For any step $t$ with $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}$ , for any $k\in\{1,2\}$ , $\left|\langle v_{1},\widetilde{x}(t+k)\rangle\right|\geq\left|\langle v_{1},\widetilde{x}(t)\rangle\right|$ .

Proof of case $k=1$ in Lemma 3.6 follows directly from plugging the assumption $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}$ into (2) (See Lemma A.5). The case of $k=2$ in Lemma 3.6 follows from Lemma A.7. We defer the complete proof of Lemma 3.6 into Appendix A. ∎

To complete the proof for Theorem 3.1, we relate the increase in the projection along $v_{1}$ at any step $t$ , $\left|\langle v_{1},\widetilde{x}(t)\rangle\right|$ , to the magnitude of the angle between $\widetilde{x}(t)$ and the top eigenspace, $\theta_{t}$ . Briefly speaking, we show that if $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}$ , $\left|\langle v_{1},\widetilde{x}(t)\rangle\right|$ has to increase by a factor of $\Theta(\theta^{2}_{t})$ in two steps. Since $\left|\langle v_{1},\widetilde{x}(t)\rangle\right|$ is bounded and monotone increases among $\{t\mid\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}\}$ by Lemma 3.4, we conclude that $\theta_{t}$ gets arbitrarily small for sufficiently large $t$ with $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2},\left\|\widetilde{x}(t+2)\right\|\leq\frac{\lambda_{1}}{2}$ satisfied. Since the one-step normalized GD update Equation 2 is continuous when bounded away from origin, with a careful analysis, we conclude $\theta_{t}\to 0$ for all iterates. Please see Section A.3 for details.

Below we show GD on loss $\sqrt{L}(x)=\sqrt{\frac{1}{2}x^{\top}Ax}$ , Equation 3, follows the same update rule as Normalized GD on $L(x)=\frac{1}{2}x^{\top}Ax$ , up to a linear transformation.

Denoting $\widetilde{x}(t)=\frac{1}{\eta}(2A)^{1/2}x(t)$ , we can easily check $\widetilde{x}(t)$ also satisfies update rule (2).

Main Results

In this section we present the main results of this paper. Section 4.1 is for preliminary and notations. In Section 4.2, we make our key assumptions that the minimizers of the loss function form a manifold. In Sections 4.3 and 4.4 we present our main results for Normalized GD and GD on $\sqrt{L}$ respectively. In Section 4.5 we show the above two settings for GD do enter the regime of Edge of Statbility.

In this section, we focus on the setting where LR $\eta$ goes to 0 and we fix the initialization $x_{\text{init}}$ and the loss function $L$ throughout this paper. We use $O(\cdot),\Omega(\cdot)$ to hide constants about $x_{\text{init}}$ and $L$ .

2 Key Assumptions on Manifold of Local Minimizers

Following Fehrman et al. (2020); Li et al. (2022b), we make the following assumption throughout the paper.

The smoothness assumption is satisfied for networks with smooth activation functions like tanh and GeLU (Hendrycks and Gimpel, 2016). The existence of manifold is due to the vast overparametrization in modern deep networks and preimage theorem. (See a discussion in section 3.1 of Li et al. (2022b)) The assumption $\operatorname{rank}\left(\nabla^{2}L(x)\right)=M$ basically say $\nabla^{2}L(x)$ always attains the maximal rank in the normal space of the manifold, which ensures the differentiability of $\Phi$ and is crucial to our current analysis, though it’s not clear if it is necessary. We also make the following assumption to ensure that $\lambda_{1}(\nabla^{2}L(\cdot))$ is differentiable, which is necessary for our main results, Theorems 4.4 and 4.6.

For any $x\in\Gamma$ , $\nabla^{2}L(x)$ has a positive eigengap, i.e., $\lambda_{1}(\nabla^{2}L(x))>\lambda_{2}(\nabla^{2}L(x))$ .

3 Results for Normalized GD

We first denote the iterates of Normalized GD with LR $\eta$ by $x_{\eta}(t)$ , with $x_{\eta}(0)\equiv x_{\text{init}}$ for all $\eta$ :

The first theorem demonstrates the movement in the manifold, when the iterate travels from $x_{\text{init}}$ to a position that is $O(\eta)$ distance closer to the manifold (more specifically, $\Phi(x_{\text{init}})$ ). Moreover, just like the result in the quadratic case, we have more fine-grained bounds on the projection of $x_{\eta}(t)-\Phi(x_{\eta}(t))$ into the bottom- $k$ eigenspace of $\nabla^{2}L(\Phi(x_{\eta}(t)))$ for every $k\in[D]$ . For convenience, we define the following quantity for all $j\in[d]$ and $x\in U$ :

In the quadratic case, Lemma 3.3 shows that $R_{j}(x)$ will eventually become non-positive for normalized GD iterates. Similarly, for the general loss, the following theorem shows that $R_{j}(x_{\eta}(t))$ eventually becomes approximately non-positive (smaller than $O(\eta^{2})$ ) in $O(\frac{1}{\eta})$ steps.

Our main contribution is the analysis for the second phase (Theorem 4.4), which says just like the quadratic case, the angle between $\widetilde{x}_{\eta}(t)$ and the top eigenspace of $\nabla^{2}L(\Phi(x_{\eta}(t)))$ , denoted by $\theta_{t}$ , will be $O(\eta)$ on average. And as a result, the dynamics of Normalized GD tracks the riemannian gradient flow with respect to $\log(\lambda_{1}(\nabla^{2}L(\cdot)))$ on manifold, that is, the unique solution of Equation 5, where $P^{\perp}_{x,\Gamma}$ is the projection matrix onto the tangent space of manifold $\Gamma$ at $x\in\Gamma$ .

Note Equation 5 is not guaranteed to have a global solution, i.e., a well-defined solution for all $\tau\geq 0$ , for the following two reasons: (1). when the multiplicity of top eigenvalue is larger than $1$ , $\lambda_{1}(\nabla^{2}L(\cdot))$ may be not differentiable and (2). the projection matrix is only defined on $\Gamma$ and the equation becomes undefined when the solution leaves $\Gamma$ , i.e., moving across the boundary of $\Gamma$ . For simplicity, we make 4.2 that every point on $\Gamma$ has a positive eigengap. Or equivalently, we can work with a slightly smaller manifold $\Gamma^{\prime}=\{x\in\Gamma\mid\lambda_{1}(x)>\lambda_{2}(x)\}$ .

Towards a mathematical rigorous characterization of the dynamics in the second phase, we need to make the following modifications: (1). we add negligible noise of magnitude $O(\eta^{100})$ every $\eta^{-0.1}$ steps, (2). we assume for each $\eta>0$ , there exist some step $t=\Theta(1/\eta)$ in phase I, except the guaranteed condition (1) and (2) (by Theorem 4.3, the additional condition (3) also holds. This assumption is mild because we only require (3) to hold for one step among $\Theta(1/\eta)$ steps from $\frac{T_{1}}{\eta}$ to $\frac{T_{1}^{\prime}}{\eta}$ , where $T_{1}$ is the constant given by Theorem 4.3 and $T_{1}^{\prime}$ is arbitrary constant larger than $T_{1}$ . This assumption also holds empirically for all our experiments in Section 6.

4 Results for GD on L𝐿\sqrt{L}

In this subsection, we denote the iterates of GD on $\sqrt{L}$ with LR $\eta$ by $x_{\eta}(t)$ , with $x_{\eta}(0)\equiv x_{\text{init}}$ for all $\eta$ :

Similar to Normalized GD, we will have two phases. The first theorem demonstrates the movement in the manifold, when the iterate travels from $x_{\text{init}}$ to a position that is $O(\eta)$ distance closer to the manifold. For convenience, we will denote the quantity $\sqrt{\sum_{i=j}^{M}\lambda_{i}(x)\langle v_{i}(x),x-\Phi(x)\rangle^{2}}-\eta\sqrt{1/2}\lambda_{j}(x)$ by $\overline{R}_{j}(x)$ for all $j\in[M]$ and $x\in U$ .

The next result demonstrates that close to the manifold, the trajectory implicitly minimizes sharpness.

5 Operating on the Edge of Stability

In this section, we show that both Normalized GD on $L$ and GD on $\sqrt{L}$ is on Edge of Stability in their phase II, that is, at least in one of every two consecutive steps, the stableness is at least $2$ and the loss oscillates in every two consecutive steps. Interestingly, the average loss over two steps decreases over time, even when operating on the edge of Stability (see Figure 1 for illustration), as indicated by the following theorems. Note that Theorems 4.4 and 4.6 ensures that the average of $\theta_{t}$ are $O(\eta)$ and $O(\sqrt{\eta})$ . We defer their proofs into Sections E.5 and G.4 respectively.

Under the setting of Theorem 4.4, by viewing Normalized GD as GD with time-varying LR $\eta_{t}:=\frac{\eta}{\left\|\nabla L(x_{\eta}(t))\right\|}$ , we have $[S_{L}(x_{\eta}(t),\eta_{t})]^{-1}+[S_{L}(x_{\eta}(t+1),\eta_{t+1})]^{-1}=1+O(\theta_{t}+\eta).$ Moreover, we have $\sqrt{L(x_{\eta}(t))}+\sqrt{L(x_{\eta}(t+1))}=\eta\sqrt{\frac{\lambda_{1}(\nabla^{2}L(x_{\eta}(t)))}{2}}+O(\eta\theta_{t})$ .

Under the setting of Theorem 4.6, we have $[S_{\sqrt{L}}(x_{\eta}(t),\eta_{t})]\geq\Omega(\frac{1}{\theta_{t}})$ . Moreover, we have $\sqrt{L(x_{\eta}(t))}+\sqrt{L(x_{\eta}(t+1))}=\eta\lambda_{1}(\nabla^{2}L(x_{\eta}(t)))+O(\eta\theta_{t})$ .

Proof Overview

We sketch the proof of the Normalized GD in phase I and II respectively in Section 5.2. Then we briefly discuss how to prove the results for GD with $\sqrt{L}$ with same analysis in Section 5.3. We start by introducing the properties of limit map of gradient flow $\Phi$ in Section 5.1, which plays a very important role in the analysis.

The limit map of gradient flow $\Phi$ lies at the core of our analysis. When LR $\eta$ is small, one can show $x_{\eta}(t)$ will be $O(\eta)$ close to manifold and $\Phi(x_{\eta}(t))$ . Therefore, $\Phi(x_{\eta}(t))$ captures the essential part of the implicit regularization of Normalized GD and characterization of the trajectory of $\Phi(x_{\eta}(t))$ immediately gives us that of $\Phi(x_{\eta}(t))$ up to $O(\eta)$ .

Below we first recap a few important properties of $\Phi$ that will be used later this section, which makes the analysis of $\Phi(x_{\eta}(t))$ convenient.

Under 4.1, $\Phi$ satisfies the following two properties:

$\partial\Phi(x)\nabla L(x)=0$ for any $x\in U.$ (Lemma B.16)

For any $x\in\Gamma$ , if $\lambda_{1}(x)>\lambda_{2}(x)$ , $\partial^{2}\Phi(x)[v_{1}(x),v_{1}(x)]=-\frac{1}{2}P^{\perp}_{x,\Gamma}\nabla\log\lambda_{1}(x)$ . (Lemmas B.18 and B.20)

Note that $x_{\eta}(t+1)-x_{\eta}(t)=-\eta\frac{\nabla L(x_{\eta}(t))}{\left\|\nabla L(x_{\eta}(t))\right\|}$ , using a second order taylor expansion of $\Phi$ , we have

where we use the first claim of Lemma 5.1 in the final step. Therefore, we have $\Phi(x_{\eta}(t+1))-\Phi(x_{\eta}(t))=O(\eta^{2})$ , which means $\Phi(x_{\eta}(t))$ moves slowly along the manifold, at a rate of at most $O(\eta^{2})$ step. The Taylor expansion of $\Phi$ , (8) plays a crucial role in our analysis for both Phase I and II and will be used repeatedly.

2 Analysis for Normalized GD

The Phase I itself can be divided into two subphases: (A). Normalized GD iterate $x_{\eta}(t)$ gets $O(\eta)$ close to manifold; (B). counterpart of preparation phase in the quadratic case: local movement in the $O(\eta)$ -neighborhood of the manifold which decreases $R_{j}(x_{\eta}(t))$ to $O(\eta^{2})$ . Below we sketch their proofs respectively:

Subphase (A): First, with a very classical result in ODE approximation theory, normalized GD with small LR will track the normalized gradient flow, which is a time-rescaled version of standard gradient flow, with $O(\eta)$ error, and enter a small neighborhoods of the manifold where Polyak-Łojasiewicz (PL) condition holds. Since then, Normalized GD decreases the fast loss with PL condition and the gradient has to be $O(\eta)$ small in $O(\frac{1}{\eta})$ steps. (See details in Section C.1).

Subphase (B): The result in subphase (B) can be viewed as a generalization of Lemma 3.3 when the loss function is $O(\eta)$ -approximately quadratic, in both space and time. More specifically, it means $\left\|\nabla^{2}L(\Phi(x_{\eta}(t)))-\nabla^{2}L(x)\right\|\leq O(\eta)$ for all $x$ which is $O(\eta)$ -close to some $\Phi(x_{\eta}(t^{\prime}))$ with $t^{\prime}-t\leq O(1/\eta)$ . This is because by Taylor expansion (8), $\left\|\Phi(x_{\eta}(t))-\Phi(x_{\eta}(t^{\prime}))\right\|=O(\eta^{2}(t^{\prime}-t))=O(\eta)$ , and again by Taylor expansion of $\nabla^{2}L$ , we know $\left\|\nabla^{2}L(x)-\nabla^{2}L(\Phi(x_{\eta}(t)))\right\|=O(\left\|x-\Phi(x_{\eta}(t))\right\|)=O(\eta)$ .

With a similar proof technique, we show $x_{\eta}(t)$ enters ainvariant set around the manifold $\Gamma$ , that is, $\{x\in U\mid R_{j}(x)\leq O(\eta^{2}),\forall j\in[D]\}$ . Formally, we show the following analog of Lemma 3.3:

Let $\{x_{\eta}(t)\}_{t\geq 0}$ be the iterates of Normalized GD (4) with LR $\eta$ . If for some step $t_{0}$ , $\left\|x_{\eta}(t_{0})-\Phi(x_{\eta}(t_{0}))\right\|=O(\eta)$ , then for sufficiently small LR $\eta$ and all steps $t\in[t_{0}+\Theta(1),\Theta(\eta^{-2})]$ steps, the iterate ${x}_{\eta}(t)$ satisfy $\max_{j\in[M]}R_{j}(x_{\eta}(t))\leq O(\eta^{2})$ .

Analysis for Phase II, Theorem 4.4:

Similar to the subphase (B) in the Phase I, the high-level idea here is again that $x_{\eta}(t)$ locally evolves like normalized GD with quadratic loss around $\Phi(x_{\eta}(t))$ and with an argument similar to the alignment phase of quadratic case (though technically more complicated), we show $x_{\eta}(t)-\Phi(x_{\eta}(t))$ approximately aligns to the top eigenvector of $\nabla^{2}L(\Phi(x_{\eta}(t)))$ , denoted by $v_{1}(t)$ and so does $\nabla L(x_{\eta}(t))$ . More specifically, it corresponds to the second claim in Theorem 4.4, that $\frac{1}{\lfloor T_{2}/\eta^{2}\rfloor}\sum_{t=0}^{\lfloor T_{2}/\eta^{2}\rfloor}\theta_{t}\leq O(\eta)$ .

We now have a more detailed look at the movement in $\Phi$ . Since $\Phi(x_{\eta}(t))$ belongs to the manifold, we have $\nabla L(\Phi(x_{\eta}(t)))=0$ and so $\nabla L(x_{\eta}(t))=\nabla^{2}L(\Phi(x_{\eta}(t)))(x_{\eta}(t)-\Phi(x_{\eta}(t)))+O(\eta^{2})$ using a Taylor expansion. This helps us derive a relation between the Normalized GD update and the top eigenvector of the hessian (simplified version of Lemma B.9):

Incorporating the above into the movement in $\Phi(x_{\eta}(t))$ from Equation 8 gives:

Applying the second property of Lemma 5.1 on Equation 10 above yields Lemma 5.3.

Under the setting in Theorem 4.4, for sufficiently small $\eta$ , we have at any step $t\leq\lfloor T_{2}/\eta^{2}\rfloor$

To complete the proof of Theorem 4.4, we show that for small enough $\eta$ , the trajectory of $\Phi(x_{\eta}(\tau/\eta^{2}))$ is $O(\eta^{3}\lfloor T_{2}/\eta^{2}\rfloor+\eta^{2}\sum_{t=0}^{\lfloor T_{2}/\eta^{2}\rfloor}\theta_{t})$ -close to $X(\tau)$ for any $\tau\leq T_{2}$ , where $X(\cdot)$ is the flow given by Equation 5. This error is $O(\eta)$ , since $\sum_{t=0}^{\lfloor T_{2}/\eta^{2}\rfloor}\theta_{t}=O(\lfloor T_{2}/\eta^{2}\rfloor\eta)$ .

One technical difficulty towards showing the average of $\eta_{t}$ is only $O(\eta)$ is that our current analysis requires $\left|\langle v_{1}(x_{\eta}(t)),x_{\eta}(t)-\Phi(x_{\eta}(t))\rangle\right|$ doesn’t vanish, that is, it remains $\Omega(\eta)$ large throughout the entire training process. This is guaranteed by Lemma 3.4 in quadratic case – since the alignment monotone increases whenever it’s smaller $\frac{\lambda_{1}}{2}$ , but the analysis breaks when the loss is only approximately quadratic and the alignment $\left|\langle v_{1}(x_{\eta}(t)),x_{\eta}(t)-\Phi(x_{\eta}(t))\rangle\right|$ could decrease decrease by $O(\theta_{t}\eta^{2})$ per step. Once the alignment becomes too small, even if the angle $\theta_{t}$ is small, the normalized GD dynamics become chaotic and super sensitive to any perturbation. Our current proof technique cannot deal with this case and that’s the main reason we have to make the additional assumption in Theorem 4.4.

Role of $\eta^{100}$ noise. Fortunately, with the additional assumption that the initial alignment is at least $\Omega(\eta)$ , we can show adding any $\text{poly}(\eta)$ perturbation (even as small as $\Omega(\eta^{100})$ ) suffices to prevent the aforementioned bad case, that is, $\left|\langle v_{1}(x_{\eta}(t)),x_{\eta}(t)-\Phi(x_{\eta}(t))\rangle\right|$ stays $\Omega(\eta)$ large. The intuition why $\Omega(\eta^{100})$ perturbation works again comes from quadratic case – it’s clear that $\widetilde{x}=cv_{1}$ for any $|c|\leq 1$ is a stationary point for two-step normalized GD updates for quadratic loss under the setting of Section 3. But if $c$ is smaller than critical value determined by the eigenvalues of the hessian, the stationary point is unstable, meaning any deviation away from the top eigenspace will be amplified until the alignment increases above the critical threshold. Based on this intuition, the formal argument, Lemma E.11 uses the techniques from the ‘escaping saddle point’ analysis (Jin et al., 2017). Adding noise is not necessary in experiments to observe the predicted behavior (see ‘Alignment’ in Figure 4 where no noise is added). On one hand, it might be because the floating point errors served the role of noise. On the other hand, we suspect it’s not necessary even for theory, just like GD gets stuck at saddle point only when initialized from a zero measure set even without noise (Lee et al., 2016, 2017).

3 Analysis for GD on L𝐿\sqrt{L}

In this subsection we will make an additional assumption that $L(x)=0$ for all $x\in\Gamma$ . The analysis then will follow a very similar strategy as the analysis for (Normalized) GD. However, the major difference from the analysis for Normalized GD comes from the update rule for $x_{\eta}(t)$ when it is $O(\eta)$ -close to the manifold:

Thus, the effective learning rate is $\sqrt{\lambda_{1}(t)}\eta$ at any step $t$ . This shows up, when we compute the change in the function $\Phi$ . Thus, we have the following lemma showcasing the movement in the function $\Phi$ with the GD update on $\sqrt{L}$ :

Under the setting in Theorem 4.6, for sufficiently small $\eta$ , we have at any step $t\leq\lfloor T_{2}/\eta^{2}\rfloor$ , $\Phi(x_{\eta}(t+1))-\Phi(x_{\eta}(t))=-\frac{\eta^{2}}{8}P^{\perp}_{t,\Gamma}\nabla\lambda_{1}(t)+O(\eta^{3}+\eta^{2}\theta_{t}).$

Experiments

Though our main theorems characterizes the dynamics of Nomalized GD and GD on $\sqrt{L}$ for sufficiently small LR, it’s not clear if the predicted phenomena is related to the training with practical LR as the function and initialization dependent constants are hard to compute and could be huge. Neverthesless, in this section we show the phenomena predicted by our theorem does occur for real-life models like VGG-16. We further verify the predicted convergence to the limiting flow for Normalized GD on a two-layer fully-connected network trained on MNIST.

We observe that the alignment function reaches close to $1$ , towards the end of training. The top eigenvalue decreases over time (as predicted byTheorem 4.4 and Theorem 4.6), and the stableness hovers around $2$ at the end of training.

Verifying Convergence to Limiting Flow on MNIST:

Conclusion

The recent discovery of Edge of Stability phenomenon in Cohen et al. (2021) calls for a reexamination of how we understand optimization in deep learning. The current paper gives two concrete settings with fairly general loss functions, where gradient updates can be shown to decrease loss over many iterations even after stableness is lost. Furthermore, in one setting the trajectory is shown to amount to reduce the sharpness (i.e., the maximum eigenvalue of the Hessian of the loss), thus rigorously establishing an effect that has been conjectured for decades in deep learning literature and was definitively documented for GD in Cohen et al. (2021). Our analysis crucially relies upon learning rate $\eta$ being finite, in contrast to many recent results on implicit bias that required an infinitesimal LR. Even the alignment analysis of Normalized GD to the top eigenvector for quadratic loss in Section 3 appears to be new.

One limitation of our analysis is that it only applies close to the manifold of local minimizers. By contrast, in experiments the EoS phenomenon, including the control of sharpness, begins much sooner. Addressing this gap, as well as analysing the EoS for the loss $L$ itself (as opposed to $\sqrt{L}$ as done here) is left for future work. Very likely this will require novel understanding of properties of deep learning losses, which we were able to circumvent by looking at $\sqrt{L}$ instead. Exploration of EoS-like effects in SGD setting would also be interesting, although we first need definitive experiments analogous to Cohen et al. (2021).

Acknowledgement

We thank Kaifeng Lyu 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 Omitted Proofs for Results for Quadratic Loss Functions

Recall the loss function $L$ is defined as $L(x)=\frac{1}{2}x^{\top}Ax$ . The Normalized GD update (LR= $\eta$ )is given by $x(t+1)=x(t)-\eta\frac{Ax(t)}{\left\|Ax(t)\right\|}$ . A substitution $\widetilde{x}(t):=\frac{Ax(t)}{\eta}$ gives the following update rule:

Now we recall the main theorem for Normalized GD on quadratic loss functions: See 3.1

We also note that GD on $\sqrt{L}$ with any LR $\eta$ can also be reduced to update rule (2), as shown in the discussion at the end of Section 3.

In this subsection, we show (1). $\mathcal{I}_{j}$ is indeed an invariant set for normalized GD $\forall j\in[D]$ and (2). from any initialization, normalized GD will eventually go into their intersection $\cap_{j=1}^{D}\mathcal{I}_{j}$ .

Note $P^{(j:D)}A=P^{(j:D)}AP^{(j:D)}$ , by definition of Normalized GD (2), we have

Note that $P^{(j:D)}A\preccurlyeq\lambda_{j}I$ , $\left\|P^{(j:D)}\widetilde{x}(t)\right\|\leq\left\|\widetilde{x}(t)\right\|$ and $\left\|P^{(j:D)}\widetilde{x}(t)\right\|\leq\lambda_{j}$ by assumption, we have

Therefore $\left\|I-\frac{P^{(j:D)}A}{\left\|\widetilde{x}(t)\right\|}\right\|\leq\frac{\lambda_{j}}{\left\|P^{(j:D)}\widetilde{x}(t)\right\|}$ and thus we conclude ${\left\|P^{(j:D)}\widetilde{x}(t+1)\right\|}\leq\lambda_{j}$ . ∎

Since $\lambda_{j}\leq\left\|P^{(j:D)}\widetilde{x}(t)\right\|\leq\left\|\widetilde{x}(t)\right\|$ , we have $0\preccurlyeq I-\frac{P^{(j:D)}A}{\left\|\widetilde{x}(t)\right\|}\preccurlyeq 1-\frac{\lambda_{D}}{\left\|\widetilde{x}(t)\right\|}$ . Therefore $\left\|I-\frac{P^{(j:D)}A}{\left\|\widetilde{x}(t)\right\|}\right\|\leq 1-\frac{\lambda_{D}}{\left\|\widetilde{x}(t)\right\|}$ . The proof is completed by plugging this into Equation 11. ∎

Lemma A.2 has the following two direct corollaries.

For any initialization $\widetilde{x}(0)$ and $t\geq\frac{\left\|\widetilde{x}(0)\right\|-\lambda_{1}}{\lambda_{D}}$ , $\left\|\widetilde{x}(t)\right\|\leq\lambda_{1}$ , that is, $\widetilde{x}(t)\in\mathcal{I}_{1}$ .

Set $j=1$ in Lemma A.2, it holds that $\left\|\widetilde{x}(t+1)\right\|\leq\left\|\widetilde{x}(t)\right\|-\lambda_{D}$ whenever $\left\|\widetilde{x}(t)\right\|\geq\lambda_{1}$ . Thus $\left\|\widetilde{x}(\big{\lceil}\frac{\left\|\widetilde{x}(0)\right\|-\lambda_{1}}{\lambda_{D}}\big{\rceil})\right\|\leq\lambda_{1}$ . The proof is completed as $\mathcal{I}_{1}$ is an invariant set by Lemma A.1. ∎

For any coordinate $j\in[D]$ and initial point $\widetilde{x}(0)\in\mathcal{I}_{1}$ , if $t\geq\frac{\lambda_{1}}{\lambda_{D}}\ln\frac{\lambda_{1}}{\lambda_{j}}$ then $\left\|P^{(j:D)}\widetilde{x}(t)\right\|\leq\lambda_{j}$ .

Since $\mathcal{I}_{1}$ is an invariant set, we have $\left\|\widetilde{x}(t)\right\|\leq\lambda_{1}$ for all $t\geq 0$ . Thus let $T=\lfloor\frac{\lambda_{1}}{\lambda_{D}}\ln\frac{\lambda_{1}}{\lambda_{j}}\rfloor$ , we have

The proof is completed since $I_{j}$ is a invariant set for any $j\in[D]$ by Lemma A.1. ∎

A.2 Proofs for Alignment Phase

In this subsection, we analyze how normalized GD align to the top eigenvector once it goes through the preparation phase, meaning $\widetilde{x}(t)\in\cap_{j=1}^{D}\mathcal{I}_{j}$ for all $t$ in alignment phase.

Let the index $k$ be the smallest integer such that $\lambda_{k+1}<2\left\|\widetilde{x}(t)\right\|-\lambda_{1}$ . If no such index exists, then one can observe that $\left\|\widetilde{x}(t+1)\right\|\leq\lambda_{1}-\left\|\widetilde{x}(t)\right\|$ . Assuming that such an index exists in $[D]$ , we have $\lambda_{k}\geq 2\left\|\widetilde{x}(t)\right\|-\lambda_{1}$ and $\left\|\widetilde{x}(t)\right\|-\lambda_{j}\leq\lambda_{1}-\left\|\widetilde{x}(t)\right\|$ , $\forall j\leq k$ . Now consider the following vectors:

By definition of $k$ , $|\left\|\widetilde{x}(t)\right\|-\lambda_{j}|\leq|\left\|\widetilde{x}(t)\right\|-\lambda_{1}|$ . Thus

By assumption, we have $\widetilde{x}(t)\in\cap_{j=1}^{D}\mathcal{I}_{j}$ . Thus

where we applied AM-GM inequality multiple times in the pre-final step.

where the final step is because $\frac{\lambda_{1}}{2}\leq\left\|\widetilde{x}(t)\right\|\leq\lambda_{1}$ and that the maximal value of a convex function is attained at the boundary of an interval.

At any step $t$ and $i\in[D]$ , if $\left\|\widetilde{x}(t)\right\|\gtreqqless\frac{\lambda_{i}}{2}$ , then $\left|\widetilde{x}_{i}(t+1)\right|\lesseqqgtr\left|\widetilde{x}_{i}(t)\right|$ , where $\gtreqqless$ denotes larger than, equal to and smaller than respectively. (Same for $\lesseqqgtr$ , but in the reverse order)

From the Normalized GD update rule, we have $\widetilde{x}_{i}(t+1)=\widetilde{x}_{i}(t)\left(1-\frac{\lambda_{i}}{\left\|\widetilde{x}(t)\right\|}\right),\text{ for all }i\in[D]$ . Thus

At any step $t$ , if $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}$ , then

where $\theta_{t}=\arctan\frac{\left\|P^{(2:D)}\widetilde{x}(t)\right\|}{\left|e_{1}^{\top}\widetilde{x}(t)\right|}$ and $\lambda=\min(\lambda_{1}-\lambda_{2},\lambda_{D})$ .

We first show that the left side inequality holds by the following update rule for $\langle e_{1},\widetilde{x}(t)\rangle$ :

Since $\left\|\widetilde{x}(t+1)\right\|\geq\left|\langle e_{1},\widetilde{x}(t+1)\rangle\right|$ and $\theta_{t}$ denotes the angle between $e_{1}$ and $\widetilde{x}(t+1)$ , we get the left side inequality.

Now, we focus on the right hand side inequality. First of all, the update in the coordinate $j\in[2,D]$ is given by

where again in the final step, we have used $\left\|\widetilde{x}(t)\right\|<\frac{\lambda_{1}}{2}$ . The above bound can be further bounded by

where we have used $\lambda=\min(\lambda_{1}-\lambda_{2},\lambda_{D}).$

If at some step $t$ , $\left\|\widetilde{x}(t+1)\right\|+\left\|\widetilde{x}(t)\right\|\leq\lambda_{1}$ , then $\left|\widetilde{x}_{1}(t+2)\right|\geq\left|\widetilde{x}_{1}(t)\right|$ , where the equality holds only when $\left\|\widetilde{x}(t+1)\right\|+\left\|\widetilde{x}(t)\right\|=\lambda_{1}$ . Therefore, by Lemma A.6, we have :

where $\theta_{t}=\arctan\frac{\left\|P^{(2:D)}\widetilde{x}(t)\right\|}{\left|e_{1}^{\top}\widetilde{x}(t)\right|},$ and $\lambda=\min(\lambda_{1}-\lambda_{2},\lambda_{D})$ .

Using the Normalized GD update rule, we have

where the equality holds only when $\left\|\widetilde{x}(t+1)\right\|+\left\|\widetilde{x}(t)\right\|=\lambda_{1}$ .

Moreover, with the additional condition that $\left\|\widetilde{x}(t)\right\|<\frac{\lambda_{1}}{2}$ , we have from Lemma A.6, $\left\|\widetilde{x}(t+1)\right\|\leq\lambda_{1}-\left\|\widetilde{x}(t)\right\|-\lambda(\lambda_{1}-\lambda)\sin^{2}\theta_{t}$ , where $\lambda=\min(\lambda_{1}-\lambda_{2},\lambda_{D})$ .

Hence, retracing the steps we followed before, we have

where the final step follows from $\left\|\widetilde{x}(t+1)\right\|\leq\lambda_{1}-\left\|\widetilde{x}(t)\right\|$ and therefore $\left\|\widetilde{x}(t+1)\right\|\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}^{2}}{4}$ . ∎

A.3 Proof of Main theorems for Quadratic Loss

Preparation phase: $\widetilde{x}(t)$ enters and stays in an invariant set around the origin, that is, $\cap_{j=1}^{D}\mathcal{I}_{j}$ , where $\mathcal{I}_{j}:=\{\widetilde{x}\mid\sum_{i=j}^{D}\langle e_{i},\widetilde{x}(t)\rangle^{2}\leq\lambda^{2}_{j}\}$ . (See Lemma 3.3, which is a direct consequence of Lemmas A.1, A.3 and A.1.)

Alignment phase: The projection of $\widetilde{x}(t)$ on the top eigenvector, $|\left\langle\widetilde{x}(t),e_{1}\right\rangle|$ , is shown to increase monotonically among the steps among the steps $\{t\mid\left\|\widetilde{x}(t)\right\|\leq 0.5\}$ , up until convergence, since it’s bounded. (Lemma 3.4)

By Lemma A.7, the convergence of $|\left\langle\widetilde{x}(t),e_{1}\right\rangle|$ would imply the convergence of $\widetilde{x}(t)$ to $e_{1}$ in direction.

Now we claim $\forall t\geq 3$ , there is some $k\in\{0,1,3\}$ such that $t-k\in S^{\prime}$ . This is because Lemma 3.5 says that if $t\notin S$ , then both $t-1,t+1\in S$ . Thus for any $t\notin S$ , $t-1\in S^{\prime}$ . Therefore, for any $t\in S/S^{\prime}$ , if $t-2\notin S$ , then $t-3\in S^{\prime}$ . Thus we conclude that $\forall t\geq 3$ , there is some $k\in\{0,1,3\}$ such that $t-k\in S^{\prime}$ , which implies $\lim\limits_{t\to\infty}\theta_{t}=0$ . Hence $\lim\limits_{t\to\infty}\left\|\widetilde{x}(t+1)-\widetilde{x}(t)\right\|=\lambda_{1}$ , meaning for sufficiently large $t$ , $\widetilde{x}_{1}(t)$ flips its sign per step and thus $\lim\limits_{t\to\infty}\widetilde{x}(t+2)-\widetilde{x}(t)=0$ , $\lim\limits_{t\to\infty}\left\|\widetilde{x}(t+1)\right\|+\left\|\widetilde{x}(t)\right\|=\lambda_{1}$ .

If $C=\frac{1}{2}$ , then we must have $\lim\limits_{t\to\infty}\left\|\widetilde{x}(t)\right\|=\frac{\lambda_{1}}{2}$ and we are done in this case. If $C<\frac{1}{2}$ , note that $\lim\limits_{t\to\infty,t\in S^{\prime}}\left|\widetilde{x}_{1}(t)\right|=C\lambda_{1}$ , it must hold that $\lim\limits_{t\to\infty,t\in S^{\prime}}\left\|\widetilde{x}(t+1)\right\|=(1-C)\lambda_{1}$ , thus there is some large $T\in S$ such that for all $t\in S,t\geq T$ , $t+1\notin S$ . By Lemma 3.5, $t+2\in S$ . Thus we conclude $\lim\limits_{t\to\infty}\widetilde{x}(T+2t)=C\lambda se_{1}$ for some $s\in\{-1,1\}$ and thus $\lim\limits_{t\to\infty}\widetilde{x}(T+2t+1)=(C-1)\lambda se_{1}$ . This completes the proof. ∎

A.4 Some Extra Lemmas (only used in the general loss case)

For a general loss function $L$ satisfying 4.1, the loss landscape looks like a strongly convex quadratic function locally around its minimizer. When sufficient small learning rate, the dynamics will be sufficiently close to the manifold and behaves like that in quadratic case with small perturbations. Thus it will be very useful to have more refined analysis for the quadratic case, as they allow us to bound the error in the approximate quadratic case quantitatively. Lemmas A.8, A.9, A.10 and A.11 are such examples. Note that they are only used in the proof of the general loss case, but not in the quadratic loss case.

Lemma A.8 is a slightly generalized version of Lemma 3.5.

Suppose at time $t$ , $\left\|P^{(j:D)}\widetilde{x}(t)\right\|\leq\lambda_{j}(1+\frac{\lambda^{2}_{D}}{\lambda^{2}_{1}}),\text{ for all }j\in[D]$ , if $\left\|\widetilde{x}(t)\right\|>\frac{\lambda_{1}}{2}$ , then $\left\|\widetilde{x}(t+1)\right\|\leq\frac{\lambda_{1}}{2}$ .

The proof is similar to the proof of Lemma 3.5. Let the index $k$ be the smallest integer such that $\lambda_{k+1}<2\left\|\widetilde{x}(t)\right\|-\lambda_{1}$ . If no such index exists, then one can observe that $\left\|\widetilde{x}(t+1)\right\|\leq\lambda_{1}-\left\|\widetilde{x}(t)\right\|$ . Assuming that such an index exists in $[D]$ , we have $\lambda_{k}\geq 2\left\|\widetilde{x}(t)\right\|-\lambda_{1}$ and $\left\|\widetilde{x}(t)\right\|-\lambda_{j}\leq\lambda_{1}-\left\|\widetilde{x}(t)\right\|$ , $\forall j\leq k$ . With the same decomposition and estimation, since $\widetilde{x}(t)\in\cap_{j=1}^{D}(1+\frac{\lambda_{D}^{2}}{\lambda_{1}^{2}})\mathcal{I}_{j}$ , we have

$\left|\langle e_{1},\widetilde{x}(t)\rangle\right|\leq(1-2c)g(\lambda_{k})$ .

$\theta_{t}\leq\sqrt{c\left|\langle e_{1},\widetilde{x}(t)\rangle\right|}$ ,

where $\theta_{t}=\arctan\frac{\left\|P^{(2:D)}(\widetilde{x}(t))\right\|}{\left|\langle e_{1},\widetilde{x}(t)\rangle\right|}$ .

From the quadratic update, we have the update rule as:

Thus, as long as, the following holds true:

We can use $(\lambda_{1}-\left\|\widetilde{x}(t)\right\|)\cos\theta_{t}\leq\left\|\widetilde{x}(t+1)\right\|\leq\lambda_{1}-\left\|\widetilde{x}(t)\right\|-\frac{\lambda}{2\lambda_{1}}\left(1-\frac{\lambda}{\lambda_{1}}\right)\lambda_{1}\sin^{2}\theta_{t}$ , where $\lambda=\min(\lambda_{1}-\lambda_{2},\lambda_{D})$ from Lemma A.6 to show the following with additional algebraic manipulation:

where the last step we use that $\left|\theta_{t}\right|\leq\sqrt{c\left|\langle e_{1},\widetilde{x}(t)\rangle\right|}$ , we only need

The above inequality is true when $\left|\langle e_{1},\widetilde{x}(t)\rangle\right|\leq\left(1-2c\right)g(\lambda_{k})$ . ∎

Then, the following must hold true at time $t$ .

where the final step holds true for any $c\in(0,1).$

The result follows after substituting this bound in Equation 12.

At any step $t$ , if $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}$ ,

$\left|\tan(\angle(\widetilde{x}(t+1),e_{1}))\right|\leq\max(\frac{\lambda_{2}}{\lambda_{1}},1-2\frac{\lambda_{D}}{\lambda_{1}})\left|\tan(\angle(\widetilde{x}(t),e_{1}))\right|$ .

$\left|\tan(\angle(\widetilde{x}(t+2),e_{1}))\right|\leq\frac{\lambda_{1}}{\left\|\widetilde{x}(t)\right\|}\left|\tan(\angle(\widetilde{x}(t),e_{1}))\right|$ .

From the Normalized GD update rule, we have

implying $\left|\widetilde{x}_{i}(t+1)\right|<\left|\left(1-\frac{1}{\left\|\widetilde{x}(t)\right\|}\right)\right|\left|\widetilde{x}_{i}(t)\right|$ for all $i\in[2,D]$ , since $\lambda_{i}<1$ .

Since $\lambda_{i}<\lambda_{1}$ and $\left\|\widetilde{x}(t)\right\|\leq\frac{\lambda_{1}}{2}$ , it holds that

Recall $\left|\tan(\angle(v,e_{1}))\right|=\frac{\left\|P^{(2:D)}v\right\|}{\left|\left\langle e_{1},v\right\rangle\right|}$ for any vector $v$ , the first claim follows from re-arranging the terms.

For the second claim, it suffices to apply the above inequality to $t+1$ , which yields that

The proof is completed by noting $\left\|\widetilde{x}(t+1)\right\|\leq\lambda_{1}-\left\|\widetilde{x}{t}\right\|$ (Lemma A.6) and $\tan(\angle(\widetilde{x}(t+1),e_{1}))\leq\tan(\angle(\widetilde{x}(t),e_{1}))$ . ∎

Appendix B Setups for General Loss Functions

Before we start the analysis for Normalized GD for general loss functions in Appendix C, we need to introduce some new notations and terminologies to complete the formal setup. We start by first recapping some core assumptions and definitions in the main paper and provide the missing proof in the main paper.

We define $\Phi$ as the limit map of gradient flow below. We summarize various properties of $\Phi$ from LABEL:ch:diffusion_on_manifold in Section B.2.

Given any two points $x,y$ , we use $\overline{xy}$ to denote the line segment between $x$ and $y$ , i.e., $\{z\mid\exists\lambda\in,z=(1-\lambda)x+\lambda y\}$ .

The main result of this chapter focuses on the trajectory of Normalized GD from fixed initialization $x_{\text{init}}$ with LR $\eta$ converges to , which can be roughly split into two phases. In the first phase, Theorem 4.3 shows that the normalized GD trajectory converges to the gradient flow trajectory, $\phi(x_{\text{init}},\cdot)$ . In second phase, Theorem 4.4 shows that the normalized GD trajectory converges to the limiting flow which decreases sharpness on $\Gamma$ , (5). Therefore, for sufficiently small $\eta$ , the entire trajectory of normalized GD will be contained in a small neighbourhood of gradient flow trajectory $Z$ and limiting flow trajectory $Y$ . The convergence rate given by our proof depends on the various local constants like smoothness of $L$ and $\Phi$ in this small neighbourhood, which intuitively can be viewed as the actual ”working zone” of the algorithm. The constants are upper bounded or lower bounded from zero because this ”working zone” is compact after fixing the stopping time of (5), which is denoted by $T_{2}$ .

We construct the ”working zone” of the second phase, $Y^{\bm{\rho}}$ and $Y^{\bm{\epsilon}}$ in Lemmas B.2 and B.5 respectively, where $0<{\bm{\epsilon}}<{\bm{\rho}}$ , implying $Y^{\bm{\epsilon}}\subset Y^{\bm{\rho}}$ . The reason that we need the two-level nested ”working zones” is that even though we can ensure all the points in $Y^{\bm{\rho}}$ have nice properties as listed in Lemma B.2, we cannot ensure the trajectory of gradient flow from $x\in Y^{\bm{\rho}}$ to $\Phi(x)$ or the line segment $\overline{x\Phi(x)}$ is in $Y^{\bm{\rho}}$ , which will be crucial for the geometric lemmas (in Section B.1) that we will heavily use in the trajectory analysis around the manifold. For this reason we further define $Y^{\bm{\epsilon}}$ and Lemma B.5 guarantees the trajectory of gradient flow from $x$ to $\Phi(x)$ or the line segment $\overline{x\Phi(x)}$ whenever $x\in Y^{\bm{\rho}}$ .

A function $L$ is said to be $\bm{\mu}$ -PL in a set $U$ iff for all $x\in U$ ,

For convenience, we define $\bm{\Delta}:=\frac{1}{2}\inf_{x\in Y}\big{(}\lambda_{1}(\nabla^{2}L(x))-\lambda_{2}(\nabla^{2}L(x)))\big{)}$ and $\bm{\mu}:=\frac{1}{4}\inf_{x\in Y}\lambda_{M}(\nabla^{2}L(x))$ . By 4.1, we have $\bm{\mu}>0$ . By 4.2, $\bm{\Delta}>0$ .

Given $Y$ , there are sufficiently small ${\bm{\rho}}>0$ such that

$L$ is $\bm{\mu}$ -PL on $Y^{\bm{\rho}}$ ; (see Definition B.1)

$\inf_{x\in Y^{\bm{\rho}}}\big{(}\lambda_{1}(\nabla^{2}L(x))-\lambda_{2}(\nabla^{2}L(x)))\big{)}\geq\bm{\Delta}>0$ ;

$\inf_{x\in Y^{\bm{\rho}}}\lambda_{M}(\nabla^{2}L(x))\geq\bm{\mu}>0$ .

We first claim for every $y\in Y$ , for all sufficiently small ${\bm{\rho}}_{y}>0$ (i.e. for all ${\bm{\rho}}_{y}$ smaller than some threshold depending on $y$ ), the following three properties hold (1) $\overline{B}_{y}({{\bm{\rho}}_{y}})\cap\Gamma$ is compact; (2) $\overline{B}_{y}({{\bm{\rho}}_{y}})\cap\Gamma\subset U$ and (3) $L$ is $\bm{\mu}$ -PL on $\overline{B}_{y}({{\bm{\rho}}_{y}}\cap\Gamma)$ .

Thus for any $c^{\prime}>0$ , for sufficiently small ${\bm{\rho}}_{y}$ , $(x-p(x))^{\top}\nabla^{2}L(p(x))(x-p(x))\geq c^{\prime}\left\|x-p(x)\right\|^{3}$ . Combining Equations 14 and 15, we conclude that for sufficiently small ${\bm{\rho}}_{y}$ ,

Again for sufficiently small ${\bm{\rho}}_{y}$ , by Taylor expansion of $L$ at $p(x)$ , we have

Meanwhile, since $\lambda_{M}(\nabla^{2}L(p(x)))$ and $\lambda_{1}(\nabla^{2}L(p(x)))-\lambda_{2}(\nabla^{2}L(p(x)))$ are continuous functions in $x$ , we can also choose a sufficiently small ${\bm{\rho}}_{y}$ such that for all $x\in\overline{B}_{y}({\bm{\rho}}_{y})$ , $\lambda_{M}(\nabla^{2}L(p(x)))\geq\frac{1}{2}\lambda_{M}(\nabla^{2}L(p(y)))=\frac{1}{2}\lambda_{M}(\nabla^{2}L(y))>\bm{\Delta}$ and $\lambda_{1}(\nabla^{2}L(p(x)))-\lambda_{2}(\nabla^{2}L(p(x)))\geq\frac{1}{2}\big{(}\lambda_{1}(\nabla^{2}L(p(y)))-\lambda_{2}(\nabla^{2}L(p(y)))\big{)}=\frac{1}{2}\big{(}\lambda_{1}(\nabla^{2}L(y))-\lambda_{2}(\nabla^{2}L(y))\big{)}\geq\bm{\mu}$ . Further note $Y\subset\cup_{y\in Y}B_{y}({\bm{\rho}}_{y})$ and $Y$ is a compact set, we can take a finite subset of $Y$ , $Y^{\prime}$ , such that $Y\subset\cup_{y\in Y^{\prime}}B_{y}({\bm{\rho}}_{y})$ . Taking ${\bm{\rho}}:=\min_{y\in Y^{\prime}}\frac{{\bm{\rho}}_{y}}{2}$ completes the proof. ∎

Here, $\mathcal{T}\left(x^{1},\cdots,x^{(k)}\right)=\sum_{i_{1}=1}^{d_{1}}\sum_{i_{2}=1}^{d_{2}}\ldots\sum_{i_{k}=1}^{d_{k}}t_{i_{1}i_{2}\ldots i_{d}}x_{i_{1}}^{(1)}x_{i_{2}}^{(2)}\ldots x_{i_{k}}^{(k)}$ .

We define the following constants regarding smoothness of $L$ and $\Phi$ of various orders over $Y^{\bm{\rho}}$ .

Given ${\bm{\rho}}$ as defined in Lemma B.2, there is an ${\bm{\epsilon}}\in(0,{\bm{\rho}})$ such that

$\sup_{x\in Y^{\bm{\epsilon}}}L(x)-\inf\limits_{x\in Y^{\bm{\epsilon}}}L(x)<\min(\frac{\bm{\mu}{\bm{\rho}}^{2}}{8},\frac{2\bm{\mu}^{5}}{\bm{\nu}^{2}\bm{\zeta}^{2}})$ ;

$\forall x\in Y^{\bm{\epsilon}}$ , $\Phi(x)\in Y^{\frac{{\bm{\rho}}}{2}}$ .

For every $y\in Y$ , there is an ${\bm{\epsilon}}_{y}$ , such that $\forall x\in B_{y}({\bm{\epsilon}}_{y})$ , it holds that $L(x)<\min(\frac{\bm{\mu}{\bm{\rho}}^{2}}{8},\frac{2\bm{\mu}^{4}}{\bm{\nu}^{2}\bm{\zeta}^{2}})$ and $\Phi(x)\in Y^{\frac{{\bm{\rho}}}{2}}$ , as both $L(x)$ and $\Phi(x)$ are continuous. Further note $Y\subset\cup_{y\in Y}B_{y}({\bm{\epsilon}}_{y})$ and $Y$ is a compact set, we can take a finite subset of $Y$ , $Y^{\prime}$ , such that $Y\subset\cup_{y\in Y^{\prime}}B_{y}({\bm{\epsilon}}_{y})$ . Taking ${\bm{\epsilon}}:=\min_{y\in Y^{\prime}}\frac{{\bm{\epsilon}}_{y}}{2}$ completes the proof. ∎

Summary for Setups:

The initial point $x_{\text{init}}$ is chosen from an open neighborhood of manifold $\Gamma$ , $U$ , where the infinite-time limit of gradient flow $\Phi$ is well-defined and for any $x\in U$ , $\Phi(x)\in\Gamma$ . We consider normalized GD with sufficiently small LR $\eta$ such that the trajectory enters a small neighborhood of limiting flow trajectory, $Y^{\bm{\rho}}$ . Moreover, $L$ is $\bm{\mu}$ -PL on $Y^{\bm{\rho}}$ and the eigengaps and smallest eigenvalues are uniformly lower bounded by positive $\bm{\Delta},\bm{\mu}$ respectively on $Y^{\bm{\rho}}$ . Finally, we consider a proper subset of $Y^{\bm{\rho}}$ , $Y^{\bm{\epsilon}}$ , as the final ”working zone” in the second phase (defined in Lemma B.5), which enjoys more properties than $Y^{\bm{\rho}}$ , including Lemmas B.7, B.8, B.9 and B.10.

B.1 Geometric Lemmas

In this subsection we present several geometric lemmas which are frequently used in the trajectory analysis of normalized GD. In this section, $O(\cdot)$ only hides absolute constants. Below is a brief summary:

Lemma B.6: Inequalities connecting various terms: the distance between $x$ and $\Phi(x)$ , the length of GF trajectory from $x$ to $\Phi(x)$ , square root of loss and gradient norm;

Lemma B.7: For any $x\in Y^{\bm{\epsilon}}$ , the gradient flow trajectory from $x$ to $\Phi(x)$ and the line segment between $x$ and $\Phi(x)$ are all contained in $Y^{\bm{\rho}}$ , so it’s ”safe” to use Taylor expansions along GF trajectory or $\overline{x\Phi(x)}$ to derive properties;

Lemmas B.8, B.9 and B.10: for any $x\in Y^{\bm{\epsilon}}$ , the normalized GD dynamics at $x$ can be roughly viewed as approximately quadratic around $\Phi(x)$ with positive definite matrix $\nabla^{2}L(\Phi(x))$ .

Lemma B.11: In the ”working zone”, $Y^{\bm{\rho}}$ , one-step normalized GD update with LR $\eta$ only changes $\Phi(x_{t})$ by $O(\eta^{2})$ .

Lemma B.13: In the ”working zone”, $Y^{\bm{\rho}}$ , one-step normalized GD update with LR $\eta$ decreases $\sqrt{L(x)-\min_{y\in Y}L(y)}$ by $\eta\frac{\sqrt{2\bm{\mu}}}{4}$ if $\left\|\nabla L(x)\right\|\geq\frac{\bm{\zeta}}{\eta}$ .

If the trajectory of gradient flow starting from $x$ , $\phi(x,t)$ , stays in $Y^{\bm{\rho}}$ for all $t\geq 0$ , then we have

Since $\Phi(x)$ is defined as $\lim_{t\to\infty}\phi(x,t)$ and $\phi(x,0)=x$ , the left-side inequality follows immediately from triangle inequality. The right-side inequality is by the definition of PL condition. Below we prove the middle inequality.

Since $\forall t\geq 0$ , $\phi(x,t)\in Y^{\bm{\rho}}$ , it holds that $\left\|\nabla L(\phi(x,t))\right\|^{2}\geq 2\bm{\mu}(L(\phi(x,t))-L(\Phi(x)))$ by the choice of ${\bm{\rho}}$ in Lemma B.2. Without loss of generality, we assume $L(y)=0,\forall y\in\Gamma$ . Thus we have

The proof is complete since $\phi(x,0)=x$ and we assume $L(\Phi(x))$ is . ∎

Let ${\bm{\rho}},{\bm{\epsilon}}$ be defined in Lemmas B.2 and B.5. For any $x\in Y^{\bm{\epsilon}}$ , we have

The entire trajectory of gradient flow starting from $x$ is contained in $Y^{\bm{\rho}}$ , i.e., $\phi(x,t)\in Y^{\bm{\rho}}$ , $\forall t\geq 0$ ;

Moreover, $\left\|\Phi(x)-\phi(x,t)\right\|\leq\min({\bm{\rho}},\frac{2\bm{\mu}^{2}}{\bm{\nu}\bm{\zeta}})$ , $\forall t\geq 0$ .

Let time $\tau^{*}\geq 0$ be the smallest time after which the trajectory of GF is completely contained in $Y^{\bm{\rho}}$ , that is, $\tau^{*}:=\inf\{t\geq 0\mid\forall t^{\prime}\geq t,\phi(x,t^{\prime})\in Y^{\bm{\rho}}\}$ . Since $Y^{\bm{\rho}}$ is closed and $\phi(x,\cdot)$ is continuous, we have $\phi(x,\tau^{*})\in Y^{\bm{\rho}}$ .

Since $\forall\tau\geq\tau^{*}$ , $\phi(x,\tau)\in Y^{\bm{\rho}}$ , by Lemma B.6, it holds that $\left\|\phi(x,\tau^{*})-\Phi(x)\right\|\leq\sqrt{\frac{2(L(\phi(x,\tau^{*}))-L(\Phi(x)))}{\bm{\mu}}}$ .

Note that loss doesn’t increase along GF, we have $L(\phi(x,\tau^{*}))-L(\Phi(x))\leq L(x)-L(\Phi(x))\leq\frac{\bm{\mu}{\bm{\rho}}^{2}}{8}$ , which implies that $\left\|\phi(x,\tau^{*})-\Phi(x)\right\|\leq\frac{{\bm{\rho}}}{2}$ . Therefore $\tau^{*}$ must be , otherwise there exists a $0<\tau^{\prime}<\tau^{*}$ such that $\left\|\phi(x,\tau)-\Phi(x)\right\|\leq{\bm{\rho}}$ for all $\tau^{\prime}<\tau<\tau^{*}$ by the continuity of $\phi(x,\cdot)$ . This proves the first claim.

Given the first claim is proved, the second claim follows directly from Lemma B.6.

The following theorem shows that the projection of $x$ in the tangent space of $\Phi(x)$ is small when $x$ is close to the manifold. In particular if we can show that in a discrete trajectory with a vanishing learning rate $\eta$ , the iterates $\{x_{\eta}(t)\}$ stay in $Y^{{\bm{\epsilon}}}$ , we can interchangeably use $\left\|x_{\eta}(t)-\Phi(x_{\eta}(t))\right\|$ with $\left\|P_{t,\Gamma}(x_{\eta}(t)-\Phi(x_{\eta}(t)))\right\|$ , with an additional error of $\mathcal{O}(\eta^{3})$ , when $\left\|P_{t,\Gamma}(x_{\eta}(t)-\Phi(x_{\eta}(t)))\right\|\leq\mathcal{O}(\eta).$

For all $x\in Y^{{\bm{\epsilon}}}$ , we have that

First of all, we can track the decrease in loss along the Gradient flow trajectory starting from $x$ . At any time $\tau$ , we have

where $\phi(x,0)=x$ . Without loss of generality, we assume $L(y)=0,\forall y\in\Gamma$ . Using the fact that $L$ is $\bm{\mu}$ -PL on $Y^{{\bm{\rho}}}$ and the GF trajectory starting from any point in $Y^{{\bm{\epsilon}}}$ stays inside $Y^{{\bm{\rho}}}$ (from Lemma B.7), we have

Moreover, we can relate $L(\phi(x,0)$ with $\left\|\Phi(x)-x\right\|$ with a second order taylor expansion:

where in the final step, we have used the fact that $L(\Phi(x))=0$ and $\nabla L(\Phi(x))=0$ . By Lemma B.7, we have $\overline{x\Phi(x)}\subset Y^{\bm{\rho}}$ . Thus $\max_{s\in}\left\|\nabla^{2}L(sx+(1-s)\Phi(x))\right\|\leq\bm{\zeta}$ from Definition B.4 and it follows that

Finally we focus on the movement in the tangent space. It holds that

By Lemma B.7, we have $\overline{\phi(x,\tau)\Phi(x)}\subset Y^{\bm{\rho}}$ for all $\tau\geq 0$ and thus

Since $P^{\perp}_{\Phi(x),\Gamma}$ is the projection matrix for the tangent space, $P^{\perp}_{\Phi(x),\Gamma}\nabla^{2}L(\Phi(x))=0$ and thus by Equation 16

Plug Equation 19 into Equation 18, we conclude that

The left-side inequality of the second inequality is proved by plugging the first claim into the above inequality Equation 20 and rearranging the terms. Note by the second claim in Lemma B.7, $\frac{\bm{\nu}\bm{\zeta}}{4\bm{\mu}^{2}}\left\|x-\Phi(x)\right\|\leq\frac{1}{2}$ , the right-side inequality is also proved. ∎

At any point $x\in Y^{{\bm{\epsilon}}}$ , we have

Moreover, the normalized gradient of $L$ can be written as

Using taylor expansion at $x$ , we have using $\nabla L(\Phi(x))=0$ :

where we use Lemma B.8 since $x\in Y^{\bm{\epsilon}}$ . Thus, the normalized gradient at any step $t$ can be written as

Consider any point $x\in Y^{{\bm{\epsilon}}}$ . Then,

where $\theta=\arctan\frac{\left\|P_{\Phi(x),\Gamma}^{(2:M)}\widetilde{x}\right\|}{\left|\langle v_{1}(x),\widetilde{x}\rangle\right|}$ , with $\widetilde{x}=\nabla^{2}L(\Phi(x))(x-\Phi(x))$ .

For any $\overline{xy}\in Y^{\bm{\epsilon}}$ where $y=x-\eta\frac{\nabla L(x)}{\left\|\nabla L(x)\right\|}$ is the one step Normalized GD update from $x$ , we have

Moreover, we must have for every $1\leq k\leq M$ ,

By Lemma B.16, we have $\partial\Phi(x)\nabla L(x)=0$ for all $x\in U$ . Thus we have

where the final step follows from using Definition B.4.

For the second claim, we have for every $1\leq k\leq M$ ,

where the first step involves Theorem F.2.

The third claim follows from using Theorem F.4. Again,

where we borrow the bound on $\left\|\nabla^{2}L(\Phi(x))-\nabla^{2}L(\Phi(y))\right\|$ from our previous calculations. The final step follows from the constants defined in Definition B.4. ∎

For any $\overline{xy}\in Y^{\bm{\epsilon}}$ where $y=x-\eta\frac{\nabla L(x)}{\left\|\nabla L(x)\right\|}$ is the one step Normalized GD update from $x$ , we have that

Here $\theta=\arctan\frac{\left\|P_{\Phi(x),\Gamma}^{(2:M)}\widetilde{x}\right\|}{\left|\langle v_{1}(x),\widetilde{x}\rangle\right|},$ with $\widetilde{x}=\nabla^{2}L(\Phi(x))(x-\Phi(x))$ . Additionally, we have that

By Taylor expansion for $\Phi$ at $x$ , we have

where in the pre-final step, we used the property of $\Phi$ from Lemma B.16. In the final step, we have used a second order taylor expansion to bound the difference between $\partial^{2}\Phi(x)$ and $\partial^{2}\Phi(\Phi(x)).$ Additionally, we have used $y-x=\eta\frac{\nabla L(x)}{\left\|\nabla L(x)\right\|}$ from the Normalized GD update rule.

Applying Taylor expansion on $\Phi$ again but at $\Phi(x)$ , we have that

Also, at $\Phi(x)$ , since $v_{1}(x)$ is the top eigenvector of the hessian $\nabla^{2}L$ , we have that from Corollary B.23,

Plug Equations 24 and 23 into Equation 22, we have that

For the second claim, continuing from Equation 22, we have that

where $\Sigma=P_{\Phi(x),\Gamma}\frac{\nabla L(x)}{\left\|\nabla L(x)\right\|}\left(P_{\Phi(x),\Gamma}\frac{\nabla L(x)}{\left\|\nabla L(\Phi(x))\right\|}\right)^{\top}$ and the last step is by Lemma B.9. Here $P_{\Phi(x),\Gamma}$ denotes the projection matrix of the subspace spanned by $v_{1}(x),\ldots,v_{M}(x)$ .

By Lemmas B.22, B.18 and B.19, we have that

Let $L_{min}=\min_{y\in U}L(y)$ . For any $\overline{xy}\in Y^{\bm{\epsilon}}$ where $y=x-\eta\frac{\nabla L(x)}{\left\|\nabla L(x)\right\|}$ is the one step Normalized GD update from $x$ , if $\left\|\nabla L(x_{\eta}(t))\right\|\geq\bm{\zeta}\eta$ , we have that

Thus for $\left\|\nabla L(x_{\eta}(t))\right\|\geq\bm{\zeta}\eta$ , we have that

where the last step is because $L$ is $\mu$ -PL on $Y^{\bm{\epsilon}}$ . In other words, we have that

where in the last step we use $L(y)-L(x)\leq 0$ . This completes the proof. ∎

B.2 Properties of limiting map of gradient flow, ΦΦ\Phi

where $\bm{\zeta}_{K}$ denotes $\sup_{x\in K}\left\|\nabla^{2}L(\phi(x,t))\right\|$ . This implies that $\left\|\nabla L(\phi(x,t))\right\|\leq e^{\bm{\zeta}_{K}T}\left\|\nabla L(x)\right\|$ and $\left\|\phi(x,t)-x\right\|\leq e^{\bm{\zeta}_{K}T}\left\|\nabla L(x)\right\|$ for all $t\in[0,T]$ .

The following results Lemmas B.16, B.17, B.18, B.19, B.20, B.21 and B.22 are from [Li et al., 2022b].

For any $x\in U$ , it holds that (1). $\partial\Phi(x)\nabla L(x)=0$ and (2). $\partial^{2}\Phi(x)[\nabla L(x),\nabla L(x)]=-\partial\Phi(x)\nabla^{2}L(x)\nabla L(x)$ .

We will also use the following two corollaries of Lemma B.22.

For any $x\in\Gamma$ , let $v_{1}$ be a top eigenvector of $\nabla^{2}L(x)$ , then

Simply note that $\mathcal{L}^{-1}_{\nabla^{2}L(x)}(v_{1}v_{1}^{\top})=\frac{1}{2\lambda_{1}(\nabla^{2}L(x))}v_{1}v_{1}^{\top}$ and apply Lemma B.22. ∎

For any $x\in\Gamma$ , let $v_{1}$ be the unit top eigenvector of $\nabla^{2}L(x)$ , then

The proof follows from using Corollary B.23 and the derivative of $\lambda_{1}$ from Theorem F.1. ∎

As a variant of Corollary B.24, we have the following lemma.

For any $x\in\Gamma$ , let $v_{1}$ be the unit top eigenvector of $\nabla^{2}L(x)$ , then

Appendix C Analysis of Normalized GD on General Loss Functions

We restate the theorem concerning Phase I for the Normalized GD algorithm. Recall the following notation for each $1\leq j\leq M$ :

See 4.3 The intuition behind the above theorem is that for sufficiently small LR $\eta$ , $x_{\eta}(t)$ will track the normalized gradient flow starting from $x_{\text{init}}$ , which is a time-rescaled version of the standard gradient flow. Thus the normalized GF will enter $Y^{\bm{\epsilon}}$ and so does normalized GD. Since $L$ satisfies PL condition in $Y^{\bm{\epsilon}}$ , the loss converges quickly and the iterate $x_{\eta}(t)$ gets $\mathcal{\eta}$ to manifold. To finish, we need the following theorem, which is the approximately-quadratic version of Lemma 3.3 when the iterate is ${O}(\eta)$ close to the manifold.

Suppose $\{x_{\eta}(t)\}_{t\geq 0}$ are iterates of Normalized GD (4) with a learning rate $\eta$ and $x_{\eta}(0)=x_{\text{init}}$ . There is a constant $C>0$ , such that for any constant $\bm{\varsigma}>1$ , if at some time $t^{\prime}$ , $x_{\eta}(t^{\prime})\in Y^{\bm{\epsilon}}$ and satisfies $\frac{\left\|x_{\eta}(t^{\prime})-\Phi(x_{\eta}(t^{\prime}))\right\|}{\eta}\leq\bm{\varsigma}$ , then for all $\bar{t}\geq t^{\prime}+C\frac{\bm{\zeta}\bm{\varsigma}}{\bm{\mu}}\log\frac{\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}$ , the following must hold true for all $1\leq j\leq M$ :

provided that for all steps $t\in\{t^{\prime},\ldots,\bar{t}-1\}$ , $\overline{x_{\eta}(t)x_{\eta}(t+1)}\subset Y^{\bm{\epsilon}}$ .

The proof of the above theorem is in Section D.1.

Let $T_{x}$ be the length of the GF trajectory starting from $x$ , and we know $\lim_{\tau\to T_{x}}\overline{\phi}(x,\tau)=\Phi(x)$ , where $\overline{\phi}(x,\tau)$ is defined as the Normalized gradient flow starting from $x$ . In Lemmas B.5 and B.2 we show there is a small neighbourhood around $\Phi(x_{\text{init}})$ , $Y^{\bm{\epsilon}}$ such that $L$ is $\bm{\mu}$ -PL in $Y^{\bm{\epsilon}}$ . Thus we can take some time $T_{0}<T_{x_{\text{init}}}$ such that $\overline{\phi}(x_{\text{init}},T_{0})\in Y^{{\bm{\epsilon}}/2}$ and $L(\overline{\phi}(x_{\text{init}}),T_{0})\leq\frac{1}{2}{L_{\text{critical}}}$ , where ${L_{\text{critical}}}:=\frac{{\bm{\epsilon}}^{2}\bm{\mu}}{8}$ . (Without loss of generality, we assume $\min_{y\in Y}L(y)=0$ ) By standard ODE approximation theory, we know there is some small $\eta_{0}$ , such that for all $\eta\leq\eta_{0}$ , $\left\|x_{\eta}(\lceil T_{0}/\eta\rceil)-\overline{\phi}(x_{\text{init}},T_{0})\right\|={O}(\eta)$ , where ${O}(\cdot)$ hides constants depending on the initialization $x_{\text{init}}$ and the loss function $L$ .

Without loss of generality, we can assume $\eta_{0}$ is small enough such that $x_{\eta}(\lceil T_{0}/\eta\rceil)\in Y^{\bm{\epsilon}}$ and $L(x_{\eta}(\lceil T_{0}/\eta\rceil))\leq{L_{\text{critical}}}$ . Now let $t_{\eta}$ be the smallest integer (yet still larger than $\lceil T_{0}/\eta\rceil$ ) such that $\overline{x_{\eta}(t_{\eta})x_{\eta}(t_{\eta}-1)}\not\subset Y^{\bm{\epsilon}}$ and we claim that there is $t\in\{\lceil T_{0}/\eta\rceil,\ldots,t_{\eta}\}$ , $\left\|\nabla L(x_{\eta}(t))\right\|<\bm{\zeta}\eta$ . By the definition of $t_{\eta}$ , we know for any $t\in\{\lceil T_{0}/\eta\rceil+1,\ldots,t_{\eta}-1\}$ , by Lemma B.11 we have $\left\|\Phi(x_{\eta}(t))-\Phi(x_{\eta}(t-1))\right\|\leq\bm{\xi}\eta^{2},$ and by Lemma B.13, $\sqrt{L(x_{\eta}(t))}-\sqrt{x_{\eta}(t-1)}\leq-\eta\frac{\sqrt{2\bm{\mu}}}{4}$ if $\left\|\nabla L(x_{\eta}(t))\right\|\geq\bm{\zeta}\eta$ . If the claim is not true, since $\sqrt{L(x_{\eta}(t))}$ decreases $\eta\frac{\sqrt{2\bm{\mu}}}{4}$ per step, we have

which implies that $t_{\eta}-\lceil T_{0}/\eta\rceil-1\leq\frac{{\bm{\epsilon}}}{\eta}$ , and therefore by Lemma B.11,

Meanwhile, by Lemma B.6, we have $\left\|\Phi(x_{\eta}(t_{\eta}-1))-x_{\eta}(t_{\eta}-1)\right\|\leq\sqrt{\frac{2L(x_{\eta}(t_{\eta}-1))}{\bm{\mu}}}\leq\sqrt{\frac{2L(x_{\eta}(\lceil T_{0}/\eta\rceil))}{\bm{\mu}}}=\frac{{\bm{\epsilon}}}{2}$ . Thus for any $\kappa\in$ , we have $\left\|\kappa x_{\eta}(t_{\eta})+(1-\kappa)x_{\eta}(t_{\eta}-1)-\Phi(x_{\text{init}})\right\|$ is upper bounded by

which is smaller than ${\bm{\epsilon}}$ since we can set $\eta_{0}$ sufficiently small. In other words, $\overline{\Phi(x_{\eta}(t_{\eta}))\Phi(x_{\eta}(t_{\eta}-1))}\subset Y^{\bm{\epsilon}}$ , which contradicts with the definition of $t_{\eta}$ . So far we have proved our claim that there is some $t^{\prime}_{\eta}\in\{\lceil T_{0}/\eta\rceil,\ldots,t_{\eta}\}$ , $\left\|\nabla L(x_{\eta}(t^{\prime}_{\eta}))\right\|<\bm{\zeta}\eta$ . Moreover, since $\sqrt{L(x_{\eta}(t))}$ decreases $\eta\frac{\sqrt{2\bm{\mu}}}{4}$ per step before $t^{\prime}_{\eta}$ , we know $t^{\prime}_{\eta}-\lceil T_{0}/\eta\rceil\leq\frac{{\bm{\epsilon}}}{\eta}$ . By Lemma B.6, we know $\left\|x_{\eta}(t^{\prime}_{\eta})-\Phi(x_{\eta}(t^{\prime}_{\eta}))\right\|\leq\frac{\bm{\zeta}\eta}{\bm{\mu}}$ .

Now we claim that for any $T_{1}^{\prime}$ , there is some sufficiently small threshold $\eta_{0}$ , $t_{\eta}\geq\frac{T_{1}^{\prime}}{\eta}+1$ if $\eta\leq\eta_{0}$ . Below we prove this claim by contradiction. If the claim is not true, that is, $t_{\eta}<\frac{T_{1}^{\prime}}{\eta}+1$ . if $t_{\eta}\leq C\frac{\bm{\zeta}\bm{\varsigma}}{\bm{\mu}}\log\frac{\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}+t^{\prime}_{\eta}$ with $\bm{\varsigma}=\frac{\bm{\zeta}}{\bm{\mu}}$ , we know $\left\|x_{\eta}(t_{\eta})-\Phi(x_{\text{init}})\right\|\leq\left\|x_{\eta}(t_{\eta})-x_{\eta}(t^{\prime}_{\eta})\right\|+\left\|x_{\eta}(t^{\prime}_{\eta})-\Phi(x_{\eta}(t^{\prime}_{\eta}))\right\|+\left\|\Phi(x_{\eta}(t^{\prime}_{\eta}))-\Phi(x_{\text{init}})\right\|={O}(\eta)$ , which implies that $\overline{x_{\eta}(t_{\eta})x_{\eta}(t_{\eta}-1)}\in Y$ . If $t_{\eta}\geq C\frac{\bm{\zeta}\bm{\varsigma}}{\bm{\mu}}\log\frac{\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}+t^{\prime}_{\eta}$ , by Lemma C.1, we have $\left\|x_{\eta}(t_{\eta})-\Phi(x_{\eta}(t_{\eta}))\right\|={O}(\eta)$ . By Lemma B.11, we have $\left\|\Phi(x_{\eta}(t_{\eta}))-\Phi(x_{\eta}(\lceil T_{0}/\eta\rceil))\right\|\leq{O}(\eta)$ . Thus again we have that $\left\|x_{\eta}(t_{\eta})-\Phi(x_{\text{init}})\right\|\leq\left\|x_{\eta}(t_{\eta})-\Phi(x_{\eta}(t_{\eta}))\right\|+\left\|\Phi(x_{\eta}(t_{\eta}))-\Phi(x_{\eta}(\lceil T_{0}/\eta\rceil))\right\|\allowbreak+\left\|\Phi(x_{\eta}(\lceil T_{0}/\eta\rceil))-\Phi(x_{\text{init}})\right\|={O}(\eta)$ , which implies that $\overline{x_{\eta}(t_{\eta})x_{\eta}(t_{\eta}-1)}\in Y$ . In both cases, the implication is in contradiction to the definition of $t_{\eta}$ .

Thus for any $T_{1}^{\prime}$ , $t_{\eta}\geq\frac{T_{1}^{\prime}}{\eta}+1$ for sufficiently small threshold $\eta_{0}$ and $\eta\leq\eta_{0}$ . To complete the proof of Theorem 4.3, we pick $T_{1}$ to be any real number strictly larger than ${\bm{\epsilon}}+T_{0}$ , as $\frac{T_{1}}{\eta}>C\frac{\bm{\zeta}\bm{\varsigma}}{\bm{\mu}}\log\frac{\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}+\frac{{\bm{\epsilon}}}{\eta}+\lceil T_{0}/\eta\rceil\geq C\frac{\bm{\zeta}\bm{\varsigma}}{\bm{\mu}}\log\frac{\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}+t^{\prime}_{\eta}$ when $\eta$ is sufficiently small with $\bm{\varsigma}=\frac{\bm{\zeta}}{\bm{\mu}}$ . By Lemma C.1 the second claim of Theorem 4.3 is proved. Using the same argument again, we know $\forall\frac{T_{1}}{\eta}\leq t\leq\frac{T_{1}^{\prime}}{\eta}$ , it holds that $\left\|\Phi(x_{\eta}(t))-\Phi(x_{\text{init}})\right\|\leq{O}(\eta)$ . ∎

C.2 Phase II, Limiting Flow

We first restate the main theorem that demonstrates that the trajectory implicitly minimizes sharpness. See 4.4

where $v^{\prime}(\tau^{\prime}+0):=\lim_{\delta\to 0}\frac{v(\tau^{\prime}+\delta)-v(\tau^{\prime})}{\delta}$ is the right time derivative of $v$ at $\tau^{\prime}$ .

To prove the first claim, we first show the movement in the manifold for the discrete trajectory for Algorithm 1 by Lemma B.12: for each step $t$ , provided $\overline{\Phi(x_{\eta}(t))\Phi(x_{\eta}(t+1))}\in Y^{\bm{\epsilon}}$ , it holds that

The high-level idea for the proof of the first claim is to bound the gap between Equation 26 and Equation 5 using Theorem C.2. And the first claim eventually boils down to upper bound the average angle by $O(\eta)$ , which is exactly the second claim.

Formally, let $t_{2}$ be the largest integer no larger than $\lfloor T_{2}/\eta^{2}\rfloor$ such that for any $0\leq t\leq t_{2}$ , it holds that $\overline{\Phi(x_{\eta}(t))\Phi(x_{\eta}(t+1))}\in Y^{\bm{\epsilon}}$ .

Since we started from a point that has $\max_{1\leq j\leq M}R_{j}(x_{\eta}(0))\leq O(\eta^{2})$ , we have from Lemma C.1, that the iterate satisfies the condition $\max_{1\leq j\leq M}R_{j}(x_{\eta}(t))\leq O(\eta^{2})$ at step $t$ as well, meaning that $\left\|x_{\eta}(t)-\Phi(x_{\eta}(t))\right\|\leq O(\eta)$ .

Therefore, for any $\tau\leq t_{2}\eta^{2}$ , note that $v^{\prime}(\tau+0)=v^{\prime}(\lfloor\tau/\eta^{2}\rfloor+0)$ and that $\left\|f(v(\lfloor\tau/\eta^{2}\rfloor+0))-f(v(\tau))\right\|=O(\left\|\Phi(x_{\eta}(\lfloor\tau/\eta^{2}\rfloor+1))-\Phi(x_{\eta}(\lfloor\tau/\eta^{2}\rfloor))\right\|)=O(\eta^{2})$ , we have that

where in the last step we use the second claim. This implies that $t_{2}$ must be equal to $\lfloor T_{2}/\eta^{2}\rfloor$ for sufficiently small $\eta$ otherwise $\overline{x_{\eta}(t_{2})x_{\eta}(t_{2}+1)}\subseteq Y^{\bm{\epsilon}}$ . This is because $\left\|x_{\eta}(t_{2}+1)-x_{\eta}(t_{2})\right\|=O(\eta)$ and $X(t_{2}\eta^{2})\in Y$ . The proof is completed by noting that $\left\|X(T_{2})-X(\lfloor T_{2}/\eta^{2}\rfloor)\right\|=O(\eta^{2})$ . ∎

Appendix D Phase I, Omitted Proofs of the Main Lemmas

The Normalized GD update at any step $t$ can be written as (from Lemma B.9)

From Lemma B.11, we have $\left\|\Phi(x_{\eta}(t))-\Phi(x_{\eta}(t+1))\right\|\leq{O}(\bm{\xi}\eta^{2})$ , which further implies, $\left\|\nabla^{2}L(\Phi(x_{\eta}(t+1)))-\nabla^{2}L(\Phi(x_{\eta}(t)))\right\|\leq{O}(\bm{\nu}\bm{\xi}\eta^{2})$ . Thus, using the notation $\widetilde{x}=\nabla^{2}L(\Phi(x))(x-\Phi(x))$ , we have

Below we will show that $\left\|x_{\eta}(t)-\Phi(x_{\eta}(t))\right\|\leq{O}(\eta)$ , and thus the trajectory of $\widetilde{x}_{\eta}$ is similar to the trajectory in the qudratic model with an ${O}(\eta^{2})$ error, with the hessian fixed at $\nabla^{2}L(\Phi(x_{\eta}(t)))$ , and hence we can apply the same techniques from Corollary A.4 and Lemma A.1.

First, we consider the norm of the vector $\widetilde{x}_{\eta}(t)$ for $t^{\prime}+1\leq t\leq\overline{t}$ . We will show the following induction hypothesis:

Base case: $(t=t^{\prime})$ . We have $\left\|\widetilde{x}_{\eta}(t^{\prime})\right\|=\left\|\nabla^{2}L(\Phi(x_{\eta}(t^{\prime})))[x_{\eta}(t^{\prime})-\Phi(x_{\eta}(t^{\prime}))]\right\|\leq\eta\lambda_{1}(t)\bm{\varsigma}\leq\eta\bm{\zeta}\bm{\varsigma}$ .

Induction case: $(t>t^{\prime})$ . Suppose the hypothesis holds true for $t-1$ . Then,

If $\left\|\widetilde{x}_{\eta}(t-1)\right\|\geq\eta\lambda_{1}(t)$ . We can directly apply Corollary A.3 on (28) to show that

where the final step follows if $\eta$ is sufficiently small. Hence, $\left\|\widetilde{x}_{\eta}(t)\right\|<\left\|\widetilde{x}_{\eta}(t-1)\right\|\leq\eta\bm{\zeta}\bm{\varsigma}$ .

If $\left\|\widetilde{x}_{\eta}(t-1)\right\|\leq\eta\lambda_{1}(t)$ . Then, we can directly apply Lemma A.1 on (28) to show that

Hence, we have shown that, $\left\|x_{\eta}(t)-\Phi(x_{\eta}(t))\right\|\leq\frac{1}{\lambda_{M}(t)}\left\|\widetilde{x}_{\eta}(t)\right\|\leq\frac{1.01\eta\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}$ for all time $t^{\prime}\leq t\leq\overline{t}$ .

We complete the proof of Lemma C.1 with a similar argument as that for the quadratic model (see Corollary A.4 and Lemma A.1). The major difference from the quadratic model is that here the hessian changes over time, along with its eigenvectors and eigenvalues. Hence, we need to take care of the errors introduced in each step by the change of hessian.

The high-level idea is to divide the eigenvalues at each step $t$ into groups such that eigenvalues in the same group are $O(\eta)$ close and eigenvalues from different groups are at least $2\eta$ far away from each other. Formally, we divide $[M]$ into disjoint subsets $S^{(t)}_{1},\cdots,S^{(t)}_{p(t)}$ (with $1\leq p(t)\leq M$ ) such that

Thus for any $t^{\prime}\leq t\leq\overline{t}-1$ and $k\in[p(t)]$ , suppose $i\in S^{(t)}_{k}$ and $j=\min S^{(t)}_{k}$ , we have that

and $\eta\lambda_{i}(t+1)\geq\eta\lambda_{i}(t)-O(\eta^{2})\geq\eta\lambda_{j}(t)-O(\eta^{2}).$

Next we will use the results from the quadratic case to upper bound $\sqrt{\sum_{h=j}^{M}\langle v_{h}(t),\widetilde{x}_{\eta}(t+1)\rangle^{2}}$ using $\sqrt{\sum_{h=j}^{M}\langle v_{h}(t),\widetilde{x}_{\eta}(t)\rangle^{2}}$ . For all $1\leq j\leq M$ , we consider the following two cases for any time $t^{\prime}+1\leq t\leq\overline{t}$ :

If $\sqrt{\sum_{h=j}^{M}\langle v_{h}(t),\widetilde{x}_{\eta}(t)\rangle^{2}}>\eta\lambda_{j}(t)$ , then we can apply Lemma A.2 on (28) to show that

If $\sqrt{\sum_{h=j}^{M}\langle v_{h}(t),\widetilde{x}_{\eta}(t)\rangle^{2}}\leq\eta\lambda_{j}(t)$ , then we can apply Lemma A.1 on (28) to show that

and therefore following the same proof of quadratic case Corollary A.4, for $\overline{t}\geq t^{\prime}+{\Omega}(\frac{\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}\log\frac{\bm{\zeta}\bm{\varsigma}}{\bm{\mu}})$ , it holds that $\forall j\in[M]$ , $\sqrt{\sum_{i=j}^{M}\langle v_{i}(\bar{t}),\widetilde{x}(\bar{t})\rangle^{2}}\leq\eta\lambda_{j}(\bar{t})+{O}(\eta^{2}).$ ∎

D.2 Properties of the condition in Equation 25

By Lemma C.1, the following condition will continue to hold true for all $1\leq j\leq M$ before $\widetilde{x}_{\eta}(t)$ leaves $Y^{\bm{\epsilon}}$ :

where $\widetilde{x}_{\eta}(t)=\nabla^{2}L(\Phi(x_{\eta}(t)))(x_{\eta}(t)-\Phi(x_{\eta}(t))).$ We will call the above condition as the alignment condition from now onwards.

From the alignment condition (29), we can derive the following property that continues to hold true throughout the trajectory, once the condition is satisfied:

The proof follows from using the noisy quadratic update for Normalized GD in Lemma B.9 (Equation 21) and the behavior in a quadratic model along the non-top eigenvectors in Lemma A.5.

Appendix E Phase II, Omitted Proofs of the Main Lemmas

The main lemma in this section is Lemma E.5 in Section E.2, which says the sum of the angles across the entire trajectory in any interval $[0,t_{2}]$ with $t_{2}=\Omega(1/\eta^{2})$ , is at most $O(\eta t_{2})$ . Before proving the main lemma, we will first recap and introduce some notations that will be used.

In Phase II, we start from a point $x_{\eta}(0)$ , such that (1) $\left\|x_{\eta}(0)-\Phi(x_{\text{init}})\right\|\leq O(\eta)$ , (2) $\max_{j\in[D]}R_{j}(x_{\eta}(t))\leq O(\eta^{2})$ , and additionally (3) $\left|\langle v_{1}(x_{\eta}(0)),x_{\eta}(0)-\Phi(x_{\eta}(0))\rangle\right|=\Omega(\eta)$ .

The condition (29) that was shown to hold true in Phase II is:

Further, we had proved in Lemma D.1 that:

The properties of the three sets $N_{0},N_{1},N_{2}$ in an interval $(\widetilde{t},\overline{t})$ given directly by the algorithm in Algorithm 3 include:

$\forall t,t\in N_{1}\Longleftrightarrow t+1\in N_{2}$ .

$N_{0}\cup N_{1}\cup N_{2}=[\widetilde{t},\overline{t}]$ , and the intersection between each pair of them is empty.

We also have the following lemmas, which is less direct:

For any step $t$ in $N_{0}$ , $t-1,t+2\in N_{2}$ and $t-2,t+1\in N_{1}$

$\tan\theta_{t+1}\leq\left(1-\frac{\min(\bm{\Delta},2\bm{\mu})}{\bm{\zeta}}\right)\tan\theta_{t}+O(\frac{\eta^{2}}{G_{t}})$

$\tan\theta_{t+2}\leq\frac{\eta\lambda_{1}}{G_{t}}\tan\theta_{t}+O(\frac{\eta^{2}}{G_{t}})$

As a direct consequence of Lemma E.3, we have the following lemma:

Given any $t$ with $\theta_{t}=\Omega(1)$ , let $\widetilde{t}=\max N_{1}\cap\{\widetilde{t}\mid\widetilde{t}\leq t\}$ . If $G_{\widetilde{t}}\geq\Omega(\eta)$ , then $\theta_{\widetilde{t}}=\Omega(1)$ .

The claim is clearly true if $t\in N_{1}$ . If $t\in N_{0}$ , then Lemma E.1 shows that $t-1\in N_{2},t-2\in N_{1}$ and thus $\widetilde{t}=t-2$ . The claim is true because of the second property of Lemma E.3. If $t\in N_{2}$ , then $\widetilde{t}=t-1\in N_{1}$ and the proof is completed by applying the first property of Lemma E.3. ∎

E.2 Time Average of Angles Against Top Eigenspace

provided $\eta$ is set sufficiently small, and for all time $0\leq t\leq t_{2}-1$ , $\overline{x_{\eta}(t)x_{\eta}(t+1)}\subset Y^{\bm{\epsilon}}$ .

We analyze the behavior of a general $\widetilde{t}_{i}$ when it falls in any of the above cases:

Case (B).

From Lemma E.6 we have that the sum of angle over this time is

Case (C).

Now it remains to upper bound the number of occurrence of (A),(B) and (C). Since our goal is to show average angle is $O(\eta)$ , which is equal to the average angle in case (A), so the number of occurrence of case (A) doesn’t matter. For case (B), if it is followed by case (A), then there is an $\Omega(1/\eta^{2})$ gap before next occurrence of (B). If (B) is followed by case (C), then by Lemma E.10, it takes at least $\Omega(1/\eta^{2})$ steps to escape from (B). Thus we can have $O(1)$ occurrence of case (B). For the same reason, there could be at most $O(1)$ occurrence of case (C).

All in all, with probability at least $1-O(\eta^{12}\cdot\eta^{12})=1-O(\eta^{10})$ , we must have

where we use $t_{2}\geq\Omega(\frac{1}{\eta^{2}})$ in the last step and and the number of occurrence of case (B) is $O(1)$ in the second to the last step. ∎

The noisy update rule for Normalized GD, as derived in Lemma B.9, which says that the Normalized GD update is very close to the update in a quadratic model with an additional $O(\eta^{2})$ error. Keeping this in mind, we then divide our trajectory in the interval $(\overline{t},t^{\prime})$ as per Algorithm 3 into three subsets $N_{0},N_{1},N_{2}$ . (Please see Section E.1 for a summary on the properties of these 3 sets.)

Consider any $t\in N_{1}$ . Using the behavior of $G_{t}$ from Lemma E.10, we can show that in each of the time-frames, $G_{t+2}\geq(1+\Omega(\sin^{2}\theta_{t}))G_{t}-O(\eta^{2}(\eta+\eta_{t}))\geq G_{t}+\Omega(\theta_{t}^{2}\eta)-O(\eta^{2}(\eta+\theta_{t}))$ .

Next we want to telescope over $G_{t+2}-G_{t}$ to get an upper bound for $\sum_{t\in N_{1}}\theta_{t}$ . If $t+2$ is also in $N_{1}$ then it’s fine. If $t+2\in N_{0}$ , then $t+3\in N_{1}$ by Lemma E.1 and we proceed in the following two cases.

If $\theta_{t+2}\leq C$ for some sufficiently small constant $C$ , since $G_{t+2}\leq\lambda_{1}(t+2)\eta/2-\Omega(\eta)$ , we have $\left\|\widetilde{x}_{\eta}(t+2)\right\|\leq\frac{G_{t+2}}{\cos\theta_{t+2}}=\lambda_{1}(t+2)\eta/2-\Omega(\eta)$ , and thus by Lemma E.9, we have $G_{t+3}\leq G_{t+2}$ and therefore, $G_{t+3}\geq G_{t}+\Omega(\theta_{t}^{2}\eta)-O(\eta^{2}(\eta+\theta_{t}))$ .

If $\theta_{t+2}\geq C$ , then by Lemma D.2, we have $\theta_{t}=\Omega(1)$ , thus $G_{t+2}\geq G_{t}+\Omega(\eta)$ by Lemma E.10. Again by Lemma E.9, we have $G_{t+3}\geq G_{t+2}-O(\eta^{2})$ . Thus again we conclude $G_{t+3}\geq G_{t}+\Omega(\eta)\geq G_{t}+\Omega(\theta_{t}^{2}\eta)-O(\eta^{2}(\eta+\theta_{t}))$ , since $\theta_{t}$ is always $O(1)$ .

Since total increase in $G_{t}$ during this interval can is most $O(\eta)$ , we conclude that $\sum_{t\in N_{1}}\theta_{t}^{2}=O(1)+\eta\sum_{t\in N_{1}}(\eta+\theta_{t})$ and thus it holds that

Moreover, by Lemma E.3, we must have $\theta_{t}<\theta_{t-1}+O(\eta)$ for any time $t\in N_{2}$ , and $t-1$ must be in $N_{1}$ . By Lemma E.9, we have $\theta_{t}\leq\Omega(\theta_{t-2})$ for any $t\in N_{0}$ and $t-2$ must be in $N_{2}$ . That implies,

Consider any coordinate $2\leq k\leq M$ . For any constants $0<\bm{\beta}$ , there is some constant $\bm{\alpha}>0$ such that for any time step $t$ where $x_{\eta}(\widetilde{t})$ is in $Y^{\bm{\epsilon}}$ , $G_{\widetilde{t}}\geq\bm{\beta}\eta$ , condition (29) holds and $\left|\left\langle v_{k}(\widetilde{t}),\widetilde{x}(\widetilde{t})\right\rangle\right|\geq\bm{\alpha}\eta^{2}$ , then there is some time $\overline{t}\leq t+O\left(\ln 1/\eta\right)$ such that if for all time $\widetilde{t}\leq t^{\prime}<\overline{t}$ , $\overline{x_{\eta}(t^{\prime})x_{\eta}(t^{\prime}+1)}\subset Y^{\bm{\epsilon}}$ , then condition (29) holds at time $\overline{t}$ and at least one of the following two conditions hold:

$G_{\overline{t}}\geq 0.99g_{t}(\lambda_{k}(\overline{t}))$ .

We will prove by contradiction. Suppose neither of the two condition happens, we will show $\theta_{t}$ grows exponentially and thus the condition (2) must be false in $O(\log 1/\eta)$ steps.

Now, we can use Equation 21 (Lemma B.9) to show that the Normalized GD update is equivalent to update in quadratic model, up to an additional $O(\eta^{2})$ error.

Similar to Lemma A.9, consider the coordinate $k$ , we have that

The third step follows from using the same argument as the one used for the quadratic update in Lemma A.9 and the assumption that $G_{t}\geq 0.99g_{t}(\lambda_{k}(t))$ . The final step holds true because we can pick $\bm{\alpha}$ as a large enough constant and by assumption $\frac{\left|\langle v_{k}(t),\widetilde{x}_{\eta}(t)\rangle\right|}{\left|\langle v_{1}(t),\widetilde{x}_{\eta}(t)\rangle\right|}\geq\bm{\alpha}\eta$ .

We then bound $\left\|v_{k}(\overline{t})-v_{k}(\widetilde{t})\right\|$ and $\left\|\Phi(x_{\eta}(\overline{t})-\Phi(x_{\eta}(\widetilde{t}))\right\|$ by $O(\eta^{2}(\overline{t}-\widetilde{t}))$ using Lemma B.11. Combining everything, we conclude that at least one of the two assumptions has to break for some $\overline{t}\leq\widetilde{t}+O(\log 1/\eta)$ . ∎

Consider any coordinate $2\leq k\leq M$ . For any constants $0<\bm{\beta}$ , suppose at time step $t$ , $x_{\eta}(t)$ is in $Y^{\bm{\epsilon}}$ , $\left(1.01\right)g_{t}(\lambda_{k}(t))\eta\leq G_{t}<0.5\eta\lambda_{1}(t)$ and condition (29) holds, then there is some time $\overline{t}\leq t+O\left(\ln 1/\eta\right)$ such that if for all time $t\leq t^{\prime}<\overline{t}$ , $\overline{x_{\eta}(t^{\prime})x_{\eta}(t^{\prime}+1)}\subset Y^{\bm{\epsilon}}$ , then the following two conditions hold:

The proof of Lemma E.8 is very similar to the proof of Lemma E.7 and thus we omit the proof. The only difference will be that we need to use Lemma A.10 in place of Lemma A.9, when we use the result for the quadratic model.

E.3 Dynamics in the Top Eigenspace

Here, we will state two important lemmas that we used for the proof of Lemma E.5, which is about the behavior of the iterate along the top eigenvector. Lemma E.10 can be viewed as perturbed version for Lemma A.7 in the quadratic case, and We assume in all the lemmas, that Equation 29 holds true for the time under consideration, which we showed in Lemma C.1, and also that we start Phase II from a point where the alignment along the top eigenvector is non negligible.

The following lemmas give the properties of dynamics in the top eigenspace in Phase II for one-step and two-step updates respectively. Recall we use $G_{t}$ to denote the quantity $\left|\left\langle v_{1}(t),\widetilde{x}(t)\right\rangle\right|$ .

provided that $\overline{x_{\eta}(t)x_{\eta}(t+1)},\overline{x_{\eta}(t+1)x_{\eta}(t+2)}\subset Y^{\bm{\epsilon}}$ .

First note that $\angle(\widetilde{x}_{\eta}(t),\nabla L(x_{\eta}(t)))=O(\frac{\left\|\widetilde{x}_{\eta}(t)-\nabla L(x_{\eta}(t))\right\|}{\left\|\widetilde{x}_{\eta}(t)\right\|})=O(\frac{\eta^{2}}{G_{t}})=O(\eta)$ , where the last step we use Lemma B.9. Let $\delta=\angle(v_{1}(t),\nabla L(x_{\eta}(t)))-\angle(v_{1}(t),\widetilde{x}_{\eta}(t))$ and we have $\left|\delta\right|\leq\angle(\widetilde{x}_{\eta}(t),\nabla L(x_{\eta}(t)))=O(\eta)$ . Therefore, it holds that

From Lemma B.11, we have $\left\|\Phi(x_{\eta}(t))-\Phi(x_{\eta}(t+1))\right\|\leq O(\bm{\xi}\eta^{2})$ , which further implies, $\left\|\nabla^{2}L(\Phi(x_{\eta}(t+1)))-\nabla^{2}L(\Phi(x_{\eta}(t)))\right\|\leq O(\eta^{2})$ . Thus, we can use Theorem F.4 to have $\left\|v_{1}(t)-v_{1}(t+1)\right\|\leq O(\frac{\bm{\nu}\bm{\xi}\eta^{2}}{\lambda_{1}(t)-\lambda_{2}(t)})=O(\eta^{2})$ . From Lemma B.12, we have $\left|\langle v_{1}(t),\Phi(x_{\eta}(t+1))-\Phi(x_{\eta}(t))\rangle\right|\leq O(\eta^{3})$ . Thus we have that

Therefore, we have the following inequality by applying the same argument above to $t+1$ :

By Lemma D.1, we know $\left\|\widetilde{x}(t+1)\right\|=\Omega(\eta)$ . By Lemma D.2, we know that $\theta_{t+1}\leq\theta_{t}+O(\eta)$ . Thus

Next we will show $\left|\left\|\widetilde{x}_{\eta}(t+1)\right\|-\left\|(I-\frac{\eta\nabla^{2}L(\Phi(x_{\eta}(t)))}{\left\|\widetilde{x}_{\eta}(t)\right\|})\widetilde{x}_{\eta}(t)\right\|\right|=O(\eta^{2}\theta_{t})$ . For convenience, we denote $\nabla^{2}L(\Phi(x_{\eta}(t)))$ by $H$ . First we have that

where $\alpha$ is the angle between $\frac{\nabla L(x_{\eta}(t))}{\left\|\nabla L(x_{\eta}(t))\right\|}-\frac{\widetilde{x}_{\eta}(t)}{\left\|\widetilde{x}_{\eta}(t)\right\|}$ and $2H\widetilde{x}_{\eta}(t)-\eta H^{2}\left(\frac{\nabla L(x_{\eta}(t))}{\left\|\nabla L(x_{\eta}(t))\right\|}+\frac{\widetilde{x}_{\eta}(t)}{\left\|\widetilde{x}_{\eta}(t)\right\|}\right)$ . Note that and that both $\angle(\widetilde{x}_{\eta}(t),v_{1}(t)),\angle(\nabla L(x_{\eta}(t)),v_{1}(t))=O(\eta_{t}+\eta)$ , we have that the angle between $\frac{\nabla L(x_{\eta}(t))}{\left\|\nabla L(x_{\eta}(t))\right\|}+\frac{\widetilde{x}_{\eta}(t)}{\left\|\widetilde{x}_{\eta}(t)\right\|}$ and $2H\widetilde{x}_{\eta}(t)-\eta H^{2}\left(\frac{\nabla L(x_{\eta}(t))}{\left\|\nabla L(x_{\eta}(t))\right\|}+\frac{\widetilde{x}_{\eta}(t)}{\left\|\widetilde{x}_{\eta}(t)\right\|}\right)$ is at most $O(\eta_{t}+\eta)$ . Further note that $\frac{\nabla L(x_{\eta}(t))}{\left\|\nabla L(x_{\eta}(t))\right\|}-\frac{\widetilde{x}_{\eta}(t)}{\left\|\widetilde{x}_{\eta}(t)\right\|}$ is perpendicular to $\frac{\nabla L(x_{\eta}(t))}{\left\|\nabla L(x_{\eta}(t))\right\|}+\frac{\widetilde{x}_{\eta}(t)}{\left\|\widetilde{x}_{\eta}(t)\right\|}$ , we know $\cos\alpha\leq O(\theta_{t}+\eta)$ . Therefore we have that

Thus we have proved a perturbed version of Lemma A.6, that is,

Therefore a perturbed version of Lemma A.7 would give us:

The proof of the first inequality is completed by plugging the above equation into Equation 32.

The second inequality is immediate by noting that $\eta\theta^{2}_{t}+C^{2}\eta^{3}\geq 2C\eta^{2}\theta_{t}$ for any $C>0$ . ∎

E.4 Dynamics in Top Eigenspace When Dropping Below Threshold

Denote $r=\eta^{100}$ . For any constant $0<\bm{\beta}$ , there is a constant $\bm{\alpha}>0$ , such that for any step $t$ and $x_{\eta}(t)\in Y^{\bm{\epsilon}}$ with the following conditions hold:

$\left|\langle v_{i}(t),\widetilde{x}_{\eta}(t)\rangle\right|\leq O(\eta^{2}),$ for all $2\leq i\leq M.$

Lemma E.11 is a direct consequence of the following lemma.

where $P_{t,\Gamma}^{(2:M)}$ denotes the subspace spanned by $v_{2}(t),\ldots,v_{M}(t)$ .

We will first prove Lemma E.11 and then we turn to the proof of Lemma E.12.

For convenience, we denote $\nabla^{2}L(\Phi(x_{\eta}(t)))[u(\widetilde{t})-\Phi(x_{\eta}(t))]$ by $\widetilde{u}(\widetilde{t})$ and $\nabla^{2}L(\Phi(x_{\eta}(t)))[w(\widetilde{t})-\Phi(x_{\eta}(t))]$ by $\widetilde{w}(\widetilde{t})$ . Suppose both $\left\|P_{t,\Gamma}^{(2:M)}\widetilde{u}(\widetilde{t})\right\|,\left\|P_{t,\Gamma}^{(2:M)}\widetilde{v}(\widetilde{t})\right\|$ are $O(\eta^{2})$ , we will show the following, which indicates the contradiction:

An important claim to note is the following:

Note that the condition has been slightly changed to use $\{v_{i}(t)\}$ as reference coordinate system and $\Phi(x_{\eta}(t))$ as reference point. The above lemma follows from the fact that both $u(0)$ and $w(0)$ are $r$ -close to $x_{\eta}(t)$ , which itself satisfies the alignment condition (Equation 29). Thus, both $u(0)$ and $w(0)$ initially follow the desired condition. Since, both the trajectories follow Normalized GD updates, the proof will follow from applying the same technique used in the proof of Lemma C.1. Another result to keep in mind is the following modified version of Corollary D.3, Lemma E.14.

The above lemma uses $\{v_{i}(t)\}$ as reference coordinate system and $\Phi(x_{\eta}(t))$ as reference point. The above lemma follows from showcasing Normalized GD updates of $u(\widetilde{t})$ and $w(\widetilde{t})$ as equivalent to the update in a quadratic model, with an additional noise of $O(\frac{\bm{\nu}\bm{\zeta}}{\bm{\mu}}\eta^{2})$ , similar to Equation 28.

Further, applying the same technique from Lemma E.10, we can show that

Initially, because $u$ was initialized close to $x_{\eta}(t)$ , we must have

Now, we use taylor expansion of $F$ around $u(\widetilde{t})$ to get

Using taylor expansion: $\nabla L(z(\bm{\gamma}))=\nabla^{2}L(\Phi(x_{\eta}(t)))(z(\bm{\gamma})-\Phi(x_{\eta}(t)))+O(\bm{\nu}\left\|z(\bm{\gamma})-\Phi(x_{\eta}(t))\right\|^{2})$ and hence, we must have $\left\|\nabla L(z(\bm{\gamma}))\right\|\geq\Omega(\eta)$ .

With $\left\|\widetilde{u}(\widetilde{t})\right\|=o(\eta)$ , we can apply Lemma E.14 to show

That implies, $\left\|\nabla L(F(z(\bm{\gamma})))\right\|\geq\Omega(\eta)$

Hence, $\mu(\widetilde{t})=\max_{\bm{\gamma}}\frac{1}{\left\|\nabla L(z(\bm{\gamma}))\right\|^{2}}+\frac{1}{\left\|\nabla L(F(z(\bm{\gamma})))\right\|^{2}}\leq\Omega(1/\eta^{2}).$

where $\partial F^{(2)}(u(\widetilde{t}))=A_{\widetilde{t}+1}A_{\widetilde{t}}$ with

Now we define $B_{\widetilde{t}}$ and claim $A_{\widetilde{t}}$ can be approximated as below with $\left\|B_{\widetilde{t}}\right\|=O(\eta)$ . Furthermore, $\left\|A_{\widetilde{t}}\right\|\leq O(1)$ .

with $\left\|A_{\widetilde{t}+1}\right\|\leq O(1)$ and $\left\|B_{\widetilde{t}+1}\right\|\leq O(\eta)$ .

Finally, we use Lemma E.15 and Lemma E.16 to handle the main and error terms in Equation 41,

which completes the proof of Lemma E.12. ∎

For simplicity of presentation, we have used $M_{\widetilde{t}}$ to define

Thus, the term under consideration can be simplified as follows,

We first recall the definition of the error term:

By Equation 37, the following property holds for all $\widetilde{t}\geq 0$ for some $\bm{\varrho}>0$

First, by induction hypothesis at $\widetilde{t}-2$ , we know

The proof is completed by picking $C$ large enough such that $\bm{\varrho}\frac{C}{\bm{\varphi}}+O(\eta)\leq C$ . ∎

E.5 Proof for Operating on Edge of Stability

According to the proof of Theorem 4.4, we know for all $t$ , it holds that $R_{j}(x_{\eta}(t))\leq O(\eta^{2})$ . Thus $S_{L}(x_{\eta}(t),\eta_{t})=\eta_{t}\cdot\sup_{0\leq s\leq\eta_{t}}\lambda_{1}(\nabla^{2}L(x_{\eta}(t)-s\nabla L(x_{\eta}(t))))=\eta_{t}(\lambda_{1}(t)+O(\eta))$ , which implies that $[S_{L}(x_{\eta}(t),\eta_{t})]^{-1}=\frac{\left\|\nabla L(x_{\eta}(t))\right\|}{\eta\lambda_{1}(t)}+O(\eta)=\frac{\left\|\widetilde{x}_{\eta}(t)\right\|}{\eta\lambda_{1}(t)}+O(\eta)$ . The proof for the first claim is completed by noting that $\frac{1}{\eta}(\left\|\widetilde{x}_{\eta}(t)\right\|+\left\|\widetilde{x}_{\eta}(t+1)\right\|)=\lambda_{1}(t)+O(\eta+\theta_{t})$ as an analog of the quadratic case.

For the second claim, it’s easy to check that $\sqrt{L(x_{\eta}(t))}=\frac{\left\|\widetilde{x}_{\eta}(t)\right\|}{\sqrt{2\lambda_{1}(t)}}+O(\eta\theta_{t})$ . Thus have $\sqrt{L(x_{\eta}(t))}+\sqrt{L(x_{\eta}(t+1))}=\frac{\left\|\widetilde{x}_{\eta}(t)\right\|}{\sqrt{2\lambda_{1}(t)}}+\frac{\left\|\widetilde{x}_{\eta}(t+1)\right\|}{\sqrt{2\lambda_{1}(t+1)}}+O(\eta(\theta_{t}+\theta_{t+1}))$ . Note that $\lambda_{1}(t)-\lambda_{1}(t+1)=O(\eta^{2})$ and $\theta_{t+1}=O(\theta_{t})$ , we conclude that $\sqrt{L(x_{\eta}(t))}+\sqrt{L(x_{\eta}(t+1))}=\eta\sqrt{\frac{\lambda_{1}(\nabla^{2}L(x_{\eta}(t))}{2}})+O(\eta\theta_{t})$ . ∎

Appendix F Some Useful Lemmas About Eigenvalues and Eigenvectors

Moreover, the functions $\lambda$ and $u$ are $\mathcal{C}^{\infty}$ on $N(X_{0})$ and the differentials at $X_{0}$ are

The next theorem is the Davis-Kahan $\sin(\theta)$ theorem, that bounds the change in the eigenvectors of a matrix on perturbation. Before presenting the theorem, we need to define the notion of unitary invariant norms. Examples of such norms include the frobenius norm and the spectral norm.

Appendix G Analysis of L𝐿\sqrt{L}

The analysis will follow the same line of proof used for the analysis of Normalized GD. Hence, we write down the main lemmas that are different from the analysis of Normalized GD. Rest of the lemmas are nearly the same and hence, we have omitted them.

The major difference between the results of Normalized GD and GD with $\sqrt{L}$ is in the behavior along the manifold $\Gamma$ (for comparison, see Lemma B.12 for Normalized GD and Lemma G.10 for GD with $\sqrt{L}$ ). Another difference between the results of Normalized GD and GD with $\sqrt{L}$ is in the error rates mentioned in Theorem 4.4 and Theorem 4.6. The difference comes from the stronger behavior of the projection along the top eigenvector that we showed for Normalized GD in Lemma E.10, but doesn’t hold for GD with $\sqrt{L}$ (see Lemma G.6). This difference shows up in the sum of angles across the trajectory (for comparison, see Lemma E.5 for Normalized GD and Lemma G.4 for GD with $\sqrt{L}$ ), and is finally reflected in the error rates.

The notations will be the same as Appendix B . However, here we will use $\widetilde{x}_{\eta}(t)$ to denote $\left(2\nabla^{2}L(\Phi(x_{\eta}(t)))\right)^{1/2}(x_{\eta}(t)-\Phi(x_{\eta}(t)))$ . We will now denote $Y$ as the limiting flow given by Equation 7.

G.2 Phase I, convergence

Here, we will show a very similar stability condition for the GD update on $\sqrt{L}$ as the one (Lemma C.1) derived for Normalized GD. Recall our notation $\widetilde{x}_{\eta}(t)=\sqrt{2\nabla^{2}L(\Phi(x_{\eta}(t)))}(x_{\eta}(t)-\Phi(x_{\eta}(t))).$

Suppose $\{x_{\eta}(t)\}_{t\geq 0}$ are iterates of GD with $\sqrt{L}$ (6) with a learning rate $\eta$ and $x_{\eta}(0)=x_{\text{init}}$ . There is a constant $C>0$ , such that for any constant $\bm{\varsigma}>0$ , if at some time $t^{\prime}$ , $x_{\eta}(t^{\prime})\in Y^{\bm{\epsilon}}$ and satisfies $\frac{\left\|x_{\eta}(t^{\prime})-\Phi(x_{\eta}(t^{\prime}))\right\|}{\eta}\leq\bm{\varsigma}$ , then for all $\bar{t}\geq t^{\prime}+C\frac{\bm{\zeta}\bm{\varsigma}}{\bm{\mu}}\log\frac{\bm{\varsigma}\bm{\zeta}}{\bm{\mu}}$ , the following must hold true for all $1\leq j\leq M$ :

provided that for all steps $t\in\{t,\ldots,\bar{t}-1\}$ , $\overline{x_{\eta}(t)x_{\eta}(t+1)}\subset Y^{\bm{\epsilon}}$ .

The proof exactly follows the strategy used in Lemma C.1. We can use the noisy update formulation from Lemma G.7 and the bound on the movement in $\Phi$ from Lemma G.10 to get for any time $t$ with $\bar{t}\geq t\geq t^{\prime}$ (similar to Equation 28):

Hence, similar to Lemma D.1, we can derive the following property that continues to hold true throughout the trajectory, once the condition Equation 42 is satisfied:

We also have the counterpart of Corollary D.3 with the same proof, which follows from using the noisy update of GD on $\sqrt{L}$ from Lemma G.7 and using the quadratic update result from Lemma A.5.

G.3 Phase II, limiting flow

Let $T_{2}$ be the time up until which solution to the limiting flow exists.

Lemma G.10 shows the movement in $\Phi$ , which can be informally given as follows: in each step $t$ ,

provided $\overline{\Phi(x_{\eta}(t))\Phi(x_{\eta}(t+1))}\in Y^{\bm{\epsilon}}$ .

Motivated by this update rule, we show that the trajectory of $\Phi(x_{\eta}(\cdot))$ is close to the limiting flow in Equation 7, for a small enough learning rate $\eta$ . The major difference from Theorem 4.4 comes from the fact that the total error introduced in Equation 43 over an interval $[0,t_{2}]$ is $\sum_{t=0}^{t_{2}}O(\eta^{2}\theta_{t}+\eta^{3})$ , which is of the order $O(\eta^{1/2})$ using the result of Lemma G.4.

The first lemma shows that the sum of the angles in an interval $[0,t_{2}]$ of length $\Omega(1/\eta^{2})$ is at most $O(t_{2}\eta^{1/2}).$

provided $\eta$ is sufficiently small and for all time $0\leq t\leq t_{2}-1$ , $\overline{x_{\eta}(t)x_{\eta}(t+1)}\subset Y^{\bm{\epsilon}}$ .

The proof is very similar to the proof of Lemma E.5, except we replace Lemma E.6 by Lemma G.5 in the analysis of case (B). Hence the final average angle becomes $O(\sqrt{\eta})$ . ∎

The proof will follow exactly as Lemma E.6, except we replace Lemma E.10 by Lemma G.6, which changes the rate into $O(\sqrt{t^{\prime}-\overline{t}}+(t^{\prime}-\overline{t})\sqrt{\eta})$ ∎

provided $G_{t}\geq\Omega(\eta)$ and $\overline{x_{\eta}(t)x_{\eta}(t+1)},\overline{x_{\eta}(t+1)x_{\eta}(t+2)}\subset Y^{\bm{\epsilon}}$ .

Here, we will follow a much simpler approach than Lemma E.10 to have a weaker error bound. The stronger error bounds in Lemma E.10 were due to the very specific update rule of Normalized GD.

First recall $\widetilde{x}_{\eta}(t)=\sqrt{2\nabla^{2}L(\Phi(x))}(x_{\eta}(t)-\Phi(x_{\eta}(t)))$ . By Lemma G.10, we have $\left\|\Phi(x_{\eta}(t+1))-\Phi(x_{\eta}(t))\right\|=O(\eta^{2})$ , thus $\widetilde{x}_{\eta}(t+1)-\widetilde{x}_{\eta}(t)=\sqrt{2\nabla^{2}L(\Phi(x))}(x_{\eta}(t+1)-x_{\eta}(t))=\eta\sqrt{2\nabla^{2}L(\Phi(x))}\frac{\nabla L(x_{\eta}(t))}{2\sqrt{L(x_{\eta}(t))}}$ . From Lemma G.7, we have

where we have used the fact that $\left\|\widetilde{x}_{\eta}(t)\right\|=O(\eta)$ .

Hence, the update is similar to the update in a quadratic model, with $\nabla^{2}L(\Phi(x_{\eta}(t)))$ guiding the updates with an additional $O(\eta^{2})$ perturbation. As a result we also get a $O(\eta^{2})$ perturbation in $G_{t}$ . Here we use the assumption $G_{t}=\Omega(\eta)$ so that GD updates are $O(1)$ -lipschitz. ∎

G.4 Omitted Proof for Operating on Edge of Stability

This proof is similar to that of Theorem 4.7.

If $M=1$ , that is, the dimension of manifold $\Gamma$ is $D-1$ , we know $\overline{x_{\eta}(t)x_{\eta}(t+1)}$ will cross $\Gamma$ , making the $\nabla^{2}\sqrt{L}$ diverges at the intersection and the first claim becomes trivial. If $M\geq 2$ , we have $\nabla^{2}\sqrt{L}=\frac{2L\nabla^{2}L-\nabla L\nabla L^{\top}}{4\sqrt{L}^{3}}$ diverges at the rate of $\frac{1}{\left\|\nabla L\right\|}$ . It turns out that using basic geometry, one can show that the distance from $\Phi(x_{\eta}(t))$ to $\overline{x_{\eta}(t)x_{\eta}(t+1)}$ is $O(\eta(\theta_{t}+\theta_{t+1}))$ , thus $\sup_{0\leq s\leq\eta}\lambda_{1}(\nabla^{2}\sqrt{L}(x_{\eta}(t)-s\nabla\sqrt{L}(x_{\eta}(t))))=\Omega(\frac{1}{\eta(\theta_{t}+\theta_{t+1})})$ . The proof of the first claim is completed by noting that $\theta_{t+1}=O(\theta_{t})$ .

For the second claim, it’s easy to check that $\sqrt{L(x_{\eta}(t))}=\left\|\widetilde{x}_{\eta}(t)\right\|+O(\eta)$ . The proof for the first claim is completed by noting that $\left\|\widetilde{x}_{\eta}(t)\right\|+\left\|\widetilde{x}_{\eta}(t+1)\right\|=\eta\lambda_{1}(t)+O(\eta+\theta_{t})$ as an analog of the quadratic case. ∎

G.5 Geometric Lemmas for L𝐿\sqrt{L}

First recall our notations, $\widetilde{x}=\sqrt{2\nabla^{2}L(\Phi(x))}(x-\Phi(x))$ and $\theta=\arctan\frac{\left\|P_{\Phi(x),\Gamma}^{(2:M)}\widetilde{x}\right\|}{\left|\langle v_{1}(x),\widetilde{x}\rangle\right|}$ .

At any point $x\in Y^{{\bm{\epsilon}}}$ , we have

Since $\Phi(x)$ is a local minimizer of zero loss, we have $\nabla L(\Phi(x))=0$ , we have that

By Lemma B.8, we know $\partial^{2}L(\Phi(x))[x-\Phi(x),x-\Phi(x)]\geq\Omega(\frac{\left\|x-\Phi(x)\right\|^{2}}{\bm{\mu}})$ and therefore

For the second claim, with $\widetilde{x}=\sqrt{2\nabla^{2}L(\Phi(x))}(x-\Phi(x))$ , we have that

By Lemma B.7, we have $\left\|x-\Phi(x)\right\|\leq\frac{2\bm{\mu}}{\bm{\zeta}\bm{\nu}}$ , thus $\frac{\bm{\zeta}^{1/2}\bm{\nu}}{\bm{\mu}}\left\|x-\Phi(x)\right\|=O(\frac{\bm{\mu}}{\bm{\zeta}^{1/2}})=O(\bm{\zeta}^{1/2})$ . ∎

The following two lemmas are direct implications of Lemma G.7.

At any point $x\in Y^{{\bm{\epsilon}}}$ , we have

Consider any point $x\in Y^{{\bm{\epsilon}}}$ . Then,

where $\theta=\arctan\frac{\left\|P_{\Phi(x),\Gamma}^{(2:M)}\widetilde{x}\right\|}{\left|\langle v_{1}(x),\widetilde{x}\rangle\right|}$ , with $\widetilde{x}=\sqrt{2\nabla^{2}L(\Phi(x))}(x-\Phi(x)).$

For any $\overline{xy}\in Y^{\bm{\epsilon}}$ where $y=x-\eta\nabla\sqrt{L(x)}$ is the one step update on $\sqrt{L}$ loss from $x$ , we have

Here $\theta=\arctan\frac{\left\|P_{\Phi(x),\Gamma}^{(2:M)}\widetilde{x}\right\|}{\left|\langle v_{1}(x),\widetilde{x}\rangle\right|}$ , with $\widetilde{x}=\sqrt{2\nabla^{2}L(\Phi(x))}(x-\Phi(x)).$

We outline the major difference from the proof of Lemma B.12. Using taylor expansion for the function $\Phi$ , we have

where in the final step, we used the property of $\Phi$ from Lemma B.16 to kill the first term and use the bound on $\frac{\nabla L(x)}{\sqrt{L(x)}}$ from Lemma G.7 for the third term.

Since the function $\Phi\in\mathcal{C}^{3}$ , hence $\partial^{2}\Phi(x)=\partial^{2}\Phi(\Phi(x))+O(\bm{\chi}\left\|x-\Phi(x)\right\|)$ .

Also, at $\Phi(x)$ , since $v_{1}(x)$ is the top eigenvector of the hessian $\nabla^{2}L$ , we have from Corollary B.23,

where recall our notation of $\theta=\arctan\frac{\left\|P_{\Phi(x),\Gamma}^{(2:M)}(x-\Phi(x))\right\|}{\left|\langle v_{1}(x),x-\Phi(x)\rangle\right|}$ .

With further simplification, it turns out that

The proof is completed by using Corollary B.25. ∎

Appendix H Additional Experimental Details

For running GD on $\sqrt{L}$ , we start from $(x,y)=(14.7,3.)$ , and use a learning rate $\eta=0.5$ . For running Normalized GD on $L$ , we start from $(x,y)=(14.7,-3)$ , and use a learning rate $\eta=5$ .

For Figure 2:

We start Normalized GD from $\langle v_{1},\widetilde{x}(0)\rangle=10^{-4},\langle v_{2},\widetilde{x}(0)\rangle=0.45$ . We use a learning rate of $1$ for the optimization updates.

H.2 Implementation Details for Simulation for the Limiting Flow of Normalized GD

We provide the code for running a single step of the riemannian flow (5) corresponding to Normalized GD. The pseudocode is given in Algorithm 4.

Riemannian gradient update details:

Each update comprises of three major steps: a) computing $\nabla^{3}L(x)[v_{1}(x),v_{1}(x)]$ , b) a projection onto the tangent space of the manifold, and c) few steps of gradient descent with small learning rate to drop back to manifold.