O N
G RADIENT
D ESCENT C ONVERGENCE BEYOND
THE
E DGE OF S TABILITY
Anonymous authors
Paper under double-blind review
A BSTRACT
Gradient Descent (GD) is a powerful workhorse of modern machine learning thanks to its scalability and efﬁciency in high-dimensional spaces.
Its ability to ﬁnd local minimisers is only guaranteed for losses with Lipschitz gradients, where it can be seen as a ‘bona-ﬁde’ discretisation of an underlying gradient ﬂow.
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.
Speciﬁcally, 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 nonasymptotic ‘Local Implicit Bias’ of GD above the EoS, whereby quasi-symmetric initializations converge to symmetric solutions — where sharpness is minimum amongst all minimisers.
I NTRODUCTION
Given a differentiable objective function f ( θ ) , where θ ∈ R
d is a high-dimensional parameter vector, the most basic and widely used optimization method is gradient descent (GD), deﬁned as
[Equation]
(1) where η is the learning rate.
For all its widespread application across many different ML setups, a 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 η ?
In the modern context of large-scale ML applications, an additional key question is not only to understand whether or not GD converges to minimisers, but to which ones, since overparametrisation deﬁnes a whole manifold of global minimisers, all potentially enjoying drastically different generalisation performance.
The sensible regime to start the analysis is η → 0 , where GD inherits the local convergence properties of the Gradient Flow ODE via standard arguments from numerical integration.
However, in the early phase of training, a large learning rate has been observed to result in better generalization (LeCun et al., 2012; Bjorck et al., 2018; Jiang et al., 2019;
Jastrzebski et al., 2021), where the extent of “large” is measured by comparing the learning rate η and the curvature of the loss landscape, measured with λ ( θ ) :
= λ max (cid:2) ∇ 2 θ f ( θ ) (cid:3) , the largest eigenvalue of the Hessian with respect to learnable parameters.
Although one requires sup θ λ ( θ ) < 2 /η to guarantee the convergence of GD (Bottou et al., 2018) to (local) minimisers 1 , the work of (Cohen et al., 2020) noticed a remarkable phenomena in the context of neural network training: even in problems where λ ( θ ) is unbounded (as in NNs), for a ﬁxed η , the curvature λ ( θ ( t ) ) increases along the training trajectory (1), bringing λ ( θ ( t ) ) ≥
/η (Cohen et al., 2020).
After that, a surprising phenomena is that λ ( θ ( t ) ) stably hovers above 2 /η 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 a 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 networks.
In this setting, GD is expected to converge to a 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 a continuous sharpnessreduction ﬂow, 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 a range of learning rates η above the EoS such that GD dynamics hover around minimisers.
Moreover, in the matrix factorization setup, where minimisers form a manifold with varying local curvature, our results give a non-asymptotic analogue of the ‘Sharpness-Minimisation’ arguments from Arora et al.
; Lyu et al. (2022);
Damian et al.
The straightforward starting point for the local convergence analysis is via Taylor approximations of the loss function.
However, in a quadratic Taylor expansion, gradient descent diverges once λ ( θ ) >
2 /η (Cohen et al., 2020), indicating that a higher order Taylor approximation is required.
By considering a 1-D function with local minima θ ∗ of curvature λ ∗ = λ ( θ ∗ ) , we show that it is possible to stably oscillate around the minima with η slightly above the threshold 2 /λ ∗ , provided its high order derivative satisﬁes mild conditions as in Theorem 1.
A typical example of such functions is f ( x ) =
14 ( x 2 − µ ) 2 with µ > 0 .
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 sufﬁciently 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 quasi-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.
R ELATED W ORK
Implicit regularization.
Due to its theoretical closeness to gradient descent with a small learning rate, gradient ﬂow is a common setting to study the training behavior of neural networks.
Barrett & Dherin (2020) suggests that gradient descent is closer to gradient ﬂow with an additional term regularizing the norm of gradients.
Through analysing the numerical error of Euler’s method, Elkabetz & Cohen (2021) provides theoretical guarantees of a small gap depending on the convexity along the training trajectory.
Neither of them ﬁts in the case of our interest, because it is hard to track the parametric gap when η > 1 /λ .
For instance, in a quadratic function, the trajectory jumps between the two sides once η > 1 /λ .
Damian et al.
(2021) shows that SGD with label noise is implicitly subjected to 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 ﬂow automatically preserves the norms’ differences between different layers of a deep homogeneous network.
(Ye & Du, 2021) shows that gradient descent on matrix factorization with a constant small learning rate still enjoys the auto-balancing property.
Also in matrix factorization, Wang et al.
(2021) proves that gradient descent with a relatively large learning rate leads to a solution with a more balanced (perhaps not perfectly balanced) solution while the initialization can be in-balanced.
In a similar spirit, we extend their ﬁnding to a 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., 2021), where they show GD with large learning rates converges to a ﬂat region in the interpolation manifold while the ﬂat region w.r.t.
our larger learing rate does not exists so GD is forced to wander around the ﬂattest minima.
Note that the implication of balancing effect is to get close to a ﬂatter solution in the global minimum manifold, which may help improve generalization in some common arguments in the community.
Edge of stability.
Cohen et al.
(2020) observes a two-stage process in gradient descent, where the ﬁrst is loss curvature grows until the sharpness touches the bound 2 /η , and the second is the curvature hovers around the bound and training loss still decreases in a 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 ﬂatter region if loss landscape around initialization is too sharp.
Some concurrent works (Ahn et al., 2022; Ma et al., 2022; Arora et al., 2022; Damian et al., 2022) are also theoretically investigating the edge of stability.
Ahn et al.
(2022) suggests that unstable convergence happens when the loss landscape of neural networks forms a 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 an assumption, shows that gradient descent with different learning rates may stay in different levels.
Arora et al.
(2022) shows the training provably enters the edge of stability with modiﬁed gradient descent or modiﬁed loss, and then its associated ﬂow goes to ﬂat regions.
Under mild conditions, Damian et al.
(2022) proves that GD beyond EoS follows an optimization trajectory subjected to a sharpness constraint so that a ﬂatter region is found.
Learning a single neuron.
Yehudai & Ohad (2020) studies necessary conditions on both the distribution and activation functions to guarantee a one-layer single student neuron aligning with the teacher neuron under gradient descent, SGD and gradient ﬂow.
Vardi et al.
(2021) extends the investigation into a neuron with a 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 a large learning rate beyond the edge of stability.
P ROBLEM S ETUP
We consider a differentiable objective function f ( θ ) with θ ∈ R d , and the GD algorithm from (1).
Deﬁnition 1.
A differentiable function f is L -gradient Lipschitz if
[Equation]
The above deﬁnition is equivalent to saying that the spectral norm of the Hessian is bounded by L , or the local curvature at each point is bounded by L .
Then η needs to be bounded by 1 /L in GD so that it is guaranteed to visit an approximate ﬁrst-order stationary point (Nesterov, 1998).
The perturbed
GD requires η = 1 /L to visit an approximate second-order stationary point (Jin et al., 2021), and stochastic variants share similar assumptions (Ghadimi & Lan, 2013; Jin et al., 2021).
However, in practice, such an assumption may be violated, or even impossible to satisfy when (cid:107)∇ 2 f (cid:107) is not uniformly bounded.
Cohen et al.
(2020) observes that, with learning rate η ﬁxed, the largest eigenvalue λ 1 of the loss Hessian of a neural network is below 2 /η at initialization, but it grows above the threshold along training.
Such a 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 λ 1 ≥ 2 /η , the training loss starts to perturb sharply.
This is not surprising because GD would diverge in a quadratic function with such a 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 /η .
A further phenomena is that, when GD is at the edge of stability, if the learning rate suddenly changes to a smaller value η s < η , then the local curvature quickly grows to 2 /η s — 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 a lot of attention.
Due to the nature of over-parameterization in modern neural networks, the global minimizers of the objective f form a manifold of solutions.
There have been active directions to understand the implicit 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 ﬂatter minima (Barrett & Dherin, 2020; Smith et al., 2021; Damian et al., 2021; Ma & Ying, 2021).
For the latter, ﬂatter minima brings better generalization (Hochreiter & Schmidhuber, 1997;
Li et al., 2018; Keskar et al., 2016; Ma & Ying, 2021; Ding et al., 2022).
It would be meaningful if ﬂatter minima could be obtained via GD with a large learning rate.
More speciﬁcally, it has been shown that the eigenvalues of the hessian of a deep homogeneous network could be manipulated to inﬁnity via rescaling the weights of each layer (Elkabetz & Cohen, 2021).
Fortunately, gradient ﬂow preserves the difference of norms across layers along the training (Du et al., 2018).
As a result, a balanced initialization induces balanced convergence, while GD would break this balancing effect due to ﬁnite learning rate.
However, recently it has been observed that GD with large learning rates enjoys a balancing effect (Wang et al., 2021), where it converges to a (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 a large learning rate, from low-dimensional to representative models.
S TABLE OSCILLATION ON 1-D FUNCTIONS
Deﬁnition 2.
(Period-2 stable oscillation.)
Consider GD on a function f in domain Ω .
Denote the update rule of GD as F ( x ) for x ∈ Ω .
A period-2 stable oscillation is ∃ x ∈ Ω such that
F ( F ( x ))
=
x and x is not a minima of f .
We initiate our analysis of the stable oscillation phenomenon in 1-D. Starting from a condition on general 1-D functions, we look into several speciﬁc 1-D functions to verify our arguments.
Then, focusing on a function in the form of f ( x ) =
( x 2 − µ ) 2 , we present the convergence analysis as a foundation for the following discussions.
Furthermore, to shed light on the multi-layer setting, we 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 a learnable parameter x ∈ R .
The parameter updates following GD with the learning rate η as
[Equation]
Assuming f is differentiable and all derivatives are bounded, the function value in the next step can be approximated by
[Equation]
If η < 2 /f (cid:48)(cid:48)
( x ( t ) ) , this approximation reveals that the function monotonically decreases for each step of GD, ignoring higher terms.
Such an assumption would guarantee the convergence to a global minimum in a convex function.
However, our interest is what happens if η > 2 /f (cid:48)(cid:48) ( x ) .
For instance, if f is a quadratic function, the second-order derivative f (cid:48)(cid:48) is constant.
As a result, once η > 2 /f (cid:48)(cid:48) ,
GD diverges except when being initialized at the optimum.
However, when trained with a large learning rate η > 2 /f (cid:48)(cid:48) (¯ x ) , there is still some hope for a function to stay around a local minima ¯ x , as stated in the following theorem.
Theorem 1.
Consider any 1-D differentiable function f ( x ) around a local minima ¯ x , satisfying (i) f (3) (¯ x ) (cid:54) = 0 , and (ii)
[ f (3) ] 2 − f (cid:48)(cid:48) f
(4) > 0 at ¯ x .
Then, there exists (cid:15) with sufﬁciently small | (cid:15) | and (cid:15) · f (3) > 0 such that: for any point x 0 between ¯ x and ¯ x − (cid:15) , there exists a learning rate η such that the update rule F η of GD satisﬁes F η ( F η ( x 0 ))
= x 0 , and
[Equation]
The details of proof are presented in the Appendix C. As stated in the Theorem 1, we provide a 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) (¯ x ) = 0 .
For instance, a quadratic function does not satisfy this condition since f (3) =
f (4) ≡
and it diverges when GD is beyond the edge of stability.
For f ( x ) = sin( x ) around ¯ x = − π 2 where f (3) (¯ x ) = 0 , it turns out the sine function allows stable oscillation.
Therefore, we extend the argument in Theorem 1 to a higher order case in Lemma 1.
Lemma 1.
Consider any 1-D differentiable function f ( x ) around a local minima ¯ x , satisfying that the lowest order non-zero derivative (except the f (cid:48)(cid:48) ) at ¯ x is f ( k ) (¯ x ) with k ≥ 4 .
Then, there exists (cid:15) with sufﬁciently small | (cid:15) | such that: for any point x 0 between ¯ x and ¯ x − (cid:15) , and 1.
if k is odd and
(cid:15) · f ( k )
¯ x
>
f
k +
) (¯ x ) < 0 , then there exists
η ∈
[Equation]
f (cid:48)(cid:48)
[Equation]
if k is even and f ( k )
(¯
x )
<
, then there exists
η ∈ ( 2 f (cid:48)(cid:48)
[Equation]
such that: the update rule F η of GD satisﬁes F η ( F η ( x 0 )) =
x 0 .
With Lemma 1, we can verify the sine function to allow stable oscillation as in Corollary 1, because its lowest order of nonzero derivative (except f (cid:48)(cid:48) ) at the local minima is f (4) (¯ x ) < 0 .
Meanwhile,
Theorem 1 provides a guarantee that squared-loss on any function g provably allows stable oscillation once g satisﬁes some mild conditions, as stated below.
Lemma 2.
Consider a 1-D function g ( x ) , and deﬁne the loss function f as f ( x )
=
( g ( x ) − y ) 2 .
Assuming (i) g (cid:48) is not zero when g (¯ x ) =
y , (ii) g (cid:48) (¯ x ) g (3) (¯ x ) <
[ g (cid:48)(cid:48) (¯ x )] 2 , then it satisﬁes the condition in Theorem 1 or Lemma 1 to allow period-2 stable oscillation around ¯ x .
This setup covers generic 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 as Corollaries 1-8 in Appendix D and E. Moreover, we provide 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 1-D functions p, q .
Assume both p ( x ) , q ( y ) at x = ¯ x, y = p (¯ x ) satisﬁes the conditions of g in Lemma 2 to allow stable oscillations.
Then q ( p ( x )) allows stable oscillation around x = ¯ x .
Proof details of the above lemmas and proposition are presented in the Appendix D and
E. Next we are going to present a careful analysis on g ( x ) = x 2 .
A special 1-D function.
Consider f ( x ) = 14 ( x 2 − µ ) 2 with µ > 0 , f (3) ( √ µ ) = 6 √ µ, f (cid:48)(cid:48) ( √ µ ) = 2 µ .
Note that this function is more special to us because it can be viewed as a symmetric scalar factorization problem subjected to the squared loss.
Later we will leverage it to gain insights for asymmetric initialization, two-layer single-neuron networks and matrix factorization.
Before that, we would like to show where it converges to when η >
[Equation]
as follows.
Theorem 2.
For f ( x ) = 14 ( x 2 − µ ) 2 , consider GD with η = K · 1 µ where 1 < K < √ 4 .
5 − 1 ≈ 1 . 121 , and initialized on any point 0 < x 0 < √ µ .
Then it converges to an orbit of period 2, except for a measure-zero initialization where it converges to √ µ .
More precisely, the period-2 orbit are the solutions x =
[Equation]
1 ∈ (0 ,
[Equation]
, x =
[Equation] [Equation]
) of solving δ in
[Equation]
The details of proof are presented in the Appendix F. As shown above, Theorem 1 and Theoremstand in two different levels: Theorem 1 restricts the discussion in a local view because of Taylor approximation, while Theorem 2 starts from local convergence and then generalizes it into a global view.
However, Theorem 1 builds a foundation for Theorem 2 because the latter would degenerate to the former when K is extremely close to 1.
A natural follow-up question is what implications Theorem 2 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 5, and high dimensions in Section 6.
O N A TWO - LAYER SINGLE - NEURON HOMOGENEOUS NETWORK
We denote a two-layer single-neuron network as f ( x ; θ ) = v · σ ( w (cid:62) x ) where v ∈ R , w ∈ R d , the set of trained parameters θ = ( v, w (cid:62) )
∈
R
d +1 , and the nonlinearity σ is ReLU.
We will keep such an order in θ to view it as a vector.
The input x ∈ R d is drawn uniformly from a unit sphere S d − 1 .
The parameters are trained by GD subjected to L 2 population loss, as
[Equation]
We generate labels from a single teacher neuron function, as y | x
= σ ( ˜ w (cid:62) x ) .
Hence ˜ w is our target neuron to learn.
We denote the angle between w and ˜ w as α ≥ 0 .
Note that α is set as non-negative because the loss function is symmetric w.r.t.
the angle.
Moreover, the rotational symmetry of the population data distribution results in a loss landscape that only depends on w through the angle α and the norm (cid:107) w (cid:107) .
Indeed, from the deﬁnition, we have
[Equation]
where we denote ˜ w ⊥ as the normalized of w − proj ˜ w w .
Consider the Hessian
[Equation]
Hence, in the global minima manifold where vw = ˜ w , the eigenvalues of the Hessian are λ
1 = (cid:107) w (cid:107)
+ v 2 d , λ 2 ...
d = v 2 d , λ d +1 = 0 .
Therefore, the largest eigenvalue λ 1 measures the imbalance ( i.e. , | (cid:107) w (cid:107) − v | ) between the two layers again as λ 1 =
( (cid:107) w (cid:107)− v )
2 +2 (cid:107) ˜ w (cid:107) d similar to the 2-D case in (15) in
Appendix A.1.
So we would like to investigate where GD converges if η >
[Equation]
that is too large even for the ﬂattest minima.
Note that a key difference between the current case and the previous 2-D analysis is that the current one includes a neuron as a vector and a nonlinear ReLU unit.
From the second row of ∇ θ L , which is ∇ w L , it is clear that updates of w always stay in the plane spanned by ˜ w and w (0) .
Hence, this problem can be simpliﬁed to three variables ( v, w x , w y ) with the target neuron ˜ w
=
The three variables stand for
[Equation]
We keep w y as nonnegative because the loss L is invariant to its sign and our previous notation α ≥requires a non-negative w y .
Then we show that w y decays to 0 as follows.
Theorem 3.
In the above setting, consider a teacher neuron ˜ w =
[1 , 0] and set the learning rate η = Kd with K ∈ (1 , 1 . 1] .
Initialize the student as (cid:13)(cid:13) w (0) (cid:13)(cid:13)
= v (0) (cid:44) (cid:15) ∈ (0 , 0 . 10] and (cid:104) w (0) , ˜ w (cid:105) ≥ 0 .
Then, for t ≥ T 1 + 4 , w ( t ) y decays as
[Equation]
Proof sketch The details of proof are presented in the Appendix I.
The proof is divided into two stages, depending on whether w y grows or not.
The key is that the change of w y follows (omitting all superscripts t )
[Equation]
where the second term in ∆ w y / w y is bounded in [0 , 12 π ] .
In stage 1 where vw x is relatively small, we show the growth ratio of w y is smaller than those of w x and vw x , resulting in an upper bound of number of iterations for vw x to reach 12 π , so max( w y ) is bounded too.
Although the initialization is balanced as v (0) =
(cid:13)(cid:13) w (0) (cid:13)(cid:13) for simplicity of proof, v − w x is also bounded at the end of stage 1.
From the beginning of stage 2, thanks to the relatively narrow range of K , we are able to compute the bounds of three variables (including v − w x , vw x and w y ) and they turn out to fall into a basin in the parameter space after four iterations.
In this basin, w y decays exponentially with a linear rate of 0.
at most.
[Equation]
With the guarantee of w y 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 a convergence result as follows.
Proposition 2.
The single-neuron model in Theorem 3 converges to a period-2 orbit where w y =and ( v, w x ) ∈ γ K with γ K = { ( δ 1 , δ 1 ) , ( δ 2 , δ 2 ) } .
Here δ 1 ∈ (0 , 1) , δ 2 ∈ (1 , 2) are the solutions δ in
[Equation]
Remark.
Actually this convergence is close to the ﬂattest minima because: if the learning rate decays to inﬁnitesimal after sufﬁcient oscillations, then the trajectory walks towards the ﬂattest minima ( v = w x = 1 , w y = 0) .
To summarize, the single-neuron model goes through three phases of training dynamics, with a intialization of the angle (cid:93) ( w, ˜ w ) as π 2 at most.
First, the angle decreases monotonically but, due to the growth of norms, the absolute deviation w y still increases.
Meanwhile, the inbalance v − w x stays in a bounded level.
Second, w y starts to decrease and the parameters fall into a basin within four steps.
Third, in the basin, w y decreases exponentially and, after w y at a reasonable low level, the model approximately follows the dynamic of the 2-D case and the inbalance v − w x decreases as well, following Theorem 5.
The model converges to a period-2 orbit as in the 1-D case in Theorem 2.
Q UASI - SYMMETRIC M ATRIX FACTORIZATION : WALKING TOWARDS FLATTEST MINIMA
Consider a matrix factorization problem, parameterized by learnable weights X ∈ R m × p , Y ∈ R q × p , and the target matrix is C ∈ R m × q .
The loss L is deﬁned as
[Equation]
Obviously { X , Y : XY (cid:62) = C } forms a minimum manifold.
In this context, the question is to describe GD dynamics in terms of a ‘descent’ phase (i.e., reaching the manifold), followed by a ‘hovering’ phase, where the dynamics evolve nearby the minimum manifold.
Although we prove that the necessary 1-D condition holds around minimum as Theorem 6 (in Appendix A.2), it is more attracting to investigate GD in high dimensions.
A straightforward subsets of the “ﬂattest" points in the manifold of minimisers are in fact given by symmetric matrices, i.e., points of the form ( X , X ) with XX (cid:62)
=
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 A.2) shows that the dynamics follows the direction of the leading eigenvector and then stably oscillates with a period-2 analogous to the 1D case in Theorem 2.
Note that, although { X : XX (cid:62) =
X 0 X
(cid:62) 0 } forms a manifold that contains inﬁnite number of minimizers, 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-sharpness is to rescale symmetric minimizers, leading to Quasi-symmetric Matrix Factorization.
Given a symmetric target C = X 0 X (cid:62) 0 , assume that we are around the (global) minima Y 1 = α X 0 + ∆ Y 1 , Z 1 = 1 α X 0 + ∆ Z 1 with α > 0 and small deviation (cid:107) ∆ Y 1 (cid:107) , (cid:107) ∆ Z 1 (cid:107) ≤ (cid:15) .
The top singular value and vectors in SVD of X 0 is σ 1 u 1 v (cid:62) 1 .
Then the EoS-learning rate at ( α X 0 , 1 α X 0 ) is 2 σ 21 ( α 2 + 1 α 2 ) , which is largest as 1 σ 21 at α = 1 .
We study the convergence of GD starting from Y 1 = α X 0 + ∆ Y 1 , Z 1 = 1 α X 0 + ∆ Z 1 with learning rate η = 1 σ 21 + β, β > 0 .
The following theorem shows that, although starting nearby a sharper minima,
GD still converges to and stably scillate around the ﬂattest one.
Theorem 4.
Consider the above quasi-symmetric matrix factorization with learning rate η = 1 σ 21 + β .
Assume 0 < βσ 21 < √ 4 .
5 − 1 ≈ 0 .
Consider a minimum ( Y 0 = α X 0 , Z 0 = 1 / α X 0 ) , α > 0 .
The initialization is around the minimum, as Y
=
Y 0 + ∆ Y 1 , Z 1 =
Z 0 + ∆ Z 1 , with the deviations satisfying u (cid:62) 1 ∆ Y 1 v 1 (cid:54) = 0 , u (cid:62) 1 ∆ Z 1 v 1 (cid:54) = 0 and (cid:107) ∆ Y 1 (cid:107) , (cid:107)
∆ Z 1 (cid:107) ≤ (cid:15) .
The second largest singular value of X 0 needs to satisfy
[Equation]
Then GD would converge to a period-2 orbit γ η approximately with error in O ( (cid:15) ) , formally written as
[Equation]
where ρ 1 ∈ (1 , 2) , ρ 2 ∈ (0 , 1) are the two solutions of solving ρ in
[Equation]
Proof sketch Details of proof can be found in Appendix J.3, which shares a similar spirit with Theorem 7.
The analysis consists of two phases, depending on whether (cid:15)
y,t (cid:44) (cid:104)
Y t − Y 0 , u 1 v (cid:62) 1 (cid:105) , (cid:15) z,t (cid:44) (cid:104)
Z t − Z 0 , u 1 v (cid:62) 1 (cid:105) are small or not.
In Phase I, all components of Y t − Y 0 and Z t − Z 0 are small due to the initialization near minima, but both (cid:15) y,t and (cid:15) z,t are growing exponentially in a rate of ησ
α 2 + η σ 21 α 2 − 1 ≥ 2 ησ 21 − 1 > 1 .
In Phase II, both (cid:15) y,t and (cid:15) z,t are much larger than other components, as long as other components are still not growing.
So the dynamics of them matches GD of 2-D function f
( y, z ) =
12 ( yz − 1) 2 with learning rate η (cid:48)
= 1+ βσ 21 .
Following the 2-D analysis in
Theorem 5, we have (cid:15) y,t and (cid:15) z,t converge to the same values, which degenerates the 2-D problem to 1-D function.
Therefore, the proof concludes with 1-D convergence analysis of f ( x ) =
14 ( x 2 − 1)as shown in Theorem 2.
(cid:3)
Remark.
Note that both Y 0 − α · σ 1 u 1 v (cid:62) 1 , Z 0 − 1 / α · σ 1 u 1 v (cid:62) 1 are residuals of Y 0 , Z 0 with the top singular value eliminated.
Then, compared with Theorem 7, we have ρ i corresponds to δ i
+ 1 , which means both symmetric and quasi-symmetric cases converge to parameters with the same top singular values and wander around the ﬂattest minima.
In other words, this convergence is close to the ﬂattest minima because: if the learning rate decays to inﬁnitesimal after sufﬁcient oscillations, then the trajectory walks towards the ﬂattest minima approximately with parameter distance in O ( (cid:15) ) .
Also note that, if η < 1 σ 21 , we anticipate it still escapes from the sharp minima and converges to a ﬂatter one (not necessarily the ﬂattest).
The result could be obtained by tracking GD on f ( x, y )
=
12 ( xy − 1)with η < 1 slightly.
But the closed form can not be expressed explicitly, because it strongly depends on initialization.
N UMERICAL 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 B.
1-D functions.
As discussed in the Section 4, we have f ( x ) =
14 ( x 2 − 1) 2 satisfying the condition in Theorem 1 and g ( x ) = 2 sin( x ) satisfying Lemma 1, so we estimate that both f and g allow stable oscillation around the local minima.
It turns out GD stably oscillates around the local minima on both functions, when η >
2 f
[Equation]
(¯ x
) slightly, as shown in Figure 1.
Two-layer single-neuron model.
As discussed in the Section 5, with a learning rate η ∈ ( d, 1 . 1 d ] , a single-neuron network f ( x ) = v · σ ( w (cid:62) 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 a sphere S 1 , as shown in Figure 2.
The student neuron is initialized orthogonal to the teacher
[Figure]
neuron.
In the end of training, w y decays to a small value before the inbalance | v − w x | decays sharply, which veriﬁes our argument in Section 5.
With a small w y , this nonlinear problem degenerates to a 2-D problem on v, w x .
Then, the balanced property makes it align with the 1-D problem where v and w x converge to a period-2 orbit.
Note that the small residuals of | v − w x | and w y are due to the difference between population loss and empirical loss
[Figure]
uniformly sampled from sphere S 1 of (cid:107) x (cid:107) =
The teacher neuron is ˜ w =
[1 , 0] and the student neuron is initialized as w (0) =
[0 , 0 . 1] with v (0) = 0 . 1 .
Symmetric and quasi-symmetric matrix factorization.
As discussed in the Section 6 and Appendix A.2, with mild assumptions, both symmetric and quasi-symmetric cases stably wanders around the ﬂattest minima.
We train GD on a matrix factorization problem with X 0 X (cid:62)
= C ∈ R 8 × 8 .
The learning rate is 1 .
02 × EoS threshold.
Following the setting in Section 6, for symmetric case, the training starts near X 0 and, for quasi-symmetric case, it starts near ( α X 0 , 1 / α X 0 ) with α = 0 . 8 , as shown in Figure 3.
Although starting with a re-scaling, the quasi-symmetric case achieves the same top singular values in Y and Z , which veriﬁes the balancing effect of 2-D functions in Theorem 5.
Then, the top singular values of both cases converge to the same period-2 orbit, supported by
Theorem 2, 4 and 7.
[Figure]
Appendices
A Additional Results
A.1 On a 2-D function . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.2 On Matrix Factorization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B Additional Experiments
B.1 2-D function . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.2 High dimension and MNIST . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C Proof of Theorem
D Proof of Lemma
E Proof of Lemma
F Proof of Theorem
G
Proof
of
Theorem
H Proof of Lemma
I Proof of Theorem
J Proof of Matrix Factorization
J.
Asymmetric Case: 1D function at the minima . . . . . . . . . . . . . . . . . . . .
J.
Symmetric Case: convergence analysis around the minima . . . . . . . . . . . . .
J.
Quasi-symmetric case: walk towards ﬂattest minima . . . . . . . . . . . . . . . .
K Useful lemmas
L Illustration of period-2 and period-4 orbits
M Discussions
M.1 Connections between theoretical results . . . . . . . . . . . . . . . . . . . . . . .
M.2 Implications from low-dimension to high-dimension . . . . . . . . . . . . . . . .
N
Conclusions
A A DDITIONAL R ESULTS A. O N A 2-D FUNCTION
Similar to f ( x ) =
14 ( x 2 − µ ) 2 , consider a 2-D function f ( x, y )
= 12 ( xy − µ ) 2 .
Apparently, if x and y initialize as the same, then ( x ( t ) , y ( t ) ) would always align with the 1-D case from the same initialization.
Therefore, it is signiﬁcant to analyze this problem under different initialization for x and y , which we would call “in-balanced” initialization.
Meanwhile, another giant difference is that all the global minima in 2-D case form a manifold { ( x, y ) | xy = µ } 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 inﬁnity as in (Elkabetz & Cohen, 2021), so the inbalance strongly affects the local curvature.
Viewing f
( x ) as a symmetric scalar factorization problem, we treat f ( x, y ) as asymmetric scalar factorization.
The update rule of
GD is
[Equation] [Equation]
Consider the Hessian as
[Equation]
When xy = µ , the eigenvalues of H are λ 1 = x 2 + y 2 , λ
2 = 0 .
Note that λ 1 = ( x − y ) 2 + 2 µ .
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 λ 1 = 2 µ when x = y = √ µ .
In other words, if the learning rate η > 2 / 2 µ , 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 to a perfect balance although they initialized differently, as follows.
Theorem 5.
For f ( x, y ) =
12 ( xy − µ ) 2 , consider GD with learning rate η = K · 1 µ .
Assume both x
and y are always positive during the whole process { x i , y i } i ≥ 0 .
In this process, denote a series of all points with xy > µ as P = { ( x i , y i )
| x
i y i > µ } .
Then | x − y | decays to 0 in P , for any 1 < K < 1 . 5 .
Proof sketch The details of proof are presented in the Appendix G. Start from a point ( x ( t ) , y ( t ) ) where x ( t ) y ( t ) > µ .
Because y ( t +1) − x ( t +1) =
( y (
[Equation]
) − x ( t ) )(1+ η ( x ( t ) y
( t ) − µ )) , it sufﬁces to show
[Equation]
Since 1 + η ( x ( t ) y ( t ) − µ ) > 1 , the analysis of 1 + η ( x ( t +1)
y ( t +1) − µ ) is divided into three cases considering the coupling of ( x ( t ) , y ( t ) ) , ( x ( t +1) , y ( t +1) ) .
[Equation]
Remark.
Actually, for a larger K ≥ 1 . 5 , it is possible for GD to converge to an inbalanced orbit.
For instance, Figure 15 in (Wang et al., 2021) shows inbalanced orbits for f ( x ) = 1 2 ( xy − 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 5 reveals x and y would converge to a perfect balance.
Note that this balancing effect is different from that of gradient ﬂow (Du et al., 2018), where the latter states that gradient ﬂow preserves the difference of norms of different layers along training.
As a result, in gradient ﬂow, inbalanced initialization induces inbalanced convergence, while in our case inbalanced-initialized weights converge to a perfect balance.
Furthermore, Theoremshows 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 3.
Numerical results are presented in Figure 4.
Proposition 3.
Following the setting in Theorem 5.
Further assume 1 < K < √ 4 .
5 − 1 ≈ 1 . 121 .
Then GD converges to an orbit of period 2.
The orbit is formally written as { ( x = y = δ i ) | i
= 1 , 2 } , with δ 1 ∈ (0 , √ µ ) , δ 2 ∈ ( √ µ, 2 √ µ ) as the solutions of solving δ in η = δ
[Equation]
(cid:16)(cid:113)
[Equation]
(cid:17) .
Remark.
Actually this convergence is close to the ﬂattest minima because: if the learning rate decays to inﬁnitesimal after sufﬁcient oscillations, then the trajectory walks towards the ﬂattest minima.
However, one thing to notice is that the inbalance at initialization needs to be bounded in Theorembecause both x
and y are assumed to stay positive along the training.
More precisely, we have
[Equation]
and then x ( t +1)
y ( t +1) < 0 when | x ( t ) − y ( t ) | is large with x ( t ) y ( t ) > µ ﬁxed.
Therefore, we provide a condition to guarantee both x, y positive as follows, with details presented in the Appendix H.
Lemma 3.
In the setting of Theorem 5, denote the initialization as m = | y
− x 0 | √ µ and x 0 y 0 > µ .
Then, during the whole process, both x and y will always stay positive, denoting p =
[Equation]
and q = (1 + p ) 2 , if
[Equation]
A. O N M ATRIX F ACTORIZATION
In this section, we present two additional results of matrix factorization.
A.2. A SYMMETRIC C ASE : 1D FUNCTION AT THE MINIMA
Before looking into the theorem, we would like to clarify the deﬁnition of the loss Hessian.
Inherently, we squeeze X , Y into a vector θ = vec ( X , Y ) ∈ R mp + pq , which vectorizes the concatnation.
As a result, we are able to represent the loss Hessian w.r.t.
θ as a matrix in R ( mp + pq ) × ( mp + pq ) .
Meanwhile, the support of the loss landscape is in R mp + pq .
Similarly, we use (∆ X , ∆ Y ) in the same shape of ( X , Y ) to denote .
In the following theorem, we are to show the leading eigenvector ∆ (cid:44) vec (∆ X ,
∆ Y ) ∈
R mp + pq of the loss Hessian.
Since the cross section of the loss landscape and ∆ forms a 1D function f ∆ , we would also show the stable-oscillation condition on 1D function holds at the minima of f ∆ .
Theorem 6.
For a matrix factorization problem, assume XY = C .
Consider SVD of both matrices as X = (cid:80)
[Equation]
σ
[Equation]
[Equation] [Equation]
and
[Equation]
σ y,i u y,i v (cid:62) y,i , where both groups of σ · ,i ’s are in descending order and both top singular values σ x, 1 and σ y, 1 are unique.
Also assume v (cid:62) x, 1 u y, 1 (cid:54) = 0 .
Then the leading eigenvector of the loss Hessian is ∆ = vec ( C 1 u x, 1 u (cid:62) y, 1 , C 2 v x, 1 v (cid:62) y, 1 ) with C 1 =
[Equation]
[Equation] [Equation]
Denote f ∆ as the 1D function at the cross section of the loss landscape and the line following the direction of ∆ passing vec (∆ X , ∆ Y ) .
Then, at the minima of f ∆ , it satisﬁes
[Equation]
The proof is provided in Appendix J.1.
This theorem aims to generalize our 1-D analysis into higher dimension, and it turns out the 1-D condition is sastisﬁed around any minima for two-layer matrix factorization.
In Theorem 1 and Lemma 1, if such 1-D condition holds, there must exist a periodorbit 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 a convergence analysis with initialization near the minima, but in any direction instead.
A.2. S YMMETRIC C ASE : CONVERGENCE ANALYSIS AROUND THE MINIMA
In this section, we focus on the symmetric case of matrix factorization where Y =
X (cid:62) .
Accordingly, we rescale the loss function as L ( X , X )
=
1 4 (cid:13)(cid:13)
XX (cid:62) − C
(cid:13)(cid:13) 2 F .
Denote the target as C = X 0 X (cid:62) 0 , and assume we are around the minima X 1 =
X 0 + ∆ X 1 with small (cid:107) ∆ X 1 (cid:107) ≤ (cid:15) .
Consider SVD as
X 1 = (cid:80) min
m,p } i =
[Equation]
u
[Equation]
v
[Equation]
i
[Equation]
v
[Equation]
1 + (cid:102) X 0 .
Then the EoS-learning-rate threshold at X =
X 0 is η = 1 σ 21 .
Therefore, we are to show the convergence of GD starting from X
= X 0 + ∆ X 1 with learning rate η = 1 σ 21 + β where β > 0 .
Theorem 7.
Consider the above symmetric matrix factorization with learning rate η = 1 σ 21 + β .
Assume 0 < βσ 21 < √ 4 .
5 − 1 ≈ 1 . 121 and ησ 22 < 1 .
The initialization is around the minimum, as
X
= X 0 + ∆ X 1 , with the deviation satisfying u (cid:62) 1 ∆ X 1 v 1 (cid:54) =
and (cid:107) ∆ X 1 (cid:107) ≤ (cid:15) bounded by a small value.
Then GD would converge to a period-2 orbit γ η approximately by a small margin in O ( (cid:15) ) , formally written as
[Equation]
where δ 1 ∈ (0 , 1) , δ 2 ∈ ( − 1 , 0) are the two solutions of solving δ in
[Equation]
Proof sketch The proof is provided in Appendix J.2.
The analysis consists of two phases, depending on whether (cid:15) t (cid:44) (cid:104)
X t − X 0 , u 1 v (cid:62) 1 (cid:105) is small or not.
In Phase I, all components of X t − X 0 are small due to the initialization near minima, but only (cid:15) t is growing exponentially in a rate of 1 + βσ 21 .
In Phase II, (cid:15) t is much larger than other components, as long as other components are still not growing.
So the dynamics of (cid:15) t matches GD of 1-D function f ( α ) =
14 (( α + σ 1 ) 2 − σ 21 ) 2 with learning rate η (cid:48)
= 1 σ 21 + β > 2 f (cid:48)(cid:48) (0)
= 1 σ 21 .
The proof concludes with 1-D convergence analysis of f ( x ) =
14 ( x 2 − µ ) 2 as shown in Theorem 2.
(cid:3)
Remark.
Theorem 7 assumes GD starts from any point in an (cid:15) -ball near the minima, except u (cid:62) 1 ∆ X 1 v 1 (cid:54) = 0 .
Note that this exception is with Lebesgue measure zero and, even if it is unfortunately satisﬁed, GD still has a chance to reach u (cid:62) 1 ∆ X t v 1 (cid:54) =
after several steps due to some higher-order small noise.
Then, this assumption could be relaxed as (cid:107) ∆ X 1 (cid:107) (cid:54) = 0 .
B A DDITIONAL E XPERIMENTS B. 2-D FUNCTION
As discussed in the Section A.1, on the function f ( x, y ) =
12 ( xy − 1) 2 , we estimate that | x − y | decays to 0 when η ∈ (1 , 1 . 5) , as shown in Figure 4.
Since it achieves a perfect balance, the two parameters follows convergence of the corresponding 1-D function f ( x ) =
14 ( x 2 − 1) 2 .
As shown in
Figure 4, xy with η = 1 .
05 converges to a period-2 orbit, as stated in the 1-D discussion of Theoremwhile xy with η = 1 .
25 converges to a period-4 orbit, which is out of our range in the theorem.
But still it falls into the range for balance in Theorem 5.
B. H IGH 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 ﬂatter minima.
Observation 2 : GD beyond EoS is in a similar style with the low dimension.
B.2. 2- LAYER HIGH - DIM HOMOGENEOUS R E LU NN S WITH PLANTED TEACHER NEURONS
We conduct a synthetic experiment in the high-dimension teacher-student framework.
The teacher network is in the form of
[Equation]
[Figure]
where x ∈ R 16 and e i is the i -th vector in the standard basis of R 16 .
The student and the loss are in forms of
[Equation]
Apparently, the global minimum manifold contains the following set M as (w.l.o.g., ignoring any permutation)
[Equation]
However, different choices of { k i } 16 i =1 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 ﬂattest minima from sharper minima .
Initialization.
We initialize all student neurons directionally aligned with the teachers as w i (cid:107) e i but choose various k i , as k i = 1 + 0 .
( i − 1) .
Obviously, such a choice of { k i } 16 i =1 is not at the ﬂattest minima, due to the isotropy of teacher neurons.
Also we add small noise to w i to make the training start closely (but not exactly) from a sharp minima, as
[Equation]
Data.
We uniformly sample 10000 data points from the unit sphere S 15 .
Training.
We run gradient descent with two learning rates η 1 = 0 .
5 , η 2 = 2 . 6 .
Later we will show with experiments that the EoS threshold of learning rate is around 2.5, so η 2 is 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 (a, b) present the gap between these two trajectories, where GD with a small learning rate stays around the sharp minima, while that with a larger one drives to ﬂatter minima.
Then GD stably oscillates around the ﬂatter minima.
Meanwhile, from Figure 5 (b), when we decrease the learning rate from 2.6 to 0.5 after 100 epochs,
GD converges to a nearby minima which is signiﬁcantly ﬂatter, compared with that of lr=0.5.
Figure 5 (c) provides a more detailed view of (cid:107) w i (cid:107) v i for all 16 neurons.
All neurons with lr=0.5 stay at the original ratio k 2 i .
But those with lr=2.6 all converge to the same ratio around k 2 = (cid:107) w (cid:107) v = 1 . 21 , as shown in Figure 5 (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, k 2 = 1 . 21 is the best choice of { k i } 16 i =1 such that the minima is the ﬂattest.
Therefore, we have shown that, in such a setting of high-dimension teacher-student network, GD beyond the edge of stability drives to the ﬂattest minima .
[Figure] [Figure]
[Figure] [Figure]
B.2. 3, 4, 5- LAYER NON - HOMOGENEOUS MLP S ON MNIST
We conduct an experiment on real data to show that our ﬁnding in the low-dimension setting in
Theorem 1 is possible to generalize to high-dimensional setting .
More precisely, our goals are to show, when GD is beyond EoS,
the oscillation direction (gradient) aligns with the top eigenvector of Hessian.
the 1D function at the cross-section of oscillation direction and high-dim loss landscape satisﬁes the conditions in Theorem 1.
Network, dataset and training.
We run 3, 4, 5-layer ReLU MLPs on MNIST LeCun et al.
The networks have 16 neurons in each layer.
To make it 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 signiﬁcant input channels where the standard deviation over the dataset is at least 110, which makes the network input dimension as 79.
We train the networks using MSE loss subjected to GD with large learning rates η = 0 .
35 and a small rate η = 0 .
1 (for 3-layer).
Note that the larger ones are beyond
EoS.
Deﬁnition 3 (line search minima) .
Consider a function f , learning rate η and a point x ∈ domain ( f ) .
We call ˜ x as the line search minima of x if
[Equation]
The line search minima ˜ x can interpreted as the lowest point on the 1D function induced by the gradient at x .
If GD is beyond EoS, ˜ x stays in the valley below the oscillation of x .
Results.
All results are presented in Figure 6, 7 and 8.
Take the 3-layer as an example.
From Figure 6 (a, b), GD is beyond EoS during epochs 10-14 and 21-60.
For these epochs, cosine similarity between the top Hessian eigenvector v 1 and the gradient is pretty close to 1, as shown in Figure 6 (c), which veriﬁes our goal 1.
In Figure 6 (d), we compute
[ f (3) ] 2 − f (2) f (4) at line search minima along training, which is required to be positive in Theorem 1 to allow stable oscillation.
Then it turns out most points have 3[ f (3) ] 2 − f (2) f
(4) > 0 except a few points, all of which are not in the EoS regime, and these few exceptional points might be due to approximation error to compute the fourth-order derivative since their negativity is quite small.
This veriﬁes our goal 2.
Both the above arguments are the same in the cases of 4 and 5 layers as shown in Figure 7 and 8.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
C P ROOF OF T HEOREM
Theorem 8 (Restatement of Theorem 1) .
Consider any 1-D differentiable function f ( x ) around a local minima ¯ x , satisfying (i) f (3) (¯ x ) (cid:54) = 0 , and (ii)
[ f (3) ] 2 − f (cid:48)(cid:48) f
(4) > 0 at ¯ x .
Then, there exists (cid:15) with sufﬁciently small | (cid:15) | and (cid:15) · f (3) > 0 such that: for any point x 0 between ¯ x and ¯ x − (cid:15) , there exists a learning rate η such that the update rule F η of GD satisﬁes F η ( F η ( x 0 ))
= x 0 , and
[Equation]
Proof.
For simplicity, we assume f (3) (¯ x ) > 0 .
Imagine a starting point x 0 = ¯ x − (cid:15), (cid:15) > 0 .
We omit f (cid:48)
(¯ x ) , f (cid:48)(cid:48) (¯ x ) , f (3) (¯ x ) , f (4) (¯ x ) as f (cid:48) , f (cid:48)(cid:48) , f (3) , f (4) .
After running two steps of gradient descent, we have
[Equation]
When η = 2 f (cid:48)(cid:48) , it holds
[Equation]
(29) which would be positive if 12 ηf (3) f (3) − 13 f (4) = 13 f
[Equation]
3[ f
[Equation]
3) ] 2 − f
[Equation]
f (4) ) > 0 and | (cid:15) | is sufﬁciently small.
When η = 2 f (cid:48)(cid:48) − (cid:15) · f (3) then ηf (cid:48)(cid:48)
[Equation]
[Equation]
(cid:15) + O ( (cid:15) 2 ) , it holds
[Equation]
η
[Equation]
(30) which is negative when | (cid:15) | is sufﬁciently small.
Therefore, there exists a learning rate η ∈ ( 2 f (cid:48)(cid:48) , 2 f (cid:48)(cid:48) − (cid:15) · f (3) ) such that x 2 = x 0 due to the continuity of ( x 2 − x 0 ) with respect to η .
The above proof can be generalized to the case of x 0 =
¯ x − (cid:15) (cid:48) with (cid:15) (cid:48) ∈ (0 , (cid:15) ] and the learning rate is still bounded as η ∈ ( 2 f (cid:48)(cid:48) , 2 f (cid:48)(cid:48) − (cid:15) · f
[Equation]
D P ROOF OF L EMMA
Lemma 4 (Restatement of Lemma 1) .
Consider any 1-D differentiable function f ( x ) around a local minima ¯ x , satisfying that the lowest order non-zero derivative (except the f (cid:48)(cid:48) ) at ¯ x is f ( k ) (¯ x ) with k ≥ 4 .
Then, there exists (cid:15) with sufﬁciently small | (cid:15) | such that: for any point x 0 between ¯ x and ¯ x − (cid:15) , and
if k is odd and (cid:15) · f ( k ) (¯ x ) > 0 , f ( k +1)
(¯ x ) < 0 , then there exists η ∈ (
[Equation]
(cid:48)(cid:48)
[Equation]
if k is even and f ( k ) (¯ x ) < 0 , then there exists η ∈ ( 2 f (cid:48)(cid:48)
[Equation]
such that: the update rule F η of GD satisﬁes F η ( F η ( x
[Equation]
) = x 0 .
Proof.
(1) If k is odd, assuming f ( k ) > 0 for simplicity, we have
[Equation]
[Equation]
When η = 2 f (cid:48)(cid:48)
[Equation]
it holds
[Equation]
When η =
[Equation]
then ηf
[Equation]
, then it holds
[Equation]
Since k is odd and (cid:15) · f ( k ) (¯ x ) > 0 , f ( k +1)
(¯ x ) < 0 , the above two estimations of x 2 − x 0 / η have one positive and one negative exactly.
Therefore, due to the continuity of x 2 − x 0 wrt η , there exists a learning rate η ∈ ( 2 f (cid:48)(cid:48) ,
[Equation]
2 ) such that x 2 = x 0 .
The above proof can be generalized to any x 0 between ¯ x and ¯ x − (cid:15) with the same bound for η .
(2) If k is even, we have
[Equation]
When η = 2 f (cid:48)(cid:48) , it holds
[Equation]
When η =
[Equation]
f
[Equation]
+
[Equation]
2 with c > 0 as some constant implying ηf
[Equation]
then it holds
[Equation]
where we then set c = 1 .
Hence, the above two estimations of x 2 − x 0 / η have one positive and one negative with sufﬁciently small | (cid:15) | .
Therefore, due to the continuity of x 2 − x 0 , there exists a learning rate η ∈ ( 2 f
[Equation]
[Equation]
such that x 2 = x 0 .
The above proof can be generalized to any x 0 between ¯ x and ¯ x − (cid:15) with the same bound for η .
Corollary 1.
f ( x ) = sin( x ) allows stable oscillation around its local minima ¯ x .
Proof.
Its lowest order nonzero derivative (expect f (cid:48)(cid:48) ) is f (4) ¯ x = sin(¯ x ) = − 1 < 0 and the orderis even.
Then Lemma 1 gives the result.
E P ROOF OF L EMMA
Lemma 5 (Restatement of Lemma 2) .
Consider a 1-D function g ( x ) , and deﬁne the loss function f as f ( x ) =
[Equation]
g
[Equation]
x ) − y ) 2 .
Assuming (i) g (cid:48) is not zero when g (¯ x ) =
y , (ii) g (cid:48) (¯ x ) g (3) (¯ x ) <
[ g (cid:48)(cid:48) (¯ x )] 2 , then it satisﬁes the condition in Theorem 1 or Lemma 1 to allow period-2 stable oscillation around ¯ x .
Proof.
From the deﬁnition, we have
[Equation]
Then at the global minima where g ( x ) =
y , we have f (cid:48)(cid:48)
( x ) =
[ g (cid:48)
( x )] 2 and f (3) ( x ) =
6 g (cid:48)(cid:48)
( x ) g (cid:48)
( x ) .
If we assume y is not a trivial value for g ( x ) , which means g (cid:48) ( x ) (cid:54) = 0 at the minima, and g is not linear around the minima (implies g (cid:48)(cid:48)
(cid:54) = 0 ), then f satisﬁes f (3) (¯ x ) (cid:54) = 0 in Theorem 1.
Meanwhile, we need 3 f
[Equation]
3) f
[Equation]
− f
[Equation]
f (4) > 0 as in Theorem 1, hence it requires
[Equation]
The remaining case is, if g (cid:48)
( x ) (cid:54) = 0 and g (cid:48)(cid:48)
=
at the minima, it satisﬁes the condition for Lemmawith k = 4 , because f (3) =
and f (4) <
due to (35, 37)
Corollary 2.
f ( x ) =
( x 2 − 1) 2 allows stable oscillation around the local minima ¯ x = 1 .
Proof.
With g ( x ) = x 2 , it has g (cid:48) (1) =
2 (cid:54) = 0 , g (cid:48)(cid:48) (1) = 2 (cid:54) = 0 .
All higher order derivatives of g are zero.
Then Lemma 2 gives the result.
Corollary 3.
f ( x ) = (sin( x ) − y ) 2 allows stable oscillation around the local minima ¯ x = arcsin( y ) with y ∈ ( − 1 , 1) .
Proof.
With g ( x ) = sin( x ) , it has g (cid:48) (¯ x ) =
cos(¯ x ) (cid:54) = 0 , g (3) (¯ x ) = − cos(¯ x ) .
We have g (3) (¯ x ) is bounded as g (cid:48) g (3) −
[ g (cid:48)(cid:48) ] 2 = − cos 2 (¯ x ) − 6 sin 2 (¯ x ) < 0 .
Then Lemma 2 gives the result.
Corollary 4.
f ( x ) = (tanh( x ) − y )
2 allows stable oscillation around the local minima ¯ x = tanh − 1 ( y ) with y ∈ ( − 1 , 1) .
Proof.
With g ( x ) = tanh( x ) , it has g (cid:48) (¯ x ) =
sech 2 (¯ x ) (cid:54) = 0 , and g (3) (¯ x ) = − 2 sech 4 (¯ x ) + 4 sech 2 (¯ x ) tanh 2 (¯ x ) is bounded as g (cid:48) g (3) −
[ g (cid:48)(cid:48) ] 2 = − 2 sech 6 + 4 sech 4 tanh 2 − 24 sech 4 tanh 2 = − 2 sech 6 − 20 sech 4 tanh 2 < 0 .
Then Lemma 2 gives the result.
Corollary 5.
f ( x ) =
( x α − y ) 2 (with k ∈ Z , k ≥ 2 ) allows stable oscillation around the local minima ¯ x = y 1 / α except y = 0 .
Proof.
With g ( x ) = x α , it has g (cid:48) (¯ x ) =
αx α − 1 , g (cid:48)(cid:48) (¯ x )
=
α ( α − 1) x
α − 2 , g (3) (¯ x )
=
α ( α − 1)( α − 2) x
α − 3 .
Then we have g (cid:48) g (3) − 6[ g (cid:48)(cid:48) ]
2 = α 2 ( α − 1)( − 5 α + 4)
x 2 α − 4 < 0 .
Then Lemma 2 gives the result.
Corollary 6.
f ( x ) =
(exp( x ) − y ) 2 allows stable oscillation around the local minima ¯ x = log y for y > 0 .
Proof.
With g ( x ) = exp x , it has g (cid:48) (¯ x ) =
g (cid:48)(cid:48) (¯ x ) = g (3) (¯ x ) =
exp(¯ x ) .
Then we have g (cid:48) g (3) − 6[ g (cid:48)(cid:48) ] 2 < 0 .
Then Lemma 2 gives the result.
Corollary 7.
f ( x ) = (log( x ) − y ) 2 allows stable oscillation around the local minima ¯ x = exp y .
Proof.
With g ( x ) = logx , it has g (cid:48) (¯ x ) =
1¯ x , g (cid:48)(cid:48) (¯ x ) = − 1¯ x 2 , g (3) (¯ x ) = − 2¯ x 3 .
Then we have g (cid:48) g (3) − 6[ g (cid:48)(cid:48) ] 2 < 0 .
Then Lemma 2 gives the result.
Corollary 8.
f ( x ) = ( 1 1+exp( − x ) − y ) 2 allows stable oscillation around the local minima ¯ x = sigmoid − 1 ( y ) for y ∈ (0 , 1) .
Proof.
With g ( x ) = 1 1+exp( − x ) , it has g (cid:48) (¯ x ) =
exp( − x ) (exp( − x )+1) 2 , g (cid:48)(cid:48) (¯ x ) = − exp( x )(exp( x ) − 1) (exp( x )+1), g (3) (¯ x ) = exp( x )( − 4exp( x )+exp(2 x )+1) (exp( x )+1).
Then we have g (cid:48) g (3) − 6[ g (cid:48)(cid:48) ]
2 ∝ − 4 exp( x )+exp(2 x )+1 − 6(exp( x ) − 1) 2 < 0 .
Then Lemma 2 gives the result.
Proposition 4 (Restatement of Prop 1) .
Consider two functions f, g .
Assume both f ( x ) , g ( y ) at x = ¯ x, y = f
(¯ x ) satisﬁes the conditions in Lemma 2 to allow stable oscillations.
Then g ( f ( x )) allows stable oscillation around x = ¯ x .
Proof.
Denote F ( x ) (cid:44) g
( f ( x )) .
Then we have
[Equation]
Thus, omitting all variables ¯ x and f (¯ x ) in the derivatives
[Equation]
it holds
[Equation]
where the inequality is due to all conditions in Lemma 2.
So the only problem is whether we can achieve g (cid:48)
For g (cid:48) g (cid:48)(cid:48)
f (cid:48)(cid:48) < 0 , construct ˆ g ( y ) (cid:44) g ( − y ) , ˆ f ( x ) (cid:44) − f ( x ) .
In this sense, it holds ˆ g ( ˆ f (¯ x ))
=
g ( f (¯ x )) .
It is easy to verify that both ˆ g, ˆ f at y = − f
(¯ x ) , x = ¯ x satisfy the conditions in Lemma 2, because ˆ g (cid:48) ( y ) =
− g (cid:48) ( − y ) =
− g (cid:48)
( f (¯ x )) , ˆ g (cid:48)(cid:48) ( y ) =
g (cid:48)(cid:48) ( − y ) =
g (cid:48)(cid:48)
( f (¯ x )) , ˆ g (3) ( y ) =
− g (3) ( − y )
= − g (3) ( f (¯ x )) ,
ˆ f (cid:48) (¯ x ) = − f (cid:48) (¯ x ) , ˆ f (cid:48)(cid:48) (¯ x ) = − f (cid:48)(cid:48) (¯ x ) , ˆ f (3) ( y )
=
− f (3) (¯ x ) .
Then, it has ˆ g (cid:48)
( y )
ˆ g (cid:48)(cid:48)
( y )
ˆ f (cid:48)(cid:48)
( x ) = − g (cid:48)
g
(cid:48)(cid:48)
f (cid:48)(cid:48) > 0 at y = − f (¯ x ) , x = ¯ x .
Therefore, we have
F (cid:48) (¯ x ) F (3) (¯ x ) − 6[ F (cid:48)(cid:48) (¯ x )] 2 < 0 and Lemma 2 gives the result.
F P ROOF OF T HEOREM
Theorem 9 (Restatement of Theorem 2) .
For f ( x ) = 14 ( x 2 − µ ) 2 , consider GD with η = K · 1 µ where 1 < K < √ 4 .
5 − 1 ≈ 1 . 121 , and initialized on any point 0 < x 0 < √ µ .
Then it converges to an orbit of period 2, except for a measure-zero initialization where it converges to √ µ .
More precisely, the period-2 orbit are the solutions x = δ 1 ∈ (0 , √ µ ) , x = δ 2 ∈ ( √ µ, 2 √ µ ) of solving δ in
[Equation]
Proof.
Assume the 2-period orbit is (¯ x 0 , ¯ x 1 ) , which means
[Equation]
First, we show the existence and uniqueness of such an orbit when K ∈ (1 , 1 . 5] via solving a high-order equation, some roots of which can be eliminated.
Then, we conduct an analysis of global convergence by deﬁning a special interval I .
GD starting from any point following our assumption will enter I 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 (¯ x 0 , ¯ x 1 ) .
Before diving into the proof, we brieﬂy show it always holds x > 0 under our assumption.
If x t − 1 > 0 and x t ≤ 0 , the GD rule reveals η ( µ − x 2 t − 1 ) ≤ − 1 which implies x 2 t − 1 ≥ µ + 1 η .
However, the maximum of x + η ( µ − x 2 ) x on x ∈ (0 , (cid:113) µ + 1 η ) is achieved when x 2 = 13 ( µ + 1 η ) so the maximum value is (cid:113) 13 ( µ + 1 η )( 23 + 23 ηµ ) ≤ 1 . 4 (cid:113) 13 ( µ + 1 η ) < (cid:113) µ + 1 η .
As a result, it always holds x > 0 .
Part I. Existence and uniqueness of (¯ x 0 , ¯ x 1 ) .
In this part, we simply denote both ¯ x 0 , ¯ x 1 as x 0 .
This means x 0 in all formulas in this part can be interpreted as ¯ x 0 and ¯ x 1 .
Then the GD update rule tells, for the orbit in two steps,
[Equation]
which means
[Equation]
Denote z := 1 + η ( µ − x 20 ) , it is equivalent to
[Equation]
If z + 1 = 0 , it means x 1 = − x
which is however out of the range of our discussion on the x >domain.
So we require − x 20 ( z − 12 ) 2 + µ − 34 x 20 = 0 .
To ensure the existence of solutions z , it is natural to require
[Equation]
Then, the solutions are
[Equation]
However, z = 12 − (cid:113) µx 20 − 34 can be ruled out.
If it holds, η ( µ − x 20 ) =
z − 1 < − 12 which means x 20 > µ + 12 η .
Since we restrict ηµ ∈ (1 , 1 . 121] , it tells x 20 > µ (1 + 1 1 . 242 ) contradicting with µ ≥ 34 x 20 .
Hence, z = 12 + (cid:113) µx 20 − 34 is the only reasonable solution, which is saying
[Equation]
Given a certain η , the above expression is a third-order equation of x 20 to solve.
Apparently x 20 = µ is 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 (¯ x 0 , ¯ x 1 ) , if the equation does have three different roots.
This also guarantees the uniqueness of such an orbit.
Assuming x 20 (cid:54) = µ , the above expression can be reformulated as
[Equation] [Equation]
One necessary condition for existence is µ ≥ 34 x 20 .
Note that here x 0 can be both ¯ x 0 , ¯ x 1 , one of which is larger than √ µ .
For simplicity, we assume ¯ x 0 < √ µ < ¯ x 1 .
Since η from Eq(39) is increasing with x 20 when µ < x 20 , let x 20 = 43 µ and achieve the upper bound as
[Equation]
[Equation]
which is satisﬁed by our assumption 1 < ηµ < √ 4 .
5 − 1 ≈ 1 . 121 .
Therefore, we have shown the existence and uniqueness of a period-2 orbit.
Part II.
Global convergence to (¯ x 0 , ¯ x 1 ) .
The proof structure is as follows:
There exists a special interval I :=
[ x s , √ µ ) such that any point in I will back to this interval surely after two steps of gradient descent.
And ¯ x 0 ∈ I .
Initialized from any point in I , the gradient descent process will converge to ¯ x 0 (every two steps of GD).
Initialized from any point between 0 and √ µ , the gradient descent process will fall into I in some steps.
(II.1) Consider a function F η ( x ) =
x
+ η ( µ − x 2 ) x performing one step of gradient descent.
Since
F (cid:48)
η ( x ) = 1 + ηµ − 3 ηx 2 , we have F (cid:48)
η ( x ) > 0 for 0 < x 2 < 13 (cid:16) µ + 1 η (cid:17) and F (cid:48)
η ( x ) < 0 otherwise.
It is obvious that the threshold has x 2 s := 13 (cid:16) µ + 1 η (cid:17) < µ .
In the other words, for any point on the right of x s , GD returns a point in a decreasing manner.
To prove anything further, we would like to restrict ¯ x 0 ≥ x s , which is
[Equation]
Solving this inequality tells
[Equation]
Consequently, by applying Eq(39), we have
[Equation]
With the above discussion of x s , we are able to deﬁne the special internal I :=
[ x s , √ µ ) .
From the deﬁnition of F η , consider a function representing two steps of gradient descent F 2 η ( x ) :
= F η ( F η ( x )) .
From previous discussion, we know F 2 η (¯ x 0 ) =
¯ x 0 .
What about F 2 η ( x s ) ?
It turns out F 2 η ( x s ) > x s : we have F η ( x s ) =
x s (1 + ηµ − ηx 2 s ) =
x s · 2 3 (1 + ηµ ) and, furthermore,
F
[Equation]
x s ) =
F
[Equation]
1 + ηµ ) 3 (cid:1) .
Then we get F 2 η ( x s ) >
x s because
[Equation]
Combining the following facts, i) F 2 η ( x ) − x is continous wrt x , ii) F 2 η ( x s ) − x s > 0 , and iii)
F 2 η (¯ x 0 ) − ¯ x 0 = 0 is the only zero point on x ∈ [ x s , ¯ x 0 ] , we can conclude that
[Equation]
[Equation]
Meanwhile, we can prove F 2 η ( x ) <
x for any x ∈ (¯ x 0 , √ µ ) .
Since F 2 η ( µ ) − µ = 0 and F 2 η (¯ x 0 ) − ¯ x 0 = 0 are the only two zero cases, we only need to show ∃ ˆ x ∈ (¯ x 0 , √ µ ) , such that F 2 η (ˆ x ) < ˆ x .
We compute the derivative of F 2 η ( x ) − x at x 2 = µ , which is ddx F 2 η ( x ) − x | x 2 = µ = − 1 +
F (cid:48)
( F ( x )) F (cid:48) (
[Equation]
) | x 2 =
[Equation]
= − 1 + [ F
[Equation]
( √ µ )] 2 = − 1 + (1 − 2 ηµ ) 2 > 0 .
Then combining it with
F 2 η (¯ x 0 ) =
¯ x 0 , there exists a point ˆ x ∈ (¯ x 0 , √ µ ) that is very close to √ µ such that F 2 η (ˆ x ) <
ˆ x .
Hence, we can conclude that
[Equation]
Since F η ( · ) is decreasing on [ x s , ∞ ) and F η ( x ) > x s for x ∈ [ x s , √ µ ] , it is fair to say F 2 η ( x ) is increasing on x ∈ [ x s , √ µ ] .
Hence, we have F 2 η ( x ) ≤ F 2 η (¯ x 0 ) =
¯ x 0 , ∀ x ∈ [ x s , ¯ x 0 ] .
And
F 2 η ( x ) ≥ F 2 η (¯
[Equation]
) =
[Equation] [Equation]
Combining the above results, we have
[Equation] [Equation]
(II.2) A consequence of Exp
, 47) is
that
any point in I will converge to ¯ x
with
even steps of gradient
descent.
For simplicity
we
provide the proof for
x
∈
[ x
s ,
¯ x
Denote a 0 ∈
[ x s , ¯ x 0 ) and a n := F 2 η ( a n − 1 ) , n ≥ 1 .
The series { a i } i ≥ 0 satisﬁes
[Equation]
Since the series is bounded and strictly increasing, it is converging.
Assume it is converging to a .
If a < ¯ x 0 , then
[Equation]
Since F 2 η ( · ) is continuous, so ∃ δ > 0 , such that, when | x − a | < δ , we have
[Equation]
Since a is the limit, so ∃ N > 0 , such that, when n > N , 0 < a − F 2 η ( a n ) < δ .
So, combining with
Exp(49), we have
[Equation]
But LHS = F 2 η ( a ) − a n +2 > F 2 η ( a ) − a , so we reach a contradiction.
Hence, we have { a i } converges to ¯ x 0 .
(II.3) Obviously, any initialization in (0 , √ µ ) will have gradient descent run into (i) the interval I , or (ii) the interval on the right of √ µ , i.e. , ( √ µ, ∞ ) .
The ﬁrst case is exactly our target.
Now consider the second case.
From the deﬁnition of x s in part III.1, we know F η ( x s ) =
max x ∈ [0 , √ µ ] F η ( x ) .
So it is fair to say this case is x n ∈ ( √ µ, F η ( x s )] .
Then the next step will go into the interval I , because
[Equation]
where the ﬁrst inequality is from the decreasing property of F η ( · ) and the second inequality is due to
F 2 η ( x ) > x on x ∈ [ x s , ¯ x 0 ) .
G P ROOF OF T HEOREM
Theorem 10 (Restatement of Theorem 5) .
For f ( x, y ) =
12 ( xy − µ ) 2 , consider GD with learning rate η = K · 1 µ .
Assume both x
and y are always positive during the whole process { x i , y i } i ≥ 0 .
In this process, denote a series of all points with xy > µ as P = { ( x i , y i )
| x
i y i > µ } .
Then | x − y | decays to 0 in P , for any 1 < K < 1 . 5 .
Proof.
Consider the current step is at ( x t , y t ) with x t y t > µ .
After two steps of gradient descent, we have
[Equation]
with which we have the difference evolve as
[Equation] [Equation]
Meanwhile, we have
[Equation] [Equation]
Note that the second term in Eq(56) vanishes when x and y are balanced.
When they are not balanced, if x t y t > µ , it holds x t +1 y t
+1 < x t y t (1 + η ( µ − x t y t )) 2 .
Incorporating this inequality into Eq(54, 55) and assuming y t − x t > 0 , it holds
[Equation]
To show that | x − y | is decaying as in the theorem, we are to show
y t
+2 − x t +2 <
y t − x t 2.
y t +2 − x t +2 > − ( y t − x t )
Note that, although x t y t > µ , it is not sure to have x t +
y t +2 > µ .
However, for any 0 < x i y i < µ
and K < 2 , we have
[Equation]
which is saying | x − y | decays until it reaches xy > µ .
So it is enough to prove the above two inequalities, whether or not x t +2 y t +2 > µ .
Part I.
To show y t +2 − x t +2 <
y t − x t
Since we wish to have y t +2 − x t +2 <
y t − x t , it is sufﬁcient to require
[Equation]
[Equation]
Since we assume x t +1 , y t +1 > 0 , Eq (50, 51) tells η ( µ − x t y t )
> − min { x t y t , y t x t } , which is equivalent to 1 − η ( µ − x t y t ) < 1 + min { x t y t , y t x t } .
(I.1)
If η
[Equation]
µ − x
[Equation]
y
[Equation]
) ≥
Then we have 1 − η ( µ − x t +1 y t +1 ) ≤ 12 .
As a result, y
[Equation]
If η ( µ − x t +1 y t +1 ) < 12 and x t +1 y t +1 ≤ x 2 s =
(cid:16)
[Equation]
The second condition reveals
[Equation]
The ﬁrst condition is equivalent to x t +1 y t +1 > µ − 12 η .
Since the second term in Eq(56) is negative, we have
[Equation]
with which we would like to ﬁnd an upper bound of x t y t .
Denoting b = x t y t , consider a function q ( b )
=
b (1 + η ( µ − b )) 2 .
Obviously q ( µ ) = µ .
Its derivative is q (cid:48)
( b ) =
(1 + ηµ − ηb ) (1 + ηµ − 3 ηb ) < 0 on the domain of our interest.
If we can show an (negative) upper bound for the derivative as q (cid:48) ( b ) < − 1 on a proper domain, then it is fair to say that, from Exp(63), x t y t < µ + 12 η .
Then we have
[Equation]
Then, combining Exp(64, 62), it tells
[Equation] [Equation]
The remaining is to show q (cid:48)
( b ) < − 1 on a proper domain.
We have q (cid:48)
( b ) =
(1+ ηµ − 2 ηb ) 2 − ( ηb ) 2 , which is equal to 1 − 2 ηµ < − 1 when b = µ .
Meanwhile, the derivative of q (cid:48)
( b ) is q (cid:48)(cid:48)
( b ) = − 2 η ( ηb + (1 + ηµ − 2 ηb ))
= − 2 η (1 + ηµ − ηb ) , which is negative when b < µ + 1 η .
As a result, it always holds q (cid:48)
( b ) < − 1 when b < µ + 1 η .
(I.3)
If x t +1 y t +1 ≥ x 2 s
Denoting again b
= x t y t , the above inequality in is saying, with b > µ ,
[Equation]
After expanding p ( · ) , we have
[Equation]
Apparently p ( µ ) = 1 .
So it is necessary to investigate whether p (cid:48) ( b ) < 0 on b > µ , as
[Equation]
Since ηb > 1 and b > µ , it is enough to require
[Equation]
It sufﬁces to show
[Equation]
Since x t +1 y t +1 ≥ x 2 s = 13 (cid:16) µ + 1 η (cid:17) , it holds
[Equation]
where the last inequality holds because: if b ≥ 3 (cid:16) µ + 1 η (cid:17) , then 1+ η ( µ − b ) ≤ − 2 ηµ − 2 < 0 , which contradicts with the assumption that both x t +1 , y t +1 are positive.
As a result, the above argument gives
[Equation]
Part II.
To show y t +2 − x t +2 > − ( y t − x t )
Since x t y t > µ , we have 1 − η ( µ − x t y t ) > 1 .
Combining with 1 − η ( µ − x t y t ) < 2 , it holds
[Equation]
So the remaining is to have y t +2 − x t +2 y t +1 − x t +1 > − 0 .
Actually it is 1 − η ( µ − x t +1 y t +1 ) ≥ 1 − ηµ = 1 − K .
Therefore, we have
[Equation]
as required.
Part III.
To show y t − x t converges to
From Exp (61, 65, 68, 69), we have for points in P , | y − x | is a monotone strictly decreasing sequence lower bounded by 0.
Hence it is convergent.
Actually it converges to 0.
If not, assuming it converges to (cid:15) > 0 , the next point will have the difference as ˜ (cid:15) < (cid:15) as well as all following points.
Hence, the contradiction gives the convergence to 0.
H P ROOF OF L EMMA
Lemma 6 (Restatement of Lemma 3) .
In the setting of Theorem 5, denote the initialization as m = | y
− x 0 | √ µ and x 0 y 0 > µ .
Then, during the whole process, both x and y will always stay positive, denoting p = 4 ( m + √ m 2 +4 ) 2 and q = (1 + p ) 2 , if
[Equation]
Proof.
Considering x t y t > µ , one step of gradient descent returns
[Equation]
To have both x t +1 > 0 , y t +1 > 0 , it sufﬁces to have
[Equation]
This inequality will be the main target we need to resolve in this proof.
First, we are to show
[Equation]
With the difference ﬁxed as m =
( y 0 − x 0 ) / √ µ , assuming y 0 > x 0 , we have m/y
= (1 − x 0 /y 0 ) / √ µ .
if x 0 y 0 increases, both x 0 and y 0 increase then m/y 0 decreases, which means x 0 /yincreases.
As a result, we have min
[Equation]
(cid:27) > min (cid:26) yx
[Equation]
Therefore, at initialization, to have positive x 1 and y 1 , it is enough to require
[Equation]
From Theorem 5, it is guaranteed that | x t − y t | < | x 0 − y 0 | with t ≥ 2 until it reaches x t y t > µ , with which r is still a good lower bound for min { y t /x t , x t /y t } .
So what remains to show is it satisﬁes η ( x t y t − µ ) <
r for the next ﬁrst time x t y t > µ .
If this holds, we can always iteratively show, for any x t y t > µ along gradient descent,
[Equation]
Note that r itself is independent of x t y t and all the history, so it is ideal to compute a uniform upper bound of η ( x t y t − µ ) with any pair of ( x t − 1 , y t − 1 ) satisfying x t − 1 y t − 1 < µ .
Actually it is possible, since we have | x t − 1 − y t − 1 | bounded as in Theorem 5.
Assume x i y i > µ
and it satisﬁes the condition of η ( x i
y i − µ ) <
r and | x i − y
i | < | x 0 − y 0 | .
As in (54), we have
[Equation]
Hence, it sufﬁces to get the maximum value of g ( z ) , with z ∈ (0 , µ ) , as
[Equation]
which is from (56).
Denote ¯ z = argmax g
( z ) .
Obviously ¯ z < 13 ( µ + 1 η ) (cid:44) z b , because the ﬁrst term of g ( z ) achieves maximum at z = 13 ( µ + 1 η ) and the second term is in a decreasing manner with z .
Then let’s take the derivative of g ( z ) as
[Equation]
where the ﬁrst term is always positive, so we have
[Equation]
which means
[Equation]
where the inequality is from ¯ z < 13 ( µ + 1 η ) .
As a result, it is safe to say
[Equation]
with which we are able to compute max η ( g ( z ) − µ ) , which is exactly the ﬁnal result.
I P ROOF OF T HEOREM
Theorem 11 (Restatement of Theorem 3) .
In the above setting, consider a teacher neuron ˜ w
=
[1 , 0] and set the learning rate η = Kd with K ∈ (1 , 1 . 1] .
Initialize the student as (cid:13) (cid:13) w (0) (cid:13) (cid:13)
= v (0) (cid:44) (cid:15) ∈ (0 , 0 . 10] and (cid:104) w (0) , ˜ w (cid:105) ≥ 0 .
Then, for t ≥ T 1 + 4 , w ( t ) y decays as w ( t ) y < 0 .
1 · (1 − 0 .
030 K ) t − T 1 − 4 ,
T
[Equation]
≤ (cid:24) log
[Equation] [Equation]
(cid:25) , β
[Equation]
(cid:18) 1 + 1 .π
(cid:19) (cid:15).
Proof.
We restate the update rules as
[Equation]
arctan
[Equation]
For simplicity, we will omit all superscripts of time t unless clariﬁcation is necessary.
From (83), if the target is to show w y decaying with a linear rate, it sufﬁces to bound the factor term in (82) (by a considerable margin) as
[Equation]
The technical part is the second inequality of (84).
If v, w x > 0 , it is equivalent to
[Equation]
where the RHS is smaller than or equal to 12 π .
Hence, 12 π is a special threshold with which we will frequently compare vw x .
Another important variable to control is v − w x
that reveals how the two layers are balanced.
If it is too large, for the iteration v ( t +1) w ( t +1) x may explode as shown in the 2-D case.
The main idea of our proof is that
• Stage 1 with vw x ≤ w x w y π (cid:107) w (cid:107) 2 : in this stage, w y grows but it grows in a smaller rate than that of v and w x .
Therefore, since we have an upper bound for vw x to stay in this stage, we are able to compute the upper bound of # iterations to ﬁnish this stage, which is T 1 in the theorem.
At the end of this stage, both of v − w x and w y are bounded under our assumption of initialization.
• Stage 2 with vw
x > w x w y π (cid:107) w (cid:107) 2 : in this stage, w y decreases.
Since our range of a large learning rate is relatively narrow (1 < K ≤ 1 . 1) , we are able to compute bounds of vw x , v − w x and w y .
After eight iterations, it falls into (and stays in) a bounded basin of these three terms, in which w y decays at least in a linear rate.
Stage 1.
We are to show that, in the last iteration of this stage, there are three facts: 1) vw x ≤ 12 π , 2) v −
w x ∈ [ − 0 .
17] , and 3) w y ≤ 0 .
At initialization, we assume v (0) =
(cid:13) (cid:13) w (0) (cid:13) (cid:13) .
Denote α 0 = arctan
( w (0)
y /w (0) x ) ∈ [0 , π/ 2] .
So for next iteration we have
[Equation]
Apparently w (1) y increases with α 0 increasing.
And
[Equation]
Since in stage 1 it holds ∆ w y > 0 which means − ( v (0) )
+ 1 π sin α 0 > 0 in (85).
So it follows
[Equation]
w
[Equation] [Equation]
x ≤ 0 .
Combining the above arguments, we have
[Equation]
Regarding vw y , it has v (0) ≥ w (0) y at initialization due to v (0) =
(cid:13)(cid:13) w (0) (cid:13)(cid:13) .
From (80, 81, 82), we have v ∆ v = w x ∆ w x + w y ∆ w y .
So it holds v ∆ v ≥ y ∆ y .
Meanwhile, ∆ w y v
[Equation]
∈ [0 , Kπ ] .
From Lemma 7, given v ( t ) ≥ w ( t ) y and ∆ w y v ∈ [0 , 1] for any t in this stage, it always holds v ( t +1) ≥ w ( t +1) y .
Therefore, it is fair to say
[Equation]
Additionally, to bound the term vw y / (cid:107) w (cid:107) 2 in ∆ w y , we would like to show it always has vw y ≤ (cid:107) w (cid:107) 2 .
At initialization, it naturally holds.
Then, for the every next iteration, given it holds in the last iteration, we have
[Equation]
where the ﬁrst equality uses v ∆ v = w x ∆ w x + w y ∆ w y , the ﬁrst inequality uses the proven v ≥ w y and v ≥ ∆ w y , the second inequality uses the assumption vw y ≤ (cid:107) w (cid:107) 2 .
Now we are to show v ∆ w y − (∆ w x ) 2 ≤ 0 .
We have
[Equation]
Since we have proven w (1) y /w (1) x ≤ 2 .
7 , it is easy to check that
[Equation]
As a result, v ∆ w y − (∆ w x ) 2 ≤ 0 at time 1.
Furthermore, by checking each term, v ∆ w y − (∆ w x )decreases with w y /w x decreasing.
We will soon show that w y /w x itself decreases, by showing the growth ratio of w x is larger than that of w y .
Our target lower bound of the growth ratio of w x is that
[Equation] [Equation]
which is larger than the growth ratio of w y bounded by 1 π due to v ∆ w y < (cid:107) w (cid:107) 2 .
So it sufﬁces to show
Kv/w x ≥ 1 .
Assuming Kv/w x ≥ 1 for the current step, we need to show Kv ( t +1) /w
( t +1)
x ≥also holds for the next step.
Let’s denote
[Equation]
Then
[Equation] [Equation]
If KA ≤ 1 , since K > 1 and A > 0 , we have (89) as positive, which is what we need.
If KA > 1 , then
[Equation]
where the ﬁrst inequality is due to A ≤ K and the assumption of Kv ( t ) /w
( t ) x ≥ 1 .
Then it sufﬁces to show ( − K + A ) v + w x ≥ ( − K + 1 K ) v + w x ≥ 0 .
Note that − K + 1 /K ∈ ( − 0 .
2 , 0] when
K ∈ (1 , 1 . 1] .
It is easy to verify that v (1) ≤ 5 w (1) x .
Then, for the next step, we need to show v ( t +1) ≤ 5 w ( t +1)
x given v ( t ) ≤ 5 w ( t +1) .
To prove this, we are to bound v − w x , as
[Equation]
where the ﬁrst inequality is due to, when w y /w x ≤ 2 . 7 ,
[Equation]
We will later show that v ( t +1) − w ( t +1) ≥ − 0 .
1( v ( t ) − w ( t ) ) .
Combining this with (90), it is safe to say
[Equation]
where the second inequality is due to v ≤ 5 w and v ≥ w y .
Since w ( t +1) ≥ w ( t ) (due to A > 0 ) in this stage, we have v ( t +1) ≤ 5 w ( t +1) x .
Combining the above discussion, we have prove (87).
Obviously, when w y /w x ≤ 2 . 7 , RHS of (87) is at least 0 .
55 , larger than 1 . 1 /π , which is the upper bound of the ∆ w y /w y .
As a result, w y /w x keeps decreasing.
The next step is to show the growing ratio of vw x is much larger than that of w y .
From (81, 82), it holds
[Equation]
where the ﬁrst inequality is due to ∆ w y ≥ 0 .
It follows v ( t +1) w ( t +1) x
[Equation]
So far, we have shown the following facts: under the deﬁned initialization at time 0, starting from time 1, we have
vw x ≤ 1 / 2 π . 2.
∆ w x /w x + 1 ≥ 1 . 55 . 3.
∆ w y /w y + 1 ≤ 1 + K/π . 4.
w y /w x ≤ 2 .
and keeps decreasing.
v ( t +1) w ( t +1) x /v ( t ) w ( t ) x ≥ 2 . 56 .
v ≥ w y .
7. v ∆ w y < (∆ w x ) 2 .
Now we are to use the above facts to bound vw x , w y and v − w x to the end of stage 1.
For vw x , in previous discussion, we have shown that vw x ≤ 12 π .
Actually, there is another special value
[Equation]
when
[Equation]
This value is slightly larger than 1 / 4 π .
Hence, we would like to split the analysis into three parts: in the ﬁrst step of stage 2 ,
vw x ≥ 1 2 π . 2.
14 π ≤ vw x < 12 π . 3.
vw x < 14 π .
Note that, although we are discussing the stage 1 in this section, investigating the lower bound of the ﬁrst step in stage 2 helps calculate the number of iterations in stage 1.
Furthermore, it helps bound several variables in stage 1.
Case (I).
If vw x ≥ 12 π in ﬁrst step of stage 2:
Since we have prove v (1) w (1) x ( w (1) y ) 2 ≥ 1 / 2 .
7 and v ( t +1) w ( t +1) x
/v ( t ) w ( t ) x ≥ 2 .
56 , the number of iterations for vw x to reach 1 / 2 π is at most
[Equation]
Meanwhile, starting from time 1, the growth ratio of w y is
[Equation]
where the ﬁrst inequality is due to vw y ≤ (cid:107) w (cid:107) 2 , the second is due to K > 1 and the third is from v ≥ w y .
Therefore, combining with (92), we can bound w y in the end of stage 1 as
[Equation]
Since it initializes as (cid:13)(cid:13) w (0) (cid:13)(cid:13) ≤ 0 .
1 , we have w (1) y ≤
[Equation]
1(1 + 1 . 1 /π )
= 0 .
Then, it can be veriﬁed that, when w (1) y ∈ (0 , 0 . 135] , it holds
[Equation]
[Equation]
The next is to bound v − w x .
Combining the update rules of v and w x in (80, 81), we have
[Equation] [Equation]
Note that
[Equation]
where the left is due to vw x >
and , the right is from ∆ w y ≥ 0 .
When w y /w x ≤ 2 . 7 , the RHS follows − 1 + arctan( w y /w x ) π ≤ − 0 .
So combining both sides tells
[Equation] [Equation]
Since ∆ w y ≥ 0 , we have 0 ≤ K w 2 y v ( − v 2 + vw y π (cid:107) w
[Equation]
2 ) ≤ Kπ w y w 2 y (cid:107) w (cid:107) 2 .
Note that at initialization w (0) x ≤ v (0) .
Then it is easy to verify that
[Equation]
Because the coefﬁcient on the positive side in (98) is larger than 0 . 4 > 0 .
1 , it is appropriate to upper bound the v − w x as
[Equation]
 
