Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability

Introduction

Deep learning has achieved great success in a variety of machine learning applications, and gradient-based algorithms are the prevailing optimization methods for training deep neural networks. However, mathematically understanding the behavior of the optimization methods for deep learning is highly challenging, due to non-convexity, over-parameterization, and complicated architectures. In particular, some recent empirical findings in deep networks contradict the traditional understandings of gradient methods. For example, Wu et al. observed that the solution found by gradient descent has sharpness approximately equal to $2/\eta$ instead of just being smaller than $2/\eta$. Also, Jastrzebski et al. observed that there is a break-even point in the SGD trajectory, and after this point, there is a regularization effect on the loss curvature.

One recent well-known example is the phenomenon called “Edge of Stability" (EOS) (Cohen et al. ). Based on the classical optimization theory, the learning rate $\eta$ of gradient-based method should be smaller than $2/\lambda$ so that the loss can decrease, where $\lambda$ is the largest eigenvalue of the Hessian of the objective, also called “sharpness” in the literature. Otherwise, the loss diverges (even for simple quadratic functions). However, the empirical findings in Cohen et al. show that under various network settings, the EOS phenomena typically occurs along the gradient descent trajectory: (1) the sharpness first increases until it reaches $2/\eta$ (called “progressive sharpening”) (2) the sharpness starts hovering around $2/\eta$ (the EOS regime) and (3) the loss non-monotonically decreases without diverging.

Although (1) seems to be consistent with the traditional beliefs about optimization, a rigorous mathematical explanation for it is still open. Moreover, phenomena (2) and (3) are more mysterious because they violate the $\eta<2/\lambda$ “rule” in traditional optimization theory, yet the training loss does not completely diverge. Instead, the loss may oscillate but still decrease in the long run, while the sharpness seems to be restrained from further increasing.

In this paper, we aim to provide a theoretical and empirical explanation for the mystery of EOS. Towards the goal, we focus on the dynamics of these key quantities when EOS happens and attempt to find out the main driving force to explain these phenomena along the gradient descent trajectory from both theoretical and empirical perspectives.

Our contributions can be summarized as follows.

(Section 3.1) We analyze the typical sharpness behavior along the gradient descent trajectory when EOS happens, and propose a four-phase division of GD trajectory, based on the dynamics of some key quantities such as the loss and the sharpness, for further understanding this phenomenon.

(Section 3.2) We empirically identify the weight norm of the output layer as an effective indicator of the sharpness dynamics. We show that analyzing the dynamics of this surrogate can qualitatively explain the dynamics of sharpness. By assuming this relation, together with some additional simplifying assumptions and approximations, we can explain the dynamics of the sharpness, the loss, and the output layer norm in each phase of EOS (Section 3.3). In this context, we also offers an interesting explanation for the non-monotonic loss decrement (also observed in Cohen et al. , Xing et al. ) (Section 3.4).

(Section 4) Following similar ideas, we provide a more rigorous proof for the progressive sharpening and EOS phenomena in a two-layer fully-connected linear neural network setting based on certain assumptions. The assumptions made here are either weaker or arguably less restrictive.

2 Related work

The structure of Hessian The Hessian matrix carries the second order information of the loss landscape. Several prior works have empirically found that the spectrum of Hessian has several “outliers” and a continuous “bulk” (Sagun et al. , Papyan ). Typically, each outlier corresponds to one class in multi-class classification. As we consider the binary classification setting, there is typically one outlier (i.e., the largest eigenvalue) that is much larger than other eigenvalues. It is consistent with our Assumption 4.1. The Gauss-Newton decomposition of the Hessian was used in several prior works (Martens , Bottou et al. , Papyan ). Papyan empirically showed that the outliers of Hessian can be attributed to a “G component”, which is also known as Fisher Information Matrix (FIM) in Karakida et al. . Also, Wu et al. analyzed the leading Hessian eigenspace by approximating the Hessian with Kronecker factorization and theoretically proved the outliers structure under some random setting assumption.

Neural Tangent Kernel A recent line of work studied the learning of over-parameterized neural networks in the so-called. “neural tangent kernel (NTK) regime or the lazy training regime (Jacot et al. , Lee et al. , Du et al. , Arora et al. , Chizat et al. ). A main result in this regime is that if the neural network is wide enough, gradient flow can find the global optimal empirical minimizer very close to the initialization. Moreover, the Hessian does not change much in the NTK regime. Our findings go beyond NTK setting to analyze the change of sharpness.

Edge of Stability regime The Edge of Stability phenomena was first formalized by Cohen et al. . Similar phenomena were also identified in Jastrzebski et al. as the existence of the “break-even” point on SGD trajectory after which loss curvature gets regularized. Xing et al. observed that gradient descent eventually enters a regime where the iterates oscillate on the leading curvature direction and the loss drops non-monotonically. Recently Ahn et al. studied the non-monotonic decreasing behavior of GD which they called unstable convergence, and discussed the possible causes of this phenomenon. Ma et al. proposed a special subquadratic landscape property and proved that EOS occurs based on this assumption. Arora et al. studied the implicit bias on the sharpness of deterministic gradient descent in the EOS regime. They proved in some specific settings with a varying learning rate (called normalized GD) or with a modified loss $\sqrt{L}$, gradient descent enters EOS and further reduces sharpness. They mainly focus on the analysis near the manifold of minimum loss, but our analysis also applies to the early stage of the training when the loss is not close to the minimum. In particular, our analysis provides an explanation of non-monotonic loss decrease that cannot be explained by their theory. Another difference is that they consider $\sqrt{L}$ (for constant learning rate) where $L$ is a fairly general MSE loss independent of any neural network structure, while our analysis is strongly tied with the MSE loss of a neural network. Very recently, Lyu et al. explained how GD enters EOS for normalized loss (e.g., neural networks with normalization layers), and analyzed the sharpness reduction effect along the training trajectory. The notion of sharpness in their work is somewhat different due to normalization. In particular, they consider the so-called spherical sharpness, that is the sharpness of the normalized weight vector. They also mainly studied the regime where the parameter is close to the manifold of minimum loss as in and proved that GD approximately tracks a continuous sharpness-reduction flow. Lewkowycz et al. proposed a similar regime called “catapult phase” where loss does not diverge even if the largest Hessian eigenvalue is larger than $2/\eta$. Our work mainly considers training in this regime and assumes that the training is not in the “divergent phase” in Lewkowycz et al. . Compared with Lewkowycz et al. , we provide a more detailed analysis in more general settings along gradient descent trajectory.

Preliminaries

Hessian, Fisher information matrix and NTK: In this part, we apply previous works to show that the largest eigenvalue of Hessian is almost the same as the largest eigenvalue of NTK. We use the latter as the definition of the sharpness in this paper. Further details can be found in Appendix F.

In this paper, we use $\bm{\theta}(t)$ to denote the parameter at iteration $t$ (or time $t$) and the sharpness at time $t$ as $\Lambda(t)=\Lambda(\bm{\theta}(t))$. We similarly define ${\bm{M}}(t),{\bm{F}}(t),{\bm{D}}(t),\mathcal{L}(t)$.

Here we show the gradient flow dynamics of the residual vector $\bm{D}(t)$:

A Four-phase Analysis of GD Dynamics

In this section, we study the dynamics of gradient descent and the change of sharpness along the optimization trajectory. We divide the whole training process into four phases, occurring repeatedly in the EOS regime. In Section 3.1, we introduce the four phases. In Section 3.2, we show empirically that the change of the norm of the output layer weight vector almost coincides with the change of the sharpness. In Section 3.3, using this observation, we attempt to explain the dynamics of each phase and provide a mathematical explanation for the changes in the sharpness. In Section 3.4, we explain why the loss decreases but non-monotonically. We admit that a completely rigorous theoretical explanation is still beyond our reach and much of our argument is based on various simplifying assumptions and is somewhat heuristic at some points. Due to space limits, we defer all the proofs in this section to Appendix E.1.

To further understand the properties along the trajectory when EOS happens, we study the behaviors of the loss and the sharpness during the training process. As illustrated in Figure 1, we train a shallow neural network by gradient descent on a subset of 1,000 samples from CIFAR-10 (Krizhevsky et al. ), using the MSE loss as the objective. Notice that the sharpness keeps increasing while the loss decreases until the sharpness reaches $2/\eta$. Then the sharpness begins to oscillate around $2/\eta$ while the loss decreases non-monotonically. This is a typical sharpness behavior in the EOS regime, and consistent with the experiments in .

We divide the training process into four phases according to the evolution of the loss, the sharpness, and their correlation, as shown in Figure 1. The four phases happen cyclically along the training trajectory. We first briefly describe the properties of each phase and explain the dynamics in more detail in Section 3.3.

