On Gradient Descent Convergence beyond the Edge of Stability
Authors: Lei Chen
Joan Bruna
Venue: Submitted to ICLR 2023 (2023)
Keywords: gradient descent
edge of stability

Abstract: Gradient Descent (GD) is a powerful workhorse of modern machine learning thanks to its scalability and efficiency in high-dimensional spaces. Its ability to find local minimisers is only guaranteed for losses with Lipschitz gradients, where it can be seen as a `bona-fide' discretisation of an underlying gradient flow. Yet, many ML setups involving overparametrised models do not fall into this problem class, which has motivated research beyond the so-called ``Edge of Stability'' (EoS), where the step-size crosses the admissibility threshold inversely proportional to the Lipschitz constant above. Perhaps surprisingly, GD has been empirically observed to still converge regardless of local instability and oscillatory behavior.

The incipient theoretical analysis of this phenomena has mainly focused in the overparametrised regime, where the effect of choosing a large learning rate may be associated to a `Sharpness-Minimisation' implicit regularisation within the manifold of minimisers, under appropriate asymptotic limits. In contrast, in this work we directly examine the conditions for such unstable convergence, focusing on simple, yet representative, learning problems. Specifically, we characterize a local condition involving third-order derivatives that stabilizes oscillations of GD above the EoS, and leverage such property in a teacher-student setting, under population loss. Finally, focusing on Matrix Factorization, we establish a non-asymptotic `Local Implicit Bias' of GD above the EoS, whereby quasi-symmetric initializations converge to symmetric solutions --- where sharpness is minimum amongst all minimisers. 

Under review as a conference paper at ICLR 2023
ON GRADIENT DESCENT CONVERGENCE BEYOND THE EDGE OF STABILITY
Anonymous authors Paper under double-blind review
ABSTRACT
Gradient Descent (GD) is powerful workhorse of modern machine learning thanks its scalability and efficiency in high-dimensional spaces. Its ability find local minimisers is only guaranteed for losses with Lipschitz gradients, where it can be seen as 'bona-fide' discretisation of underlying gradient flow. Yet, many ML setups involving overparametrised models do not fall into this problem class, which has motivated research beyond the so-called "Edge of Stability" (EoS), where the step-size crosses the admissibility threshold inversely proportional to the Lipschitz. constant above. Perhaps surprisingly, GD has been empirically observed to still converge regardless of local instability and oscillatory behavior.
The incipient theoretical analysis of this phenomena has mainly focused in the overparametrised regime, where the effect of choosing a large learning rate may be associated 'Sharpness-Minmiine Minimisation' implicit regularisation within the manifold of minimisers, under appropriate asymptotic limits. In contrast, in this work we directly examine the conditions for such unstable convergence, focusing simple, yet representative, learning problems. Specifically, we characterize a local condition involving third-order derivatives that stabilizes oscillations GD above the EoS, and leverage such property in teacher-student setting, under population loss. Finally, focusing on Matrix Factorization, we establish non- asymptotic 'Local Implicit Bias' of GD above the EoS, whereby quasi-symmetric initializations converge to symmetric solutions - where sharpness is minimum amongst all minimisers.
INTRODUCTION
Given differentiable objective function f(), where Rd is high- digensionni parameter vector, the most basic and widely used optimization method is gradient descent (GD), defined as
= g()) g(
where y is the learning rate. For all its widespread application across many different ML. setups, basic question remains: what are the convergence guarantees (even to a local minimiser) under typical objective functions, and how they depend on the (only) hyperaparameter n)? the modern context of large-scale ML applications, an additional key question is not only understand whether or not GD converges to minimisers, but to which ones, since overparametrisation defines a whole manifold of global minimisers, all potentially enjoying drastically different generalisation performance.
The sensible regime to start the analysis is 77 - O, where GD inherits the local convergence properties of the Gradient Flow ODE via standard arguments from numerical integration. However, in the carly phase of training, a large learning rate has been observed to result in better generalization (LeCun al. 2012 Bjorck al. 2018 Jiang al. 2019 Jastrzebski al. 2021), where the extent of "large" is measured by comparing the learning rate 1j and the curvature of the loss landscape, measured with
X(0) := Amax [O)], the largest eigenvalue of the Hessian with respect to learnable parameters.
Although requires sup, X(0) < 2/n to guarantee the convergence of GD (Bottou et al. 2018) to (local) minimisers!!
| the work (Cohen et al. 2020) noticed a remarkable phenomena in the context
One replace the uniform curvature hound by suPo:1(0) A(0).

---

Under review as a conference paper at ICLR 2023
of neural network training: even in problems where X(0) is unbounded (as in NNs), for fixed 17. the curvature X(O(t)) increases along the training trajectory ( bringing X(o(t) 2 2/n (Cohen et al. 2020). After that, surprising phenomena is that X(O(0) stably hovers above 2/11 and the neural network still eventually achieves a decreasing training loss - the so-called "Edge of Stability". We would like to understand and analyse the conditions of such convergence with large learning rate under a variety models that capture such observed empirical behavior.
Recently, some works have built connections between EoS and implicit bias (Arora et al. 2022 Lyu et al. 2022 Damian et al.) 2021 2022) in the context of large, overparametrised models such as neural manifold of minimisers, and the question is to what extent a large learning rate 'favors' solutions with small curvature. In essence, these works show that under certain structural assumptions, GD is asymptotically tracking continuous sharpness- reduction flow, in the limit of small learning rates. Compared with these, we study non-asymptotic properties of GD beyond EoS, by focusing on certain learning problems (e.g., single-neuron ReLU networks and matrix factorization). In particular, we characterize range of learning rates 1) above the EoS such that GD dynamics hover around minimisers. Moreover, in the matrix factorization setup, where minimisers form manifold with varying local curvature, our results give non-asymptotic analogue of the Sharpness-Minimisatee -iiiimisttiin' arguments from Arora et al. (2022): y: Lyu et al. (0022): )
networks. In this setting, GD is expected to converge
Damian al. (2022).
The straightforward starting point for the local convergence analysis is via Taylor approximations of the loss function. However, in quadratic Taylor expansion, gradient descent diverges once X(0) > 2/n (Cohen et al. 2020), indicating that a higher order Taylor approximation is required. By considering 1-D function with local minima 0 of curvature X A(0)). we show that it is possible to stably oscillate around the minima with rj slightly above the threshold 2/X", provided its high order derivative satisfies mild conditions as in Theorem I A typical example of such functions is f()) = (x²² n))² with μ > O. Furthermore, we prove that it converges to an orbit of period 2 from a more global initialization rather than the analysis of high-order local approximation.
As it turns out, the analysis of such stable one-dimensional oscillations is sufficiently intrinsic to become useful in higher-dimensional problems. First, we leverage the analysis to a two-layer single-neuron ReLU network, where the task is to learn a teacher neuron with data on a uniform high-dimensional sphere. We show a convergence result under population loss with GD beyond EoS, where the direction of the teacher neuron can be learnt and the norms of two-layer weights stably oscillate. We then focus on matrix factorization, a canonical non-convex problem whose geometry is characterized by a manifold of minimisers having different local curvature. Our techniques allow us to establish a local, non-asymptotic implicit bias of GD beyond EoS, around certain guami-symmetric initialization, by which the large learning rate regime 'attracts' the dynamics towards symmetric minimisers precisely those where the local curvature is minimal. A further discussion is provided in Appendix M
2 RELATED WORK
Implicit regularization. Due to its theoretical closeness gradient descent with small learning rate, gradient flow a common setting to study the training behavior of neural networks. [Brrett & Dherin (2020) suggests that gradient descent is closer to gradient flow with an additional term regularizing the norm of gradicnts. Through analysing the numerical error of Euler's method, Elkaaeetz & Cohen ()2T)) provides theoretical guarantees of small gap depending on the convexity along the training trajectory. Neither of them fits in the case of our interest, because it is hard to track the parametric gap when 1) > 1/2. For instance,
quadratic function, the trajectory jumps between the two sides once > 1/A. Damian et al.) (2021) shows that SGD with label noise is implicitly subjected a regularizer penalizing sharp minimizers but the learning rate is constraint strictly below the edge of stability threshold.
Balancing effect. Du et al. (2018) proves that gradient flow automatically preserves the norms'
differences between different layers of a deep homogeneous network. (Ye & Du 2021) shows that gradient descent on matrix factorization with constant small learning rate still enjoys the auto-balande bolancanniing property. Also in matrix factorization, Wang et al. (202T) proves that gradient descent with relatively large learning rate leads solution with more balanced (perhaps not perfectly
2

---

Under review as a conference paper at ICLR 2023
balanced) solution while the initialization can be in-balanced. In similar spirit, we extend their finding to larger learning rate, with which the perfect balance may be achieved in our setting. We estimate our learning rate is strictly larger than theirs (Wang et al., 202I), where they show GD with large learning rates converges to a flat region in the interpolation manifold while the flat region w.r.t. larger learing does not exists GD is forced to wander around the flattest minima. Note that the implication of balancing effect is to get close to flatter solution in the global minimum manifold, which may help improve generalization some common arguments the community.
Edge of stability. Cohen et al. (2020) observes two-stage process in gradient descent, where the first is loss curvature grows until the sharpness touches the bound 2/7,, and the second is the curvature hovers around the bound and training loss still decreases in macro view regardless of local instability. Gilmer et al. (2021) reports similar observations in stochastic gradient descent and conducts comprehensive experiments of loss sharpness on learning rates, architecture choices and initialization. Lewkowycz et al. (2020) argues that gradient descent would "catapult" into a flatter region if loss landscape around imitialization too sharp.
Some concurrent works (Ahn et al., 20022] Ma et al. 2022; Arora ct al. [2022 Damian et al. 2022)
are also theoretically investigating the edge of stability. AAhn et al. (2022) suggests that unstable
convergence happens when the loss landscape of neural networks forms local forward-invariant set
near the minima due to some ingredients, such as tanh as the nonlinear activation. Ma et al. (2022)
empirically observes a multi-scale structure of loss landscape and, with it as assumption, shows that gradient descent with different learning rates may stay in different levels. Arooa et al. (2022)
shows the training provably enters the edge of stability with modified gradient descent modified
loss, and then its associated flow goes to flat regions. Under mild conditions, Damian et al. (2022)
proves that GD beyond EoS follows optimization trajectory subjected to sharpness constraint so that a flatter region is found.
Learning a single neuron. Yehudai & Ohad (2020) studies necessary conditions on both the
distribution and activation functions to guarantee one-layer single student neuron aligning with the teacher neuron under gradient descent, SGD and gradient flow. Vardi et al. (202T) extends the investigation into neuron with bias term. Vardi & Shamir (2021) empirically studies the training dynamics of a two-layer single neuron, focusing on its implicit bias. In this work, we present a convergence analysis of a two-layer single-neuron ReLU network trained with population loss in large learning rate beyond the edge of stability.
3 PROBLEM SETUP
We consider differentiable objective function S(O) with 0 E Rd, and the GD algorithm from Definition 1. A differentiable function f is L-gradiemt Lipschitz if
[Vv)(0;) f(2))l s L [0] - 0211. A/ 02.02.
(2)
The above definition is equivalent to saying that the spectral norm of the Hessian is bounded by L. or the local curvature each point is bounded by L. Then 7) needs to be bounded by 1/L in GD so that it is guaranteed to visit an approximate first-order stationary point (Nesterov 1998). The perturbed
GD requires 7) 1/L L to visit an approximate second-order stationary point (Jn al. 2021]. and stochastic variants share similar assumptions (Ghadimi & Lan 2013 Jin et al. (021))
However, in practice, such an assumption may be violated, or even impossible satisfy when || / is not uniformly bounded. Cohen et al. (2020) observes that, with learning rate 1) fixed, the largest eigenvalue A of the loss Hessian of neural network is below 2/1 at initialization, but grows above the threshold along training. Such phenomena is more obvious when the network is deeper or narrower. This reveals the non-smooth nature of the loss landscape of neural networks.
Furthermore, another observation from Cohen et al. (2020) is that once A > 2/7, the training loss starts to perturb sharply. This is not surprising because GD would diverge in a quadratic function with such large curvature. However, despite of local instability, the training loss still decreases in a longer range of steps, during which the local curvature stays around 2/17. A further phenomena is that, when GD is the edge of stability, if the learning rate suddenly changes smaller value
3

---

