PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding Gradient Descent on Edge of Stability in Deep Learning

Sanjeev Arora, Zhiyuan Li, Abhishek Panigrahi

arXiv:2205.09745v3cs.LGmath.OC

TL;DR

Edge-of-Stability dynamics challenge traditional assumptions that gradient descent requires η < 2/λ for stable loss decrease. This paper analyzes two finite-step settings and proves loss reduction with dynamics that regulate sharpness along minimum-loss manifolds.

  • Problem

    Existing analyses struggle to explain how finite-step gradient descent can decrease loss and regulate sharpness after violating the traditional η < 2/λ condition.

    Source: §1 · pp. 1–2

  • Method

    The paper analyzes Normalized Gradient Descent and constant-rate gradient descent on transformed loss, under regularity assumptions near a manifold of minimizers.

    Source: §1.1 · p. 2, §7 · p. 13

  • Results

    Theoretical results show both settings enter Edge of Stability, while Normalized Gradient Descent tracks a manifold flow that minimizes the Hessian’s largest eigenvalue.

    Source: §4.3 · p. 7, §7 · p. 13

  • Takeaways & Limitations

    Finite learning rates can produce implicit regularization through deterministic motion on minimum-loss manifolds, including rigorous sharpness reduction in one setting.

    Source: §7 · p. 13

    The analysis applies only near manifolds of local minimizers, whereas experiments observe Edge-of-Stability behavior beginning substantially earlier.

    Source: §7 · p. 13

Abstract

from arXiv · show

Deep learning experiments by Cohen et al. [2021] using deterministic Gradient Descent (GD) revealed an Edge of Stability (EoS) phase when learning rate (LR) and sharpness (i.e., the largest eigenvalue of Hessian) no longer behave as in traditional optimization. Sharpness stabilizes around $2/$LR and loss goes up and down across iterations, yet still with an overall downward trend. The current paper mathematically analyzes a new mechanism of implicit regularization in the EoS phase, whereby GD updates due to non-smooth loss landscape turn out to evolve along some deterministic flow on the manifold of minimum loss. This is in contrast to many previous results about implicit bias either relying on infinitesimal updates or noise in gradient. Formally, for any smooth function $L$ with certain regularity condition, this effect is demonstrated for (1) Normalized GD, i.e., GD with a varying LR $η_t =\fracη{\| \nabla L(x(t)) \|}$ and loss $L$; (2) GD with constant LR and loss $\sqrt{L- \min_x L(x)}$. Both provably enter the Edge of Stability, with the associated flow on the manifold minimizing $λ_{1}(\nabla^2 L)$. The above theoretical results have been corroborated by an experimental study.

1 Introduction

The paper challenges traditional stability assumptions by analyzing Edge of Stability (EoS), where sharpness exceeds 2/η and loss decreases non-monotonically. It develops two provable EoS mechanisms—non-smooth losses and adaptive learning rates—and shows that GD induces a limiting flow toward lower-sharpness minimizers.

1 Introduction

  • [Cohen et al., 2021, Ahn et al., 2022] found EoS behavior in which sharpness exceeds 2/η, later stabilizes or decreases, and loss oscillates locally while declining over longer timescales.These observations contradict the traditional condition η < 2/λ under which the Descent Lemma guarantees iteration-wise loss decrease.Source: §1 · p. 1
  • The paper argues that these phenomena require new analyses because GD can keep reducing loss despite violating η < 2/λ while regulating sharpness.Li et al. [2020b] had already noted that normalization with weight decay can force arbitrarily high sharpness near zero loss, whereas Cohen et al. [2021] observed EoS even without normalization.Source: §1 · p. 1
  • Stableness evaluates whether a point–learning-rate pair satisfies a local Descent-Lemma condition by taking the maximum Hessian sharpness along the proposed update segment.It is defined as η times the supremum of λ_1(∇^2L(x − s∇L(x))) for 0 ≤ s ≤ η, with stability threshold 2.Source: §1 · p. 2