Phase I: Sharpness $\Lambda<2/\eta$. In this stage, all the eigenvalues of Gram matrix $\bm{M}$ are below the threshold $2/\eta$. In particular, using standard initialization, the training typically starts from this phase, and during this phase the loss keeps decreasing and the sharpness keeps growing along the trajectory. This initial phase is called progressive sharpening (PS) in prior work Cohen et al. . Empirically, the behavior of GD trajectory (as well as the loss and the sharpness) is very similar to that of gradient flow, until the sharpness reaches $2/\eta$ (this phenomena is also observed in Cohen et al. . See Figure 5 or Appendix J.1 in their paper). We note that GD may come back to this phase from Phase IV later.

Phase II: Sharpness $\Lambda>2/\eta$. In this phase, the sharpness exceeds $2/\eta$ and may keep increasing. We will show shortly that the fact that $\Lambda>2/\eta$ causes $|\bm{D}^{\top}\bm{v}_{1}|$ (where $\bm{v}_{1}$ the first eigenvector of ${\bm{M}}$) to increase exponentially (Lemma 3.2). This would quickly lead $\|\bm{D}\|$ to exceed $\|{\bm{Y}}\|$ in a few iterations, which leads the sharpness to start decreasing by Proposition 3.1, hence the training process enters Phase III.

Phase III: Sharpness $\Lambda>2/\eta$ yet begins to gradually drop. Before $\Lambda$ drops below $2/\eta$, Lemma 3.2 still holds, so $|\bm{D}^{\top}\bm{v}_{1}|$ keeps increasing. Proposition 3.1 still holds and thus the sharpness keeps decreasing until it is below $2/\eta$, at which point we enter Phase IV. A distinctive feature of this phase is that the loss may increase due to the exponential increase of $|\bm{D}^{\top}\bm{v}_{1}|$.

Phase IV: Sharpness $\Lambda<2/\eta$. When the sharpness is below $2/\eta$, $|\bm{D}^{\top}\bm{v}_{1}|$ begins to decrease quickly, leading the loss to decrease quickly. At the same time, the sharpness keeps oscillating and gradually decreasing for some iterations. This lasts until the loss decrease to a level that is around its value right before Phase III. The sharpness is still below $2/\eta$ and our training process gets back to Phase I.

2 The Norm of the Output Layer Weight

It is difficult to rigorously analyze the dynamics of the sharpness $\Lambda(t)$. In this subsection, we make an interesting observation, that the change of the norm of the output layer of the network (usually a fully-connected linear layer) is consistent with the change of the sharpness most of the time.

Note that $\bm{M}=(\frac{\partial\bm{F}}{\partial\bm{\theta}})^{\top}(\frac{\partial\bm{F}}{\partial\bm{\theta}})$ can be decomposed as follows:

where the $(i,j)-$entry of $\bm{M_{W}}$ is $(\bm{M_{W}})_{ij}=\left\langle\frac{\partial f(\mathbf{x}_{i})}{\partial\bm{W}},\frac{\partial f(\mathbf{x}_{j})}{\partial\bm{W}}\right\rangle=\bm{A}^{\top}\frac{\partial\bm{h}(\bm{W},\mathbf{x}_{i})}{\partial\bm{W}}\frac{\partial\bm{h}(\bm{W},\mathbf{x}_{j})}{\partial\bm{W}}^{\top}\bm{A}.$

In this expression, intuitively $\|\bm{A}\|$ should be positively related to $\|{\bm{M}}_{\bm{W}}\|$. We empirically observe that the part $\bm{M}_{\bm{A}}=(\frac{\partial\bm{F}}{\partial\bm{A}})(\frac{\partial\bm{F}}{\partial\bm{A}})^{\top}$ has a much smaller spectral norm compared to the whole Gram matrix $\bm{M}$ (see Figure 3(a) and Appendix D), which means $\|\bm{M}_{\bm{W}}\|$ dominates $\|\bm{M}_{\bm{A}}\|$. Therefore, $\|\bm{A}\|$ should be positively correlated with $\|\bm{M}\|$.

The benefit of analyzing $\|\bm{A}\|^{2}$ is that the gradient flow of $\|\bm{A}\|^{2}$ enjoys the following clean formula:

In this work, we do experiments on two-layer linear networks, fully connected deep neural networks, and Resnet18, and all of them have such output layer structures. From Figure 3(a), we can observe that the output layer norm $\|\bm{A}\|^{2}$ and the sharpness $\Lambda$ change in the same direction most of the time along the gradient descent trajectory, i.e., they both increase or decrease at the same time. We note that they may change in different directions very occasionally around the time when $\|\bm{A}(t+1)\|^{2}-\|\bm{A}(t)\|^{2}$ changes its sign (see the experiments in Figure 2).

3 Detailed Analysis of Each Phase

In this section, we explain the dynamics of each phase in more detail. For clarity, we first list the assumptions we need in this section. For different phases, we may need some different assumptions to simplify the arguments. Most of the assumptions are consistent with the experiments or the findings in the literature. Some of them are somewhat stronger, and we also discuss how to relax them.

It is the key observation that we have discussed in Section 3.2. The following are two assumptions about the gradient descent trajectory. The first one assumes that $\bm{D}(t)$ and $\|\bm{A}\|^{2}$ are updated according their first order approximations. Empirical justification of this approximation can be found in Appendix D.1.3.

Assumption 3.3 holds empirically, especially in the progressive sharpening phase (see Figure 5 or Appendix J.1 in Cohen et al. ) when the networks are continuously differentiable. We include these experimental details in Appendix D. See also (Theorems 4.3 and 4.5) in Arora et al. for further theoretical justification. We need this assumption for the proof in the progressive sharpening phase.

Then we state an assumption on the upper bound of the sharpness to restrict the regime we discuss:

(Sharpness upper bound) If the training does not diverge, there exists some constant $B_{\Lambda}$, such that $0<\Lambda(t)\leq\frac{B_{\Lambda}}{\eta}$ for all $t$.

This assumption states that there is an upper bound of the sharpness throughout the optimization process. Actually, in Lewkowycz et al. , they proved that $4/\eta$ is an upper bound of the sharpness in a two-layer linear network with one datapoint, otherwise the training process (loss) would diverge. They empirically found that similar upper bounds exist also for nonlinear activations, albeit with somewhat larger constant $B_{\Lambda}$. In the work, We focus on the case when the loss does not diverge and hence we make Assumption 3.4.

The main set of assumptions we need is about the change of ${\bm{M}}$’s eigendirections.

Denote $\{\bm{v}_{i}\}_{i=1}^{n}$ to be the set of eigenvectors of ${\bm{M}}(t)$. We have three levels of assumptions on ${\bm{M}}$’s eigenspace.

(fixed eigendirections) the set $\{\bm{v}_{i}\}_{i=1}^{n}$ is fixed throughout the phase under consideration;

(principal directions moves slowly) at all time $t$, there is a small constant $\epsilon_{2}\geq 0$ such that $\langle{\bm{v}}_{1}(t),{\bm{v}}_{1}(t+1)\rangle\geq 1-\epsilon_{2}$.

Clearly, these three assumptions are increasingly weaker from (i) to (iii). Assumption 3.5 (i) on the eigenvectors is somewhat strong, and the eigenvectors corresponding to small eigenvalues may change notably in our experiments. We use it to illustrate a basic proof idea of the progressive sharpening phase, but later we relax this assumption to Assumption 3.5 (ii).

Moreover, for the proof in Phase II and III, Assumption 3.5 (iii), which only assume that the main direction changes slightly, is sufficient for our proof. Actually, we note that ${\bm{v}}_{1}(t)$ (the eigenvector corresponding to the largest eigenvalue) changes slowly and the inner product of its initial direction and its direction at the end of the phase is also large (see Appendix D for the empirical verification).

For the proof in Phase II, we need another small technical assumption:

Assume $\bm{D}(t)^{\top}\bm{v}_{1}(t)\geq c\epsilon_{2}\|\bm{D}(t)\|$ for some $c>1$ for some $t=t_{0}$ at the beginning of this phase. Here $\epsilon_{2}$ is defined in Assumption 3.5 (iii).

Assumption 3.6 says that $\bm{D}(t)$ has a non-negligible component in the direction of ${\bm{v}}_{1}$. Since $\epsilon_{2}>0$ is a small constant, this is not a strong assumption as some small perturbation (due to discrete updates) would make the assumption hold for some $c>1$.

3.2 Detailed Analysis

In each phase, we attempt to explain the main driving force of the change of the sharpness and the loss.

Phase I: In this phase, we show that ${\bm{D}}(t)^{\top}{\bm{F}}(t)<0$ under certain assumptions (detailed shortly) on the spectral properties of ${\bm{M}}(t)$ (see Lemma 3.1 below). By Assumption 3.2, we have $\|\bm{A}(t+1)\|^{2}-\|\bm{A}(t)\|^{2}>0$, implying that the sharpness $\Lambda(t)$ also increases based on Assumption 3.1. This phase stops if $\Lambda(t)$ grows larger than $2/\eta$.