Under review as a conference paper at ICLR 2023
?e 17, then the local curvature quickly grows to 2/n, indicating the ability to 'manipulate" the local curvature by adjusting the learning rate.
Besides the analysis of GD, the local curvature itself has also received lot of attention. Due to the nature of oveo-vvereee weer-paraetterizttion in modern neural networks, the global minimizers of the objective f form manifold of solutions. There have been active directions understand the implicis bias of GD methods, namely where do they converge to in the manifold, and why some points in the manifold are more preferable than others. For the former question, it is believed that (stochastic) GD prefers flatter minima (Barrett & Dherin 2020 Smith et al., 2021 Damian et al., 2021] Ma & Ying 202T) For the latter, flatter minima brings better generalization (Hochreiter Schmidhuber 1997 et al.) 2018; Keskar et al.) 2016 Ma & Ying, 2021; Ding et al. 2022). It would be meaningful if flatter
minima could obtained via GD with a large learning rate.
More specifically, it has been shown that the eigenvalues of the hessian of deep homogeneous network could be manipulated to infinity via rescaling the weights of each layer (Elkabetz Cohen 2021). Fortunately, gradient flow preserves the difference of norms across layers along the train- ing (Du et al. 2018). As result, balanced initialization induces balanced convergence, while GD would break this balancing effect due to finite learning rate. However, recently it has been observed that GD with large learning rates enjoys a balancing effect <Wang et al. 202I), where it converges to (not perfect) balanced result despite of inbalanced initialization.
Motivated by the connections of optimization, loss landscape and generalization, we would like to understand the training behavior of gradient descent with large learning rate, from low-dimensional to representative models.
STABLE OSCILLATION ON 1-D FUNCTIONS
Definition 2.
(Period-2 stable oscillation.) Consider GD on a function f in domain St Denote the update rule of GD as F(x) for x E Q. A period-2 stable oscillation is 3 x E S such that F(F(z)) and x not a minima of f.
We initiate our analysis of the stable oscillation phenomenon in 1-D. Starting from a condition on general I-D functions, we look into several specific 1-D functions verify our arguments. Then, focusing on a function in the form of f(x) = (x²² 2))², we present the convergence analysis as foundation for the following discussions. Furthermore, shed light the multi-layer setting, propose a balancing effect to make a connection to the 1-D analysis. as shown in Appendix A.1
General 1-D functions. Consider a 1-D function f(x) with learnable parameter æ R. The parameter updates following GD with the learning rate as
((++ = (0 (x:()))
(3)
Assuming f is differentiable and all derivatives are bounded, the function value in the next step can be approximated by
= (1- o((x
-2()))
If < (x.) this approximation reveals that the function monotonically decreases for each step of GI), ignoring higher terms. Such assumption would guarantee the convergence to a global minimum in convex function. However, our interest is what happens if 77 > 2/f"(x). For instance, if i is quadratic function, the second-order derivative f" is constant. As a result, once 7) > 2/ GD diverges except when being initialized at the optimum. However, when trained with a large learning rate 7 > (z)" there is still some hope for a function to stay around local minima T, as stated in the following theorem.
Theorem 1. Consider any 1-D differentiable function f(x) around local minima J., satisfying (i)
(33 T + and (ii) f1 f" f(4) at Then, there exists e with sufficiently small |e] and E- f-f(3) (3) such that: for any point z0 between and e, there exists a learning rate n such that the update rule Fy, of GD satisfies (F,(F()))) = Fp. and
2
<
f"(T)
(*( (z)
4

---

Under review as a conference paper at ICLR 2023
The details of proof are presented in the Appendix C As stated in the Theorem we provide necessary condition that allows GD to stably oscillate around a local minima. But still we cannot tell whether or not some functions allow it with f(3) (>) = 0. For instance, a quadratic function does not satisfy this condition since f(²) = f(4) = and diverges when GD is beyond the edge of stability. For f() (z) around
turns out the sine function allows stable oscillation. Therefore, we extend the argument in Theorem to a higher order case in Lemma !|
- where f(
Lemma 1. Consider any I-D differentiable function f(x) around local minima T, satisfying that the lowest order non-zero derivative (except the fr) at is (()()) with k 2 4. Then, there exists with sufficiently small |<] such that: for any point zo between E and T-- - 6, and
1. if is odd and ((0) >> O, ((++) < then there exists vj E (Gr, 3" fl t).
2. if is even and f(b) (T) < O, then there exists 1j E (ir: p² f(n,-))))))
such that: the update rule F,j of GD satisfies (x)(F)())))) Tu.
With Lemma]) ! we can verify the sine function to allow stable oscillation as in Corollary D because its lowest order of nonzero derivative (except f") at the local minima (z) O. Meanwhile, Theorem I provides a guarantee that squared-loss on any function y provably allows stable oscillation once g satisfies some mild conditions, as stated below.
Lemma 2. Consider a I-D function g(x) and define the loss function f as f(x) = (g(x) - y)?. Assuming (i) g' is not zero when g(z) = y. (ü) g ( (ω) < 6[g"( then it satisfies the condition in Theorem oor Lemma] 7 to allow period-2 stable oscillation around T.
This setup covers generie non-linear least squares problems, including the base model g being sine, tanh, high-order monomial, exponential, logarithm, sigmoid, softplus, gaussian, etc. The proof details for these settings of g(x) are provided Corollaries I S in Appendix and E Moreover, we provide Dandd
a straightforward method to build a more complicated model from two simple base models, as follows.
Proposition 1 (Composition Rule for Stable Oscillation), Consider two I-D functions p.q. Assume both p(z),q()) qQy) at 5. p(z) satisfies the conditions of g in Lemma 2 allow stable oscillations. Then q(p(z)) allows stable oscillation around x = Σ.
Proof details of the above lemmas and proposition are presented in the Appendix D and E are going to present careful analysis on g(x) = x.2
Next we
A special 1-D function. Consider f(x) (22 11)2 with μ 0, (3) ( (vp) ((V)
2p. Note that this function is more special to us because it can be viewed as a symmetric scalar factorization problem subjected the squared loss. Later we will leverage to gain insights for asymmetric initialization, two-layer neglonn networks and matrix factorization. Before that, we would like to show where it converges to when 17 > ("(V)) as follows.
Theorem 2. For f(x) ]x= 2x)² consider GD with 17 K- where << v4.-11 1.121.
and initialized on any point 0 < Ho < VH. Then it converges to an orbit of period 2, except for a measure-zero initialization where it converges to vT. More precisely, the period-2 orbit are the solutions x = 51 E (0. vP) 82 (VF.2 2 /p) of solving 8 in
1
(5)
82
3 + 2
The details of proof are presented in the Appendix F.. As shown above, Theorem aad Theorem 2 stand in two different levels: Theorem restricts the discussion in a local view because of Taylor approximation, while Theorem 2 starts from local convergence and then generalizes into a global view. However, Theorem T builds foundation for Theorem because the latter would degenerate to the former when K is extremely close to 1.
A natural follow-up question is what implications Theorem brings, because 1-D is far from the practice of neural networks that contain multi-layer structures, nonlinearity and high dimensions. We precisely incorporate two layers and nonlinearity in Section and high dimensions in Section 6
5

---

Under review as a conference paper at ICLR 2023
5 ON A TWO-LAYER SINGLE-NEURON HOMOGENEOUS NETWORK
o We denote a two-layer single-neuron network as f(z;0) = v. oww x) where v E R, w E Rd, the of trained parameters 0 (u, w [ and the nonlinearity o is ReLU, We will keep such an order in 0 to view it as a vector. The input z E R is drawn uniformly from a unit sphere Sd-1 The parameters are trained by GD subjected L2 population loss, as
01+1 = - (oL(()).
L(O,) = Ezesa- (f(:; 0c) - v)²
We generate labels from single teacher neuron function, as y|æ = a z). Hence w is our target neuron to learn. We denote the angle between w and T as 2 0. Note that a is set as non-negative because the loss function symmetric w.r.t. the angle. Moreover, the rotational symmetry of the population data distribution results in loss landscape that only depends on w through the angle and the norm || Indeed, from the definition, we have
(wTx)
v l| |ell (sin a + (æ a) cos (x) ||0||
V&L = v ( - a + 1/ sin 2a)- v (- cos 2a + 5) ||ع||| 01
where we denote w as the normalized of proj w. Consider the Hessian
H &
O.L L,
|||| On vf vo d vw e
vic
E R(d+1)x(4+1)
(6)
Hence, in the global minima manifold where ow = w. the eigenvalues of the Hessian are λ | A2...d
2² di Ad+1 = O. Therefore, the largest eigenvalue À1 measures the imbalance (i.e..
|| v between the two layers again À [[l - Will similar to the 2-D case in 115 in
Appendix A.1 So we would like investigate where GD converges r श d/ / that
too large even for the flattest minima. Note that key difference between the current case and the previous 2-D analysis is that the current one includes a neuron vector and nonlinear ReLU unit.
From the second row of VoL, which V L, it clear that updates of ur always stay in the plane spanned by and (0) Hence, this problem can be simplified three variables (u, Way wy) with the target neuron e0 = [1, 0]. The three variables stand for
(0) = (t)) (t) == -proja (0) w = proja w(t) - (w )
We keep Wy as nonnegative because the loss I is invariant to its sign and our previous notation a 2 0 requires non-negative Wy. Then we show that wy decays
follows.
Theorem 3. In the above setting, consider a teacher neuron w
[l, 0] and set the learning
rate 17 Kd with K E (1, Initialize the student as [| (()) & (0,0.10) and
(0) w) > 0. Then, for 2 Ti + 4,, wy decays
1.35
1.1 1 1+
w < 0.1 0.030 Ti < log2.56
G ==
c.
nT22
The details of proof are presented in the Appendix ! The proof is divided into two stages, depending on whether Wy grows or not. The key is that the change of wy follows (omitting all superscripts t)
Proof sketch
Awy α -U Wy
+
Mi
(++1) = Iuy, + Aw,l- W
10x
where the second term in Awy/w, is bounded in [0, 2n In stage 1 where vw,, is relatively small, we show the growth ratio of wy is smaller than those of w. and ow,, resulting in an upper bound of number of iterations for ow,, to reach 2. so max(w,) bounded too. Although the initialization
balanced as ()))) = [0) for simplicity of proof, e - wy also bounded at the end of stage 1. From the beginning of stage 2, thanks to the relatively narrow range of K, are able to compute the

---

Under review as a conference paper at ICLR 2023
bounds of three variables (including Was twt, and wyy and they turn out to fall into basin in the parameter space after four iterations. In this basin, wy decays exponentially with a linear rate of 0.97 at most.
With the guarantee of wy decaying in the above theorem, the dynamics of the single-neuron ReLU network follow the convergence of the 2-D case in Appendix A.1 with convergence result as follows.
The single-neuron model in Theorem 3 converges to period-2 orbit where wy and (V, Wx) "KK with 7 = {(61, 51), (62, 62)). Here d1 E (0, (0,1),62 E (1,2) are the solutions in
Proposition 2.
K
82 (V 82 1 +
Remark. Actually this convergence is close the flattest minima because: if the learning rate decays to infinitesimal after sufficient oscillations, then the trajectory walks towards the flattest minima (v W2 w, 0).
To summarize, the single-neuron model goes through three phases of training dynamics, with intialization of the angle <(w, ()) as ž at most. First, the angle decreases monotonically but, due to the growth of norms, the absolute deviation "y still increases. Meanwhile, the inbalance v wx stays in a bounded level. Second, By starts to decrease and the parameters fall into a basin within four steps. Third, the basin, Wy decreases exponentially and, after Wy at a reasonable low level, the model approximately follows the dynamic of the 2-D case and the inbalance
decreases as The model converges to a period-2 orbit as in the 1-D case in Theorem
well, following Theorem 5
QUASI-SYMMETRIC MATRIX FACTORIZATION: WALKING TOWARDS FLATTEST MINIMA
Consider matrix factorization problem, parameterized by learnable weights X E Rmxp, Y E RTxp, and the target matrix is C pm,g The loss L is defined
I(X = XY - CII
Obviously (X, Y XYT = C} forms a minimum manifold. In this context, the question is to deseribe GD dynamics in terms of a 'descent' phase (i.e., reaching the manifold), followed by 'hovering' phase, where the dynamics evolve nearby the minimum manifold. Although we prove that the necessary 1-D condition holds around minimum Theorem (in Appendix A.. is more attracting to investigate GD in high dimensions.
A straightforward subsets of the "flattest" points in the manifold of minimisers are in fact given by symmetric matrices, i.e. points of the form (X, X) with XXT = C. As it turns out, the local behavior of GD beyond EoS in this symmetric submanifold of minimisers can be explicitly analysed. Indeed, Theorem 7 (in Appendix
) shows that the dynamics follows the direction of the leading eigenvector and then stably oscillates with period-2 analogous to the 1D case in Theorem 2 Note that, although (X XX = X, } forms manifold that contains infinite number of minimizers F all of them have the same sharpness due to the same leading singular values. So a natural follow-up question is to analyse minimizers with different sharpness.
The simplest setting that contains minimizers of varying-sharpnen sharpness rescale symmetric minimizers, leading to Pnami-symmetric Matrix Factorization. Given a symmetric target C = X assume that we are around the (global) minima Y1 = aXo + AY, Z, _ Xo + AZ, with a > 0 and small deviation || Y |I AZ, < <. The top singular value and vectors in SVD of X. is a Ty 1 xX,) is
Then the EoS- EoS-learning rate at (o:Xn, a
We study the convergence of GD starting from Y, = aX,- AY, Z = a Xo + AZ, with learning rate of + 3, 3 > 0. The following theorem shows that, although starting nearby sharper minima, GD still converges and stably scillate around the flattest one.
of 11 which is largest π? at
in particular, it contains the orhit (XoU;U U O(p)
7

---

Under review as a conference paper at ICLR 2023
Theorem 4. Consider the above quasi-symmetric matrix factorization with learning rate n
oj B.
Assume O < Bot V4.5 - 1 0.121. Consider a minimum (Yo axXn, Zo 1/a //X > The initialization is around the minimum, as Y1 = Yo + AY, Zi = Zo + AZ, with the deviations satisfying on 1Y * O, uj 2Z v1 # and |AY A value of Xo needs to satisfy
S e. The second largest singular
max
of a²
1 1+
of
nojia²
of ato7 (+t
} S².
(10)
Then GD would converge to period-2 orbit 71 approximately with error in
written as
(Y,, 2.) + (AY,AZ)
MAYII,IIAZ = =(E),
(11)
(x= (pa-a) Zo + (pi - 1/a ( ))
In = (Yo (p:
(i = 1,2)
(12)
where E (1, E (0, are the two solutions solving in
3o
(13)
p² ( +
Proof sketch Details of proof can be found in Appendix J.3 which shares similar spirit with The- orem The analysis consists of two phases, depending on whether Ey,, A (Y, Yo, Yo,u ". ) A (Z, Zo, u1v]) small not. In Phase I, all components of Y, Yo and Z, Zo small due to the initialization near minima, but both Ey,t and Ez,t are growing exponentially in a rate of
gofa² + 2 - 2 f - 1 > 1. In Phase II, both Cy,t and 62,t are much larger than other components, long other components of function fw,, 2) = $ (yz 1)² with learning rate n' = 1+ 3oî. Following the 2-D analysis in Theorem we have c,,z and Es,t converge the same values, which degenerates the 2-D problem to
still not growing. So the dynamics of them matches GD
1-D function. Therefore, the proof concludes with 1-D convergence analysis of f [(x)) [x2 1)2 as shown in Theorem
0
Remark. Note that both Y,- - a- "auno Zo 1/a 0 Li are residuals of Yo. Zo with the top singular value eliminated. Then, compared with Theorem we have p: corresponds to 6; + 1, which means both symmetric and quasi-symmetric cases converge to parameters with the same top singular values and wander around the flattest minima. In other words, this convergence is close to the flattest minima because: if the learning rate decays to infinitesimal after sufficient oscillations, then the trajectory walks towards the flattest minima approximately with parameter distance in O Also note that, if 1) 7 we anticipate it still escapes from the sharp minima and converges
flatter one (not necessarily the flattest). The result could be obtained by tracking GD f(r:, v) (ay 1)2 with < slightly. But the closed form can not be expressed explicitly, because it strongly depends on initialization.
NUMERICAL EXPERIMENTS
In this section, we provide numerical experiments to verify our theorems. Additional experiments on 2-D functions, MLP and MNIST can be found in Appendix
1-D functions. As discussed in the Section 4 we have (x) = (..2 - 1)2 satisfying the condition in Theorem and g(x) = 2 (z) satisfying Lemma]] [ so we estimate that both f and allow stable oscillation around the local minima. It turns out GD stably oscillates around the local minima on both functions, when 17 > (' slightly, as shown in Figure
2
Two-layer single-neuron model. As discussed in the Section 5 with a learning rate 17 E (d, 1.1d], single-neuron network f(x) v o(w" x) is able to align with the direction of the teacher neuron under population loss. We train such a model in empirical loss on 1000 data points uniformly sampled from sphere S1, shown in Figure 2 The student neuron is initialized orthogonal to the teacher
8

---

Under review as a conference paper at ICLR 2023
1.05
-1.925
1.00
1,950
0.005
-1,975
0.000
2,000
1.0
10
-1.75
-1.50
apochs
op90Ts
Figure 1: Running GD around the local minima of f(x) [ 1)² (left) and f(x) (
(right) with learning rate 7) 1.01 > - (*) 1. Stars denote the start points. It turns both functions allow stable oscillation around the local minima.
neuron. In the end of training, wy decays small value before the inbalance - w, decays sharply, which verifies our argument in Section With small Wy. this nonlinear problem degenerates to 2-D problem on e, W/. Then, the balanced property makes align with the 1-D problem where and Wc converge to a period-2 orbit. Note that the small residuals of | War and wy are due to the difference between population loss and empirical loss.
x|
1.-
10"
0.30
222
www.w
327
Figure 2: Running GD in the teacher-student setting with learning rate 17 2.2 = 1.1d, trained on 1000 points uniformly sampled from sphere S of ||>||| 1. The teacher neuron is ú [1, 0] and the student neuron is initialized (0) j0, 0.1] with (0)) = 0.1.
Symmetric and quasi-symmetric matrix factorization. As discussed the Section and Ap- pendix A.22, with mild assumptions, both symmetric and quasi-symmette cases stably wanders around the flattest minima. We train GD on matrix factorization problem with X,X, = C E 8 The learning rate is 1.02x EoS threshold. Following the setting in Section 6 for symmetric case, the training starts near X. and, for quasi-symmetric case, it starts near (aXo. /aXo) with a = 0.8, shown in Figure 3 Although starting with re-scaling, the qumet-symmetric
achieves the same top singular values Y and Z, which verifies the balancing effect of 2-D functions in Theo- rem 5 Then, the top singular values of both cases converge the period-2 orbit, supported by Theorem and
201
10-1
3 30
10
5.4
Figure 3: Symmetric and Quasi-symmetric Matrix factorization: running GD around flat (a 1) and sharp (a 0.8) minima. In both cases, their leading singular values converge to the same period-2 orbit (about 6.1 and 5.3). (Left: Training loss. Middle: Largest singular value of symmetric case. Right: Largest singular values of quasi-symmetric case.)

---

Under review as a conference paper at ICLR 2023
REFERENCES
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. arXiv preprint arXiv:2204, 01050, 2022.
Sanjeey Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arXiv preprint arXiv:2205.09775: 2022.
David Barrett and Benoit Dherin. Implicit gradient regularization. In International Conference on Learning Representations, 2020.
Nils Bjorek, Carla P Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normaliza- tion. Advances in neural information processing systems, 31, 2018.
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223 311, 2018.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representations, 2020.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34, 2021.
Alex Damian, Eshaan Nichani, and Jason D) Lee. Self-stabilization: The implicit bias of gradient descent the edge of stability. arXiv preprint arXiv:2209. 15594, 2022.
Lijun Ding, Dmitriy Drusvyatskiy, and Maryam Fazel. Flat minima generalize for low-rank matrix recovery. arXiv preprint ar arXiv:2203.03756, 2022.
Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. Advances in Neural Information Processing Systems, 31, 2018.
Omer Elkabetz and Nadav Cohen. Continuous vs. discrete optimization of deep neural networks. Advances in Neural Information Processing Systems, 34, 2021.
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4)):341-2366 2013.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instability in deep learning. arXiv preprint arXiv:2110.0433
2021.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Stanislaw Jastrzehski, Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof J Geras. Catastrophic fisher explosion: Early phase fisher matrix impacts generalization. In International Conference on Machine Learning, pp. 4772-4784. PMLR, 2021.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint 1912v2112.022778 2019.
Chi Jin, Praneeth Netrapalli, Rong Gc, Sham M Kakade, and Michael I Jordan. On nonconvex optimization for machine learning: Gradients, stochasticity, and saddle points. Journal of the ACM (JACM), 68(2):1-29, 2021.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On latch-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint 1609 04836, 2016.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, 86(11):222 (6(11):2278-2222 1998.
10

---

Under review as a conference paper at ICLR 2023
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-48. Springer, 2012.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.0221 2020.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances neural information processing systems, 31, 2018.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:2/ 2021.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normaliza- tion layers: Sharpness reduction. arXiv preprint arXiv:2206. 07085, 2022.
Chao Ma and Lexing Ying. The sobolev regularization effect of stochastic gradient descent. arXiv preprint arXiv:2 105. 13462, 2021.
Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure neural network loss functions: The effect optimization and origin. arXiv preprint arXiv:2204.11322 2022.
Yu Nesterov. Introductory lectures convex programming, 1998.
Samuel Smith, Benoit Dherin, David GT Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. arXiv preprint arXiv:2101.12176. 2021.
Gal Vardi and Ohad Shamir. Implicit regularization in relu networks with the square loss. In Conference on Learning Theory, pp. 1224-4255 PMLR, 2021.
Gal Vardi, Gilad Yehudai, and Ohad Shamir. Learning a single neuron with bias using gradient descent. Advances in Neural Information Processing Systems, 34, 2021.
Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect. arXiv preprint arXivv2110.03677, 2021.
Tian Ye and Simon S Du. Global convergence of gradient descent for asymmetric low-rank matrix factorization. Advances Neural Information Processing Systems, 34, 2021.
Gilad Yehudai and Shamir Ohad, Learning a single neuron with gradient methods. In Conference on Learning Theory, pp. 3756-3786. PMLR, 2020.

---

Under review as a conference paper at ICLR 2023
Appendices
A Additional Results
13
A.1 On a 2-D function
13
A.2 On Matrix Factorization
14
Additional Experiments
15
B. 2-D function
15
B.2 High dimension and MNIST
15
Proof of Theorem
21
D Proof of Lemmal
21
Proof of Lemmal
23
F Proof of Theorem 2
25
Proof of Theorem 5
28
Proof of Lemmal
31
Proof of Theorem 3
32
Proof of Matrix Factorization
47
U.1
Asymmetric Case: ID function at the minima
48
0.2 Symmetric Case: convergence analysis around the minima
49
J.3
Quasi-symmetric case: walk towards flattest minima
54
K Useful lemmas
58
Illustration of period-2 and period-4 orbits
59
M Discussions
59
M. Connections between theoretical results
59
M.2 Implications from low-dimension to high-dimension
Conclusions
61
12

---

Under review as a conference paper at ICLR 2023
A ADDITIONAL RESULTS
A.1 ON A 2-D FUNCTION
Similar f(x) [x²² pi)2, consider 2-D) function f(a,, y) = (ay 2.)2. Apparently, if and y initialize as the same, then (x(0), (() would always align with the 1-D case from the same initialization. Therefore, it is significant to analyze this problem under different initialization for and y. which we would call "in- balaaccaddd initialization. Mcanwhile, another giant difference is that all the global minima in 2-D case form manifold {(z, y)] = (u} while the 1-D case only has two points of global minima. It would be great if we could understand which points in the global minima manifold, or in the whole parameter space, are preferable by GD.
Note that reweighting the two parameters would manipulate the curvature to infinity as in (Elkabetz & Cohen 2021), so the inbalance strongly affects the local curvature. Viewing f(r) symmetric scalar factorization problem, we treat [I,, u) as asymmetric scalar factorization. The update rule of GD is
(1+1)) = 101) - 7(2(0))() - -1y g((t1) = y(t) m(x(0)()) - u)z(0.
(14)
Consider the Hessian as
oef 0,0,5
0,0af
y12
2xy 1 x²
11
(15)
2ry 1
When zy = μr, the eigenvalues of H are d = Hence, in the global minima manifold, the local curvature of each point is larger if its two parameters are more inbalanced. Among all these points, the smallest curvature appears to be A = 2/4 when y vII. In other words, if the learning rate 1) 2/2ju, all points in the manifold would be too sharp for GD to converge. We would like to investigate the behavior of GD in this case. It turns out the two parameters are driven perfect balance although they initialized differently, follows.
O. Note that d = (x - y)² + 2p
Theorem 5. For [x, y) 1/2 (ay- -ryy-Ho 2)² consider GD with learning rate n K -. Assume both and y are always positive during the whole process {I,, y 1220 In this process, denote a series of all points with xy > as P {a ys) [o.y, 11). Then |a: yl decays to 0 in P. any 1 < K < 1.5.
The details of proof are presented in the Appendix G Start from a point (zx()), (()) where (0)) (1) > p. Because (+++1) x (++1) = (t(() - x (1)) + 77 (xxx (y
Proof sketch
k)). it suffices to show
((+2)) - (((+2)
1(1 +11+10 n(((0))000 n))(1 + (1(x(+11))) (1+1)) n0)l < 1.
(16)
y(() x())
Since + ()((0)(()))))))))))))))))))))( - 1) > 1, the analysis of + (1x((++1) considering the coupling of (z()),(()) (z(++)),, g(t+)))
14) is divided into three cases
Remark. Actually, for a larger K > 1.5, is possible for GD to converge to an inbalanced orbit.
For instance, Figure 15 in (Wang et al. 2021 shows inbalanced orbits for f(z) 1 (zy 1)2 with K = 1.9.
Combining with the fact that the probability of GD converging to a stationary point that has sharpness beyond the edge of stability is zero (Ahn et al. 2022), Theorem (eteorem[]revese reveals and y would converge to perfect balance. Note that this balancing effect is different from that of gradient flow (Du et al. 2018). where the latter states that gradient flow preserves the difference of norms of different layers along training. As a result, in gradient flow, inbalanced initialization induces inbalanced convergence, while in our case aabalanced-iiitallian
weights converge to perfect balance. Furthermore, Theorem 5 shows an effect that the two parameters are squeezed to a single variable, which re-directs to our 1-D analysis in Theorem 2 Therefore, actually both cases converge to the same orbit when 1 < K < 1.121, as stated in Prop Numerical results are presented in Figure
Proposition 3. Following the setting in Theorem 5 Further assume < K < V4.5 155-1- 1 1.121.
Then GD converges an orbit of period 2. The orbit is formally written as {(x = y = 6,)li 1,2).
with & E (0, /P), 82 E (VP, 2/) as the solutions of solving 8 in
1
82
(
#
3
+
13

---

Under review as a conference paper at ICLR 2023
Remark. Actually this convergence is close the flattest minima because: if the learning rate decays to infinitesimal after sufficient oscillations, then the trajectory walks towards the flattest minima.
However, one thing to notice is that the inbalance at initialization needs to be bounded in Theorem because both x and y assumed stay positive along the training. More precisely, have
x((+1)y((+)) x! t0 n n))² n x)(x(()) y())
(17)
and then (((+t))((t)) when (z( y is large with 20) (()) fixed. Therefore, we provide a condition to guarantee both z; y positive as follows, with details presented in the Appendix
Lemma 3. In the setting of Theorem 5 denote the initialization as m Iæe-xol and zOYO vp
Then, during the whole process, both x and y will always stay positive, denoting (m+vm++4)
and q = (1+p)
4 (1 K)² + -
(2
1 + 3
qk²
max TT -H); 27
2(K + 1) m² m² qm² K <p.
A.2
ON MATRIX FACTORIZATION
In this section, we present two additional results of matrix factorization.
A.2.1 ASYMMETRIC CASE: ID FUNCTION AT THE MINIMA
Before looking into the theorem, we would like to clarify the definition of the loss Hessian. Inherently, we squeeze X, into a vector 0 = vec(X, Y) E mmpp which vectorizes the concatnation. As a result, we are able to represent the loss Hessian w.r.t. 0 as a matrix in (mp++p x (mp++p)) Meanwhile, the support of the loss landscape is in mpp+y Similarly, we use (AX, AY) in the same shape of (X, Y) to denote. In the following theorem, we are to show the leading eigenvector - vec AX, AY) E mpp--- of the loss Hessian. Since the cross section of the loss landscape and A forms ID function fA. would also show the oteila-oscillation condition on 1D function holds at the minima of Sa.
Theorem 6. For matrix factorization problem, assume XY C. Consider SVD of both matrices as X Emiuim-p3 oz,iug,,"! and Y = Emin(p-9) uy. where both groups of a., 'ss in descending order and both top singular values dx,1 and Jy,1 unique. Also assume 2T, Uyy11 # O. Then the leading eigenvector of the loss Hessian is A vec( u 12 Denote fa as the 1D ID function at the cross section of the Fqyl ,22 to with C1 V
loss landscape and the line following the direction of A passing vec(AX, AY). Then, at the minima
of fz. it satisfies
3
>.
(18)
The proof is provided in Appendix J.I This theorem aims generalize our 1-D analysis into higher dimension, and it turns out the 1-D condition is sastisfied around any minima for two-layer matrix factorization. In Theorem and Lemma I if such 1-D condition holds, there must exist a period-2 orbit around the minima for GD beyond EoS. However, this is not straightforward to generalize to high dimensions, because 1) directions of leading eigenvectors and (nearby) gradient are not necessarily aligned, and 2) it is more natural and practical to consider initialization in any direction around the minima instead of strictly along leading eigenvectors. Therefore, below we present convergence analysis with initialization near the minima, but in any direction instead.
A.2.2 SYMMETRIC CASE: CONVERGENCE ANALYSIS AROUND THE MINIMA
In this section, we focus on the symmetric case of matrix factorization where Y = x Accordingly, rescale the loss function L(X, X CII Denote the target as C X,XJ. and
14

---

Under review as a conference paper at ICLR 2023
assume we around the minima X; Xo + AX₁ with small ||Xi|| Consider SVD as X1 2i=1 (mi ",uiy A dannto + Xo. Then the GoS-learning-rate threshold at X = Xg is
7) = o. Therefore, we are to show the convergence of GD starting from X Xo + AXi with
learning rate 1) = oं + where O.
Theorem 7. Consider the above symmetric matrix factorization with learning rate 1j of + 8. 1.121 and no < 1. The initialization is around the minimum, as
Assume 0 < 3o9 < < =1= = x1=X+ AX,, with the deviation satisfying in AX10 & 0 and ||AX,] || S e bounded by a small value. Then GD would converge period-2 orbit "nn approximately by a small margin in O
formally written
X, 74 + AX,
| O()),
(19)
V/- (Xo + puureee Xo + 8201 (onuure]
(20)
where 61 E (0,1),82 E (-1,0) are the two solutions of solving in
1 + 3o
(21)
(6+ 1)3²
(V (8+1)
+
Proof The proof provided in Appendix 0.2) The analysis consists of two phases, depending whether 6 - (X, Xo, uivi) is small not. In Phase I, all components of X, X, small due to the initialization near minima, but only < is growing exponentially in a rate of 1 + Bof. In Phase II, 4 is much larger than other components, as long other components are still not growing. So the dynamics of et matches GD of 1-D function f(a = +((a + o1)² - af)² with learning rate r/ = of + P"(0 og. The proof concludes with 1-D convergence analysis of f(x) = (x² 11)² as shown in Theorem 2
2
Remark. Theorem 7 assumes GD starts from any point in an c-ball near the minima, except un AXiii f O. Note that this exception is with Lebesgue measure zero and, even if it is fortunately satisfied, GD still has a chance to reach u, AX,", * 0 after several steps due to some higher-order small noise. Then, this assumption could be relaxed as ||AX]] |l * 0.
B ADDITIONAL EXPERIMENTS
B. 2-D FUNCTION
As discussed in the Section A.1 on the function f(z, y) ) (ry 1)2, we estimate that |z yl decays to 0 when / E 1.5), as shown in Figure 4 Since it achieves a perfect balance, the two parameters follows convergence of the corresponding T-D function f(x) = ( - 1)2. shown in Figure 4 ry with 77 1.05 converges to period-2 orbit, stated in the 1-D discussion of Theorem while zy with 7 = 1.25 converges to a period-4 orbit, which is out of our range in the thcorem. But still it falls into the range for balance in Theorem 5
B.2 HIGH DIMENSION AND MNIST
We perform two experiments in relatively higher dimension settings. We are to show two observations that coincides with our discussions in the low dimension:
Observation 1: GD beyond EoS drives to flatter minima.
Observation 2: GD beyond EoS is in similar style with the low dimension.
B.2.1 2-LAYER HIGH-DIL HOMOGENEOUS RELU NNs WITH PLANTED TEACHER NEURONS
We conduct a synthetic experiment in the high-dimension teacher-student framework. The teacher network is in the form of
o freachee (x; 0)
ReLU(e TT),
(22)
15

