# A Rod Flow Model for Adam at the Edge of Stability

arXiv: https://arxiv.org/abs/2605.06821 · HTML: https://mlsanity.com/paper/2605.06821/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 6 of its 52 references that are in the map.
- First posted on arXiv 7 May 2026.

## Builds on (most cited first, 6 of 6)

- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Adaptive Gradient Methods at the Edge of Stability](https://mlsanity.com/paper/2207.14484/) (2022)
- [ODE approximation for the Adam algorithm: General and overparametrized setting](https://mlsanity.com/paper/2511.04622/) (2025)
- [Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2602.01480/) (2026)
- [Modeling AdaGrad, RMSProp, and Adam with Integro-Differential Equations](https://mlsanity.com/paper/2411.09734/) (2024)
- [Continuous-Time Analysis of Adaptive Optimization and Normalization](https://mlsanity.com/paper/2411.05746/) (2024)

## Built on by (most cited first, 3 of 3)

- [The Road Taken: The Role of Optimizers at the Edge of Stability](https://mlsanity.com/paper/2608.18415/) (2026)
- [Adam at the Edge of Stability: Adaptive Feedback, Provable Oscillation, and Gradient Reversal](https://mlsanity.com/paper/2608.20638/) (2026)
- [Variational Structure at the Edge of Stability](https://mlsanity.com/paper/2608.21660/) (2026)

## Abstract (arXiv)

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 nu 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.