Convergence Rates for Gradient Descent on the Edge of Stability for Overparametrised Least Squares
Authors: Lachlan Ewen MacDonald
Hancheng Min
Leandro Palma
Salma Tarmoun
Ziqing Xu
Rene Vidal
Venue: NeurIPS 2025 poster ()
Keywords: optimization
edge of stability
large learning rate
gradient descent
theory
manifold
sharpness
flat minima
Primary area: theory

Abstract: Classical optimisation theory guarantees monotonic objective decrease for gradient descent (GD) when employed in a small step size, or "stable", regime. In contrast, gradient descent on neural networks is frequently performed in a large step size regime called the "edge of stability", in which the objective decreases non-monotonically with an observed implicit bias towards flat minima. In this paper, we take a step toward quantifying this phenomenon by providing convergence rates for gradient descent with large learning rates in an overparametrised least squares setting. The key insight behind our analysis is that, as a consequence of overparametrisation, the set of global minimisers forms a Riemannian manifold $M$, which enables the decomposition of the GD dynamics into components parallel and orthogonal to $M$. The parallel component corresponds to Riemannian gradient descent on the objective sharpness, while the orthogonal component corresponds to a quadratic dynamical system. This insight allows us to derive convergence rates in three regimes characterised by the learning rate size: the subcritical regime, in which transient instability is overcome in finite time before linear convergence to a suboptimally flat global minimum; the critical regime, in which instability persists for all time with a power-law convergence toward the optimally flat global minimum; the supercritical regime, in which instability persists for all time with linear convergence to an oscillation of period two centred on the optimally flat global minimum.

Convergence Rates for Gradient Descent on the Edge of Stability for Overparametrised Least Squares
Lachlan Ewen MacDonald
Hancheng Min"
Leandro Palma" Salma Tarmoun'
Ziqing Xu" René Vidal
Innovation Data Engineering and Science (IDEAS)
University of Pennsylvania
Pennsylvania, PA 19104
lemaccoonale@@rroommal.ccooo
Abstract
Classical optimisation theory guarantees monotonic objective decrease for gra- dient descent (GD) when employed in a small step size, or "stable", regime. In contrast, gradient descent on neural networks is frequently performed in large step size regime called the "edge of stability". in which the objective decreases non-moottonicalyy with observed implicit bias towards flat minima. In this pa- per, we take a step toward quantifying this phenomenon by providing convergence rates for gradient descent with large learning rates in overparametrissed least squares setting. The key insight behind our analysis is that, as a consequence overparametrisation, the set of global minimisers forms Riemannian manifold M, which enables the decomposition of the GD dynamics into components par- allel and orthogonal M. The parallel component corresponds Riemannian gradient descent on the objective sharpness, while the orthogonal component is a bifurcating dynamical system. This insight allows
derive convergence rates in three regimes characterised by the learning rate size: (a) the suberitical regime, in which transient instability overcome in finite time before linear convergence a suboptimally flat global minimum; (b) the critical regime, in which instability persists for all time with power-law convergence toward the optimally flat global minimum; and (c) the supercritical regime, in which instability persists for all time with linear convergence to an orbit of period two centred on the optimally flat global minimum.
Introduction
The well-known gradient descent (GD) iteration
+ GD' = --
E IP
for minimisation of an objective function É: IP - R. with learning rate y > 0, is the foundation of most practical algorithms for training deep neural networks (DNNs). Despite the apparent simplicity of the algorithm and practical ease with which DNNs can be trained using GD, theoretically GD DNNs remains poorly understood. In addition to the well-known non-convexity of DNN loss functions L, our theoretical understanding is hampered by the fact that in practice DNNs tend to be trained with larger learning rates than permissible according to standard optimisation theory.
Specifically, classical optimisation theory requires that the learning rate 1j be less than twice the reciprocal of the largest eigenvalue À (the sharpness) of the Hessian 2g The reason for this is easily
Equal contribution, alphabetical order
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

seen by considering the simplest quadratic objective E(0) := A of a single variable 8: if > 2/X, then the iterates 0c+1 = (1 7) of GD' (n,-) grow in magnitude like |1 1 leading rapid divergence. In contrast, DNNs are typically trained in a regime called the edge of stability (EOS) which 7 is often (pointwise) strictly larger than 2/A. Miraculously, despite this large learning rate, GD is capable of converging monooonotyyyylll to global minimum of the loss, exhibiting an apparent bias toward less sharp (flatter) minima.
1
Despite significant theoretical progress the past few years, GD at the edge of stability is still poorly understood when compared with its small learning rate counterpart. In partieular, while convergence rates for GD in the small learning rate regime can be obtained for DNNs using standard methods (25), convergence rates for GD the edge stability have so far only been obtained in settings wherein monotonic loss decrease can be eventually guaranteed [42] 41 11 32).
Our purpose in this paper is provide convergence rates in a least squares setting, wherein sustained monotonic decrease of the loss may never occur. To our knowledge, these are the first rates of this kind to be provided in the literature.
Our analysis synthesises Riemannian geometry and dynamical systems theory to formalise an intuition that has been folklore in the literature for some time [16] 12]. namely that GD with large step size implicitly favours flatter global minima. This formalisation is achieved for codimension I, overparametrised least squares problems, including deep scalar factorisation, in which the solutions (global minima) are guaranteed to form a p 1-dimensional Riemannian submanifold M of the p-dimensional Euclidean parameter space IR.. Inspired by the "self-stabilliation" idea proposed in [16), we introduce new coordinates (0l1, ,) for neighbourhood of M, where Øll is point in M and 0 coordinatises the direction orthogonal to M, with respect to which the GD map (0ll, A + (GD ((ll, ()), GD (0l1, ()))) has approximately the following form:
GD (0)), 04, GD n(O0)²
2c(n, A1l) all
GD (t)),0))+ (1 (1-m)()))) + (0-)0. (2)
The left map GD is Riemannian gradient descent on À along M (see 4). with step size m(0-)² (2c(1, ) for some strictly positive number c(n, 0 the right is the normal form for a period-doubling bifurcation 27 Our analysis thereby clarifies the implicit bias of GD toward flat minima:
In a neighbourhood of the solution manifold, GD with a large learning rate implicitly performs Riemannian GD on the sharpness along the solution manifold, and oscillates as bifurcating dynamical system in the direction orthogonal to the solution manifold.
Paper contributions. In the context of codimension least squares problems, we:
1. Make the implicit bias of GD toward flat minima explicit: overparametrisation enables the decomposition of GD into an explicitly sharpness-minimisingg component along the solution manifold, controlled by an oscillating component orthogonal to the solution manifold.
Prove that a class of non-convex scalar factorisation problems fit into our framework. In particular, we prove that despite the concontexyty of the landscape itself, the sharpness in such examples is geodesically strongly convex along geodesic ball in the solution manifold. To our knowledge, this observation has not previously been made.
3. Identify and prove convergence and implicit bias results for large step size GD in three regimes:
(a) the suberitical regime, in which the iterates can be guaranteed to (eventually) decrease monotonically at a linear rate until convergence to a suboptimally flat minimum; (b) the critical regime, in which the iterates converge non-monotonically with a power law rate to the sharpness-ninimising global minimum;
(c) the supereritical regime, in which the iterates converge linearly stable period-two orbit orthogonal to the solution manifold at the point of minimum sharpness.
Although the former and the latter have already been observed in some examples, to our knowledge the middle regime has not yet been identified. Moreover, our work is the first to be able to theoretically quantify the rate of convergence and implicit bias in all three regimes.
4. Verify our theory against numerical experiments.

---

2
Related work
Overparametrisation and convergence guarantees for DNNs. Convergence analyses of GD for non-convex loss landscapes such as DNNs are by now well-established 19 18 3
35 34 36 10]
Although none of these works are able account for non-monotonic convergence of GD using large step size, they do point to overparametrisation as a key feature of DNN loss landscapes which enables (local) Polyak-E.jasiewwn Specifically, overparametrisation enables the map f sending parameters to network outputs to be submersive, in the sense that the derivative matrix DJ of f'is pointwise full-rank. Much of the analysis involved in the aforementioned papers is aimed at quantifying this submersivity by lower-bounding the smallest eigenvalue of the 'neural tangent kernel" Df f [0]. Overparamettristioon in this sense plays a key role in our work too, since it guarantees that the global minima form a smooth manifold with a well-defined normal direction, which is the foundation for Equation (2) and all the analysis that follows.
inequality and thus convergence guarantees in the absence of convexity.
GD with large learning rates. Recent work in convex optimisation has demonstrated surprising benefits of large step sizes [21 22
5], however the mechanism in these works appears to be distinct from the EOS phenomenon in deep learning. Concerning deep learning, although local stability analysis had already hinted at a relation between step size and sharpness prior to 2020 [43] the seminal work of [13 in 2021 was the first conduct systematic empirical study of large step size GD. In particular. [13 exposed the fact that GD on DNNs typically uses a larger step size than admissible by classical theory. Moreover, [13] made the empirical observation that this larger step size implicitly regularises sharpness during the late stages training, and coined the term "edge of stability" for this phase. An explosion of empirical and theoretical work has since been produced in an attempt to
account for these empirical observations 2 7
40
38,
39
16 30
12
26 42
32, 20).
At a high level, work up to this point on large step size GD can be classified into one of two categories: general analysis attempting to outline the essence of the mechanism of stability in this regime while
abstracting from specific architectures [16 14]; and specific analysis focusing on precise architectural
details in an attempt to derive more fine-grained results on convergence [42 41 11 (32] or implicit
bias [88 39 12 [7]. The advantage of the former approach is generality and clarity of insight, with
the disadvantage of less facility in the proof of quantitative results. The ad advantage of the latter is the wealth of tools for the proof of quantitative results, with the disadvantage of obfuscating essential mathematical structures with inessential architectural detail.
Our work attempts to achieve the strengths of both approaches with middle ground of abstraction. On the one hand, in common with and inspired by the seminal work [16). our formulation identifies and hinges on minimal mathematical structures which underlie the dynamics of GD with large step size, independent of architectural details; on the other hand, in common with the latter works, we retain enough mathematical structure in our assumptions to cover some of the specific examples already studied in prior work and obtain quantitative convergence rates and implicit bias characterisations. In particular, we are able to provide convergence rates outside of the "eventually monotonically stable" regimes considered in prior work [42 33], well rigorously quantify oscillatory implicit biases of large step size GD observed empirically and partially accounted for theoretically in [12
3 Theoretical setting
3.1 Notation
Given smooth (i.e., infinitely differentiable, C*) manifold M without boundary, and point 0 E M, we will use TyM to denote the tangent space to M at 0, and TM == Loo ToM to denote its tangent bundle. Given a smooth map g : M - N of smooth manifolds, we will denote by DMg TM > TN its derivative, acting at each point 0 E M as a linear map (m T - N
between tangent spaces.
Riemannian metric on M is a vaooth inner product (+, 1/ on each Mr we will usually drop the subscript and denote the metric by simply (.-).. Given Riemannian metrics on M and N, any higher derivative Digg= T TM + TN is also defined, and at each point 0 E M acts a symmetric, multilinear map D TyMS - g00 N whose value on vectors U1, Ui,...,0% 2% E To M we denote Dj, g (f))[1....... Ux] (respectively, (8) ( if v; v Vi). This Då may also act on vector fields Vi...., Vv: M - TM to give map DiugIVi..... Vil M + TNN defined by the formula
3

---

DirgiVi Vi.... Vil(0)= := (M) Vi(0)] (resp. D O == D (0))))) if V; V vi) for M.
If gg:: M + is scalar-valued, then Ving TM + TM will be used to denote the map obtained by dualising one of the inputs of D using the Riemannian metric. Specifically, given any E M and tangent vectors U2,....;U% E ToM, V (f)).() uxy is the unique element of T& satisfying
(V (o2..), v1) = (0)
Ukl,
Ves E ToM,
(3)
where (: is the Riemannian metric in MgM. The objects VMg and Vig are the Riemannian gradient and Riemannian Hessian respectively. In particular, D and Re- will be denoted simply D*g 9 and V g: g and Våg then are the ordinary Euclidean gradient and Hessian respectively.
A geodesic is a locally length-minimising curve in M (for instance, geodesics in Euclidean space just straight lines). We will use (0r 0') to denote the geodesic distance between 0, 0' M, which is the length of the shortest geodesic connecting 0 to 0'; das makes M into a metric space. For any 0 E .M and any sufficiently small v E M, there that y(0) 0 and (0) = u, the exponential map on (0, u) is then by definition expo(e) == >(1); in Euclidean space exp
unique geodesic y [0, 1] - M such
+ v. See Appendix A for more details. Any smooth function M - IR is associated to Riemannian gradient descent map (n, 0) + GD (n, 0) E M defined an open neighbourhood of (0) x M C Rz2 x M by
GD ( == expo ( 17 Mg()))
This reduces to the familiar (7, 0) + 0 V g(0) in the Euclidean setting, wherein it is denoted simply GD with the subscript dropped.
Recall that a regular value of a function f IP + Rd a point y E IRd such that f f (2) is nonempty and ʃ is C in a neighbourhood of (y} with derivative DS(0)= To + oo Ra surjective for all E f-1 [y). Moreover, if y is a regular value of f. then M == f [y} is a smooth submanifold of IR,, and if is also analytic in neighbourhood of M, then M is an analytic manifold.
3.2 Problem setting
Given a natural number p > 1. function f : R2 + E and target value y R. we aim solve the codimension 1 least squares problem
min ((0) 2
(f(0) - -)²
(5)
using gradient descent with constant step size 1j > 0:
0 GD 0) :- v)()
(6)
We will make the following assumptions on f and y. Assumption 3.1 (Regularity and analyticity). The point y R is a regular value of f, and f is analytic in neighbourhood of the pre-image M = f [y).
The assumption that y is regular value of f is strong version of overparametrisationn: it is equivalent to assuming that the "neural tangent kernel" Df Df" = || is non-vanishing along M. as is often insisted upon in overparametrriation theory for deep learning (4 [0]. Analyticity of f in a neighbourhood of M := f-)())) (y) will be used to invoke powerful results from holomorphic dynamics (6] in our convergence theorems.
The key consequence of Assumption is that, by the regular value theorem, the solution set M of (5) a (p- dimeenssiaa analytie submanifold of Rº, Points in M will be denoted @li in what follows. For any All e M,, the tangent space Tgn M M at All is the kernel of D Df(ll) and any singular value decomposition of ((ll)) admits precisely one nonzero singular value corresponding vector orthogonal to M.
singular M. This singular vector is the normal vector n0lll) := (.(0l f (00)) []-
follows from the chain rule and the fact that f = y along M that the Hessian 2 of l satisfies the identity
V&llM (V/VF" + +f f
IM

---

along M. We denote by À the largest eigenvalue of V which is equal to |V We assume M to be equipped with the Riemannian metric inherited from its embedding into the Euclidean space I."..
Our new coordinate representation (2) for gradient descent is essentially a consequence of coordinatis- ing tubular neighbourhood N of M by ||-corrdinate along M, and L--cordiinae orthogonal to M (Figure []. However, the rigorous form of (2) requires two additional assumptions on the solution manifold M. The first of these provides an invariant line segment about which GD can oscillate, and enables the correct decay rate of the error terms in the ||-component of GD.
Assumption 3.2 (Orthogonal stability). There is e_ . and line segment L through 6!! orthogonal to M which is invariant under gradient descent on (see Figure I).
The second assumption necessary for the rigorous form of (2) concerning the L-equation, is a standard assumption from bifurcation theory enabling the realisation of the L-equation in well- known normal form which is easily analysed [27] Theorem 4.3].
Assumption 3.3 (Genericity). Recall the normal vector field n = V /f/ll At each point f E M, for all 17 in neighbourhood of 2/X(01). ). one has
c(n. 0 = D2 enn ((0l))
" D (1(ll) > 0.
(8)
Assumptions 3.1 3.2 and 3.3 are sufficient to obtain the rigorous form of (Theorem 4.1). As is evident from Equation 2 GD with large step size behaves like gradient descent on the sharpness along M. To utilise this fact in our convergence theorems, we will make the following additional assumptions on M and the sharpness A.
"''
6
6!
0+
Assumption 3.4 (Strong convexity of A). There is a ball By (0.,r) M, centred on 0! and of radius 0 with respect to the geodesic distance, which is geodesically conver in the sense that any two points in (O!, r) are connected by a unique distance-minimising geodesic, over which the sharpness A M + R is u-geodesically strongly conver and has L-geodesically Lipschitz gradi- ents, in the sense that the Riemannian Hessian V M of satisfies
Figure 1: M {(x, : zy = 1} with tubular neighbourhood N (shaded) and line L (dotted). Inside N, any point is closest to a unique point Oil on M, with 8- @l = A nn(ll) orthogonal to M for some 0+ RR. Assumption 3.2 says that Vi(0') is parallel to L at any point O' E C.
plTM I VarA LTTM
uniformly over Bas (B!!,)), r). where ITM is the identity map TMM We assume moreover that \ achieves its minimum value X(0!) in Bar uniquely at e!!, where one moreover has
=
(10)
for some > O, where ITM (e!) is the identity map on the tangent space Toy M of M at o!!
These assumptions need not be satisfied for general functions f. The geodesic strong convexity of in particular seems at first to be alarmingly strong assumption, since it is well-known that DNN loss landscapes are non-convex. Surprisingly, these non-trivial assumptions can be guaranteed to hold for the following class of toy examples of deep learning non-convex loss landscapes.
Multilayer scalar factorisation. The map ʃ RpP - IR defined by ʃ(O',.. ".."....,'"" OP) == 0P 0²
corresponds to a linear network of depth p and width 1 on single input datum. Any nonzero target value y * 0 is a regular value for f. and f {y} is then a union of hypersurfaces (Proposition B. Assuming without loss of generality that > 0 and setting 0. y Ip one takes M to be the connected component of e!! The line L is given by the span of the ones-vector 1 (Proposition B.2). Assumption 3.3 holds by Proposition B.3 while Proposition B.6 demonstrssss that À is geodesically strongly convex on a geodesic ball centred at 0!!, with
((,,( 4-2- 4/p TTMM (e!))
(11)
5

---

so that Assumption holds. One may also take "(0'.....,6) := 8" 10"(((p110p)) .. (o(0')) )). ....,, with é any nonlinearity that is the identity on neighbourhood of y²/r and still satisfy all assumptions.
Some thought reveals that these assumptions can be relaxed in more-or-less straightforward ways to deep linear networks of sufficient constant width on multi-point datasets. However, the scalar case we consider in this paper is already sufficiently instructive and non-trivial that we leave a more general elaboration of this framework to future work.
4 A normal form for gradient descent about the solution manifold
In this section we state rigorous form (Theorem 4.1 of the approximate update equations alluded to in Equation (2). and give an idea of its proof. This section makes use only of Assumptions 3.1 3.2
and 3.3 3.3) The geodesic convexity of Assumption
not needed to derive the normal form of GD about M, and is necessary only for the convergence theorems to be given in the next section.
Our analysis of the dynamics of GD is inspired by that of [6]; it proceeds from a Taylor expansion However, while [16) Taylor expands around a set of points with sharpness S 2/1,, Taylor expand around the solution manifold M. Our derivation is novel and provides number of advantages: the assumptions are more transparent and easily checked, and we are guaranteed that M is Riemannian manifold, while the set of points with sharpness S 2/11 considered in [6] need not be.
Specifically, every point sufficiently close to M admits a unique nearest point øll on M. The difference 0 All then a scalar multiple A$ nn of the normal vector n(@ll) f
01. Thus, pairs (0ll.*)) 0 with Ø) E .M and sufficiently small 0 R suffice completely coordinatise a neighbourhood of M in IPP known as tubular neighbourhood N [28] p. 147]. Our first result is then the following characterisation of the GD dynamics in terms of 0ll, 0-, the sharpness and its Riemannian gradient descent map GD
Theorem 4.1 (Normal form for GD about M). There is an analytic change of coordinates (n, 0ll, 00) - (17 0) for a tubular neighbourhood N of M in which the gradient descent map GD (n, 0) + (n. A v]V(()))) takes the form (17, 00) + (n, GD (ll, ).), GD (l, 0 )), where
m(0-)2
GD (n, (,001,00) G
2c(1, 011) (l1 + 0((0)) min(1, dol ,o!
(12)
GD 001, 0t) = (1 (1-n)\(l)))) + (0")3 + +0((1)))
(13)
with c(n, ) defined in Assumption
3.3
Outline of proof. Recalling the normal vector field n == /f/l| we approximate the gradient (l(()) 0 E N by the Taylor series
(0+)²
Ve(0) = Ve(ell) + 8- (eeell
V ('e((l)((((e))) n )] 0((0))²).
(14)
2
Since M consists of global minima for e, ve(el) = 0. Additionally, since r2(lll) is the sole eigenvector along which v²e(Oll)= = (olll nontrivial, with eigenvalue X(Ol),, one has (2E((0))((( = x(O)) )n()). One thus has:
0 12 ((0)))(( m(ol)) + 0((0 )3)
ve(0) = 00 (0ll)
(15)
Plugging this into the gradient descent update formula, applying the orthogonal projection Pras of TRP onto TM and n² of TR? onto the normal direction respectively, and noting that
PrM ((, n) = VM (V2 = VM
(16)
"The idea of higher order terms inducing implicit bias is also present in literature on Sharpness- Aware Minimisation (8).

---

since v(n"n) = (1) = 0, yields the following formulae for the GD update in (0ll, ))
7(0 )²
@lه - 0
+0((0)
(17)
2
0- + (1-1>1(0l))) n7(0+)2
D3 D2 ( 0((0))))
(18)
The O(00) error in the ||-update is tightened to (( )a min (1, d (tll, )!) using Assumption
3.2 and 0 - n( - V A(O)) ) is the first order approximation of expet (- '0(0")²² V VarA(ol))) by the
Tay
Taylor lor expansion of the exponential map (33. The final form obtained by invoking Assumption
3.3
and applying standard transformation from bifurcation theory [27 Theorem 4.3]. See Theorem
C.6 for a rigorous proof.
We can immediately make the following qualitative observations of the dynamics in Theorem with which we will be able to give an idea of how our convergence theorems work. Let (0, 0t) denote the th iterate of gradient descent in the (oll, 0-) coordinates of Theorem 4.1
The |ccoompoeent of the GD dynamics is perturbed version of Riemannian gradient descent on the sharpness along the solution manifold M, with time-varying step size determined by the magnitude of the L-iterates 0,' Recalling that assume À to be geodesically strongly convex as in Assumption 3.4 we can be assured of descent guarantees for À provided |0z-1 can be controlled.
The behaviour of 10:11 is governed by the map simpler "flip bifurcation"
13p. which is a perturbed, time-varying version of the
I (1- -))m)))== x3
(19)
from bifurcation theory
27 The iterates x, of (19 admit the following dynamics.
1. When 7) < 2/X', is hyperbolic attractor, and the iterates |ate| go zero exponentially fast. Were this to hold also for the iterates |0,-1 of (33, 0! would evolve like gradient descent
Ả with exponentially decaying step size, hence would converge exponentially fast to suboptimally flat point.
2. When 77 2/X',
parabolic attractor, and the iterates |æz| go zero like @(t-1/2). Were this to hold also for the iterates 10,¹ of (13), 8!! would evolve like gradient descent on À with power-aw-ddeaying stepsize, hence would converge with a power-law rate the optimally flat point.
3. When > 2/X, there is hyperbolically attracting orbit of period two with amplitude & Vnd - 2 to which the iterates o, converge exponentially fast. Were the same to be true for the iterates 0,- of (33 o!! would evolve like gradient descent on À with step size + ((n>' - 2)), hence would converge exponentially to the optimally flat point.
In the next section, we provide theorems showing that these intuitions gathered from the unperturbed, time-invarian (19) can be rigorously carried over the perturbed, time-varying system defined by Theorem
5 Convergence theorems
In this section, we state our convergence theorems and provide numerical demonstrations of them on a multilayer scalar factorisation problem. In addition to the Assumptions 3.1 3.2 and reqqured for the normal form (Theorem 4.1), now also invoke Assumption 3.4 to provide rates of convergence. Iterates will be denoted 0, (respectively (o!! 0F)) for E N U (0). For convenience, our 10,- plots concern the orthogonal coordinate of the intermediate Proposition C.5 rather than that of the final Theorem since the transformation going between them analytic (Theorem) C.6), this causes no difference in the convergence rate.
7

---

5.1 Suberitical regime
The suberitical regime is when 2/X(0!) 11 > 2/\(0)) In this regime, considered also in certain examples in prior works 226 32]. the orthogonal component 0F of the iterates exhibits initial, transient oscillation driven by 13 of Theorem
during this time, the sharpness values x(0!) are monotonically decreasing according to Theorem 4.1 to the point of stability 7 < 2/x(0!)) in finite number of steps. Following the achievement of stability, 0z decays exponentially fast, resulting in exponentially less aggressive steps in e!! decrease ). ultimately resulting in convergence to a suboptimally- global minimum of f. See Figure The theorem below appears as Theorem D.4 in the appendix.
Theorem 5.1. Assume that y, < 2/)(().). Then there is a constant ? > O such that for all f sufficiently close zero and all 0g # o! E (0,( if 2/1(0)) ) is sufficientll small, then there is
T T r50 (0t)
( - )²)
->(0!) - 2(o!))
(A(0))
(20)
2/1
such that 7 < 2/)(0)) for all 2 T and n 2 2/\(e!)) for all t < T. Consequently, setting == (2-n )) < 1, the iterates (8!! 0.+) converge global minimum (0z,;)) 0) of & in M
with suboptimal sharpness
x(0 x 2 exp (-0()))911333)-)))))
( Xx).
(21)
at a rate )
10! 18e
10
A(0,1) -
10
10 11
10 0.66
10-077
10-1a
10-12
50
100
50
100
50 Itteratoon (t)
100
Iteration
Iteration (t)
|0,1 ..
10,"
A(o ) xAae
10-049
10-1.1
10-5.66
10-0.79
10-77
10-12
5000
10000
5000 Iteration (t)
10000
0
5000 Iteration (t)
10000
Iteration (t)
Figure 2: Log-scale plots of distance from 8. to e!! (left), magnitude of 0,' (centre) and sharpness suboptimality gap (right) for gradient descent on depth scalar factorisation the suberitical regime. Trajectories from five different initialisations shown. Initial instability in 10-1 (top) is overcome in finite time with rapid convergence to a suboptimally flat global minimum (bottom).
5.2 Critical regime
The critical regime happens when 77 2/X(e!)). To our knowledge, this regime has not been observed previously in the literature. In this regime, the dynamics of GD are non-hyperbolic, with 0,- exhibiting 2-periodic decay to zero with rate ( Consequently, the dynamics of 0, are essentially those of gradient descent on À with a step size decaying according to power-law. Ultimately then, O, does converge to the optimally flat global minimum of é, but does so at a slow rate. See Figure The theorem below appears Theorem D.9 in the appendix.
8

---

Theorem 5.2. Assume that 14 2/x(0!) and that // v/(c((/X,,00) 0!)2,)) < 1, where and .) are
defined as in Assumptions 3.3 and 3.4 respectively. Then for all 0f sufficiently close to 0! and all 0e + 0 sufficiently small, one has d (0), 0..) (t= and 10 = 0
|0,1 0.1
10,1
A(0 A(8.)
10
10
10
10
10 2.2
10 2.6
5000
10000
5000 Iteration (t)
10000
5000 Iteration (t)
10000
Iteration (t)
Figure 3: Log-y-scale plots of distance from 6, to e!! (left), magnitude of fJ (centre) and sharpness suboptimality gap (right) on depth 5 scalar factorisation in the critical regime. Trajectories from five different initialisations shown. The iterates 10, may or may not initially increase; in both cases, asymptotic, power law decrease of 10,4| and ||0, - e!ll to zero indicates power law convergence to the optimally flat global minimum.
Some further remarks about the critical regime are necessary. First, the additional hypothesis that v/(c(2/>_, e. )A.) < 1 (which is satisfied for multilayer scalar factorisation, see Proposition B.7 not strictly necessary for convergence; power-law convergence still occurs the same minimum without this assumption, but is in that case faster for dır (0) 0..) than the ( rate given in the theorem statement. However, we are unaware of examples having this faster rate. Second, the 10, iterates may initially increase substantially in Figure depending the initialisation; our theorem accommodates this behaviour explicitly (see Theorem D.8
5.3 Supereritical regime
The supercritical regime is when > 2/A(O!). In this case, when > 2/\(0!)) is sufficiently small, Equation (13) exhibits linear convergence to a stable orbit of period two with amplitude = (nA((!)) - 2).. follows that the step sizes in 8!! in its descent on À are of approximately constant size, that 0/ will converge linearly to the optimally flat o!. However, the complete iterates O of GD do not converge to a global minimum, but asymptotically oscillate orthogonally to M about the point e!! along the line L of Assumption 3.2 See Figure 4 Although this regime has been observed in prior works [12 [7 [7[20] 20], no general quantitative convergence theorem in this regime has yet been proved. The theorem below appears Theorem i.10 the appendix.
Theorem 5.3. Assume that > 2/2(0l) is sufficiently small. Then there are positive constants C2, C3 and stable orbit of period two of the form (o!!, (m)(0))) 22 + O(n2(o!!) - 22) to which the iterates (0l), 0+) starting from any (0), 0t)) satisfying 0 < |0t| s Ci(AA()! - 2) and
d (0%, o.) S Ca(r))(o) - 2)2/2 converge with rute - (C(0))C()))) -
Limitations, discussion and conclusion
Our work has a number of limitations, all of which point toward avenues for further exploration along the lines we have developed.
Higher order orthogonal dynamics. Prior work has indicated that beyond the 2-periodic behaviour we studied in this work, hidher-order periodicity and chaos emerge the learning rate increases [77] Our Theorem 4.1 provides new insight into this behaviour as manifestation of bifurcating system oscillating orthogonally to the solution manifold. Our analysis only handles the simplest (namely 2-periodic) case of this supercritical behaviour. We expect our framework to be able to admit convergence proofs of higher oscillatory and chaotic behaviour also, using ideas from dynamical systems theory, however we do not expect this to be easy and have not attempted it in this paper.

