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

• URL: http://arxiv.org/abs/2602.01480v1
• Date: Sun, 01 Feb 2026 23:10:44 GMT
• Title: Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stability
• Abstract summary: gradient descent with large step sizes often diverges away from the gradient flow.<n>"Central Flow", proposed in Cohen et al., provides an accurate approximation to the GD dynamics dynamics.<n>We propose Rod Flow, which better captures GD dynamics for simple toy examples.
• Score: 9.11122093402205
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• A Lagrangian View of Flow Matching [17.972940829317455]
  We present an alternative, bottom-up mechanical derivation grounded in a Lagrangian (particle-centric) perspective.<n>We demonstrate that solving this Partial Differential Equation analytically yields the straight-line trajectories of Flow Matching.
  arXiv  Detail & Related papers  (2026-08-31T18:13:33Z)
• Entropy Across the Bridge: Conditional-Marginal Discretization for Flow and Schrödinger Samplers [68.9946633641494]
  Flow matching and Schrdinger bridges define probability, yet their inference grids are usually or inherited from one-endpoint diffusion.<n>We derive a conditional-marginal entropy-rate objective for bridge-aware discretization, separating endpoint-conditioned bridge geometry from marginal flow evolution.<n>On EDM/CIFAR-10, the entropic time-discretization gives the best tested five-step FID (186.3 pm 4.0 versus 200.5 pm 2.9 for linear and 238.0 pm 5.3 for cosine)
  arXiv  Detail & Related papers  (2026-05-15T16:11:10Z)
• Gradient Flow Drifting: Generative Modeling via Wasserstein Gradient Flows of KDE-Approximated Divergences [1.7415937218905126]
  We prove an equivalence between the recently proposed Drifting Model and the Wasserstein gradient flow of the forward KL divergence under kernel density estimation (KDE) approximation.<n>We provide a concise identifiability proof, and a theoretically grounded mixed-divergence strategy.
  arXiv  Detail & Related papers  (2026-03-11T09:48:42Z)
• Gradient flow for deep equilibrium single-index models [32.2015869030351]
  Deep equilibrium models (DEQs) have emerged as a powerful paradigm for training infinitely deep weight-tied neural networks.<n>We rigorously study the gradient descent dynamics for DEQs in the simple setting of linear models and single-index models.<n>We then prove linear convergence of gradient descent to a global minimizer for linear DEQs and deep equilibrium single-index models.
  arXiv  Detail & Related papers  (2025-11-21T06:14:41Z)
• Optimization Insights into Deep Diagonal Linear Networks [10.395029724463672]
  We study the implicit regularization properties of the gradient flow "algorithm" for estimating the parameters of a deep diagonal neural network.<n>Our main contribution is showing that this gradient flow induces a mirror flow dynamic on the model, meaning that it is biased towards a specific solution of the problem.
  arXiv  Detail & Related papers  (2024-12-21T20:23:47Z)
• Kernel Approximation of Fisher-Rao Gradient Flows [52.154685604660465]
  We present a rigorous investigation of Fisher-Rao and Wasserstein type gradient flows concerning their gradient structures, flow equations, and their kernel approximations. Specifically, we focus on the Fisher-Rao geometry and its various kernel-based approximations, developing a principled theoretical framework.
  arXiv  Detail & Related papers  (2024-10-27T22:52:08Z)
• On the Wasserstein Convergence and Straightness of Rectified Flow [54.580605276017096]
  Rectified Flow (RF) is a generative model that aims to learn straight flow trajectories from noise to data.<n>We provide a theoretical analysis of the Wasserstein distance between the sampling distribution of RF and the target distribution.<n>We present general conditions guaranteeing uniqueness and straightness of 1-RF, which is in line with previous empirical findings.
  arXiv  Detail & Related papers  (2024-10-19T02:36:11Z)
• Neural Sinkhorn Gradient Flow [11.4522103360875]
  We introduce the Neural Sinkhorn Gradient Flow (NSGF) model, which parametrizes the time-varying velocity field of the Wasserstein gradient flow. Our theoretical analyses show that as the sample size increases to infinity, the mean-field limit of the empirical approximation converges to the true underlying velocity field. To further enhance model efficiency on high-dimensional tasks, a two-phase NSGF++ model is devised.
  arXiv  Detail & Related papers  (2024-01-25T10:44:50Z)
• Particle-based Variational Inference with Generalized Wasserstein Gradient Flow [32.37056212527921]
  We propose a ParVI framework, called generalized Wasserstein gradient descent (GWG) We show that GWG exhibits strong convergence guarantees. We also provide an adaptive version that automatically chooses Wasserstein metric to accelerate convergence.
  arXiv  Detail & Related papers  (2023-10-25T10:05:42Z)
• Neural Gradient Learning and Optimization for Oriented Point Normal Estimation [53.611206368815125]
  We propose a deep learning approach to learn gradient vectors with consistent orientation from 3D point clouds for normal estimation. We learn an angular distance field based on local plane geometry to refine the coarse gradient vectors. Our method efficiently conducts global gradient approximation while achieving better accuracy and ability generalization of local feature description.
  arXiv  Detail & Related papers  (2023-09-17T08:35:11Z)
• Implicit Bias in Leaky ReLU Networks Trained on High-Dimensional Data [63.34506218832164]
  In this work, we investigate the implicit bias of gradient flow and gradient descent in two-layer fully-connected neural networks with ReLU activations. For gradient flow, we leverage recent work on the implicit bias for homogeneous neural networks to show that leakyally, gradient flow produces a neural network with rank at most two. For gradient descent, provided the random variance is small enough, we show that a single step of gradient descent suffices to drastically reduce the rank of the network, and that the rank remains small throughout training.
  arXiv  Detail & Related papers  (2022-10-13T15:09:54Z)
• Continuous vs. Discrete Optimization of Deep Neural Networks [15.508460240818575]
  We show that over deep neural networks with homogeneous activations, gradient flow trajectories enjoy favorable curvature. This finding allows us to translate an analysis of gradient flow over deep linear neural networks into a guarantee that gradient descent efficiently converges to global minimum. We hypothesize that the theory of gradient flows will be central to unraveling mysteries behind deep learning.
  arXiv  Detail & Related papers  (2021-07-14T10:59:57Z)
• Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks [78.76880041670904]
  In neural networks with binary activations and or binary weights the training by gradient descent is complicated. We propose a new method for this estimation problem combining sampling and analytic approximation steps. We experimentally show higher accuracy in gradient estimation and demonstrate a more stable and better performing training in deep convolutional models.
  arXiv  Detail & Related papers  (2020-06-04T21:51:21Z)
• A Near-Optimal Gradient Flow for Learning Neural Energy-Based Models [93.24030378630175]
  We propose a novel numerical scheme to optimize the gradient flows for learning energy-based models (EBMs) We derive a second-order Wasserstein gradient flow of the global relative entropy from Fokker-Planck equation. Compared with existing schemes, Wasserstein gradient flow is a smoother and near-optimal numerical scheme to approximate real data densities.
  arXiv  Detail & Related papers  (2019-10-31T02:26:20Z)

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.