1.1 Two Provable Mechanisms for Edge of Stability: Non-smoothness and Adaptivity

  • The paper studies two EoS mechanisms: non-smooth transformed losses and smooth losses trained with effectively adaptive learning rates.The latter is exemplified by Normalized GD, x ← x − η∇L/∥∇L∥, whose effective learning rate grows as ∇L approaches zero.Source: §1.1 · p. 2
  • Under the stated manifold and Hessian-rank assumptions, sufficiently small-step GD first follows gradient flow, then enters an EoS phase with non-monotone loss and oscillations across the minimum-loss manifold.As η → 0, the GD trajectory converges to a limiting flow along the manifold’s tangent direction that decreases sharpness.Source: §1.1 · pp. 2–3
  • In the second phase, GD operates at EoS: average stableness over two consecutive steps is at least 2, while the induced manifold flow decreases sharpness.The paper states that this provides a formal deterministic mechanism for implicit selection of lower-sharpness solutions, unlike analyses relying on SGD stochasticity [Li et al., 2022b].Source: §1.1 · p. 3
  • Although the proofs inject occasional O(η^100) perturbations, experiments show the theoretical predictions directly for deterministic GD without perturbations.The authors attribute the perturbation requirement to a technical limitation of the current analysis.Source: §1.1 · p. 3
  • The analysis establishes gradient–top-eigenvector alignment for Normalized GD on quadratic losses using a novel potential, yielding a chaotic power-iteration-like dynamic.This alignment supports the proof that the manifold projection moves in the direction that reduces sharpness.Source: §1.1 · p. 3

2 Related Works

Prior work links sharpness to generalization, analyzes limiting dynamics near minimizer manifolds, and studies implicit bias and non-monotone GD convergence. This paper differs by identifying finite-learning-rate sharpness reduction along minimizers at the Edge of Stability.

Sharpness

  • Sharpness has been associated with flat minima and generalization [Hochreiter and Schmidhuber, 1997, Keskar et al., 2016], but parameter rescaling can alter sharpness without changing the represented function [Dinh et al., 2017].This ambiguity motivated scale-invariant sharpness notions [Yi et al., 2019a,b, Tsuzuku et al., 2020, Rangamani et al., 2021] and corresponding regularization algorithms [Yi et al., 2021, Kwon et al., 2021].Source: §2 · p. 3

Limiting Diffusion/Flow around Manifold of Minimizers

  • Prior analyses explain SGD’s limiting flow near minimizer manifolds through label noise and stochastic differential equations [Blanc et al., 2020, Damian et al., 2021, Li et al., 2022b].The stochastic framework builds on convergence results for SDEs with large drift.Source: §2 · pp. 3–4

Implicit Bias

  • Unlike many implicit-bias results that extend to infinitesimal learning rates, this paper’s sharpness reduction along the minimizer manifold requires finite learning rates.Related finite-rate mechanisms include stability analysis [Wu et al., 2017, Ma and Ying, 2021] and implicit gradient regularization [Barrett and Dherin, 2021].Source: §2 · p. 4

Non-monotone Convergence of Gradient Descent

3 Warm-up: Quadratic Loss Functions

For positive-definite quadratic loss, Normalized GD iterates align with the top Hessian eigenvector, producing non-monotonic loss oscillations characteristic of the Edge of Stability. The proof uses invariant-set preparation followed by monotonic top-eigenspace alignment.

Main result

  • Normalized GD converges in direction to the top eigenvector v1, establishing the quadratic Edge-of-Stability behavior.Theorem 3.1 gives alternating even- and odd-step limits proportional to v1 under the nonzero top-component condition.Source: §3 · p. 4

Assumptions

  • The theorem assumes the top-eigenvector component remains nonzero; if a component vanishes, it remains zero and analysis can restrict to the corresponding subspace.This condition is required for the stated convergence result.Source: §3 · pp. 4–5

Main result

  • The loss oscillates between consecutive iterates rather than decreasing monotonically, indicating Edge-of-Stability dynamics.This follows directly from the alternating asymptotic iterate behavior.Source: §3 · pp. 4–5

Proof strategy

  • The trajectory first enters invariant sets, then aligns with v1 as its top-eigenvector projection increases monotonically on low-norm steps.The analysis separates preparation from alignment; once the iterate norm is at most 0.5λ1, the top projection is monotone and bounded.Source: §3 · p. 5
  • When the iterate norm is at most λ1/2, the top-eigenvector projection grows by Θ(θ_t^2) over two steps, forcing the angle θ_t to vanish.Continuity of the normalized update extends convergence from low-norm steps to all iterates.Source: §3 · p. 6

4 Main Results

Under a smooth-loss setting with a manifold of minimizers and a positive top-eigenvalue eigengap, the paper proves two-phase dynamics for Normalized GD and GD on L: after approaching the manifold, updates follow sharpness-reducing flows and operate at the Edge of Stability.

