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

2026-06-16

AI Summary

This work addresses the inadequacy of conventional continuous models, such as gradient flow, in capturing the dynamics of gradient descent operating at the Edge of Stability (EoS) in deep learning. The authors propose Edge Flow, a continuous-time model comprising three coupled ordinary differential equations that decompose EoS dynamics into a central trajectory, an oscillation direction, and an amplitude component. Edge Flow naturally exhibits sharpness stabilization through an intrinsic self-stabilizing feedback mechanism and analytically describes the initial oscillatory behavior of sharpness. The model integrates a modified gradient flow, Rayleigh quotient dynamics to track the leading Hessian eigenvector, and an exponential model for amplitude evolution. Its discretization requires only two gradient evaluations and one Hessian-vector product per step. Experiments demonstrate that Edge Flow substantially outperforms existing continuous models in faithfully tracking optimization trajectories, capturing sharpness oscillations, and offering principled guidance for managing instability.

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.

Problem

• Edge of Stability
• Gradient Descent
• Continuous-Time Model
• Hessian Eigenvalues
• Sharpness

Innovation

• Edge Flow
• Edge of Stability
• continuous-time model
• sharpness stabilization
• Hessian dynamics