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

The paper addresses how deterministic GD can continue reducing loss in the Edge of Stability regime despite violating traditional stability conditions. It analyzes two finite-learning-rate mechanisms and proves that their trajectories approach flows on minimizer manifolds that reduce sharpness, while identifying scope limitations near those manifolds.

  • Problem

    Traditional analyses do not explain EoS behavior, where sharpness exceeds 2/η and loss oscillates while continuing to decrease.

  • Method

    The paper analyzes non-smooth transformed losses and Normalized GD under manifold and Hessian-rank assumptions, using stableness and limiting-flow arguments.

  • Results

    The two settings provably exhibit EoS dynamics, with GD trajectories approaching manifold flows that reduce sharpness or its logarithm-related objective.

  • Takeaways & Limitations

    Finite learning rates can induce deterministic implicit regularization toward flatter minimizers without relying on stochastic gradient noise.

  • Takeaways & Limitations

    The analysis applies only near manifolds of local minimizers, whereas experiments observe EoS and sharpness control beginning substantially earlier.

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 analyzes Edge of Stability dynamics, where finite-step GD can continue reducing loss despite violating the traditional stability condition. It identifies non-smoothness and effective learning-rate adaptivity as two mechanisms producing motion toward flatter minimizers.

  • Motivation: Traditional analyses require η < 2/λ, but EoS experiments show sharpness exceeding 2/η while loss continues decreasing non-monotonically.Sharpness may stabilize above 2/η, and loss can oscillate while retaining an overall downward trend.
  • Motivation: The paper defines stableness using the maximum Hessian eigenvalue along the entire gradient step, rather than only at its starting point.This captures oscillatory behavior that pointwise sharpness can miss.
  • Two mechanisms: The paper studies modified non-smooth losses f(L), including a square-root transformation, whose Hessian can diverge near minimizers.The setting assumes a smooth base loss with zero minimum and nonzero Hessian at minimizers.
  • Two mechanisms: Normalized GD provides a second mechanism: although L is smooth, its effective learning rate grows as the gradient norm approaches zero.The analysis assumes the zero-loss set contains a manifold and the Hessian has fixed rank on it.
  • Main dynamics: In the EoS phase, GD repeatedly overshoots across the minimizer manifold while moving tangentially toward lower sharpness.The trajectory converges, as η → 0, to a limiting flow on the manifold, and average loss decreases despite non-monotonic iterates.
  • Contribution: The analysis establishes deterministic implicit regularization without SGD noise, although occasional perturbations of magnitude O(η^100) are used in the proofs.Experiments reportedly match the theoretical predictions using deterministic GD without perturbations.

2 Related Works

Related work connects sharpness to flat minima and generalization, studies limiting dynamics near minimizer manifolds, and characterizes algorithm-dependent implicit bias. Prior non-monotone GD analyses provide related two-phase convergence results.

  • Sharpness and generalization: Sharpness has been associated with flat minima and generalization, while SAM explicitly controls a sharpness-related quantity.Parameter rescaling can nevertheless change sharpness without changing the network function.
  • Manifold dynamics: Prior work analyzes stochastic optimization near manifolds of minimizers, including label-noise dynamics and Hessian-related regularization.Later work established a broader stochastic differential-equation framework for these limiting dynamics.
  • Implicit bias: Implicit-bias research studies how training algorithms select among multiple solutions, often through trajectory analyses that extend to infinitesimal-learning-rate limits.The paper contrasts this literature with finite-learning-rate dynamics.
  • Non-monotone convergence: Related non-monotone convergence results use two-phase analyses in which loss oscillation accompanies sharpness reduction before monotone descent resumes.These results provide context for the paper’s EoS analysis.

3 Warm-up: Quadratic Loss Functions

