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

1.1 Gradient Descent at the Edge of Stability

• The descent lemma guarantees loss decrease for learning rates below 2/ℓ, making 2/ℓ the classical stability cutoff.For quadratic losses, learning rates above 2/ℓ lead to exponentially increasing loss.Source: §1.1 · p. 2
• Cohen et al. [8] observed progressive sharpening to 2/η followed by sharpness hovering near that threshold while loss continues decreasing non-monotonically, despite the descent lemma no longer guaranteeing decrease.The introduction notes similar observations for SGD by Jastrz˛ebski et al..Source: §1.1 · p. 2

1.2 Self-stabilization: The Implicit Bias of Instability

• The authors present self-stabilization as a general property of gradient descent, explaining how sharpness remains bounded around 2/η despite progressive sharpening and instability.They provide predictions for loss, sharpness, and deviation from the constrained trajectory, prove those predictions track gradient descent under mild conditions up to higher-order error terms, and verify them empirically.Source: §1.2 · pp. 2–4
• A cubic Taylor term links motion along the top Hessian eigenvector to decreasing sharpness, restoring stability as sharpness falls below 2/η and that motion shrinks.When the top eigenvalue is unique, ∇S(θ) = ∇3L(θ)(u(θ), u(θ)); this term supplies negative feedback that limits sharpness and movement along u.Source: §1.2 · p. 3
• Self-stabilization makes gradient descent implicitly track projected gradient descent constrained to S(θ) ≤ 2/η, with ∇L(θ) · u(θ) = 0 defining the stable set.The orthogonality condition is included to ensure the constrained trajectory is not unstable; it does not affect the stationary points.Source: §1.2 · pp. 3–4

2 Related Work

• Prior observations report nonmonotone full-batch loss [32], inability to converge above 2/η, catapulting from initially excessive sharpness [20], and EOS oscillation around 2/η after sharpening [8].Cohen et al. [8] observed that the loss continues decreasing non-monotonically in the EOS regime.Source: §2 · p. 4
• Along SGD, Jastrzębski et al. found progressive sharpening followed by decline and top-eigenvector relationships, while conjectured sharpness growth bounded by a stability criterion.For full-batch training, their criterion reduces to S(θ) ≤ 2/η.Source: §2 · p. 4
• Prior theory addresses EOS under growth, invariant-set, near-minimizer, normalization, and toy or network convergence settings, [2],, [24],.Concurrent work Li et al. [23] proves progressive sharpening for a two-layer network and analyzes EOS in four stages, using output-layer norm as a sharpness proxy.Source: §2 · p. 4
• Large-step or small-batch SGD has been linked to decreasing sharpness [18], while Gilmer et al. [13] describe connections between EOS, learning-rate warm-up, and gradient clipping.Source: §2 · p. 5
• Earlier implicit-regularization analyses use oscillations arising from stochasticity, whereas this proof uses oscillations caused by instability and their longer-term third-order Taylor effect [5] [10] [22].Source: §2 · p. 5

• The loss is assumed to be C3, and the setup studies its gradient-descent iterates.Source: §3 · p. 5
• Time is shifted so θ0 is the first iterate whose projection onto the stable-point set M has sharpness 2/η.proj_M denotes the orthogonal projection onto M.Source: §3 · p. 5
• The key assumption is progressive sharpening along the constrained trajectory, represented by the coefficient α(θ).Source: §3 · p. 5
• The analysis focuses on a single unstable eigenvalue, assumes an eigengap, and leaves multiple unstable eigenvalues to future work.Source: §3 · p. 5

4 The Self-stabilization Property of Gradient Descent

• The derived equations predict sharpness, loss, and displacement from the constrained trajectory, with negative feedback preventing divergence in sharpness and the unstable direction.The section uses these equations to describe gradient descent as a dynamical system.Source: §4 · p. 5

4.1 The Four Stages of Edge of Stability: A Heuristic Derivation

• A cubic Taylor term first drives progressive sharpening, then sends iterates growing along the top eigenvector when sharpness exceeds 2/η; its effect lowers sharpness below 2/η, after which the unstable-direction movement shrinks.These are the four stages: progressive sharpening, blowup, self-stabilization, and return to stability.Source: §4.1 · pp. 6–7

4.2 Analyzing the simplified dynamics

• The simplified dynamics form a negative-feedback system with fixed point (δ, 0), where δ = 2α/β, and bound sharpness fluctuations by ~O(√α) and unstable-direction and total displacement by ~O(δ).Fluctuations increase with progressive-sharpening coefficient α and decrease with self-stabilization force β.Source: §4.2 · p. 8