4.2 Key Assumptions on Manifold of Local Minimizers

  • The analysis assumes a C4 loss whose zero-loss local minimizers form a C2 manifold, with maximal normal-space Hessian rank and a positive top-eigenvalue eigengap.These conditions ensure the gradient-flow limiting map is sufficiently regular and λ1(∇2L) is differentiable for the main theorems.Source: §4.2 · p. 7

4.3 Results for Normalized GD

  • Normalized GD first moves O(η) closer to the minimizer manifold, then maintains an O(η)-average angle to the top Hessian eigenspace while tracking Riemannian gradient flow for log(λ1(∇2L)) on the manifold.The Phase II result holds under the stated initialization and residual conditions, with probability at least 1 − O(η10) while the limiting flow exists.Source: §4.3 · pp. 7–8

4.4 Results for GD on L

  • GD on L likewise has an initial O(η)-approach phase, followed by a Phase II trajectory that implicitly minimizes sharpness near the minimizer manifold.The Phase II approximation holds under Assumptions 4.1–4.2 and the theorem’s initialization and residual conditions, with probability at least 1 − O(η10) while the limiting flow exists.Source: §4.4 · pp. 8–9

4.5 Operating on the Edge of Stability

  • Both methods enter the Edge of Stability: stability is at least 2 in at least one of every two consecutive steps, loss oscillates across two steps, and its two-step average decreases over time.For Normalized GD, the reciprocal-stableness relation across consecutive steps equals 1 + O(θt + η).Source: §4.5 · p. 9

5 Proof Overview

The proof centers on the gradient-flow limit map Φ, which tracks Normalized GD near the minimum-loss manifold and exposes its slow tangential motion. It then analyzes preparation, top-eigenvector alignment, manifold-flow convergence, and the analogous GD-on-L dynamics.

5.1 Properties of Φ

  • For small η, Φ(xη(t)) captures Normalized GD’s implicit regularization because iterates remain O(η)-close to the minimum-loss manifold, while Φ moves at most O(η^2) per step.A second-order Taylor expansion of Φ is reused throughout both phases of the proof.Source: §5.1 · pp. 9–10

5.2 Analysis for Normalized GD

  • Normalized GD first enters an O(η)-neighborhood, then reaches an invariant region where every residual R_j is O(η^2) for Θ(η^-2) steps.The first subphase uses ODE approximation and the PL condition; the second generalizes the approximately quadratic preparation argument.Source: §5.2 · p. 10
  • During Phase II, the iterate and gradient approximately align with the Hessian’s top eigenvector, yielding convergence of Φ(xη(τ/η^2)) to the flow X(τ).The proof derives this alignment from a Taylor expansion around Φ and then establishes trajectory closeness over bounded τ.Source: §5.2 · pp. 10–11
  • The Phase II proof requires an initial alignment of at least Ω(η), while poly(η) perturbations can prevent alignment collapse; experiments need no added noise.The proof uses instability near the top eigenspace, though the authors note floating-point errors may provide sufficient perturbations.Source: §5.2 · p. 11

5.3 Analysis for GD on L

  • For GD on L under L(x)=0 on Γ, the analogous analysis yields a manifold movement lemma, with an effective learning rate determined by λ1(t)η.The update differs from Normalized GD near the manifold, but the proof follows a similar strategy.Source: §5.3 · p. 11

6 Experiments

Experiments on VGG-16 and a two-layer MNIST network support the predicted Edge-of-Stability phenomena and limiting-flow convergence. The results show decreasing sharpness, near-perfect gradient alignment, stableness near 2, and close agreement between Normalized GD and Riemannian flow.

Verification for Predicted Phenomena on Real-life Models

  • The VGG-16 experiment used 5000 randomly selected training examples, GeLU and average pooling, and ℓ2 loss instead of softmax loss.The study tracked top Hessian eigenvalue, alignment, stableness, and test accuracy.Source: §6 · p. 12
  • Normalized GD on VGG-16 trained with CIFAR-10 exhibits decreasing top Hessian eigenvalue, alignment near 1, and stableness near 2.The experiment used GeLU, average pooling, and ℓ2 loss to satisfy the theory’s smoothness and minimizer requirements.Source: §6 · p. 12