We assume the output vector $\bm{F}(t)$ is initialized to be small (this is true if we use very small initial weights). For simplicity, we assume $\bm{F}(0)=0$ in the following argument.

For all $t$ in Phase I, under Assumption 3.5 (i) and 3.3, it holds that ${\bm{D}}(t)^{\top}{\bm{F}}(t)<0$.

From this lemma, $\|{\bm{A}}\|$ keeps increasing by Assumption 3.2; hence the sharpness keeps increasing by Assumption 3.1 until it reaches $2/\eta$ or the loss converges to 0. In the former case, the training process enters Phase II, while the latter case is also possible when $\eta$ is very small (e.g., even the largest possible sharpness value is less than $2/\eta$). We admit that Assumption 3.5 (i) is somewhat strong. In fact, the assumption can be relaxed significantly to Assumption 3.5 (ii). We show in Appendix E.2 that under Assumption 3.5 (ii) and Assumption 3.3, we can still guarantee ${\bm{D}}(t)^{\top}{\bm{F}}(t)<0$. Moreover, we provide a dynamical system view of the dynamics of ${\bm{D}}(t)^{\top}{\bm{F}}(t)$ in that Appendix.

Phase II: When the training process just enters Phase II, the sharpness keeps increasing. We show shortly that $\bm{D}(t)^{\top}\bm{v}_{1}(t)$ starts to increase geometrically, and this causes the sharpness to stop increasing at some point, thus entering Phase III.

In this phase, we adopt a weaker assumption on the sharpness direction ${\bm{v}}_{1}$: Assumption 3.5 (iii). This assumption holds in our experiments (See Figure 17). Also, Assumption 3.6 is necessary.

Suppose Assumption 3.5 (iii) and 3.6 hold during this phase (with constants $\epsilon_{2}>0$ and $c>1$). If $\Lambda(t)=(2+\tau)/\eta$ and $\tau>\frac{1}{1-\epsilon_{2}-1/c}-1$, then $\bm{D}(t)^{\top}\bm{v}_{1}(t)$ increases geometrically with factor $(1+\tau)(1-\epsilon_{2}-1/c)>1$ for $t\geq t_{0}$ in this phase.

Since $\bm{D}(t)^{\top}\bm{v}_{1}(t)$ increases geometrically, $\|{\bm{D}}\|\geq\bm{D}(t)^{\top}\bm{v}_{1}(t)$ will exceed $\|{\bm{Y}}\|$ eventually. Next, the following proposition states that when this happens, $\bm{D}(t)^{\top}\bm{F}(t)>0$. Consequently, $\|\bm{A}\|$ decreases by Assumption 3.2, leading to the decrement of the sharpness based on our Assumption 3.1.

If $\|\bm{D}(t)\|>\|\bm{Y}\|$, then $\bm{D}(t)^{\top}\bm{F}(t)>0$.

Phase III: The sharpness is still larger than $2/\eta$, but it starts decreasing. Meanwhile, the loss continues to increase rapidly due to Lemma 3.2. Eventually, the sharpness will fall below $2/\eta$ and then the training process enters phase IV.

By Lemma 3.2, if the sharpness stays above $2/\eta$, then we can have an arbitrarily large loss. According to Proposition 3.1, if the loss is large enough, the sharpness keeps decreasing.

Now we show that if the sharpness stays above $2/\eta$, $\|{\bm{A}}(t)\|^{2}$ will decrease by a significant amount. This partially explains that the sharpness should also decrease significantly until it drops below $2/\eta$ (instead of decreasingly converging to a value above $2/\eta$ without ever entering the next phase).

Under Assumption 3.2, if $\|{\bm{D}}(t)\|>\|{\bm{Y}}\|$, then $\|\bm{A}(t+1)\|^{2}-\|\bm{A}(t)\|^{2}<-\frac{4\eta}{n}(\|{\bm{D}}(t)\|-\|{\bm{Y}}\|)^{2}$.

From the above argument, we can see that if $\bm{D}(t)^{\top}\bm{v}_{1}(t)$ is larger than $\|{\bm{Y}}\|$, then $\bm{D}(t)^{\top}\bm{v}_{1}(t)$ does not decrease in Phase III, and according to Proposition 3.2, $\|{\bm{A}}(t)\|^{2}$ decreases significantly, implying the sharpness drops below $2/\eta$ eventually.

Remark: The fact that the sharpness can provably drop below $2/\eta$ in this phase can be proved more rigorously in Section 4 for the two-layer linear setting. See Theorem 2.

Phase IV: First, since the training process has just left phase III, $\bm{D}(t)^{\top}\bm{F}(t)$ is still positive and large, hence $\|\bm{A}(t)\|^{2}$ keeps decreasing and the sharpness decreases as well. Since the sharpness stays below $2/\eta$, the loss decreases due to the following descent lemma (with $\bm{u}$ replaced by ${\bm{D}}(t)$).

Next we argue that $\bm{D}(t)^{\top}\bm{F}(t)$ will become negative eventually, which indicates that $\|{\bm{A}}(t)\|^{2}$ and hence the sharpness will grow again. Since the sharpness is below $2/\eta$, $\bm{D}(t)^{\top}\bm{v}_{1}(t)$ decreases geometrically due to Lemma 3.3 (replacing $\bm{u}$ with ${\bm{D}}(t){\bm{v}}_{1}(t){\bm{v}}_{1}(t)^{\top}$).

In fact, ${\bm{D}}$ can be decomposed into the ${\bm{v}}_{1}$-component ${\bm{v}}_{1}{\bm{v}}_{1}^{\top}{\bm{D}}$ and the remaining part ${\bm{R}}$ defined as $\bm{R}(t):=(\bm{I}-\bm{v}_{1}(t)\bm{v}_{1}(t)^{\top})\bm{D}(t)$. Then we have

As shown in the next subsection, ${\bm{R}}(t)$ almost follows a similar gradient descent trajectory ${\bm{R}}^{\prime}(t)$ (Lemma 3.5). More precisely, $\bm{R}^{\prime}(t)$ is defined as $\bm{R}^{\prime}(t+1)=(\bm{I}-\eta\bm{M}(t)(\bm{I}-{\bm{v}}_{1}(t){\bm{v}}_{1}(t)^{\top}))\bm{R}^{\prime}(t)$ (Lemma 3.4). While ${\bm{D}}$’s dynamics is ${\bm{D}}(t+1)=(\bm{I}-\eta\bm{M}(t)){\bm{D}}(t)$, ${\bm{R}}^{\prime}$ follows a similar dynamics ${\bm{R}}(t+1)=(\bm{I}-\eta\bm{M}^{\prime}(t)){\bm{R}}^{\prime}(t)$, where $\bm{M}^{\prime}(t)=\bm{M}(t)(\bm{I}-{\bm{v}}_{1}(t){\bm{v}}_{1}(t)^{\top})$. Note that $\bm{M}^{\prime}(t)$ has eigenvalues smaller than $1/\eta$ for any time $t$ (by Assumption 3.7), hence with an assumption similar to Assumption 3.5 (i) (or a similar version of our relaxed assumption in Appendix E.2 for $\bm{M}^{\prime}$), we can prove that ${\bm{R}}^{\prime}(t)^{\top}({\bm{R}}^{\prime}(t)+\bm{Y})<0$ for any time $t$ (See Appendix E.1 for the rigorous proof).

Since ${\bm{R}}(t)\approx{\bm{R}}^{\prime}(t)$, the second term in the decomposition (4) is always negative and the first term (${\bm{v}}_{1}$ direction term) is decreasing geometrically. Therefore, there are only two possible cases. The first possibility is that the first term decreases to a small value near 0 and the second term remains largely negative. Then their sum will be negative, which is ${\bm{D}}(t)^{\top}{\bm{F}}(t)<0$, thus implying the training enters Phase I. The second possibility is that when the first term decreases to a small value near 0, the second term is also a small negative value. In this case both ${\bm{R}}(t)$ and ${\bm{D}}(t)^{\top}{\bm{v}}_{1}(t)$ are small, implying the loss is almost 0, which is indeed the end of the training.

4 Explaining Non-monotonic Loss Decrement

In this subsection, we attempt to explain the non-monotonic decrement of the loss during the entire GD trajectory. See Figure 3(b). As defined in the last section, we decompose ${\bm{D}}$ into the ${\bm{v}}_{1}$-component ${\bm{v}}_{1}{\bm{v}}_{1}^{\top}{\bm{D}}$ and the remaining part ${\bm{R}}$. Below, we prove that $\bm{R}(t)$ is not affected much by the exponential growth of the loss (Proposition 3.5) in Phase II and III, and almost follows a converging trajectory (which is defined as ${\bm{R}}^{\prime}(t)$ later in this section).