4.3 Relationship with the constrained trajectory θ†t

• In directions orthogonal to the unstable and sharpness-changing directions, gradient descent follows a projected gradient update; its distance from the constrained trajectory is bounded by ~O(δ).The projected update removes the ∇S direction to keep sharpness unchanged, up to higher-order terms.Source: §4.3 · p. 9

5.1 Notation

• The notation distinguishes progressive sharpening α_t, stabilization strength β_t, and deviation size δ_t; β_s→t approximates sharpness change after a propagated displacement.The displacement changes at each step under the Hessian, which is why β_s→t tracks its accumulated effect.Source: §5.1 · pp. 9–10

5.2 The equations governing edge of stability

• The loss and sharpness are determined by the x and y directions alone, yielding a two-dimensional dynamical system with time-dependent coefficients.When the coefficients β_s→t are constant, the update reduces to the simpler case analyzed in Section 4.Source: §5.2 · p. 10

5.3 Coupling Theorem

• Under mild assumptions, the true dynamics are accurately governed by the predicted dynamics; the assumptions include well-behaved Hessian and third-derivative quantities orthogonal to the top eigenvector.The paper verifies the relevant orthogonal-direction assumptions empirically; its global bounds on third derivatives are needed only near the constrained trajectory in the proof.Source: §5.3 · pp. 10–11
• Theorem 1 characterizes gradient descent as following the constrained PGD trajectory plus a rapidly oscillating process, with predicted dynamics capturing deviations and predicting loss and sharpness.The theorem’s predictions hold up to higher-order terms; the paper describes sharpness as controlled by a slowly evolving quantity combined with rapid periodic behavior.Source: §5.3 · p. 11

6 Experiments

• The evaluation replicates [8] on a 3-layer MLP and CNN trained on CIFAR10 and a 2-layer Transformer trained on SST2.The MLP and CNN use MSE on a 5k CIFAR10 subset; the Transformer uses MSE on SST2, Socher et al..Source: §6 · p. 12
• Gradient descent and the constrained trajectories remain close throughout training, while gradient flow diverges from gradient descent at a linear rate.The dynamics converge to (|x_t|, y_t) = (δ_t, 0), with ∥θ_t − θ†_t∥ → δ_t, supporting the claim that gradient descent implicitly follows the constrained trajectory.Source: §6 · p. 12
• Appendix E numerically verifies the model assumptions used to derive edge-of-stability behavior.Additional experiments vary the learning rate and loss function using the generalized predicted dynamics described in Section 7.2.Source: §6 · p. 12

7.1 Takeaways from the Predicted Dynamics

• Loss and sharpness are governed by a two-dimensional dynamical system that cycles with a conserved potential under constant coefficients and is generally well behaved.After a few cycles, the system's state approaches a fixed point, whose location the analysis predicts accurately.Source: §7.1 · p. 12

7.2 Generalized Predicted Dynamics

• When a fourth-derivative bound fails, notably in some logistic-loss experiments, generalized predictions track gradient descent accurately across a wider range of settings.They remain qualitatively similar to the original predictions; empirical verification is reported in Appendix G.Source: §7.2 · p. 14

7.3 Implications for Neural Network Training

• Larger learning rates can accelerate convergence because the constrained trajectory can move quickly in directions orthogonal to the two restricted directions.The authors propose that gradually increasing the learning rate lets self-stabilization reduce sharpness without divergence, whereas an excessively large initial rate may cause divergence. This offers a plausible explanation for warmup [13].Source: §7.3 · pp. 14–15
• Gradient descent implicitly keeps sharpness near 2/η, and the authors report better generalization with large learning rates in full-batch training.The analysis suggests larger learning rates may improve generalization by reducing sharpness.Source: §7.3 · p. 15
• For fixed η, gradient descent implicitly minimizes sharpness while seeking to minimize loss subject to S(θ) = 2/η, unlike SAM's explicit sharpness minimization [12].For scale-invariant losses, the sharpness-reduction mechanism of gradient descent with weight decay described by Lyu et al. [24] is self-stabilization in that setting.Source: §7.3 · pp. 14–15

8.1 Towards A Complete Theory of Optimization

• Theorem 1 gives sufficient, verifiable conditions for coupling GD and constrained PGD for O(ε^-1) steps, but they are not necessary and do not establish global convergence.Future work asks what prerequisites on the loss and learning rate enable self-stabilization, and whether longer or repeated coupling can prove convergence to KKT points.Source: §8.1 · p. 15

8.2 Multiple Unstable Eigenvalues