[Equation]
 
[Equation]
where the second inequality is from the different growth ratios of w x and w y .
Note that here we take all T ≥ 1 and pick the largest value of RHS to bound w y .
It turns out
[Equation]
Furthermore, to lower bound v − w x , since obviously | v − w x | ≤ 0 .
17 , it follows
[Equation]
Case (II).
If 14 π ≤ vw x < 12 π in ﬁrst step of stage 2:
Similar to the discussion in Case (I), we are able to compute the number of iterations for vw x to reach 1 / 4 π .
It is at most
[Equation]
Accordingly, w y is bounded as
[Equation]
For simplicity, we just keep the bounds for v − w x as in Case (I), as
[Equation]
Case (III).
If vw x < 14 π in ﬁrst step of stage 2:
From the condition, we know vw x < 14 π as well in the last step of stage 1.
Since ∆ w y > 0 in stage 1, it tells
[Equation]
which means
[Equation]
Since 2 + √ 3 > 2 . 7 , if w y /w x ≥ 2 + √ 3 , then for time 1, ( v (1) , w (1) x , w (1) y ) is already in the stage 2.
However, it is not possible because (cid:13)(cid:13) w (0) (cid:13)(cid:13)
= v (0) ≤ 0 . 1 , which means v (1) w (1) x can not reach
1 π 2 .
7 1+2 .
Therefore, the only possible is w x w y ≥ 2 + √ 3 .
In this case, we are able to bound w y as
[Equation]
where the second inequality is due to vw x ≤ 14 π and v − w x ≥ − 0 .
Note that here we still use the bound of v − w x from Case (I), although it is loose somehow but it is enough for our analysis.
We leave the analysis of the bound of number of iterations to the end of this section.
Stage 2.
In the case (I) of stage 1, where the ﬁrst step in stage 2 is with vw x ≥ 12 π , it has v − w x ∈ [ − 0 .
17] and w y ≤ 0 .
In the case (II), where the ﬁrst step of stage 2 is with vw x ∈ [ 14 π , 12 π ] , it has v − w x ∈ [ − 0 .
17] and w y ≤ 0 .
In the case (III), where the ﬁrst step of stage 2 is with vw x ∈ [ 1 4 π , 1 2 π ] , it has v − w x ∈ [ − 0 .
17] and w y ≤ 0 .
[Equation]
[Equation] [Equation]
[Table]
[Equation]
[Table]
[Equation]
where we
One
[Equation]
≤ 0 .
[Equation]
where we use vw x ≤ 1 / 4 π , x/ (1 + x 2 ) ≤ 0 .
5 for any x .
Therefore, we can see that, in the ﬁrst step of stage 2, vw x ≤ 0 .
Next we are going to show how the iteration goes in the stage 2.
In Case (I), there are three facts:
w y ≤ 0 .
v − w x ∈ [ − 0 .
vw x ∈ [ 12 π , 0 . 668] .
Similarly, in Case (II), there are three facts as well:
w y ≤ 0 .
v − w x ∈ [ − 0 .
vw x ∈ [ 14 π , 12 π ] .
The main idea is to ﬁnd a 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 vw x ≥ 1 .
Consider any iteration with vw x ∈ [ m 1 , m 2 ] , v − w x ∈
[ d 1 , d 2 ] , w y ≤ e , we compute the bounds of v ( t +1) w ( t +1) x , v ( t +1) − w ( t +1) x , w ( t +1)
y in the following process (naturally assuming d 1 < 0 < d 2 )
If
m 1 ≥ 1 :
(a) Compute w x ≥ (cid:112)
[Equation]
(b) Compute w y w x ≤
[Equation]
(c) Compute arctan
w y /w x ) − wxwy
w
[Equation]
) Compute v ( t +1)
w
1) x
[Equation] [Equation]
h ) max {| d 1 | , | d
[Equation] [Equation]
[Equation]
1 e 2 m 2 .
This is from
[Equation]
(e) Compute v ( t +1) w ( t +1) x ≤ m 1 (1 + 1 .
0(1 − m 1 )) 2 .
This is due to x (1 + K (1 − x ))decreases with x increasing when x ≥ 1 .
(f) Compute v ( t +1)
− w ( t +1) x ∈
[Equation] [Equation]
+
[Equation] [Equation]
[Equation] [Equation]
[Equation]
d 2 / 2)) , d 2 (1 + 1 . 1( m 2 − 1 + h ))] .
This is due to ∆
[Equation]
where vw x ≥ 1 , the last term is between − Kvw 2 y and 0.
(g) Compute w ( t +1)
y
[Equation]
e
[Equation]
1 | , | j 2 |} , where
[Equation]
This is due to
[Equation]
then we would like to have the smallest value as j 1 − 1 and the largest value as j 2 −
[Equation]
Since w y is always non-negative, taking the maximum absolute value gives the upper bound.
If m 2 < 1 : (a) Compute w x ≥ (cid:112) m 1 + ( d 2 / 2) 2 − d 2 / 2 (cid:44) f .
(b) Compute w y w x ≤ e/f (cid:44) g .
(c) Compute arctan( w y /w x ) − wxwy (cid:107) w (cid:107) 2 π
[Equation]
(d) Compute v ( t +1) w ( t +1) x ≥ min x ∈
[ m 1 ,m 2 ] x (1 + 1 .
0(1 − x − h ))
2 − 1 .
1 e 2 x .
Compared with the case of m 1 ≥ 1 , we drop the term 1 .
1(1 − m 2 − h ) max {| d 1 | , | d 2 |} 2 because it is possible to have v − w x = 0 in some iterations.
(e) Compute v ( t +1) w ( t +1) x ≤ max x ∈ [ m 1 ,m 2 ] x
(1 + 1 . 1(1 − x ))
+ 1 .
1(1 − x ) max {| d 1 | , | d 2 |} 2 .
Compared with the case of m 1 ≥ 1 , we add a term depending on the | v − w x | max because it enlarges vw x in the in-balanced case.
(f) Compute v ( t +1)
− w ( t +1) x
[Equation]
d 2 / 2)) , d 2 (1 + 1 . 1( m 2 − 1 + h ))] .
In fact, a rigorous left bound should include more terms to select a minimum from.
Here it is simple because it keeps 1+ K ( m 1 − 1) ≥
in the following computing so we do not need to worry about the ﬂipping sign of dand d 2 .
(g) Compute w ( t +1)
y ≤ e · max {| j 1 | , | j 2 |} , where j 1 , j 2 are the same with those in the case of m 1 ≥ 1 .
Therefore, with the above process, we are able to brutally compute the ranges of v ( t +1) w ( t +1)
x , v ( t +1) − w ( t +1) x , w ( t +1)
y from the current ranges.
Note that this process plays a role of building a mapping from one interval to another interval, which covers all points from the source interval.
However, it is loose to some extent because gradient descent is a mapping from a point to another point.
The advantage of such a loose method is feasibility of obtaining bounds while losing tightness.
To achieve tightness, later we will also include some wisdom in a point-to-point style.
Also note that, a nice way to combine tightness and efﬁciency in this method is to split and to merge intervals when necessary.
For Case (I):
Now we are to compute the ranges starting from the interval where I = { w y ≤ 0 .
44 , v − w x ∈ [ − 0 .
17] , vw x ∈ [ 12 π , 0 .
First, we split it into three intervals:
I 1 = { w y ≤ 0 .
44 , v − w x ∈ [ − 0 .
17] , vw x ∈ [0 . 213 , 0 . 4] } .
I 2 = { w y ≤ 0 .
44 , v − w x ∈ [ − 0 .
17] , vw x ∈ [0 . 4 , 0 .
I 30 = { w y ≤ 0 .
44 , v − w x ∈ [ − 0 .
17] , vw x ∈
[Equation]
12 π , 0 . 213] } .
Then, following the above method with splitting and merging intervals, we have
Starting from
I
1 , (a) Step 1: I 1 mapps to I 3 = { w y ≤ 0 .
416 , v − w x ∈ [ − 0 .
162 , 0 . 068] , vw x ∈ [0 .
(b) Step 2: Splitting I 3 , we have i. I 4 = { w y ≤ 0 .
416 , v − w x ∈ [ − 0 .
162 , 0 . 068] , vw x ∈ [0 .
ii.
I 5 = { w y ≤ 0 .
416 , v − w x ∈ [ − 0 .
162 , 0 . 068] , vw x ∈ [0 . 8 , 0 . 9] } .
iii.
I 6 = { w y ≤ 0 .
416 , v − w x ∈ [ − 0 .
162 , 0 . 068] , vw x ∈ [0 . 9 , 1 . 0] } .
iv.
I 7 = { w y ≤ 0 .
416 , v − w x ∈ [ − 0 .
162 , 0 . 068] , vw x ∈ [1 . 0 , 1 . 12131] } .
Then, we
have
i. I 4 mapps to I 8 = { w y ≤ 0 .
214 , v − w x ∈ [ − 0 .
309 , 0 . 0545] , vw x ∈ [0 . 942 , 1 . 25786] } .
ii.
I 5 mapps to I 9 = { w y ≤ 0 .
0966 , v − w x ∈ [ − 0 .
335 , 0 . 0613] , vw x ∈ [0 . 880 , 1 . 19649] } .
iii.
I 6 mapps to I 10 = { w y ≤ 0 .
0756 , v − w x ∈ [ − 0 .
362 , 0 . 068] , vw x ∈ [0 . 777894 , 1 . 11178] } .
iv.
I 7 mapps to I 11 = { w y ≤ 0 .
134 , v − w x ∈ [ − 0 .
394 , 0 . 0782] , vw x ∈ [0 .
(c) Step 3:
Splitting and merging I 8 , I 9 , I 10 , I 11 , we have i. I 12 = { w y ≤ 0 .
134 , v − w x ∈ [ − 0 .
394 , 0 . 078] , vw x ∈ [0 . 595 , 0 .
ii.
I 13 = { w y ≤ 0 .
214 , v − w x ∈ [ − 0 .
394 , 0 . 078] , vw x ∈ [0 .
iii.
I 14 = { w y ≤ 0 .
214 , v − w x ∈ [ − 0 .
362 , 0 . 068] , vw x ∈ [1 , 1 . 11178] } .
iv.
I 15 = { w y ≤ 0 .
214 , v − w x ∈ [ − 0 .
309 , 0 . 061] , vw x ∈ [1 . 11178 , 1 . 25786] } .
Then we have
i. I 12 mapps to I 16 = { w y ≤ 0 . 0372 , v − w x ∈ [ − 0 .
317 , 0 . 061] , vw x ∈ [1 . 14493 , 1 . 31246] } .
ii.
I 13 mapps to I 17 = { w y ≤ 0 . 0432 , v − w x ∈ [ − 0 .
448 , 0 . 078] , vw x ∈ [0 . 943633 , 1 . 24393] } .
iii.
I 14 mapps to I 18 = { w y ≤ 0 .
v − w x ∈ [ − 0 .
462 , 0 . 077] , vw x ∈ [0 .
iv.
I 15 mapps to I 20 = { w y ≤ 0 .
, v − w x ∈ [ − 0 .
456 , 0 . 0785] , vw x ∈ [0 . 550 , 0 . 878] } .
Starting from I 2 , (a) Step 1: I 2 mapps to I 21 = { w y ≤
332 , v − w x ∈ [ −
114] , vw x
∈
] } (b) Step 2: Splitting I 21 , we have
i
I 22 = { w y ≤ 0 .
332 , v − w x ∈ [ − 0 .
205 , 0 . 114] , vw x ∈ [0 . 864 , 1] } .
ii
I
23 = { w y ≤ 0 .
332 , v − w x ∈ [ − 0 .
205 , 0 . 114] , vw x ∈ [1 , 1 . 125894] } .
Then, we have i. I 22 mapps to
I 24 = { w y ≤ 0 .
081 , v − w x ∈ [ − 0 .
336 , 0 . 114] , vw x ∈ [0 . 858 , 1 . 14813] } .
ii.
I 23 mapps to
I 25 = { w y ≤ 0 .
184 , v − w x ∈ [ − 0 .
409 , 0 . 148] , vw x ∈ [0 . 463 , 1] } .
(c) Step 3: Splitting and merging I 24 , I 25 , we have
i. I 26 = {
w y ≤ 0 .
184 , v − w x ∈ [ − 0 .
409 , 0 . 148] , vw x ∈ [0 . 463 , 1] } .
ii.
I
= { w y ≤ 0 .
081 , v − w x ∈ [ − 0 .
336 , 0 . 114] , vw x ∈ [1 , 1 . 14813] } .
Then, we have i. I 26 mapps to
I 28 =
w y ≤ 0 .
083 , v − w x ∈ [ − 0 .
452 , 0 . 148] , vw x ∈ [0 .
ii.
I 27 mapps to
I 29 = { w y ≤ 0 .
034 , v − w x ∈ [ − 0 .
399 , 0 . 133] , vw x ∈ [0 .
Starting from I 30 , (a) Step 1: I 30 mapps to I 31 = { w y ≤ 0 .
44 , v − w x ∈ [ − 0 .
037] , vw x ∈ [0 . 422 , 0 .
767] } (b) Step 2: Splitting I 31 , we have i. I
= { w y ≤ 0 . 44 , v − w x ∈ [ − 0 .
037] , vw x ∈ [0 . 422 , 0 .
ii.
I 33 = { w y ≤ 0 .
44 , v − w x ∈ [ − 0 .
037] , vw x ∈ [0 . 5 , 0 .
iii.
I 34 = { w y ≤ 0 .
44 , v − w x ∈ [ − 0 .
037] , vw x ∈ [0 . 6 , 0 .
Then, we have i. I 32 mapps to I 35 = { w y ≤ 0 .
301 , v − w x ∈ [ − 0 .
218 , 0 . 0185] , vw x ∈ [0 .
ii.
I 33 mapps to I 36 = { w y ≤ 0 . 262 , v − w x ∈ [ − 0 .
245 , 0 . 023] , vw x ∈ [0 . 96322 , 1 . 25093] } .
iii.
I 34 mapps to I 37 = { w y ≤ 0 .
213 , v − w x ∈ [ − 0 .
288 , 0 . 029] , vw x ∈ [0 .
(c) Step 3:
Splitting and merging I 35 , I 36 , I 37 , we have i. I 38 = { w y ≤ 0 .
301 , v − w x ∈ [ − 0 .
288 , 0 . 029] , vw x ∈ [0 .
ii.
I
= { w y ≤ 0 .
301 , v − w x ∈ [ − 0 .
288 , 0 . 029] , vw x ∈ [1 , 1 . 1] } .
iii.
I 40 = { w y ≤ 0 .
301 , v − w x ∈ [ − 0 .
288 , 0 . 029] , vw x ∈ [1 . 1 , 1 . 25093] } .
iv.
I 41 =
{ w y ≤ 0 .
262 , v − w x ∈ [ − 0 .
245 , 0 . 029] , vw x ∈ [1 . 25093 , 1 . 25345] } .
Then, we have i. I 38 mapps to I 42 = { w y ≤ 0 .
0404 , v − w x ∈ [ − 0 .
392 , 0 . 029] , vw x ∈ [0 . 888 , 1 . 11696] } .
ii.
I 39 mapps to I 43 = { w y ≤ 0 . 0740 , v − w x ∈ [ − 0 .
428 , 0 . 033] , vw x ∈ [0 .
iii.
I 40 mapps to I 44 = { w y ≤ 0 . 125 , v − w x ∈ [ − 0 .
482 , 0 . 038] , vw x ∈ [0 .
iv.
I 41 mapps to I 45 = { w y ≤ 0 .
109 , v − w x ∈ [ − 0 .
400 , 0 . 038] , vw x ∈ [0 .
(d) Step 4: Splitting and merging I 42 , I 43 , I 44 , I 45 , we have i. I 46 = { w y ≤ 0 .
125 , v − w x ∈ [ − 0 .
482 , 0 . 038] , vw x ∈ [0 .
ii.
I 47 = { w y ≤ 0 . 074 , v − w x ∈ [ − 0 .
428 , 0 . 033] , vw x ∈ [0 .
iii.
I 48 =
{ w y ≤ 0 .
041 , v − w x ∈ [ − 0 .
40 , 0 . 029] , vw x ∈ [1 , 1 . 11696] } .
Then,
we
have
i. I 46 mapps to I 49 =
w
y ≤ v − w x
∈
[ − vw x
∈
ii.
I 47 mapps to I 50 =
w y ≤
v − w x
∈
[ − 0 .
vw x ∈ [0 . 993 , 1 . 13943] } .
iii.
I 48 mapps to I
=
w
y ≤ 0 .
v − w x ∈ [ − 0 .
0 . 033] , vw x ∈ [0 .
For Case (II):
Now we are to compute the ranges starting from the interval where I = { w y ≤ 0 .
37 , v − w x ∈ [ − 0 .
17] , vw x ∈ [ 14 π , 12 π ] } .
First, we denote it as
[Equation]
Then, following the above method with splitting and merging intervals, we have
Starting from I 52 , (a) Step 1: I 52 mapps to I 53 = { w y ≤ 0 .
37 , v − w x ∈ [ − 0 .
079 , 0 . 0271] , vw x ∈ [0 . 222 , 0 .
(b) Step 2: I 53 mapps to I 54 = { w y ≤ 0 .
343 , v − w x ∈ [ − 0 .
017] , vw x ∈ [0 .
(c) Step 3:
Splitting I 54 , we have i. I
= { w y ≤ 0 .
343 , v − w x ∈ [ − 0 .
017] , vw x ∈ [0 .
ii.
I
= { w y ≤ 0 .
343 , v − w x ∈ [ − 0 .
017] , vw x ∈ [1 , 1 . 24894] } .
Then
we have i. I 55 mapps to I 57 = { w y ≤ 0 .
150 , v − w x ∈ [ − 0 .
305 , 0 . 017] , vw x ∈ [0 . 840 , 1 . 25908] } .
ii.
I 56 mapps to I 58 = { w y ≤ 0 .
137 , v − w x ∈
[ − 0 .
367 , 0 . 022] , vw x ∈ [0 .
(d) Step 4: Splitting and merging I 57 , I 58 , we have i. ii.
I 59 =
{ w y ≤ 0 .
150 , v − w x ∈ [ − 0 .
367 , 0 . 022] , vw x ∈ [0 .
iii.
iv.
I
= { w y ≤ 0 .
150 , v − w x ∈ [ − 0 .
305 , 0 . 017] , vw x ∈ [1 , 1 . 25908 } .
Then, we have i. I 59 mapps to I 61 = { w y ≤ 0 . 0705 , v − w x ∈ [ − 0 .
393 , 0 . 022] , vw x ∈ [0 . 971 , 1 . 304] } .
ii.
I 60 mapps to I 62 = { w y ≤ 0 .
, v − w x ∈ [ − 0 . 421 , 0 . 0219] , vw x ∈ [0 .
For both Cases (I, II):
From I 16 − 20 , I 28 , I 29 , I 49 − 51 , I 61 , I 62 , we can see that it has fallen into an interval I f
= { w y < 0 .
1 , v − w x ∈ [ − 0 .
148] , vw x ∈ [0 .
Something special here is that w y has been much smaller than w x .
More broadly, let’s deﬁne an interval I s generated by I
g = { w y = 0 , v − w x ∈ [ − 0 .
464 , 0 . 148] , vw x ∈ [1 , 1 . 5] } .
Here “generated” means
[Equation]
Then each element ( v, w x , w y ) ∈
I s has the following properties:
w y = 0 .
vw x ∈ [0 . 181 , 1 . 5] .
If vw x ≤ 1 , then v − w x ∈ [ − 0 . If vw x > 1 , then v − w x ∈ [ − 0 .
The ﬁrst property is obvious.
The third can be proven as follows: for each element ( v, w x , w y
) ∈
I g , it has v ( t +1) − w ( t +1) x
= ( v − w x ) (1 + K ( vw x − 1)) , where the ratio 1 + K ( vw x − 1) ∈ [1 , 1 + 1 . 1(1 .
5 − 1)] when vw x ∈ [1 , 1 . 5] .
Furthermore, in the proven 2-D case, we have shown that “if vw x > 1 with some mild conditions, then v ( t +2) − w ( t +2) x v − w x ∈ ( − 1 , 1) ”.
Actually it can be tighter as v ( t +2) − w ( t +2)
x v − w x ∈ ( − 0 .
2 , 1) because here K ≤ 1 .
1 while the original bound is for K ≤ 1 . 5 .
The condition of bounded | v − w x | can also be veriﬁed, the purpose of which is to keep v, w x always positive.
Then the bound [ − 0 .
2 , 1] will tell v − w x ∈ [ − 0 .
474 , 0 . 148] on vw x ≥ 1 , because 0 .0 . 474 > 0 .
148 > 0 .
For the second property, the left bound can be veriﬁed as min x ∈ [1 , 1 . 5] x (1 + 1 .
1(1 − x ))
+ 1 .
1(1 − x ) · 0 . 474
= (cid:18) x
(1 + 1 . 1(1 − x ))
+ 1 .
1(1 − x ) · 0 . 474 2 (cid:19)(cid:12)(cid:12)(cid:12)(cid:12) x =1 .≥ 0 .
The right bound can be veriﬁed as max
[Equation]
735 2 < 1 .
After proving these three properties, we would like to bound how far I f is away from I s .
More precisely, the distance is measured by w y .
We are going to show w y decays exponentially.
Remind the update rules in (80, 81).
Denote γ = 1 π (arctan( α ) − w x w y (cid:107) w (cid:107) 2 ) again and δ = K w 2 y v ( − v 2 + vw y π (cid:107) w (cid:107) 2 ) , then it is
[Equation] [Equation]
Note that both γ and δ are very small, so we are to show their effects separately, which is enough to be a good approximation.
Consider an iteration where v ( t ) w ( t ) x > 1 and the corresponding γ ( t ) .
Let’s denote v ( t +1) , w ( t +1) x as the next parameters with no corruption from γ ( t ) .
Similarly, we denote ˆ v ( t +1) , ˆ w x ( t +1) are corrupted with γ ( t ) .
From the 2-D analysis, we know
[Equation]
We would like to show, with a small γ ( t ) and ignoring δ , ˆ v ( t +2) − ˆ w x
( t +2)
[Equation] [Equation]
where γ ( t +1) is in time ( t + 1) accordingly.
The difference of LHS of the above two expressions turns out to be
[Equation]
Since ∆ w x w x = K vw x ( − vw x + 1 − γ ) , we have
[Equation]
Also we have
[Equation]
Since w ( t +1)
y ≤ w ( t ) y and arctan(
[Equation]
−
[Equation]
arctan
[Equation]
+
[Equation]
≤ m 3 , for any m > 0 , x > 0 , (122) we have
[Equation]
For general vw x ∈ (1 , 1 . 5] , (123) holds as
[Equation] [Equation]
5 +
[Equation] [Equation]
Since 1 + K ( v ( t ) w ( t ) x − 1) ≤ 1 + 1 . 1 ∗ 0 .
5 = 1 .
55 , it is fair to say (118) − (117)
(cid:47) Kγ ( t ) (1 + 1 .
55 ∗ ( − 2 + 22))
+ O ( γ 2 ) = 35 .
[Equation]
Actually γ ( t ) is bounded by
[Equation]
As a result,
[Equation] [Equation]
Note that this small value is very easy to cover in (117), requiring
[Equation] [Equation]
except when vw x is pretty close to 1.
When vw x −→ 1 , from the analysis of 2-D case, (derived from the case of x t +1 y t
+1 ≥ x
[Equation]
s )
[Equation]
For (118) − (117) , denote a function p ( x ) as
[Equation]
where x = v ( t ) w ( t ) x − 1 in (119, 123).
It is obvious that p (0)
=
1 + ( − 2 K + 1) < 0 .
When x is small, it turns out
[Equation]
As a result, (118) − (117) < 0 when vw x − 1 =
o ( K − 1) .
What if vw x − 1 =
Ω
( K − 1) ?
Actually, we can get a better bound by a more care analysis, as
[Equation]
where the last term is due to v ( t +1) w ( t +1) x ≤ v ( t ) w ( t )
x
(1+ K (1 − v ( t ) w ( t ) x )) 2 .
Hence, with this bound, by expanding the last term, (132) becomes
[Equation]
which is deﬁnitely negative because
[Equation]
Meanwhile, we are to prove the δ in (114) will not make ˜
I s make v − w x < − 0 .
474 starting from v − w x ≥ − 0 .
First, in the region of { vw x ∈ [1 , 1 . 5] , v − w x ≤ 0 . 148 } , we have
Kvw 2 y ≤ 1 . 1 · ( √ 1 . 5 + 0 . 074
+ 0 . 074) ∗ 0 .
1 2 ≤ 0 . 0144 .
Also note that in this region with v − w x ≥ − 0 . 462 , we have
[Equation]
Hence Kv ( w ( t ) y ) 2 + Kv ( w ( t +1) y ) 2 ≤ 0 .
0144 ∗ (1 + 0 .
37 2 ) = 0 .
Since | v ( t +2) − w ( t +2) | < | v ( t ) − w ( t ) | if there is no δ , we shall see that there is no need to discuss the case of v − w x ≥ − 0 .
462 + 0 .
=
− 0 .
4456 because it still holds v ( t +1)
− w ( t +1) x > − 0 . 462 .
When v ( t ) − w ( t ) x ∈ [ − 0 .
462 , − 0 .
4456] , we shall see that in (59), after adding the term of δ in v , v ( t +2)
[Equation]
t +2)
[Equation]
(138) which means the absolute value of v − w x decays at least by a margin depending on δ .
After multiplying the current difference v ( t ) − w ( t ) x
on both side, it gives
[Equation]
Note that here v ( t +2) − w ( t +2)
x does not include δ ( t ) and δ ( t +1) .
As stated above, we have δ ( t +1) δ ( t ) ≤ 0 .
37 2 ≤ 0 .
16 due to the decay of w y .
So it is safe to say δ ( t )
+ δ ( t +1) ≥ 1 .
16 δ ( t ) .
Combining with the above inequality, it gives
[Equation]
where
[Equation]
Furthermore, from our previous discussion, w ( t +2)
y
[Equation]
w ( t +
[Equation]
) y gives that the sum of (140) is bounded by
[Equation] [Equation]
Since − 0 .
474 − ( − 0 . 462) < − 0 . 0103 , we shall see that the term of δ cannot drive v − w x < − 0 .
Note that (140) shall include a factor ( < 1) in front of δ ( t ) , but we have ignored it to show a more aggressive bound.
Therefore, we are able to say an Interval ˆ I s generated by I f also has the following properties: for each element ( v, w x , w y )
∈
ˆ I s ,
vw x ∈ [0 . 181 , 1 . 5] .
If vw x ≤ 1 , then v − w x ∈ [ − 0 .
If vw x > 1 , then v − w x ∈ [ − 0 .
Then the decreasing ratio of ∆ w y /w
y is bounded by
[Equation] [Equation]
144) =
[Equation]
Hence, w y decays with a linear ratio of 0 .
97 (or 1 − 0 .
030 K ) at most for Cases (I, II) in stage 2.
For Case (III), in the ﬁrst step of stage 2, it already has w y ≤ 0 .
078 and v − w x ∈ [ − 0 .
So surely it will also converge to I s .
Here we present the time analysis for Case (III) of both stages.
The number of iterations in the ﬁrst stage is apparently similar to that of case (I, II), as
[Equation]
where ψ < 14 π is the value of vw x in the ﬁrst step of stage 2.
In stage 2, since our target is to ﬁnd how many steps are necessary to get vw x ≥ 0 .
181 , so it is
[Equation]
where obviously it still holds w y w x ≤ 2 − √ 3 and w 2 y < 0 .
1 2 in stage 2.
Since 3 . 28 > 2 .
56 , we have the total number of steps to have vw x > 0 .
181 bounded as
[Equation]
which is not beyond the bound for Cases (I, II).
J P ROOF OF M ATRIX F ACTORIZATION
Consider a two-layer matrix factorization problem, parameterized by learnable weights X ∈ R m × p ,
Y ∈ R p × q , and the target matrix is C ∈ R m × q .
The loss L is deﬁned as
[Equation]
Obviously { X , Y : XY = C } forms a minimum manifold.
Focusing on this manifold, our targets are: 1) to prove our condition for stable oscillation on 1D functions holds at the minimum of L for any setting of dimensions, and 2) to prove a convergence result with initialization close to a minimum beyond the edge of stability for the symmetric case Y =
X (cid:62) .
J. A SYMMETRIC C ASE : 1D FUNCTION AT THE MINIMA
Before looking into the theorem, we would like to clarify the deﬁnition of the loss Hessian.
Inherently, we squeeze X , Y into a vector θ = vec ( X , Y ) ∈ R mp + pq , which vectorizes the concatnation.
As a result, we are able to represent the loss Hessian w.r.t.
θ as a matrix in R ( mp + pq ) × ( mp + pq ) .
Meanwhile, the support of the loss landscape is in R mp + pq .
In the following theorem, we are to show the leading eigenvector ∆ (cid:44) vec (∆ X ,
∆ Y ) ∈
R mp + pq of the loss Hessian.
Since the cross section of the loss landscape and ∆ forms a 1D function f ∆ , we would also show the stable-oscillation condition on 1D function holds at the minima of f ∆ .
Theorem 12.
For a matrix factorization problem, assume XY = C .
Consider SVD of both matrices as X = (cid:80) min { m,p } i =
σ x,i u x,i v (cid:62) x,i and Y = (cid:80) min { p,q } i = σ y,i u
y,i v (cid:62) y,i , where both groups of σ · ,i ’s are in descending order and both top singular values σ x, 1 and σ y, 1 are unique.
Also assume v (cid:62) x, 1 u y, 1 (cid:54) = 0 .
Then the leading eigenvector of the loss Hessian is ∆ = vec ( C 1 u x, 1 u (cid:62) y, 1 , C 2 v x, 1 v
[Equation]
y, 1 ) with C 1 = σ y,
[Equation]
y
[Equation]
σ x, 1 √ σ 2 x, 1 + σ 2 y, 1 .
Denote f ∆ as the 1D function at the cross section of the loss landscape and the line following the direction of ∆ passing vec (∆ X , ∆ Y ) .
Then, at the minima of f ∆ , it satisﬁes
[Equation]
Proof.
To obtain the direction of the leading Hessian eigenvector at parameters ( X , Y ) , consider a small deviation of the parameters as ( X +∆ X , Y +∆ Y ) .
With XY = C , evaluate the loss function as
[Equation]
Expand these terms and split them by orders of ∆ X , ∆ Y as follows:
[Equation]
From the second-order terms, the leading eigenvector of ∇ 2 L is the solution of
[Equation]
Since both the top singular values of X , Y are unique, the solution shall have both ∆ X , ∆ Y of rank 1.
Actually the solution is (here for simplicity we eliminate the sign of both)
[Equation]
Equipped with the top eigenvector of Hessian, vec (∆ X , ∆ Y ) , we consider the 1-D function f ∆ generated by the cross-section of the loss landscape and the eigenvector, passing the minima ( X , Y ) .
Deﬁne the function as
[Equation]
Then, around µ = 0 , we have
[Equation]
Therefore, the several order derivatives of f ∆ ( µ ) at µ = 0 can be obtained from Taylor expansion as
[Equation]
Then we compute the condition of stable oscillation of 1-D function as
[Equation]
because all of ∆ XY , X ∆ Y , ∆ X ∆ Y are parallel to u x, 1 v (cid:62) y, 1 and v (cid:62) x, 1 u y, 1 (cid:54) = 0 .
J. S YMMETRIC C ASE : CONVERGENCE ANALYSIS AROUND THE MINIMA
In this section, we focus on the symmetric case of matrix factorization where Y =
X (cid:62) .
Accordingly, we rescale the loss function as L ( X , Y ) =
14 (cid:13)(cid:13) XX (cid:62) − C
(cid:13)(cid:13) 2 F .
Denote the target as C = X 0 X (cid:62) 0 , and assume we are around the minima X 1 =
X 0 + ∆ X 1 with small (cid:107) ∆ X 1 (cid:107) .
Consider SVD as
X
= (cid:80) min
[Equation]
m,p
[Equation]
i =
[Equation]
σ i u i v
[Equation]
i (cid:44)
[Equation]
u
[Equation]
(cid:62) 1 + (cid:102) X 0 .
Then the EoS-learning-rate threshold at X =
X 0 is η = 1 σ 21 .
Therefore, we are to show the convergence initializing from X
= X 0 + ∆ X 1 with learning rate η = 1 σ 21 + β with β > 0 .
Theorem 13.
Consider the above symmetric matrix factorization with learning rate η = 1 σ 21 + β .
Assume 0 < βσ 21 < √ 4 .
5 − 1 ≈ 1 . 121 and ησ 22 < 1 .
The initialization is around the minimum, as
X
= X 0 + ∆ X 1 , with the deviation satisfying u (cid:62) 1 ∆ X 1 v 1 (cid:54) =
and (cid:107) ∆ X 1 (cid:107) ≤ (cid:15) bounded by a small value.
Then GD would converge to a period-2 orbit γ η by a small margin in O ( (cid:15) ) , as
[Equation]
165) where δ 1 ∈ (0 , 1) , δ 2 ∈ ( − 1 , 0) are the two solutions of
[Equation]
Proof.
The update rule of gradient descent gives
[Equation]
Denoting ∆ X t =
X t − X 0 , the update rule is equivalent to
[Equation]
0 ∆ X (cid:62) t + ∆ X t ∆ X (cid:62) t (cid:1) ( X 0 + ∆ X t ) (168)
Consider a decomposition of ∆ X t = (cid:15) t ˜ u t v (cid:62) 1 + (cid:93) ∆ X t where (cid:93)
∆ X t v 1 = 0 and (cid:107) ˜
u t (cid:107) = 1 .
We also control the sign of ˜ u t by claiming (cid:104)
˜
u t , u 1 (cid:105) > 0 .
Then, the update rule is again equivalent to ∆ X t +1 =
(cid:15) t ˜ u t v (cid:62)
+ (cid:93) ∆ X t − η (cid:34) (cid:16) (cid:15) t ˜ u t v (cid:62) 1 + (cid:93)
∆ X t (cid:17) X (cid:62)
+ X 0 (cid:16) (cid:15) t ˜ u t v (cid:62)
+ (cid:93) ∆ X t (cid:17) (cid:62) + (169) (cid:16) (cid:15) t ˜ u t v (cid:62) 1 + (cid:93)
∆ X t (cid:17) (cid:16) (cid:15) t ˜ u t v (cid:62) 1 + (cid:93) ∆ X t (cid:17) (cid:62) (cid:35) (cid:16)
X 0 + (cid:15) t ˜ u t v (cid:62)
+ (cid:93) ∆ X t (cid:17) .
After expanding X 0 = σ 1 u 1 v
[Equation]
1 +
[Equation]
X 0 and projecting ∆ X t +1 onto v 1 , we have
[Equation]
Phase I: ˜
u t gets close to u 1 sharply from a random direction
For now, let’s assume (cid:13)(cid:13)(cid:13) (cid:93)
[Equation]
X
[Equation]
stays small (which we will show later), then ignoring high-order small values gives
[Equation]
where (174) is the projection of ∆ X t +1 onto u 1 v (cid:62)
[Equation]
and (175) is the orthogonal-to- u 1 component of ∆ X t +1 v 1 .
Meanwhile
[Equation]
we have the following iteration of (cid:93) ∆ X t following the update rule,
[Equation] [Equation]
Ignoring higher-order small values, it is
[Equation]
Now we are to verify two facts:
The orthogonal-to- u 1 component of ∆ X t v 1 , denoted as m t , stays small.
Then combining the exponential growth of parallel-to- u 1 component in (174) gives (cid:104)
˜
u t , u 1 (cid:105) → 1 quickly.
(cid:13)(cid:13)(cid:13) (cid:93) ∆ X t (cid:13)(cid:13)(cid:13) stays small.
First is the bound of (cid:107) m t (cid:107) .
From (175), we have m t +
=
(1 − ησ 21 ) m t −
[Equation]
X
[Equation]
∆ X
[Equation]
1 + O ( (cid:15) 2 ) .
Combining with (179) and noticing u (cid:62) 1 (cid:102) X
=
0 , it holds
[Equation]
Furthermore, we have (cid:15) t
[Equation]
X
[Equation]
˜ u t =
[Equation]
X
[Equation]
m t due to (cid:102)
X
[Equation]
u 1 = 0 , so (182) can be rewritten as
[Equation]
Combining (180, 183) gives a form of m t +1 as a combination of all previous terms as
[Equation]
Since our goal is to verify that (cid:107) m t (cid:107) is bounded
[Equation]
pick any eigenvector v p of (cid:102) X 0 (cid:102) X 0 (cid:62) with associated eigenvalue λ p .
Then we have
[Equation]
Obviously (cid:104) m
t +1 , v p (cid:105) shall converge to a series with exponential growth.
Assume the ratio as (cid:104) m t +1 , v p (cid:105) / (cid:104) m t , v p (cid:105)
=
r .The
above equation is equivalent to
[Equation]
The solutions for this equation are r = 1 or r = 1 − ηλ p − ησ 21 , both of which are in [ − 1 , 1] once λ p ≤ (1 − βσ 21 ) /η .
Hence, it is safe to say (cid:104) m t +1 , v p (cid:105) is bounded as non-increasing.
In fact, |(cid:104) m t +1 , v p (cid:105)| is bounded by |(cid:104) m 0 , v p (cid:105)| because |(cid:104) m 1 , v p (cid:105)| < |(cid:104) m 0 , v p (cid:105)| due to the scaling factor βσ 21 < 1 .
Therefore, after picking any eigenvector v p of (cid:102) X 0 (cid:102) X 0 (cid:62) , we can conclude (cid:107) m t (cid:107) ≤ (cid:107) m 0 (cid:107) ≤ | (cid:15) 0 | .
A further result is that |(cid:104)
[Equation]
Notice that
[Equation]
where RHS is a constant due to the proven non-increasing |(cid:104) m t , v p (cid:105)| .
But this constant is proportional to λ p because (cid:104) m 2 , v p (cid:105) + βσ 2 1 (cid:104) m 1 , v p (cid:105) =
ησ 2 1 λ p (cid:104) m 0 , v p (cid:105) ∝ λ p and all further iterations follow a similar factor.
Therefore, we have |(cid:104) m 1 , v p (cid:105)| ∝ λ p .
Now let’s show (cid:13)(cid:13)(cid:13) (cid:93) ∆ X t (cid:13)(cid:13)(cid:13) stays small.
Consider u (cid:62) 1 (cid:93) ∆ X t , (179) gives
[Equation]
so the norm is bounded as, for some σ of singular value of (cid:102) X 0 ,
[Equation]
where Cσ 2 are from the previous discussion of |(cid:104) m 1 , v p (cid:105)| ∝ λ p , which follows σ = (cid:112) λ p .
Hence, it is fair to say (cid:13)(cid:13)(cid:13) u (cid:62) 1 (cid:93)
∆ X t (cid:13)(cid:13)(cid:13) stays small.
Meanwhile, the residual component of (cid:93) ∆ X t that is orthogonal to u 1 on the left, denoted as (cid:94) ∆ X t, ⊥ , iterates following
[Equation]
As a result, due to (cid:13)(cid:13)(cid:13)
[Equation]
X 0 (cid:102) X 0 (cid:62) (cid:13)(cid:13)(cid:13) ≤ σ 22 < 1 /η , the following norm is recursively bounded as
[Equation]
[Equation]
Since (cid:94) ∆ X t +1 , ⊥ is a polynomial of (cid:102) X 0 and its transpose, the above bound tells that (cid:13)(cid:13)(cid:13) (cid:94) ∆ X t +1 , ⊥ (cid:13)(cid:13)(cid:13) does not grow, which means
[Equation]
(cid:94) ∆ X t, ⊥ (cid:13)(cid:13)(cid:13) .
Combining (192, 199), we can conclude that (cid:13)(cid:13)(cid:13) (cid:93) ∆ X t (cid:13)(cid:13)(cid:13) stays small.
After proving that both m t and (cid:13)(cid:13)(cid:13) (cid:93) ∆ X t (cid:13)(cid:13)(cid:13) stay small, from (174), the only term growing fast is (cid:104) (cid:15) t ˜ u t , u 1 (cid:105) exponentially, which means the projection of , ˜ u t onto u 1 is also growing sharply.
Phase II: after ˜ u 1 (cid:107) u 1 approximately
In the ﬁrst phase, (cid:104) (cid:15) t ˜ u t , u 1 (cid:105) grows exponentially with all other components in (cid:93) ∆ X t keeping small.
The consequences of such a growth are
˜
u 1 gets close to u 1 in direction, with the cosine similarity between them growing like cos(arctan(exp( t ))) , where t is a time variable starting from some constant.
while ˜ u t is a unit vector, the growth of (cid:104) (cid:15) t ˜ u t , u 1 (cid:105) makes (cid:15) t grow sharply as well.
In the phase I, the proof is strongly dependent of the assumption that (cid:15) t is small.
But in the phase II, (cid:15) t is not small any more.
We would like to assume that the second consequence comes later than the ﬁrst one, which is feasible once we make the initialization smaller.
Then, we would like to verify the dynamics of (cid:104) (cid:15) t ˜ u t , u 1 (cid:105) when (cid:15) t is relatively large.
After keeping higher-order terms of (cid:15) t in (172) and considering u (cid:62) 1 ˜ u t ≈ 1 , we have
[Equation]
˜
[Equation]
We are to show that u (cid:62) 1 (cid:93) ∆ X t (cid:102) X 0 (cid:62)
˜ u t is small, so that (cid:104) (cid:15) t +1 ˜ u t , u 1 (cid:105) is a function approximately of only η, (cid:15) t , σ 1 .
After re-introducing higher-order terms around (178)
[Equation]
we have
[Equation]
Then, it holds
[Equation]
So we have to ensure that (cid:107) m t (cid:107) keeps small.
To obtain this, we start to re-write the update of m t by m t = (cid:15) t ˜
u t − (cid:104) (cid:15) t ˜ u t , u 1 (cid:105) u 1 , which is
[Equation]
Compared with (180), the above equation has two differences, the ﬁrst is the coefﬁcient of m t −being (cid:0) 1 − η ( σ 2 + (cid:15) 2 t ) (cid:1) instead of (cid:0) 1 − ησ 2 (cid:1) , and the second is additional two terms in (208).
Let’s ﬁrst discuss the second difference.
Actually these two terms somehow act as the coefﬁcient as well, so if (cid:13)(cid:13)(cid:13) (cid:102) X 0 (cid:13)(cid:13)(cid:13) (cid:28)
then it is safe to neglect them and only consider (cid:0) 1 − ησ 2 (cid:1) .
To show this, we again look into the stability of (cid:13)(cid:13)(cid:13) u (cid:62) (cid:94) ∆ X t +1 (cid:13)(cid:13)(cid:13) and (cid:13)(cid:13)(cid:13) (cid:94) ∆ X t +1 , ⊥ (cid:13)(cid:13)(cid:13) .
We have
[Equation]
where the ﬁrst line simply differs from (191) with the coefﬁcient as η ( σ 1 + (cid:15) t ) instead of ησ 1 and the second is the same as (193).
Therefore, following the same arguments in the phase I, if both of (cid:107) m t − 1 (cid:107) , (cid:13)(cid:13)(cid:13) (cid:93) ∆ X t (cid:13)(cid:13)(cid:13) =
O ( (cid:15) ) are small, we have the following step they are still in O ( (cid:15) ) .
Therefore, we have the dynamics of (cid:15) t as
[Equation]
which corresponds to the update rule of gradient descent on a 1D function f ( (cid:15) )
=
14 (cid:16) ( (cid:15) + σ 1 ) 2 − σ 21 (cid:17) 2 with learning rate η .
Since f (cid:48)(cid:48) (0) = 2 σ 21 , the learning rate η = 1 σ 21 + β > 22 σby a margin βσ 21 .
Since for 1D function f ( x ) = 14 ( x 2 − µ ) 2 with learning rate η , it converges to the two positive solutions, (cid:15) , in (39)
[Equation]
Therefore, for 1D function f ( (cid:15) ) =
14 (cid:16) ( (cid:15) + σ 1 ) 2 − σ 21 (cid:17) 2 with learning rate η , it converges to the solutions (one in ( − σ 1 , 0) , one in (0 , σ 2 ) ) of
[Equation]
Along with the above argument of stability (cid:13)(cid:13)(cid:13) (cid:93) ∆ X t (cid:13)(cid:13)(cid:13) =
O ( (cid:15) ) , it concludes that the trajectory converge to the above solution with deviation upper-bounded by O ( (cid:15) ) .
J. Q UASI - SYMMETRIC CASE : WALK TOWARDS FLATTEST MINIMA
Theorem 14 (Restatement of Theorem 4) .
Consider the above quasi-symmetric matrix factorization with learning rate η = 1 σ 21 + β .
Assume 0 < βσ 21 < √ 4 .
5 − 1 ≈ 1 . 121 .
Consider a minimum ( Y 0 = α X 0 , Z 0 = 1 / α X 0 ) , α > 0 .
The initialization is around the minimum, as
Y
= Y
+ ∆ Y 1 , Z 1 =
Z 0 + ∆ Z 1 , with the deviations satisfying u (cid:62)
[Equation]
=
[Equation]
and (cid:107) ∆ Y 1 (cid:107) , (cid:107) ∆ Z 1 (cid:107) ≤ (cid:15) .
The second largest singular value of X 0 needs to satisfy
[Equation] [Equation]
Then GD would converge to a period-2 orbit γ η approximately with error in O ( (cid:15) ) , formally written as ( Y t , Z t
[Equation]
where ρ 1 ∈ (1 , 2) , ρ 2 ∈ (0 , 1) are the two solutions of solving ρ in
[Equation]
Proof.
The update rule of gradient descent gives (for t ≥ 1 )
[Equation]
Denoting ∆ Y t
=
Y t − Y 0 , ∆ Z t =
Z t − Z 0 , for t ≥ 1 , the update rule is equivalent to
[Equation]
Consider the decompositions ∆ Y t = (cid:15) y,t ˜ u y,t v (cid:62) 1 + (cid:93) ∆ Y t , ∆ Z t = (cid:15) z,t ˜ u
z,t v (cid:62) 1 + (cid:103) ∆ Z t where we assume (cid:93)
∆ Y t v 1 = 0 , (cid:103) ∆ Z t v 1 = 0 , (cid:107) ˜
u y,t (cid:107) = 1 , (cid:107) ˜
u z,t (cid:107) = 1 .
We also control the sign of ˜ u y,t , ˜ u z,t by claiming (cid:104) ˜
u y,t , u 1 (cid:105) > 0 , (cid:104) ˜
u z,t , u 1 (cid:105) > 0 .
Then, the update rule is again equivalent to
[Equation]
Phase I: ˜ u y,t and ˜ u z,t get close to u 1 sharply from random directions
Decompose the three matrices at the around-initialization minimum X 0 =
σ 1 u 1 v (cid:62) 1 + (cid:102) X 0 , Y 0 =
ασ 1 u 1 v (cid:62) 1 + (cid:102) Y 0 , Z
= 1 α σ 1 u 1 v (cid:62) 1 + (cid:102) Z 0 .
Obviously they satisfy (cid:102) X
= 1 α (cid:102) Y
= α (cid:102)
Z 0 .
For now, let’s assume
[Equation]
stay small, then ignoring high-order small values gives
[Equation]
then we have
[Equation]
From the above two, we would like to ﬁnd a lower bound of P y,t , P z,t .
Note that
[Equation]
Since 1 − ησ 21 α 2 − η σ 21 α 2 ≤ 1 − 2 ησ 21 < − 1 , we can see | P y,t | is growing exponentially with the ratio of 2 ησ 21 − 1 at least.
Meanwhile, with P z,t − α 2 P y,t ﬁxed along time, it holds | P z,t | is growing exponentially with the same ratio as well.
Now let’s see how other things stay small when P y,t and P z,t are the only two terms growing exponentially.
If that holds, we can conclude that ˜ u y,t and ˜ u z,t get close to u 1 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)
[Equation]
Take any eigenvector v p of (cid:102) X 0 (cid:102) X 0 (cid:62) with associated eigenvalue σ 2 p with σ p > 0 .
Then the above system can be written as
[Equation] [Equation]
Then the above system can be re-written as matrix
A ∈
R 4 × 4 , which maps from
[Equation]
to
[Equation]
The explicit form of A is
[Equation]
So we have a rank-2 matrix B with all elements as non-negative.
Our target is to show A n is bounded with n → ∞ .
Hence, we require the spectral norm (cid:107) B (cid:107) ≤ 2 by Lemma
[Equation]
which gives max
[Equation]
Therefore
[Equation]
this gives an upper bound for σ 2 , as
[Equation]
Hence, with the above discussion
[Equation]
we have shown (cid:107) m y
[Equation]
, (cid:13)(cid:13)(cid:13) u (cid:62) 1 (cid:103) ∆ Z t (cid:13)(cid:13)(cid:13) stay small in this phase.
Meanwhile, the residual components of (cid:93) ∆ Y t , (cid:103) ∆ Z t that are orthogonal to u 1 on the left, denoted as (cid:94) ∆ Y t, ⊥ , (cid:94)
∆ Z t, ⊥ , iterate following
[Equation] [Equation]
Then we have
[Equation]
Then, by reorganizing terms, (cid:107) M t +1 (cid:107) can be well bounded as (cid:107) M t +1 (cid:107) ≤ (cid:107) M t (cid:107) .
Similarly, we need to the bound the norm of the following term
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Equation]
Regarding the above equation, there are several key observations:
(a) N t is symmetric.
(b) For any eigenvector v of (cid:102) X 0 (cid:102) X 0 (cid:62) , we have v (cid:62) N t +1 v = v (cid:62) N t v .
(c) For any distinct eigenvectors v p , v q of (cid:102) X 0 (cid:102)
X 0 (cid:62) , we have v (cid:62) p N t +1 v q + v (cid:62) q N t +1 v p = v
[Equation]
p N t
[Equation]
q +
[Equation]
(cid:62) q N t
[Equation]
Combining these three observations, we have v (cid:62) N t +1 v = v (cid:62) N t v for any v with v decomposed as a linear combination of the eigenvectors of (cid:102) X 0 (cid:102)
X 0 (cid:62) .
Therefore, it is fair to say (cid:107)
N t
+1 (cid:107) =
(cid:107) N t (cid:107) .
So combining (cid:107) M t
+1 (cid:107) ≤ (cid:107) M t (cid:107) and (cid:107)
N t +1 (cid:107) =
(cid:107) N t (cid:107) tells that both (cid:107) (cid:94) ∆ Y t +1 , ⊥ (cid:102) X
0 (cid:62) + (cid:102) X 0 (cid:94)
∆ Y t +1 , ⊥ (cid:62) (cid:107) , (cid:107) (cid:94) ∆ Z t +1 , ⊥ (cid:102) X 0 (cid:62) + (cid:102) X 0 (cid:94) ∆ Z t +1 ,
⊥ (cid:62) (cid:107) stay small, which tells (cid:13)(cid:13)(cid:13) (cid:94) ∆ Y t, ⊥ (cid:13)(cid:13)(cid:13) , (cid:13)(cid:13)(cid:13) (cid:94) ∆ Z t, ⊥ (cid:13)(cid:13)(cid:13) also stay small.
Therefore, we can conclude that both (cid:13)(cid:13)(cid:13) (cid:93) ∆ Y t (cid:13)(cid:13)(cid:13) , (cid:13)(cid:13)(cid:13) (cid:103) ∆ Z t (cid:13)(cid:13)(cid:13) stay small.
After proving that all of (cid:107) m y,t (cid:107) , (cid:107) m z,t (cid:107) , (cid:13)(cid:13) (cid:13) (cid:93)
∆ Y t (cid:13)(cid:13) (cid:13) , (cid:13)(cid:13) (cid:13) (cid:103) ∆ Z t (cid:13)(cid:13) (cid:13) stay small, from (174), the only terms growing fast are (cid:104) (cid:15) t ˜ u
y,t , u 1 (cid:105) , (cid:104) (cid:15) t ˜ u z,t , u 1 (cid:105) exponentially, which means the projections of ˜ u y,t , ˜ u
z,t onto u 1 is also growing sharply.
Phase II: after ˜ u y , ˜ u
z (cid:107) u 1 approximately
Following the same spirit of everything in the symmetric case, we re-introduce higher-order terms, as ∆ Y t +1 v 1 (cid:44) (cid:15)
y,t +1 ˜
u
y,t = (cid:15) y,t ˜ u y,t − η (cid:20) (cid:15) y,t ˜ u y,t
[Equation]
Then we have
[Equation]
Note that the shared factor (cid:0)
[Equation]
Hence, the above system is equivalent to
[Equation]
Furthermore, this is equivalent to f ( y, z ) =
12 ( yz − 1) 2 with learning rate η (cid:48)
= 1 + βσ 21 .
Therefore, from Theorem
[Equation]
we know y, z will converge to the same values
[Equation]
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
[Equation]
where ρ 1 ∈ (1 , 2) , ρ 2 ∈ (0 , 1) are the two solutions of solving ρ in
[Equation]
Note that (cid:107) ∆ Y (cid:107) , (cid:107) ∆ Z (cid:107)
=
O ( (cid:15) ) 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 U SEFUL LEMMAS
Lemma 7.
Assume a · ∆ a ≥ b · ∆ b and a ≥ b .
All of a, b, ∆ a, ∆ b are positive.
If ∆ b ≤ a , then a + ∆ a ≥ b + ∆ b .
Proof.
[Equation]
a
[Equation]
Lemma 8.
In (242), if we have (cid:107) B (cid:107) ≤ 2 and with σ 2 < σ 1 , then A t is bounded entry-wise, for any t ≥ 1 .
Proof.
Denote u
[Equation]
v
[Equation]
(cid:62) , where z 1 , z 2 , z 3 , z 4 are positive normalization terms to ensure (cid:107)
u i (cid:107) = 1 , (cid:107) v i (cid:107) =
1 for i = 1 , 2 .
Then A can be re-written as
[Equation]
Let’s denote a 1 = η σ 21 α
[Equation]
+
[Equation]
α
[Equation] [Equation]
, a
[Equation]
= ησ 21 α
[Equation]
1 + (cid:16) σ p
[Equation]
1 α 2 (cid:17) 2 (cid:19) .
Obviously, for any t ≥ 1 , the t -th power of A , can be represented as
[Equation]
and the update rule of k ( t ) i ’s is
[Equation]
Since A t is bounded ∀ t if and only if k ( t ) i is bounded for any t ≥
[Equation]
we only need to show the below matrix C with (cid:107) C (cid:107) spectral < 1 , with C deﬁned as
[Equation]
where
[Equation]
To obtain (cid:107) C (cid:107) 2 < 1 , we only need to compute (cid:107) C i, : (cid:107) < 1 for i ∈ [4] .
Taking i = 1 as an example, we have
[Equation]
If α < 1 , the above RHS < (1 − η σ 21 α 2 ) 2 + ( η σ 2 p α 2 ) 2 < 1 where the second inequality is due to σ 2 p ≤ σ 22 < σ 21 .
If α ≥ 1 , the condition of a 1 ≤ 2 (from (cid:107) B (cid:107) ≤ 2 ) gives σ 2 p ≤ 2 /η − σ 21 /α 2 α, which is further σ 4 p < 2 /η − σ 21 /α 2 ασ 21 due to σ 2 p < σ 21 .
As a result, it holds ησ 41 − 2 α 2 σ 21 + ησ 4 p α 4 < 0 , which means (1 − ησ 21 /α 2 ) 2 + ( ησ 2 p ) 2 < 1 .
Noting the above RHS ≤ (1 − ησ 21 /α 2 )
+ ( ησ 2 p /α 2 · α 2 ) 2 = (1 − ησ 21 /α 2 )
+ ( ησ 2 p ) 2 when α ≥ 1 , it ﬁnishes the proof.
L I LLUSTRATION 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 √ 5 − 1 ≈ 1 . 236 .
Local convergence is guaranteed if the learning rate is smaller than it.
Conversely, if the learning rate is larger than it, although the period-2 orbit still exists, GD starting from a point inﬁnitesimally close to the orbit still escapes from it.
This is when GD converges to a higher-order orbit.
Figure 9 precisely shows the effectiveness of such a bound where GD converges to the period-2 orbit when η = 1 . 235 < √ 5 − 1 and a period-4 orbit when η = 1 . 237 > √ 5 − 1 .
M D ISCUSSIONS
First, we provide a general roadmap of our theoretical results in Section M.1, as illustrated in Figure 10.
Then, in Section M.2 we discuss three implications from our current low-dimensional settings to more complicated models for future understanding of EoS in pratical NNs, where low-dimension theorems are enhanced with high-dim experiments.
M. C ONNECTIONS BETWEEN THEORETICAL RESULTS
In this section, we discuss the connections between our presented theoretical results, as illustrated in
Figure 10.
Theorem 1 and Lemma 1 present (local) intrinsic geometric properties for a 1-D function to allow stable oscillations.
Such properties provide us the 1-D function f ( x ) =
( µ − x 2 ) 2 and, furthermore,
1.0 x 0.
f ( x ) lr=1.
1.0 x 0.
f ( x ) lr=1.
1.0 x 0.
f ( x ) lr=1.
Local Geometry (Thm. 1)
High
[Figure]
LG (Lem. 1)
[Figure]
Composition rule of g (Prop 1)
[Figure]
we generalize the local property to a global convergence result in Theorem 2.
Then we are to generalize the 1-D analysis to cases of i) multi-parameter, ii) nonlinear and iii) high-dimension.
(a) Multi-parameter .
Compared with 1-D f ( x ) = ( µ − x 2 ) 2 , the 2-D function f ( x, y ) =
( µ − xy ) 2 can be viewed as the simplest setting of two-layer models.
We prove that the 2-D
case converges to the region of x = y in Theorem 5 in Appendix A.1, which means it shares the same convergence as the 1-D model.
Also, x = y means its sharpness is the ﬂattest.
(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 w y → 0 ) in Theorem 3.
Then the problem degenerates to the above 2-D analysis, which means it shares the same convergence with the 2-D, where ( v, w x ) 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 a sharp minima, GD still walks towards the ﬂattest 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 1-D case.
Meanwhile, from Theorem 1 and Lemma 1, we prove a condition for base models g in regression tasks to allow stable oscillation in Lemma 2.
Furthermore, we provide a composition rule of two base models to ﬁnd a more complicated model that allows stable oscillation in Prop 1.
M. I MPLICATIONS 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 ﬂat 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 1, especially its second condition, we wish to discover an 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.
(cid:117) Theoretical: we prove the 1-D condition holds at any minima for non-trivial matrix factorization, shown as Theorem 6 in Appendix A.2.
(cid:93) 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(d), 8(d) in Appendix B.2.2.
(b) With the above intrinsic geometric property, the next question is whether the training trajectory utilizes this property.
(cid:117) Theoretical: in the case of quasi-symmetric matrix factorization, we prove that the training trajectory follows the leading eigenvector of the Hessian (i.e. the leading component of X 0 ) in Theorem 4, where the only top components of weights are changing in ω ( (cid:15) ) .
(cid:93) Empirical: for MLPs on MNIST, we show the almost perfect alignment of the gradient and the top Hessian eigenvector in Figure 6(c), 7(c), 8(c).
(c)
The ﬁnal implication is the implicit bias of EoS after such oscillation.
It turns out GD is driven to ﬂatter minima from sharper minima.
In the 1-D case, obviously there is 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.
(cid:117) Theoretical 1: in the 2-D case in Theorem 5, we prove the two learnable parameters x, y will converge to the same values after oscillations of their product xy .
Actually in the minimum manifold, smaller | x − y | means a ﬂatter minimizer.
(cid:117) Theoretical 2: in the single-neuron ReLU network in Theorem 3 and Prop 2, we show the model degenerates to the 2-D case since w y → 0 .
The 2-D argument tells that this nonlinear model also walks towards the balanced situation, veriﬁed with experiments in Figure 2.
(cid:117) Theoretical 3: in the quasi-symmetric MF in Theorem 4, although the initialization is around a sharp minima, GD is still driven towards the ﬂattest minima where σ max ( Y ) =
σ max ( Z ) .
(cid:93) Empirical 1: for 2-layer 16-neuron ReLU network in a student-teacher setting, it turns out learning rate decay after beyond-EoS oscillations drives the model very close to the ﬂattest minima, as shown in Figure 5 and in Appendix B.2.1.
(cid:93) Empirical 2: for 3,4,5-layer MLPs on MNIST, larger learning rate drives to a ﬂatter minima, as shown in Figure 6(b).
N C ONCLUSIONS
In this work, we investigate gradient descent with a 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, we prove that the necessary 1-D condition holds around any minima.
Furthermore, we conduct an analysis of GD in symmetric matrix factorization, which converges to a period-2 orbit aligned with the 1-D convergence.
Moreover, we generalize the analysis to quasi-symmetric cases where GD walks towards the ﬂattest minimiser although initialized near sharp ones.
A further discussion is provided in Appendix M.
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.
An important item for future work is therefore to 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 be complete without integrating the noise, either in the classic SGD sense or in the labels, as done in Damian et al.
Li et al.