Verifying Convergence to Limiting Flow on MNIST

  • The MNIST convergence test used a 1-hidden-layer, 784-unit GeLU network trained with ℓ2 loss on a random subset of 1000 examples.Both runs started from a checkpoint with loss of order 10^-3.Source: §6 · p. 13
  • Normalized GD and Riemannian gradient flow show an exact match across test loss, test accuracy, top Hessian eigenvalue, and Hessian trace on MNIST.Their parameter trajectories also remain close throughout training under the specified continuous-time scalings.Source: Figure 5 · p. 12, §6 · p. 13, Figure 6 · p. 13

7 Conclusion

The paper establishes Edge of Stability behavior in two finite-learning-rate settings, including loss decrease after stability is lost and sharpness reduction in one setting. Its analysis is local to minimizer manifolds, leaving broader EoS behavior and SGD as future work.

Conclusion

  • Two general-loss settings show gradient updates can decrease loss over many iterations after stability is lost, with finite learning rate η rather than infinitesimal updates.The analysis contrasts with recent implicit-bias results requiring infinitesimal learning rates.Source: §7 · p. 13
  • In one setting, the trajectory reduces sharpness, rigorously establishing an effect conjectured in deep-learning literature and documented for GD by Cohen et al. [2021].Sharpness is the maximum Hessian eigenvalue.Source: §7 · p. 13
  • The analysis applies only close to the manifold of local minimizers, whereas experiments show EoS and sharpness control beginning much sooner.Bridging this gap and analyzing EoS for L itself rather than the transformed loss used here is left for future work.Source: §7 · p. 13
  • Exploring EoS-like effects in SGD remains future work, pending definitive experiments analogous to Cohen et al. [2021].The authors anticipate that broader analysis may require new understanding of deep-learning loss properties.Source: §7 · p. 13

A Omitted Proofs for Results for Quadratic Loss Functions · A.1 Proofs for Preparation Phase

The appendix fixes a diagonal positive-definite quadratic setting and proves the preparation properties underlying Normalized GD. Under a nonzero top-eigenvector component, the iterates converge along even and odd subsequences to signed top-eigenvector states, while invariant-set arguments establish entry into the preparation region.

A Omitted Proofs for Results for Quadratic Loss Functions

A.1 Proofs for Preparation Phase

  • Normalized GD’s coordinate-subspace sets Ij are invariant, and every initialization eventually enters their intersection, providing the preparation-phase framework.Source: §A.1 · p. 19
  • Corollary A.3 establishes finite-time entry into I1 from any initialization using the norm decrease bound while ||x̃(t)|| remains at least λ1.Source: §A.1 · p. 20
  • Corollary A.4 shows that starting in I1, each coordinate subspace Ij is reached after a finite bound, with invariance preserving membership thereafter.Source: §A.1 · p. 20

A.2 Proofs for Alignment Phase

This subsection proves that normalized GD aligns with the top eigenvector during the alignment phase after preparation. The analysis establishes bounds on the alignment angle and shows monotonic growth of the top-coordinate magnitude under a norm condition.

Alignment-phase setup

  • Normalized GD is analyzed after preparation, when the iterate remains in the alignment-phase domains and is shown to align with the top eigenvector.The proof specifically assumes ˜x(t) belongs to the stated intersection of domains throughout the alignment phase.Source: §A.2 · p. 20

Alignment bounds

  • The alignment proof bounds the evolution of the top-eigenvector component and the angle θt, using the spectral gap λ1−λ2 and the smallest eigenvalue through λ = min(λ1−λ2, λD).The coordinate-wise update analysis derives the right-hand bound using λ1−λ2, λD, and the condition ∥˜x(t)∥ < λ1/2.Source: §A.2 · p. 22

Monotonic top-coordinate growth

  • If ∥˜x(t+1)∥ + ∥˜x(t)∥ ≤ λ1, then Lemma A.7 guarantees |˜x1(t+2)| ≥ |˜x1(t)|, with equality only at the boundary.This monotonicity result is obtained by combining the normalized GD updates and the preceding alignment bound.Source: §A.2 · pp. 22–23

Norm contraction estimates

  • The estimates give ∥˜x(t + 1)∥ ≤ λ1 − ∥˜x(t)∥ − λ(λ1 − λ) sin² θt, with λ = min(λ1 − λ2, λD).They also imply ∥˜x(t + 1)∥∥˜x(t)∥ ≤ λ2.Source: §A.2 · p. 23

