Edge Flow: A Tractable and Predictive Continuous-Time Model for Gradient Descent at the Edge of Stability

• URL: http://arxiv.org/abs/2606.18080v1
• Date: Tue, 16 Jun 2026 15:45:47 GMT
• Title: Edge Flow: A Tractable and Predictive Continuous-Time Model for Gradient Descent at the Edge of Stability
• Abstract summary: We propose a tractable, faithful, and predictive model of gradient descent dynamics at the edge of stability (EoS)<n>We show that Edge Flow tracks the dynamics of gradient descent at least as faithfully as previously proposed continuous-time EoS models.
• Score: 11.598128347617957
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Gradient descent in deep learning may operate at the edge of stability (EoS), a regime in which the largest eigenvalue of the loss Hessian hovers near the stability threshold $2/η$, where $η$ is the learning rate. Classical analysis tools such as gradient flow and the descent lemma do not apply here, motivating the search for a continuous-time model valid at EoS. We propose Edge Flow, a system of three coupled ordinary differential equations that provides a tractable, faithful, and predictive model of gradient descent dynamics at EoS. Edge Flow decomposes the dynamics into a center, an oscillation direction, and an oscillation magnitude. The center follows a modified gradient flow on a symmetrized loss; the direction tracks a top eigenvector of the Hessian via Rayleigh quotient dynamics; and the magnitude grows or decays exponentially depending on whether the sharpness exceeds or falls below the threshold $2/η$. Crucially, sharpness stabilization emerges from the coupled dynamics via a self-stabilization feedback loop. Discretizing Edge Flow only requires two gradient evaluations and one Hessian--vector product at each iteration. We demonstrate empirically that Edge Flow tracks the dynamics of gradient descent at least as faithfully as previously proposed continuous-time EoS models, while in addition resolving the oscillation of the sharpness at the onset of EoS, and that it provides a principled framework for understanding and mitigating instabilities in this regime.

Related papers

• The Multiple Timescales of Gradient Descent on the Edge of Stability: A Perturbative Derivation of the Central Flow [8.426855646402238]
  We derive a perturbative regime in which the central flow is the limit of gradient descent.<n>Three timescales emerge: a fast timescale of oscillations along the sharpest direction, an intermediate timescale of the self-stabilization mechanism, and a slow timescale of the dynamics along the minimizers of $g$-the central flow.
  arXiv  Detail & Related papers  (2026-09-01T10:32:48Z)
• The Road Taken: The Role of Optimizers at the Edge of Stability [53.81268610971847]
  edge of stability refers to a phenomenon in deep learning where the Hessian eigenvalues of the loss remain stable above a threshold that the classical descent predicts to be unstable.<n>Previous works formulate the edge of stability with respect to the maximum Hessian eigenvalue and the learning rate.<n>We observe that many first-order methods, including gradient descent, significantly violate the stability bound predicted by these theories by a factor as large as $times 21.1$.
  arXiv  Detail & Related papers  (2026-08-19T01:03:34Z)
• The Map Behind the Flow: Finite-Step Gradient Descent as a Dynamical System [42.117840303613164]
  Edge-of-stability behavior, sharpness oscillations, catapult phases, balancing, and movement toward flatter representations are effects of the training map itself.<n>This paper studies fixed-step gradient descent as a discrete dynamical system in a hierarchy of exactly solvable models.
  arXiv  Detail & Related papers  (2026-07-06T12:28:24Z)
• A Bifurcation Theory Framework for Gradient Descent on the Edge of Stability [3.4773470589069473]
  We develop a bifurcation theory framework for gradient descent on the edge of stability.<n>We show that stable EoS training arises from a flip bifurcation in the normal direction.
  arXiv  Detail & Related papers  (2026-06-14T02:39:54Z)
• Gradient descent at the Edge of Stability: free energy model and kinetic description of the two-layer network [0.0]
  We study the dynamics of gradient descent in the Edge of Stability regime.<n>We propose a continuous-time effective model that tracks the evolution of the average trajectory.<n>We show that this equation can be interpreted as a Wasserstein-2 gradient flow of a macroscopic free energy.
  arXiv  Detail & Related papers  (2026-06-03T18:11:33Z)
• A Rod Flow Model for Adam at the Edge of Stability [9.11122093402205]
  We extend rod flow to Adam by working in the joint phase space of parameters and first moment $(w, m).<n>We also develop rod flows for heavy ball momentum.<n>For all eight$s, we empirically evaluate rod flow on representative machine learning architectures.
  arXiv  Detail & Related papers  (2026-05-07T18:21:59Z)
• Stochastic Gradient Descent in the Saddle-to-Saddle Regime of Deep Linear Networks [74.46751089984072]
  gradient descent (SGD) noise on descent in deep linear networks is poorly understood.<n>We investigate the dynamics of SGD during training of DLNs in the saddle-to-saddle regime.<n>Our results establish that SGD noise encodes information about the progression of feature learning but does not fundamentally alter the saddle-to-saddle dynamics.
  arXiv  Detail & Related papers  (2026-04-07T18:43:08Z)
• Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality [54.20763128054692]
  We study the dynamics of gradient flow for training a multi-head softmax attention model for in-context learning of multi-task linear regression. We prove that an interesting "task allocation" phenomenon emerges during the gradient flow dynamics.
  arXiv  Detail & Related papers  (2024-02-29T18:43:52Z)
• On the Convergence of Gradient Descent for Large Learning Rates [55.33626480243135]
  We show that convergence is impossible when a fixed step size is used.<n>We provide a proof of this in the case of linear neural networks with a squared loss.<n>We also prove the impossibility of convergence for more general losses without requiring strong assumptions such as Lipschitz continuity for the gradient.
  arXiv  Detail & Related papers  (2024-02-20T16:01:42Z)
• Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stability [69.01076284478151]
  In machine learning optimization, gradient descent (GD) often operates at the edge of stability (EoS) This paper studies the convergence and implicit bias of constant-stepsize GD for logistic regression on linearly separable data in the EoS regime.
  arXiv  Detail & Related papers  (2023-05-19T16:24:47Z)
• Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability [40.17821914923602]
  We show that gradient descent at the edge of stability implicitly follows projected gradient descent (PGD) under the constraint $S(theta) le 2/eta$. Our analysis provides precise predictions for the loss, sharpness, and deviation from the PGD trajectory throughout training.
  arXiv  Detail & Related papers  (2022-09-30T17:15:12Z)
• Beyond the Edge of Stability via Two-step Gradient Updates [49.03389279816152]
  Gradient Descent (GD) is a powerful workhorse of modern machine learning. GD's ability to find local minimisers is only guaranteed for losses with Lipschitz gradients. This work focuses on simple, yet representative, learning problems via analysis of two-step gradient updates.
  arXiv  Detail & Related papers  (2022-06-08T21:32:50Z)
• Fine-Grained Analysis of Stability and Generalization for Stochastic Gradient Descent [55.85456985750134]
  We introduce a new stability measure called on-average model stability, for which we develop novel bounds controlled by the risks of SGD iterates. This yields generalization bounds depending on the behavior of the best model, and leads to the first-ever-known fast bounds in the low-noise setting. To our best knowledge, this gives the firstever-known stability and generalization for SGD with even non-differentiable loss functions.
  arXiv  Detail & Related papers  (2020-06-15T06:30:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.