---

Under review as a conference paper at ICLR 2023
Joy
10-
10
10-1
0.8
10
104
10
109
10
10
107
10:
19
spacts
esochs
<pects
1.05
1.50
12-
1.0
10"5
08-
0.6
10-13
107
(b) 1.25
Figure 4: Rumning GD on = 1)2 with learning rate 17 1.05 (top) and 1) 1.25 (bottom). When 1.05, it converges to period-2 orbit. When 17 1.25, it converges period-4 orbit. both cases, |x y decays sharply.
where x E IR. and e; the i-th vector in the standard basis of 16 The student and the loss are in forms of
f(x;0) =
Σ ReLU(wJx))
(23)
==1
L(0:0) 1 16
Σ
(f(a; - -))z))
(24)
Apparently, the global minimum manifold contains the following set M as (w.loog... ignoring any permutation)
M (( 1 Vi E [16), w, = k, e;
k; k 0}.
(25)
However, different choices of {kg 116 induce different extents of sharpness around each minima. Our aim is to show that GD with a large learning rate beyond the edge of stability drives to the flattest minima from sharper minima.
Initialization. We initialize all student neurons directionally aligned with the teachers as w; || e; but choose various k; as + 0.0625(1 - 1). Obviously, such a choice of [ flattest minima, due to the isotropy of teacher neurons. Also we add small noise to wy make the training start closely (but not exactly) from a sharp minima, as
16 not the
k; - (e: + 0.01e), e e~ N(0,0 I). w.
(26)
Data.
We uniformly sample 10000 data points from the unit sphere $15
Training. We run gradient descent with two learning rates 7 = 0.5, 7/2 = 2.6. Later we will show with experiments that the FoS threshold of learning rate around 2.5, so 1)2 beyond the edge of stability. GD with these two learning rates starts from the same initialization for 100 epochs. Then we extend another 20 epochs with learning rate decay to 0.5 from 2.6 for the learning-rate case.
Results.
All results are provided in Figure 5 Both Figure 5 two trajectories, where GD with a small learning rate stays around the sharp minima, while that with larger drives to flatter minima. Then GD stably oscillates around the flatter minima.
b) present the gap between these
16

---

Under review as a conference paper at ICLR 2023
Mcanwhile, from Figure (b), when decrease the learning rate from 2.6 to 0.5 after 100 epochs, GD converges to nearby minima which is significantly flatter, compared with that Ir-0.5.
Figure 5 (c) provides more detailed view of ||wl for all 16 neurons. All neurons with Ir=0.5 stay at
the original ratio k: But those with Ir=2.6 all converge to the same ratio around k2 |=|| = 1.21, as shown (d)).. We compute the relationship between the sharpness of global minima in M and different choices of k, as shown in Figure 5 (e, f). Actually, h²2 = 1.21 is the best choice of ] such that the minima is the flattest.
Therefore, we have shown that, in such setting of high-dimension teacher-student network, GD beyond the edge of stability drives to the flattest minima.
17

---

Under review as a conference paper at ICLR 2023
100
14
Ir-2.6
Irw0.55
10
Change
1855
Ir:05
1.0
change Irw0.5
0.6
10
epach
opocts
(a) Training loss
(b) Sharpness
neurons: 1r-26
1.24
neuron
neurons: 0r-05
predicted
3.0
1.20
1.5
1.18
0.5
20
80
100
120
12
14
epachs
No. neurons
(c) ratio of ||ws||| and v; during training
(d) final ratio of |w: and o when 1r=2.6
1.4
0.815
1.3
Sharpess
0.805
0.800
1.4
1.8
1.000 1,025 1,050 1.075 1,100 1.125 1,150 1.175
(e) sharpness for different ratio
(f) sharpness for different ratio (zoom-in)
Figure 5: Result of 2-layer 16-neuron teacher-student experiment.
B.2.2 3, 4, 5-LAYER HON-HOMOGEN MLPs MNIST
conduct an experiment on real data to show that our finding in the low-dimension setting in Theorem
possible to generalize to high-dimensionnl setting. More precisely, goals are to is is
show, when GD is beyond EoS,
1. the oscillation direction (gradient) aligns with the top eigenvector Hessian.
2. the ID function at the cross-section of oscillation direction and high-dim loss landscape satisfies the conditions in Theorem I
Network, dataset and training. We run 3, 4, 5-layer ReLU MLPs MNIST MNISTLeCun et al. (1998). The networks have 16 neurons in each layer. make easier to compute high-order derivatives, we simplify the dataset by 1) only using 2000 images from class 0 and 1, and 2) only using significant input channels where the standard deviation over the dataset is at least 1 100 which makes the network input dimension 79, We train the networks using MSE loss subjected to GD with large learning rates = 0.5, 0.4, 0.35 and a small rate 7 = 0.1 (for 3-layer). Note that the larger ones are beyond EoS.
18

---