A.3 Proof of Main theorems for Quadratic Loss … C.1 Phase I, Convergence

This section develops the dynamical and geometric tools underlying the main results, from quadratic-loss convergence and control of general-loss deviations to the local structure of normalized GD near the minimum-loss manifold. It then characterizes the limiting map and shows how Phase I tracks normalized gradient flow toward the low-loss manifold while identifying the projected descent direction for the top Hessian eigenvalue.

A.3 Proof of Main theorems for Quadratic Loss

  • The quadratic-loss proof separates preparation, where iterates enter an invariant set, from alignment, where the top-eigenvector projection increases monotonically until convergence.The limiting projection implies convergence in direction to the top eigenvector.Source: §A.3 · p. 23
  • As t→∞, when C = 1/2, ||x̃(t)|| converges to λ1, and for sufficiently large t the top coordinate x̃1(t) flips sign at every step.Source: §A.3 · p. 23

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

  • For general losses, auxiliary lemmas provide quantitative perturbation controls by extending quadratic-case invariant-set, coordinate, and angle-contraction arguments.These lemmas are used only for the general-loss proof, not the quadratic-loss proof.Source: §A.4 · pp. 24, 25, 27

B Setups for General Loss Functions

  • Under the general-loss setup, normalized GD is proved to follow gradient flow in phase one and a sharpness-decreasing limiting flow on the minimizer manifold in phase two.For sufficiently small η, the trajectory remains in neighborhoods of the gradient-flow and limiting-flow trajectories, with rates determined by local smoothness and Φ constants.Source: B Setups for General Loss Functions · p. 28
  • The analysis assumes an open neighborhood where gradient flow converges to the minimum-loss manifold, with normalized GD entering a controlled working zone satisfying PL, eigengap, and eigenvalue bounds.A smaller working zone provides additional properties used in the later analysis.Source: B Setups for General Loss Functions · p. 30

B.1 Geometric Lemmas

  • The geometric lemmas relate distance to the limiting manifold, gradient-flow path length, loss, and gradient norm, while keeping trajectories and connecting segments inside the valid neighborhood.Within this zone, normalized GD is approximately quadratic around Φ(x), its one-step projection movement is O(η^2), and its loss decreases.Source: §B.1 · pp. 30, 36

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

  • The limiting map Φ is defined on an open set U, is locally C^2 there, and satisfies ∂Φ(x)∇L(x)=0, linking gradient-flow convergence to tangent-space projection.At points on Γ, ∂Φ is the projection onto the tangent space, while tangent directions have zero Hessian action.Source: §B.2 · pp. 36–37
  • At minimum-manifold points, the nonzero Hessian eigenvectors lie in the normal space, separating normal curvature directions from tangent directions.This structure supports the local quadratic description of normalized GD around Φ(x).Source: §B.2 · p. 37
  • The limiting-map flow follows −1/2 times the tangent-space projection of ∇log(λ1(∇2L(x))), thereby reducing the top Hessian eigenvalue along the manifold.A variant gives the corresponding projected gradient of λ1 itself.Source: §B.2 · p. 37

C.1 Phase I, Convergence

  • For sufficiently small η, Phase I Normalized GD tracks normalized gradient flow from xinit, enters a neighborhood Y_ϵ, and reaches the minimizing manifold while the loss converges quickly.The argument uses time-rescaled gradient-flow dynamics, a Polyak–Łojasiewicz condition in Y_ϵ, and O(η) ODE-approximation error.Source: §C.1 · pp. 37–38
  • Once iterates are O(η)-close to the limiting map within Y_ϵ, Lemma C.1 controls their subsequent evolution provided all intervening steps remain in that neighborhood.This approximate-quadratic control completes the Phase I convergence argument for Theorem 4.3.Source: §C.1 · pp. 38–39

C.2 Phase II, Limiting Flow · D Phase I, Omitted Proofs of the Main Lemmas

Theorem 4.4 shows that perturbed Normalized GD implicitly minimizes sharpness by tracking a limiting flow under specified initialization and regularity conditions. Its proof controls discrete–continuous error through a differential inequality and bounds distance from the minimum-loss manifold.