For quadratic loss, Normalized GD evolves through a preparation phase and then aligns with the top Hessian eigenvector, producing persistent loss oscillations characteristic of the Edge of Stability.

  • Setup: Normalized GD on a positive-definite quadratic loss is analyzed through the transformed iterate ˜x(t)=Ax(t), whose dynamics follow a normalized linear update.The quadratic loss has eigenvalues λ1 > λ2 ≥ ... ≥ λD and corresponding eigenvectors v1,...,vD.
  • Main result: Theorem 3.1 states that the transformed iterates converge in direction to v1, with even and odd iterates approaching distinct scalar multiples of that direction.There exist 0 < C < 1 and s ∈ {±1} such that ˜x(2t) → Csλ1v1 and ˜x(2t+1) → (C−1)sλ1v1.
  • Main result: The resulting loss oscillates between consecutive even and odd steps rather than decreasing monotonically, indicating Edge-of-Stability behavior for the quadratic loss.The quadratic analysis also supplies the alignment observation and proof technique used to derive sharpness-reduction implicit bias for general loss functions.
  • Preparation phase: The trajectory first enters invariant sets during a preparation phase before reaching the regime where the iterate norm is at most λ1/2.The norm remains above λ1/2 for at most one time-step once the relevant regime is reached.
  • Alignment phase: When ∥˜x(t)∥≤λ1/2, the magnitude of the projection onto the top eigenvector v1 increases monotonically across subsequent eligible steps.This projection is bounded, so its sustained increase supports convergence toward the top eigendirection.
  • Alignment phase: If the top-eigenvector projection never vanishes, the angle to the top eigenspace becomes arbitrarily small and ultimately converges to zero for all iterates.The projection increases by a factor of Θ(θ_t^2) over two steps in the low-norm regime.

4 Main Results

Under smoothness, manifold, rank, and eigengap assumptions, the paper analyzes two-phase dynamics for Normalized GD and GD on √L near a zero-loss manifold. In Phase II, both methods follow a limiting tangent flow that decreases sharpness while operating at the Edge of Stability.

  • Assumptions: The analysis assumes a C4 loss, a (D−M)-dimensional C2 manifold Γ of local zero-loss minimizers, maximal normal Hessian rank, and a positive top-eigenvalue gap.These conditions ensure the gradient-flow projection is sufficiently regular and the sharpness function is differentiable along the manifold.
  • Phase II: In Phase II, the iterate leaves gradient-flow behavior, repeatedly overshoots across Γ, and moves tangentially in a direction that decreases sharpness.The limiting dynamics are formulated through the gradient-flow projection onto Γ and its tangent-space projection.
  • Normalized GD: The Phase-II trajectory of Normalized GD tracks Riemannian gradient flow for log(λ1(∇^2L)) on Γ, with the angle to the top Hessian eigenspace O(η) on average.The theorem applies to perturbed Normalized GD under initialization, residual, and top-direction alignment conditions.
  • Two-phase dynamics: For sufficiently small η, Phase I tracks gradient flow, decreases loss monotonically, approaches Γ within O(η), and lasts Θ(η^-1) steps.The corresponding Phase-I results are stated for both Normalized GD and GD on √L.
  • GD on √L: The analogous Phase-II theorem applies to perturbed GD on √L, with the projected iterates approximating the associated limiting solution while the dynamics implicitly minimize sharpness.The result holds with probability at least 1−O(η^10) while the limiting solution exists.
  • Operating on the Edge of Stability: During the Edge-of-Stability phase, stableness reaches at least 2 in at least one of every two steps, loss oscillates over two steps, and average loss decreases.Theorem 4.7 gives a two-step inverse-stableness relation for Normalized GD; an analogous result is identified for GD on √L.

5 Proof Overview

The proof analyzes Normalized GD through its projection onto the manifold of minimizers. It shows an initial approach to the manifold followed by slow manifold motion associated with the top Hessian eigendirection and sharpness reduction.

  • Proof framework: The limiting map Φ projects iterates toward the manifold and captures Normalized GD’s implicit regularization up to O(η) error.For small η, the iterate remains O(η)-close to the manifold, while Φ(xη(t)) describes its essential slow dynamics.
  • Phase I: Normalized GD first enters an O(η)-neighborhood of the minimizer manifold and an invariant set where Rj(xη(t)) ≤ O(η2).The preparation phase reaches this set after an initial approach governed by gradient-flow and PL-condition arguments.
  • Phase II: During Phase II, the iterate and gradient approximately align with the top eigenvector of the Hessian at the projected manifold point.This extends the quadratic-loss alignment mechanism to locally approximately quadratic neighborhoods of the manifold.
  • Limiting flow: On the η^-2 timescale, the projected trajectory Φ(xη(τ/η2)) approaches a limiting flow X(τ) on the manifold.The movement of Φ is only O(η2) per step, enabling convergence to the continuous flow.
  • Proof assumptions: The proof requires an alignment condition because approximate quadratic dynamics can make normalized GD chaotic and highly sensitive when alignment becomes too small.Adding poly(η) perturbations, including noise as small as Ω(η100), can prevent this case by maintaining Ω(η) alignment.
  • GD with transformed loss: For GD with the transformed loss, the analysis assumes zero loss on the minimizer manifold and derives a corresponding O(η2)-scale manifold movement.The update rule differs from Normalized GD when iterates lie O(η)-close to the manifold.

6 Experiments

Experiments test the predicted EoS behavior in modified VGG-16 training on CIFAR-10 and compare Normalized GD with its limiting flow on MNIST. The observed dynamics support both the theoretical phase behavior and convergence to the predicted flow.

  • Experimental scope: The experiments are intended to verify theoretical phenomena in practical models, although the authors note uncertainty about whether sufficiently-small-learning-rate predictions transfer directly to practical learning rates.The VGG-16 study modifies the architecture and loss to meet theoretical assumptions.
  • Real-life models: VGG-16 experiments on CIFAR-10 use GeLU, average pooling, and ℓ2 loss to satisfy the theory’s smoothness and minimizer requirements.Training uses a randomly selected subset of 5000 examples for efficient full-batch experiments.
  • Real-life models: The observed alignment approaches 1, the top Hessian eigenvalue decreases, and stableness approaches 2 near the end of training.These trends match the predicted second-phase behavior.
  • Limiting-flow verification: The parameter trajectories of Normalized GD and the Riemannian flow remain close throughout MNIST training.Figure 6 measures both absolute and relative parameter differences at matched continuous times.
  • Limiting-flow verification: On MNIST, Normalized GD and Riemannian gradient flow show an exact match across test loss, test accuracy, top Hessian eigenvalue, and Hessian trace under their respective time scalings.The comparison uses a one-hidden-layer network initialized from a checkpoint with loss of order 10^-3.

7 Conclusion

The paper identifies deterministic EoS dynamics that can reduce loss after stableness is lost and, in one setting, reduce sharpness. Its analysis remains local to minimizer manifolds and leaves earlier experimental onset and related settings for future work.

  • The paper gives two general loss-function settings where gradient updates decrease loss over many iterations after stableness is lost.
  • One setting shows that the trajectory reduces the maximum Hessian eigenvalue, formalizing sharpness reduction during EoS.
  • The analysis applies only close to the manifold of local minimizers, whereas experiments show EoS and sharpness control beginning sooner.
  • Extending the analysis to the original loss L and understanding deep-learning loss properties are left for future work.
  • EoS-like effects in stochastic gradient descent remain an open direction pending definitive analogous experiments.

A Omitted Proofs for Results for Quadratic Loss Functions

For quadratic loss, the appendix formulates Normalized GD in diagonal eigen-coordinates and proves preparation, invariance, and convergence behavior under a nonzero top-eigenvector component.

  • The quadratic loss uses a positive definite diagonal matrix A with ordered eigenvalues and coordinate-subspace projections.
  • Normalized GD updates x by subtracting ηAx normalized by ||Ax||, with a transformed iterate used to analyze the dynamics.
  • A nonzero projection onto the top eigenvector is necessary because zero alignment remains zero thereafter.
  • Under this condition, transformed iterates converge on alternating steps to signed multiples of the top eigenvector.
  • The proof establishes invariant sets and finite-time entry into them, supporting the preparation and alignment analysis.

A.2 Proofs for Alignment Phase

The alignment-phase proof tracks transformed Normalized GD iterates through coordinate updates and bounds, showing that the top-eigenvector component grows under the stated conditions.

  • The alignment analysis assumes the transformed iterate has entered the relevant invariant-set intersection and tracks its direction relative to the top eigenvector.
  • The proof concludes by combining these update bounds to establish the alignment behavior.
  • The proof derives coordinate-wise update relations and bounds involving the eigenvalue gap and the angle to the top eigenvector.
  • When two consecutive transformed norms sum to at most λ1, the top-coordinate magnitude cannot decrease, with equality only at the boundary.
  • The resulting inequalities bound the next norm and control angular change through the minimum of the top eigengap and smallest eigenvalue.

A.3 Proof of Main theorems for Quadratic Loss

The quadratic-loss proof separates preparation and alignment, establishes convergence toward the top eigenvector, and supplies the local assumptions and phase-I setup used for the general-loss analysis.

  • Proof structure: The proof divides quadratic Normalized GD into a preparation phase entering invariant sets and an alignment phase increasing the top-eigenvector projection.
  • Alignment phase: The top-eigenvector projection increases monotonically over designated steps, and convergence of that projection implies convergence in direction.
  • Limit behavior: The limiting dynamics include sign flipping of the top coordinate and alternating transformed-iterate behavior.
  • Technical bounds: The appendix develops quadratic inequalities and ratio bounds to control coordinate evolution across eigenvalue regimes.
  • General-loss setup: For general losses, the analysis assumes smoothness, a local-minimizer manifold, positive Hessian rank there, and a positive top-eigenvalue gap.
  • General-loss setup: Phase I shows sufficiently small-step Normalized GD converges toward gradient flow over a finite interval from an initialization in the flow basin.

C.2 Phase II, Limiting Flow

Phase II shows that perturbed Normalized GD remains close to a minimum-loss manifold while its projected trajectory follows a deterministic limiting flow. This flow is obtained by controlling discrete-to-continuous errors over time scales of order η^-2.

  • Phase II theorem: The main theorem states that, under initialization and alignment conditions, perturbed Normalized GD follows the limiting flow for sufficiently small η.The guarantee holds with probability at least 1 − O(η^10) until the limiting solution exists.
  • Discrete-to-continuous comparison: The proof compares the discrete trajectory with a piecewise-linear interpolation and bounds its deviation from the continuous flow using a Lipschitz differential inequality.The interpolation connects iterates sampled at times tη^2, while the projected dynamics provide the vector field.
  • Time horizon: The limiting-flow approximation is established over intervals of length proportional to η^-2, provided the projected iterates remain inside the regular neighborhood.The stopping time is chosen so that consecutive projections stay in Y_ϵ, after which the approximation extends through the full interval.
  • Manifold proximity: The trajectory remains O(η)-close to the minimum-loss manifold when the iterates satisfy the stated residual bounds.The initial residual condition is propagated, yielding ∥xη(t) − Φ(xη(t))∥ ≤ O(η).

D.1 Proof of Lemma C.1

The proof of Lemma C.1 controls Normalized GD by reducing its transverse dynamics to a slowly changing quadratic model. It maintains bounded residuals, groups nearby eigenspaces, and shows persistent alignment with the top eigenspace.

  • Quadratic reduction: The transverse trajectory is treated as a quadratic-model update with O(η^2) errors caused by the changing Hessian.This reduction permits reuse of quadratic-model bounds while accounting for evolving eigenvalues and eigenvectors.
  • Residual control: The iterates remain O(η)-close to the manifold under the propagated residual conditions.A representative bound gives ∥xη(t) − Φ(xη(t))∥ ≤ 1.01η^ζ.
  • Eigenspace grouping: Eigenvalues are partitioned into groups whose members are O(η)-close, while distinct groups are separated by at least 2η.This grouping supports stable tracking of changing eigenspaces.
  • Alignment dynamics: Once the alignment condition holds, the transverse iterate cannot remain above 0.5ηλ1(t) + Ψnormη^2 for more than one step.This bound controls the amplitude of the component normal to the manifold.
  • Top-eigenspace alignment: The average angle between the transverse trajectory and the top eigenspace is at most O(η) over intervals of length Ω(1/η^2).The result is formalized through the average-of-angles lemma under the Phase II initialization assumptions.
  • Time-step grouping: The proof organizes time steps into 1-cycles and 2-cycles to propagate norm and angle bounds across the trajectory.The sets N0, N1, and N2 partition the analyzed interval and support the inductive argument.

E.5 Proof for Operating on Edge of Stability

The EoS proof links the stability measure along Normalized GD to the top Hessian eigenvalue. The connection follows from controlling residuals and the angle to the top eigenspace.

  • Stability estimate: The stability measure satisfies SL(xη(t), ηt) = ηt(λ1(t) + O(η)) along the analyzed trajectory.This uses the residual bounds established for Phase II.
  • EoS connection: The EoS relation is obtained by combining the stability expression with the two-step transverse-norm estimate.The proof’s remaining step checks the corresponding algebraic relation.
  • Top-eigenvalue tracking: The average two-step transverse norm satisfies η^-1(∥x̃η(t)∥ + ∥x̃η(t + 1)∥) = λ1(t) + O(η + θt).Thus, small angle θt makes the stability behavior track the top eigenvalue.

F Some Useful Lemmas About Eigenvalues and Eigenvectors

This section supplies eigenvalue, eigenvector, and perturbation tools used to analyze the second loss setting. It also identifies weaker alignment and approximation rates for GD with the transformed loss compared with Normalized GD.

  • Perturbation tools: Matrix perturbation results provide smooth eigenvalue and eigenvector dependence, eigenvalue bounds, and Davis–Kahan control of eigenspace changes.These tools require symmetric matrices and, for the smooth-eigenvector statement, a simple eigenvalue.
  • Second loss setting: The proof strategy for GD with the transformed loss follows the Normalized GD analysis but changes the transverse scaling and stability conditions.The two analyses share most lemmas, while several key bounds differ.
  • Flow approximation: The accumulated approximation error over an interval of length Ω(η^-2) is O(η^1/2) for the transformed-loss GD analysis.The bound follows from summing per-step errors involving the angle and η^3 terms.
  • Assumptions: The transformed-loss result assumes residual control and a nontrivial initial projection onto the top eigenvector.The stated theorem also holds with probability at least 1 − η^10 under its initialization conditions.
  • Alignment rate: GD with the transformed loss has a weaker average-angle rate of O(√η), compared with the stronger Normalized GD rate.The weaker rate results from replacing the stronger top-eigenvector behavior lemma with a simpler bound.

G.4 Omitted Proof for Operating on Edge of Stability

The omitted proof establishes geometric and Hessian bounds for GD updates near the minimum-loss manifold, completing the theorem’s Edge-of-Stability claims. It combines local manifold properties, Taylor expansions, and prior lemmas.

  • The proof distinguishes the one-dimensional normal case, where consecutive iterates cross the minimum-loss manifold, from higher-dimensional cases requiring geometric bounds.For M = 1, the intersection makes the first claim immediate; for M ≥ 2, the proof bounds the distance between the manifold projection and the next iterate.
  • The top Hessian eigenvalue along the relevant segment grows at least as 1 divided by η(θ_t + θ_{t+1}), establishing the required sharpness lower bound.The argument uses the distance estimate and the divergence of the Hessian of the square-root loss near the manifold.
  • The proof controls successive normal components through the quadratic-case analogue, yielding ||x̃_η(t)|| + ||x̃_η(t + 1)|| = ηλ_1(t) + O(η + θ_t).This estimate supports completion of the first claim after noting θ_{t+1} = O(θ_t).
  • For the second claim, smoothness of Φ and the top-eigenvector identity reduce the required derivative expression to a term involving ∂²(∇L)(Φ(x))[v_1(x), v_1(x)].The proof then invokes Lemmas G.9 and Corollary B.25 to finish the argument.
  • Local minimizer and curvature properties of Φ bound the displacement from x to Φ(x), while prior lemmas control the gradient and higher-order terms.The proof uses ∇L(Φ(x)) = 0, a quadratic lower bound in the displacement, and ||x − Φ(x)|| = O(ζ^1/2).

H.1 Experimental details

The experimental appendix specifies initialization, learning rates, projection procedures, and loss constructions for reproducing the paper’s GD, Normalized GD, and limiting-flow experiments. It also describes the numerical implementation of the Riemannian flow.

  • For a two-dimensional comparison, GD starts at (14.7, 3.) with η = 0.5, while Normalized GD starts at (14.7, −3) with η = 5.
  • The limiting-flow simulation computes the top Hessian eigenvector, projects ∇λ_1 onto the manifold tangent space, then performs inner gradient steps to return to the manifold.Algorithm 4 uses power-iteration eigenvector computation, a least-squares tangent projection, and projected gradient descent.
  • The loss is an average of n nonnegative per-example losses, each zero exactly when its prediction equals the true label.The construction covers regression and binary classification, and extends to multiclass classification through squared prediction errors for class scores.
  • Each Riemannian update has three stages: computing ∇³L(x)[v_1(x), v_1(x)], projecting onto the tangent space, and taking small-step gradient updates toward the manifold.
  • For Figures 5 and 6, the flow uses η = 10^-2, η_proj = 10^-2, T_proj = 10^3, and starts from loss 3.803 × 10^-3; Normalized GD uses learning rate 10^-2.The top Hessian eigenvector is computed with 100 power iterations at each point.
Loading 2205.09745v3…