The arguments in this subsection need Assumption 3.4 and Assumption 3.5 (iii), both very consistent with the experiments. We need an additional assumption on the spectrum of ${\bm{M}}$.

All ${\bm{M}}(t)$’s eigenvalues except $\Lambda(t)=\lambda_{\max}({\bm{M}}(t))$ are smaller than $1/\eta$ for all $t$.

Recall that the largest eigenvalue is at most $\frac{B_{\Lambda}}{\eta}$ by Assumption 3.4. Empirically, the largest eigenvalue is an outlier in the spectrum, i.e., it is much larger than the other eigenvalues. Hence, we make Assumption 3.7 which states that all other eigenvalues are at most $1/\eta$, which is consistent with our experiments. See Figure 3(b). Similar fact is also mentioned in .

First, we let $B_{D}$ be an upper bound of $\bm{D}(t)$, i.e., for all $t$, $\|\bm{D}(t)\|\leq B_{D}$. In the two-layer linear network case, we can have an explicit form of $B_{D}$. (see Lemma C.9 in Appendix C.) Recall that in Assumption 3.4, $B_{\Lambda}$ is the upper bound of $\eta\Lambda$.

Suppose Assumption 3.5 (iii) holds. $\bm{R}(t)$ satisfies the following:

Define an auxiliary sequence $\bm{R}^{\prime}(t)$ by $\bm{R}^{\prime}(0)=\bm{R}(0)$, and $\bm{R}^{\prime}(t+1)=(\bm{I}-\eta\bm{M}(t)(\bm{I}-{\bm{v}}_{1}(t){\bm{v}}_{1}(t)^{\top}))\bm{R}^{\prime}(t)$. If Assumption 3.4, Assumption 3.5 (iii), Assumption 3.7 hold, and for any time $t$ there exists a quantity $\lambda_{r}>0$, such that the smallest eigenvalue of ${\bm{M}}(t)$, i.e. $\lambda_{\min}({\bm{M}}(t))>\lambda_{r}$, then there exists a constant $c_{r}>0$ such that $\|\bm{R}(t)-\bm{R}^{\prime}(t)\|\leq c_{r}\frac{B_{D}(B_{\Lambda}-1)\sqrt{\epsilon_{2}}}{\eta\lambda_{r}}$.

Now, in light of Lemma 3.2 and Lemma 3.5, we arrive at an interesting explanation of the phenomena of non-monotonic decrease of the loss. Basically, ${\bm{D}}$ can be decomposed into the ${\bm{v}}_{1}$-component ${\bm{v}}_{1}{\bm{v}}_{1}^{\top}{\bm{D}}$ and the remaining part ${\bm{R}}=(\bm{I}-\bm{v}_{1}\bm{v}_{1}^{\top})\bm{D}$. The ${\bm{v}}_{1}$-component may increase geometrically during the EOS (Lemma 3.2), but the behavior of the remaining part ${\bm{R}}(t)$ is close to $\bm{R}^{\prime}(t)$, which follows the simple updating rule $\bm{R}^{\prime}(t+1)=(\bm{I}-\eta\bm{M}(t))\bm{R}^{\prime}(t)$, so Lemma 3.3 implies that the ${\bm{R}}$ part almost keeps decreasing during the entire trajectory (here Lemma 3.3 applies with $\bm{u}$ replaced by $\bm{R}^{\prime}(t)$, noticing that the eigenvalues except the first are well below $2/\eta$). Hence, the non-monotonicity of the loss is mainly due to the ${\bm{v}}_{1}$-component of ${\bm{D}}$, and the rest part ${\bm{R}}$ is optimized in the classical regime (step size well below 2/(the operator norm)) and hence steadily decreases. See Figure 3(b).

A Theoretical Analysis for 2-Layer Linear NN

In this section, we aim to provide a more rigorous explanation of the EOS phenomenon in two-layer linear networks. The proof ideas follow similar high-level intuition as the proofs in Section 3.3. In particular, we can remove or replace the assumptions in Section 3.3 with arguably weaker assumptions. Due to space limit, we state our main theoretical results and elaborate their relation with the proofs in Section 3.3. The detailed settings and proof are more tedious and can be found in Appendix C.

Dataset: For simplicity, we assume $y_{i}=\pm 1$ for all $i\in[n]$, and $\|\mathbf{X}^{\top}\mathbf{X}\|_{2}=\Theta(n)$. We assume $\mathbf{\mathbf{X}^{\top}\mathbf{X}}$ has rank $r$, and we decompose $\mathbf{\mathbf{X}^{\top}\mathbf{X}}$ and ${\bm{Y}}$ according to the orthonormal basis $\{{\bm{v}}_{i}\}$, the eigenvectors of $\mathbf{\mathbf{X}^{\top}\mathbf{X}}$: $\mathbf{\mathbf{X}^{\top}\mathbf{X}}=\sum_{i=1}^{r}\lambda_{i}\bm{v}_{i}\bm{v}_{i}^{\top},\ {\bm{Y}}=\sum_{i=1}^{r}({\bm{Y}}^{\top}\bm{v}_{i})\bm{v}_{i}:=\sum_{i=1}^{r}z_{i}\bm{v}_{i}$ where ${\bm{v}}_{i}$ is the eigenvector corresponding to the $i$-th largest eigenvalue $\lambda_{i}$ of $\mathbf{\mathbf{X}^{\top}\mathbf{X}}$. $z_{i}={\bm{Y}}^{\top}\bm{v}_{i}$ is the projection of ${\bm{Y}}$ onto the direction ${\bm{v}}_{i}$. Here we suppose $n\gg r$ and the global minimum $(\bm{A}^{*},\bm{W}^{*})$ exists.

Update rule: We write explicitly the GD dynamics of $\bm{D}(t)$: $\bm{D}(t+1)=(\bm{I}-\eta\bm{M}^{*}(t))\bm{D}(t)$, where $\bm{M}^{*}(t)=\frac{2}{mn}(\|\bm{A}(t)\|^{2}\mathbf{X}^{\top}\mathbf{X}+\mathbf{X}^{\top}\bm{W}^{\top}(t)\bm{W}(t)\mathbf{X})-\frac{4\eta}{n^{2}m}(\bm{D}(t)^{\top}\bm{F}(t))\mathbf{X}^{\top}\mathbf{X}$ is the Gram matrix combined with second order terms.

2 Main Theorem and The Proof Sketch

There exists some constant $\chi>1$, s.t. for all $i\in[r-1]$, $\lambda_{i}(\mathbf{X}^{\top}\mathbf{X})\leq\chi\lambda_{i+1}(\mathbf{X}^{\top}\mathbf{X})$. Moreover, $\lambda_{1}(\mathbf{X}^{\top}\mathbf{X})\geq 2\lambda_{2}(\mathbf{X}^{\top}\mathbf{X}).$

There exists $\kappa=\Omega(r^{-1})$ such that $\min_{i\in[r]}\{z_{i}/\sqrt{n}\}\geq\kappa$.

The first assumption is about the eigenvalue spectrum of $\mathbf{X}^{\top}\mathbf{X}$. It guarantees the gap between two adjacent eigenvalues is not very large, and there is a gap between the largest and the second largest eigenvalue. Note the second part of the assumption is a relaxed version of Assumption 3.7. In our CIFAR-10 1k-subset with samples’ mean subtracted, $\lambda_{1}/\lambda_{2}=\chi\approx 3$ (See Figure 19). The second assumes that all component $z_{i}={\bm{Y}}^{\top}{\bm{v}}_{i}$ are not too small.

Suppose Assumption 4.1, Assumption 4.2 hold, the smallest nonzero eigenvalue $\lambda_{r}=\lambda_{r}(\mathbf{X}^{\top}\mathbf{X})>0$ and $\lambda_{1}=\lambda_{\max}(\mathbf{X}^{\top}\mathbf{X})=c_{1}n$. Then for any $\epsilon>0$, if $m=\Omega(\frac{c_{1}n^{2}}{\lambda_{r}^{2}})$, and $n=\Omega(\frac{\lambda_{r}^{2}}{\kappa^{4}\epsilon^{2}})$, we have the progressive sharpening property: $\Lambda(t+1)-\Lambda(t)>0$ for $t=1,2,...,t_{0}-1$ where $t_{0}$ is the time when $\|{\bm{D}}(t)\|^{2}\leq O(\epsilon^{2})$ or $\lambda_{\max}({\bm{M}}^{*}(t))>1/\eta$ for the first time.