C.2 Phase II, Limiting Flow

  • Theorem 4.4 establishes that perturbed Normalized GD implicitly minimizes sharpness under Assumptions 4.1–4.2 and specified initialization conditions.The theorem holds for sufficiently small η, with probability at least 1 − O(η^10), while the limiting-flow solution exists.Source: §C.2 · p. 39
  • The theorem requires O(η)-closeness to the manifold, O(η^2) residual bounds, and an Ω(η) top-eigenspace component or negative residual.These conditions constrain initialization and maintain the regime needed for the Phase II approximation.Source: §C.2 · p. 39
  • The proof closes by bounding the average angle between the trajectory’s displacement and the top Hessian eigenspace by O(η), enabling uniform tracking through time T2.The stopping-time argument shows the iterates remain in the required neighborhood for sufficiently small η.Source: §C.2 · pp. 39–40
  • The proof compares the discrete manifold movement with the limiting flow using a piecewise-linear interpolation and a Lipschitz differential inequality.The interpolation runs through Φ(xη(t)) at times tη^2, while the vector field is the projected sharpness-gradient flow.Source: §C.2 · pp. 39–40
  • Residual bounds propagate across iterations, ensuring O(η) distance from the minimum-loss manifold throughout the analyzed time interval.The argument also shifts the analysis by one step when the initial norm condition is not immediately satisfied.Source: §C.2 · pp. 39–40
  • The discrete trajectory moves along the manifold according to the projected gradient of log sharpness, up to errors controlled by η^2 and off-manifold deviation.This movement estimate is the discrete counterpart of the limiting flow used in the convergence argument.Source: §C.2 · p. 40

D.1 Proof of Lemma C.1 … E.4 Dynamics in Top Eigenspace When Dropping Below Threshold

These sections analyze Normalized GD near a quadratic model, establishing alignment, cycle structure, and perturbed dynamics in the top Hessian eigenspace. They also prove that random perturbations escape low-alignment regions and complete Lemma E.16 through error control, induction, and lower bounds.

D.1 Proof of Lemma C.1

  • The nonlinear trajectory remains O(η)-close to its projection onto the minimum-loss manifold, with O(η²) quadratic-model error despite changing Hessian eigenspaces.The proof accounts for errors introduced by time-varying Hessians and groups nearby eigenvalues into stable eigenspaces.Source: §D.1 · pp. 41–42

D.2 Properties of the condition in Equation (25)

  • Once the alignment condition holds, it persists, while the iterate cannot remain above 0.5ηλ1(t) + Ψnormη² for more than one step.The same condition also controls changes in the angle to the top eigenvector when the norm is near this threshold.Source: §D.2 · pp. 43–44

E Phase II, Omitted Proofs of the Main Lemmas

  • Phase II assumes O(η)-initial proximity, O(η²) residuals, and an Ω(η) top-eigenvector component, enabling long-horizon angle analysis over Ω(1/η²) steps.The analysis introduces top-eigenspace alignment measures and divides the trajectory into cases based on their magnitudes.Source: E Phase II, Omitted Proofs of the Main… · p. 44, §E.2 · pp. 45–46

E.2 Time Average of Angles Against Top Eigenspace

  • Over intervals of length at least Ω(1/η²), the cumulative angle to the top eigenspace is O(ηt2) with probability at least 1 − O(η¹⁰).The proof analyzes alignment, intermediate decay, and perturbation-driven recovery cases; only O(1) intermediate or recovery episodes occur.Source: E Phase II, Omitted Proofs of the Main… · p. 44, §E.2 · pp. 46–47

E.1 Dividing Time Steps Into Cycles of Length 1 or 2

  • The time partition into N0, N1, and N2 forms disjoint one- or two-step cycles, with N0 and N1 satisfying the sub-threshold norm bound.For a step in N0, neighboring memberships are constrained by Lemma E.1; Lemma E.4 transfers a nontrivial top-component condition to the latest preceding N1 step.Source: §E.1 · pp. 44–45

E.2 Time Average of Angles Against Top Eigenspace

  • If alignment conditions persist, the angle against a non-top eigendirection must grow exponentially, forcing one condition to fail within O(log 1/η) steps.The argument uses a quadratic-model approximation with an additional O(η^2) error.Source: §E.2 · p. 48

E.3 Dynamics in the Top Eigenspace

  • Under Equation (29), small η, bounded iterates, and trajectory containment in Y_ϵ, one- and two-step lemmas characterize perturbed dynamics along the top eigenvector.The two-step analysis controls angle changes and eigenvector movement using O(η) and O(η^2) perturbation bounds.Source: §E.3 · p. 49