---

|8,1 -
A(8 -
10-240
10
10-5.66
10-5
10
10-9.6
5000 Iteration (t)
10000
5000 Iteration (t)
10000
5000
10000
Iteration
Figure 4: Log-scale plots of distance from 6, to l! (left), magnitude of 0, (centre) and sharpness suboptimality gap (right) on depth 5 scalar factorisation in the supereritical regime. Trajectories from five different initialisations shown. The iterates |0,- converge stable, period-two orbit driving linear convergence of 6! to the optimally flat global minimum e!!
Higher codimension least squares. Our theory only allows us to treat codimension 1 problems. This is clearly severe limitation, ruling for instance overparametrised regression multiple datapoints or multiple outputs. It is not difficult to extend our definitions to higher codimension, however proving anything in this setting seems very challenging, as it would require understanding of higher-dimensional bifurcating dynamical systems. We leave this question to future work.
Geodesic convexity of sharpness for more general solution manifolds. One of the most surprising results of our work is that the sharpness of the loss is geodesically strongly convex over geodesic ball in the solution manifold for number of overparametrised non-convex problems. Should this prove to be a more general fact, which we suspect may be the case, it would go a long way to explaining the still-mysterious ease of optimisation and implicit bias of GD DNNs. We leave more thorough investigation of this question to future work.
Finding the tubular neighbourhood. Our theory is premised on the assumption that GD is initialised tubular neighbourhood of the solution manifold. While there no reason to think this the case with standard initialisation schemes used in practice, we conjecture that GD does converge rapidly to a tubular neighbourhood from a standard initialisation during the phase known as progressive sharpening [13), after which the familiar dynamics described here would apply. We leave exploration of this question to future work.
Relation to stochasticity. Prior works strated an implicit bias toward flat minimisers arising from stochasticity. All of these works consider smaller learning rate than those considered in our work; consequently, the mechanism behind the fundamentally different than that considered in our work. Studying how these distinct implicit biases interact is an important direction for future research.
15 31 studying stochastic gradient descent have demon-
implicit bias considered in 2 15 31
Acknowledgments and Disclosure of Funding
This research was supported by NSF 2031985 and Simons Foundation 814201 (Theorinet), ONR MURI 503405-7805i and University of Pennsylvania Startup Funds.
References
[1] Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. In ICML, 2023.
[2] K. Ahn, J. Zhang, and S. Sra. Understanding the unstable convergence of gradient descent. In ICML, 2022.
[3] Z. Allen-Zhi Y. Li, and Z. Song. A Convergence Theory for Deep Learning via Over- Parameterization. ICML, pages 242-252, 2019.
[4] M. Altschuler and Parrilo. Acceleration by Stepsize Hedging: Multi-Step Descent and the Silver Stepsize Schedule. Journal of the ACM, 2023.
10

---

[5] J. M. Altschuler and P. Parrilo. Acceleration by stepsize hedging: Silver Stepsize Schedule for smooth convex optimization. Mathematical Programming, pages 1-14, 2024.
[6] M. Arizzi and J. Raissy. On Eccle-hakim' theorems on holomorphic dynamics. Frontiers in Complex Dynamics, 2014.
[7] Arora, Z. Li, and A. Panigrahi. Understanding Gradient Descent on Edge of Stability in Deep Learning. In ICML, 2022.
[8] P. L. Bartlett, P. M. Long. and O. Bousquet. The Dynamics of Sharpness-Awan Minimization: Bouncing Across Ravines and Drifting Towards Wide Minima. Journal of Machine Learning Research, 24:1-36, 2023.
[9] G. Blanc, N. Gupta. G. Valiant, and P. Valiant. Implicit regularization for deep neural networks driven by an Prostein-Uhlenbee like process. In COLY, 2020.
[10] S. Bombari, M. H. Amani, and M. Mondelli. Memorization and Optimization in Deep Neural Networks with Minimum Over-parameterrizaion. In NeurlPS, 2022.
[11] Y. Cai, Wu, S. Mei, M. Lindsey, and L. Bartlett. Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization. In NeurIPS, 2024.
[12] Lei Chen and Joan Bruna. Beyond the Edge of Stability via Two-step Gradient Updates. In ICML, 2023.
[13] J. Cohen, S. Kaur, Y. J. Zico Kolter, and A. Talwalkar. Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability. In ICLR, 2021.
[14] Jeremy Cohen, Alex Damian, Ameet Talwalkar, J Zico Kolter, and Jason D. Lee. Understanding Optimization in Deep Learning with Central Flows. In ICLR, 2025.
[15] A. Damian, T. Ma, and J. Lee. Label Noise SGD Provably Prefers Flat Global Minimizers, In NeurIPS, 2021.
[16] A. Damian, E. Nichani, and J. Lee. Self-Sttailizztiion: The Implicit Bias of Gradient Descent at the Edge of Stability. In ICLR, 2023.
[17] Dayal Singh Kalra and Tianyu He and Maissam Barkeshli. Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos. In ICLR, 2025.
[18] S. Du, Lee, Li, Wang, and X. Zhai. Gradient Descent Finds Global Minima of Deep Neural Networks. In ICMI,, pages 1675-1685, 2019.
[19] S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient Descent Provably Optimizes Over- parameterized Neural Networks. In ICLR, 2019.
[20] Avrajit Ghosh, Soo Min Kwon, Rongrong Wang, Saiprasad Ravishankar, and Qing Qu. Learning dynamics of deep matrix factorization beyond the edge of stability. In ICLR, 2025.
[21] B. Grimmer. Provably faster gradient descent via long steps. SIAM Journal Optimization, 34:2588-2608, 2024.
[22] B. Grimmer, K. Shu, and A. L. Wang. Accelerated gradient descent via long steps. arXiv:2309.09661. 2023.
[23] B. Grimmer, K. Shu, and A. Wang, Accelerated objective gap and gradient norm convergence for gradient descent via long steps. INFORMS Journal on Optimization, 7:156-169, 2025.
[24] A. Jacot, F. Gabriel, and C. Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In NeurIPS, pages 8571-8580, 2018.
[25] H. Karimi, J. Nutini, and M. Schmidt. Linear Convergence of Gradient and Proximal-Gradies Methods Under the Dojyak-Lojasiewicz Condition. In ECML PKDD, pages 795- 995-8811, 2016.

---

[26] Itai Kreisler, Mor Shpigel Nacson, Daniel Soudry, and Yair Carmon. Gradient descent mono- tonically decreases the sharpness of gradient flow solutions in scalar networks and beyond. In ICMI., 2023.
[27] Y. A. Kuznetsov. Elements of Applied Bifurcation Theory, 4th Edition. Springer, 2023.
[28] J. M. Lee. Introduction Smooth Manifolds, Second Edition. Springer, 2013.
[29] J. M. Lee. Introduction Riemannian Manifolds, Second Edition. Springer, 2018.
[30] Lee and C. Jang. new characterization of the edge of stability based on sharpness measure aware of batch gradient distribution. In ICLR, 2023.
[31] Z. Li, T. Wang, and S. Arora. What Happens after SGD Reaches Zero Loss? -A Mathematical Framework. ICLR, 2022.
[32] Liming Liu, Zixuan Zhang, Simon Du, and Tuo Zhao. A minimalist example of edge-of-stability and progressive sharpening, 2025.
[33] Maria G Monera, A Montesinos-Amilibi and Esther Sanabria-Codesal The taylor expansion of the exponential map and geometric applications. Revista de Real Academia de Ciencias Exactas, Fisicas y Naturales. Serie A. Matematicas, 108:881-906, 2014.
[34] Q. Nguyen. On the Proof of Global Convergence of Gradient Descent for Deep ReLU Networks with Linear Widths. In NeurlPS, 2021.
[35] Q Nguyen and M. Mondelli. Global Convergence of Deep Networks with One Wide Layer Followed by Pyramidal Topology. In NeurIPS, 2020.
[36] Q. Nguyen, M. Mondelli, and G. Montufar. Tight Bounds on the Smallest Eigenvalue of the Neural Tangent Kernel for Deep ReLU Networks. In ICML, 2021.
[37] P. Petersen. Riemannian Geometry. Springer, 2016
[38] Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large Learning Rate Tames Homo- geneity: Convergence and Balancing Effect. In ICLR, 2022.
[39] Yuqing Wang, Zhenghao Xu, Tuo Zhao, and Molei Tao. Good regularity creates large learning rate implicit biases: edge of stability, balancing, and catapult. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
[40] Z. Wang. Z. Li, and J. Li. Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability. In NeurlPS, 2022.
[41] J. Wu, P. L. Bartlett, M. Telgarsky, and B. Yu. Large Stepsize Gradient Descent for Logistic Loss: Non-Monotonicity of the Loss Improves Optimization Efficiency. In COLT, 2024.
[42] Wu, V. Braverman, and J. Lee. Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stability. In NeurIPS, 2023.
[43] L. Wu, C. Ma, and W. E. How SGD Selects the Global Minima in Over-parameterrzzed Learning: Dynamical Stability Perspective. In NeurIPS, 2018.
[44] X. Zhu, Z. Wang, X. Wang, M. Zhou, and R. Ge. Understanding Edge-of-Stability Training Dynamics with Minimalist Example. In ICLR, 2023.
12

---

A Differential geometry background
The purpose of this section is to give a brief overview of essential notions from Riemannian geometry. A good source for the first subsection is [29]. We have not been able to locate a good source for the latter subsection; in particular, to our knowledge no attempt has yet been made consider the Riemannian Hessian of the square gradient norm of a function defining a hypersurface, however the calculation elementary.
A.1 Riemannian metrics, connections and geodesics
Recall that a Riemannian metric on a manifold M is a smoothly-varying family of positive-definite inner products (-, 1o on the tangent spaces TyM of M, with associated norm l1 Any submanifold M of a Riemannian manifold (N, gNv inherits a Riemannian metric from N; precisely, given tangent vectors v, w E M C ToN, one defines
gM(v, w) = gN(u, w).
(22)
Since we work with submanifolds of Euclidean space, whose tangent spaces are all themselves subspaces of Euclidean space equipped with the Euclidean inner product, we will often use the notation (; and || without further decoration when considering the metric of a submanifold of Euclidean space.
Denote by T*M the cotangent bundle of M, whose fibre over A E M is the space of linear functionals T' + IR. Given (k, N U [0} (k, ()-ensor at the point 0 E M is an element of ToM* @ TV M and (ks ()-tensor field is smooth section of the bundle TM @ T M that is, a smooth assignment of a tensor T()) E T&M &k @ T M @l to each 0 E M. The space of such sections will be denoted [ @k @ T M
Whether or not M carries a Riemannian metric, the derivative of a map f M - R makes sense as a map D f TM + TR which is fibrewise- -reeare-linearr in the sense that its restriction f TeM- + (T) R each fibre is linear map. The derivative Darf f can in this sense thought of as a (0, 1)-tensor field. Higher derivatives D of f are thus derivatives of tensor fields, and can be defined using a Riemannian metric. Specifically, assuming that M has a Riemannian metric, there is distinguished derivative operator D : I T TMM M + [TM @ T* M+++1 defined for all (k, 1) E NU U {0} called the Levi-Civita connection. The additional T M-slot gained by applying Dar to a tensor T is to be thought of as a direction in which T is differentiated by DM. In particular, given T E [[T (TM T (S) and X E (TM)), the notation Das T[X] M@T & T M is the directional derivative of in the direction X. Note that when M = RRP D == DR= is the usual multivariate derivative, and the higher derivatives D D coincide with the usual higher-order derivatives in Euclidean space. If g M + R is a smooth, scalar-valued function, then v rig E I 1 M@k-1) will be used to denote D with one of the T M-slots dualised to a T'M-slot as in 33. In particular, V and Vig are the Riemannian gradient and Riemannian Hessian respectively; they coincide with the usual gradient Vg and Hessian V 9 respectively when M- = RR..
A special case of the Levi-Civita connection is that inherited by submanifold. If M is submanifold of a Riemannian manifold N, then the Levi-Civita connection associated to the induced metric (22) is Da Prus Dx where Pr is the orthogonal projection TN TM induced by the metric N.
Given (0, v) E TM, classical ordinary differential equation (ODE) theory guarantees the existence of a unique M-valued solution y to the TM- second order ODE
(D = 0, 70) = 8 700)
(23)
on some interval e0. This unique solution is called the geodesic through (0, In particular, for all v E To M sufficiently small, 711) makes sense and is the value of the exponential map expg The exponential map exp, : BT.M (0, R) + M is invertible onto its image expy (BT. M (0, R)) C M for all R > 0 sufficiently small, and its inverse, denoted logs expe (BT,M(0, (0, R)) + BT, (0, R),, is called the logarithm. When M = Ip with the Euclidean metric, (x) is defined for all (0, v) TRp and is equal to 0 + The length of geodesic y on [0. is given by fo dt)ldt and the geodesic distance a (0, 0') between two points 0, 0' E M is the length of the shortest
13

---

geodesic connecting & to 6'. Associated to any geodesic 7 [0. bj - M and any t E [0, b] is unique orthogonal map (T), To M + + M such that (I()) : H[())) satisfies the ODE
(DML( 0, 11(7) to) M-
(24)
This map IT(7), is called the parallel transport morphism, and extends map Ty(o) M % ((() M - -
M @ Ticyy M of arbitrary tensors. In particular, if 0, . admit unique geodesic 7 [0, + M with 7(0) = 0 and 76) = 0', (.()) TeM + Ter M will be denoted simply 18
Tensor fields admit covariant Taylor expansions on M defined as follows. Fix E M and suppose that 0' MM is sufficiently close to 9 that there is unique geodesic Y [0, - M such that 7(0) and 7(1) = 0' Consider then the function T(t) == (T>))=((((())) To 60 Tg M Taylor's & theorem in a single variable implies that
T(1) = T()) DT + 1 D +
(25)
Since (+) = and (b = 0, one has:
(*) (T(((T)) (1)
(26)
(IT()), ((()))))))))))))
(27)
((O)
(28)
Substituting these formulae into 25 yields the covariant Taylor expansion
1
T(0') = -oo (())) Dar Dz (0)) + (0) (0') + O(I loga (0')
(29)
A.2
Riemannian geometry of submanifolds and hypersurfaces in Euclidean space
When M is hypersurface Euclidean space IR,, with normal vector field the Hessian of smooth function takes the following form.
Lemma A.1. Let M be smooth hypersurface in TR with normal vector field The Riemannian Hessian V2 À of a smooth map X RP - IR along M is the p x matrix-valued function
PV (n, DAMD
(30)
on M.
Proof. When M is an embedded submanifold of RP, the action of the Levi-Civita connection function À (i.c.. the Riemannian gradient operator) is given by Pr VA, where V is the ordinary
Euclidean gradient. The action of the Levi-Civita connection a vector field X Rp R' is given by PrM DX PrM- Thus:
V3 A
PrM
PVM - Dn(n, VA) n((Dn) - nn m Prm
(31)
(PrM (n, VA)Dn - ((D)) PrM (VPA (n, VA)Dn)
(32)
(33)
where the final line follows from
PrM (nw")) = nn nn- (nn=) not not =0
(34)
for any vector
We will be particularly interested in the case hypersurface M = f {g} CR, where f R IR is smooth submersion, with metric inherited from the ambient Euclidean space. Specifically, when n is the unit normal vector field defined by
(()))
(0)) =
E M,
(35)
14

---

the Riemannian metric on M is defined by
( = P0) n())(00)
OE M
(36)
where we use Pr (0) denote the projection onto T.M. We will be particularly interested in computing the Riemannian Hessian of the function À = [V which admits the following formula entirely in terms of f and its derivatives.
Lemma A.2. For hypersurface M f- {y} of Euclidean space defined by smooth function f: IP 7 R, the Riemannian Hessian of X == [V is given by
TM=2PPM ((V) IV (V 5) (n, fn)
(37)
where f is regarded as a map R + RexP, or as the map sending a direction vector v to the directional derivative of the Hessian matrix 2 in the direction v.
Proof. By Lemma A.1 we have that
P = PRM (v>> (n. VA)
(38)
so it remains only to compute the various quantities. The following identities are elementary to verify:
= = T
(39)
V² / (V
(40)
and
A |
Dn= PTM
(41)
The result now follows.
Our next lemma allows us lower-bound the convexity radius conv(0) at 0 E M which is the largest r > 0 such that the geodesic ball (9, r) is geodesically convex. This will be used in verifying Assumption for multilayer scalar factorisation. If M is hypersurface in R2 with normal vector field n, then the second fundamental form M is the map h : TM @ TM - RR defined by
(O)) w| = - (D)
(42)
for all 0 E M and w E T.M.
Lemma A.3. Let M be smooth hypersurface in IP with normal vector field n, and fix 0, E M. Assume that for all R > 0, there exists CR > 0 such that
[h(O) rll Cr
(43)
for all 0 (0,,0., 2R) and all orthonormal vectors u, v ToM. Then for all R 0, one has
1 min
conv(O.) 2
CR R RJ
(44)
Proof. The bound follows from a well-known relationship between the convexity radius and the injectivity radius, which, at point 0 M, the largest such that expy : BT, (0,r) M is injective, We thus first bound the injectivity radius. Fix R > O. By the Gauss equation and the hypothesis, at any point B(e., 2R). the sectional curvature ((#) of M satisfies
K (0) [u, w]] = [[(0)[u, u| M(0),,, o] - h()) j,, w| < C₁
(45)
for all unit, orthogonal pairs u, v E M. Consequently, by Klingenberg's 's formula 37 Lemma 6.4.7].
the injectivity radius (0) at 0, namely the largest value of such that exp, BT.M (0,r) Al is injective, satisfies
inj(@) 2 min
π
I(O)
(46)
CR
15

---

where L(0) is the length of the shortest geodesic loop based at A. If ( 2 2R, then this simply
yields ijj 2 min(=/CR- R). If L(Ø) < 2R, then the shortest geodesic loop Y [0. L(>]] + M at
is contained entirely in B(0,, 2R). Fenchel's theorem then combines with the hypothesis to give
L.
L(0)
2x
CR
(AAO)) 2 T(7))
(47)
so that (1/) 2 x/Cn. Consequently, one has the bound
Cn m
inf inj(0) 2 Cn (EB
(48)
in general.
Now, setting
Pn min
Cn RJ.
(49)
2
one sees that the hypotheses of Theorem 6.4.8] hold, implying that ((0,, rn) is geodesically convex.
B Case study: deep scalar factorisation
In this section, we verify that all of our assumptions are satisfied for deep scalar factorisation, where f: IR + R is defined by
f(e1..... 0p) 01 ....ppp.. 8pp = -f1,...,,
... 0p) E R..
(50)
For RP having all entries nonzero, will be convenient to denote
v(0) = (
0 27
(51)
We first demonstrate that Assumption 3.1 satisfied.
Proposition B.1. Any y + a regular value of f. Consequently, Assumption 3.1
/ is satisfied.
Proof. Since y * any point All E {y) has all coordinates being nonzero. Hence Df( y ((ll)) + 0 making y regular value of f.
The pre-image manifold fyy of any y * 0 has several connected components. Without loss of generality, will assume from hereon that y > 0 and denote by M the component of (y) contained in the positive orthant. We next demonstrate that Assumption 3.2 is satisfied.
Proposition B.2. Set 8!! := /11 where 1 is the vector of Is. Then the line spanned by the normal vector n(!!!) ) is invariant under gradient descent on f = (1/2)(f - y) for any y, so that Assumption 3.2 is satisfied.
Proof. At 00 = I one has n((!!) = p 1/2 For any a E IR, one sees that
v((an((!)) (f(ap-1/2) y) /f(pp-11111 (f(ap-1/2) 1/2 y) a=- p (0- 1)/2 1p
(52)
(53)
is a multiple of n(!!),, implying that the span of n n(!!) is invariant under gradient descent on l.
We will continue to denote o! : 1 /p I, which is often referred in the literature as the "balanced"
solution. We next come to Assumption 3.3
Proposition B.3, Assumption 3.3 holds along M. In particular, at o!!
<2/ o!!
32 p
y
2/P
(54)
3
2
pa
where the convention (3) 0 i 3.
16

---

Proof. Given symmetric, multilinear functionals A, B taking & and inputs respectively, let A @ B denote their symmetric product:
(A o eI......... UK+il =
Σ Alva(1):..... VO(R) (l Bo(k+tj) (55)
(A D)!
oEPerm(k+-))
where Perm(k: (k + 1) is the group of permutations of ++ elements. Since & = (1/2) - v)² one has
DE = (J- = De =0,
(56)
Df + (f- f) = DP Df @ Df 3V>=3 DS 1 + (J- = DA 3 DJ
(57)
(58)
D'L D' Dzf (f (f y D' = 1D' Df f D
(59)
One then sees that
c(n, 0 =
1
p$
83
(01)
6 "'
(60)
2
3m D² f(n
Df
2
(el) o/ AD fln 33 + 3(D f(n (8) (61)
2
? 4 ID5(0)))122 ") (D² ((0l))(((()))))
4n DI )| D2 ((0l)) (((l))
By continuity, it suffices to prove the claim at 17 = =/ 1/ )| 2 substituting this into the above yields
c(2 )(2/X(0)),
8
D2 f(o) )[((0))8992
8 G1)DS(0111
D2 f(0) Iin(ol) soy
(63)
DDSO
One computes
Dfl((ll) 0.'
(64)
D2 f1112 =(1 11-6i -611)) 01, 012
(65)
y [(1 - 61,12)( ō1,1,)(1 digts 01, 01g
if 3
D3 (01tt1
(66)
L
Consequently,
D2 (((()))
D2 ((
(67)
(01))
(1
8
y 01,
(68)
|DF
01, Or O1p
h1.2
+ O.
(69)
!DSON
02
4
If p = 2, then, since DdJ = O, Assumption 3.3 trivially holds for all Øll E M and all close to 2/A(01). Let us assume then that p 2 3. One has:
Dd D00( (1
D3 s(oll MD [
(70)
Σ (1 - (1 -811) - (ists) 02, 01, Pi, Oi, Ot, Or,
y y (71)
!DS[O1) 4 helaa
Σ
02
(72)
17

---

Thus one has
d2/A(("),
32y
of oF }
Σ
2
8y²
Σ
(73)
of 00
4
Substituting
P Σ
y²
12100))2
(74)
09
1=1
then yields
c(2/X(01)), 32
Σ
2-20 0020²
0
(
L 012)
-3
-8 8
20100
20
11
(75)
-3
2
(Σ 8,²
Σ
07²
32
Σ
20/
Σ
0 20
112/2
(76)
(0S) 3
32 P
(0(l)) -8p
S
(ell) (ell)) ss(ell)))
(77)
where S, ( are the elementary symmetric means in the variables 0, Newton's inequality for the elementary symmetric means gives S2(0l) )a Si (0ll ( so that ((2/X(0ll),00 All) > 0 since
(p / y²(p 1)(p 2)/6 = p())
To complete the proof, we evaluate 0! = 1 / 1p. One has:
Σ = py 2/p
Σ 0 20=2 2
p
A/V
Σ
O
1,-00--00 20
(
-//p
12
2
3 y/
(78)
from which it follows that
<2/x(0!),00)=
32
(p) 1 2
2/p
8
2/p
(79)
12
thus yielding the claimed formula.
We finally come to verifying Assumption 3.4 This requires us both to demonstrate the existence of geodesically convex ball containing 0!, and to demonstrate that À is geodesically strongly convex in this ball.
We first give an explicit ball about a! which is geodesically convex using Lemma A.3
Lemma B.4. The geodesic ball of radius
xC
if 2
r =
1/1
4m 10:+1-1166 16 16t -8
if 2 3
(80)
centred on 0!. is geodesically convex.
Proof. That r can taken be oo p = follows from the fact that M this case is a copy the real line.
18

---

Suppose that 2 3. Observe that for any = (61,.... 0p) E M and any E Ton M, one has
(u, r] (nn (ul)
(81)
-II s(o1))11 (u, V2 s(0l)))
(82)
-
Σ 0.²)
(u, )((l)) doiag)(00
(83)
Σ 0,²
Σ uui05²
(84)
since n(lll)TTw w = -0 for any vector w E Toi M. Consequently, for any R E [0, (1/2)y /r), since Bar (0., 2R) C Bg. 2R) by the distance minimising property of geodesics, if @ll Bar ( ( 2R) then 0g [y //2 - 2R, // + 2R) for all that
h(() [u, wil & (y 2R) 2R) =: CR.
(85)
By Lemma
A.3 for all R E [0, (1/22y//// the ball of radius
R)
TR = min
(86)
Cn
is geodesically convex; it thus remains only to optimise this rg. Since R + R is monotonically increasing and R + /T/ T/CRR is monotonically decreasing, it suffices to find the first smallest value of R for which π //Cn R. This equality gives quadratic whose solution is
4T
V16T +
R 1/
(87)
8T 4
from which the result follows.
We next Lemma A.2 compute the Riemannian Hessian of along M.
Lemma B.S. For 0ll = (0I...., 0p) E M, define si (0ll) = 0. and s2(01) = E. Then the function has Riemannian Hessian
PrM
3 (wot) s2 diag(v02) Pras
ifp>
V2 P
(88)
2y² PrM
diag(u@4) (s) 31
diag(v02) PrM
i
In particular, at 0.. U,
vir = 4y² 2-4rr (e.).
(89)
Moreover, p = 2, then À is 2-geodesically strongly convex over all of M; ifp > 2, then for any E [0. 2 - v3), X is 2yy22///1122 4/p (1 5) - over the closed geodesic- ball Ba (e!! /9) of radius // about 0!! in M.
(1+5)3 1-))))))edddsiccall geedesicaall strongly convex
Proof. For notational convenience, set Vi d and V2 94), where & denotes the Hadamard (componentwise) product and is the vector field defined in (T). From the identities (64 (65) and (66). one sees that along M one has
f = y
V = v(ve" Vi),
= 8
1/2
(90)
19

---

The pointwise- pointwise-bilinear map V is zero if p = 2, and p > 2 it acts on vector to give the matrix with components
(V f(z)), =y(1 - (61)) E (1 - 5-5m)))(1 5 1
(91)
- =y(1 -by) Σ (1: - - + +)s041011)14)))
(92)
K
=y(1 ((u) z)vo =((22)== (Vi2)e² diag(v o 2))
(93)
-u((v²²)) z)o"" -VVi)= - (V/z)=² diag(v o z)+
(94)
(('2)) + diag(u 03 2) + diag(r 2) diag(
0 z))
(95)
=y((v (ou Vi) e (Viz)² (Viz)²² 2diag( (u
2 2) 17
(96)
where the fourth line follows from the fact that (1 - Ai = (A diag( ())); for any matrix A and the fact that di22()(Vi2)") = diag dizg(((2))))) one sees that
= diag(1 (v c 2). Substituting = = Vf = y
8² (s1(10 - 1/1) - -v vv - Viu0" + +2V)
if >
(97)
0
pf = 2
(V² = (s100= vo Vi + V12)
(98)
and
y²
(n, V..2'n)VV
(0 -V10)(00 - Vi) = y32 82/ -V)),
(99)
81
so that, by Lemma A.2
V \ 2PTM 8 fV + (v (n, (n
(100)
= 2yy² Pras (3V- (#2/8) /1 V1 + (81 + (52/ $1 yooo 200 -V/ Prm
(101)
if p > 2 and
Vir = Pr
P1 M V/2 + (s1 - (s2/s1)) - 2o V1 Vi + (82/81)00") Pras (102)
if p = 2. Now, let 7 be any tangent vector field to M. Since is normal to M one has u 7 = 0, so that
((s1 + (82/81))ou7 - 2v0 V1 2V100")= 2200
(103)
and
V-vo"Vi Vive + (82/s1)0v") -vo ViT
(104)
are both multiples of the normal vector e; it follows that
(2y²Pra (3V-- (s2/s1)[i) PrM 22y² PrM (V2 + (51 - (32 $1 1/) PrM
if p > 2
V2 A
(105)
if 2
as claimed. In either case, evaluating this matrix 0. y /01, gives the claimed identity
((O!) = 2-40 (r
(106)
In particular, if p 2, evaluating Vigd on the unit tangent vector field 7
02) +
(-0,,0,) /87 82 yields
2y² (0R + 02)
20 2
(V)A)))(()))))
0 2
0₂
2
(107)
0
+ 0
2y²
0 + 02
4
8y²
= 2 +
2 2
(108)
07 + 02
0/02
0 +
02
(0
02)
20

---

for any All = (2) E M, so that À is 2-geodesically strongly convex over all of M as claimed.
We assume that p > 2 and come to determining neighbourhood of e! on which V is positive definite. It is clear that a sufficient condition for V. À to be positive-definit is to have all entries of the diagonal matrix 3V2 (s2/s) VVi be positive; we will show that this sufficient condition can be guaranteed over a geodesic ball in M centred
at e!! Since the geodesic ball of radius r at a point in M is contained in the Euclidean ball of radius at the same point by the mintmiiiminiiisnng property of geodesics, it suffices to show that this sufficient condition can be guaranteed over a Euclidean ball centred at e!!
0p) in the Euclidean ball BR (0!,, 51," one has 0g [(1 - 5)y /2, (1+ + 8)y Vr]. Consequently, for any 0 E Bg (el, 8y!/m) one has
Fix 8 E [0, 1): for any 0 (01....,0p)
min 0. 2 N (1 + y 2/p
82
(0) S max 0: 2 81 (0) S (1 - 55 2/10
(109)
It follows that for all Bg (6!! By /r) one has
Amin (3V2- - VV) 2 min 0- 2 (3 min -
(52/ //2/s)\(0))
(110)
2 (1+ + 6) ④y 2/8 (3(1 2/1
+ 8) py 2/p 61 2y 2/P)
(111)
( + 6) (11 6) 2 1/2 (3(1 - 5)2 (1 + 5))),
(112)
which is strictly greater than zero provided that 3(1 8)² (1 + 8)2 > 0, which holds provided that 0 8 < 2 v3. This proves the result.
Note that À is not globally convex over M in general. For instance, when p = 3, along the family O == (e, 1, ) E M, elementary calculations show that A00 (e)) exhibits negative eigenvalues
0+.
Combining Lemmas 4 and B.5 one obtains explicit bounds under which Assumption 3.4 holds for multilayer scalar factorisation.
Proposition B.6. pfp then X is 2-geodesically strongly convex over the geodesically conver set M. If 2 3, then À is --looeesccl strongly convex over the geodesically convex ball Bar (0., 19). In all cases, one has
V xO0) I1)
(113)
Proof. Numerical computation reveals that the radius of Lemma is lower-bounded by 0.15y /p. Substituting ō = 0. into the strong convexity bound of Lemma B.5 and rounding down to two decimal places yields the claimed result.
We conclude the section by proving that the additional hypothesis required for Theorem 5.2 is satisfied.
Proposition B.7, For any 2 2, with v = =4y²- 2-/2-// as in Proposition B.6 one has
1,
(114)
c(2 c22 ((
so that the additional hypothesis of Theorem 5.2 holds.
Proof. Since x(e!) |IDf(e!))?? 2-22221 2/v, calculates using Proposition B.3 that
(!!
22-4// 4/p
32
8
8PP)
(115)
p2 2
p 3
1 2-4// 8(p 1)(p - (1/6) 2))
(116)
2 y² 4/p 8(p - ()(p - (p (0-2)) 2))
(117)
2 2-10
(118)
for any p 2 2. The result follows.
21

---

C The normal form of gradient descent
To prove Theorem 4.1 rigorously requires number of steps. First, one must give leading order formulae for the expression of gradient descent in tubular neighbourhood coordinates. It is essential here that the errors these leading order formulae are appropriately sharpened from their naive esti- mates using the structure of the tubular neighbourhood map (Lemma C.2) and the orthogonal stability assumption (Lemma C.4): without this sharpening our convergence theorems are impossible. These
considerations culminate in an intermediate coordinate expression for gradient descent in Proposition
C.5 Second, one must invoke Assumption 3.3 to transform the L-dynamics of Proposition C.5 into
the easily-analysed normal form quoted in Theorem
4.1 This is formalised in Theorem C.6
and
concludes this section.
Throughout this section, given functions
g, h,hh....hh, hg RR RC we will use the notation
g = ( hw))
(119)
to mean that there exist constants C₁
C. > 0 such that
||((w)] s Cill
(120)
for all sufficiently close to zero
In addition to this asymptotic notation, we will continue the notation of the main body of the paper,
denoting by .M := f-1(y) the submanifold of R2 obtained as the preimage of regular value y E I of an analytic map f: RP + R. We will also denote by n == f the unit normal along M.
Assumptions 3.1 3.2 and 3.3 are assumed in all of what follows. Assumption 3.4 is not needed for
this section.
Consider the map E M x R + R2 defined by
E(0ll, 0+) == -1l + (
(121)
Observe that
DE(O 0+) (OTM(())) ot Dn(0)) ), (,))))
(122)
where (fll) Ton M - Toi M is the identity map. The derivative DE( 0²) acts the linear map
M @ ToLR: (will wー) + wll + g+ Dn(0) |oll| + v-n(0l)) Te RP
(123)
between tangent spaces, and can equivalently be written in terms of the shape operator S(0))
)
To M 3 u + - [v) ToM as
DEOO 0+) (LTM(O)) o+ s(oll) n(
(124)
Observe that since n2@ll) is orthogonal to Ten M. DE( 01) ) is invertible whenever | ||S(0))]]=) where [[S(@]]]] ]| denotes the operator norm of S(oll): Toi M + Toi M with respect to the Riemannian metric. Applying the inverse function theorem proves the following.
<
Proposition C.1. The map F: M x R. > TRP is invertible on open neighbourhood U of M (0)- The image of U R.P a tubular neighhourhood N of M, and we denote the inverse of F. thereon by F-1.
The coordinates E ()) E ,0¹) are called tubular neighbourhood coordinates. The next step is to derive a formula for E and thereafter derive formula for the conjugate of gradient descent in the coordinates (0), 0-) induced by E. The following technical lemma will be necessary get the correct asymptotics for the error terms in our formula for E-1
Lemma C.2. Let U C. 1 x _ be an open set and let U e (x, v) + g(I,2) E R2 be Ca function. Assume that D = 0 on U for k- = 2.3 and that is imvertible on U. Then for any (z. E U and all E RP sufficiently small, setting (Uz, vyy = Dg(æ, 8) " and 2 == g(ø, V), the point z + v is contained in the domain of g-1 and one has
g (z + u)=+v------ (x) + Dg 10n--ggg9990010000 (()]o] + z - + O(1ea 13 ||. Pl II. |/ I11 12)
(125)
22

---

Proof. Since g is invertible on U and U is open, g(U) = domain(g is also open. Fix (z, y) E U and set 2 = gzz,)) y). Then since g(U) is open, + E g(U) for all sufficiently small so that g= + v) makes sense. Moreover, for all such e. there exists E IR x mR such that g((2, y) + u) = 2 + U so that 9 g- + v) = (a, v) + u = g (2) + u. Hence one has the following formula for the remainder R(v) of the Taylor expansion of (= + v) about
R(u)= : (= u)
(z) - Dg [z)[c]
マ Ddg (=)(,,)
(126)
= D Dg - [z][c]]
و2 (z),,
(127)
To bound R(v), therefore, we seek formula (w)) for in terms of e.
From g((2, y) + u) = 2 + v, one obtains
(uu)) == vu] + D2 g(x, y)uu, u] + S(u),
(128)
2
where the remainder ++ Sū) satisfies S(ū) = O (1 11 11= (1 since D = 0 for k = 2,3 by hypothesis. Since Dg(x, y) is invertible by the invertibility of g on U, G is smoothly invertible on neighbourhood V of by the inverse function theorem. Consider now the second-order approximation
= Dg(a, (0- (0- 1
1 D2 g(r, u) [Dg(2, |w]
(129)
u. Since u' = ([]e)),, taking smaller if necessary we are guaranteed that u' V. Thus there is C > 0 such that
| - u/l s ClIG(u') - G(w)l = CTIG((') v/|.
(130)
We will show that the right side of (130) admits sufficiently tight bound b(v) that u(v) u.' (v) Ob((v) suffices to give the desired decay in R(v).
For notational convenience, set w = (1/2) Dg(a, -) -D) v) [Dg(x, v) [M]] Dg(æ, y)- [1 w, and substituting this expression into (128) yields
Then v'
1 Dd D2g((1)))))) w) + S Do(a,
1
G(u') DD' (Dg(x,)))))))))))))))))))))))))) [v].
2
2
(131)
Since Dig 0 by hypothesis, one has w O w2 1/ || [l0l))))) as well as
1²g(z, v) Dg(x, [[o] w| = O |._ 1| || || 11)),,,1)),
(132)
D3 = [[[ O f1 || (ea.1111 1 1I
(133)
and
1
s Dg(a, y) [c]]
011 11,11
(134)
2
Thus, by 130p.
= + 01 10 1| 11 Ilyllll 11 || 112)
(135)
1
Df(xz y) ['o]
(x, f(x, v) -|]][]+
(136)
2
+ 11
1| 11)
(137)
Finally, applying the chain rule to the equation y -ty = Identity shows that Dz) = Dy(a, y1 and D2
(2) -Dg(e, y) D2 g(. D)( v)= Do(a,u) =1]. It thus follows from substitut-
ing (37 into 127 that
R(v) = u' DJ D 1 D2 = 0 11,, |I 11 |1<,1111,,,,13)
(11. || 11 1l..11111100 11. 111)
(138)
(139)
as claimed.
23

---

A formula can now be given for E-1 using Taylor expansion as follows.
Proposition C.3. For all Oll E M and all v all + v-n(lll) Toi M @ span(n((0)) sufficiently small, one has
TRP RP
[lll E (0) +2) (61l + all Dn((l))(olTot) O(1111), || 1 || ||l
(140)
v
Proof. The result follows from a second-order Taylor expansion with remainder:
E E + +) - E 1 (0ll) DE- [o] _ D' v] error(0 v).
(141)
Since the 0 derivaatiees of F(Bll, 0-) = fl + - ( vanish for all orders > by Lemma error(0l), v), satisfies
C.2
error(@ll, u) = 01!|' " |1 ||-l-)|
(142)
It thus remains only to calculate the leading order terms. One easily computes
will (111) ((:11)
DE(O"I o+))
w ot Dn(O ((. + w=n(())
(143)
Dd D²E( 8t)
blll [
82
=D (F((1,80)
2 [) = Den (vll +0- Dn(lll) + v- n(oll) 1o
211
(144)
v1
(145)
Do (ull 0+ + Dnll)) + wt u jut]
(146)
=0 DDn(( wlll + vt Dn(ell) + at Dr("ll)
(147)
=0 D2n(00) [ul wlll 2v [oo")
(148)
R. Recalling that is invertible when restricted to an open neighbourhood U of M x (0) in M x R. by differentiating the identity E-1 o E Identity one computes
any (01,0' 0-) E M x R and (ull, o-) E Tm
(Prr (ell) (ol)(( m(0lly7 (
Prrs
DE-1(Ol) DE- (DE(Ol.)
(149)
where PrM (fl) R" - Ton M is the orthogonal projection, so that
DE DE-
(150)
Again differentiating the identity E-1 o F = Identity & yields
D2 D² - Dー b E(O 0) DF 0) Dd r(oll ())E(())))))))).))) 0) DE(O!! .0)
(151)
so that from (18)) one obtains
( (2 Dn((ll)(oo))) (2v On(01)(e!1)
D2 P [v.
(152)
with the final equality being consequence of n Dn (Dm) -n Dn 0 as follows from differentiating the identity n." T = 1. The result follows.
The next result characterises the scaling of the |l-component of Vi in neighbourhood of e!!. It essential for obtaining the correct asymptotic decay in the error terms.
Lemma C.4. Assumption 3.2 implies that for all Øll M sufficiently close to の!! and all 01 e R
sufficiently small, one has PTM(Oll ( E( (00(0l,0-))) 0-)) = Odd (Oll, 0!!)) where Pr T RP To M is the orthogonal projection onto Ten M.
24

---

Proof. By Assumption 3.2 one has
(+) =0
(153)
for all 0t E R sufficiently small. Since for all sufficiently small 0² the map Øll ( (e((E( 0-)) is smooth, the result follows.
We are now able to give the tubular neighbourhood coordinate expression for gradient descent.
Proposition C.5. In the tubular neighbourhood coordinates E (0 e0) ") E(Oll, 0 ) for tubular neighhourhood of M, the gradient descent map GD (1. -) 0 +> 17 v)()( is conjugated the analytic map
(0t) (1- (1- - (00) / D ((0)) (0+)) o((0-)9))
Oll - mi 10
d V MMOl)) + -000)) min min(1, (00, o!))))
(154)
Proof. Since f is analytic by Assumption 3.1, so too are GD (n, :) and and hence, since the inverse function theorem preserves analyticity, so too E-1, Hence the conjugate E (11, ) E is analytic wherever defined. Recalling the open neighbourhood U of M {0} in M x R from Proposition C.I on which E is diffeomorphism, given (0ll, ,+) E
GD (n. E(O)) ,0)) = E(lll, 0+) - - o+))
(155)
0 000 - 17((( + (0-((0l))
(156)
Applying Proposition
C.3 with 0+ ) 1/ v)(1 + 0~n(0l)), one then obtains
Øll + vll O(1
|00
| I)
E EGD (n, E(oll, 0+1)) = د. O | 11 ||e |1
(157)
here: were there be a pure 11] term in the error for some k 2 1, then the same error term would have appeared in the the formula (157) and, as shall now see, would have prevented the ||-component of the dynamics from having an O(d )!!)) error term.
Note the crucial role played by the estimate on the error in Proposition C.3
We now compute formulae for wlll and vt. Using the fact that Vε(Oll) 0, one has the Taylor expansion
V((E(Oll, = 8- 2 + (0-)2 V (2 + (0-)2
+ (( (((()))
6
(158)
(0-)2
v (02 (all) + (0-)
= =0-Xn(0l))
V (0ll) ((())))) (159)
where the second line follows from the fact that (0ll) ) is an eigenvector of V with eigenvalue A(ll)). Taking the inner-product of (159) with n(ll)) then reveals that
-- = (n(ll)), n)
(160)
= -1-17(01) - 2
E #3 (oll)
7 100)
(161)
6
One obtains wlll by projecting 159
onto Tou M:
ا PTM ( (0) 0+1)
(162)
n(O0)2
Pras v (001) ((( min dar dar o!13)
(163)
2
_n(00)a
V (v + ((+))3 min da (ell, e!!))
(164)
2
VM<((ol) + 0(0 1 '3 1, das (0ll o!!
(165)
2
25

---

where the second line follows from Lemma C.4 the third line follows from the definition VM P of the Levi-Civita connection on M R", and the final line follows from the identity (l n) = À that holds along M. Substituting these expressions into 157) yields the result. 17
By finally invoking Assumption 3.3 we may now transform the tubular neighbourhood coordinate C.5 into the normal form expression of Theorem
expression of gradient descent from Proposition
Theorem C.6. Recall the function c: R x M + R30
c(1,0)= " D' 1e - " D' (n > all E M3 in
(166)
from Assumption 3.3 There is an open neighbourhood W of {(2/A(0)),
R x M x R on which up: W IR x M R defined by
p(n, 001, 8+) (n, all 8t), (n, 0ll, 0 (n. 01. 8+))
(167)
0
D 863
(0))2
:= 17, e
(168)
dn, 0
(1- ( (1 - m)((01) c(n 01l)
is an analytic diffeomorphism onto its image. Moreover, the composite (Ig x E) 4 W - R x R, where Ta is the identity map, conjugates (Jz, GD IR R2 + R x Rp to the map
"
|1 0
+
GD
n(0t2 20(n.01):
011
0(1 j dv o!]1)
(169)
(1- ()(ll)))) + 10-)3 (((4)))
Proof. That an open neighbourhood W of ((2/X 0) All E M) exists on which op is diffeomorphism follows from the fact that c(17, A ) * for all Ø E M and rj in a neighbourhood of 2/ A(Ol)) and the inverse function theorem. That yp is analytic a consequence of the analyticity of I and the functions used in defining in terms of = (1/2)( u)².
The formula 1699) pertains to the conjugation
(/R E GD (IR E) op
(170)
of GD Proposition C.5 already gives formula for (Iz x E GD (1z x E), so suffices merely to conjugate this formula by yo
Since (n, 0ll, 0-) = 0l1, the |-component of (77) is given by substituting * 0 0-) in place of 6- in the |l-component of (144. which yields the map
1(01)2 2c(n, 01) V + 0((0) y (8
gll 0-
+ All
(171)
As demonstrated in 33 however, one has
expoll (-
n(0->2 2c(n, 01)
Var2(6l)) All
n2
vard(()) ( (172)
2c(1, 0ll)
101122
All
V + (((0)) min(1, (1, des (A!l, (!!)?))
2c(17 All
(173)
since v = 0. Thus the ||-copponent of (170 be written as
n10-22
01 ((0-)² min(1,dxx(B)) o!)))
(174)
20(n, Olly
0+
as claimed.
Finally, the L-component of (169) follows from the same argument used in 27 Theorem 4.3).
26

---

D Convergence theorems
The purpose of this section is to provide proofs of the main convergence theorems presented in the paper. addition to assuming Assumptions [3. 3.2 3.3 in the previous section, in this section we will also assume Assumption 3.4, which states that there is geodesically convex ball (xr centred on !!! over which is geodesically smooth, geodesically strongly convex, and has unique global minimum at 0_ with Riemannian Hessian at 0! being a multiple of the identity.
D.1 Suberitical regime
Our convergence theorem in the subcritical regime has two components. First, it is proved that gradient descent implicitly descends x(o!!) from being initially 2/1 eventually being < 2/n- For this, is necessary to have a descent lemma for the perturbed Riemannian gradient descent component of Theorem 4.1 (Lemma D.1 well upper and lower bounds on |0t1 during this phase (Lemma D.3). Following this transient phase, convergence is easily proved using the fact that [0+]] contracts exponentially zero.
Lemma D.1 (Descent lemma for perturbed Riemannian gradient descent). For the u-geodesically
strongly convex, L-geodesically smooth smoothfunction function À on the geodesically convex ball Byd (x((0.)) centred
on 0. in M, consider the map
1n((1)² 2c(11, 0 VM VarACol)) + 1000))* (1, o!!
Gn. = expoil
(175)
of Theorem C.6
Then for any > and all 0 sufficiently close to zero, GD (n, 0²) maps BMr(0,, into BM(O.,7) Moreover, for any nj O, any Oll Bar(O!!,)) and all 0+ sufficiently close to zero one has
(1, - -(O0)) s (1- pn(())² 4e(n, 011) (A(0l)) -((!)))
(176)
Proof. For 0+ sufficiently small, there is unique 5(0), 0((0))3 min(1, (Oll, (oll,o!))) ,!!))) that
such
1(0122 GD (n. = expol (- 2c(n, 0ll) ) + 8(0) ((01100))
(177)
To that (D - 0-) preserves the ball (0., consider the function p(0ll) (1 (1/2)dar(01, o By the Gauss lemma (29 Corollary 6.10]. VMp(Oll) = logar (o!!) Fixing @ll By Bol, r), for any vector w E Tat M one therefore has
p (xpxpl (- (-VVMAA(H)) + w)) 8!! = - (V (VMMp((ll), ((ell)))))
(178)
dt 1=
= (Logon (e!!), vas((e)))
(179)
S ( < O.
(180)
where the third line follows from Lemma E.2 Thus, for all 0 sufficiently small, we are assured that p(GD)) (n, Ul1, 8-)) p(ell), implying that GD (n, All 0+) e (ol.)
27

---

We now come to the descent lemma. Since BaS Bas( r) is geodesically convex, we may perform covariant Taylor expansion and invoke the geodesic smoothness of À to obtain
(T, )) =A(ell) + VIN n(!1) 0e(1)011) VMr(oo)) + 5(0) )
5(01).1))+
(181)
+
VMrd0ll) - - .0.)
-800),
(182)
2c(17, )
n(0
1 7 1
SS(ell)
4c(n, 01l |1 (10 )|2
(183)
2c(1, 01)
)2
(V (Vxs)(ol),, 8(01, 0²)) + _ $10 0 0²)1?.
+
(184)
2c(n, All)
From Cauchy-Schwart and the estimate 0 S (E!/3a-c 1/))/ (implying that ab (1/2)
-212)) with 1)(0))// 0 )), one obtains
(Var)(0)),8(0 0-)) - | ~00 0-1|
(185)
V MA00) )|12 + 2c(n. fll) 8c(n, 0
770 Y²
70012² |8(00), 0-))1
(186)
Taking |0-| sufficiently small that ((0 Bl1)) 1/4, therefore, one obtains
X(GD) 04,0) s 2(01))
n(0-)2
V > (2c(1, 0ll)
n(0 } + 2/2 (10(011, 00)1 (187)
4c(n, 01)
710-)²
= =
V M || + )( min(1, dn (oll, o!l))). (188)
4e(11, 01)
Subtracting x(e!)) from both sides and applying geodesic strong convexity (I ( z 2p(A(00)) x(e!!)) by Lemma E.3) yields
(T, 0+)) x(O!) s (1
pm((-)2 2c(1), 011 (AcOl) x(o!)) 0((0-)* min dv (0!))))))
(189)
Finally, invoking geodesic strong convexity once more to estimate d. (All, o!!) < (2/1))((0)) x(!!)) as in Lemma E. and taking 0- sufficiently small allows one to absorb the remainder term on the left side into the A(ll) x(e!!) factor and yields the result..
The next result will be used in giving uniform bounds on the magnitudes of the 0 iterates.
Lemma D.2. For R, define fa R IR by
fa(x) == -(1 + a) + x² 0(z*).
(190)
For 0g, on E TR, consider the composite fajao = sar o Sao Then for all y > sufficiently small and all ctp, ai E 70;
1. faroo monotonically increasing on 1-2v7.2
2. Jajag admits the sole fixed points O,
<- 5--1 00 + + O(an + a1), &- =-=10
On + a1
+ +O(an + +a))
(191)
2
in the interval 2-277.2/7)
Proof. For > O small and all org, [0, 7), one computes
foro (x) (1 + 0j + 0g)x 2xd + 0(a0a a022
(192)
28

---

af = +1 + ag 6x + O(ava1, a1x" " 2²).
(193)
In particular, for all |x| s 2 one has
Dfa (
+0((7/ >
(194)
for all y sufficiently small, so that 1-2.5.2 is monotonically increasing, Set == (ao + /2 for notational convenience, and denote
$ (x) == Sajan (1) - x = 23x 2x.h + 0(82x, ,x²,
(195)
Fixed points of fag are precisely roots of 4, one of which is clearly I = O. Moreover, if 3 = 0. then ((x) -2x + O(z*) that in this case
is the only root of ①b [-2 2 for all y sufficiently small. Suppose now that 3 > O. We will demonstrate the existence of unique root ++ of & in (0, 2 with the claimed formula. Since 2,xx - 2x.3 > 0 for all < x - 2 1/2 /B. we see that for all y sufficiently small, ((x) is strictly positive on (0, 2 10.2-1/2// /3] and so cannot admit any roots therein. However, observe that for all y sufficiently small one has
(+(x) 2,3 6z O 3a. x²)
(196)
for all 3-1/2/33 S x S 2/7, implying that & is monotonically decreasing on 23,
and thus admits most root therein. Since moreover 4 (B2-1/2//) > 0 and ((2/B)
- -12,3/22 ( < 0, by the intermediate value theorem & admits precisely one root & in 12-1/2 /3. By the mean value theorem, there is z between vB and & S 2/B such that
0 = +(+)) = (VVB) + (+ VB),
(197)
so that, since [D$(2)] 2 43 + 0(88/2)= =S(() and @ (V3) = 0
$ (((B = v 0(3)
++ = VB
(198)
(2
as claimed. A similar argument can be used show there exists unique root E- -- + ( of + in [-2 0), thus proving the lemma.
Lemma D.3 (Upper and lower bounds for T iterates). Let (Aeltes be a monotonically decreasing sequence of numbers. Then for all > 0 such that 1/ao > 2 is sufficiently small, there is constant C O such that for any to sufficiently small, the iterates defined by
0 = (1 -1-m)000 + (0/)3 + (((0t)))
(199)
satisfy the bounds
10,1
I 10/11 s 22 nx(e!)) 2
(200)
V1 + 1(0) )2
for all t E such that T) 2 2.
Proof. We prove the upper bound first. For each s E N such that 17d,, > 2, denote a, = 1/1s and denote by Ja the map
Jo 0+ - (11-n)2a)0- (0+)3 + +0((0)))) =- + a+ (0+)3 + (((0-))).
(201)
(202)
Supposing that s E N is such that both 17 114 2 since {eeHEN is monotonically decreas- ing Lemma D.2 guarantees that fa+1.- : fst1 fa admits attractive fixed points Es+1.e.t =
+ (a++1 + /2 + O(a.+1 + a.) and is monotonically increasing on 1-2 2 provided ag 77A 2 is sufficiently small. Suppose now that 0% E 2--2voo 2/aol {0}. The iterates 6,
as defined in (199) satisfy
0 = f2.+1.22(
B = f20.2-11
(8
t E N.
(203)
29

---

Assuming without loss of generality that 0g > O, we will prove by induction that 02. S 2vo for all & E N such that as 2 0 for all I 2f - 1. Clearly fg - 2 by assumption. Suppose as an inductive hypothesis that O < Ozi S 2//0 and that 1 arr 2 O. Using the monotonicity of 120+1,21 from Lemma D.2 and the fact that 22++1.11 is its sole fixed point (0, 2 /oo either
0z. S 1,1-1,24++ = 82 I f2141,200 = 02142 S E20+1,21,+ I 2Vo0.
(204)
or
02 2 E20+1,2,+ = 220+1,21, S [2141.21 0 S 022 S 20000
(205)
by the inductive hypothesis; in either case 0 S 2 2 thus proving the claim. A symmetric
argument shows that the odd iterates satisfy 0 - 2/00 for all /. E N such that a₂ 2 0 for all
S 21, thus proving the upper-bound in (200).
We now prove the lower bound. Invoking the upper bound, og can be taken sufficiently small such that the recursive estimate
10 | = 10 + a, (0/1)² + (((0t))))
(206)
2
- 10+1(1-2(07))
(207)
holds as long a; 2 0. By taking og even smaller if necessary, squaring, taking the reciprocal of both sides and applying 1/(1 - x)² +1+32 3r for x sufficiently small one deduces that
1 $ (0-)² (1 6(0,))) = (0t
(0//) S 6,
(208)
(0
from which it follows by a telescoping sum that
10
10 V 1 + 6(0) )t
(209)
for all such that a₁ 2 O. This proves the claim.
Theorem D.4. Assume that n < 2/ (0!), and let 0 c, C < xo satisfy e S c(n, 01) for all All E Bar(fl, r). Then for all 0g sufficiently close to zero and all o! + l!! E BM r), if n 2/X(0)) is sufficiently small, then there is N, with
1 2/n x(o!)
24C
TS r
(210)
6(0g y2 ((08)) x(0!!)
such that 7 < 2/)(0.) for all 2 t and n 2 2/X(0!) for all t < T. Consequently, setting
1= (2 - (\((! < 1, the iterates (e!!, 0 converge global minimum (0x., 0) of & in M with suboptimal sharpness
x(BL.) - x(B!) 2 exp (- (-AnL² (0+)2
4nL² cj 1 1- 32 (A(6!)) Am).
(211)
at rate ((9t).
Proof. For any 0!! + A!!, take 0f * 0 sufficiently close to zero that the hypotheses of Lemma D.1 and
Lemma D.3 are satisfied, with Ao Lemma D.3 taken to be equal X(B)). and assume that 17i0
is sufficiently small as in Lemma D.3 Then we may assume that X == x(e!) is monotonically decreasing and the bounds of Lemma D.3 apply to 0.- Since (At is monotonically decreasing, any satisfying 1) < 2/!(e!!) for all 2 and 17 2: 2/X(0,) for all t < is necessarily unique. We claim that this t satisfies the upper bound T 1 7(8), 0.)). Indeed, suppose for a contradiction that (00) 0.). Then for all r. denoting X, X(0!), ), one would have
de41 h, S <(1-
4c(n, 0!! (1 + 6(0)) )2t) (Ae Aa))
pm
(0+)²
(212)
30

---

by Lemmas D. and D.3
Invoking c(n, c(v,0)) S C, taking logarithms, applying the upper-bound (n(1 - z) S = and applying a telescoping sum then yield
In(Ae) - 1m(A)) -
pn
(213)
+ 24C ==0 6(0))²
s
In( [ 6(0 )²()
(214)
24C
by Lemma E.4 that
A S (No )1) + 6(8))) )2t)
an + Aa-
(215)
However, plugging = 7(8 og) then reveals that A < 2/1,, thus contradicting the assumption that > =(6) ) and implying that S =(6) e))) as claimed.
We now come to the final convergence and sharpness suboptimality claim. For all 2 T, Lemma D.I continues to imply that X < 2/1 decreases monotonically, while 0% contracts exponentially according to the estimate
16 = 0, || + (n1 - 2) - (0.
10,111 - (2 n)) S 10, 1 - (2 -))))
(216)
This proves the claim that the iterates (0!, 0 converge to global minimum 0% of & at rate ((11- - (2 )))) We now turn to proving the suboptimality of the sharpness at 0l.. The geodesic strong convexity of À in the form of Lemma E.2 implies that
pdar(0l) , S (TI sol V M A(Bll). log,! (e")) s V do
(217)
for any Oll BBM (0!, r), that the l1-update can be written
11(0)²²
GD 001 -00 = expat
V v
o!!
(218)
2c(n, All)
n(01)²
=
exPal
2c(n, 0ll)
v +
Vas (
(219)
Consequently, invoking the geodesic strong convexity of À once again the form of Lemma
Aet1 2d:- + (V (V log, (6!+1)): ~ー
log (0!4
(220)
=>-- Vard(o!),
n00,1 )a
VMrA(O)) ((0,)) min(1, IIV N (cAO))I))) >
(221)
20(n, )!)
+ 11 (0t))
M (L
(222)
2e(1.0) 0.)
)0t))2
\ 0(+
(223)
2c(1) )!!)
)2 )2
2\: c(1, )
V ((!!)(!
(224)
for sufficiently small 2 Ana - 2> 10 Subtracting X, from both sides and invoking the Lipschitz
gradients ((0ll( < Ld (lll, )!)) and strong convexity (1/22)das (0", ,0!)2 < x(e1l X.
from Lemma E. properties then yield
m0,t)2
Aet1 X, 2 d - X.
1,2 I22d (0) 0.)2
(225)
c(n. 0!!)
2 ()e - 2.)) 2n( 2
(226)
c(1, c(n,0!! 0!!
(0 g26
2 (he - X.)
(227)
cp
31

---

where in the final line we have invoked c(1,. 0.) 2 It follows that
[-1
2
A d II
g20
(A, - A+).
(228)
cp
Taking 2/11A0 - 2> 10-1 smaller if necessary and using the bound In -x) 2 -2x valid for small has
II
2mL)
820) 820))
(n, t-1
82a
(229)
cp
cp
I.2 (0,))
(230)
- -2
from which the result follows.
D.2 Critical regime
The convergence theorem in the critical regime is the most difficult to prove. This is because in this regime, the fixed point to which the iterates converge merely parabolic rather than hyperbolic as is (at least asymptotically) true in the suberitical and supercritical regimes. The convergence theorem in the critical regime will follow from a careful analysis of the following parabolic system.
Fix 0 < functions, and consider the analytic functions
and b >> O. Let A(x, v) = Olx, v). B(x) O(x) and C(yy O be analytic
T2 (z, v) = (1 (a + A(I, )))"))
T3 y) = (1+ + (6 + B(x))=² - (c C(()))22) y.
(231)
Denote by = (T,, R² + R² the corresponding analytic map. Clearly 0 is a fixed point of T; will be concerned with proving convergence of the iterates of T to this fixed point. The following result guarantees that admits an invariant manifold approaching the origin. This result is essential, since our final convergence proof is two-stage proof; first, convergence invariant manifold is proved, following which, second, convergence to the origin is easy.
neighbourhood of the
Lemma D.5. There exists r > and a function {x : lx² r| < r)
R whose graph
r. = (1x, u(x)) dom( } is invariant under T, for which u(x) cma x + 0
Proof. The analytic map T IR22 - R2 extends uniquely to a holomorphic map / defined on neighbourhood in C2, Like T, at the origin the map T is the identity to first order, vanishes at order two, and has third order Taylor coefficients given by
ayy (0,0) = 22, Yzur T (0,0) = 2b,
To
= -6c
(232)
with all other third order coefficients being zero. Observe that on the vector given by
1,
(233)
c
one has
3 BC c-a
b
(0
6a
6b
(234)
c c-a
3
-660
(235)
so that v is a non-degenerate characteristic direction of T 6. Definition 4.1], which is moreover
attracting since -6a /b/l///(e a) < 0. By (6, Section 6), therefore, there exists r > 0 and a
holomorphic function {x E C | rl < r) + C such that u(z)) - - + (xx²) whose graph invariant under Restricting this the intersection of its domain with the positive real axis yields the result.
32