• Cohen et al. [8] observed progressive sharpening across eigenvalues; dynamics appear more chaotic with multiple unstable eigenvalues, whose nondifferentiable sharpness complicates extending the analysis.The authors leave multiple-eigenvalue dynamics for future work; their discussion suggests coupling to PGD under a nondifferentiable sharpness constraint and analyzing the relevant subspace.Source: §8.2 · pp. 15–16, Figure 5 · p. 16

8.3 The Mystery of Progressive Sharpening

• The analysis assumes progressive sharpening, leaving its cause in neural networks open; concurrent work by Li et al. [23], Zhu et al., and Agarwala et al. offers partial progress.Source: §8.3 · p. 16

8.4 Connections to Stochasticity

• Understanding SGD requires identifying its appropriate stability notion and self-stabilization force; a proposed connection is implicit regularization that acts like a log-barrier enforcing S(θ) ≤ 2/η.The authors call for clarifying how this mechanism interacts with SGD regularization described by Blanc et al. [5], Damian et al. [10], and Li et al. [22]; related works discuss Hessian-trace effects and label-noise SGD.Source: §8.4 · pp. 16–17

A Notation

• The notation appendix defines kth-order loss derivatives, tensor contractions and orthogonal projections, and specifies the matrix-product convention A_k := A_t … A_1.Source: A Notation · p. 22

B A Toy Model for Self-Stabilization

• In the toy model, the process follows the stated dynamics in the x and y directions while tracking the constrained trajectory.The constrained trajectory follows x†_t = −ηt because it cannot decrease y without raising sharpness above 2/η.Source: B A Toy Model for Self-Stabilization · pp. 22–23

C Definition of the Predicted Dynamics

• The appendix presents the full predicted-dynamics definition and introduces a timestep map and formulas for unfolding its recursion.Source: C Definition of the Predicted Dynamics · p. 23

D.1 Architectures; D.2 Data

• The experiments use four architectures—MLP, CNN, ResNet18, and a two-layer Transformer—and three tasks: CIFAR10 multiclass and binary classification, and SST2.The listed settings include CIFAR10 multiclass with categorical MSE or cross-entropy, and CIFAR10 binary and SST2 with binary MSE or logistic loss.Source: §D.1 · p. 24, §D.2 · p. 24

D.3 Experimental Setup

• Experiments track gradient descent through instability, then compare constrained trajectory, gradient descent, gradient flow, and both predicted dynamics while monitoring assumptions and switching to 64-bit gradients.Training stops when the second eigenvalue along the constrained trajectory reaches 1.9/η, preserving a single unstable eigenvalue; 64-bit precision addresses sensitivity when |x| is small.Source: §D.3 · p. 24
• The constrained trajectory uses a linearized projection approximation and a Newton step enforcing ∇L·u = 0; each projection step uses three linearized steps and three eigenvalue computations.Eigenvalues are computed with JAX’s LOBPCG sparse solver. Experiments use JAX, replicate the setup in [8], and code is available online.Source: §D.3 · p. 24

E Empirical Verification of the Assumptions

• Across six model-and-loss settings, plots along the constrained trajectory verify the stated assumptions, with training stopped when the second eigenvalue reaches 1.9/η.The plots also track ε and quantities associated with Assumptions 1, 3, 4, and 5; the stopping rule ensures Assumption 2 holds with c = 1.9.Source: E Empirical Verification of the… · p. 25

CNN+MSE on CIFAR10

• The listed empirical settings include CNN+Logistic on CIFAR10 cats versus dogs, ResNet18+MSE on CIFAR10 cats versus dogs, and Transformer+Logistic on SST2.These are the model-and-dataset captions provided for this subsection.Source: CNN+MSE on CIFAR10 · pp. 28–30

F The Generalized Predicted Dynamics

• When the fourth-derivative bound needed for the cubic Taylor approximation is often violated, generalized predicted dynamics capture edge-of-stability behavior for tested MSE and cross-entropy settings.The top-eigenvector loss may be sub-quadratic, where the quadratic approximation overestimates loss and sharpness, or super-quadratic, where it underestimates them.Source: F The Generalized Predicted Dynamics · p. 31

F.1 Deriving the Generalized Predicted Dynamics

• In the generalized update, F_t describes departures from quadratic behavior and movement orthogonal to the sharpness gradient changes sharpness by y_t; when F_t(x) = 0, the update reduces to the standard dynamics.The update for y is unchanged.Source: §F.1 · p. 31

F.2 Properties of the Generalized Predicted Dynamics