E.4 Dynamics in Top Eigenspace When Dropping Below Threshold

  • After uniform perturbation in B_xη(t)(r), Normalized GD escapes below-threshold top alignment within Θ(log 1/η) steps with probability at least 1−η^12.The stuck set occupies only a thin slice of width at most η^1/2r, with volume ratio O(η^1/2).Source: §E.4 · p. 51
  • The proof begins by recalling the definition of the error term and invoking Equation (37) to establish a property for all ˜t ≥ 0.The property holds for some ϱ > 0.Source: §E.4 · p. 55
  • An induction argument propagates the required bound from the base case t = 0 to successive even iterates.The induction step considers ˜t + 2 after assuming the claim for all even 0 ≤ t′ ≤ ˜t.Source: §E.4 · p. 55
  • Defining ϕ = min˜t≤tesc ∥M˜t∥, the proof uses the previously established fact that ϕ > 1 to derive the needed inequality.The argument concludes by choosing C sufficiently large so that the required condition involving ϱ and C holds.Source: §E.4 · p. 56

E.5 Proof for Operating on Edge of Stability

The proof establishes that the local sharpness quantity satisfies SL(x_η(t), η_t) = η_t(λ_1(t) + O(η)), using the residual bound R_j(x_η(t)) ≤ O(η^2). It also derives an analogous quadratic-case relation involving the scaled iterates.

E.5 Proof for Operating on Edge of Stability

  • The residual bound R_j(x_η(t)) ≤ O(η^2) yields SL(x_η(t), η_t) = η_t(λ_1(t) + O(η)).The proof obtains this by expanding the supremum of λ_1 along the gradient step.Source: §E.5 · p. 57
  • Together, these steps complete the proof’s first claim and set up the argument for its second claim.The passage explicitly marks the first claim as completed before introducing the second claim.Source: §E.5 · p. 57
  • The proof then states the corresponding relation η^-1(∥x̃_η(t)∥ + ∥x̃_η(t + 1)∥) = λ_1(t) + O(η + θ_t), analogous to the quadratic case.This identifies the error terms controlling the comparison with the quadratic analysis.Source: §E.5 · p. 57

F Some Useful Lemmas About Eigenvalues and Eigenvectors … G.3 Phase II, limiting flow

The appendix develops eigenvalue/eigenvector perturbation tools and applies them to GD with transformed notation, first establishing Phase I convergence and then showing Phase II follows a limiting flow with controlled error. The analysis also quantifies angle accumulation and distinguishes GD’s weaker top-eigenvector control from Normalized GD.

F Some Useful Lemmas About Eigenvalues and Eigenvectors

  • The appendix collects smooth eigenvalue/eigenvector derivative results, symmetric-matrix eigenvalue perturbation bounds, and Davis–Kahan eigenvector stability tools.The Davis–Kahan result is formulated for eigenspaces under unitary invariant norms, including Frobenius and spectral norms.Source: F Some Useful Lemmas About Eigenvalues… · p. 57
  • Compared with Normalized GD, GD lacks the stronger top-eigenvector projection behavior, producing differences in angle accumulation and error rates.The appendix notes that the stronger bounds for Normalized GD rely on its specific update rule and do not hold for GD with the transformed loss.Source: F Some Useful Lemmas About Eigenvalues… · p. 58, §G.3 · p. 60

G.1 Notations

  • The appendix introduces notation for the projected residual ˜xη(t) and denotes Y as the limiting flow from Equation (7).These definitions provide the common notation for the subsequent convergence and limiting-flow analysis.Source: §G.1 · p. 58

G.2 Phase I, convergence

  • Phase I establishes a stability property that persists along the trajectory once the stated stability condition and proximity requirements hold.The proof reuses the noisy GD update and bounds on movement in the projection map Φ.Source: §G.2 · p. 58

G.3 Phase II, limiting flow

  • GD’s Phase II trajectory stays close to the limiting flow, with cumulative update error O(η^1/2) over the analyzed interval.The per-step movement follows the projected gradient of λ1, while the accumulated error is O(η^1/2).Source: §G.3 · p. 59
  • For intervals of length Ω(η^-2), the average angle is O(√η), under small-learning-rate and trajectory-stability conditions.Lemma G.4 assumes controlled initial residuals and a nontrivial top-eigenvector component, with probability at least 1 − η^10.Source: §G.3 · p. 59
  • The Phase II update is governed by the Hessian at the projected point, with an O(η^2) perturbation when the top-eigenvector component is Ω(η).This quadratic-model approximation makes the GD updates O(1)-Lipschitz under the stated lower bound.Source: §G.3 · pp. 59–60