---

The next lemma will be key in establishing convergence to the invariant manifold. Lemma D.6. Let u(x) = V - + O(x²) be the function from Lemma D.5 and define A v) == y - u(x). Then (xx,y) = A (x, (x,1)/ u)/x satisfies
Á o 2)) S ((
----"²² 1
(236)
for all sufficiently small 0, 0.
Proof. We estimate AoT Define
D(a, y) (z)
A.T(r.y)
(z,
(2,y) v) E dom(u) Rzo.
(237)
x
A(x,y)
Observe that both D and N are analytic on the domain of u. Observe moreover that
Â o T(x,y) Ao AoT(x, I
(xr D(x, v) N(,,y) D(r,y)
(238)
(xx,y)
TT(x, v) A(x,V) D(x,y)
so that estimation of ĀOT is reduced to estimating D= We first estimate D N. For this, observe A that since Txx 0) (a, O), one has (xx, 0) (x, 0) 0 for all x dom(u). Thus, by analyticity of D and N, there is unique analytic function H/ on dom(u) x IR such that
D(x, v) N (z. y) = yH y)
(239)
We compute (xx, y) second order. For notational convenience, set
K
(240)
c-a a
so that u(x) KZ + O(²)). Differentiate (239 with respect y and evaluate at y = 0 to obtain
= =(, 0,
(241)
Observe that
1x, xr 0) 02(1 - az²² + Axxx
(242)
while, from N = (AoT)/A, A, one has
(x, 0) Ad, o T) (A o a)
(z.
(243)
2
From A(z,8) 9) u(x), is easily seen that
A 0) = -u(x), 0, 0) 1,
= - = - -u(x)
(244)
while
0 o () 0) = 0, v' (zx
ba b +
(245)
so that 243 yields
bx + +0zi)
0) = =0,
(c a) ka + O
(246)
u(x)
To compute the second order component of H, differentiate 2339 twice with respect to y and evaluate at 0 to obtain
(0 -)
0,H - 1
(247)
One sees that
-2a 0zD(x, 0) = 02(1 - + A(x, y) = 0(x),
(248)
33

---

while
(x, (8 TT)
(A (A
2 T)
(a
(z, 0).
A
A
A
A3
(249)
Since
(A))T)(2
(x, - "/ (3)) 2aKT 0(x²)
(250)
one obtains
(z,
20KT + ((22)
-
1 1x + 0xx
1
(251)
K. + O(2²)
x²1² + 0(5))
R2x² + 0(x))
-a 26/n 122 + 0(x)) + 00z).
(252)
Thus
0y 1 0) 1/2 2a + 2c 0(x)) = (c ac) + Ozx).
(253)
We deduce from 246
253 and 239 that
D (z c v)
N (z, y) (c ay²² + kyz) + Oyzz y²x, x²x)))))
(254)
follows that for all sufficiently small x > 0, y 2 0, one has D y) - N (x, v) 2 (c-e) y² and 1 2 D(r, v) 0 that
2
A T(r,V) =1 D(x, ((2 D(x,y)
31 yp
(255)
Axx,y)
as claimed.
Lemma D.7 below gives uniform upper-bounds on the magnitudes of iterates, which are used in enabling certain estimates in the convergence theorem.
Lemma D.7. Define (z, y) := bx² + y Then for all (zo. yo) in a sufficiently small neighbour- hood 0, with yo 0. the sequence x; is monotonically non-increasing and the sequence ye satisfies the uniform bound yf s 2W (zo, /o)/a for all E N.
Proof. In any sufficiently small neighbourhood U of R2 one has
W(T(x, v)) = hx2 (1 ay² + Olry² y/
ay² + bx - es² + 0 1²1)²
(256)
= =ba² 20y2aay))) 0 2 ay 2 - + 0 2 3
(257)
= =W(x,y) - abx 22 acy
acy + 0(r² 23,""""))
(258)
< (x, " (aba² + acy²²)
(259)
for all (z. v) U, so that W is non- aan-increasing as function of the iterates of the map close to zero.
Now fix (ro: Ho) sufficiently small that, with Wo = W (wo- Ho) one has (ro, / 2Wo/a contained in U. Then in particular (Ip, Ho) is contained in U since yo S /2Wa/a. Suppose as an inductive hypothesis that yt S /2Wo aa, 8t S xo and W, < Wo. Shrinking U yet further if necessary, one sees that
1> di + O vi) >
(260)
thus yielding I1+1 < कe S To. Moreover, the inductive hypothesis implies that (x,, Mc) U: thus, denoting W. = W (Xx,, Wh) and W = W Wx M1++)) ), one sees that
W+1 S W.
2 avi (bz + cyi) < Wo.
(261)
so that yt+1 S 22W441/ S /2Wo/a. This completes the proof.
34

---

Finally we can prove convergence of the iterates of T toward zero.
Theorem D.8. Set Vbl(c- a). For all sufficiently small zo 0 and yo 0 and all 0 < < there is
== 7 r) =O(v0² 20:002330/(c-00 -3/((caa)
(262)
so that for all 2 T, one has
2x S
(263)
V1+ 3a c)2 -TT)
V + 2a(n 2)2 7)
and
(K- c/2x) 1/ (x + c/2)aar-
(264)
In particular, Ft, y/: ((--1/2).
Proof. Using Lemma D.7 we may take (zo: yo) sufficiently small that the quantity A. = (xrt, Vt)
of
the recursion
A S 3. (1- 2 Vi
(265)
for all / N. By taking (Zo, yo) yet smaller necessary we may also assume that y/2 satisfies the lower-bound
Mt+1 2 ye(1 203²²) )
(266)
Squaring and taking the reciprocal yields
(267)
Vi+1
Vi
77
2cy
and, taking (No, yo) yet smaller if necessary so that 1/(1 - -2cy)) S 1 -yy for all obtains the recursive estimate
S 6c,
(268)
v2+1
y²
from which one obtains the bound
y/
2
(269)
+
6cypt
for all E N by telescoping. Plugging this into the recursion for Á and taking logarithms, one has
c-a a yī
In mm In Ad S
(270)
2
1 1 6cy
Taking a telescoping sum yields
In A In Aol <
C
1
<
C a
In(1 GCyat)
(271)
12c
+
12c
==0
6eyg
by Lemma E.4 Thus
Adl S [ + Gcyyt) (c-a))))
(272)
and for any 0 < bbllc a) one sees that for any greater than equal to
(20/(c-o)
-1)]
<(vo,) = 6cyg [|
(273)
one has A.l < e/4. It follows from the definition of ÁZx, v) == (y-xx+0(22)) that T E)
implies
V/ E [KT (c/4)x + mz (c/4)= (6/2 C [[x </2)x, (n + c/2)x),
(274)
35

---

where, shrinking 20, Yo if necessary, the final containment is obtained using Lemma D.7 t 2 T(yo e) and the bounds (Is + e/2)z, 2 y 2 (x e/2)x,, in hand, one obtains
With
(1 1--a(t a( + c) xx 2z - 2141 - (1 a(k - <)²x )
(275)
again shrinking 20, Yo and invoking Lemma D.7 if necessary to deal with the higher-order terms. Inverting and squaring yield the bounds
1
S 5 (1 3a(k + {)
(276)
It [1 alk + c)2o))2
and
1
2
2
2 (1+ 2a(n - =)²=2), zt
(277)
2
2j (1 a(k e))
2
1++1
again shrinking To, yo and invoking Lemma D.7 if necessary to obtain the final upper bound in (276). so that
1
3a + c)²
2a(n c) c)2
(278)
x
from which telescoping summation yields the bounds 263). follow from the limits of (274). This completes the proof.
The corresponding bounds for ye now
Finally, we may prove our convergence theorem in the critical regime using Theorem D.8.
Theorem D.9. Assume that 1) 2/)(O!).. Assume in addition that // (c(2/>(0!), (o!!))
where v is defined in Assumption 3.4 and e is defined in Assumption 3.3 Then for all Bl sufficiently close to 0 and all 0+ + sufficiently small, one has das (0!! 0!!) O 1/2) and |0 (t--1/2).
Proof. Denote A, == x(e!!) and c, = c(2/>., B!!) for notational convenience. By Theorem 4.1 for all (0), e sufficiently close to (el, 0), GD has the coordinate expression
expal (- a(/X(ONN
(0-y2
Var>(Ol)) 0((00)8d
GD (1l GD (n, GD+ (1,00000)))
(279)
(1- X.
2A(0"))
)o (3 + O((A1)))
We first derive formula for d (GD (11, 0ll, (11,00)))) 0!!) || |log, (GD (n, All 0 0-)) Using covari- ant Taylor expansion (29 observe that:
V M !وT! + (Vrr((o!) + v v2.2(0)) (0))) O )) II -at (u log (0l) + Odd (el), (!)
(280)
(281)
by Assumption 3.4 where II -6!-001 is parallel transport from o!! to all. Similarly, observe that
(0
(0 + +0((0L)ddr00)) o!)).
(282)
(2/A(0))))
cod.
We may thus write
expo1 (0-)2
logo! (GDI 011, 0+)) =logg!
c,7. Hol son (logo) (81)
(283)
0 P (Bll, o!!), (0 ) (el1, (!!)2)
(284)
codo (04)²) (o%,!(1)
(285)
+ (((0)))3 (0ll, o!!), 22 (011, e!!)))
(286)
36

---

Hence: dm (GD) (n, (1, )0)
||ogop! (GD (,,
0-))]|
(287)
(1 - -
-1-=000 +0 + (04)ddm(0), o!)) dm 0!).
co2a
(288)
We now derive a formula for |GD (11, 011 0-)|. Perform covariant Taylor expansion 29 on X(0ll)
to obtain
A(oll) = X, + (V (Van2(0)) log (0 )} + ( (0)) + Odas (00) (!
(289)
A. 2 d (0 0!!)2 + O(dsr (0d )!!)))
1
(290)
by Assumption (34] Then one sees that
(0ll + O(d (0ll, _o!)))))) )or (0t)s + -((04))) (291)
(D,0n,l 00)
1
X,
v
dM doll 0!!12 - (0-)² (ddr .00)).000000)))) (0459)) 0t.
1
(292)
Xa
so that
olly (0-)22 + O(dm(0 (00) ..))).(000)))) (01)3) 104
(293)
provided dıs (0ll,!!)) D.
o!!) and 10+ are sufficiently small. The result now follows by invoking Theorem
D.3
Supercritical regime
Our convergence theorem in the supercritical regime is the casiest of the three. The orbit to which the iterates are attracted is hyperbolic unlike in the critical case, and the transient phenomena are relatively simple compared with those of the subcritical case.
Theorem D.10. Assume that 2/ 2/X(0!) is sufficiently small. Then there are positive constants C1 C2,Cr, C3 and stable orbit of period two of the form (6!! (n))(()) 2) + (nz)(0!) - 22) to which the iterates (e!! 0,+) starting from any (8), 0 satisfying 0 < 10/ < (n 2) and dxs (0/, e.) s C2(n>(e!) 2) converge with rate ((1_))())(e)) -
2))*). -
Proof. Arguing using covariant Taylor expansions as in the proof of Theorem D.9 one sees that
dm (GD (1, 0ll, 0-), = (1 (1- T
2c(n, 0!!) (0-)² (000)+0(000,000)))110 (0
.0))) dar o!)
(294)
and
GD+ 00) --(1- - (n>(o!))) - 2) +
nv dv (0l) 01)32 - -0~22 + Odd o1! (01, (3) 0-.
2
(295)
Consequently, the result follows from Theorem D.11 below.
Theorem D.11. Given a > 0 and constants a > 0 and b > 0, consider the functions
(z = (1 ay² + O xy))))x,
F, v) = -(1 + a + bx y² + 0x23, "
(296)
Then, there are constants C1 Cz, C3 0 such that for all sufficiently small, the map F Fy) admits a stable, period-2 2 orbit {0, {+), (0, with & = + and the iterates (#141, (t+1) = F(ri, y1) starting from any (Iu, yo) with 0 < |yol CI and |zo| S C2va converge to this orbit at rate of ((1 Cza)").
37

---

Proof. The line z = 0 is preserved by F, and thereon one sees that Fy takes the form
Fg(0,,) :y -(1 +1++00 aJyy v² + O(y²)
(297)
which, by Lemma D.2, admits period-two orbit {É+,E-} of the form &= = =vā O(a). We
will show that the iterates of the square F of F converge to (0. for all initial points (Fo, vo) sufficiently small, Yo * O. This will be achieved by showing that ( y)ll < uniformly over a
neighbourhood of (0, }, followed by proving guarantee of convergence to this neighbourhood in finite time.
Via a routine calculation one sees that
)²1) (z, (u(( (v((1+))++ + - (1+ (1 a) + 0(²), x² x. .221) (298)
x a1-a(11 + +0zr"
-
The derivative of F²2 is then given by
DF (1 a(1 (1 (1+0)) , 4(1 + abyry
-2a + (1+ a)+
(1
- 3(1 + +a) + (1 (1+)))))
(299)
+ +0 ury, xn)² ,1²). In particular, for all satisfying
(300)
a1+a) (1 + a) S S 20,,
(301)
the matrix DF7 y) has entries with magnitudes upper-bounded by
( (olzva) 1- (1/2)a + +(x2)) + O(x", Var* oz, az,2
T
(302)
for all a sufficiently small. Using the bound ||Al|2 S v |A for a matrix A, where || || and || oo are the maximum column 1 1-norm and maximum row 1-norm respectively, one sees then that there are C2. C2 such that for all a sufficiently small, all æ satisfying
| C2va.
(303)
and all y satisfying (30T), one has |I 2l12 S 11-C0 130a.
Letting V be the neighbourhood of the two- -ro-point>>> set {=) enclosed by the bounds (301) and (303, one sees that the iterates (I22: 1/2) of F2 starting at any point (zo. yo) V converge toward (0, Esign(yo)) with rate [| 1/2-- - { < (1 C3a) (x Vo Esign(vo))
We complete the proof by showing that the iterates starting from any ane satisfying 303 and any yo satisfying
v0 v
a (1 (1+ (1+
(304)
are eventually drawn into V. This, however, is easy: for any x satisfying (303) and any y satisfying
y² I < ((1 a) (1 + +))) one has
1(F ) (z,v) 2 [(1 + a) a) + (1 (1+aa) )) + 0(a312))
(305)
2 [1 (1+ (1/2)a)
(306)
for all sufficiently small. It follows that the iterates 1/21 = (120)((o Mo) satisfy the recursion 1y2(+ ryl 2 (1+ (1/2)a) //2 as long as v < a/ (1+a) (1+ a1+a))))) We claim now that there
exists 7 E satisfying
TW In
v1 + a) (1 (1+))3)) In 1 20)
(307)
such that y/2 2 a/(11 a) (1 +)))) Suppose for a contradiction that this were not the case; then
v + (1/2)o) >
(308)
(1+a)(++
(1 (1+a)(1 + (1 (1+))
which is contradiction. This proves the result.
38

---

E Elementary lemmas
In this section, for ease of reference, we collect several well-known facts and their elementary proofs. Recall that subset M' of a Riemannian manifold M is geodesically corvex if any two points of M' can be connected by a unique geodesic contained entirely M'.
Lemma E.1. Let X M' IR. be C² function geodesically convex subset M' of Riemannian manifold M. If X is p-geodesically strongly convex in the sense that Vi & plTM at all points on M', then
(v)) + 1 dM (x, v)
( 2 A(x) + (V,
(309)
2
for all x,y E M'.
Proof. Fix x, y E MM.. Since M' geodesically convex, there exists a unique geodesic y [0, 1] +
M' such that 7(0) = x and (1) y. This satisfies
+()) = (0))))(()) = II(7), (v),
[[(()]I = dM( (x,y)
(310)
for all t E [0. where 11(7): To A - Tacej-M M is parallel transport. Define cpt)) = ((7(t)).
Then using (23) and the chain rule,
((() =
(()),
(()) = (+(1, V3(((()))())
(311)
for all E [0, 1]. Thus, by Taylor's theorem, there is s [0, 1] such that
X(v) (11 (0)) + (0) (s))
(312)
1
= A(x) + < Azz) log, (v)) + 2 (1()), V(rA(((s))(())
(313)
2 A(x) + (V.A(2)) 1o),(1)) + 4
|i())]1
(314)
2
= A(x) + (V Azx), lo))(w)) + dm (x,
(315)
2
as claimed.
Lemma E.2. Let X M + IR be a C2 function on geodesically convex subset M' of Riemannian manifold M. If X p-geodesically strongly convex in the sense that V 2 plTM all points on M', then
(12 VM)(2) V V Mog, 2 (d y)
(316)
for all x. y E M', where Tg T.M - TyM is parallel transport.
Proof. Fix x, y E M'. From Lemma E.1 one has the estimates
X(y) 2 (xx) + (V Ar A(r), log,, (v)) + dar y)2
(317)
X(x) 2 X V as ((y), log, (x)) + das y) 2
(318)
Since parallel transport is orthogonal transformation, one has
(V M log., (3)) = (V 11 log, = - (11a-r V
(319)
Substituting this into 317 and rearranging yields
m Mor), (z)) 2 (x, +(xx) -\y)
(320)
while 318 may be rearranged to give
aVar log,, (z)) > M (I, + x(y) - A(x).
(321)
Adding these estimates then yields the result.
39

---

Lemma E.3. Let X M' + I be C² function geodesicalty convex subset M' of Riemannian manifold M. If X p-geodesically strongly convex in the sense that V > = plTM at all points on M' and admits critical point v, in M', then x. is the unique global minimum of X in M' and
M 2 -
(322)
for all x E M'.
Proof.
That the global minimum is unique follows from Lemma E.1 for any x + in M, one has
(xx) 2 M(x,) + 2/2 das (1,, xx.) ((x.).
(323)
thus proving the first claim. To prove the second claim, fix x E M'. Lemma E.1 then applies to yield
X(x) Mxx.) log,
1 (a.))l?
M
(324)
2
L - (Mr(2))
1
log +08
+
(325)
2
2/1
$ 2/
\vr\(>)1
(326)
from which the result follows.
Lemma E.4. For any e > O, and any integers 0 S 8 < t, one has
c c
c c+t
+
(327)
c+8 8
c+c k c+
c c+
Proof. Since I + (ctx)
is monotonically decreasing along the positive reals, one has
f(z)dz < S. n+1
f(n)dr= f(n) = [" f(n)dx <
1m f(z)dz
(328)
nn
for any natural number n, from which it follows that
f (c+x) dx
Σ (c+ k)
(329)
(c+ k)
(c 8)
k=
k=ot1
(c+ {' 1 dr.
(330)
xc+
Integrating both sides now yields the result.
Experimental supplement
Our experiments were conducted for a depth 5, width scalar factorisation problem, corresponding to the function f : RD + R. defined by
0,) _(.......
0. (0...., E RRP,
0p))
(331)
with target value y = Our code runs gradient descent on this problem, and can be executed in seconds on single cpu. At each iterate, the projection onto the solution manifold M [-] is computed as follows.
The KKT conditions for the constrained optimisation problem
min _ ||0|| - 0112
1
p I i=1
such that
=
(332)
01 こリ 2
40

---

are given by
0. - + = 0,
i= 1=1,...51 Pb
(333)
0,
where a is the constraint parameter. These equations admit the quadratic solutions
0 (a) 0; + 03 4oy
i= 1.1.....P
(334)
2
It is casily verified 8 (a), is monotonically decreasing (when (334) is taken with the sign) or monotonically increasing (when (334) taken with the sign); consequently, the map I 0 (ai is also monotonic provided the signs are consistent among the (a).
The sign conventions we use are as follows. If [I 0 < y. then we must take the + sign for all (a). If IT= 0g y, then, noting that P I1 0g is the minimum value possible for II (o)
when all 6! (a) are taken with the positive sign in 334), either:
2 2 II O,, in which case we take the positive sign (334) for all v.
2. y < 2 II= 0g in which case we take the negative sign 334 for all i.
Having determined the sign
uniformly all 1, ..... p (334
assured that
+ ②(a) TT= 0!! (a) monotonic, hence that p(a) = admits a unique solution which we obtain using the bisection method.
41

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: See the "Paper contributions" section in the Introduction, which outlines our contributions in the context of codimension least squares problems.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The limitations are discussed in detail in the Discussion/CComcuusio section under several bold headings.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
42

---

Answer: [Yes]
Justification: Complete proofs of all theoretical results are provided in several appendices with appropriate referencing.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We included the code to produce our plots with the submission.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
43

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification:
We provide open access our code; see experimental supplement.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ "cipps.c/ppulli//guu /ttoe//uuuddeeoooouuuubmiiiiia ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All relevant hyperparameters etc are contained the code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification:
Our experiments are deterministic.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
44

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See the "Experimental supplement" appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp:///aarrrss.ccc/oooo//cccee
Answer: [Yes]
Justification:
We have reviewed and conformed with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work primarily theoretical and concerned with understanding in simple settings what practitioners have already been doing scale for over 10 years. No forseeable societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
45

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We considered only toy models in our experiments.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: No existing assets used.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
46

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets introduced.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooell as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at Ieast the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No crowdsourcing research with human subjects.
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
47

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: LLMs were not used outside of creation of code for toy experiments and minimal assistance with proving theorems. All material was thoroughly checked and revised if necessary.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy lattps://neurij: http://arrpp.ccccooooeerree///// for what should or should not be described.
48


Reviews:
[{"review_id": "5kYRp3AUqR", "reviewer": "Reviewer_Qn8i", "summary": "This paper studies the convergence of gradient descent in the \"edge-of-stability,\" or large step size, regime for a single data point and the square loss. The authors consider three learning rate regimes -- the subcritical regime, where $\\eta$ is small and the model converges to a minimizer with suboptimal sharpness; the supercritical regime, where $\\eta$ is large and the model converges linearly to the flattest minimizer, and a critical regime. The proof relies on the observation that GD on the loss projected onto the tangent space approximates reimannian GD on the sharpness, along with an assumption that the sharpness is strongly convex.", "questions": "- Theorem 4.4 further assumes that $\\eta$ is only very slightly larger than $2/\\lambda(\\theta_*^\\parallel)$. This also seems like a potentially restrictive assumption. Can you comment on what happens for larger $\\eta$?\n- What is meant by \"Given a tangent vector $v \\in T_{\\theta}\\mathbb{R}^p$, let us also denote by $v^\\parallel$ and $v^\\perp$ the tangent and orthogonal components of $v$, respectively\" in line 217? Isn't $v$ already a tangent vector?\n- There appears to be a typo in Theorem 4.2 (line 290) for the conditions on $\\eta$.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- It is of much interest in the community to better understand the edge-of-stability phenomenon. Prior works with convergence guarantees often rely on specific models/assumptions. This paper presents a very generic analysis framework for proving convergence to the flattest minimizer, which is to my knowledge novel.\n- The introduction of strong convexity of sharpness as an assumption to prove convergence, and observation that common toy models have such strong convexity, is also interesting.\n- The paper is well written and easy to understand. I have skimmed the proofs and they appear to be sound.\n\n## Weaknesses\n\n- The main limitation of this paper is that the analysis is \"local,\" i.e the main theorems require $\\theta_0$ to be sufficiently close to the manifold of minimizers, and in the case of Theorem 4.4, $\\theta_0^\\parallel$ to also be sufficiently close to the flattest minimizer $\\theta_*^\\parallel$. This is a similar style of analysis to prior works on label noise SGD [1,2,3], with the difference being that the oscillations from the edge of stability, rather than label noise, drive down the sharpness. As such, I find the claims that the paper analyzes the non-monotonic decrease of GD to be slightly overstated, as the theorems can only handle the case where the loss is already small. I also recommend the authors to add discussion on these prior works.\n\n- As the authors indeed mention, the analysis is limited to one training example (codimension 1)\n\n[1] Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeck like process. Guy Blanc, Neha Gupta, Gregory Valiant, Paul Valiant. COLT 2020.\n[2] Label noise SGD provably prefers flat global minimizers. Alex Damian, Tengyu Ma, Jason D. Lee. NeurIPS 2021.\n[3] What happens after SGD reaches zero loss? –a mathematical framework. Zhiyuan Li, Tianhao Wang, Sanjeev Arora. ICLR 2022.", "was_revised": true, "final_justification": "I believe that this paper makes valuable contributions towards understanding the dynamics of gradient descent at the edge of stability. The authors responded to my questions on the local nature of the analysis, and will provide more clarification on the specific conditions required in the next revision. I thus recommend acceptance for this paper.", "rebuttal": "Thank you for your considered reading and critique of our work. In summary it seems that your primary critique is about the local nature of the analysis. Our analysis for Theorem 4.2 is local only in a weak sense that is also true for state of the art results in Riemannian optimization more generally; moreover, we can demonstrate empirically that a much stronger kind of locality is *necessary* for convergence in Theorems 4.3 and 4.4. More details below.\n\n**Weaknesses**\n\nWe are happy to include discussion of the works you mentioned in the next version.\n\n> The main limitation of this paper is that the analysis is \"local,\" i.e the main theorems require $\\theta_0$ to be sufficiently close to the manifold of minimizers and in the case of Theorem 4.4, $\\theta_0^{\\parallel}$ to also be sufficiently close to the flattest minimizer\n\n(a) Concerning closeness of $\\theta_0$ to the manifold of minimizers: this is indeed a limitation in the broader context of optimization, but in the context solely of edge of stability (EOS) dynamics (which is our concern in this paper) we do not believe this to be a limitation since EOS dynamics typically don't begin until the loss has already decreased substantially from initialisation [1].\n\n(b) Concerning assumptions on $d_0:=d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$: analysis of Riemannian gradient descent on a geodesically strongly convex objective on a curved manifold generally requires that the initial distance to the objective is bounded by some finite constant $D$; the constants appearing in the convergence analysis then depend on $D$ and grow arbitrarily large as $D$ grows large due to the curvature [2,3,4]. In this sense all Riemannian gradient descent convergence theorems require an \"initial distance to optimum is sufficiently small\" (i.e. $\\leq D$) assumption, although $D$ can be taken as large as one likes (for the price of worsening convergence rate) provided that the objective remains geodesically strongly convex and the neighbourhood remains uniquely geodesic. Our own analysis, since it boils down to Riemannian gradient descent in the parallel direction, must also make this assumption. For Theorem 4.2 specifically, this is the *only* assumption required on $d_0:=d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$, so \"sufficiently small\" in Theorem 4.2 means $O(1)$ and can be taken as large as any finite-diameter, uniquely geodesic neighbourhood on which the strong convexity of $\\lambda$ holds. Thus Theorem 4.2 is no more local than state of the art results in Riemannian optimisation generally. We will make this point clearer in the next version.\n\n(c) Compared with Theorem 4.2, the assumptions on $d_0$ must be strengthened in Theorems 4.3 and 4.4 due to the fact that we *in addition* need to control the $\\theta^{\\perp}$-dynamics at each point relative to the limiting $\\theta^{\\perp}$-dynamics (which, unlike in Theorem 4.2, do not convergence exponentially to zero).  This requires $d_0$ to be small relative to these limiting $\\theta^{\\perp}$-dynamics (e.g. $d_0 = O((\\eta\\lambda(\\theta_*^{\\parallel})-2)^{3/2})$ in Theorem 4.4). Moreover, we can demonstrate empirically that this locality is *necessary*: even in the simplest case of a two layer scalar linear network, having learning rate $\\eta\\geq 2/\\lambda(\\theta_*^{\\parallel})$ means that if $d_0$ is sufficiently large, so that $\\eta>2/\\lambda(\\theta_0^{\\parallel})$ is sufficiently large, optimisation diverges for any $\\|\\theta^{\\perp}_0\\|\\neq 0$. Thus while it may be possible to quantitatively weaken the $d_0$ assumption in Theorems 4.3 and 4.4, it is not possible to remove the assumption entirely. We are happy to provide code for this experiment if requested.\n\n\n**Questions**\n\n>Theorem 4.4 further assumes that $\\eta$ is only very slightly larger than $2/\\lambda(\\theta_*^{\\parallel})$. This also seems like a potentially restrictive assumption. Can you comment on what happens for larger $\\eta$?\n\nGreat question. For larger $\\eta$, the periodicity of the orbit to which the dynamics converge increase in a \"period-doubling\" fashion (i.e., to an orbit of period 4, then to an orbit of period 8, and so on) before eventually becoming chaotic (see Lines 324-330). This \"period doubling\" phenomenon has already been observed empirically for edge of stability dynamics in prior work [5], and is well-known to occur for perturbed quadratic systems of the type we consider ([6], Section 10.3). We did not include analysis of this case because it would have required a much lengthier analysis of the quadratic dynamical system than what we have supplied, and it seemed to us that the analysis would have to be performed on a case-by-case basis (a separate theorem for each period, with determining the $2^n$th-periodic oscillation requiring the solution of a $2^n$th degree polynomial followed by a perturbation analysis). With the paper already being quite long and technical, we felt it best to leave investigation of this question to future work.\n\n>What is meant by \"Given a tangent vector $v\\in T_{\\theta}\\mathbb{R}^p$, let us also denote by $v^{\\parallel}$ and $v^{\\perp}$ the tangent and orthogonal components of $v$, respectively\" in line 217? Isn't already a tangent vector?\n\nThanks for picking up on this, we should have been clearer. $v\\in T_{\\theta}\\mathbb{R}^p$ is indeed a tangent vector *to Euclidean space*, but $v^{\\parallel}$ is its projection along *the tangent space to $M$*, while $v^{\\perp}$ is its projection *orthogonal* to $M$. We will clarify this in the final version.\n\n>There appears to be a typo in Theorem 4.2 (line 290) for the conditions on $\\eta$.\n\nThank you for spotting this, you are correct: the condition should read $1/\\lambda(\\theta_*^{\\parallel})<\\eta<2/\\lambda(\\theta_*^{\\parallel})$ rather than $1/\\lambda(\\theta_*^{\\parallel})<\\eta<2/\\lambda(\\theta_0^{\\parallel})$.\n\nPlease let us know if you would like any further clarifications.\n\n1. Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, Cohen et al, ICLR2021\n2. First order methods for geodesically convex optimization, Zhang et al, COLT, 2016\n3. Convergence and Trade-Offs in Riemannian Gradient Descent and Riemannian Proximal Point, Rubio et al, ICML2024\n4. On the Convergence of Gradient Descent for Finding the Riemannian Center of Mass, Afsari et al, SIAM Journal on Control and Optimization, Vol. 51,  Issue 3, 2013\n5. Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos, Kalra et al, ICLR2025\n6. Nonlinear Dynamics and Chaos: With applications to Physics, Biology, Chemistry and Engineering, Steven Strogatz, CRC Press, 2015\n7. Normal Forms, Differentiable Conjugacies, and Elementary Bifurcations of Maps, Glendinning et al, SIAM Journal on Applied Mathematics, Volume 83, Issue 2, 2023"}, {"review_id": "Ve41r0c6fk", "reviewer": "Reviewer_SHZE", "summary": "The paper investigates the edge of stability behavior of large step gradient descent (GD) in a codimension one overparameterized least squares problem (with a single input data point). As the authors argue, the setting strikes a balance between being overly minimal and overly general, allowing for clear insights into the dynamics. Moreover, unlike several prior works that focus on logistic regression, which show that GD iterates eventually enter a stable regime where the loss decreases monotonically, this work presents a case that the iterates never reach such a stable regime. This observation could be of independent interest.\n\nUnder this paper's setting, the set of global minimizers forms a Riemannian manifold, enabling a decomposition of the GD dynamics. The authors show that the sharpness $\\lambda$ is geodesically strongly convex along the manifold for two instances. Assuming that the manifold is symmetric with respect to a reference point $\\theta_*^{\\|\\|}$, they analyze the convergence behavior of GD when the iterates remain within a neighborhood of $\\theta_*^{\\|\\|}$. In particular, the authors identify three regimes exhibiting distinct behaviors depending on whether the step size is smaller than, equal to, or larger than $2/\\lambda(\\theta_*^{\\|\\|})$.", "questions": "- Lines 39-40: Unlike [31,30,6], which shows that large step GD eventually enters a regime where the logistic loss decreases monotonically, the paper [21] focuses on the least squares loss and does not appear to provide a similar analysis. Could you clarify this distinction?\n- Line 51: I suggest clarifying that the GD update approximations in equation (2) will be formally derived later.\n- Line 71: Could you cite relevant papers to support this point?\n- Lines 112-113: Where was the oscillatory implicit bias observed? How does it align with the technical findings in this paper?\n- Line 147: I recommend rephrasing to make it clear that this is an assumption, later shown to be satisfied by several examples. As written, the statement may initially make the assumption seem artificial, even though justifications follow soon after.\n- Lines 183-184: Could you clarify why this is the case?\n- Line 220: I suggest stating Assumption 3.1 explicitly here for clarity.\n- Line 239: Why is the sharpness $\\lambda$ characterized through the geodesic strong convexity of itself?\n- Line 245-246: It would be helpful to briefly summarize the main message of prior works involving bifurcation diagrams here.\n- Line 248: I do not fully follow why the observations above motivate the consideration of the nominal dynamical system described in equations (15-16). A clearer explanation would be appreciated.\n- Line 281: $\\lambda(\\theta_0^{\\|\\|}) > \\lambda(\\theta_*^{\\|\\|})$ appears to implicitly assumed here. It would be helpful if the authors could make this explicit and provide some explanation or justification. Additionally, is this assumption also made in the analysis of the other two regimes?\n- Line 299: This line seems to treat $\\theta_*^{\\|\\|}$ as the optimally flat optimal minimizer. If that is indeed the case, it would be beneficial to elaborate on this interpretation earlier in the paper.", "limitations": "The paper clearly list several limitations at the end. However, I believe the assumption that the iterates remain near the reference point $\\theta_*^{\\|\\|}$ constitutes an additional non-negligible limitation that deserves explicit discussion.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n- The decomposition of GD dynamics into components tangential and orthogonal to the manifold, as presented in Theorem 4.1, and their characterization in terms of the sharpness $\\lambda$ is novel and insightful. (While I did not have sufficient time or expertise to fully verify the proof of this theorem, the proof sketch appears sound.)\n- Demonstrating that Assumptions 3.1 and 3.2, which initially seemed quite strong, hold in settings such as multilayer scalar factorization and scalar factorization by two vectors is interesting. This helps illuminate the structure of the manifold and the property of the sharpness $\\lambda$ along it.\n\nWeaknesses\n- Although the study of edge of stability in least squares problems is valuable, the paper does not clearly motivate why it is important to extend this analysis beyond the logistic loss setting.\n- The main conclusion relies on the assumption that the iterates remain within a neighborhood of the reference point  $\\theta_*^{\\|\\|}$, which may limit the generalizability of the results to broader regions of the manifold.\n- This paper is technically dense, which makes it challenging to extract the main message. A clearer high-level overview of intuitive summary (especially, of Sections 4.2-4.4) would improve accessibility for a broader audience,", "was_revised": true, "final_justification": "Many of my questions were requests for clarification, and the authors responded well, which helped me better position this paper.\nWhile the local nature of the results was initially a concern, I now understand that only the initialization needs to be close to the reference point—not all iterates—which makes the assumption less restrictive.\nOther reviewers also raised concerns about this local restriction, but like them, I do not see it as a critical limitation. Overall, I believe this paper may offer valuable insight into understanding the EoS phenomenon", "rebuttal": "Thank you for your detailed reading and critique. In summary, it seems your primary critiques concern lack of clarity and an assumption that the iterates remain near the fixed point. We have attempted to clarify the points you raised, as well as clarify that we only assume the *initialisation* is near the fixed point, and explain why this assumption is necessary. See below for details.\n\n**Weaknesses**\n\n> the paper does not clearly motivate why it is important to extend this analysis beyond the logistic loss setting\n\nThe square loss is frequently used in applications to scene synthesis [1] and has been shown to be equally performant to logistic losses for classification [2]. Whether one uses the square loss or logistic loss, one observes edge of stability convergence with large learning rates [3]. The logistic loss setting is just one side of a more general set of phenomena which are still poorly understood; our work is a step towards a more complete understanding.\n\n> The main conclusion relies on the assumption that the iterates remain within a neighborhood of the reference point\n\nWe only assume that the *initial* iterate is in a neighborhood of the reference point; our theorems then *prove* that all further iterates remain in the neighborhood (see Lines 261 through 272 where this is identified explicitly). Moreover:\n\n(a) Concerning assumptions on $d_0:=d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$: analysis of Riemannian gradient descent on a geodesically strongly convex objective on a curved manifold generally requires that the initial distance to the objective is bounded by some finite constant $D$; the constants appearing in the convergence analysis then depend on $D$ and grow arbitrarily large as $D$ grows large due to the curvature [4]. In this sense all Riemannian gradient descent convergence theorems require an \"initial distance to optimum is sufficiently small\" (i.e. $\\leq D$) assumption, although $D$ can be taken as large as one likes (for the price of worsening convergence rate) provided that the objective remains geodesically strongly convex and the neighborhood remains uniquely geodesic. Our own analysis, since it boils down to Riemannian gradient descent in the parallel direction, must also make this assumption. For Theorem 4.2 specifically, this is the *only* assumption required on $d_0:=d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$, so \"sufficiently small\" in Theorem 4.2 means $O(1)$ and can be taken as large as any finite-diameter, uniquely geodesic neighborhood on which the strong convexity of $\\lambda$ holds. Thus Theorem 4.2 is no more local than state of the art results in Riemannian optimisation generally. We will make this point clearer in the next version.\n\n(b) Compared with Theorem 4.2, the assumptions on $d_0$ must be strengthened in Theorems 4.3 and 4.4 due to the fact that we *in addition* need to control the $\\theta^{\\perp}$-dynamics at each point relative to the limiting $\\theta^{\\perp}$-dynamics (which, unlike in Theorem 4.2, do not convergence exponentially to zero).  This requires $d_0$ to be small relative to these limiting $\\theta^{\\perp}$-dynamics (e.g. $d_0 = O((\\eta\\lambda(\\theta_*^{\\parallel})-2)^{3/2})$ in Theorem 4.4). Moreover, we can demonstrate empirically that this locality is *necessary*: even in the simplest case of a two layer scalar linear network, having learning rate $\\eta\\geq 2/\\lambda(\\theta_*^{\\parallel})$ means that if $d_0$ is sufficiently large, so that $\\eta>2/\\lambda(\\theta_0^{\\parallel})$ is sufficiently large, optimisation diverges for any $\\|\\theta^{\\perp}_0\\|\\neq 0$. We are happy to provide code for this experiment if requested.\n\n> This paper is technically dense...\n\nThis is a fair critique. We will include a higher-level overview of these sections at the beginning of Section 4 in the next version.\n\n**Questions**\n\nWe are happy to make the changes you have suggested. Concerning your questions:\n\n> Lines 39-40...Could you clarify this distinction?\n\nThe logistic loss in the papers [30,31,6] has sharpness tending to zero as one approaches infinity in parameter space in the direction of perfect classification. The analysis in these papers demonstrates an implicit bias in this direction implying that eventually, *irrespective of how large the learning rate $\\eta$ is*, the algorithm will find a point at which the sharpness is $<2/\\eta$ (stable regime). For the square loss considered in [21] the sharpness is nonvanishing, having minimal value $\\lambda_*>0$, and the learning rate $\\eta$ *must be taken small enough* that $\\eta<2/\\lambda_*$ to ensure eventual stability. The setting in [21] thus corresponds to the subcritical regime in our paper (Theorem 4.2). We also consider the square loss in our paper, but consider regimes (the critical and supercritical regime) in which $\\eta\\geq 2/\\lambda_*$, which are not covered by [21].\n\n\n> Line 71: Could you cite relevant papers to support this point?\n\nYes: the subcritical regime (our Theorem 4.2) has been observed empirically and analysed theoretically in [5]. Their convergence theorem is similar to ours, but they work with a specific (width 2, scalar output) model rather than a *class* of models as we do. The supercritical regime (our Theorem 4.4) has been observed empirically in [6] for 2 layer linear neural nets and in [7] for deep scalar linear neural nets, but neither work supplies a convergence theorem comparable to ours. The critical regime (our Theorem 4.3) has not been identified or analysed in any other works to our knowledge.\n\n> Lines 112-113: Where was the oscillatory implicit bias observed?\n\nThe oscillatory implicit bias has been observed in [6], where it is shown that the sharpness itself oscillates during training and exhibits a bifurcation diagram as the learning rate is increased. Our theory demonstrates theoretically that these oscillations are a consequence of the quadratic dynamical system on $\\theta^{\\perp}$ (Eq. (12)); such systems are well-known to exhibit such oscillations and bifurcation diagrams [9]. Our theory covers the first \"fork\" in the bifurcation diagram. We will attempt to make this point clearer in the final version.\n\n> Lines 183-184: Could you clarify why this is the case?\n\nThe reason is that ReLUs do not change the loss landscape in a neighborhood of the solution manifold. This is most easily seen in the depth 2 case. The solutions of $\\theta_2\\theta_1 = 1$ form a hyperbola admitting two connected components: one in the $\\theta_1,\\theta_2>0$ quadrant and one in the $\\theta_1,\\theta_2<0$ quadrant. The solutions of $\\theta_2\\phi(\\theta_1) = 1$ form *exactly the same* hyperbola component in the $\\theta_1,\\theta_2>0$ quadrant, with no solutions occurring in the $\\theta_1,\\theta_2<0$ quadrant. Since we restrict our analysis to only the component in the positive quadrant, the solution manifold is the same for our analysis in each case.\n\n> Line 239: Why is the sharpness $\\lambda$ characterized through the geodesic strong convexity of itself?\n\nWe need to word this more clearly. What we intend to say is that the geodesic strong convexity can be used to control the sharpness of the solution to which one converges: in the subcritical regime, the $\\theta^{\\parallel}$ dynamics are Riemannian gradient descent on $\\lambda$ with exponentially decreasing stepsize, so they converge to a suboptimally flat minimum. In the critical and supercritical regimes, the $\\theta^{\\parallel}$ dynamics are Riemannian gradient descent on $\\lambda$ with polynomially decreasing and constant step size respectively, which, since $\\lambda$ is geodesically strongly convex, is sufficient to guarantee convergence to the minimiser of $\\lambda$.\n\n> Line 248: I do not fully follow...\n\nThe nominal systems are motivated by looking at Equations (11) and (12). Equation (11) looks like gradient descent on $\\lambda$ with time-varying step size, but since it is supposed to occur along the solution manifold $M$, we are motivated to compare it with Riemannian gradient descent on $\\lambda$, which is the natural generalisation of gradient descent to manifolds. Equation (12) is a quadratic dynamical system but with time-varying coefficients; the nominal system just fixes the coefficients to make the system easier to study. We will attempt to be clearer in the final version.\n\n> Line 281: $\\lambda(\\theta_0^{\\parallel})>\\lambda(\\theta_*^{\\parallel})$ appears to implicitly assumed here\n\nYes, $\\lambda(\\theta_0^{\\parallel})>\\lambda(\\theta_*^{\\parallel})$ is implicitly assumed throughout because (as you have identified) $\\theta_*^{\\parallel}$ is indeed the optimally flat optimal minimzer. We have made this explicit at Line 164, however we will attempt to make this clearer in the next version.\n\nPlease let us know if you would like any further clarification.\n\n1. NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis, Mildenhall et al, ECCV2020.\n2. Evaluation of Neural Architectures Trained with Square Loss vs Cross-Entropy in Classification Tasks, Hui et al, ICLR2021.\n3. Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, Cohen et al, ICLR2021\n4. Convergence and Trade-Offs in Riemannian Gradient Descent and Riemannian Proximal Point, Rubio et al, ICML2024\n5. A Minimalist Example of Edge-of-Stability and Progressive Sharpening, Liu et al, arXiv, 2025.\n6. Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos, Kalra et al, ICLR2025.\n7. Gradient Descent Monotonically Decreases the Sharpness of Gradient Flow Solutions in Scalar Networks and Beyond, Kreisler et al, ICML2023.\n8. Nonlinear Dynamics and Chaos: With applications to Physics, Biology, Chemistry and Engineering, Steven Strogatz, CRC Press, 2015."}, {"review_id": "XPdnQzemZD", "reviewer": "Reviewer_CvCh", "summary": "This paper proves a generalized convergence result for gradient descent on the edge of stability in some “overparametrized least-squares models” with some generic assumptions, which generalizes several previous papers on certain deep scalar factorization. Based on the assumptions, the paper can rigorously prove local convergence in sub-critical, critical, and supercritical regimes (characterized by the sharpness of the flattest minimizer) where the initialization and flattest convergence point $\\theta_*^{||}$ is sufficiently close.", "questions": "Questions:\n1. Is there any more general case satisfying the assumptions? \n2. Is it possible to remove any part of the assumptions in the theorem, e.g. $d_M(\\theta_0^{||},\\theta_*^{||})$ sufficiently small or dynamics-related assumptions?\n3. Why Theorem 4.2 need to assume the $d_M(\\theta_0^{||},\\theta_*^{||})$ is sufficiently small? The final convergence point is some sub-optimal solution with larger sharpness.\n4. A minor question for the title: why stress overparameterization least squares? I am not sure where “overparameterization” is technically mentioned in the paper, though it might be a necessary condition for the assumptions to hold.", "limitations": "See weakness.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n\n- This paper generalizes some toy models in the previous works with generic assumptions on the solution manifold. It offers a more generalized local convergence view of deep scalar multiplication. The paper also unifies the different regimes of EoS by characterizing the sharpness of the flattest minima. The technical contribution is solid.\n\nWeakness: \n- The scope of the results are a bit limited. The Assumptions are generalized, but still very strong. For example, it cannot go beyond scalar-valued least squares.\n- Theorem 4.2-4.4 are mostly local convergence with the conditions that $d_M(\\theta_0^{||},\\theta_*^{||})$ is sufficiently small, which weakens the result. Theorem 4.3 and 4.4 even additionally assume the dynamics have sufficiently small higher-order terms, which is a very strong assumption that needs to hold along the training trajectories.", "was_revised": true, "final_justification": "The authors settled most of my concerns in the rebuttal. Specifically, the authors managed to lift the strong assumptions (training dynamics related) in the main theorems. Though the scope is limited, I believe this paper makes its contribution in rigorous analysis within the EoS regime. Now I lean slightly toward acceptance if (1) Theorem 4.2 can be restated as the rebuttal claimed (2) the third-order term assumptions are lifted. I appreciate the authors' effort after submission, and I have increased the score to 4.", "rebuttal": "Thank you for your attention to our paper. In summary, it seems your main critiques regard our strong assumptions and local analysis. We have been able to weaken some of the assumptions. Moreover, the analysis for Theorem 4.2 is local only in a weak sense which holds also for state of the art results in Riemannian optimization generally, while for Theorems 4.3 and 4.4 we can demonstrate empirically that the locality is necessary. See below for details.\n\n**Weaknesses**\n\n> The scope of the results are a bit limited. The Assumptions are generalized, but still very strong. For example, it cannot go beyond scalar-valued least squares.\n\nWe agree that our assumptions are strong. However, we request that the current state of knowledge be kept in mind. Edge of stability was only empirically identified in 2021, and theoretical understanding of it is still extremely limited. The phenomenon cannot be explained by classical optimisation theory and its theoretical explanation even in toy examples such as scalar nets is very much an open and active area of research [7,8]. Our work is a significant step forward compared with existing literature since it is the only one that can prove quantitative convergence theorems in the unstable regimes considered.\n\n> Theorem 4.2-4.4 are mostly local convergence with the conditions that $d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$ is sufficiently small, which weakens the result. Theorem 4.3 and 4.4 even additionally assume the dynamics have sufficiently small higher-order terms, which is a very strong assumption that needs to hold along the training trajectories\n\n(a) Concerning the locality of the analysis: please see the answer to your question below.\n\n(b) Concerning the \"third order terms\" assumption, your critique is entirely valid and we completely agree that this is too strong. Since submitting, we have discovered the assumption can be removed entirely using a \"normal form\" from bifurcation theory (Theorem 3.4, [6]), whose hypotheses are all satisfied by the examples we consider. Using this result, all our theorems continue to hold and retain identical quantitative rates.\n\n**Questions**\n\n> Is it possible to remove any part of the assumptions in the theorem, e.g. $d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$ sufficiently small or dynamics-related assumptions? Why Theorem 4.2 need to assume the $d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$ is sufficiently small? The final convergence point is some sub-optimal solution with larger sharpness.\n\nThese are great questions.\n\n(a) Analysis of Riemannian gradient descent on a geodesically strongly convex objective on a curved manifold generally requires that the initial distance to the objective is bounded by some finite constant $D$; the constants appearing in the convergence analysis then depend on $D$ and grow arbitrarily large as $D$ grows large due to the curvature [2,3,4]. In this sense all Riemannian gradient descent convergence theorems require an \"initial distance to optimum is sufficiently small\" (i.e. $\\leq D$) assumption, although $D$ can be taken as large as one likes (for the price of worsening convergence rate) provided that the objective remains geodesically strongly convex and the neighbourhood remains uniquely geodesic. Our own analysis, since it boils down to Riemannian gradient descent in the parallel direction, must also make this assumption. For Theorem 4.2 specifically, this is the *only* assumption required on $d_0:=d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$, so \"sufficiently small\" in Theorem 4.2 means $O(1)$ and can be taken as large as any finite-diameter, uniquely geodesic neighbourhood on which the strong convexity of $\\lambda$ holds. Thus Theorem 4.2 is no more local than the state of the art results in Riemannian optimisation generally. We will make this point clearer in the next version.\n\n(b) Compared with Theorem 4.2, the assumptions on $d_0$ must be strengthened in Theorems 4.3 and 4.4 due to the fact that we *in addition* need to control the $\\theta^{\\perp}$-dynamics at each point relative to the limiting $\\theta^{\\perp}$-dynamics (which, unlike in Theorem 4.2, do not convergence exponentially to zero).  This requires $d_0$ to be small relative to these limiting $\\theta^{\\perp}$-dynamics (e.g. $d_0 = O((\\eta\\lambda(\\theta_*^{\\parallel})-2)^{3/2})$ in Theorem 4.4). Moreover, we can demonstrate empirically that this locality is *necessary*: even in the simplest case of a two layer scalar linear network, having learning rate $\\eta\\geq 2/\\lambda(\\theta_*^{\\parallel})$ means that if $d_0$ is sufficiently large, so that $\\eta>2/\\lambda(\\theta_0^{\\parallel})$ is sufficiently large, optimisation diverges for any $\\|\\theta^{\\perp}_0\\|\\neq 0$. Thus while it may be possible to quantitatively weaken the $d_0$ assumption in Theorems 4.3 and 4.4, it is not possible to remove the assumption entirely. We are happy to provide code for this experiment if requested.\n\n(c) We can remove the \"third order terms sufficiently small\" assumption as detailed above without changing the convergence theorems or the validity of our examples.\n\n(d) We can relax Assumption 3.1 to the assumption that at every point on the line orthogonal to $M$ at $\\theta_*^{\\parallel}$, the loss gradient is orthogonal to $M$ at $\\theta_*^{\\parallel}$. We can prove that this condition is *necessary* for Theorem 4.4 (the proof is easy: if gradient descent oscillates between two points on the line orthogonal to $M$ at $\\theta_*^{\\parallel}$, then the gradient itself, being parallel to this line, must also be orthogonal to $M$ at $\\theta_*^{\\parallel}$). The examples we cite in the paper all satisfy this relaxed assumption. We can also empirically demonstrate the necessity of this relaxed assumption (see the paragraph that follows).\n\n> Is there any more general case satisfying the assumptions?\n\nIn place of ReLU on Line 186 (multilayer scalar factorisation of depth $L$) one may take any nonlinearity which is the identity on a neighbourhood of $y^{1/L}$, where $y>0$ is the target value. A nice *counterexample* is the 2 layer nonlinear scalar network $f(\\theta_1,\\theta_2) = \\theta_2\\tanh(\\theta_1)$ when the target value $y>0$; this example does *not* satisfy the weakened assumption in (d) above (but satisfies all the rest of our assumptions), and it can be seen by numerically plotting the trajectory that, in the supercritical regime, convergence to a stable period-2 orbit centred on $\\theta_*^{\\parallel}$ *does not hold*; convergence is instead to a stable period-2 orbit centred on a *different* point, of *suboptimal* sharpness. This is to be contrasted with the linear and ReLU case, in which convergence to a stable period-2 orbit about the point of *optimal sharpness* $\\theta_*^{\\parallel}$ *does* occur, as predicted by our theory. We will cite this $\\tanh$ counterexample as an additional limitation of the theory, however this counterexample is itself of interest: it suggests that the flatness-seeking implicit bias of edge of stability dynamics is *mediated* by the choice of nonlinearity in such a way that it can hold for ReLUs but not for other activations (such as $\\tanh$). To our knowledge, this is itself a new observation that will be of great interest to the community studying the optimization of deep nets. We are happy to supply code to run this experiment if required.\n\n> A minor question for the title: why stress overparameterization least squares?\n\nThis is a point worth clarifying. For a model $f$, overparametrisation is often used in the optimisation literature to guarantee that the neural tangent kernel $df df^T$ is positive definite [5], which is the same as requiring that $df$ is full rank. We insist on precisely this full-rank property at lines 132-134 by asking that $y$ be a regular value. Thus, we regard this assumption as an overparametrisation assumption. Without it, the solutions need not form a smooth manifold and our analysis need not apply. We will attempt to make this clearer in the next version.\n\nPlease let us know if you would like any further clarifications.\n\n1. A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks, Arora et al, ICLR 2019.\n2. First order methods for geodesically convex optimization, Zhang et al, COLT, 2016.\n3. Convergence and Trade-Offs in Riemannian Gradient Descent and Riemannian Proximal Point, Rubio et al, ICML2024\n4. On the Convergence of Gradient Descent for Finding the Riemannian Center of Mass, Afsari et al, SIAM Journal on Control and Optimization, Vol. 51,  Issue 3, 2013.\n5. Memorization and Optimization in Deep Neural Networks with Minimum Over-parameterization, Bombari et al, NeurIPS2022.\n6. Normal Forms, Differentiable Conjugacies, and Elementary Bifurcations of Maps, Glendinning et al, SIAM Journal on Applied Mathematics, Volume 83, Issue 2, 2023.\n7. Gradient Descent Monotonically Decreases the Sharpness of Gradient Flow Solutions in Scalar Networks and Beyond, Kreisler et al, ICML 2023.\n8. Beyond the Edge of Stability via Two-step Gradient Updates, Chen et al, ICML 2023."}, {"review_id": "C4KbMeQVEF", "reviewer": "Reviewer_QFWU", "summary": "This paper studies the effect of step sizes in gradient descent through the phenomenon named edge of stability. The authors analyze three regimes: subcritical, critical, and supercritical. The analysis is based on decomposing the training dynamics into the parallel and orthogonal directions with respect to the solution manifold. Under the assumption of certain symmetry, isometry, and convexity, the authors prove the convergence results in each regime.", "questions": "The questions follow from \"Weakness\":\n1. Could you give examples where all your assumptions are satisfied?\n2. Could you discuss why those assumptions are necessary (not just sufficient) and provide examples demonstrating such necessity?\n3. Could you conduct experiments that are more closely related to machine learning in practice?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\n\n1. The intuitions are well explained.\n2. The authors state the assumptions that the theorems are based on.\n3. The proof of theoretical results are given coherently in the appendix.\n\nWeakness:\n\n1. The assumptions are so strong (symmetry, isometry, strong convexity) that it is difficult to find a reasonable setting where all assumptions are satisfied.\n2. The experiments are based on simple settings such as scalar network (depth-5, width-1) and shallow linear network (depth-2, width-5). It is not convincing that that this will be a general behaviors for machine learning models.", "was_revised": true, "final_justification": "I appreciate the fact that authors are honest with the limitation of their work, and upon critique they improve their results by loosening the assumptions. I still believe that the simplification of the model and the strong assumptions are strong limitations, which are also pointed out by other reviewers, but at the same time I agree that theory on edge of stability is so sparse that every progress counts.", "rebuttal": "Thank you for your time reviewing and critiquing our paper. In summary, it seems your main critiques concern the strong assumptions and whether they are necessary. While our assumptions are strong and our examples are toy models only, they are no stronger than those of comparable works in the literature (which also deal only with toy models), while our results are much stronger. We are able to both prove and empirically demonstrate that one of our assumptions (suitably weakened) is *necessary* for our supercritical convergence theorem. See below for details.\n\n**Weaknesses**\n\n> The assumptions are so strong (symmetry, isometry, strong convexity) that it is difficult to find a reasonable setting where all assumptions are satisfied.\n\nPlease note that we only insist on convexity of *the sharpness along the solution manifold*; we do not insist on convexity of the loss itself. We give two classes of neural network examples, with *non-convex loss landscapes*, where all our assumptions are satisfied (namely, factorization of a scalar either by other scalars or by two vectors, Lines 175-191). While we agree these are toy models, the present state of understanding of edge of stability dynamics is so limited that theoretical analysis is always restricted to simple models such as scalar linear networks [1] and two layer linear networks [2]. Unlike these other works, however, our assumptions are formulated in general mathematical terms which do not rule out possibly more complex models in the future and are suggestive of the mathematical features that should be considered when studying more complex models. Thus, even though our assumptions are restrictive, they are still an advance over prior works.\n\n> The experiments are based on simple settings such as scalar network (depth-5, width-1) and shallow linear network (depth-2, width-5). It is not convincing that that this will be a general behaviors for machine learning models.\n\nWe make no claim that the precise assumptions we make, or the empirical observations we make, would extend to practical models like ResNets or Transformers. However, the *theory* of edge of stability is still primarily restricted to toy models [1,2], which must be understood before more complex models can be understood. By formulating our theory in general mathematical terms rather than in terms of specific toy architectures like in prior works, our theory is still an advance over prior works and a step towards understanding more complex models.\n\n**Questions**\n\n> Could you give examples where all your assumptions are satisfied?\n\nPlease see Lines 175-191 for examples satisfying the assumptions we make (namely, factorization of a scalar either by other scalars or by two vectors). While we agree these are toy examples, we emphasise again that the theory of edge of stability is still so undeveloped that it has not progressed past such toy examples [1,2]. We cannot expect to have a theory of practical models without first having a theory of toy models, but by formulating our theory in general mathematical terms our paper is a step beyond prior works and a step toward bridging this gap.\n\n> Could you discuss why those assumptions are necessary (not just sufficient) and provide examples demonstrating such necessity?\n\nGreat question.\n\n(a) We have discovered since submission that the symmetry assumption (Assumption 3.1, Lines 151-152) can be weakened as follows. Recall that we use $\\theta_*^{\\parallel}$ to denote the minimiser of the sharpness along the solution manifold $M$, and $n$ to denote the unit normal vector field along $M$. Then Assumption 3.1 can be weakened to the assumption that at each point $\\theta:=\\theta_*^{\\parallel}+\\alpha n(\\theta_*^{\\parallel})$ orthogonal to $M$ at $\\theta_*^{\\parallel}$, the loss gradient $\\nabla\\ell(\\theta)$ is also orthogonal to $M$ at $\\theta_*^{\\parallel}$. The examples we cite satisfy this weakened assumption, and this weakened assumption *is* necessary for convergence to a stable period-2 orbit orthogonal to $M$ at the point $\\theta_*^{\\parallel}$. The proof is immediate: if gradient descent oscillates between two points on a line orthogonal to $M$ at $\\theta_*^{\\parallel}$, then necessarily the loss gradient must be parallel to this line, hence orthogonal to $M$. We are able to produce an empirical demonstration of this necessity, comparing the iterates of gradient descent (after approximate convergence to a stable orbit) for the 2 layer scalar linear network $f(\\theta_1,\\theta_2) = \\theta_2\\theta_1$ with a 2 layer scalar ReLU network $f(\\theta_1,\\theta_2) = \\theta_2 ReLU(\\theta_1)$ and a 2 layer scalar $\\tanh$ network $f(\\theta_1,\\theta_2) = \\theta_2\\tanh(\\theta_1)$ to fit the target value $4$. In the linear and ReLU cases, the assumption is satisfied, leading to oscillations orthogonal to $M$ at exactly the point of minimum sharpness $\\theta_*^{\\parallel}$; in the $\\tanh$ case, the assumption is not satisfied, and it is clear from the plot that the oscillations do *not* occur about the point of minimum sharpness $\\theta_*^{\\parallel}$. This is a novel finding, suggesting that the properties of the nonlinearity play a key role in determining the implicit bias of edge of stability dynamics.\n\n(b) We can demonstrate empirically that Theorems 4.3 and 4.4 do not hold if $d_0:=d(\\theta_0^{\\parallel},\\theta_*^{\\parallel})$ is too large. Specifically, for a two layer scalar linear network, if $\\eta\\geq 2/\\lambda(\\theta_*^{\\parallel})$ and $d_0$ is too large, gradient descent starting from $\\|\\theta^{\\perp}_0\\|\\neq0$ diverges, violating Theorem 4.3 and 4.4. While it may be possible to quantitatively weaken the assumptions on $d_0$ in these theorems, this experiment demonstrates that theorems are *necessarily* local. We are happy to provide code for these simple experiments upon request.\n\n> Could you conduct experiments that are more closely related to machine learning in practice?\n\nEdge of stability experiments for more practical models such as ResNets, VGGs and Transformers on more practical datasets such as CIFAR10 and WikiText have been conducted in prior work [3]. However, these experiments only plot the sharpness at each iterate $\\theta_t$, rather than at the *projection* $\\theta^{\\parallel}_t$ of each iterate to the solution manifold as we require. To conduct experiments reflective of our theory on practical models, we would need a scalable way of projecting iterates to the solution manifold as well as a method for finding the minimum value of the sharpness along the solution manifold. To our knowledge neither of these calculations have been attempted in the literature before. While we agree that these experiments would be desirable, we do not know how to perform them and believe they would be sufficiently difficult to warrant their own paper.\n\nPlease let us know if you would like any further clarifications.\n\n1. Gradient Descent Monotonically Decreases the Sharpness of Gradient Flow Solutions in Scalar Networks and Beyond, Kreisler et al, ICML2023.\n2. Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos, Kalra et al, ICLR2025.\n3. Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, Cohen et al, ICLR2021."}]

Author rebuttal:

Meta-review: The paper investigates the edge-of-stability phenomenon of gradient descent under large stepsize. By focusing on the overparametrized least square problem, the paper provides a detailed analysis of the phenomenon. Based on the local characterization of the solution set which is a Riemannian manifold, the authors manage to decompose the gradient descent dynamics, which further allows to  identify three regimes exhibiting distinct behaviors depending on  the stepsize: subcritical, critical, and supercritical. 


**Strength:** 
 - The paper provides an insightful perspective to understand the edge-of-stability phenomenon. 
 - The paper is well presented.

**Weakness:** 
 - The analysis is only for "local", and limited on one training example.
 - Relative strong assumptions, which limits the applicability of the result. 

**Reason of decision:**  
The topic of the paper is important and has broad interests, the obtained result delivers an interesting and insightful perspective to understand the edge-of-stability phenomenon. 

**Discussion summary:** The authors are suggested to adopt the reviewers' suggestions to revise the final version of the paper, not limited to the following points
 - Better justification and discussion on the result only for "local" analysis.
 - Re-organization of the paper for better presentation.
 - Modification of Theorem 4.2, and weaken assumption by lifting third-order term assumption.
Decision: Accept (poster)