• The analysis assumes F_t is even, since its odd part affects the dynamics only through additional oscillations of period 2.The text says F can otherwise be redefined as its even part.Source: §F.2 · p. 32
• The fixed point remains x_t = δ_t regardless of F_t's shape, unlike earlier one-dimensional predictions in which it depends strongly on the loss along the top eigenvector.This fixed point is required to stabilize the ∇S⊥ direction.Source: §F.2 · p. 32
• For a period-2 orbit of the generalized update, the limiting y value can be read from the update for x_t.Source: §F.2 · p. 32
• Sharpness no longer follows the approximation S(θ_t) ≈ 2/η + y_t because it also changes along u_t according to F′′.When F_t = 0 and the loss is exactly quadratic in the u direction, the update reduces to the fixed-point predictions in Section 4.1.Source: §F.2 · p. 32
• Cohen et al. [8] observed cross-entropy sharpness hovering above 2/η, and the generalized predicted dynamics provide an explanation for this departure from standard predictions.Standard predicted dynamics instead predict a fixed point with sharpness 0.Source: §F.2 · p. 32
• For a super-quadratic loss, sharpness converges to a value above 2/η. More generally, whether the loss converges to a value above or below 2/η depends on the sign of F ′′.Source: §F.2 · p. 33
• Experiments found loss usually sub-quadratic when instability first occurs, then gradually super-quadratic as training continues at the edge of stability.Source: §F.2 · p. 33

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

• The large-learning-rate experiments also compare multiple rates with the schedule η_t := 1/S(θ_t).The supplied passage does not state a separate outcome for that schedule.Source: §G.1 · p. 33
• On ResNet18/CIFAR10, larger learning rates converge faster and generalize better, with final sharpness bounded by 2/η.The experiments use full-batch gradient descent on the 50k-example training set and compare multiple learning rates with η_t := 1/S(θ_t).Source: Figure 6 · p. 33, §G.1 · p. 33

CNN+MSE on CIFAR10

• The comparison panels cover five model-task settings and show gradient descent, predicted dynamics, constrained trajectories, gradient flow, and predicted fixed points.The listed settings include CNNs and ResNet18 on CIFAR10, plus Transformers on SST2, with MSE or logistic losses.Source: CNN+MSE on CIFAR10 · pp. 35–38

H Scale Invariant Losses

• The paper’s primary result is that gradient descent solves the constrained problem of minimizing L(θ) subject to S(θ) ≤ 2/η.The scale-invariance argument invokes a lemma from [24] relating the sharpness of f across parameter rescalings.Source: H Scale Invariant Losses · p. 39
• When ηλ is sufficiently small and initialized near a global minimizer, the dynamics converge to minimizing Sf(θ) over unit-norm θ subject to f(θ) = 0.The stated constrained objective is min_{||θ||=1} Sf(θ) subject to f(θ) = 0.Source: H Scale Invariant Losses · p. 39

H.1 Scale Invariant Lemmas

• The scale-invariant lemmas state that the top eigenvector of ∇²f is perpendicular to θ and relate Hessian sharpness across parameter normalization.These lemmas are used in the proof of Lemma 6.Source: H Scale Invariant Losses · p. 39, §H.1 · p. 40

I.1 Properties of the Constrained Trajectory

• Bounded third derivatives yield local Lipschitz bounds for Hessian eigenvalues, the top eigenvector, and sharpness near the constrained trajectory.The proof invokes the third-derivative bound and the derivative-of-eigenvector formula.Source: §I.1 · pp. 42–43
• The constrained trajectory stays at sharpness 2/η: the proof uses induction and KKT conditions to show the sharpness constraint is active after projection.If the multiplier were zero, the next sharpness would exceed 2/η, contradicting feasibility; complementary slackness then gives equality.Source: §I.1 · pp. 43–44

I.2 Proof of Theorem 1

• The proof of Theorem 1 invokes three auxiliary lemmas, then uses Taylor expansions to calculate sharpness and loss.The final loss-expansion step follows from Assumption 5.Source: §I.2 · p. 46

I.3 Proof of Auxiliary Lemmas

• The auxiliary-lemma proof derives the one-step update in the top-eigenvector direction and bounds its error terms.It first subtracts two equations and expands the third-derivative term using non-worst-case bounds.Source: §I.3 · pp. 47–48
• The proof also derives a two-step update for x_t by bounding its error terms and the change in the sharpness gradient’s top-eigenvector component.The displayed bound decomposes that change into contributions from the sharpness-gradient difference and the eigenvector difference.Source: §I.3 · pp. 49–50
• For Lemma 16, an auxiliary trajectory is constructed to match the original trajectory at every even time up to T.The proof establishes the even-time agreement by induction.Source: §I.3 · pp. 51–52