In the proof of this theorem, we show that the Gram matrix ${\bm{M}}(t)\approx\frac{2}{mn}(\|\bm{A}(t)\|^{2}+\frac{m}{d})\mathbf{X}^{\top}\mathbf{X}$, which serves as a justification of Assumption 3.5 we made in Section 3.3. That shows all ${\bm{M}}(t)$ approximately share the same set of eigenvectors as $\mathbf{X}^{\top}\mathbf{X}$. In our proof, we also prove more rigorously that $\|\bm{A}(t)\|^{2}$ is an indicator of the sharpness in this simpler setting.

There exists some constant $c_{2}>0$, such that $\|\bm{\Gamma}(t)\|\leq\frac{c_{2}}{m}.$

This assumption is based on Theorem 1. In Theorem 1, we state that in the progressive sharpening phase, $\|\bm{\Gamma}(t)\|$ has an upper bound of $O(1/m)$. Now in the EOS phase, we assume that $\|\bm{\Gamma}(t)\|$ grows larger by at most a constant factor. Further discussions refer to Appendix D.2.2.

There exists some constant $\beta>0$, such that $\Lambda\leq\frac{4}{\eta}(1-\beta).$

This assumption is consistent with Assumption 3.4, which assumes an upper bound of the sharpness.

There exist some constant $c_{3}$ such that $|{\bm{D}}(t)^{\top}{\bm{v}}_{1}|>c_{3}\sqrt{n}/m$ for some $t=t_{0}$ at the beginning of phase II.

This assumption is in the same spirit of Assumption 3.6 with the only change of the bound in terms of $m$ and $n$. Now, we are ready to state our theorem in this stage.

Denote the smallest nonzero eigenvalue as $\lambda_{r}\triangleq\lambda_{r}(\mathbf{X}^{\top}\mathbf{X})>0$ and the largest eigenvalue as $\lambda_{1}\triangleq\lambda_{1}(\mathbf{X}^{\top}\mathbf{X})$. Under Assumption 4.3, 4.4, 4.5, and $\lambda_{1}(\mathbf{X}^{\top}\mathbf{X})\geq 2\lambda_{2}(\mathbf{X}^{\top}\mathbf{X})$ in Assumption 4.1, there exist constants $c_{4},c_{5},c_{6}$ such that if $n>c_{6}\lambda_{r}\eta,m>\max\{\frac{c_{4}d^{2}n^{2}}{\lambda_{r}^{2}},c_{5}\eta\}$, then

There exists $\rho=O(1)$ which depends on $c_{3}$ such that if $\Lambda(t_{0})>\frac{2}{\eta}(1+\rho)$ for some $t_{0}$, there must exist some $t_{1}>t_{0}$ such that $\Lambda(t_{1})<\frac{2}{\eta}(1+\rho)$.

If $\Lambda(t),\Lambda(t+1)>\frac{2}{\eta}(1+\rho)$, then there is a constant $c_{7}>0$ (depending on $c_{3}$) such that $|{\bm{D}}(t+1)^{\top}{\bm{v}}_{1}|>|{\bm{D}}(t)^{\top}{\bm{v}}_{1}|(1+c_{7})$.

Define ${\bm{R}}(t):=(\bm{I}-{\bm{v}}_{1}{\bm{v}}_{1}^{\top}){\bm{D}}(t)$, and ${\bm{R}}^{\prime}(t):=(\bm{I}-\eta{\bm{M}}^{*}(t)(\bm{I}-{\bm{v}}_{1}{\bm{v}}_{1}^{\top})){\bm{R}}^{\prime}(t-1)$. It holds that $\|{\bm{R}}(t)-{\bm{R}}^{\prime}(t)\|=O(\frac{\sqrt{n^{3}}d}{\lambda_{r}\sqrt{m}})$.

We can conclude the following from Theorem 2: (1) The first statement of the theorem states that if the progressive sharpening phase causes the sharpness to grow over $2/\eta$, then the sharpness eventually goes below $2/\eta$. This illustrates the regularization effect of gradient descent on the sharpness (this is consistent with the analysis of Phase III in Section 3.3). (2) The second states that $|{\bm{D}}(t)^{\top}{\bm{v}}_{1}|$ geometrically increases in Phase II and III. Note that we proved a similar Lemma 3.2 for Phase II in the more general setting in Section 3.3. (3) The third conclusion gives an upper bound for the distance between ${\bm{R}}(t)$’s trajectory and ${\bm{R}}^{\prime}(t)$’s. This bound helps illustrate why ${\bm{R}}(t)$’s trajectory is similar with ${\bm{R}}^{\prime}(t)$ in Phase IV of Section 3.3.

Discussions and Open Problems

In this section, we discuss the limitation of our theory and some related findings. First, our argument crucially relies on the assumption that $\|{\bm{A}}\|$ changes in the same direction as $\Lambda$ does most of the time. Here, we elaborate more on this point. Seeing from a longer time scale, $\|{\bm{A}}\|^{2}$ and the sharpness may have very different overall trends (See Figure (c) in 2), i.e., the sharpness oscillates around $2/\eta$ but $\|{\bm{A}}\|^{2}$ increases. Moreover, the sharpness may oscillate more frequently than $\|{\bm{A}}\|^{2}$, while the low-frequency trends seem to match well (See the late training phases in Figure (b) in 2). Currently, our theory cannot explain the high-frequency oscillation of the sharpness in Figure (b).

While we still believe the change of $\|{\bm{A}}\|$ is a major driving force of the change of the sharpness, other factors (such as other layers) must be taken into consideration for a complete understanding and explanation of the sharpness dynamics. We also carry out some experiments that reveal some interesting relation between the inner layers and the sharpness, which is not yet reflected in our theory. Due to space limit, we defer it to Appendix D.3.

We conclude with some open problems. It would be very interesting to remove some of our assumptions or replace them (especially those related to the spectrum of ${\bm{M}}$) by weaker or more natural assumptions on the data or architectures, or make some of the heuristic argument more rigorous (e.g., first order approximation of the dynamics (3)). Extending our results in Section 4 to deeper neural networks with nonlinear activation function is an intriguing and challenging open problem.

References

Appendix A Experimental Setup

We provide detailed experimental setup in this section.

