Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stability

Eric Regis, Sinho Chewi

TL;DR

Rod Flow introduces a principled continuous-time model for gradient descent at the edge of stability by treating GD iterates as a moving rod with center $\bar{w}$ and extent $\Sigma$. Derived from exact difference equations and backward error analysis, the model yields explicit ODEs that correctly predict the $2/\eta$ sharpness threshold and explain self-stabilization in quartic potentials. It is shown to match, and in some cases outperform, prior Central Flow in capturing GD trajectories on toy problems and representative neural networks, while remaining cheaper to compute thanks to a low-rank representation and Hessian-vector products. Empirically, Rod Flow tracks gradient descent more closely than gradient flow and provides a simpler, more interpretable alternative to Central Flow with broad applicability to non-convex optimization dynamics.

Abstract

How can we understand gradient-based training over non-convex landscapes? The edge of stability phenomenon, introduced in Cohen et al. (2021), indicates that the answer is not so simple: namely, gradient descent (GD) with large step sizes often diverges away from the gradient flow. In this regime, the "Central Flow", recently proposed in Cohen et al. (2025), provides an accurate ODE approximation to the GD dynamics over many architectures. In this work, we propose Rod Flow, an alternative ODE approximation, which carries the following advantages: (1) it rests on a principled derivation stemming from a physical picture of GD iterates as an extended one-dimensional object -- a "rod"; (2) it better captures GD dynamics for simple toy examples and matches the accuracy of Central Flow for representative neural network architectures, and (3) is explicit and cheap to compute. Theoretically, we prove that Rod Flow correctly predicts the critical sharpness threshold and explains self-stabilization in quartic potentials. We validate our theory with a range of numerical experiments.

Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stability

TL;DR

Rod Flow introduces a principled continuous-time model for gradient descent at the edge of stability by treating GD iterates as a moving rod with center and extent . Derived from exact difference equations and backward error analysis, the model yields explicit ODEs that correctly predict the sharpness threshold and explain self-stabilization in quartic potentials. It is shown to match, and in some cases outperform, prior Central Flow in capturing GD trajectories on toy problems and representative neural networks, while remaining cheaper to compute thanks to a low-rank representation and Hessian-vector products. Empirically, Rod Flow tracks gradient descent more closely than gradient flow and provides a simpler, more interpretable alternative to Central Flow with broad applicability to non-convex optimization dynamics.

Abstract

How can we understand gradient-based training over non-convex landscapes? The edge of stability phenomenon, introduced in Cohen et al. (2021), indicates that the answer is not so simple: namely, gradient descent (GD) with large step sizes often diverges away from the gradient flow. In this regime, the "Central Flow", recently proposed in Cohen et al. (2025), provides an accurate ODE approximation to the GD dynamics over many architectures. In this work, we propose Rod Flow, an alternative ODE approximation, which carries the following advantages: (1) it rests on a principled derivation stemming from a physical picture of GD iterates as an extended one-dimensional object -- a "rod"; (2) it better captures GD dynamics for simple toy examples and matches the accuracy of Central Flow for representative neural network architectures, and (3) is explicit and cheap to compute. Theoretically, we prove that Rod Flow correctly predicts the critical sharpness threshold and explains self-stabilization in quartic potentials. We validate our theory with a range of numerical experiments.
Paper Structure (66 sections, 1 theorem, 84 equations, 12 figures, 5 tables)

This paper contains 66 sections, 1 theorem, 84 equations, 12 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Contributions
  3. Related Work
  4. Gradient Descent Dynamics
  5. Rod Flow
  6. Setup and Notation
  7. The Rod Flow Equations
  8. Exact Difference Equations
  9. Difference Equation for the Center
  10. Difference Equation for the Extent
  11. Why Track $\delta \otimes \delta$ Instead of $\delta$?
  12. Properties of the Difference Equations
  13. Backward Error Analysis
  14. A Motivating Example
  15. General Backward Error Analysis
  16. ...and 51 more sections

Key Result

Lemma 3.1

If $L$ has $S$-Lipschitz gradients ($\|\nabla^2 L\| \leq S$ everywhere), then for $\eta < 2/S$:

Figures (12)

  • Figure 1: Illustration of Rod Flow.Left: A rod moving down a flat loss landscape, where the two endpoints represent consecutive GD iterates and the midpoint represents their average. Center Left: Comparison of gradient descent (top row) and Rod Flow (bottom row) on a quadratic loss landscape under convergent (left column) and divergent (right column) step sizes. Right: Comparison of gradient descent, gradient flow, and Rod Flow on a 3-layer MLP. The left plot shows sharpness over time: gradient descent and rod flow remain at the edge of stability while gradient flow continues to sharpen. The right plot shows the distance between the GD iterates and the two flows, demonstrating that Rod Flow stays closer to the GD trajectory than gradient flow.
  • Figure 2: Edge of Stability.Left: Edge of stability exhibited for three step sizes. For all step sizes, we observe progressive sharpening followed by stabilization of the sharpness at $2/\eta$. Right: The three phases of edge of stability. Phase I (steady descent): progressive sharpening with decreasing gradient norms. Phase II (entering edge of stability): gradient norm spikes and fluctuates wildly. Phase III (steady-state edge of stability): sharpness remains at $2/\eta$ while gradient norm becomes smooth.
  • Figure 3: Backward Error Analysis. The ODE that passes through discrete points differs from the ODE whose discretization generates them. Left: Exponential growth. Middle: Exponential decay. Right: GD iterates alternate in sign and no first-order ODE passes through them.
  • Figure 4: Comparison on 2D square root loss$L(x,y) = \sqrt{1 + (xy)^2}$. This loss has minima along the axes, with sharpness increasing away from the origin. When the local minimum is too sharp, GD bounces until the sharpness decreases to $2/\eta$. Gradient flow converges to the local minimum irrespective of the sharpness. Rod Flow tracks GD precisely. Central Flow follows the sharpness manifold.
  • Figure 5: Comparison of First-Order Rod Flow and Central Flow.
  • ...and 7 more figures

Theorems & Definitions (1)

  • Lemma 3.1: Descent Lemma