Under review as a conference paper at ICLR 2023
Definition (line search minima). Consider function f. learning rate 17 and point domain f) We call as the line search minima of x: if
Vf(), e" argmin (0(0.)) I (x-c.
(27)
(28)
The line search minima & can interpreted as the lowest point on the ID function induced by the gradient at æ. GD is beyond EoS, & stays in the valley below the oscillation of æ.
Results. All results presented in Figure 6 and
Take the 3-layer as an example. From Figure (a, b), GD is beyond EoS during epochs 10-14 and 21-60. For these epochs, cosine similarity between the top Hessian eigenvector vn and the gradient is pretty close to 1, as shown Figure(6)c.. which verifies our goal 1.
In Figure (d), we compute ([ f(2 f(4) at line search minima along training, which is required to be positive in Theorem Io allow stable oscillation. Then it turns out most points have ([ f() > except few points, all of which are not in the EoS regime, and these few exceptional points might due approximation error to compute the fourth-order derivative since their negativity is quite small. This verifies goal 2.
Both above arguments are the same in the cases of and layers as shown Figure and
Ire0.5 search minima In-0.5
Ire0.5
2
10-1
1055
10-4
20
time opocis x
time
opochs x
(a) Training loss
(b) Sharpness
10P
[costvl, I.]]
175
3rrri p
1- [cnev. god 11
non-positive
125
100
10-1
25
25
50
75
100
125
130
175
200
0
25
50
75
100
125
150
175
200
epochs
(c) similarity of gradient and top eig-vector U1
(d) (3) f(2) (22 (()) at line search minima
Figure 6: Result of 3-layer ReLU MLPs on MNIST. Both (c) and (d) are for learning rate 0.5.
19

---

Under review as a conference paper at ICLR 2023
real
10
sharpness
Iiee seanch minima
2
16-4
sharpness
10
25
100
175
50
125
175
cpochs
(a) Training loss
(b) Sharpness
10
1 [costvl, grad Lil
1200
non-positive
-[oo//, grad
10-3
800
600
10
400
200
1000 spochs
125
175
25
1.00
125
150
175
200
epachs
(c) similarity of gradient and top eig- eig-vector
(d) 3[ f(2) f(4) at line search minima
Figure 7: Result of 4-layer ReLU MLPs on MNIST.
16a
real
15.0
sharpness
--.
line search minima
12.5
sharpowes
1022
10-1
7.5
10-2
100
175
50
125
175
epochs
(a) Training loss
(b) Sharpness
30P
1- [costvl, grad LJ
3hrup (
--JossVV, grad bil
1200
10-4
1000
800
3rrbpp
800
10-2
400
100
175
100
spochs
opochs
(c) similarity of gradient and top eig- vectvector
(d) (d) (3)
(() f(4) at line search minima
Figure 8: Result of 5-layer ReLU MLPs on MNIST.
20

---

Under review as a conference paper at ICLR 2023
C PROOF OF THEOREM 1
Theorem (Restatement of Theorem Consider any 1-D differentiable function f(z) around
local minima म, satisfying (i) ((8) (i) + 0, and (i) 3] f" f(() > at T. Then, there exists e with sufficiently small |e| and f(3) > such that: for any point xo between and x 6, there
exists a learning rate n such that the update rule Fnj of GD satisfies (z())((z)) To, and
2 ("(()
<
(*( f()
Proof. For simplicity, we assume (x) > 0. Imagine starting point To 3 -, 0. We omit (z), (" (z), f( (z) f', f", f(8), f(4), After running two steps of gradient descent, have
50 = 3 - E,
(((x)) f' - f + 2
1 (3)
1 +
+(Ee)
0
1 g0me - -O(E),
= ff (3)
6
T1 = 2o nf (zo) = ã - c - n( - ff' 1 11:02-1-000
2 - + O(()),
f((x1) = 1d (x) - 8) 2
(3) (क) - =) plt) (0) - +0)-+)++00
+ +O(E),
32 = 81 (I1),
222
30
8
3 f(3)
1
)w02) - (-c-m( -0v--fee fecc 1
2 g(3)2 1 )())2)
6
7
2
1 ((3)
(-e-n( - f"
1
p(3)
1
g())))
2
-(1)
1
(-c- n =0(6))
2
2
6
(2f" (
1
7
1
(()
1
1(3) - m 6
2
6 ((1)
1
f(1)
1
(-1
nf
(3)
1/1 --1 - )5))8((0) e e O(e)).
2
When 17 = 7 it holds
22 20
(1 ((3) 1
1
3 ((4) 0 + O
2
(29)
which would be positive if (()) small.
(1)(3)(()))) f(3) - (((( उ 3/ (3 fr (()(()) > 0 and |e| is sufficiently
2 When = FloT
then n = f c + O(22), it holds
22 5n = (3) + (-1 y(3) + p(3) ((3) 2 + - +
(30)
which is negative when |c| is sufficiently small.
Therefore, there exists learning rate 77 E (//,, fe such that 22 Zo due to the continuity of (x2 zo) with respect to r7.
The above proof can be generalized to the case of xo I e' with e' E (0, e] and the learning rate is still bounded as 17 E ( (o))
D PROOF OF LEMMA
Lemma 4 (Restatement of Lemmal 1. Consider any 1-D differentiable function (xx around local
minima T. satisfying that the lowest order non-zero derivative (except the f") at T. is ((k) with
21

---

Under review as a conference paper at ICLR 2023
k 2 4. Then, there exists with sufficiently small |c] such that: for any point zo between E and v-
and
1. if is odd and c. (z) > O, 1 (T) < O, then there exists E (Tm -
2. if k is even and f(k) (z) O, then there exists 1j E (fr: f1 )
such that: the update rule F.g of GD satisfies (,((F((20)) = Tu-
Proof. (1) If is odd, assuming f(1) > 0 for simplicity, we have
Zo e,
(x0) -S" ( + (k: 1)! (k) k-1 (k)
1
p(att)) + 0(6
31 00 - n1 (no) 6 =
(I: - 1}1 171 ( k- + p O((
ks
(x11 8" (I1 8) +
1
p(h) .x1 -2)--1)-11 z)(-1+ + 1
) (x) - x) t 0((
(k: 1)1
bt
22 - 20
(x1) -I0 -f(o ((xo)- (x1
17
11
(28" - nf" f
+
(k- - 1)1 p(h) +
(k - 1)! g(t))
-k- 1o--1)!- 1)1 (-1 + + ch-
f(h)
+ M! k!!" x1" g(kt1) p((+1) p (-1+ e 0((
sktl
ki!
2 When 77 pr it holds
22 20
2 [k+1)) + 066-1))
(31)
"1
When 17 f"- then nf = 2 + 2 )(t)
k- + 0(22h-4), then it holds
pr
I2
Fo 11
x + +(E)).
(32)
Since is odd and c. (x) > 0, (z) < 0, the above two estimations of 02-20/u have one positive and one negative exactly. Therefore, due to the continuity of x22 Io wrt 7), there exists learning rate 7) E (fm: -" =) such that 2i Zo.
The above proof can be generalized to any 2o between T and T c with the same bound for 7).
(2) If l: is even, we have
To = I -,
f((a)) -f"' 1 p(k) +
(k - 11]!
I1 xo 11f((2)) = --E-11 e
e +
(k: 1}! 11 (k))))+ k-1 + 0(Et),
1
1
fx( = 5 (x) 3) +
|(B) (x) - =0(Et)
+
[I: - 1}!
32 00
31 n (x1) FD (2(()) - (m(x)
17
"1
= (2f" nl" fr)
(k- 1)1 p(k)
(k 1)1 n1" ((t)
(k 1)! (-1 + ) x-1 +c(0))
1
22

---

Under review as a conference paper at ICLR 2023
When 7) p. it holds
D22 30
2
p(k) +0(et).
7
(k - 1)!
When 7) = f000 with c > some constant implying n," 2((1 -c ( -2 +0((2-1)))) then it holds
2 2(c- c- (k- 1)1 )
D2 - 20
1(0)
11
where we then set c 1.
Hence, the above two estimations of =2-///n have one positive and one negative with sufficiently small Je]. Therefore, due to the continuity 22 - Io, there exists learning rate 17 E (ir) yd f such that 22 -0.
The above proof be generalized any xn between T and with the same bound for v).
Corollary 1. f(z) = sin(x) allows stable oscillation around its local minima T.
Proof. Its lowest order nonzero derivative (expect ( is f(A) sin(#) -1 and the order 4 even. Then Lemma] gives the result.
E PROOF OF LEMMA 2
Lemma (Restatement of Lemma 2). Consider a I-D function g(z), and define the loss function f
as f(x) = (g(x) - y)?. Assuming (i) g' is not zero when g(ø) = y. (ii) g' g (x) < 6[g" (z)
then satisfies the condition in Theorem 7 or Lemma 7 to allow period-2 stable oscillation around
Proof. From the definition, we have
f"(x) = 2[g(x) -
(33)
(z) 2[g(x) (2)
(34)
(x) 2[g(x) - ())
(=)""""(""="))))""))""""))"""" 6g"m)))))))) 8g)))))()))))))))))))))) 8g))
(35)
Then at the global minima where g(z) we have (=)x)=2 2[g'(x))) (z)]² and (z)
If assume y trivial value for g(x), which means g' (x) + at the minima, and g is linear around the minima (implies g" f 0), then f satisfies f(3] (x) in Theorem Meanwhile,
we need (3) f(3) p(4) as Theorem | hence requires
(x)(())(()) 36g (2)g (1)))
3 (6g" (g)" 8g (z)g
(36)
(z)g") (x) >g (
(37)
The remaining case is, if gx())) + and g" = 0 at the minima, it satisfies the condition for Lemma
with l: = 4, because (33 = and (()) < 0 due to 35 37
Corollary 2. f(x) = (x. 1)2 allows stable oscillation around the local minima T
Proof. With ( 2 it has g((1) 2 (1) 2 0. All higher order derivatives of g zero. Then Lemma gives the result.
L
Corollary 3. f(x) = (sin(x) y) allows stable oscillation around the local minima & arcsin (v) with 1).
Proof. With g(x) sin(x)) has g'()) cos(F) A (z) cos(F). We have (x) is
bounded g' g3g(3) '[g"]2 cos? (7) 6sin (T) < 0. Then gives the result.
-sinn
23

---

Under review as a conference paper at ICLR 2023
Corollary 4. f(x) (tanh(x) y)² allows stable oscillation around the local minima T tanh (v) with (-1,1).
Proof. With g(x) tanh(æ), it has g())) sech²(x) * 0, and (z) -2sech (z) 4sech" 2 (x) Lanh" (x)
) is bounded as
g'y 6|/"]² 2sech 6 + 4sech* tanh² 24sech tanh? = 2sech 20sech lanl? < 0.
Then Lemma 2 gives the result.
Corollary 5. f(x) = (o - y)² (with E Z, 2 2) allows stable oscillation around the local minima I = y except y = O.
Proof. With g(x) x" it has g' (z) ar- ,g" a(a 1)x0-2 (g( ala - 1)(a
2)x"-3 Then we have g' (0) 6[g"|² (aaaa)(--1 a 2)x < O. Then Lemma 2 gives the result.
Corollary 6. [(x (exp(x) y²² allows stable oscillation around the local minima for > O.
log
Proof. With gax) exp z. it has g((x) (z) (x) exp(x). Then we have g' g3) 6' < 0. Then Lemma gives the result.
Corollary 7. f(x) (log(z) y)² allows stable oscillation around the local minima exp y.
Proof. With ( logr, it has g' (x) +. g" = 00 2 Then we have g'g(3) - 6/g" Then Lemma gives the result.
Corollary 8. (xx (1 a) y)² allows stable oscillation around the local minima sigmoid for 1).
Proof. With g(x) T+rxpt- z) it has (x) expl -2)
[expl 1)+ . exp(=)(exp(a>-1)
(exp(e)+1] g (z))
Then we have g'g 61g"12 x -1
(Ixp(=)+11)
1)² < 0. Then Lemma 2 gives the result.
Proposition 4 (Restatement of Prop [p. Consider two functions f. g. Assume both f(x), g(y) at Σ,y = (() satisfies the conditions in Lemma 2 to allow stable oscillations. Then (z(()))
allows stable oscillation around x Z.
Proof. Denote F(x) 2 )(f())). Then we have
F(x) (2)
F"(x) = -"
(x) = - (0UU")))
Thus, omitting all variables F and f(i) in the derivatives, it holds
F((z) (F) g' (3)((())))))) 3g" f' + + +
< 5
where the inequality is due to all conditions in Lemma 2 So the only problem is whether we can achieve g'g" f" >> O. The good news is that, even if holds g g" f" < 0, we can still find functions to re-represent (((()) as (x)))) such that "' j" < 0 and all other conditions in Lemma are satisfied by ğ,
24

---

Under review as a conference paper at ICLR 2023
For g' f" < O, construct ((y) A g((-)),(()) - -x()). In this sense, it holds (f (z) g(f())) satisfy the conditions in Lemma because
It is easy to verify that both 9. f
g'(v) - - = - (F)( (2(v) = (-v) = =,((())), g00 = --g 1 = - (()))
(1) (z))
Then, it has (v) fz)x = -g'g"" 0 at y = -f()), = ã. Therefore, we have
F(())(()))) - 6[)" -t) < 0 and Lemma 2 gives the result.
F PROOF OF THEOREM 2
Theorem (Restatement of Theorem 2 For f(x) = -2²² n2),, consider GD with vj = K 1
where 1 < K < V4.5 R 1.121, and initialized on any point O < xo < Vu. Then converges to orbit of period 2, except for a measure-zer initialization where it converges to yi. More precisely, the period-2 orbit are the solutions x = 81 E (0, v = 82 E (VT1, of solving in
1
(38)
82
3/ 3
V
Proof. Assume the 2-period orbit is (Zo, I ). which means
31 Fg 10-11 1) f (3o) To + 7 (1r 20)
io = - f(1) = 5i + n (μ - 27)21.
First, show the existence and uniqueness of such an orbit when K E (1, via solving high-order equation, some roots of which can be eliminated. Then, we conduct an analysis of global convergence by defining special interval J. GD starting from any point following our assumption will enter in some steps, and any point in I will back to this interval after two steps of iteration. Finally, any point in I will converge to the orbit T1).
Before diving into the proof, we briefly show it always holds x > 0 under our assumption. If Tt-1 > O and 24 S 0, the GD rule reveals 7(p x²-1) S -1 which implies 2² 2 + 1 However,
")"
the maximum of + n(u )²)x on x (0, /u + (/) is achieved when x² = 1/2 (1 + !)) so the
maximum value is Vi(u= + 1 m))) < 1.4, 1 2 < vu+
J. n. As result, it always
holds r > O.
Part I. Existence and uniqueness of T1).
In this part, we simply denote both To. ā as zo. This means xo in all formulas in this part can be interpreted To and Tt. Then the GD update rule tells, for the orbit in two steps,
In = xn + n11 =o)
x zu 2 + n74 xf)=1,
which means
0 = nu 2z)20 + 17 [2- (zo + n(p - -0))==)0)))0 (2o + +nnu =2)=0)
0 = | 2/ + (n- (ro + πuu 2 =2)=0)²) )² (1 + n(u =1²))=
Denote == 1 7(1 x;1), it is equivalent
(2 0 = μ 2/ + nμ =² (2+
-22 + f =8))
==++1 + 1) (-zğ(= 1 |
3 2
50
25

---

Under review as a conference paper at ICLR 2023
2 #1++1= 0, it means 21 zo which is however out of the range of our discussion the
domain. So we require -zz(2 ) + p natural to require
3
1 o
= O. To ensure the existence of solutions 2, is
3
11
0 2 0.
2
Then, the solutions are
2 2
20
4
However, 2 = S > +
22
can be ruled out. If holds, 774 2z)) = - < - which means
1 2n'
Since we restrict ny E (1, (111,, it tells x > 101 + 1.212
contradicting with
1 2 T
Hence, 2 2
+
P is the only reasonable solution, which is saying
23
4
3
2
0 4
Given a certain 1), the above expression is a third-order equation of zÿ to solve. Apparently x3 one trivial solution, since for any learning rate, the gradient descent stays at the global minimum. Then the two other solutions are exactly the orbit (To, F1), if the equation does have three different roots. This also guarantees the uniqueness of such an orbit.
Assuming z f Je, the above expression can be reformulated as
7
(39)
One necessary condition for existence is 2 3 4 2 Note that here zo can be both Tp, Ti, one of which
larger than vp. For simplicity, we assume To < / < T1. Since 17 from 39 is increasing with T when μ zn. let a 3 d and achieve the upper bound as
39)
3 my 2
(40)
which is satisfied by our assumption < mp < 4.5 R 1.121. Therefore, we have shown the existence and uniqueness of period-2 orbit.
Part II. Global convergence T1).
The proof structure is as follows:
1. There exists special interval = [I., /T) such that any point in I will back this interval surely after two steps of gradient descent. And Zo E I.
2. Initialized from any point in 7, the gradient descent process will converge to Zo (every two steps of GD).
3. Initialized from any point between and vp. the gradient descent process will fall into in some steps.
(II.1) Consider function Fx(x)) + 111 x²)rr performing one step gradient descent. Since (x) nu - 3rp=", we have (x) > for x²
( ( + and (x(x) < otherwise.
is obvious that the threshold has x² = (μ ( < μ. In the other words, for any point on the right of æs, GD returns point in decreasing manner.
26

---

Under review as a conference paper at ICLR 2023
To prove anything further, would like to restrict By 2 2s, which is
3
1
z ス 3
1 2
Solving this inequality tells
+ V2
a 2 7 μ.
(41)
Consequently, by applying Eq 39 we have
np S V4.5 - 1 = 1.121.
(42)
With the above discussion of x,, are able to define the special internal / == [Ta: vT.. From the definition of Fy. consider a function representing two steps of gradient descent 1x = (z)) From previous discussion, we know F7j (mo) Tu- What about F2 (2xx)?
turns out (1ra) > Zgg we have Fy = N. + TH nor2) = ix (1 + mpp) and, furthermore, F2 (xx)) = F, (1+) = Is 3 (1+ - (1 + mp (1 + ))))) Then we get (2 > T,
because
3 [1+ nu)
2
1+ 17μ1 27 (1 + nujs) >
qu E (1, V 1).
(43)
Combining the following facts, i) (2) - x is continous wrt x, ii) F2 (7)))) zo 0 is the only zero point x E [x,, Tol, we can conclude that
I.s > 0, and iii)
(x) > x, Ya E [x,, To).
(44)
Meanwhile, we can prove (x) < x for any x (To, vp). Since F7 (n) p 0 and F2 (Zo)
To are the only two zero cases, we only need to show E à E (To, vTJ, such that F < at x² = μ. which is
f. We compute the derivative of (x) F"(F(z))
d F (x) x -1 =-1-
de
-1+ + [F (āo)) To there exists point å (Tu. /p) that is very close to v/ such that (2) < f.
-1 (1 -1-120 2nj)² > 0. Then combining it with
Hence, we can conclude that
(x) < x, Vr E (āo, VP).
(45)
Since (-) is decreasing on [I,, 0o) and (x) > I, for x E [I v/ll. it is fair to say Fx(x)
is increasing on E [r,,, vpl. Hence, we have (x(x) S Fz(0o) Tu. Vo E [ol. And (x) 2 (ān) To. Var(To, vTa)
Combining the above results, we have
(x) E (x, 2ol. Vx E [r,, Ro):
(46)
(x())) E ( Vx E (Zo, VT).
(47)
(II.2) A consequence of Exp 46 47 is that any point in I will converge to To with even steps of gradient descent. For simplicity, we provide the proof for x E (o).
Denote ag E [I,, To) and a, = F? (a, ()--1).2 n 2 1. The series {a,)20 satisfies
in a1 an ag.
(48)
Since the series is bounded and strictly increasing, it is converging. Assume it is converging to If a < To. then
(a) a F (an)
27

---

Under review as a conference paper at ICLR 2023
Since ()) is continuous, so 3 0, such that, when |x al < 6, we have
F2)()) - < (a((a) a.
(49)
Since is the limit, 3 N 0, such that, when N, 0 < a F (an) <.. So, combining with Exp(49 we have
IF Fa) F7 a.
But LHS = (a) an++ > (a) so we reach contradiction.
Hence, we have {a,} converges to žо.
(II.3) Obviously, any initialization in (0, vP) will have gradient descent run into (i) the interval I, or (ii) the interval on the right of //,, i..,, (v/, oc). The first case is exactly our target.
Now consider the second case. From the definition of 2s in part III,, we know (x, max (z). So it is fair to say this case is In (Vp. Fy(2,)) Then the next step will go into the interval I, because
(2,) 2 (F,((()))) Fxxx)) > Is.
where the first inequality is from the decreasing property (.)()) and the second inequality is due to F2 (x) > x on To)-
G PROOF OF THEOREM 5
Theorem 10 (Restatement of Theorem 5 For f(x, y) (xy (zy-w1) w)² consider GD with learning rate = K 1. Assume both x and are always positive during the whole process { Yi 120- In i.'
this process, denote series of all points with xy μ P (()i,,))aa; p.). Then Ix yl decays to in P. for any 1 < 1.5.
Proof. Consider the current step (Ie, Mc) with 24y4 > μ. After two steps of gradient descent, we have
Tt+1 = + 7(p - 1zHi)81
(50)
Yt+1 y: + 7/p xzy()
(51)
J++2 = Jt+1 + niu - Fe+146+1)
(52)
Vt+2 = Vt++ + n7 11+111++))4-11
(53)
with which have the difference evolve as
Vt+1 21+1 (yr 22) (1 17 (μ - x(31e)
(54)
V T142 = (yt+1 - T141) (1 - 17 (μ re+141+1)).
(55)
Meanwhile, we have
2141//+1 = Zili (μ 2(341 (x, + + n2 (μ 21y11 xe
= Z.ye (1 + 77 (1 xx33)) + (μ x.3() (xt - y1)
(56)
Note that the second term in Eq(56) vanishes when x and y are balanced. When they are not balanced,
if FeVt > it holds r2+199+11 < ZzZe (1 + = Incorporating this inequality into Eq Eq
55 and assuming 3t 24 0. it holds
M4+2 It+2 < (y: xx) (1 n ( x.y:)) (1 (1- - 17 (1 zeye (1 + ny (u =eae))))
(57)
To show that Jxr yl is decaying as in the theorem, we are to show
1. Vt+2 Jt+2 3/
2. 1/+2 Fe+2 -(y 2+t)
28

---

Under review as a conference paper at ICLR 2023
Note that, although ziti > M, it is not sure to have 22428642 μ. However, for any 2,ys μ and < 2, we have
Fi+1 Vi+il
=11 17 (1+ --yy:) 1,
(58)
24 3/l
which is saying Jz yl decays until it reaches ry μ. So it is enough to prove the above two inequalities, whether or not T4420t+2 > J..
Part I. To show y1+2 2442 1/1 It
Since we wish to have yl+2 - 2142 < Yt - Ie, it is sufficient to require
(1 (1-18 (1-7 (w- z181e) (1 17 (A= maye (1 + 17 (# - -xx)))))))) < 1.
(59)
Since we assume T( Ut+1
> 0, Eq 50
tells 17 (p xz.l) > min =e,
y #) which is
equivalent to - (1 T(9tt)
< minl
V
(I.1) If 7(p 11-181+1)
Then we have 71 Ft- 10+1) $ 2.2 As result,
(1 - 11 (n -x(8))) (1 - n (u 1++14++1)) < ( 1 min
Vt+2 Ft+2
2t St }
x 2
(60)
y1
Y1 xz
1
yt
min
(61)
2
y/ 1t
(I.2) If 71 13+1))+1) < and 1++14++1 2n = {1(
The second condition reveals
Vi+2 Fi+2
= 1 - 17 (μ x++10++1) S 17
-+
V/+1 T1+1
3
2
7 K.
(62)
3 3
The first condition is equivalent to 1e+14t+1 > we have
1 2n)
Since the second term in 5q
56 is negative,
1
Telh (1+ 1 (p x23r1)2 > 1 21j
(63)
with which would like to find an upper bound of 3ttht-
Denoting Telit. consider a function q(b) 611+ (m- b)) Obviously q(x) p. Its derivative is qb(b) = (1 + mp 17)) (1 + 7リ 3rb) < 0 on the domain of our interest. If we can on proper domain, then it is fair to
show an (negative) upper bound for the derivative as g'(b)
say that, from Exp 63). Telle < f + 2n Then we have
1
Ft+1
7u Iz9e) < - 17
11
1
3
(64)
2m
2
Then, combining Exp(64
62) it tells
1t+ Ft+2
< 2 - K.
(65)
y1
The remaining is to show (bb < -1 on a proper domain. We have (b( = (1 + nリ - (m6)).
which is equal to 2np < -1 when b μ. Meanwhile, the derivative of q'(b) is (b) -2m (1 + mр 2mb)) = =-2y(1 + np 16), which negative when 6 < μ + 4. As a result, always holds (b(b) -1 when b / + 1
29

---

Under review as a conference paper at ICLR 2023
(I.3) If 2441/++1 x² Denoting again Zt,i,, the above inequality in is saying, with > f4,
p(b) = (1 - ( b)) (1 17
7 (N b1 + n bb))² < 1.
(66)
After expanding p(-), we have
p(b) - = 1y (n - b) (-2+ + ( (n- - + 2nyb 17 (1r - b) - b7²b (μx b)²)
Apparently p(u) = 1. So it is necessary to investigate whether pb( < 0 on b > μc,
(b(b) = 2np + (μ b) (17² (1 + 17 17 6)) 3b) + n (pr b))
Since 176 > and > it is enough to require
(1 + (11 b)) (- + + 1/b (1+ - b) >
(1 + (eln(uub)(((-))) ())( nb(u b) + 2b(1 + n - b)) 0.
suffices to show
"i - b) + 2(1 + , - b1) = 2 3n(1 - b) > 0.
(67)
1 + +). holds
Since Ze+ 18+++ 2 x
b nn - b))²
1
+ 1
17. )
3
3 1 2 + 37( b)
->> 0,
where the last inequality holds because: i 2 3 ( + 1)) then 1 1+7( ([(4u-b), b) S 2mp < 0, which contradicts with the assumption that both 1e+ M441 are positive. As a result, the above argument gives
Dt+2 T++2
<(b() 2(K 1)(b 1n)
(68)
V1 24
Part II. To show 36+2 I142 -(y; Izt)
Since Zeltt > 14, we have 1 - 7144 - 2z3/:) 1. Combining with 77 Ie3it) 2, it holds
V1+1 I++1 1 7pp xx9e) (1,
1/1 Tt
So the remaining is to have M1+2-2142 Therefore, have
-0.5.
it 11 2 1-nj K.
Jt+1-Zest
V1+2
T142 > -1 -1 + (3-2K),
(69)
y1
as required.
Part III. To show Vt 21, converges to
we have for points in P, - is monotone strictly decreasing sequence lower bounded by O. Hence it is convergent. Actually it converges to 0. If not, assuming it converges to e > O, the next point will have the difference as é < as well as all following points. Hence, the contradiction gives the convergence 0.
From Exp 61 65
68 69
30

---

Under review as a conference paper at ICLR 2023
H PROOF OF LEMMA 3
Lemma 6 (Restatement of Lemma 3 In the setting of Theorem 5 denote the initialization as
990 zol and Fofoo > p.. Then, during the whole process, both x and y will always stay positive,
VP
denoting (m+vm244)) and (1 p)?. if
max TTT030 - 11), 27
4 (1 K)² +
2 K2
K
qK² m²
qm² - [k] <pp
3
2(K
Proof. Considering Ttyh > nr one step of gradient descent returns
1++1 I2 + n( Fz3t)
M++1 y/ + n1 Fz Ja
To have both r++1 > 0, M > 0, suffices to have
yt St
n( p)
min
(70)
yr
This inequality will be the main target need to resolve in this proof.
First, we are to show
Yuu
To
min
20 yo
(m
vm2 +4)
With the difference fixed as m = (yo ro)/ VP. assuming yo > 2o. we have m/yo = (1 zo/ /o///o///// vp. if 2030 increases, both zo and yo increase then m/ //o decreases, which means 20/30 yo
increases. As a result, we have
yo
zo
Mo
2o
4
min
min
Z0
yo
20
yo
m + Vmz + 4)
Therefore, at initialization, to have positive 21 and yı, it is enough to require
T) - 14) (m + Vm + +4)
r.
From Theorem 5 is guaranteed that [x, yt 1000 yo with / until it reaches wxys >> μe, with xz.
which r is still a good lower bound for min(fr. / x///e) So what remains to show is it satisfies ((2 - μ)) for the next first time 2tZt > p. If this holds, can always iteratively show, for any Tzlee > along gradient descent,
1/ F4
( 14) <r min
Us
Note that itself is independent 2,34 and all the history, so it is ideal to compute uniform upper bound of ( (4) with any pair of (.:-1. satisfying 18-- μ.. Actually it is possible, since we have [.,-1 V/-1/ bounded as Theorem 5
Assume 3, > and it satisfies the condition of n7 μ)) < r and I - Vil < |20 Vol As in 54), we have
Tit1
Vi+1
(u I.yIi E (1,
(71)
3; y/
Hence, it suffices get the maximum value of 9(2), with z E (0, 14) as
o(2) z (1 + n7 =)))² niu 21 r) 2 yo)
(72)
31

---

Under review as a conference paper at ICLR 2023
which is from (56). Denote 2 argmax g(2). Obviously 2 < ( + 1,) _ 2b, because the first term of g(x) achieves maximum at = (14 + ,) and the second term is in a decreasing manner with 2. Then let's take the derivative of g(2) as
d(()) - (1 + 7lu =)) (1 + mp 31/²) 17(1 + ()) - y0)²
1711 (x0 1/0
= (1 + n1 2)) 1 m/4 -372
+ ni 2)
where the first term is always positive, so we have
7711 + +)) - -y0) 0,
1 nu 3nz
(73)
+ n(p 2)
which means
(1 + (x0 - yo)²
1 + m
(74)
3n
+ n74
1711 + r) - yn0
+ n/
(75)
3m
+ 71 [ + h
1
3(1+
2 (g0 yo)²
(76)
11
2
(77)
where the inequality is from E < 3 + 1.). As result, it is safe to say
g(2) S (1 =)))
+ +uuu 2)
(0(1+110 + 2 yo)
(78)
5=3b
==2,
(2
27 (1 + nT)²3
+ 7 + r)2
T/ 3n +
(xo 300 (xo - yo)
(79)
nu +
with which we are able to compute max 7(((z)) 11), which is exactly the final result.
I PROOF OF THEOREM 3
Theorem 11 (Restatement of Theorem 3). the above setting, consider a teacher neuron = [1,
and set the learning rate 1 = Kd with K 1111. Initialize the student as [l: = v(0) _ E
(0, 0.10) and (w) w) > 0. Then, for 2 Ti 4. w decays as
1.35
(1) 3
0.1 (1 0.030K)(-7-4
S log2. 5G
B
n
32

---

Under review as a conference paper at ICLR 2023
Proof. We restate the update rules
(c = (lt+1) v(t)) !w!
(-v10) + - (0
wy
wy
arctan
w [0)
w
10
Kw! (1)
(-v1))
+
arclan
Wy Wx
+ K [wy 12
v0(0) w(())
- (2(11))
(80)
u(t)
(8)
(w = w (I+1) w = Ke(1)
(-v(t) w0l
+ 1)
1
arctan
wy
we
(8)
to
(81)
Au = - K (((llt)22 π [w(t
(82)
((+1) 10! + Aw wy
(83)
For simplicity, we will omit all superscripts of time t unless clarification necessary. From 83
the target is to show w, decaying with linear rate, it suffices to bound the factor term in 82 considerable margin) as
(by
-2 < K
-2²
viy
0.
(84)
||li|
The technical part is the second inequality of
84
If wx > O. it is equivalent to W.
Wr Wy
ग (w + m
where the RHS is smaller than equal 2 Hence, 1 is special threshold with which will frequently compare owr. Another important variable to control is - wa that reveals how the two ((+1)
layers are balanced. If it is too large, for the iteration ((+11) ((+1) case.
may explode as shown in the 2-D
The main idea of our proof is that
Stage with vw,, S wywy, in this stage, my grows but it grows smaller rate than that x[|æa]= of and wx. Therefore, since we have an upper bound for ow to stay in this stage, we are able to compute the upper bound of #iterations to finish this stage, which is T₁ in the theorem. At the end of this stage, both of w2 and Wy are bounded under our assumption initialization.
Stage with vws, wwwy. : in this stage, Wh decreases. Since our range large learning x||
rate is relatively narrow (1 K < 1.1), are able to compute bounds of tWaa w2 and Wy. After cight iterations, falls into (and stays in) a bounded basin of these three terms, in which Wy decays at least in linear rate.
Stage 1.
We are to show that, in the last iteration of this stage, there are three facts: 1) oulz S 2 2) W. E [-0.017, 0.7. and 3) wy S 0.44.
At initialization, we assume (()) = ||] Denote ag = arctan (0) /wi (0) E T0.r So for next iteration we have
(1 (0))) 1 K (((0100)2 sin on ))
(85)
w((') (0) cos og K
11 cos
cos og sin 00
(86)
π
33

---

Under review as a conference paper at ICLR 2023
Apparently wy (1)
increases with oo increasing. And
Oap w/! w:0))
sin ox + K (1o(0))2s sin on
sin dg + cos Olp
(1110)) sin ag
sin sin 0g + -in 000
sin 0g + K
π
Since in stage 1 it holds Amy 0 which means -00(0))2 Oag W
sin on
in 885
So it follows
x
(1) < 0. Combining the above arguments, we have
K
w 2 w
v(())
a==
(
S w (1)
m==
1
K
K(o10)))
v10)
K 010) =(11 1+ π
π
(1)
2 2+ 2K
wig (1)
< 2.7.
K
Regarding J it has (0) 2 wy at initialization due to r()) = || From 80
81 82
we have
vAv = Ay + wy Aw So it holds vAv 2 yAy. Meanwhile, Aw = K vuy + ||| E
10, 11. From Lemma given v(() 2 my and Aw, E [0, 1] for any t this stage, always holds
7 v(++1)= wy ((+1)
Therefore, is fair to say
(1) (1) 1
wz
2
(w)) 9 2)2
2.7
Additionally, to bound the term vnyy/ ||@ in Awy. we would like to show it always has umy I || At initialization, naturally holds. Then, for the every next iteration, given it holds in the last iteration, we have
(u + Av) (my + Awyy - Iw + Awa)² + (my + Aw,y)
( Wz Aw, wy Awy (my + Amry) - [[ + Awx)) 2 + (wy + Amyy)
Wy
+ Awy + (wyAuu + (Awe) ) Wy - [[ww + Awz)² 2 + (w, + Aw,)
-wy + vAw + Wo Awz
$ Uwyy + AAw + wy Au,
Wy
(w² + wy + 2w,Au, + (2
v
S SdAwy + wy Amy Wy 2n, - (Awz)²
=dAw, (1 Wyyz (Aw,)²
VdW,, (Aw,)²
where the first equality uses vAr Ar, + WyD the first inequality uses the proven wy
and v 2 Awy. the second inequality uses the assumption ow,, < ||]| Now we are to show vAw, (2 < We have
w
2013) 2
Adw, - (2wa)² < Kv²
22²
π
|!lm|
2x
(()
1
arctan
(t)
(K) w
wy
w
34

---

Under review as a conference paper at ICLR 2023
Since have proven wy /we 2 2.7, it is easy to check that
(1) (1)
1
(1
1
(0)
2T
2
π
As result, vAw, - m at time 1. Furthermore, by checking each term, vAmyy - (Aw,) decreases with wy/ / decreasing. We will soon show that wyy wr itself decreases, by showing the growth ratio of wy is larger than that of Wy.
Our target lower bound of the growth ratio of wy is that
Aw
(87)
which is larger than the growth ratio of wy bounded by due to vAwr Ko/w,, w2, 2 1. Assuming Kv/w,, 2 1 for the current step, we need to show A Ko(t+1) also holds for the next step. Let's denote
|w So suffices to show
(ooetl//ul /ww (++1)
2 1
(
(2) (1) W
A K
(-o +1) π
aarctan
wy w!
my |010
(88)
Then
(u Av)
1 K
(w2 + Aww)) 2 Aw
W1
Au K
K
(v
Wx
KA)
+ o(K
R
1
)A
(89)
K
If KKA < 1, since K >> 1 and A > O, we have
89 as positive, which is what we need. If > 1,
then
89
2 (v * (1 - K2) v(K
}
1
K
1 )
((-- + A)U w. (K
where the first inequality is due A < K and the assumption of K Kv(t) (o(t)//u, /w/ 2 1. Then it suffices
(t)
to show (-K + A)u + w,, 2 (-K + K )u + we 2 O. Note that - K + 1/K E (-0.2, 0] when
K E 11, is easy to verify that p(1) < 5w Then, for the next step, we need to show ((++1) 5m ((+1) given p(t) & (w+1+)) To prove this, are to bound Wx.
(++1) w ((+1) (1 - A)(u - w) + K
||wl"
w
w) + Kwy
S 0.4(v w) + Kw,,
S
0.4(v
(90)
n ||w|
2
where the first inequality is due to, when wy/ws, I 2.7,
(t) wy
(1)
(t)
1
wy
A K
K
arctan
n
Wx
(c)
1
wy
2 K
arctan
0.6.
(t)
π
wr
We will later show that (((+1) (c+11 > -0.1 we (0). Combining this with 90p. is safe to
say
Kwy < 0.4 x 4u
K x 5w
(1+1) - (0+1) < <0.4(v w)
< 4w,
35

---

Under review as a conference paper at ICLR 2023
where the second inequality is due to I 5w and v 2 wy. Since w(++1) 1) 2 w((t) (due to A > 0) in this stage, we have y)(++) I b
Combining the above discussion, we have prove (8 Obviously, when wy/wx I 2.7, RHS of (87 is at least 0.55, larger than 1T, which is the upper bound of the Awy/wy As a result, Wyy Wx keeps decreasing.
The next step is to show the growing ratio of wwe, is much larger than that of wy. From 81 holds
82
(1+1))
wl(+1) (v + (e) + Aws) vw KA(e² + w + 22A uw (1 A)2,
where the first inequality is due to Auy 2 0. It follows (((+1) wx (t-1) /v(0 2 1.6² = 2.56.
So far, have shown the following facts: under the defined initialization time 0, starting from time 1, we have
1. ww, 1 1/2TT.
2. Awz/Wz + 2 1.55.
3. AWy/Wi + $ 1 + K/n.
4. My//w S 2.7 and keeps decreasing.
((+1) / 2 2.56. 5. ((++1) Wx
6. v 2 Wy.
7. vAwy (
Now we are to use the above facts bound unz, wy and e wy to the end of stage
For uw,, in previous discussion, have shown that var, < 2 A.tuallly there is another special value
W.2 0.104 when w.y/ W,,, 2.7. 7 (w2 w2)
(91)
This value is slightly larger than 1/4T. Hence, we would like to split the analysis into three parts: in the first step of stage 2.
1. ww 2 1
2
2. 1 v 0i < 2
3. 0w2 < 4n
Note that, although we are discussing the stage 1 in this section, investigating the lower bound of the first step stage 2 helps calculate the number of iterations in stage 1, Furthermore, it helps bound several variables in stage 1.
Case (I). If ww,
first step of stage 2:
2x
(0 (1) (wn)"
2 1/2.7 and (1+1) (++1) 100000)0000 2 2.56, the number of iterations
Since we have prove
for awa to reach 1/2m is at most
T
log2.56
2
(92)
(w (2) /2.7 2
36

---

Under review as a conference paper at ICLR 2023
Meanwhile, starting from time 1, the growth ratio of Wy is
(wy + Aww,)/w, S K(- + + 11+111/- / (w ) (93)
where the first inequality is due to ww I ||ol the second is due K 1 and the third is from v 2 Wy. Therefore, combining with (92), we can bound wy in the end of stage 1 as
1082.38
2r.
"'y S (1 + 1.1/2 (1)
(wy 12/22
(94)
Since it initializes as
| - 0.1, we have wy !1) S 0.1(1 + 1.1/7) 0.135. Then, it can be
verified that, when wy (1)
E (0, 0.135.. it holds
W S 0.44.
(95)
The next to bound - Wg. Combining the update rules of o and w. in 80
81
we have
A v - = (v(t+1) w((+!)) (v(() - (
arctan( w/y/ //2)
WewWy Twil"
2
= K(U - wx)
Dwlz - 1
+ K
(
y
(-1"
v
+
vwy
π [
(96)
Note that
arctan (yy ww))
10 T
S - (retan(wy//un
S wiCr
(97)
where the left is due to wwa > and the right is from Awy 2 0. When wy/we S 2.7, the RHS
follows -1 11+ awetan(w,,/vl /ww] 0.6. So combining both sides tells
x
arctan(w,, /wx)
w²,
1+ K
UWC - 1
(98)
Since Awy 2 0, we have 0 < K (- x|
vwyz)
" s Wy Tw
Note that initialization
wz S v(0) Then it easy to verify that
K
K (0(0))
0.01 S -0.1(v ()) ( S (1
S 0.082.
(99)
π
2
Because the coefficient the positive side in 98 is larger than 0.4 > bound the 1 W as
0.1, it is appropriate upper
T (n
(w
1 Wr S max
0.082, 0.082 0.4
+
0.4
1=1
K
< max
0.082, 0.082 0.4 +
Σ
0.4
w
2(t-1)
π
1 1
1.55
2.7
I+K/n
T
1.1
1.1
4.4
< max 0.082, 0.080 - 0.4"
+
0.4'
2(t-1)
+
1
1.55
2.7
1+1.1/m
where the second inequality is from the different growth ratios of we and Wy. Note that here we take all T > and pick the largest value RHS to bound wy- turns out
W. S 0.17.
(100)
37

---

Under review as a conference paper at ICLR 2023
Furthermore, lower bound
W2z, since obviously | Wx S 0.17, it follows
v w 2
00.1
We Imax -0.017.
(101)
Case (II). 4n S vw
<
2x
in first step of stage 2:
Similar to the discussion in Case (I), we are able compute the number of iterations for ow,, to reach 1/47. It at most
Ti & [log2.56
(102)
(wly 2
Accordingly, wy is bounded as
wly S (1+11117- (ww()))) ((w)))))
1lgg.50
(w) (2/27
< 0.37.
(103)
For simplicity, we just keep the bounds for v - Wy as in Case (I), as
Wx [-0.017, 0.17].
(104)
Case (III). If uw., 4m in first step of stage 2:
From the condition, know uwlz 1, it tells
as well in the last step of stage 1. Since Auty 0 in stage 4n
Wi Wy
ww S
(105)
| |
4m
which means
max
Wr
Wy
2 2 v3.
(106)
Wy
Wx
11) wy ()) is already in the stage Since + 133 > 2.7. wy//u W2 2 2 + v3, then for time 1, (v(1), w
2. However, it is not possible because 2.7
|| = ()))) s 0.1, which means (()) (1) can not reach
1+2.72-
Therefore, the only possible is
Mx wy
2 2 V3. this case, we are able to bound Wy as
(V 4
wy (2 (2-v33 < 52
v3)
+ 0.00852 0.0085
< 0.078,
(107)
where the second inequality due vw, < the bound of v w. from Case (I), although it is loose somehow but it is enough for our analysis.
4y and v w. N 0.017. Note that here we still
We leave the analysis of the bound of number of iterations the end of this section.
Stage 2.
In the case (I) of stage 1, where the first step in stage is with vw, 2 2 it has v WT E 1-017, 0.17] and wy < 0.44. In the case (II), where the first step of stage 2 is with vw,, T Tm has wx [-0.017, 0.17] and wy < 0.37. In the case (III), where the first step of stage is with vw,, | 2 has wx -0.017, 0.17 and wy < 0.078.
1 11
38

---

Under review as a conference paper at ICLR 2023
To upper bound vwl in the first step of stage 2, there are two candidates. One is from the case (I),
((++)) vwx (1 (1+ ((1 ow,
arctan(
Meu wy
/ 1+(1,, / π
K W W
-0²
w
|w:1|
arctan Wy
Wy Me//ex (mmy/Ww))
-K(v w.)? 1 vw
sww, (1+ [1+K(1= -vw=))² + K
-
wy
Wx
x |0:11
K(v -we) (1+ K(1 - vw=))
1 1.+1.111
_
2
+ 1.1 0.44
(-4
1
+
1
+ 1.1-11111 0.172
(1+11(11) -
2T
2
4m
2n
2m (108)
$ 0.668,
where we use ow, 2 1/4T, x
//111111/1111 (1 + x²) S 0.5 for any æ.
One is from the case (II),
y (1+1) (1+1)
S war (1 + K(1 vw,))) + K
-w0 +
Wy
Wsx
||
+ K(u - Wo (1
(1+ +
- K(1-ow,))
1 1.1 11
2
1-1- -..7-..77
1.1 0.172
(1+
1
4π
4T
2
4x
< 0.48,
(109)
where we use ow, S 1/ LT, x (1 + x²) < 0.5 for any Therefore, we can see that, in the first step of stage 2,
tw < 0.668.
(110)
Next are going show how the iteration goes in the stage 2. In Case (I), there three facts:
1. " < 0.44.
2. V Wr E 11000077,001
3. ww2 E 2n 0.668).
Similarly, in Case (II), there are three facts well:
1. wy S 0.37.
2. V Wr [-0.017,0.17] 1 In 2
3. ow E
The main idea is to find basin that any iteration with the above properties (i.e., in the interval) will converge to and then stay in. The method is to iteratively compute the ranges of the variables for several steps, thanks to the narrow range of K. Before explicitly computing the ranges, let's write down the computing method, depending on whether or not uWa 2 1.
Consider any iteration with ewe [m,, w, [di, wy S compute the bounds of g(++1),, (++1) (u(++)) - (t+1) (t+1) in the following process (naturally assuming dj < 0 < d2) wy
1. If mm 2 1:
39

---

This is due to Under review as a conference paper at ICLR 2023
(a) Compute W/ 2 V mm1-- + (d2/2)² d2/2 = f.
(b) Compute My S e 9. we
aretam(w, //=) mame
arctan (rctan(1)-g/(Laa))) A h.
(c) Compute
(d) Compute v((+1) 11 (1+1)
m2(1 + 1.1(1 m2 hi)² + 1.1 m2
h) max([d1], (d2)² 1.1e2mg. This is from
v(+11 (++1) 2 ww2 (1 + K(1 - vwl - -1)) + K W
-12 12 + wil
π |10 1?
v
+ K( - wzz) )² (1 + K(1 - 1wx h))
uw, (1 K[1 vtūz h))² K Kwg y vw
((v w/ )² (1 + K /w1 h))
(e) Compute w0++1) (t+1) (t+1)) S m1 (1 + 1.0(1 - m: m1))² This is due to s01 + K(1 - ==))²
W2
decreases with x increasing when 2 1.
(f) Compute (tt-1) we (l+1) E [d1 (1 + 1.1(m2 - 11+1 h) 1.1e? (V (d2 d2/2)), d2 + 1.1(m2- - 1+ h))]. This is due to
Ar Aw, K(v- wz) uw, 1+ (arctan(a)
1 π
wally
K
wyy
-pp + wwy
||0||
v
||
where ww, 2 1, the last term is between v and O.
gg) Compute wy (t+1) I e max max( [2ll). where
V m2 + (d2/2)² + d2/2 (-m2), j1 1.1
(111)
V mi + (d2/2)2 d2//
m1 + (d//2)2 d1/2
1
j2 = 1+ 1.0
(-m1 +
) 2m
(112)
v /ml + (21 + di d1/2
Awy Wy
uw, + W. Wy )
Wa
7 ||l||
then we would like to have the smallest value as f1 - 1 and the largest value as je 1. Since wy is always non-negative, taking the maximum absolute value gives the upper bound.
2. If m2 1:
(a) Compute
Wy 2 /m1 + (d2/2)² d f
(b) Compute
Me S e/f 2 g. We
wr, /w2) Tl
arcten())=g((zt9) 2 h.
(c) Compute
7
(d) Compute (1++1) wx ((-1) 2 min (m) x01 + 1.0(1 - h))² - 1.1 Compared
with the case of m_ 2 1, we drop the term 1.1(1 - m2 - hi) (max(1d] |d because
it is possible have ww 0 in some iterations.
(e) Compute 1((+1) (0 S max.e[m,,m) x=1 + 1.1(1 x))² + 1.1(1
a) max{Idr), |d21] Compared with the case of m, 2 1, we add term depend- ing on the |V Wa Imax because it enlarges aw,, in the in-balanced case.
(f) Compute af (1-1)) w (t+1) E [d, + 1.1(m2
+ h) -1.122 (Vm2 (d2 2)
d2/2)), (1+111 1 - 1 1 h))]. In fact, a rigorous left bound should include more
terms to select a minimum from. Here it is simple because it keeps + K K(m, - 1) > 0
in the following computing, so we do not need to worry about the flipping sign of d, and dg.
40

---

Under review as a conference paper at ICLR 2023
(1+1) (g) Compute wy S (1a where J1, are the same with those in the case of mm 2 1.
Therefore, with the above process, w(t+)) (++1) (((+1) w (++1), Wy (1+1) from the current ranges. Note that this process plays role of building a mapping from interval to another interval, which covers all points from the source interval. However, it is loose to some extent because gradient descent a mapping from a point to another point. The advantage of such loose method feasibility of obtaining bounds while losing tightness. To achieve tightness, later we will also include some wisdom a point-to-poing style.
able to brutally compute the ranges of
Also note that, nice way combine tightness and efficiency in this method is to split and merge intervals when necessary.
For Case (I):
Now we are to compute the ranges starting from the interval where / = {wy S 0.44, v - we E [-0.017, 0.17], ow E 12x 0.668]] First, we split it into three intervals:
1. 11 (w, - 0.44, wa [-0007, 0.17], vw, [0.213, 0.4]).
2. I2 {w, < 0.44, wx 0.017, 0.17 vw, [0.4, 0.668))
3. I00 = (wy < 0.44, v - w2 E 0.010.0170000 0.17]. ww E 2a 0.213]].
Then, following the above method with splitting and merging intervals, we have
Starting from Z1.
(a) Step 1: 11 mapps to 13 = {wy Σ 0.416, - we [-0.162, ww, [0.55, 1.12131]].
(b) Step 2: Splitting 13. we have
44 = {w, S 0.416, v wy 0..162 0.068), vw,, 6 [0.55, 0.8]]. ii Is {wy I 0.4166,
wy -0.162, 0.068). ow,, E [0.8, 0.9]). iii. I (wy S 0.416, w, -.. 0.068). uw, E [0.9, 1.0)).
iv. 17 =[wy I 0.416, w. 10.162 0.068), vw, E [100, 1.113311
Then, have
i. L4 mapps
Ig [wy < 0.214, v |- 0.309.3090000 0.0545), ow,, [0.942, 1.25786))
ii I; mapps to
Ig (wy < 0.0966, wx E 0.335 0.0613) ow, E [0.880, 1.19649]].
iii. TO mapps
110 {wy S 0.0756, v wa E 0-.360002222 0.068), ow, E [0.777894, 1.11178))
iv. I, mapps to
I1 {wy I 0.134, v w2 | 0.494, 0.0782), ow, [0.595, 1]].
(c) Step 3: Splitting and merging 13, 49 I10- 41, we have
I12 = {w, < 0.134, v Wr E -0.394, 0.078) uw, E [0.595, 0.777])
ii I13 {wy $ 0.214, v wr -0.394, 0.078) ow, E [0.777, 1]}.
iii. 714 {wy S 0..114 v w. 00.362, 0.068 ow, E [1,1,11178]].
iv. I1s {wy I 0..144, v W. 0.3309, 0.061], uw, E 1.2578)).
Then, have
/12 mapps to /18 (wy < 0.0372, w. - 00..0.317, 0.061), ow,, E [1.14493, 1.31246)..
i. 413 mapps to
J17 (wy & 0.0432, v u2 0.448, 0.078) ow, E [0.943633, 1.24393))
iñ.. 414 mapps 118 {wy < 0.0662, v wx E 0-. 462, 0.077], vw, E [0.77846, 1]}-
iv. Is mapps to 720 (wy 0.0998, w2 0.456, 0.0785) ve0, [0.550, 0.878)
41

---

Under review as a conference paper at ICLR 2023
2. Starting from Iz.
(a) Step 1: I2 mapps to I21
{uy S 0.332, wx [-0.205, 0.114]. ow
[0.864, 1.25894]]
(b) Step 2: Splitting /21, we have
i. I22 {wy Σ 0.332, we -0.205 0.114]. uw2 E [0.864, 1])-
ii I23 {wy S 0.332, v Wh E 0.200 0.114), vw,, E [1, 1.128998.
Then, have
i. I22 mapps to
124 (yy I 0.081, V Wx E [- 0.336, 0.114), uw,, E [0.858, 1. 1.14813)).
ii /23 mapps to
125 {wy I 0.184, v wr E 0.499, 0.148), uw, E [0.463, 1]].
(c) Step 3: Splitting and merging 221, 125, we have
i. /26 {wy - 0.184, v - we E 4099 0. 0.148] uw, E [0.463, 1]}. ii I27 = fuy & 0.081, V - we E [ 0.3366 0.114] ww, E [1, 1.1413))
Then, we have
i. I28 mapps to
128 = {uy $ 0.083, V - we E 0-4452 0.148) uw, E [0.952783, 1.37778)
ii T27 mapps to
T29 (uy & 0.034, v wi E | 0.399, 0. 0.133] un E (0.777, 1]).
3. Starting from 130.
(a) Step 1: Iao mapps I31 {wy 0.44, Wr (-0.124, 0.037] vw, [0.422, 0.767]}
(b) Step 2: Splitting J31. we have
i. Ia2 {wy 00.4, W2 E -0.124, 0.037 ow,, [0.422, 0.5]].
ii 133 =(w, < 00.44 Wx E -0.1244 0.037 ow, [0.5, 0.6]
iii. I34 {wy S 0.44, wx E -0.124, 0.037 vw, [0.6, 0.767]}
Then, have
132 mapps
135 {wy I 0.301, V w2 E | 0.2188 0.0185], ow, E [0.901, 1.20971]]
ii. Ja3 mapps to
736 (wy 0.262, 1 w 0.245, 0.023). un,, E [0.96322, 1. 1.25093))
iii. Jas mapps to 137 {wy & 0.213, V wr E | 0. 288, 0.029], uw, E [0.947, 1.25345]).
(c) Step 3: Splitting and merging I35, 136, 137, we have
i. Las {wy < 0.301, v W2 E -0.288, 0.029]. vw, E [0.901, 1]}.
ii. I39 {wy < 0.301, v Wr E -0.288, 0.029]. uw, E [1,]]
iii. 1r0 {[uy < 0.301, v Wx E 00.288, 0.029] vw, E (.21,93)
iv. Ia = =( I 0.262, v w. E [ 0.245, 0.029]. uw, E [1.25093, 1.25345)).
Then, we have
i. Las mapps to
Is = (wy I 0.0404, v - w₂ E 1.392,00029) ow,, E [0.888, 1.1 1.11696]].
ii J39 mapps to
LT3 {wy < 0.0740, w2 |- | 0-.428 0.033], ow, [0.741, 1]].
iii. Jao mapps to Ia (wy < 0.125, v wr E 0. 4.48 0.038), uw, E [0.497, 0.891])
iv. In mapps
Is {uy & 0.109, v -- w. E 0.400, 0.038), aw, E 00.534, 0.702)
(d) Step 4: Splitting and merging L1,, 443, /44 445, we have
i. 748 (w, 0.125, v W. -0.482, 0.038). uw,, E (0.497, 0.891]). ii. T47 {wy 0.074, w. -0.428, 0.033], uwz E (0.891, 1]). ii.. L48 {wy < 00.041 v W. E 00.40, 0.029), ow, [1, 1.1166])
Then, have
42

---

Under review as a conference paper at ICLR 2023
i. La mapps to
I19 = ( $ 0.0424, V - W2 E 0 ow, E [1.07853, 1.34708))
ii [47 mapps to T50 (my I 0.0110, M2 | 0.435, 0.033). rw2 E [0.993, 1.13943)).
ii.. 4s mapps to Is1 (My & 0.0109, v w2 0. 0.454, 0.033] owr E [0.497, 0.891]].
For Case (II:
Now compute the ranges starting from the interval where I {wg 0.37, v
w
[-0.017, 0.17], ows E L 2nl First, denote
1. 132 {wy I 0.37, w E [- 0.017, 0.17], ow, [
2
Then, following the above method with splitting and merging intervals, we have
1. Starting from 152.
(a) Step 1: Is2 mapps to Is3 {wy 1 0.37, wx [-0.079,0.024) vww
[0.222, 0.616]]-
(b) Step 2: I53 mapps to 154 = {wy < 0.343, wx E |-0.171,0.017, UWC E
[0.621, 1.24894]].
(c) Step 3: Splitting J54, we have
Iss - {wy 5 0.343, v - Wi E 0.017] vw, E [0.621, 1}. nw, E [1, 1.24894]]
ii Iss {wy S 0.343, v w2
Then, we have
i. Iss mapps to
um? 157 (wy I 0.150, Wa 0..305 0.017], untz E [0.840, 1. 1.25908).
ii Ise mapps to
158 = {wy S 0.137, v - wi E 0.366, 0.022], uw, E [0.472, 1]}.
(d) Step 4: Splitting and merging 157, 5s. we have
i.
ii I59 = (w, $ 0.150, v we E - 0.3367 0.022], ww, E [0.472, 1}.
iii.
iv. 160 -w, S 0.150, v we 0.3000 0.017], vw E [1, 1.25908).
Then, have
159 mapps to
161 {wy I 0.0705, w2 1..393, 0.022]. ow, [0.971, 1.304]]
ii Teo mapps to 162 (wy 5 0.0613, v w. E 0.02,,, uwe, [0.583, 1]}.
For both Cases (I, II):
From /16-20. 20, 128, 129, 44- 51, 161 462, we can see that it has fallen into an interval 1, {wy < 0.1,v - wx E [-0.462, 0.148), ow [0.497, 1.34078)] Something special here is that wy has been much smaller than Wx- More broadly, let's define an interval I, generated by Ig fwy 0, v w. E [-0.464, 0.148), vw. [1, 1.5]]. Here "generated" means
I. ((((2)) (T) w (( ((() w(() Ig).
(113)
T2t
Then each element (u, Wxy wyy) E I has the following properties:
1. wy = 0.
2. uw, 10.181,155..
3. If vic < 1, then wy E [-0.735, 0.23). If ene, > then w. [-0.477,0.1488)
43

---

Under review as a conference paper at ICLR 2023
The first property is obvious. The third can be proven follows: for each element (u, Wxi Wy) E Lg. it has v(t+1) wz (1+1) = [v - wxr) (1 + K(ww> - 1)), where the ratio + K(ww, - 1) E [1, 1.1 - 1)] when vw [1, Furthermore, in the proven 2-D case, we have shown that "if 2w with some mild conditions, then v(t+))- (-1,1)). Actually it can be tighter as 1-#
y((+2) E (-0.2,1) because here K 1. while the original bound is for K S 1.5. The condition of bounded IV- Wxx can also be verified, the purpose of which is to keep e, wy always positive. Then the bound [-0.2, 1] will tell e wx E [,.0.474, 0. 0.148) on ww, 2 1, because
0.148
0.474
0.474 0.2, 0.148
0.2.
For the second property, the left bound can be verified as
min sx11 + 1.1(1 - x))² + 1.1 - x) 0.474 = #E(1.13]
10(1 + 1.1(1 (1==)))) + +11 - 2). -1) x
1.5
2 0.181.
The right bound can be verified as
max s011 + 1.1(1 - =))² + 1.1(1 2) 0.735* < 1.5. ==
After proving these three properties, we would like to bound how far 1 is away from 7., More
precisely, the distance is measured by Wg. We are going to show wy decays exponentially.
Remind the update rules in (80 81). Denote = n (arctan((a)
" - - ||l
again and 5 = 10 (-u2+
ewy ). then it
>||
Av Kw,( (-v + 1) Kw +
(114)
Aw, = Ko(- -w + 1) - Kvy,
(115)
8 E [- K vw 0].
(116)
Note that both y and are very small, so we are to show their effects separately, which is enough to be a good approximation.
w ( (t) > 1 and the corresponding ( Let's denote (((+1), ((+1) Consider an iteration where (1)
wr
as the next parameters with no corruption from 7(0), Similarly, we denote ((++1), th (++1) are corrupted with (0) From the 2-D analysis, we know
((++)) w (++2)
(4) (1 (w 1)) K(u(t)) w(t+)) 11) < 1.
(117)
(1() _
We would like to show, with small 7(()) and ignoring 8, ((++2) - (1+2) (1)) (
(1 K(u -1 1+
+ K(glt+)) (1+1) + ((+1)) E 1,
(118)
where 7 (t+1) is in time (t + 1) accordingly. The difference of LHS of the above two expressions
turns out to be
118
117
(1 + (1(((1)) ((+1) 1))
+11+ K(v w!! - 1))K((((+))))))))) - m(+t)) (++1) + ((+1))
+ 0(72)
=K,(0((1 + k(( ((+1) -1))
+ +K(11 + ((v(()) w -)))(-K(r())))) (2))0))))) K(W())27(1) + ((+11) + (662) K(n(())2
1 + (1 K(r 2w 1)) K(0 (((0))-((ud)))) + 3(0h))) + 0603)
< (0000)
11+(1 (1 K(v 2 1))( 2kv w0 + ) + +0(7)
< k2(1))
(119)
44

---

Under review as a conference paper at ICLR 2023
Since
Aws Wix
K v we
(-oWs + 1 - 7),
have
(t-1)
(()
Wx
+ K
w (2)
+ 11-110)
< 1.
(120)
2x
W
Also we have
2 (++1)) [e++y
w(r ww
y(+11)
arctan
[|,(e+3)
(121)
y(0)
Wi (1) arctan
20 my
(9)
||:(.) |
Since Wy
(1+1)
I (t) and
ma arctan(mx) 1+
S m for any m > 0,
(122)
arctan( (x)
T4.22
we have
(1+1)
(123)
700
(-vlt 10 +1 - -110) 3.
K
plt) [1T
For general ww, (1, 1.5).
(13
holds
((+1)
S 22
(124)
7(0
VL+0.077-000771 (-1.5 + 1))3
140.07444000077
Since + K(n(0 (1)
uiz
1) S
1.55, is fair to say
118
117
Kg
2 0 35.2 0(73))
(4)
1
1.55
(125)
Actually (0) is bounded by
(t)
wi
0.099
0.1066,
(126)
(t)
v1 0.0743 0.074
Wi
arctan(r) Tt.
(0)
< 2.6 10 10-4
(127)
As result,
118
117
0.0084.
(128)
Note that this small value is very easy to cover
117 requiring
(((+2)
(t+2)
wzz
0.0084,
(129)
(11)
wzz
except when ww is pretty close to 1. When ww + 1. from the analysis of 2-D case, (derived from the case of T2+191+11 2 x
((+2)
(t+2)
wz
2 (2K - 20)x(((( (0) 1).
(130)
(())
w(()
For 118
117
denote a function p(x:) as
p(z) + (1 Kx)
-2K(x+1) +1
(131)
1 K = (-2)
45

---

Under review as a conference paper at ICLR 2023
where x y(O! 1 in 119 123p. It is obvious that p(0) + (-2K 1) < 0. When is small, it turns out
((()
p(z) 2K + -2K - 1 +
x + O
(132)
W
As a result, (18
117
< when vwly - 1 = o(K - 1). Wf - 1 Q(K - 1)? Actually,
we can get better bound by more care analysis,
118
117)
S1+ (1+ K K66 - 1))( K(6(0)2 - K(w + 2(0+1)
Kn(t)
y(1))
K
(1))) w0 (1 + K(1 v101o(1)) -1
(133)
where the last term is due to ((++1) (+11) by expanding the last term, (322 becomes
< v10 uz (1 + ((1 - u0 (1)) )²² Hence, with this bound,
(1)
Wx
v(())
p(x) = 2K +2+ K
-2K 1
x+ K(1 - 2k)++ 0x²
(134)
We
(1))
-K + K
-4K +3
2 + O
(135)
(1)
w
which is definitely negative because
(
/1 + 0.074 0.074
4
1.16
(136)
(8)
v1 + 0.074 - 0.074
3
W
Meanwhile, we are to prove the 8 in 1144 will not make I, make w/ < -0.474 starting
from V Wx 2 -0.462. First, in the region of {ew, E [1,]. S 0.148), we have
K Kow 2 < 1.1
+ 0.074) 0.12 < 0.0144. Also note that in this region with
U we 2 -0.462, we have
(++1) Wy
vI 0.231 - 0.231
0.37.
(137)
Wy
vI + 0.231 + 0.231
Hence Ku(ug) (2) + Ke(w, (( 22 < 0.0144 * (1 + 0.37*) = 0.0164. Since [b((+2) (+(++))
(o(()) (()) if there no 6, we shall see that there is no need to discuss the case of Wy 2
-0.462 + 0.0164 -0. 00.4556 because it still holds (c(11) - w [1+1) > -0.4 -0.462. When ((() w (1)
[-0.462, -0.4456) we shall see that in (59). after adding the term of ő in v,
(1+2) w (1+2)
-1- (1+ K(ow, - 1)) Kw2ō
(138)
(())
(!)
which means the absolute value of w, decays at least by margin depending on 6. After multiplying the current difference (1) w((e on both side, it gives
(o(0+2) w ((+2)) (e(() w(() 2 w00) w (0)
(139)
Note that here ((t-2) wh! 0.372 < 0.16 due to the decay of Wy. So it safe to say g(t) + g(++1) > 1.168(t), Combining with the above inequality, it gives
(1+2) does not include g(t) and 5(++1), As stated above, we have g(++1)
80c S
(e(0+2) - w((+))) - (v(() - w(()) + 8(1) + g(++1) > 10 (0) + 1.16)8(9,
(140)
where
(1) (0) (t) + 1.16 S vow. (vrvz + 0.4456 )2 )2
0.4456
0.4456.
) + 1.16 < 0.6.
(141)
2
2
46

---

Under review as a conference paper at ICLR 2023
Furthermore, from our previous discussion, my wy by
(++
(++2)
gives that the of 140) bounded
0.6 g(0)
0.6
0.16 (-0.0144) 2 00.0103.
(142)
0.16 1
Since -0.474 (4744(-00462) < 0.0103,, we shall see that the term of cannot drive v w < <0.447.
Note that (140) shall include factor (< 1) in front of 8(0) but we have ignored show a more aggressive bound.
Therefore, we are able to an Interval I, generated by 1 also has the following properties: for each element (v, W2 wyy E L..
1. ww2 E [0.181,1.5]
2. If une, I 1, then v - wx E (10.733,0.22 0.23). . wo > 1, then v wx E [- 0.472, 0.148).
Then the decreasing ratio of Awy/ Wy is bounded by
Awy
1 Myy
K
UW'
(143)
wyy
Wx
π
||
1-1111 1.5 + 0.074 + 0.074) -0.030K
(144)
[-8.8, - 0.030K].
(145)
Hence, Wy decays with a linear ratio 0.97 (or 0.030 at most for Cases (I, II) in stage 2.
For Case (III), in the first step of stage 2, it already has wy 0.078 and wa [-0.017, 0.17].
So surely it will also converge to I,.
Here present the time analysis for Case (III) of both stages. The number of iterations in the first stage is apparently similar to that of case (1, II), as
2.7
T 20g2.56
(146)
g2
where vi < 48
1 the value uw, in the first step of stage 2. In stage 2, since our target is to find how many steps necessary to get uw, 0.181, so it is
1 - 0.181 +1 arctan (2 v3)
2-v3
y (1+1) (++1) 2 w(
1+ 1+(2- v3y
1.1
(147)
2 3.28 w0
(148)
where obviously still holds Wy 2 - v3 and w < 0.12 in stage 2. Since 3.28 > 2.56, we have uc
the total number of steps to have ow, > 0.181 bounded as
2.7
0.181
0.675
0.181
log2.56
loga.28
log2.56
loga.28
+
g2
+
1
4
0.675
log2.56
+ 3
n,82
1.35
log2.56
+ 4,
TT
which is not beyond the bound for Cases (I, II).
J PROOF OF MATRIX FACTORIZATION
Consider two-layer matrix factorization problem, parameterized by learnable weights X E R=xp, E RPX4, and the target matrix is E R x The loss is defined as
L(X,Y)
XY C
(149)
2
47

---

Under review as a conference paper at ICLR 2023
Obviously (X, Y xY forms minimum manifold. Focusing this manifold, our targets are: 1) to prove our condition for stable oscillation on ID functions holds at the minimum L for any setting of dimensions, and 2) to prove convergence result with initialization close to minimum beyond the edge of stability for the symmetric case = X
J.1 ASYMMETRIC CASE: ID FUNCTION AT THE MINIMA
Before looking into the theorem, we would like clarify the definition the loss Hessian. Inherently, we squeeze X, Y into vector 0 = vec X, Y) E Rmptry, which vectorizes the concatnation. As result, we are able to represent the loss Hessian w.r.t. 0 as matrix in ((mp++o)) (mp-+70) Meanwhile, the support of the loss landscape is in mPpPp the following theorem, we are to show the leading eigenvector A - vec(AX, AY) Rm of the loss Hessian. Since the cross section of the loss landscape and A forms ID function Sa. we would also show the stable-oscillation condition on 1D function holds at the minima of fa-
Theorem 12. For matrix factorization problem, assume XY C. Consider SVD of both matrices as X Lim1 pmin Pu.uUgiizg. and Y = - Ty Uy..uyy..y.. where both groups of T.. s's are in descending order and both top singular values Jn,1 and Jy,1 unique. Also assume T 22,1 uy 1.11 # 0. Then the leading eigenvector of the loss Hessian is A = =vec(Cn 1e 12, 1211 12v Denote fa the ID function at the cross section of the +0 loss landscape and the line following the direction of passing (ec(AX, AY). Then, at the minima
with C1
Gv..
of Sa. it satisfies
31 (2) >0
(150)
Proof. To obtain the direction of the leading Hessian eigenvector at parameters (X, Y), consider small deviation of the parameters as (X + AX, Y + AY). With X = C, evaluate the loss function
as
LXX
|AXY + XAY AXAY
(151)
Expand these terms and split them by orders of AX, AY as follows:
@(||X| 2 |
2 AXY XAYII
(152)
O(||A + + [): (AXY + XAY, AXAY),
(153)
+|XXX++
2 IAAAY F.
(154)
From the second-order terms, the leading cigenvector 2 the solution of vec AX,AXY =
arg max
|AXY + XAYI
(155)
Since both the top singular values of X. Y unique, the solution shall have both AX, AY of rank Actually the solution (here for simplicity we eliminate the sign of both)
0,,1
AX
Uz 1 9.1
AY
1 2.1
(156)
+ 9.1
V V
2.1
Equipped with the top eigenvector of Hessian, vec(AX, AY), we consider the 1-D function fs generated by the cross-section of the loss landscape and the eigenvector, passing the minima (X, Y). Define the function as
fa(w) L(X L(X+ + Y PAY), HER.
(157)
Then, around O, we have
(u) 1 |AXY XAY
P . (AXY XAY, AXAY) 113 11+ " P
(158)
48

---

Under review as a conference paper at ICLR 2023
Therefore, the several order derivatives of fa (²) μl can be obtained from Taylor expansion as
(0) |XY I
(159)
(1) (0) = 6(XXY + XAY, AXAY),
(160)
( (0) = AAXXXI
(161)
Then we compute the condition of stable oscillation of 1-D function as
[3 12 (0) = 108/AXY + XAY, AXAY) - 12 + |AXAY| 13
(162)
-66| + XAY |AAA >0, O.
(163)
because all of AXY, XAY, AXAY are parallel to uz, and V.j Ny + O.
J.2 SYMMETRIC CASE: CONVERGENCE ANALYSIS AROUND THE MINIMA
In this section, we focus on the symmetric case of matrix factorization where Y = T Accordingly, rescale the loss function L(X, Y) X - ClI Denote the target C XxX and assume
around the minima X Xo + AX, with small ||AX|]|- Consider SVD as ||.
X = Ex -mia(mp)) a T - d + Xg. Then the EoS-learning-rate threshold at X = Xo is of Therefore, we are to show the convergence initializing from X; = Xo + AX1 with learning rate 77 = o/4 + 3 with 3 > O.
Theorem 13. Consider the above symmetric matrix factoriiaonn with learning rate r J.- + B.
oi
Assume 0 < 3o} < <.455 1.121 and nos 1. The initialization is around the minimum, as X, = X, + AX,, with the deviation satisfying u. AAx1) vn # 0 and || |XX || S bounded by a small value. Then GD would converge to a period-2 orbit ?'n by a small margin in O as
X, + 77 AX,
| = =(e),
(164)
In = (Xo + S X. + δ20ru1e])
(165)
where 61 E (),1),82 E (-1,0) are the two solutions of
1 + 3o
(166)
(6+ 1)2 (v
3 (8+1)²
Proof. The update rule gradient descent gives
X X, (X,X, Xxxx,) X..
(167)
Denoting AX, X, Xx,, the update rule is equivalent to
AX - AX, (XX,X; + X AX,AX)) (X. + AX
(168)
Consider a decomposition of AX, = t + AX, where AX,0 and | = 1. We also control the sign of ū, by claiming (ūe, u1) > 0. Then, the update rule is again equivalent to
AX quæT + AX, - /
(anieT + +XX,) X Xx (aunn + +AX.)
7
(169)
7|
(E@TN + +XX,) (auee) + Axi)
(Xo (Xo+ AX,)
(170)
49

---

Under review as a conference paper at ICLR 2023
After expanding Xu 01u10 + Xo and projecting AX onto V1, we have
AX Gg/ 11 oleccs uu â ee + ofeund in aieu o1e u1
(171)
to AXX, +X #1 AAX ã, GAXXXX, in +xx,XX, dt
(172)
Phase üh gets close to uj sharply from a random direction
For now, let's assume AX stays small (which we will show later), then ignoring high-order small values gives
E++14++1 E,0i: ofeeie + ofEeun in to, 41 +0())),
(173)
(c+110+1,,1) u1) 2n (esūi, 01) +0()) =(16) Boi) U1) + O (174)
m++1 A 44+1u++1 (c++1ū++1, u1 - u1)1) n11 X U1 + +c²),
(175)
where 174) is the projection of AX onto 11 and 175 is the orthogonal-o--u component of AX41144 Meanwhile, we have the following iteration of AX, following the update rule,
AX = AX,
o AX, + AX,X, Xo + AXXX AX, Oieuin X
(176)
toleuuū AX, X X + XAX, AX, + G X,
(177)
+eūū AX, AXAX, X AXAX AX,
(178)
Ignoring higher-order small values, is
AX AX; (0ICuuuu,X AX,X, X XXAX, X,) +0(2).
(179)
Now we are to verify two facts:
1. The orthogonal-to-u component of AX, denoted as m, stays small. Then combining the exponential growth of parallel-to-u; component in (17) gives (we. u1) 1 quickly.
2. stays small.
First is the bound of me/] From 175). we have m++1 = (1 - nof)m - no X,X, 21 + (E2)).
Combining with (19 and noticing in Xo O. it holds
m4-1 = (1 nof)m - non - X,XX un
(180)
T1 XX tn = XX,XXXX 41 - not Xo in + XoXo XOAX, "1
(181)
=II-nXoXo' OxXoX d1 No in.
(182)
50

---

Under review as a conference paper at ICLR 2023
T
Furthermore, we have &g Us Xo my due to Xo 0, so 182 can be rewritten as
X1XaAX++1 u1 (I Xxxo oxX u1 no Xo my
(183)
Combining (88 183 gives a form of m4+1 as combination of all previous terms as
m++1 1--oomm + Σ of (I ( X1-n
) XoXo ma-6-
6-1
(184)
Since our goal is to verify that ||mm is bounded, pick any eigenvector up of XoXo with associated eigenvalue Ap- Then have
(m²+1. = (1-nof) m Up) + Σ
oojc T (1=nxoxo))
XoXo my-t
(185)
i=1
= (1 (11-no7)) (m,,, Up) + o Σ n\-)
(186)
Obviously (m4+1, Up) shall converge to series with exponential growth. Assume the ratio as (me+1: (m,, up) = r. The above equation is equivalent to
1 m/pp//r
r -Bo + 1220 oī Σ
1-1 1
-Bo
(187)
1/ oooR19
= -Bf + -- 1 1 1
(188)
The solutions for this equation are r or 1- n)p - -nos. both of which are in [-1,1] 1] once Ap S (1 (11--03]) /11 Hence, it safe to say (m++1, Up) bounded as non-increasing. In fact,
(m Up) is bounded by Kmo, uphl because ((m1, Up) < (mo, up) due to the scaling factor
Go < 1. Therefore, after picking any eigenvector vp of X we can conclude ||m || |moll s
|col. A further result is that (m1, up)l x Ap. Notice that
(m up) + Bo (m,, up) = 'o E 21-11 (1p)))p) (m4
(189)
1=1
where RHS a constant due to the proven non-increasing (ra,, up>il But this constant is proportional Ap because (m2, up) + (m,, vp) not (mo, Up) x Ap and all further iterations follow similar factor. Therefore, we have [m11, up) x Ap-
Now let's show
AX;
stays small. Consider u
AX 179 gives
AA 11 uf AX, n-noo X1)
n01m, Xo
(190)
n) T momm Xo (I Xo)
(191)
so the norm is bounded as, for some of singular value of X
S no i 1- no²) Co2 =co
(192)
where Co" are from the previous discussion of [(m1, up>i x Ap. which follows a AAp Hence,
is fair to say u AX,
stays small.
51

---

Under review as a conference paper at ICLR 2023
Meanwhile, the residual component of AX, that is orthogonal to Un on the left, denoted as A
iterates following
AX AX, -1 (X, Xo X,AX,,1 X
(193)
AX AX Xo X,AX
AX,, + 1XAXI11 - /x, XoXo
(194)
+ +X,XX, X +X X
+XoX,
X
(195)
AX,, Xo + -XAX..1 ")
(0.5I-1XX08) -ixon
(196)
+ (0.II- xxXXo
(AX + xXAXX.I))
(197)
As result, due to XoXo
Soz < 1/1,, the following norm is recursively bounded as
+ +XAAX 7| s | XXAAX
XX
(198)
Since AX++ is polynomial of Xo and its transpose, the above bound tells that AX not grow, which means
does
A
S AX
(199)
Combining (192, 199h, we can conclude that AX stays small.
After proving that both m, and ÁX
stay small, from 174). the only term growing fast is (e,ūe, u1)
exponentially, which means the projection of ūz onto is also growing sharply.
Phase II: after ū1 u1 approximately
In the first phase, (egūe, u1) grows exponentially with all other components in AX, keeping small. The consequences of such growth are
1. ú₁ gets close to ij in direction, with the cosine similarity between them growing like cosl (ac)aaneexp(()). where f is time variable starting from some constant.
2. while ü, is a unit vector, the growth of (ceūe, u1) makes e, grow sharply as well. In the phase I, the proof is strongly dependent of the assumption that e, is small. But in the phase & is not small any more.
We would like to assume that the second consequence comes later than the first one, which is feasible make the initialization smaller. Then, would like to verify the dynamics of (Epūe, u1) when 6, is relatively large.
After keeping higher-order terms of c, in 172 and considering uJū, ü, zs 1, we have
(e++1ū, u1) = (e, u1) (ofer o + et + off + o + o + euu, AX,X, ū
(200)
(q u1) - (20 + 3 + + quu Ax ut).
(201)
52

---

Under review as a conference paper at ICLR 2023
We are to show that u/ XX iy is small, so that (E++1ūe, w1) is function approximately of only
7) Ct 01.
After re- re-introducing higher-order terms around 178 we have
AX AX;
01EEUIi Xo AX Xo + XOAX, Xo ue Xo)
(202)
AX 11 01tellii1y X AX,X, Xo XOAX, Xo (mem, (c)1m)
Xo
(203)
Then, holds " X
ir T AX,
(204)
u/ AX
(1-nxxxo)) mh-1 7/01 + +.)m, XoXo my
(205)
S u Ax m + n + a)m, X, m
(206)
So we have to ensure that [|m=]] keeps small. To obtain this, we start to re-write the update of my by my egit - (erūy u1 u1)u1, which
m, = 1(1 100 + m4-1 1101 AX W1
(207)
T
X m4-1 - 1) (I 11 AX,-1X0 7--1-
(208)
Compared with (180). the above equation has two differences, the first is the coefficient of m-
being (1 - 17(02 + instead of (1 (1-11) njo²), and the second is additional two terms in (208). Let's
first discuss the second difference. Actually these two terms somehow act as the coefficient as well,
so if X <<
then it is safe to neglect them and only consider - -0=)))). To show this, we again
look into the stability of 1 T TTXX+11
and
|
||
We have
AX = in AX,
(I (1-1/a
11 Xo)
nai + +e)m Xo
(209)
AX AX.1 XX Xo AX
Xo
(210)
where the first line simply differs from 191) with the coefficient as 77(1 + 44) instead of nj and the second is the same as (193). Therefore, following the same arguments in the phase if both of
||m;-I|1,
Ax,
O()) are small, we have the following step they are still in O(e).
Therefore, have the dynamics of é, as (c+-1ū; un) = = -1 (207a + 301 + <?) which corresponds to the update rule of gradient descent (c - of with learning rate y, Since (0) 2a the learning rate = , + 3 2n
(211)
ID function f(e)
by a margin Sof. Since for ID function f(x) = 1x² - 11)2 with learning rate 17. it converges to the two positive solutions, c, in (39
1
17 x2
(212)
11 1t
2 +
Therefore, for ID function f(e) = ((e (le+ of) 2
with learning rate n, it converges to the
solutions (one in (-01, 0), one in (0, of
1 Bo
(213)
( + 112
3 +
(+)+))²
Along with the above argument stability to the above solution with deviation bpper-b by O(e).
AX
O(e), it concludes that the trajectory converge
53

---

Under review as a conference paper at ICLR 2023
J.3 QUASI-SYMMETRIC CASE: WALK TOWARDS FLATTEST MINIMA
Theorem 14 (Restatement of Theorem 4). Consider the above quasi-symmetric matrix factor-
ization with learning rate ll = o7 + G. Assume O < fof < v4.5 R 1.121. Consider a minimum (Yo aXo, Zo 1/aXg).
1
The initialization is around the minimum, as Y1 = Yo + AY, D1 = Zn + AZ,, with the deviations satisfying un 1Y1 # 0, uj AZ # 0 and || Y |Z Sc. 6. The second largest singular value of Xo needs to satisfy
(o (1+010 10 of )
(1
of (1+
}
max
02
aAo
S2.
(214)
Then GD would converge a period-2 orbit "hy approximately with error in (c)), formally written as Z.) + 7/n (AY, AZ), || ||AZ| O(e),
(215)
In { (Yo + (p: - a) 1 Zo + (pi - 1/a) 10a-)aaauuu)) }.
(i 1,2)
(216)
where P1 E (1,2), P2 E (0,1) are the two solutions of solving in
+ 3o
(217)
p² - -
Proof. The update rule gradient descent gives (for t 1 1)
Y Y, 7) (Y,Z, - -xoX,] Zt
)
(218)
Y41 Z. 1) ( -xxX,) Y..
(219)
Denoting AY, = Y, Yo, AZ, Z2 Zo. for 1, the update rule is equivalent to
AY AY, 1 + YoAZ, + AY,AZ,) (Zo + AZ,),
(220)
AZ AZ. -) YA AY, + AZ,YY, (Yo AY,)
(221)
Consider the decompositions AY, = En, AY AZ, = ". + AZ, where we assume YY,u 0, AZ 0, / || 1, || 1. We also control the sign of uy uz,e by claiming (ug,,, u1) > 0, (ūz,t, u1) > 0. Then, the update rule is again equivalent to
/f (Sunthero +Y,) 2 + Yo (Ex, + AZ.)² +
AY Ey + AY, -
(222)
(Iyetineeo) + AY.) (ISniZinN) +Z,))
Zo +
(2 T + AZ;),
(223)
AZ+1 = E2, + ÅZ, - (SrAVENNT + AZ.) Yo + Zo (cy.tūy.Au) + AY.)
1+
+
(224)
]
(Ir.eigeee) + +YY.) (cx.(uenUT + AZ.))
(Yo (Yu + + AYY).
(225)
Phase I: uy. and ūz,t get close to in sharply from random directions
Decompose the three matrices at the ground-initialization minimum Xg = 014141 + X,, Yo =
Laniieo + Yo, Zo Laqnio Zg- Obviously they satisfy X, Yo aZz- For now, let's
54

---

Under review as a conference paper at ICLR 2023
assume
AY
AZ
stay small, then ignoring high-order small values gives
6g,+100,441 AY = Cg 17 Cu. of
a + a YoAZ. u1 + t,,
(226)
62,4-10,441 = 14 = E2,000, - (c,,02,,0faa + + "1 + f vo
v.
(227)
then we have
Pu.t11 & (cy,41100441,,1)) = 1 11 of a²
py. Pz,t
(228)
of
Pet+1 A (ta,0t
1 1|
P2.--110 - -111 pu.
(229)
a2
From the above two, we would like to find lower bound Py,t: P2,t- Note that
Pee++ - aa = P2, P2 k,
(230)
Py.e+1 =
of 02
-11 -1 - of py.t n0o (a?Pa.t Pott +k) = (1 -o-oooaa=
Pg.e -non k. (231)
Since - n0ja² - 17 o < 1 2no < 1, ratio of 2nj least. Meanwhile, with P2,a a² uue fixed along time, it holds |P,,(] is growing
see |Py.el is growing exponentially with the
exponentially with the same ratio as well.
Now let's how other things stay small when pyte and P... the only two terms growing exponentially. If that holds, we can conclude that uy and ūz.t get close to sharply from random directions. Similar to the discussion of the symmetric case, we have two remaining components as follows (which are wished to be bounded)
of a²
my A Gy.+10 - (cy,441 (yutt10y,tti unjun = (1 1 -
My,t - 01 YoAZe U1
(232)
m2,+11 (2,411021411 ((2,1410 u1)u1 (1 - not a²) m2,t 1701 Zn AY,
11 (233)
u AY (+1 = u AY, - 17 u AZ, Yo
Yo + tomm X1)
(234)
" AZ = ul AZ, u AY, Zo
Zo + Tim Xo)
(235)
Take any eigenvector up of XoXo with associated eigenvalue of with op > 0. Then the above system can be written
(mg,441, Up) = 1 11 x2
(my,e Up) 10 AZ., up)
(236)
(mz,441,%) U/)) = =(1 - =)) (mz,2, Up) - 10 (u) AY c, upp)
(237)
(uT AY Up) = (u AY, es upp) - 2 AZ,, Up) + 110 (m
(238)
of AY ( + (00 (10p(ma,111; (uj
[u] AZ Up) = (I AZ, up)) - a2
(239)
Then the above
system which maps from
be re-written (myy 12), (m Up).
can
matrix A E R4x4,
Au up), (u, AZ,, up))
to
(mmy,t vp), (m vp), (MT AY V4t vp), (u] ÁZ v,)).
The explicit form of A is
55

---

Under review as a conference paper at ICLR 2023
-
11 a
10
0
nofa
pp
0
A
pp
0
na?o
(240)
0
0
11 n
1
0
I
Q²
01
a
0
1
0
d a 21
61
noja²
1
1
010
0]
(241)
0
di
Op or1
A I.B.
(242)
So have rank matrix B with all elements non-negative. Our target is to show A " is bounded with n - 0c. Hence, we require the spectral norm [
S 2 by Lemma & which gives
of
of 1+
) ato9
max
n a²
1 1
. moja²
=2
(243)
Therefore, this gives an upper bound for d2.
{ of a
02
noia² (1+
}
max
1) a2
adoo
2.
(244)
of
Hence, with the above discussion, we have shown || ||mm
u AY
2nj
AZ:
stay
small this phase. Mcanwhile, the residual components of AY, AZ, that are orthogonal to ún on the left, denoted AY, f, 1, AZi, iterate following
AY AY - AZ Yo Yo ZoYY. Yo
(245)
AZ41. = AZ -1 AYY Zo Zu YoAZ..1 Zo
(246)
Then we have
M A AYH Xo T
+ +XAYHI + 2AZIH, Xo + +X,AZHH,1
(247)
AY + XDAY +O X + + AZ
(248)
-AZZ Yo YOX, -IXOY YOUL, NYAYY X,XI
(249)
-X AY,X, X T -n00 Zo ZIXD na X ZOY
(250)
A. X,X, 110 X AZ,1X
(251)
AY + XAYY 61 + 1 Xo + AZ
(252)
na? AZ, X X,X
110 X X.
- X,Xo
(253)
X AY, -nAY,
I. X X X X,YY,
(254)
- Xo XoXo -na XoxX AZ.
(255)
MM, (0..5IT7xXx0000 + (0.5I - Xoxo 1) M..
(256)
Then, by reorganizing terms, M || can be well bounded as |M++ ||,]]]
56

---

Under review as a conference paper at ICLR 2023
Similarly, need the bound the norm of the following term
N++1 - (AY Xo X,AYHIA ) + a (AZ441, X + +X0AZ++1.1
(257)
- XOAY +o AZ Xo AZ
(258)
+ na² AZ, Xo X
t X Xo LZAZ,,1 A X
(259)
+ Xo A - - X
XOXO - X X,YY
(260)
X0X0 -T -NJ -na X0
n X AZ,1
AZ,
(261)
= o² AZ X
I [111XoXX0
T aZAZI,,1X. (
(262)
(11)X
(oxo
d'
(263)
I-
X-nXoxx
AYHLIIX
(264)
(-
(I-nX
X,AYHLL
(265)
Regarding the above equation, there several key observations:
(a) N2 is symmetric.
(b) For any eigenvector of X,Xo have o NN+1U= u N.v.
(c) For any distinct eigenvectors Up: Pg of X,Xo we have 1i Net1 + v/ No+1"p =
1/ New. + 1/ N.up.
Combining these three observations, we have "NN441 Ue Neo for any with decomposed as linear combination of the eigenvectors of Xo Therefore, it fair to say ||N++1|| ||
So combining |M++ll S | and |N++1|| |N tells that both || (+1, _ T
Xo4Y+ |l, | Xo
AZ
|| stay small, which tells AY AZ.
also stay small.
Therefore, conclude that both
AY
AZ,
stay small.
stay small, from (174), the only terms growing fast are (e u1), (cpūz,t uz) exponentially, which means the projections of Ug ūz,t onto Uj is also growing sharply.
After proving that all of ||mayeell II ||
AZ,
Phase II: after un üz || u1 approximately
Following the same spirit of everything in the symmetric case, we re-introduce higher-order terms, as
AY+1101 & cg.tiing &y 17 cy of a² + E T1
(266)
tofe,ru 2 2 + σ Egg Egnee, Tyyeo
(267)
a
6,,41 (, = (,
u1) ) -
o2
Ett -
σ tydtza + ofe,, + ar + 2,1
a
a
(268)
6,,41 (c,,4420,,141. un) = (,,,, u1) n) ofa 20100y,66s. + togc, a ve + 2 Ey
2
(269)
57

---

Under review as a conference paper at ICLR 2023
Then we have
01 + g,t+1 = 010 + Eytt -1 ("a
01 + (z,)
{
01
Cut + turte + 01062s)
(270)
Φ1 + +,441
01 + Czf (0 + Cy,e)
CN + Co,te, + 010c2.1)
(271)
Note that the shared factor (: y. + C3,C2,t + 010(2.1) the above system is equivalent
= ( (n,) (a + - of.. Hence,
y y 7yz o2)z,
(272)
2 = : 7yyz o² }y--
(273)
Furthermore, this is equivalent to f(y. { 1)² with learning rate n/ = 11 + Boq. Therefore, from Theorem [ we know y, 2 will converge to the same values, which make the problem with the same solution of 1D functions. Fortunately, it follows the solution for (39). In summary, after straightforward computation, it converges to
(Y,, Z.) + 7/p + (AY, AZ),
|AY|| |AZI = =((),
(274)
iiy --= (Yo (p: - a) 1 Zo + (p - 1/4) ( }. (t= 1,2)
(275)
where A1 E (1,2), (0, 1) the solutions of solving p in
3o1
(276)
3
Note that || ||Z|| (c) can achieved by the already presented discussion in Phase I, and here we omit similar discussion for Phase II, very close to the symmetric case.
K USEFUL LEMMAS
Lemma 7. Assume a Aa 2 b Ab and a 2 All of Da, Ab are positive. If Ab S a, then a + Aa 2 b Ab.
Proof. (a + Aa) (b + Ab)
+ 444 4
Ab ( 1)(b - a) 20
Lemma 8. In 242 if we have
B < and with 02 < Tn, then A' is bounded entry-wise, for any
t21 1.
Proof. Denote d1 1, (1,0,022,,00 U2 = 1 0.1.0. - .1 = 1 01.0.0, o
U2 = [0 1, on , where 21: 22, #3, are positive normalization terms to ensure |u:||
||0:|| for i 1,2. Then A can be re-written
A = I
11
of a2
1+
a2
2 iiti - 1ooça²
Tp 1+
T Obviously, for any 2 1,
U2
(277)
σ1
Let's denote 01 = o the t-th power of A, can be represented as
=
of 1+
2
022 = (1+
-
2 )²)
I'= (1))
Tn + 10 42 + kg 22 + 110 uzen
(278)
and the update rule of R(
(t)
s is
k ((-1) 1100 (1 anve u1)
(0) 2 41 01
kl (t-1) 11 (1 U2 U2)
(1)
a2u U2 02
k(e-1) 1) 13 (1 02 T U2)
(1 a221 U2
k ((--1) (000 anve u1
(1)
0102 u1.
58

---

Under review as a conference paper at ICLR 2023
Since A is bounded V / if and only if k is bounded for any / 1, we only need to show the below < 1, with C defined as
matrix C with C
spectaal
al u1
0
012/
C
02 U2
22
d2
(279)
02 1 U2
02 u2
a 02
u1
a1
11
where
2
Zz
U u1
T
1/2
(280)
2
V
1
[ea?)
V 1+
aea?
1+
Ee
a²
+
01
Z2
1
U/ U2
11 u2
(281)
+
1+
xp o
1
ov oja
0.
To obtain |C|||2 < 1, we only need to compute we have
||C,,:|
for [4]. Taking = 1 as example,
|C,,,11== (1 anu u1 + (a2n)) u1)² (a2nJu1)² (1---1110)
of
1+
a
11
+
(282)
a2
a
1
a
01a²
2 If a < 1, the above RHS < (1 - ) + (1
)
where the second inequality is due to
of S o3 < o.. If 2 1, the condition of a1 S 2 (from |B|| - 2) gives o7 v 2/n- 2/n-oj/o" which is further o < 2/m-oi/a" /a² of due to of < off As result, holds you - of + no < 0. which
a²
means (1 - / (n < 1. Noting the above RHS (1 - /o/ + (no / a2)2
(1 / + (n)2)² when 1, it finishes the proof.
L ILLUSTRATION OF PERIOD-2 AND PERIOD-4 ORBITS
In the proof of local convergence to the period-2 orbit in (??), we give a bound of learning rate as v5 - ss 1.236. Local convergence is guaranteed if the learning rate is smaller than it. Conversely, the learning rate is larger than it, although the period-2 orbit still exists, GD starting from a point infinitesimally close to the orbit still escapes from it. This is when GD converges to higher-order orbit.
Figure 2 precisely shows the effectiveness of such bound where GD converges to the period-2 orbit
when 1) = 1.235 < 55 - 1 and period-4 orbit when 1.237 > v5 -
M DISCUSSIONS
First, we provide a general roadmap of our theoretical results in Section M.I as illustrated in Figure I0 Then, in Section M.2 we discuss three implications from our current low-dimensional settings to more complicated models for future understanding EoS in pratical NNs, where low-dimension theorems are enhanced with high-dim experiments.
M.1 CONNECTIONS BET WEEWEEN THEORETICAL RESULTS
In this section, we discuss the connections between our presented theoretical results, as illustrated in Figure T0
Theorem and Lemma D present (local) intrinsic geometric properties for a 1-D function to allow
stable oscillations. Such properties provide us the 1-D function f(x:) (n 2.22 and, furthermore,
59

---

Under review as a conference paper at ICLR 2023
0.08
Ir=1.05
0.08
lr=1.235
0.08
Ir=1.237
0.06
0.06
0.06
0.04
0.04
0.02
0.02
0.00
0.00
0.00
0.8
1.0
0.8
1.0
0.8
1.0
Figure 9: The convergent orbits GD on f(x) = 1)2 with learning rate=1.05 1.235 and 1.237. The first two smaller learning rates drive to period-2 orbits while the last one goes to an period-4 orbit. The significant bound between period-2 and period-4 predicted by our proof in (??),
as v5- 1 1.236.
Local Geometry (Thm.
High-order LG (Lem.
1-D case (Thm.
LG for MF (Thm. 6
(g(x) y)² (Lem.
2-D case (Prop. B
Balancing effect (Thm.
Composition rule of (Prop
Single neuron (Prop.
Quasi-sym MF (Thm.)
Wy - 0 (Thm.)
Figure 10: Connections between our presented theoretical results. The arrows stand for "implies". LG stands for Local Geometry. MF stands for Matrix Factorization.
we generalize the local property to global convergence result in Theorem Then we are to generalize the 1-D analysis cases of multi-parameter, ii) nonlinear and iii) high-dimension.
(a) Multi-parameter. Compared with 1-D f()) (μ - 22)², the 2-D function f(x, y) (μ xyy)² can be viewed as the simplest setting of two-layer models. We prove that the 2-D case converges to the region of = in Theorem 5 in Appendix A.. which means it shares the same convergence the 1-D model. Also, y means its sharpness is the flattest.
(b) Nonlinear. We extend the 2-D model to a two-layer single-neuron ReLU model in Section 5 Although the student neuron can be initialized far from the direction of the teacher neuron, we prove the student neuron converges to the correct direction (as Wy + 0) in Theorem 3 Then the problem degenerates the above 2-D analysis, which means it shares the same convergence with the 2-D, where (v, w.)) corresponds to (x,, y) in 2-D.
(c) High-dimension. We extend the 2-D model to quasi-symmetric matrix factorization in Section
6 Although the parameters are initialized near sharp minima, GD still walks towards the flattest minima, as shown in Theorem 4 At convergence, the top singular values of Y. Z are the same, following the 2-D analysis. So the singular values are in the same period-2 orbit as the I-D case.
Meanwhile, from Theorem and Lemma we prove condition for base models g in regression tasks to allow stable oscillation Lemma 2 Furthermore, we provide a composition rule of two base models to find complicated model that allows stable oscillation in Prop T
60

---

Under review as a conference paper at ICLR 2023
M.2 IMPLICATIONS FROM LOW-DIMENSION TO HIGH DIMENSION
We would like to emphasize that, although our current simple settings are a little far from practical NNs, it still helps understand the ability of GD at large LRs to discover flat minima in three steps as follows. We include more experiments in Appendix B.2 to present the following hopes for complicated networks:
(a) By Theorem | especially its second condition, wish to discover intrinsic geometric property around local minima of more complicated models. The key is to investigate the 1-D function at the cross-section of the leading eigenvector and the loss landscape.
Theoretical: prove the 1-D condition holds at any minima for non-trivial matrix factorization, shown as Theorem c in Appendix A.2
Empirical: we show the 1-D condition holds around minima of 3,4,5-layer ReLU MLPs on MNIST, shown in Figure 6 d). 7 Bid) in Appendix B.2.2
(b) With the above intrinsic geometric property, the next question is whether the training trajectory utilizes this property.
Theoretical: in the case of quasi- symmmtii matrix factorization, we prove that the training trajectory follows the leading eigenvector of the Hessian (i.e. the leading component of Xo) in Theorem where the only top components weights are changing in (())-
* Empirical: for MLPs MNIST, we show the almost perfect alignment of the gradient and the top Hessian cigenvector in Figure 6 c). 7 [c), ().).
(c) The final implication the implicit bias of EoS after such oscillation. It turns out GD is driven flatter minima from sharper minima. In the 1-D case, obviously there nothing about implicit bias since the only thing GD is doing is to approximate the target value. However, an implicit bias from the oscillation appears starting from the 2-D case.
Theoretical 1: in the 2-D case in Theorem 5 we prove the two learnable parameters I, will converge to the same values after oscillations of their product zy. Actually in the minimum manifold, smaller x y means a flatter minimizer.
Theoretical 2: the single-neuron ReLU network in Theorem and Prop 2
we show the model degenerates to the 2-D case since wy - O. The 2-D argument tells that this nonlinear model also walks towards the balanced situation, verified with experiments
in Figure 2
Theoretical 3: in the quasi-symmetric MF in Theorem g although the initialization is around a sharp minima, GD is still driven towards the flattest minima where Fmax (Y)
m ZZ).
* Empirical 1: for 2-layer 16-neuron ReLU network in student-teacher setting, it turns out learning rate decay after beyond-Eos oscillations drives the model very close to the flattest minima, shown in Figure 5 and Appendix B.2.1
* Empirical 2: for 3,4,5-layer MLPs on MNIST, larger learning rate drives a flatter minima, as shown in Figure 6 b).
N CONCLUSIONS
In this work, we investigate gradient descent with large step size that crosses the threshold of local stability. In the low dimensional setting. we provide conditions on high-order derivatives that allow stable oscillation around local minima. For a two-layer single-neuron ReLU network, we prove its convergence to align with the teacher neuron under population loss. For matrix factorization, prove that the necessary 1 1-D condition holds around any minima. Furthermore, analysis of GD in symmetric matrix factorization, which converges to a period-2 orbit aligned with the 1-D convergence. Moreover, generalize the analysis to quami-symmetric cases where GD walks towards the flattest minimiser although initialized near sharp ones. further discussion is provided in Appendix(N)
conduct an
While these are encouraging results that contribute to the growing understanding of gradient descent beyond the Edge of Stability, our analysis suffers from important limitations that require further work.
61

---

Under review as a conference paper at ICLR 2023
An important item for future work is therefore extend it to general dimensions with nonlinearity,
which will enable the analysis of empirical landscapes as well as multiple neurons. Next. the understanding of the implicit bias of GD in the large-learning rate regime won't 't be complete without integrating the noise, either in the classic SGD sense or in the labels, as done Damian et al. (202T) : Li al. 202TT.
62


Reviews:
[{"review_id": "xdVQMvUyce", "reviewer": "Reviewer_WSV5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the dynamics of gradient descent with a large step size.  In particular, the authors are interested in the study beyond the edge of stability, i.e., the step-size regime above the inverse of the Lipschitz constant. \nIn that regime, local convergence is not guaranteed anymore. However, they show that under some additional assumptions on the higher derivative, one can actually show convergence to a cycle close to the optimum. They then apply the techniques to a one hidden neuron student teach problem and to a matrix factorization problem, showing that such a notion of convergence to a cycle can be obtained in that situation. \n", "summary_of_the_review": "This paper is providing results regarding an important question in our field. \nThe quality and clarity of the paper are okay, but they could be improved.\n\nMy main concern is the significance of the results which are very toyish to a certain extent. It is not clear that these results can be related to the behavior of neural network training", "strength_and_weaknesses": "## strenght: \n- This paper tries to tackle a complex and important problem: explaining the experimental evidence that GD convergence for NN training even with a step-size that is too large according to the optimization literature. \n- The presentation is relatively clear, and the (simple) experiments illustrate the theory. \n## Weaknesses:\n- The setting is simplistic. It is unclear whether such a setting captures the phenomenon that occurs in deep learning. \n- Lemma 1 and Theorem 1 seem very weak results, relatively disconnected from the rest of the paper. (the key 1D result seems to be Theorem 2).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "x7MZbVHWdmB", "reviewer": "Reviewer_Hb8V", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the (non)convergence of gd on general 1d function, $f(x)=\\frac{1}{4}(x^2-\\mu)^2$, two-layer single-neuron homogeneous network and matrix factorization, and shows a similar patten among then, i.e. GD converges to a period-2 orbit around the minimizer when step size slightly exceed the critical value for establishing local stability at the minimizer.", "summary_of_the_review": "This paper demonstrate EoS phenomenon in several specific examples. Some confusing notations, missing definitions and unclear writing affect the quality of the article, but I believe they could be fixed by a slight revision and I am willing to re-evaluate during review process.", "strength_and_weaknesses": "Strength:\n\nAlthough it has been known that GD on a function with non-uniform curvature (or specifically, step size guarantee local stability on some region but not for other region) could lead to period-2 orbit, e.g. in Ma et al. (2022), this paper provide a more detailed analysis on several specific examples.\n\nWeakness:\n\n1. I cannot verify the soundness of Theorem 4/14 due to confusing notation and missing definition.\n\nThe first conflict I observe is that authors define $\\Delta \\mathbf{Y}_t=\\mathbf{Y}_t-\\mathbf{Y}_0$ below eq.(218), which implies $\\Delta \\mathbf{Y}_0=0$, but in the statement of Theorem 4/14, authors require $u_1^{\\top} \\Delta \\mathbf{Y}_0 v_1 \\neq0$.\n\nThe second difficulty is that $\\widetilde{\\mathbf{Y}}_t$, which shows up multiple times in the proof, is not defined.\n\nMaybe these problems are just due to some typos, and I am willing to re-check the proof if authors could fix it during review process.\n\n2. The writing need to be improved.\n\nTheorem 3 doesn't directly state anything about edge of stability. Although the oscillating behaviors is discussed later informally, it would be better to express them rigorously in a theorem.\n\n$\\sigma_1$ is used in Section 6 without definition. Although in appendix J.2, authors relate it with SVD, it is better to give an explicit definition in the main paper.\n\n> we study non-asymptotic properties of GD beyond EoS\n\nIt is not clear to me why this paper is \"beyond\" EoS.\n\n3. The EoS in this paper is established by carefully selecting a step size that slightly break the local stability at the flattest minimizer. This is not how EoS works in DNN, where progressive sharpening drives curvature up until local stability is violated. Thus it is not clear how the results in this paper could contribute to the understanding of EoS in deep learning.\n\n4. Figure 2 right shows oscillating w_y. However, according to Theorem 3, $w_y^{(t)}$ should exponentially decay to 0. Why they are different? Is it a consequence of finite number of points instead of explicit expectation?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "s99wU00BC16", "reviewer": "Reviewer_zoGM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work investigates gradient descent with a step size larger than the $\\frac{2}{L}$ predicted by the standard theory, a regime known as the edge of stability. The authors provided sufficient conditions for ensuring an oscillatory (instead of diverging) behavior around the local minima for 1-dimensional functions. For higher-dimensional problems, the authors considered two scenarios: a two-layer single-neuron ReLU network in a teacher-student setup, and quasi-symmetric matrix factorization. The main results are that these objectives can still converge or stabilize for some step size $\\eta = \\frac{2K}{L}$, $K>1$.\n\n", "summary_of_the_review": "This work studies the EoS phenomenon and provides several interesting examples. However, the lack of a global picture, clarity of the presentation, and mathematical rigor makes me feel that a major revision is required before accepting it. ", "strength_and_weaknesses": "Strength:\n\n- The authors provided many intersting observations on specific examples.\n\nWeaknesses:\n\n- A major issue is that the work seems like collecting a number of scattered examples where GD can converge with step size larger than $\\frac{1}{L}$, without any coherent observation.\n\n- The results on the 1-diemsnional functions are qualitatively different from the high-dimensional ones: In Section 4 (barring the special case of Theorem 2), the statements are \"There **exists** a large step size such that XXX happens\", whereas for higher-dimensional ones, they change to \"**For all** step size in this region, XXX happens.\" Moreover, the existence proofs in Section 4 produce step sizes that are of Lebesgue measure 0, which shows that it is almost impossible for any practical algorithm to satisfy such conditions. In this sense, contrary to what the authors claimed, there is a disconnection between the low-dimensional and high-dimensional analysis.\n\n- All the proofs, especially that of Theorem 3, rely on tedious and highly-specific calculations that do not reveal much about the general behavior of GD.\n\n- The presentation of Section 5 is quite opaque. Can the authors explain what is meant by *\"$\\tilde{w}_\\perp$ as the unit-length orthogonal residual of $w$ after projecting onto $\\tilde{w}$\"*? I understood it as the normalized version of $w- \\text{proj}_\\tilde{w} w$, but then this would violate the followup statement *\"it is clear that updates of $w$ always stay in the plane spanned by $\\tilde{w}$ and $w^{(0)}$\"* since $\\tilde{w}_\\perp$ is not in the span of these two vectors. In addition, unless I missed something, the meaning of Theorem 3 is not clearly stated: it states that the orthogonal component of $w_t$ will decrease to 0, but it does not prove convergence (i.e., $vw\\rightarrow \\tilde{w}$). Can the authors provide a more clear interpretation than *\"... the dynamics of the single-neuron ReLU network is getting closer to the 2-D case in Section A.1\"*?\n\n- Theorem 5 in Appendix A is not correct. The authors seem to have assumed that $xy > \\mu$ always happens in the GD process, which is not guaranteed. If the authors wish to investigate the case of only those iterates satisfying $x^{(t)}y^{(t)} > \\mu$, then one cannot assume an GD update from $x^{(t)}y^{(t)}$ to $x^{(t+1)}y^{(t+1)}, which the authors implicitly did.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VFZSjHB1WLa", "reviewer": "Reviewer_RMoF", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper is concerned with gradient descent beyond the so-called Edge of Stability (EoS). That is, the step size is large than what standard theory would suggest (using the Lipschitz constant of the loss function).  The EoS is interesting to the (theoretical) deep learning community because very often neural networks are trained in this regime.\n\nThis paper studies this regime by focusing on the following three simple settings.\n\n1. One-dimensional functions\n2. Learning a single neutron\n3. Non-convex matrix factorization\n\nSeveral results are established. For example: In Setting 1 it is shown that there exists a step size in the EoS regime such that gradient descent converges to an 2-periodic orbit. In Setting 2 it is shown that even for large learning rate, the single neuron can be learned. In the symmetric matrix factorization setting, it is shown that the gradient descent iterates converge to a 2-periodic orbit.\n\nMoreover, there are several numerical experiments provide which support the theoretical findings.\n", "summary_of_the_review": "It is not clear how the theoretical results can amplify our current understand of the EoS regime in deep learning practice. Moreover, the paper misses a clear storyline and the quality of writing should be improved in my opinion. For this reason,  I cannot recommend acceptance.\n\n--------------------------\n\nI have read all the reviews and the rebuttals by the authors. While the changes made by the authors improve the paper in my opinion, I am still not convinced that the toy examples can give sufficient into the high-dimensional phenomena which appear in modern machine learning. Also the storyline still needs to be more coherent. (Although this would rather amount to a major revision and conference rebuttals are not for major revisions of papers.)\n\nThus, I can only increase my score to 5.", "strength_and_weaknesses": "Strengths: The topic is very interesting and important to the theoretical deep learning community and has not been addressed properly. The paper contains several interesting observations. The numerical experiments support the theoretical findings and are presented in a clear manner.\n\nWeaknesses: \n1. The paper misses a clear storyline. It feels that the paper presents several examples in the EoS regime but it is not clear how they are connected.\n2. It is not clear how the findings or the setting under consideration is important to understand deep learning theoretically. The high-learning rate is often used in the early stages of neural network training and then the step size is decreased. It is not at all clear to me how this connects to the findings in the paper which seem rather to appear very locally, i.e. in the later stages of training.\n\nFurthermore, I have the following list of specific comments.\n\n1. page 4: \"$x_0$ will back to $x_0$ in two steps\"; Can you make this statement more formal\n2. page 5, Lemma 1: \"stably oscillate\" has not been defined yet\n3. page 5, Lemma 1: \"except the $f''$\"; what is meant by that?\n4. page 6: $v^{(t)}=v^{(t)}$; Is this a typo?!\n5. page 6; proof sketch; I cannot find where $\\Delta w_y$ is defined. \n6. page 7: typo: \"...condition holds around (the) minimum\"\n7. page 7: \"period-$2$ orbit $\\gamma_\\eta$; I cannot find the definition of a period-$2$ orbit\n8. page 7: What is $Z_0$ in eq. (11)?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject