PaperlayerView exact revision on arXiv

Source-linked AI summary

Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability

Alex Damian, Eshaan Nichani, Jason D. Lee

arXiv:2209.15594v2cs.LGcs.ITmath.OCstat.ML

TL;DR

The paper addresses why neural-network gradient descent can remain near the instability threshold even when classical descent guarantees fail. Using cubic Taylor dynamics, it identifies self-stabilization as the mechanism that keeps sharpness near 2/η and makes gradient descent follow projected gradient descent, with theoretical and empirical support. The analysis is limited by its focus on a single unstable eigenvalue and its assumption of progressive sharpening.

  • Problem

    Classical descent analysis does not explain neural-network edge-of-stability dynamics, where sharpness reaches 2/η while loss continues decreasing non-monotonically.

  • Method

    The paper uses a cubic Taylor expansion and analyzes instability-driven oscillations to model self-stabilization and coupling to projected gradient descent.

  • Results

    The analysis predicts loss, sharpness, and deviation from the constrained trajectory and matches gradient descent dynamics theoretically and empirically.

  • Takeaways & Limitations

    Self-stabilization is a general gradient-descent property that keeps sharpness bounded near 2/η and prevents divergence despite instability.

  • Takeaways & Limitations

    The analysis assumes progressive sharpening and focuses on a single unstable eigenvalue, leaving its causes and multiple-eigenvalue dynamics open.

Abstract

from arXiv · show

Traditional analyses of gradient descent show that when the largest eigenvalue of the Hessian, also known as the sharpness $S(θ)$, is bounded by $2/η$, training is "stable" and the training loss decreases monotonically. Recent works, however, have observed that this assumption does not hold when training modern neural networks with full batch or large batch gradient descent. Most recently, Cohen et al. (2021) observed two important phenomena. The first, dubbed progressive sharpening, is that the sharpness steadily increases throughout training until it reaches the instability cutoff $2/η$. The second, dubbed edge of stability, is that the sharpness hovers at $2/η$ for the remainder of training while the loss continues decreasing, albeit non-monotonically. We demonstrate that, far from being chaotic, the dynamics of gradient descent at the edge of stability can be captured by a cubic Taylor expansion: as the iterates diverge in direction of the top eigenvector of the Hessian due to instability, the cubic term in the local Taylor expansion of the loss function causes the curvature to decrease until stability is restored. This property, which we call self-stabilization, is a general property of gradient descent and explains its behavior at the edge of stability. A key consequence of self-stabilization is that gradient descent at the edge of stability implicitly follows projected gradient descent (PGD) under the constraint $S(θ) \le 2/η$. Our analysis provides precise predictions for the loss, sharpness, and deviation from the PGD trajectory throughout training, which we verify both empirically in a number of standard settings and theoretically under mild conditions. Our analysis uncovers the mechanism for gradient descent's implicit bias towards stability.

1 Introduction

Classical stability criteria fail to describe neural-network gradient descent at the edge of stability. The paper explains this behavior through self-stabilization, showing that cubic dynamics keep sharpness near 2/η and make gradient descent track a constrained trajectory.

  • Classical stability: The descent lemma guarantees loss decrease when η < 2/ℓ, motivating stability thresholds based on the largest Hessian eigenvalue.For quadratic losses, exceeding the threshold leads to exponentially increasing loss.
  • Observed dynamics: Neural-network gradient descent exhibits progressive sharpening until sharpness reaches 2/η, after which it remains near that threshold while loss decreases non-monotonically.These observations motivate the paper’s focus on edge-of-stability dynamics.
  • Self-stabilization: Self-stabilization keeps sharpness bounded near 2/η and prevents gradient descent from diverging despite instability.The paper presents self-stabilization as a general property of gradient descent, unlike progressive sharpening, which is tied to specific loss functions.
  • Mechanism: A cubic Taylor expansion captures how motion along the top eigenvector produces a sharpness-decreasing term through ∇3L(θ)(u, u) = ∇S(θ).As instability grows, this negative feedback lowers sharpness below 2/η and returns the dynamics to the quadratic regime.

2 Related Work

Related work documents edge-of-stability behavior and develops several theoretical explanations, while also connecting sharpness dynamics to stochastic optimization and other training practices. The paper’s analysis emphasizes instability-driven oscillations and their third-order effects.

  • Edge-of-stability observations: Prior studies report non-monotone loss, catapulting into stable regions, and difficulty converging to minima whose sharpness exceeds 2/η.These observations establish the empirical and theoretical context for analyzing edge-of-stability dynamics.
  • Theoretical analyses: Existing theories analyze edge of stability under assumptions including subquadratic growth, forward-invariant subsets, normalized gradient descent, and proximity to minimizers.The cited approaches cover several specialized settings rather than the paper’s general self-stabilization framing.
  • Sharpness and training practices: Large step sizes or small batches can decrease sharpness, while other work connects edge of stability with learning-rate warm-up and gradient clipping.These studies broaden the context beyond full-batch gradient descent.
  • Analytical perspective: The paper analyzes oscillations caused by instability and their longer-term third-order effects, contrasting them with stochasticity-driven oscillations in SGD analyses.The shared analytical idea is implicit regularization emerging from oscillatory dynamics.

3 Setup

The setup assumes a thrice-continuously differentiable loss, progressive sharpening along a constrained trajectory, and a regime with one unstable eigenvalue separated by an eigengap. Multiple unstable eigenvalues are left for future work.

  • Loss and iterates: The loss function is assumed to satisfy L ∈ C3(R^d), enabling the third-order analysis used for self-stabilization.The setup studies gradient descent iterates in R^d.
  • Constrained trajectory: The constrained trajectory is initialized by shifting time to the first point whose projected sharpness reaches 2/η.The stable set and orthogonal projection define the constrained reference trajectory.
  • Progressive sharpening: The analysis assumes progressive sharpening along the constrained trajectory through a progressive-sharpening coefficient α(θ).This assumption supplies the sharpening force modeled in the dynamics.
  • Spectral assumptions: The theoretical regime contains a single unstable eigenvalue and imposes an eigengap condition on the Hessian.The treatment of multiple unstable eigenvalues is explicitly deferred.

4 The Self-stabilization Property of Gradient Descent

The paper models edge-of-stability dynamics through a cubic Taylor expansion and a four-stage negative-feedback process. This process bounds oscillations and keeps gradient descent close to a constrained projected trajectory.

  • 4.1 The Four Stages of Edge of Stability: The analysis tracks displacement from the constrained trajectory in the unstable eigenvector and sharpness-changing directions.The coordinates are x_t := u · (θ_t − θ⋆) and y_t := ∇S · (θ_t − θ⋆), with y_t approximately measuring sharpness change.
  • 4.1 The Four Stages of Edge of Stability: Edge of stability proceeds through progressive sharpening, blowup, self-stabilization, and return to stability.When sharpness exceeds 2/η, the unstable-direction magnitude grows; the cubic term then decreases sharpness until stability returns.
  • 4.1 The Four Stages of Edge of Stability: The cubic Taylor term is proportional to ∇S x_t^2/2 and biases updates in the −∇S direction, decreasing sharpness.As the unstable displacement grows, this negative feedback strengthens until sharpness falls below 2/η.
  • 4.2 Analyzing the simplified dynamics: The simplified dynamics form a negative-feedback system with fixed point (X, Y) = (δ, 0), where δ = 2α/β.The associated ODE can be written as a Lotka–Volterra predator-prey model with a conserved quantity.
  • 4.2 Analyzing the simplified dynamics: Sharpness fluctuations are Õ(√α), unstable-direction fluctuations are Õ(δ), and the full process remains bounded by Õ(δ).The fluctuations increase with progressive sharpening α and decrease with self-stabilization strength β.
  • 4.2 Analyzing the simplified dynamics: In directions orthogonal to u and ∇S, gradient descent follows a projected gradient update and remains within Õ(δ) of the constrained trajectory.The self-stabilizing dynamics control the deviation from the constrained trajectory up to higher-order terms.

5 The Predicted Dynamics and Theoretical Results

The paper derives predicted edge-of-stability dynamics by Taylor expanding around a constrained trajectory. Under mild assumptions, these dynamics predict loss, sharpness, and trajectory deviations up to higher-order terms.

  • 5.1 General dynamics: The general analysis Taylor expands each iterate’s gradient around the corresponding point on the constrained trajectory.The predicted displacement dynamics depend on the remaining directions through the coefficients β_s→t.
  • 5.1 General dynamics: The coefficient β_s→t measures the sharpness change caused by a displaced gradient-sharpness direction after Hessian-mediated evolution.When the relevant quantities are constant and orthogonal, β_s→t reduces to a constant β.
  • 5.2 The equations governing edge of stability: The edge-of-stability dynamics can be fully captured by a two-dimensional system in the x and y directions with time-dependent coefficients.These directions determine the loss and sharpness values.
  • 5.3 Coupling Theorem: The predicted displacement follows the constrained projected-gradient trajectory plus a rapidly oscillating process.The predicted dynamics capture the loss, sharpness, and deviation from the constrained trajectory.
  • 5.3 Coupling Theorem: Theorem 1 states that the predicted dynamics determine loss, sharpness, and trajectory deviation up to higher-order terms.The sharpness is controlled by a slowly evolving quantity while the dynamics combine gradual and rapid periodic behavior.

6 Experiments

The experiments test the predicted edge-of-stability dynamics across multilayer perceptrons, convolutional networks, and Transformers. The predictions track loss, sharpness, and proximity to the constrained trajectory.

  • Experimental setup: The experiments replicate prior edge-of-stability studies on a 3-layer MLP, a 3-layer CNN, and a 2-layer Transformer.The models use MSE on a 5k subset of CIFAR10 or SST2.
  • Results: Figure 4 confirms that the predicted dynamics accurately predict loss, sharpness, and distance from the constrained trajectory.Additional experiments vary learning rate and loss function in Appendix G.
  • Results: Gradient flow diverges linearly from gradient descent, whereas gradient descent and the constrained trajectories remain close throughout training.The dynamics converge to the fixed point (|x_t|, y_t) = (δ_t, 0), with ∥θ_t − θ†_t∥ → δ_t.
  • Assumption checks: The experiments also numerically verify the assumptions used to establish the theoretical edge-of-stability results.The authors track the relevant quantities throughout training and stop before a second unstable Hessian direction emerges.

7 Discussion

The predicted two-dimensional dynamics accurately track edge-of-stability training and explain its convergence, loss behavior, and sharpness control. These results connect gradient descent’s implicit stability bias to generalization, training speed, and related sharpness-control methods.

  • 7.1 Takeaways from the Predicted Dynamics: The predicted dynamics accurately track loss, sharpness, and deviation from the constrained trajectory, while gradient descent remains close to the constrained trajectory throughout training.The dynamics converge to the fixed point (|x_t|, y_t) = (δ_t, 0).
  • 7.1 Takeaways from the Predicted Dynamics: The edge-of-stability behavior reduces to a generally well-behaved two-dimensional dynamical system governing the quantities that determine loss and sharpness.With constant coefficients, the system cycles and has a conserved potential.
  • 7.1 Takeaways from the Predicted Dynamics: After a few cycles, the dynamics approach the system’s fixed point, whose location and loss increase relative to the constrained trajectory can be predicted.The fixed point is described as (±δ_t, 0).
  • 7.1 Takeaways from the Predicted Dynamics: The cubic Taylor analysis requires a bound on the fourth derivative, but generalized predicted dynamics address experiments where this condition, especially under logistic loss, is violated.The generalized formulation is deferred to Appendix F.
  • 7.3 Implications for Neural Network Training: Although loss fluctuates non-monotonically over short periods, the constrained trajectory’s descent eventually dominates bounded fluctuations, producing long-term loss decrease.This behavior is reflected in the experiments in Section 6.
  • 7.3 Implications for Neural Network Training: Gradient descent implicitly constrains sharpness near 2/η, and for fixed η seeks to minimize loss subject to S(θ) = 2/η.This provides a connection to sharpness-aware minimization, which explicitly targets sharpness.

8 Future Work

The paper identifies several boundaries for its current theory and outlines extensions toward a more complete account of optimization dynamics. These include longer coupling guarantees, multiple unstable eigenvalues, progressive sharpening, and stochastic gradient descent.

  • 8.1 Towards A Complete Theory of Optimization: Theorem 1 couples gradient descent and projected gradient descent for O(ϵ^-1) steps, but this local result does not establish global convergence to a stationary point.The authors leave longer coupling periods and convergence to KKT points as open questions.
  • 8.2 Multiple Unstable Eigenvalues: The analysis assumes a single unstable eigenvalue, whereas multiple eigenvalues can successively become unstable and produce more chaotic dynamics.Figure 5 marks when each eigenvalue becomes unstable; the dynamics are described as relatively well behaved with one unstable eigenvalue but significantly more chaotic with several.
  • 8.2 Multiple Unstable Eigenvalues: Extending the analysis to multiple unstable eigenvalues is technically difficult because sharpness is nondifferentiable when the top eigenvalue is not unique.The authors suggest that projected-gradient-descent coupling may still hold under a nondifferentiable sharpness constraint, but leave the analysis for future work.
  • 8.3 The Mystery of Progressive Sharpening: The analysis assumes progressive sharpening rather than explaining why neural networks exhibit it.The authors identify this mechanism as an important direction for future work and cite concurrent partial progress.
  • 8.4 Stochastic Gradient Descent: The paper studies full-batch gradient descent, leaving connections between self-stabilization and analogous edge-of-stability phenomena in SGD unresolved.The authors note that SGD phenomena lack as simple a characterization and discuss possible implicit-regularization links involving the Hessian trace and sharpness.

D.3 Experimental Setup

The experiments track gradient descent, constrained trajectories, and predicted dynamics across several neural architectures, tasks, losses, and learning rates. They also test the assumptions underlying the theory and identify settings where the cubic approximation is insufficient.

  • Experimental tracking: The experiments track constrained trajectories, gradient descent, gradient flow, and standard and generalized predicted dynamics after instability.They also monitor quantities tied to the theoretical assumptions and the second Hessian eigenvalue along the constrained trajectory.
  • Implementation: The constrained trajectory is computed with linearized projections and a Newton step enforcing ∇L·u = 0, requiring three eigenvalue computations per projection step.Sharpness and the top eigenvector are recomputed for each linearized approximation.
  • Approximation limits: The cubic Taylor approximation requires a bound on the fourth derivative, and this condition is often violated for ResNet18 and logistic-loss models.The generalized predicted dynamics are introduced to capture sub-quadratic or super-quadratic behavior in such cases.

F.2 Properties of the Generalized Predicted Dynamics

The generalized predicted dynamics account for non-quadratic behavior in the top-eigenvector direction while preserving the mechanism that fixes the displacement from the constrained trajectory. They predict that limiting sharpness can lie above or below 2/η depending on local curvature.

  • Fixed-point structure: The generalized dynamics retain the fixed point x_t = δ_t regardless of the top-direction loss shape, because stabilization of the ∇S⊥ direction remains necessary.This differs from earlier one-dimensional analyses in which the fixed point depended strongly on the top-direction loss shape.
  • Limiting sharpness: The limiting value of y is inferred from the x update when the fixed point forms a period-2 orbit.The sharpness must additionally account for changes along the top-eigenvector direction through F′′(x).
  • Cross-entropy behavior: The generalized dynamics explain why cross-entropy experiments can hover above 2/η, unlike the standard dynamics' prediction of a zero-sharpness fixed point.The generalized formulation is introduced specifically to account for this discrepancy.
  • Limiting sharpness: Sub-quadratic loss yields limiting sharpness below 2/η, whereas super-quadratic loss yields limiting sharpness above 2/η.More generally, the limiting side of 2/η depends on the sign of F′′.
  • Experimental behavior: Experiments observed losses that were usually sub-quadratic at instability but gradually became super-quadratic during training at the edge of stability.This behavior helps explain why the generalized dynamics can describe sharpness values on either side of 2/η.

G.1 The Benefit of Large Learning Rates: Training Time and Generalization

Large learning rates accelerate convergence and improve generalization in full-batch ResNet18 training, while self-stabilization keeps final sharpness bounded by 2/η. The analysis connects this behavior to an implicit sharpness-constrained optimization problem.

  • Empirical benefit: Large learning rates converge faster and generalize better for ResNet18 trained with full-batch gradient descent on CIFAR10.The experiment compares various fixed learning rates with the schedule η_t = 1/S(θ_t).
  • Empirical benefit: Self-stabilization bounds the final sharpness by 2/η, which is smaller for larger learning rates and is associated with better generalization.The paper reports this relationship in the ResNet18 and CIFAR10 setting.
  • Implicit constraint: The analysis identifies gradient descent with the constrained problem min_θ L(θ) such that S(θ) ≤ 2/η.For scale-invariant losses, the paper further relates this to minimizing S_f(θ) over unit-norm parameters subject to f(θ) = 0.
  • Implicit constraint: Under sufficiently small ηλ and initialization near a global minimizer, the dynamics converge to the scale-invariant constrained solution.The stated constrained form is min ||θ||=1 S_f(θ) such that f(θ) = 0.

I Proofs

The proofs establish regularity properties around the constrained trajectory and use Taylor expansions, auxiliary lemmas, and induction to control the dynamics. They then derive the predicted updates and convergence result under the paper’s assumptions.

  • Auxiliary bounds: The proof framework establishes Lipschitz control for the Hessian eigenvalues, top eigenvector, and sharpness gradient near the constrained trajectory.These bounds follow from assumptions on higher derivatives and eigenvector perturbation formulas.
  • Constrained trajectory: Induction and KKT conditions show that the constrained trajectory remains on the boundary S(θ†_t) = 2/η.The argument uses complementary slackness after projecting each update onto the constraint set.
  • Convergence result: Corollary 4 guarantees that an iterate reaches the stated loss threshold by some time t ≤ T.The corollary follows by repeatedly applying the preceding lemma under the theorem’s assumptions.
  • Dynamics comparison: Taylor expansions of sharpness and loss produce the two-step updates used to compare gradient descent with the constrained dynamics.The proof separately bounds the resulting error terms and tracks components along and orthogonal to the top eigenvector.
  • Dynamics comparison: The proof introduces auxiliary trajectories and induction arguments to control the top-eigenvector component and associated error terms over time.These constructions support the one-step and two-step update estimates.
Loading 2209.15594v2…