Under GPU memory constraints when computing Gram matrix of our models, we limit the size of dataset we use. The dataset is a 1,000-sample subset from CIFAR-10 (Krizhevsky et al. ) (https://www.cs.toronto.edu/~kriz/cifar.html). To make it a binary dataset, we constructed it by selecting the first 500 samples of class 0 and 1, respectively.Cohen et al. selects the first 5,000 examples from CIFAR-10. Our subset is smaller because of the computation limitation when calculating the Gram matrix. Experiments show that the properties along GD trajectory (e.g. the loss, the sharpness) is similar on both datasets. Then we label the samples by $\pm 1$. In the experiments in Appendix B and D, we fix the objective as training on the 1000-example subset of CIFAR-10.

A.2 Network Architecture

In general settings, we experiment with three architectures from simple models to more complicated models: one-hidden-layer linear neural network, four-hidden-layer fully-connected network, convolutional network and a ResNet18 (He et al. ) model. The initialization of each layer follows the default initialization of PyTorch (Paszke et al. ).

Linear Network We first use a simple two-layer linear neural network. The hidden layer has 200 neurons. We empirically show that even a simple linear network can enter the EOS regime.

Fully-connected Network We conduct further experiments on several different fully-connected networks with 4 hidden layers with various activation functions. We consider tanh, ReLU, ELU activations. For example, the structure of a fully-connected tanh network is shown in Table 2.

Convolutional Network We also conduct experiments on several different convolutional networks with two convolutional layers and two max-pooling layers. Like the fully-connected network experiments, we consider tanh, ReLU and ELU activations. For example, the structure of a convolutional tanh network is shown in Table 3.

ResNet18 We also conduct experiment on the ResNet18 architecture proposed by He et al. . We use the default architecture implemented in PyTorch (Paszke et al. ). When calculating the sharpness, we use the numerical methods in the package (Golmant et al. ) to calculate the top eigenvalue of the Hessian matrix.

Appendix B Further Experiments and Discussions

In this appendix, we use the 1000-sized subset of CIFAR-10 introduced in Appendix A to conduct further experiments on various architectures. We verify our main observation about the correlation between the sharpness and the weight norm of the output layer (A-norm) of the neural network through the following experiments.

We consider simple linear networks, fully-connected networks, convolutional networks in this appendix. For nonlinear networks, we choose tanh, ReLU, ELU as activation functions. We train the networks with MSE loss. Here we exclude ResNet18 experiment since it is not feasible to compute the Gram matrix or the leading Gram matrix eigenvector due to GPU memory limitation. The sharpness and A-norm correlation of ResNet18 are included in Figure 2 in Section 3.2.

Here we run full-batch gradient descent with a selected step size $\eta$ such that the sharpness at initialization is smaller than $2/\eta$.

We first verify our four-phase division of EOS phenomena in a simple linear network. This experiment shows that even linear networks can also enter EOS regime and the four-phase division of the gradient descent trajectory is quite apparent in this setting. The following Figure 4 illustrates the positive correlation between the sharpness and the A-norm, and the relationship between the loss $\|{\bm{D}}(t)\|^{2}$ and $\|{\bm{R}}(t)\|^{2}$ along the trajectory.

B.1.2 Fully-connected networks

We train three 5-layer fully-connected networks with different activation functions: tanh, ReLU, and ELU activations. In Figure 5, 6, and 7, we verify that in these fully-connected networks, the sharpness is positively correlated to the dynamics output layer norm (A-norm) most of the time in the progressive sharpening stage and the first few oscillations.

Note that anomaly points appear much more frequently after a few oscillations. Meanwhile, the sharpness oscillates more frequently around $2/\eta$ in every few iterations. We further discuss the phenomenon in Appendix B.2, in which we elaborate the complicated relationship of the sharpness, A-norm and parameters of other layers.

B.1.3 Convolutional Networks

We train three convolutional networks with different activation functions: ReLU, tanh, and ELU activations. In Figure 8, 9, and 10, we verify that in convolutional networks, the positive correlation between the sharpness and the output layer norm (A norm) is still correct in the training process. In the first few oscillations of the sharpness, the four-phase division is also valid. On the other hand, we notice that the same anomaly appears in this convolutional setting as in the fully-connected examples. We defer the discussion to Section B.2.

B.1.4 Gaussian Data on Linear Networks

In this section, we train two-layer fully-connected linear networks with datapoints sampled from Gaussian distribution and different label vectors.

In Figure 11, we verify that even with simple two-layer linear network and Gaussian data, progressive sharpening and EOS can still be observed. However, because Gaussian data is easy for the network to learn, the convergence is so fast that within tens of epochs the training loss converges to zero. In this case, the EOS phenomenon is not quite typical.

In Figure 12(a), the label ${\bm{Y}}$ is uniformly sampled from $\text{Unif}\{-1,1\}^{n}$. Comparing the experimental results with those obtained using the same network trained with 500 CIFAR-10 data points, we find out that both have a similar degree of progressive sharpening. In Figure 12(b) and Figure 12(c), we let the label ${\bm{Y}}$ be $\sqrt{n}v_{1}$ and $\sqrt{n}v_{400}$ respectively. The result shows that when the label vector is aligned with the top eigenvector $v_{1}$, the degree of progressive sharpening is relatively small, and when the label vector is aligned with a bottom eigenvector, the degree of progressive sharpening is relatively large.

This phenomenon is consistent with our intuition. Empirically we found that a dataset that is easier to learn leads to faster convergence rate, which then leads to a smaller degree of progressive sharpening. For example, training standard Gaussian distributed dataset converges faster than that with CIFAR-10 dataset, hence the degree of progressive sharpening of the former is much smaller; as another example (Figure 12(b) and Figure 12(c)), if the label vector is aligned with the top eigenvector, the convergence in the first phase (the PS phase) is faster, which leads to a shorter first phase and thus a smaller degree of progressive sharpening compared to the other case.

B.2 Further Discussions on the Relation between A-norm and Sharpness

In our paper, we use the observation that $\|\bm{A}\|^{2}$ shares the same trend as the sharpness to explain the dynamics of the sharpness. However, as shown in the experiments in this section (see Figure 5, 6, 7, 8, 9, 10), anomaly points exist during the training process. Here we briefly discussion these anomaly points further. We divide them into three kinds.

At the time that $\|\bm{A}\|^{2}$ changes its trend: This kind of anomaly points appear when $\|\bm{A}\|^{2}$ changes its trend. When $\|\bm{A}\|^{2}$ changes its trend, its changing rate (or differential, the $\bm{D}^{\top}\bm{F}$ term) changes its sign, hence the changing rate’s absolute value is small. Now because the dynamics of the sharpness is effected by both $\|\bm{A}\|^{2}$ and the inner layers, when changing rate of $\|\bm{A}\|^{2}$ is small, the inner layers may play a larger role in the direction of sharpness, which may cause the anomaly.

When the sharpness oscillates more frequently: We notice that in some cases, $\|\bm{A}(t)\|^{2}$ and the sharpness $\Lambda(t)$ have very similar overall trend, but the sharpness oscillates more frequently but the magnitude is small (i.e., the sharpness curve has higher frequency oscillations. See Figure 6 or Figure 10). In this case, while we believe the change of $\|\bm{A}\|^{2}$ is a major driving force of the change of the sharpness, other layers must be taken into consideration for understanding such small and frequent oscillations of sharpness.

In late training phases: We notice that in the late training phases in most settings, sharpness oscillates and crosses $2/\eta$ more frequently (it changes directions in a few iterations). In this case, our four-phase division does not strictly apply. At the same time, $\|\bm{A}\|^{2}$ may also change direction more frequently, resulting more anomaly points during this period of time. Hence, understanding the behavior of the late stages is beyond our current analysis and requires new insights.

Appendix C Proof for Section 4

Model: In this section, we study a two-layer neural network with linear activation, i.e.

Also, we suppose there exists some parameter ${\bm{A}}^{*}$, $\bm{W}^{*}$, s.t.

Note when $n\gg d$, the matrix $\mathbf{X}^{\top}\mathbf{X}$ cannot be full rank. This condition guarantees that Gradient Descent (GD) only travels in the column space of $\mathbf{X}^{\top}\mathbf{X}$ according to Lemma E.10 at Appendix E.

Define the output vector as ${\bm{F}}$ and the residual vector as ${\bm{D}}$ in the training process.

(Note that they are functions of time $t$. So we use ${\bm{D}}(t),{\bm{F}}(t),{\bm{A}}(t),{\bm{W}}(t)$ to denote ${\bm{D}},{\bm{F}},{\bm{A}},{\bm{W}}$ at time $t$.)

Consider the mean square error MSE loss during the training process:

Initialization: We run GD on the loss and start from symmetric initialization for the weights, which guarantees ${\bm{F}}(0)=0.$

Sharpness: Recall the definition of the sharpness: $\Lambda(t)=\lambda_{\max}({\bm{M}}(t))$, where ${\bm{M}}(t)$ is the Gram matrix (8).

Learning rate selection: We select a learning rate $\eta$ such that $\Lambda(0)<2/\eta$. Specially, based on our method of initialization, $\Lambda(0)=\frac{m\lambda_{1}(d+1)}{d}\cdot\frac{2}{mn}$, hence we have

Gradient Descent Update Rule: Following GD, the training dynamics is as follows:

Write them into matrix forms and we have:

Then we have the $\|{\bm{A}}(t)\|$’s dynamics according to $\eqref{C1}$:

We can also show the dynamics of ${\bm{D}}(t)$ according to (6) and (8). See Appendix E for a proof.

(See Lemma 21) The update rule of the residual vector of ${\bm{D}}(t)$:

We define some extra notations for preparation.

Remark. Here we explain the notations above and what they are for.

${\bm{M}}^{*}(t)$ is the Gram matrix ${\bm{M}}(t)$ plus a second order term ($-\frac{4\eta}{n^{2}m}({\bm{D}}(t)^{\top}{\bm{F}}(t))\mathbf{X}^{\top}\mathbf{X}$). Since the second order term is small, we have ${\bm{M}}(t)\approx{\bm{M}}^{*}(t)$. Also, in Section C.3 we prove ${\bm{M}}^{*}(t)$ is almost a linear interpolation of ${\bm{M}}(t)$ and ${\bm{M}}(t+1)$, which corroborates this argument.

$\bm{\Gamma}(t)$ is the difference between $\frac{2}{mn}\mathbf{X}^{\top}{\bm{W}}(t)^{\top}{\bm{W}}(t)\mathbf{X}$ and $\frac{2}{mn}\mathbf{X}^{\top}{\bm{W}}(0)^{\top}{\bm{W}}(0)\mathbf{X}$. Note that $\bm{\Gamma}(0)=0$ because of the initialization condition. In the following proof, we will show a small upper bound of the norm of $\bm{\Gamma}(t)$. Hence, we can show that the eigenvectors of ${\bm{M}}(t)$ are approximately aligned with the eigenvectors of $\mathbf{X}^{\top}\mathbf{X}$.

$\widetilde{\bm{M}}(t)$ is the Gram matrix with the second order term, yet with $\bm{\Gamma}(t)$ excluded. $\widetilde{\bm{M}}(t)$ is used for bounding the main part of $\bm{\Gamma}(t)$ (i.e., the terms without the noise ${\bm{e}}(t)$, e.g. (A1),(A5)) which is defined in Theorem 3) in the proof of Theorem 3. Note that there exists some number $\gamma$ s.t. $\widetilde{\bm{M}}(t)=\gamma\mathbf{X}^{\top}\mathbf{X}$. This is because