G.4 Omitted Proof for Operating on Edge of Stability · G.5 Geometric Lemmas for L · H Additional Experimental Details

The appendices complete the Edge-of-Stability proof using geometric bounds and develop lemmas controlling the loss landscape near the minimizer manifold. The supplied passages provide no substantive experimental details for H Additional Experimental Details.

G.4 Omitted Proof for Operating on Edge of Stability

  • For M≥2, the distance from Φ(xη(t)) to the consecutive iterates is O(η(θt+θt+1)), yielding Hessian growth of order 1/[η(θt+θt+1)].The argument concludes using θt+1=O(θt).Source: §G.4 · p. 60
  • The omitted proof also relates the transverse component to the loss through L(xη(t))=||x̃η(t)||+O(η), paralleling the quadratic-case estimate.The consecutive transverse norms satisfy ||x̃η(t)||+||x̃η(t+1)||=ηλ1(t)+O(η+θt).Source: §G.4 · p. 60

G.5 Geometric Lemmas for L

  • G.5 uses that Φ(x) is a zero-loss local minimizer, so ∇L(Φ(x))=0, and invokes curvature lower bounds around Φ(x).The supplied derivation bounds the displacement by O(ζ^1/2) after scaling by µ.Source: §G.5 · pp. 60–61
  • The geometric lemmas analyze one-step updates y=x−η∇L(x) using Taylor expansion of Φ and bounds on ∇L(x)/√L(x).The proof uses Φ∈C3, a property of Φ that cancels the first term, and Lemma G.7 for the remaining gradient bound.Source: §G.5 · pp. 61–62
  • The analysis further connects the curvature of Φ at Φ(x) to the top Hessian eigenvector v1(x) and derivatives of ∇L.It uses the C3 regularity of Φ and a relation involving λ1(x), ∂Φ(Φ(x)), and ∂2(∇L)(Φ(x))[v1(x),v1(x)].Source: §G.5 · p. 62

H.1 Experimental details · H.2 Implementation Details for Simulation for the Limiting Flow of Normalized GD

H.1 specifies initialization and learning-rate settings for the GD and Normalized GD experiments, while H.2 provides an Algorithm 4 simulation of the limiting Riemannian flow for Normalized GD. The implementation covers admissible averaged losses, tangent-space projection, top-eigenvector computation, and gradient-descent retraction to the minimum-loss manifold.

H.1 Experimental details

  • H.1 initializes GD at (x, y) = (14.7, 3.) with η = 0.5 and Normalized GD at (x, y) = (14.7, −3) with η = 5.Source: §H.1 · p. 62
  • For Figure 2, Normalized GD starts with ⟨v1, ˜x(0)⟩ = 10−4 and ⟨v2, ˜x(0)⟩ = 0.45, using learning rate 1.Source: §H.1 · p. 62

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

  • Algorithm 4 simulates one step of the Riemannian flow (5) corresponding to Normalized GD.Source: §H.2 · pp. 62–63
  • The algorithm assumes L is the average of n nonnegative per-example losses ℓi(x) = ℓ(fi(x), bi), where ℓ(y, b) = 0 iff y = b.This setting covers regression, binary classification through 0, 1 labels, and multi-class classification via squared class-score errors.Source: §H.2 · pp. 62–63
  • Each Riemannian update computes ∇3L(x)[v1(x), v1(x)], projects onto the tangent space, and uses small-step gradient descent to return to the manifold.Source: §H.2 · p. 63
  • Algorithm 4 computes the top eigenvector v1 of ∇2L, obtains ∇λ1 from ∇3L[v1, v1], projects its gradient component, and takes the update y(0) ← x(t) − η λ1(x(t))(I − Px(t),Γ)∇λ1(x(t)).Source: §H.2 · p. 63
  • The inner loop applies y(˜t + 1) = y(˜t) − ηproj∇L(y(˜t)) for Tproj iterations, then sets x(t + 1) = y(Tproj).For Figures 5 and 6, the implementation uses 100 power iterations for v1, L(x(0)) = 3.803 × 10−3, η = 10−2, ηproj = 10−2, Tproj = 103, and Normalized GD learning rate 10−2.Source: §H.2 · p. 63

Related papers

Loading 2205.09745v3…