A Rod Flow Model for Adam at the Edge of Stability

• URL: http://arxiv.org/abs/2605.06821v1
• Date: Thu, 07 May 2026 18:21:59 GMT
• Title: A Rod Flow Model for Adam at the Edge of Stability
• Abstract summary: 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.
• Score: 9.11122093402205
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Cohen et al. (arXiv:2207.14484) observed that adaptive gradient methods such as Adam operate at the edge of stability. While there has been significant work on continuous-time modeling of gradient descent at the edge of stability, extending these models to momentum methods remains underdeveloped. In the gradient descent setting, Regis et al. (arXiv:2602.01480) introduced rod flow, which models consecutive iterates as an extended one-dimensional object -- a "rod." Here we extend rod flow to Adam by working in the joint phase space of parameters and first moment $(w, m)$ and treating the second moment $ν$ as a smooth auxiliary variable. We also develop rod flows for heavy ball momentum, Nesterov momentum, and scalar and per-component versions of RMSProp, Adam, and NAdam. For all eight optimizers, we empirically evaluate rod flow on representative machine learning architectures, where it tracks the discrete iterates through the edge-of-stability regime significantly more accurately than the corresponding stable flow.

Related papers

• Provable Edge-of-Stability for Adam on a One-Dimensional Quadratic [4.5845031477639155]
  We study uncorrected Adam on a one-dimensional quadratic, a clean setting where constant curvature isolates the edge-induced dynamics behind the EoS.<n>We identify settings in which this edge-seeking mechanism breaks down, including strictly subcritical periodic orbits and specially tuned trajectories that converge to the optimum.<n>These results give a concrete dynamical explanation for Adam's EoS in a setting free of evolving loss geometry, while also exposing its limitations.
  arXiv  Detail & Related papers  (2026-08-21T00:22:11Z)
• 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)
• Edge Flow: A Tractable and Predictive Continuous-Time Model for Gradient Descent at the Edge of Stability [11.598128347617957]
  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.
  arXiv  Detail & Related papers  (2026-06-16T15:45:47Z)
• Taming Curvature: Architecture Warm-Up for Stable Transformer Training [52.6838951211596]
  We introduce a fast online estimator of the largest (preconditioned) Hessian eigenvalue (i.e., curvature) based on a warm-started variant for power with Hessian-vector products.<n>We find that training instabilities coincide with surges in preconditioned curvature and that curvature grows with depth.<n> Experiments on large Transformers validate that our approach enables efficient curvature tracking and reduces instabilities.
  arXiv  Detail & Related papers  (2026-06-15T14:16:56Z)
• Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training [56.323119575322146]
  Extra-Merge consistently outperforms standard merging baselines in experiments across GPT-2 and LLaMA families.<n>It yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon citepjordan2024muon.
  arXiv  Detail & Related papers  (2026-05-26T02:48:34Z)
• Stabilizing Streaming Video Geometry via Dynamic Feature Normalization [45.91399406203578]
  Dynamic Feature Normalization (DyFN) is a causal recurrent module that modulates feature statistics to maintain stable geometry over time.<n>DyFN achieves state-of-the-art temporal stability, improving over prior streaming methods by up to 14% and even outperforming heavier non-causal video baselines.
  arXiv  Detail & Related papers  (2026-05-25T00:13:15Z)
• 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)
• Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stability [9.11122093402205]
  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.
  arXiv  Detail & Related papers  (2026-02-01T23:10:44Z)
• Flow marching for a generative PDE foundation model [0.0]
  We propose Flow Marching, an algorithm that bridges neural operator learning with flow matching motivated by an analysis of error accumulation in physical dynamical systems.<n>We also introduce a Physics-Pretrained Variational Autoencoder (P2E) to embed physical trajectories into a compact latent space.<n>We curate a corpus of 2.5M trajectories across 12 distinct PDE families and train suites of P2Es and FMTs at multiple scales.
  arXiv  Detail & Related papers  (2025-09-23T04:00:41Z)
• A Langevin sampling algorithm inspired by the Adam optimizer [0.0]
  We present a framework for adaptive-stepsize MCMC sampling based on time-rescaled Langevin dynamics.<n>Our algorithm is straightforward to implement and can be readily combined with any off-the-peg fixed-stepsize Langevin integrator.
  arXiv  Detail & Related papers  (2025-04-26T12:57:57Z)
• Generative Modeling with Phase Stochastic Bridges [49.4474628881673]
  Diffusion models (DMs) represent state-of-the-art generative models for continuous inputs. We introduce a novel generative modeling framework grounded in textbfphase space dynamics Our framework demonstrates the capability to generate realistic data points at an early stage of dynamics propagation.
  arXiv  Detail & Related papers  (2023-10-11T18:38:28Z)
• Data-driven low-dimensional dynamic model of Kolmogorov flow [0.0]
  Reduced order models (ROMs) that capture flow dynamics are of interest for decreasing computational costs for simulation. This work presents a data-driven framework for minimal-dimensional models that effectively capture the dynamics and properties of the flow. We apply this to Kolmogorov flow in a regime consisting of chaotic and intermittent behavior.
  arXiv  Detail & Related papers  (2022-10-29T23:05:39Z)
• Deep Equilibrium Optical Flow Estimation [80.80992684796566]
  Recent state-of-the-art (SOTA) optical flow models use finite-step recurrent update operations to emulate traditional algorithms. These RNNs impose large computation and memory overheads, and are not directly trained to model such stable estimation. We propose deep equilibrium (DEQ) flow estimators, an approach that directly solves for the flow as the infinite-level fixed point of an implicit layer.
  arXiv  Detail & Related papers  (2022-04-18T17:53:44Z)

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.