A Rod Flow Model for Adam at the Edge of Stability

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.

• The paper introduces an ODE-based rod flow model that accurately captures the oscillatory dynamics of Adam at the edge of stability.
• It generalizes the methodology to include momentum and adaptive preconditioning by modeling dynamics in a joint phase space.
• Empirical evaluations on CIFAR-10 architectures validate the model’s prediction of stability thresholds and oscillation scales, aligning with theoretical analysis.

A Rod Flow Model for Adam at the Edge of Stability

Context and Motivation

The dynamics of deep neural network optimization have attracted intensive scrutiny, especially in regimes where large learning rates place gradient-based optimizers at the so-called "edge of stability" (EoS). Prior work has revealed that optimizers such as SGD, Adam, and RMSProp exhibit nontrivial oscillatory behaviors when training at high learning rates or with aggressive hyperparameters; these are not captured by classical continuous-time (ODE) limits, which ignore the finite-step-size and fail to model the characteristic oscillations about the stability threshold. While the continuous-time theory for plain gradient descent at EoS is increasingly mature, an analogous framework for Adam and its relatives—momentum and adaptive methods—has until recently been absent. Addressing this gap, this work generalizes the rod flow model to momentum and adaptive preconditioned optimizers, with an emphasis on Adam, the de facto optimizer of modern deep learning.

Rod Flow Formalism and Its Generalization

The rod flow framework, introduced previously for gradient descent, substitutes the pointwise trajectory of iterates with a trajectory for the "center" (midpoint) of oscillations and an associated "extent tensor," encapsulating both the amplitude and orientation of the typical period-2 oscillatory pattern at EoS. This low-rank description allows for a smooth ODE approximation to inherently nonsmooth, oscillatory discrete-time dynamics, successfully interpolating the trajectory of the centers, even in high-sharpness regimes unsuited to naive ODE limits.

Generalizing to momentum and adaptive preconditioned methods requires two crucial extensions:

1. Phase Space Rods: For momentum methods (heavy-ball, Nesterov, Adam, NAdam), the oscillatory dynamics synchronize not just the parameters $w$, but also the momentum $m$. The rod must thus live in the joint phase space $(w, m)$; the center and half-difference are computed on this lifted space.
2. Smooth Preconditioners: For adaptive methods (RMSProp, Adam, etc.), the second moment estimate $v$ (the preconditioner) has fundamentally different dynamics, remaining smooth, as squaring the gradients annihilates sign changes. The rod framework only tracks the smooth center of $v$.

This unifies the treatment of a broad class of optimizers (SGD, heavy-ball, Nesterov, RMSProp, Adam, NAdam, and their variants), endowing each with a rod flow ODE description that accurately models their iterates at the EoS regime.

Theoretical Framework

For each algorithm, the authors derive explicit discrete-time difference equations for the phase-space centers and extent tensor, and obtain the associated ODEs by formal continuous-time promotion. This process is exact for the rod variables (center and extent) except for the continuous interpolation step; no additional approximations are introduced. The authors then detail the derivations for heavy-ball and Nesterov momentum, establishing sharpness thresholds for stability as a function of learning rate and momentum:

• Heavy-ball: For step size $\eta$ and momentum $\beta$, the sharpness threshold is $\lambda^* = \frac{2}{\eta}\frac{1+\beta}{1-\beta}$.
• Nesterov: The threshold is further reduced due to the look-ahead; $\lambda^* = \frac{2}{\eta (1-\beta)(1+2\beta)}$.

For RMSProp and Adam, the analysis demonstrates that the relevant stability criterion is the preconditioned sharpness: $\lambda^* = \frac{2}{\eta}$ for RMSProp, and $m$0 for Adam (with bias-corrected moments). The paper highlights the "acceleration via regularization" property: the preconditioner can absorb increasing sharpness, raising the raw Hessian eigenvalues that are sustainable before instability manifests.

Empirical Evaluation

The rod flow ODEs are validated on full-batch training of representative architectures (MLP, CNN, ViT) on CIFAR-10. For each optimizer, the rod flow solution is shown to closely track the centers of the discrete iterates and maintain correct amplitude and frequency of EoS oscillations, dramatically outperforming the naive continuous or "stable" flow ODE, which fails to capture the EoS regime.

Notably, the rod flow achieves self-consistency of preconditioned sharpness near the predicted thresholds, both numerically and analytically. The discrepancies between the rod flow and the discrete iterates plateau after a short transient, with persistent agreement in oscillation scale and alignment. For Adam, the rod flow precisely locks in on the correct value of preconditioned sharpness (as predicted by theory), even as raw sharpness continues to rise—a feature absent from any non-adaptive or non-momentum method.

The principal limitation identified is sensitivity to the momentum coefficient: at high $m$1, period-doubling and higher order oscillations can invalidate the period-2 rod framework, leading to degraded accuracy in those regimes. Mini-batch (stochastic) dynamics and transient (non-steady-state or progressive sharpening) phenomena are also beyond the current scope.

Implications

Practical Implications

The rod flow model for Adam at EoS establishes, for the first time, a principled ODE-based tool for precise tracking of high-amplitude, oscillatory optimization dynamics in finite-step-size adaptive and momentum methods. This provides a foundation for understanding:

• The stabilizing effects of momentum and preconditioning in large learning-rate training, including why training does not diverge even in flat-plateaued sharpness regimes.
• The implicit bias and regularization effects induced by the EoS regime, with potential implications for generalization and loss landscape navigation.
• The possible design of new optimizers or learning-rate schedules that exploit the robust self-stabilizing regime characterized by the rod flow.

Theoretical Implications

• The rod flow formalism rigorously rationalizes the existence of an "EoS attractor" in the learning dynamics, where the optimizer locks into a nonlinear oscillator phase governed by higher-order geometry of the loss landscape and the choice of optimizer hyperparameters.
• The preconditioned sharpness threshold ties together the mechanics of adaptive and momentum optimizers, unifying prior disparate stability analyses.
• The analysis of central rod dynamics invites further work on connecting the rod flow to higher-order or stochastic corrections that may account for the breakdown in period-doubling or multidirectional sharpness.

Future Directions

Research directions suggested include:

• Extension to stochastic (mini-batch) regimes, via diffusion-augmented ODEs or SDE analogues.
• Formal conditions for the onset of higher-period/doubling bifurcations in high-momentum regimes.
• Rod flows for non-Adam-family optimizers, including recently proposed architectures (e.g., Lion, Shampoo, Muon).
• Deeper connection to implicit regularization and generalization performance at scale.

Conclusion

This work provides a comprehensive theoretical and empirical framework for understanding the oscillatory regime of Adam and its relatives at the edge of stability, generalizing the rod flow to encompass both momentum and adaptive preconditioned structures. The proposed ODE formalism accurately models discrete optimizer iterates, predicts correct stability thresholds, and reveals new mechanistic insights into the self-stabilizing phenomena underpinning large-step-size deep learning.

The rod flow ODE, unifying Adam, RMSProp, heavy-ball, and Nesterov momentum, forms a robust foundation for future theoretical development and may underpin both practical diagnostic tools and principled optimizer design for deep networks at scale.