Hence for any eigenvector ${\bm{v}}_{i}$ of $\mathbf{X}^{\top}\mathbf{X}$, $\lambda_{i}(\widetilde{\bm{M}}(t))={\bm{v}}_{i}^{\top}\widetilde{\bm{M}}(t){\bm{v}}_{i}.$

Then the dynamics of ${\bm{D}}(t)$ can be written as:

Update Rule of ${\bm{M}}(t)$: Using the update rule of ${\bm{A}}(t)$ and ${\bm{W}}(t)$, we can also derive the update rule of the Gram matrix ${\bm{M}}(t)$ (see Appendix E for the proof detail):

(See Lemma E.9) The update rule of the Gram matrix ${\bm{M}}(t)$ is,

C.2 Phase I and Progressive Sharpening

We suppose $\mathbf{X}^{\top}\mathbf{X}$ has rank $r$, and we decompose $\mathbf{X}^{\top}\mathbf{X}$ and ${\bm{Y}}$ into a basis $\{{\bm{v}}_{i}\}$ composed of orthonormal eigenvectors of $\mathbf{X}^{\top}\mathbf{X}$.

where ${\bm{v}}_{i}$ is the eigenvector corresponding to the $i$-th largest eigenvalue $\lambda_{i}$ of $\mathbf{X}^{\top}\mathbf{X}$. $z_{i}={\bm{Y}}^{\top}{\bm{v}}_{i}$ is the ${\bm{Y}}$’s projection onto the direction ${\bm{v}}_{i}$. In particular, ${\bm{v}}_{1}$ is the eigenvector corresponding to the largest eigenvalue of $\mathbf{X}^{\top}\mathbf{X}$.

Recall that the sharpness $\Lambda(t)=\lambda_{\max}({\bm{M}}(t))$. Here we propose an approximation of sharpness $\Lambda$ by the eigenvector ${\bm{v}}_{1}$:

The following corollary rationalizes this approximation. The proof is deferred to Corollary C.2 in Section C.3.

There exists some constant $c_{2},c_{4}$, s.t. for all time t $\langle{\bm{v}}_{1},\bm{u}\rangle^{2}\geq 1-c_{4}\sqrt{\frac{1}{nm}}$ and

where $\bm{u}$ is the top eigenvector of ${\bm{M}}(t)$.

By this corollary, we can see that the top eigenvector of ${\bm{M}}(t)$ is approximately ${\bm{v}}_{1}$, and the approximation $\Lambda^{*}(t)$ is very close to the real sharpness $\Lambda(t)$. Thus, we consider $\Lambda^{*}(t)$ as an approximation of the sharpness and analyze its dynamics.

We first show the main theorem of progressive sharpening under Assumption C.1 and C.2 below.

There exists some constant $\chi>1$, s.t. for all $i\in[r-1]$, $\lambda_{i}(\mathbf{X}^{\top}\mathbf{X})\leq\chi\lambda_{i+1}(\mathbf{X}^{\top}\mathbf{X})$. Moreover, $\lambda_{1}(\mathbf{X}^{\top}\mathbf{X})\geq 2\lambda_{2}(\mathbf{X}^{\top}\mathbf{X}).$

There exists $\kappa=\Omega(r^{-1})$ such that $\min_{i\in[r]}\{z_{i}/\sqrt{n}\}\geq\kappa$.

The first assumption is about the eigenvalue spectrum of $\mathbf{X}^{\top}\mathbf{X}$. It guarantees the gap between two adjacent eigenvalues is not very large, and there is a gap between the largest and the second largest eigenvalue. Note the second part of the assumption is a relaxed version of Assumption 3.7. In our CIFAR-10 1k-subset with samples’ mean subtracted, $\lambda_{1}/\lambda_{2}=\chi\approx 3$ (See Appendix D.2.1, Figure 19). The second assumes that all components $z_{i}={\bm{Y}}^{\top}{\bm{v}}_{i}$ are not too small.

Under the two assumptions, we have the following theorem:

Suppose Assumption C.1, C.2 hold. Suppose $\lambda_{r}=\lambda_{\min}(\mathbf{X}^{\top}\mathbf{X})>0,\lambda_{1}=\lambda_{\max}(\mathbf{X}^{\top}\mathbf{X})=\Theta(n)$. For any $\epsilon>0$, if $m=\Omega(\frac{n^{2}}{\lambda_{r}^{2}})$, and $n=\max\{\Omega(\frac{\lambda_{r}}{\kappa^{2}}),\Omega(\frac{\lambda_{r}^{2}}{\kappa^{2}}),\Omega(\frac{\lambda_{r}^{2}}{\kappa^{4}\epsilon^{2}}),\Omega(\frac{\lambda_{r}^{4}}{\kappa^{4}\epsilon^{2}})\}$, we have the following properties for $t=1,2,...,t_{0}-1$ where $t_{0}$ is the time when $\|{\bm{D}}(t)\|^{2}=O(\epsilon^{2})$ or $\lambda_{\max}({\bm{M}}^{*}(t))>1/\eta$ for the first time.

(Progressive Sharpening) (Lemma C.6) $\Lambda^{*}(t+1)-\Lambda^{*}(t)>0.$

We first prove Lemma C.3, which implies that ${\bm{D}}(t)$ approximately converges independently in each direction of $\mathbf{X}^{\top}\mathbf{X}$’s non-zero eigenspace. Lemma C.3 also proves the second point $\|\bm{\Gamma}(t)\|\leq O(1/m)$. (Recall that $\|\bm{\Gamma}(t)\|$ is defined in (10).)

Then, based on the conclusion in Lemma C.3, Assumption C.1 and C.2, we prove that $\|{\bm{A}}(t)\|^{2}$ grows in Lemma C.4 and Lemma C.5. Specifically, here we prove a sufficient condition of $-{\bm{D}}(t)^{\top}{\bm{F}}(t)>0$.

Finally, we use the dynamics of $\Lambda^{*}(t)$ and its dependence on $\|{\bm{A}}(t)\|^{2}$ dynamics to prove the sharpness grows (Lemma C.6), thus finishing the proof.

Remark 1. In the theorem, $\lambda_{\max}({\bm{M}}^{*}(t))>1/\eta$ is the termination condition for this theorem. Here, ${\bm{M}}^{*}(t)$ is the Gram matrix ${\bm{M}}(t)$ plus the second order term ($-\frac{4\eta}{n^{2}m}({\bm{D}}(t)^{\top}{\bm{F}}(t))\mathbf{X}^{\top}\mathbf{X}$). Since the second order term is small, we have ${\bm{M}}(t)\approx{\bm{M}}^{*}(t)$. Also, in Section C.3, we prove ${\bm{M}}^{*}(t)$ is almost a linear interpolation of ${\bm{M}}(t)$ and ${\bm{M}}(t+1)$. Thus, $\lambda_{\max}({\bm{M}}^{*}(t))$ can be seen as an approximation of the sharpness.

Remark 2. From the experiments, one can see that gradient descent is in progressive sharpening phase until the sharpness crosses the threshold $2/\eta$. Right now, our proof only works till $\lambda_{\max}({\bm{M}}^{*}(t))$ reaches $1/\eta$. It would be interesting to extend our result to $2/\eta$.

