Technical Summary: Edge Flow

Problem Statement

Gradient Descent (GD) in deep learning frequently operates in a regime known as the Edge of Stability (EoS). In this regime, the largest eigenvalue of the loss Hessian (sharpness, $S(w_t)$) hovers near the stability threshold $2/\eta$, where $\eta$ is the learning rate. While the training loss generally decreases, the iterates exhibit quasi-periodic oscillations, and the sharpness stabilizes through a self-regulating mechanism rather than diverging.

Classical continuous-time analysis tools, such as gradient flow ($\dot{w} = -\nabla L(w)$) and the descent lemma, fail in this regime because they are oblivious to the learning rate and cannot capture phenomena dependent on the ratio $\eta \cdot S(w)$. Existing continuous-time models for EoS have limitations:

• Central Flow (Cohen et al., 2025): Models the center of oscillations via a projected gradient flow. While it stabilizes, the projection constraint makes the system difficult to simulate and analyze theoretically.
• Rod Flow (Regis and Chewi, 2026): Models iterates as endpoints of a rigid rod. While conceptually appealing, its update rule for oscillation magnitude fails to capture the self-stabilization dynamics at the onset of EoS.

The central open question addressed by this paper is: What is a tractable, faithful, and predictive continuous-time model of GD dynamics at the edge of stability?

Methodology: Edge Flow

The authors propose Edge Flow, a system of three coupled ordinary differential equations (ODEs) that decomposes GD dynamics into three slowly varying quantities:

1. The Center ($\bar{w}_t$): The slowly drifting midpoint of the oscillations.
2. The Oscillation Direction ($u_t$): A unit vector tracking the top eigenvector of the Hessian.
3. The Oscillation Magnitude ($x_t$): The half-length of the oscillation.

The GD iterate is modeled as $w_t \approx \bar{w}_t + (-1)^t x_t u_t$. The Edge Flow system is defined as:

1. Center Dynamics:
   $$\frac{d\bar{w}_t}{dt} = -\frac{1}{2} \left( \nabla L(\bar{w}_t + x_t u_t) + \nabla L(\bar{w}_t - x_t u_t) \right)$$
   This represents a gradient flow on a symmetrized loss (the average of gradients at the two endpoints). It differs from standard gradient flow by a third-derivative correction term quadratic in $x_t$, which biases the center toward regions of lower curvature.

2. Direction Dynamics:
   $$\frac{du_t}{dt} = 2 \left( \nabla^2 L(\bar{w}_t) - (u_t^\top \nabla^2 L(\bar{w}_t) u_t) I \right) u_t$$
   This is the Rayleigh quotient gradient flow on the unit sphere. It drives $u_t$ toward the top eigenvector of the Hessian at $\bar{w}_t$, analogous to the continuous-time limit of the power method.

3. Magnitude Dynamics:
   $$\frac{dx_t}{dt} = \left( S(\bar{w}_t) - \frac{2}{\eta} \right) x_t$$
   This scalar linear ODE governs the growth or decay of the oscillation magnitude. If sharpness $S(\bar{w}_t) > 2/\eta$, $x_t$ grows exponentially; if $S(\bar{w}_t) < 2/\eta$, $x_t$ decays.

Discretization: Edge Gradient Descent (EGD)

The authors propose a natural discretization called Edge Gradient Descent (EGD). Each iteration requires:

• Two gradient evaluations: $\nabla L(\bar{w}_k \pm x_k u_k)$.
• One Hessian-vector product: $\nabla^2 L(\bar{w}_k) u_k$.
• No eigenvalue decomposition or optimization subroutines.

The algorithm uses two distinct stepsizes: $\eta$ (the original learning rate, defining the sharpness threshold) and $\rho$ (the integration step for the ODEs, where $\rho \le \eta$).

Key Contributions

1. A Simple, Tractable Model: Edge Flow is derived from first principles and elementary observations of GD structure at EoS. It is a system of three coupled ODEs that is simple to simulate and analyze.
2. Emergent Self-Stabilization: Unlike previous models that impose constraints to enforce stability, Edge Flow reproduces the self-stabilization mechanism (identified by Damian et al., 2023) as a natural consequence of the coupled dynamics. When sharpness exceeds the threshold, the oscillation magnitude grows, pushing the center toward lower curvature regions via the symmetrized gradient, which in turn reduces sharpness.
3. Faithful Tracking and Prediction: Empirically, Edge Flow tracks full-batch GD trajectories, loss, and sharpness across various architectures (MLP, CNN, ResNet, ViT) and losses. Crucially, it resolves the overshoot and damped oscillation of sharpness at the onset of EoS, a feature missed by Central Flow's instantaneous projection.
4. Insight into Instabilities: The model identifies the source of loss spikes at the onset of EoS as a discretization artifact of the stiff center dynamics. It suggests two principled remedies:
   • Refining discretization: Using a smaller integration step $\rho$ (larger $K = \eta/\rho$) suppresses loss spikes.
   • Increasing base oscillation: Injecting a small initial magnitude $\epsilon$ accelerates the self-stabilization feedback loop, reducing overshoot.

Results

• Tracking Performance: In experiments on CIFAR-10 subsets, Edge Flow matches the bulk dynamics of GD and Central Flow. However, Edge Flow uniquely reproduces the transient oscillatory behavior of sharpness and the specific "ringing" observed in GD at the onset of EoS (Figure 4).
• Loss Prediction: The model accurately predicts the training loss trajectory, including the characteristic spikes.
• Stabilization Mechanism: The coupled dynamics successfully demonstrate how the system naturally settles near the $2/\eta$ threshold without explicit projection constraints.
• Discretization Effects: Experiments show that increasing the discretization resolution ($K$) eliminates loss spikes without altering the self-stabilization mechanism, confirming the spikes are numerical artifacts rather than intrinsic to the EoS phenomenon.

Significance and Claims

The paper claims that Edge Flow provides a principled framework for understanding and mitigating instabilities in the large-stepsize regime. By decomposing the dynamics into center, direction, and magnitude, it offers a tractable alternative to complex projected flows or rigid rod models.

The authors emphasize that Edge Flow is faithful (tracking individual trajectories including transients), predictive (suggesting interventions like noise injection or adaptive stepsizes), and tractable (cheap to simulate). They posit that this model opens realistic pathways for:

• Proving implicit regularization results for large learning rates.
• Understanding the interplay between loss curvature, architecture, and generalization.
• Designing more efficient discretization schemes tailored to the stiff nature of EoS dynamics.

The paper remains modest regarding its scope, noting that the derivation neglects higher-order coupled corrections and that rigorous error analysis and extension to stochastic settings (SGD) and adaptive methods are open directions for future work.