Remark 3. In this theorem, we require mild over-parameterization ($m=\Omega(n^{2})$, assuming $\lambda_{r}=\Theta(1)$, to prove the direction guarantee of the Gram matrix ${\bm{M}}(t)$ and the monotone increment of $\|{\bm{A}}(t)\|^{2}$. Astute readers may find it similar to the NTK regime, where the parameters do not move far from the initialization. However, we stress that our analysis does not necessarily require NTK regime, and can go beyond NTK. We defer the discussion on the difference between our results and the NTK regime to Appendix C.4.

Then we break the proof of theorem into lemmas. The first lemma (Lemma C.3) proves that the Gram matrix ${\bm{M}}(t)\approx\frac{2}{mn}(\|\bm{A}(t)\|^{2}+\frac{m}{d})\mathbf{X}^{\top}\mathbf{X}$ by bounding the movement of $\|\bm{\Gamma}(t)\|$. In this way, ${\bm{D}}(t)\approx-\prod_{j=0}^{t-1}(\bm{I}_{n}-\eta\widetilde{\bm{M}}(j)){\bm{Y}}$ can approximately descend in each eigenvector ${\bm{v}}_{i}$ of $\mathbf{X}^{\top}\mathbf{X}$ independently. Also, by proving ${\bm{M}}(t)\approx\frac{2}{mn}(\|\bm{A}(t)\|^{2}+\frac{m}{d})\mathbf{X}^{\top}\mathbf{X}$, we justify $\|{\bm{A}}(t)\|^{2}$ as an indicator of the sharpness.

Along GD training trajectory, if $m\geq\frac{112c_{1}n^{2}}{\lambda_{r}^{2}}$ , we have the following properties until $\lambda_{\max}({\bm{M}}^{*}(t))>1/\eta$:

$\|{\bm{A}}(t)\|^{2}\geq m/2$ and $\|\bm{\Gamma}(t)\|\leq R_{w}:=40/m$.

${\bm{D}}(t)=-\prod_{j=0}^{t-1}(\bm{I}_{n}-\eta\widetilde{\bm{M}}(j)){\bm{Y}}+{\bm{e}(t)}$, where $\|\bm{e}(t)\|_{2}\leq\frac{40n^{3/2}}{\lambda_{r}m}$

$\lambda_{r}(\widetilde{\bm{M}}(t))\geq\lambda_{r}/n,\ \lambda_{r}({\bm{M}}^{*}(t))\geq\lambda_{r}/n$.

We first consider the base case. For property 1, $\|{\bm{A}}(0)\|^{2}=m$ and $\|\bm{\Gamma}(0)\|=0$. For property 2, with symmetric initialization, ${\bm{F}}(0)=0$, hence ${\bm{D}}(0)=-{\bm{Y}}$, $\|\bm{e}(0)\|=0$. For property 3,

The minimal eigenvalue of $\frac{2}{nm}\|{\bm{A}}(0)\|^{2}\mathbf{X}^{\top}\mathbf{X}$ is $2\lambda_{r}/n>\lambda_{r}/n$. Thus all the properties hold at iteration $t=0$.

Suppose for all $k\leq t$, these properties hold. Then we consider the case in iteration $t+1$.

We first show a worst-case upper bound for $\|{\bm{D}}(t)\|$.

The second inequality uses property 3 in $k\leq t$. Note $\|{\bm{Y}}\|=\sqrt{n}$ in our setting, so the last equality holds. That means $\|{\bm{D}}(t)\|\leq\sqrt{n}$ and $\|{\bm{F}}(t)\|=\|{\bm{D}}(t)+{\bm{Y}}\|\leq 2\sqrt{n}.$

Now, we show the error $\bm{e}(t)$ is bounded. We have

$\begin{aligned} {\bm{D}}(t+1)&=(\bm{I}-\eta\widetilde{{\bm{M}}}(t)){\bm{D}}(t)-\eta\bm{\Gamma}(t){\bm{D}}(t)\\ &=-\prod_{j=0}^{t}(\bm{I}-\eta\widetilde{{\bm{M}}}(j)){\bm{Y}}+(\bm{I}-\eta\widetilde{{\bm{M}}}(t))\bm{e}(t)-\eta\bm{\Gamma}(t){\bm{D}}(t)\\ &=-\prod_{j=0}^{t}(\bm{I}-\eta\widetilde{{\bm{M}}}(j)){\bm{Y}}+(\bm{I}-\eta{\bm{M}}^{*}(t))\bm{e}(t)-\eta\bm{\Gamma}(t)\prod_{j=0}^{t-1}(\bm{I}-\eta\widetilde{{\bm{M}}}({j})){\bm{D}}(0).\end{aligned}$

Hence we have $\bm{e}(t+1)=\bm{e}(t)(\bm{I}-\eta{\bm{M}}^{*}(t))-\eta\bm{\Gamma}(t)\prod_{j=0}^{t-1}(\bm{I}-\eta\widetilde{{\bm{M}}}(j)){\bm{D}}(0)$.

Then we can have the following bound by this recursion:

Here the third inequality holds because $\|\bm{e}(0)\|=0$. Thus the second property holds for $t+1$.

Then we consider the lower bound of $\|{\bm{A}}(t)\|^{2}$.

Consider the dynamics of $\|{\bm{A}}(t)\|^{2}$ in (7) and sum the difference up from 0 to $t$. Recall that we proved $\|{\bm{D}}(t)\|\leq\sqrt{n}$ by (D) above.

Since $\|{\bm{A}}(0)\|^{2}=m$, we have the lower bound $\|{\bm{A}}(t)\|^{2}\geq m/2$.

Next, we lower bound the minimal non-zero eigenvalue (i.e. the $r$-th largest eigenvalue since $\mathbf{X}^{\top}\mathbf{X}$ is rank $r$) of ${\bm{M}}^{*}(t+1)$ and $\widetilde{\bm{M}}(t+1)$.

For ${\bm{M}}^{*}(t+1)$ and $\widetilde{\bm{M}}(t+1)$, since the $\mathbf{X}^{\top}{\bm{W}}(t)^{\top}{\bm{W}}(t)\mathbf{X}$ part is PSD, we have

From the inequality in (# ‣ C.2), and from (D) $\|{\bm{D}}(t)\|\leq\sqrt{n},\|{\bm{F}}(t)\|\leq 2\sqrt{n}$, we know

Thus ${\bm{M}}^{*}(t+1)$ (and also $\widetilde{\bm{M}}(t+1)$) has its smallest eigenvalue larger than $\lambda_{r}/n$. Property 3 holds. To extend this conclusion, we actually have $\lambda_{i}(\widetilde{\bm{M}}(t+1))\geq\lambda_{i}/n$ for all $i\in[r]$ from the argument above.

Finally we bound $\|\bm{\Gamma}(t+1)\|$. We first write down the dynamics of $\bm{\Gamma}(t)$ according to Lemma C.2.

Then we bound these three terms one by one.

Term (G1): $\frac{4}{n^{2}m}\eta\left\|\sum_{k=0}^{t}\left({\bm{D}}(t){\bm{D}}(t)^{\top}\mathbf{X}^{\top}\mathbf{X}+\mathbf{X}^{\top}\mathbf{X}{\bm{D}}(t){\bm{D}}(t)^{\top}\right)\right\|_{2}$

By symmetry, we just consider the first term ${\bm{D}}^{\top}(t){\bm{D}}(t)\mathbf{X}^{\top}\mathbf{X}.$

The first inequality comes from $\|\cdot\|_{2}\leq\|\cdot\|_{F}$. The last three inequalities require that $\eta\lambda_{\max{}}({\bm{M}}^{*})<1.$

Combined with its symmetric counterpart in $\mathbf{X}^{\top}\mathbf{X}{\bm{D}}(t){\bm{D}}(t)^{\top}$, the sum of (A1) is smaller than $\frac{8}{m}$.

Similarly, we have their symmetric counterpart added and get a bound of $\frac{16c_{1}n^{2}}{m\lambda_{r}^{2}}R_{w}$.

Similarly, we can combine its symmetric counterpart and get a bound of $\frac{16c_{1}n^{2}}{m\lambda_{r}^{2}}R_{w}$.

Add them up, and we can get the bound of the first part.

The last inequality is the same one as in the term (G1) bound.

Adding the three terms together and using the induction hypothesis, we have

Property 1 holds. Therefore the proof is completed. ∎

Lemma C.3 tells us that ${\bm{D}}(t)$ decreases along GD trajectory in some fixed directions independently depending on $\mathbf{X}^{\top}\mathbf{X}$. After we have this GD trajectory, we can have the following Lemma C.4 about the dynamics of $\|{\bm{A}}(t)\|^{2}$ under the condition in Lemma C.3. It shows a sufficient condition for $\|{\bm{A}}(t)\|^{2}$ to grow.

Under the Assumption C.1 and Assumption C.2, if

and as long as there exist two number $a_{1},a_{2}$ (to be determined) and some $i\in[r]$ at time t s.t.

we have $-{\bm{D}}(t)^{\top}{\bm{F}}(t)>0$.

We use property 2 of Lemma C.3 to obtain the following expression:

Notice that all inequalities hold since $\lambda_{i}({\bm{M}}^{*}(t))<1/\eta$ for all $i$. In the first inequality, we use the $\|{\bm{e}}(t)\|$ bound in (E1), and in the second we just replace $m$ with its lower bound $\frac{112c_{1}n^{2}}{\lambda_{r}^{2}}.$ Then we use $n>(70\lambda_{r}+25\lambda_{r}^{2})/(196\kappa^{2}\min\{a_{1}-a_{1}^{2},a_{2}-a_{2}^{2}\})$ and complete the proof.

Then, we use Assumption C.1 to prove the next lemma. It tells that condition (14) is satisfied under this assumption in a time interval. That means $-{\bm{D}}(t)^{\top}{\bm{F}}(t)>0$ during this period.

Under Assumption C.1, we have condition (14) satisfied with $a_{2}=e^{(a_{1}-1)/\chi}$ in the time interval $[t_{1},t_{2})$. Here, $t_{1}$ is the iteration that $\prod_{j=0}^{t-1}(1-\eta\lambda_{1}(\widetilde{{\bm{M}}}(j)))<a_{2}$ for the first time, and $t_{2}$ is the iteration when $\prod_{j=0}^{t-1}(1-\eta\lambda_{r}(\widetilde{{\bm{M}}}(j)))<a_{1}$ for the first time.