Unifying Central Flow, Rod Flow, and the Edge Coupling

Eric Regis; Sinho Chewi

stat.MLNov 5, 2026

Unifying Central Flow, Rod Flow, and the Edge Coupling

Authors: Eric Regis, Sinho Chewi

Organizations: Yale University

Abstract

At the edge of stability, the iterates of gradient descent oscillate along the sharp directions of the Hessian. This poses a significant challenge for continuous-time models as the usual continuous-time approximation -- the gradient flow -- admits no oscillations. Several continuous-time models have been proposed to capture the behavior at the edge of stability, among them central flow and rod flow. Although both models can be interpreted as performing a coarse-graining in time, the connection between them remains poorly understood. Here we show that both flows can be derived from a third object, the edge coupling: a functional on consecutive iterates whose critical points are the fixed points and period-two orbits of gradient descent. It is revealed that the two flows correspond to coarse-grainings at different time scales. Rod flow coarse-grains over only the fastest time scale, while central flow coarse-grains over both the fast and intermediate time scales. Furthermore, using the edge coupling, we derive rod flows and central flows for momentum methods such as heavy ball and Adam. We verify the derived flows experimentally on representative neural network architectures, where they track the discrete iterates throughout all phases of training -- including the transition from the stable regime to the edge of stability.

Figures & tables

Figure 1: The Edge Coupling Generates Both Rod Flow and Central Flow. Rod flow is gradient descent–ascent on the centered edge coupling, with the oscillations tracked by an explicit state variable. Central flow is constrained descent on the same objective, with the covariance matrix acting as a Lagrange multiplier. Central flow is rod flow averaged over the intermediate time scale of self-stabilization.

Figure 2: Rod Flow and Central Flow Track AdamW. AdamW on an MLP (

\eta=10^{-5}

\beta_{1}=0.9

\beta_{2}=0.999

\zeta=10

). Left: With both rod flow and central flow, the preconditioned sharpness self-stabilizes at the correct value. Right: Both flows track the discrete iterates more accurately than the stable flow does.

Figure 3: Second Moment Versus Covariance. Second-moment and covariance rod flows on an MLP (

\eta=0.05

). Left: Second-moment rod flow fails to self-stabilize at the transition, whereas covariance rod flow succeeds. Middle and right: Alignment of each flow’s extent with the mean gradient

\bar{g}_{t}

and with the gradient half-difference

\llbracket g\rrbracket_{t}

of the discrete iterates. The second-moment rod stays aligned with the mean gradient, while the covariance rod correctly aligns with the gradient half-difference at the onset.

Figure 4: The Alternating Term at the Onset. Gradient descent on an MLP (

\eta=0.05

). Top: Norms of the mean

\bar{g}_{t}

and the half-difference

\llbracket g\rrbracket_{t}

of consecutive gradients. Before the edge of stability,

\|\bar{g}_{t}\|\gg\|\llbracket g\rrbracket_{t}\|

. In the steady-state edge-of-stability regime,

\|\bar{g}_{t}\|\ll\|\llbracket g\rrbracket_{t}\|

. Bottom: Norm of the alternating term

\bar{g}_{t}\otimes\llbracket g\rrbracket_{t}+\llbracket g\rrbracket_{t}\otimes\bar{g}_{t}

\hat{\delta}_{t}\otimes\hat{\delta}_{t}

, relative to that of the non-alternating terms

\bar{g}_{t}^{\otimes 2}+\llbracket g\rrbracket_{t}^{\otimes 2}

. The alternating term is negligible before the crossing but peaks just after the onset, when

\bar{g}_{t}

and

\llbracket g\rrbracket_{t}

are similar in magnitude.

Figure 5: Heavy-Ball Central Flow. The gray region is the feasible region for heavy ball. While the stable flow leaves the feasible region, the central flow turns a corner at the boundary and follows it to a constrained minimum. Inset: The heavy-ball update points out of the feasible region. The sharpness-reduction force cancels its outward component, and what remains is the velocity of the central flow.

Figure 6: Central Flow as Time-Averaged Rod Flow. Top: The extent matrix of rod flow has rank one at any given instant. In the presence of multiple sharp eigenvalues, rod flow self-stabilizes by turning in the critical subspace. The covariance matrix of central flow can be interpreted as the time average of the extent matrix. It has rank

k

, where

k

is the number of critical eigenvalues. Bottom: Gradient descent on an MLP (

\eta=0.05

). There are two sharp eigenvalues. The extent matrix of rod flow and the discrete iterates are averaged over windows of

T

iterations. Left: As the width of the averaging window increases, their participation ratio rises from one to that of

\Sigma_{\mathrm{central}}

. Right: Their normalized Bures–Wasserstein distance to

\Sigma_{\mathrm{central}}

falls toward zero.

Figure 7: Turning Rate of the Rod. Gradient descent on an MLP (

\eta=0.05

). Left: Time series of the turning rate for rod flow and for the discrete oscillation, compared with the prediction of the linear turning law. Middle: Scatter plot of measured versus predicted turning rates for the discrete oscillation and for rod flow, showing strong agreement. Right: Participation ratio of the time-averaged extent matrix as a function of the swept angle.

Figure 8: Position and Momentum Extents of Heavy Ball. Heavy ball on an MLP (

\eta=0.3

\beta=0.9

). There are two critical eigenvalues. Left: The amplitude ratio

\lVert\gamma\rVert/\lVert\delta\rVert

over time. For period-two orbits, the amplitude ratio is

2/\eta

(dashed line). The ratio hovers around this value but shows non-trivial fluctuations. Middle: The alignment

\cos(-\gamma,\delta)

between the position and momentum extents. Right: The normalized Bures–Wasserstein distance between the time averages

\langle\tfrac{\eta^{2}}{4}\gamma\gamma^{\top}\rangle_{T}

and

\langle\delta\delta^{\top}\rangle_{T}

. As the averaging window

T

grows, the time-averaged position and momentum extent matrices increasingly agree with the relation they satisfy on a period-two cycle.

Figure 9: Experimental Results for Gradient Descent. Left: MLP (

\eta=0.05

). Center: CNN (

\eta=0.05

). Right: ViT (

\eta=0.1

). The flows are handed off from the discrete iterates

1\,000

iterations before the crossing (

500

for the ViT). All flows except the stable flow self-stabilize at the sharpness threshold, and they track the discrete iterates orders of magnitude more closely than the stable flow does.

Figure 10: Tracking Through Three Critical Eigenvalues. Gradient descent on an MLP (

\eta=0.05

). Top: The top six eigenvalues of the Hessian at the discrete midpoint.

\lambda_{1}

\lambda_{2}

, and

\lambda_{3}

reach the threshold one after another. Bottom: The distance from each flow’s center to the discrete midpoint. Rod flow and central flow stay close to the iterates through all three transitions, while the stable flow drifts away.

Figure 11: Experimental Results for Adam (with bias correction). Left: MLP (

\eta=10^{-5}

). Center: CNN (

\eta=5\times 10^{-5}

). Right: ViT (

\eta=10^{-5}

). All three with

\beta_{1}=0.9

\beta_{2}=0.999

. All flows except the stable flow self-stabilize the preconditioned sharpness at its threshold

38/\eta

, and they track the discrete iterates more closely than the stable flow does.

Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$\eta$	Step size
$\beta$	Momentum coefficient
$\beta_{1}$	First-moment decay rate
$\beta_{2}$	Second-moment decay rate
$\epsilon$	Numerical stability constant
$d$	Number of parameters

Appendix

Table 1: Notation used throughout the main body.

Figure 12: Round-Off Sets the Floor. Gradient descent on an MLP at

\eta=0.05

( left ) and

\eta=0.075

( right ), run in single and double precision from the same initialization. Shown is the sharp component

|\langle a_{t},u_{1}\rangle|

of the oscillation amplitude. Before the crossing, the sharp component sits at the floor predicted by the rounding model, which is about eight orders of magnitude lower in double precision. As a result, the double-precision runs self-stabilize later than the single-precision runs, though eventually they settle on the same plateau.

Optimizer	State	$A$
Gradient descent	$w$	$I_{d}-\eta H$
RMSProp	$w$	$I_{d}-\eta P^{-1}H$
Heavy ball	$(w,m)$	$\begin{pmatrix}I_{d}-\eta(1-\beta)H&-\eta\beta I_{d}\\ (1-\beta)H&\beta I_{d}\end{pmatrix}$
Nesterov	$(w,m)$	$\begin{pmatrix}I_{d}-\eta(1-\beta)H&-\eta\beta\bigl(I_{d}-\eta(1-\beta)H\bigr)\\ (1-\beta)H&\beta\bigl(I_{d}-\eta(1-\beta)H\bigr)\end{pmatrix}$
Adam	$(w,m)$	$\begin{pmatrix}I_{d}-\eta(1-\beta_{1})P^{-1}H&-\eta\beta_{1}P^{-1}\\ (1-\beta_{1})H&\beta_{1}I_{d}\end{pmatrix}$
AdamW	$(w,m)$	$\begin{pmatrix}\kappa I_{d}-\eta(1-\beta_{1})P^{-1}H&-\eta\beta_{1}P^{-1}\\ (1-\beta_{1})H&\beta_{1}I_{d}\end{pmatrix}$

Appendix

Table 2: Linearized Update Maps. The linear part

A

of each optimizer’s update under the quadratic approximation about the midpoint. The preconditioner

P

is held fixed and bias correction is omitted.

Figure 13: A Floor for Each Precision. Midpoint rod flow with the single- and double-precision rounding floors, compared against gradient descent run in each precision from the same initialization (MLP,

\eta=0.05

). Top: The magnitude of the oscillation along the top eigenvector

\hat{u}_{1}

over time. Gradient descent and rod flow agree closely at each precision. Bottom: The sharpness at the onset of self-stabilization. Self-stabilization lags behind the crossing of the sharpness threshold, and the lag is longer in double precision for both gradient descent and rod flow.

Figure 14: Floor Versus No Floor. Comparison of midpoint rod flow, endpoint rod flow, edge flow, and the mean–fluctuation model against gradient descent on an MLP (learning rate

\eta=0.05

). Each model is shown with a noise floor (solid) and without it (dotted). The floor of the mean–fluctuation model is ours ( Section G.4 ). Top and middle: Oscillation along the top eigenvector

\hat{u}_{1}

, shown at two zoom levels. Without a floor, every flow’s oscillation decays by more than fifty orders of magnitude before the crossing, and endpoint rod flow’s collapses to exactly zero. Bottom: Without a floor, the flows arrest late. With it, every flow arrests at the same time as the gradient descent iterates.

Optimizer	$\mathsf{A}$	$\mathsf{B}[X]$
Gradient descent	$U^{\top}\mathcal{T}^{*}[\nabla L]\,U$	$\tfrac{1}{2}\,U^{\top}\mathcal{T}^{*}\mathcal{T}[UXU^{\top}]\,U$
Heavy ball	$U^{\top}\mathcal{T}^{*}\bigl[\beta\bar{m}+(1-\beta)\nabla L\bigr]U$	$\tfrac{1-\beta}{2}\,U^{\top}\mathcal{T}^{*}\mathcal{T}[UXU^{\top}]\,U$
Nesterov	$U^{\top}\mathcal{T}^{*}\bigl[\beta^{2}\bar{m}+(1-\beta^{2})\nabla L(\bar{\theta})\bigr]U$	$\tfrac{1-\beta^{2}}{2}\,U^{\top}\mathcal{T}^{*}\mathcal{T}[UXU^{\top}]\,U$
RMSProp	$U^{\top}\mathcal{T}^{*}\bigl[P^{-1}\nabla L\bigr]U$	$\tfrac{1}{2}\,U^{\top}\mathcal{T}^{*}\bigl[P^{-1}\mathcal{T}[UXU^{\top}]\bigr]U$
	${}+\tfrac{1-\beta_{2}}{\eta}\,U^{\top}\nabla_{\bar{\nu}}\Lambda\bigl[(\nabla L)^{\odot 2}-\bar{\nu}\bigr]U$	${}+\tfrac{1-\beta_{2}}{\eta}\,U^{\top}\nabla_{\bar{\nu}}\Lambda\bigl[\operatorname{diag}(HUXU^{\top}H)\bigr]U$
Adam	$U^{\top}\mathcal{T}^{*}\bigl[P^{-1}\bigl(\beta_{1}\bar{m}+(1-\beta_{1})\nabla L\bigr)\bigr]U$	$\tfrac{1-\beta_{1}}{2}\,U^{\top}\mathcal{T}^{*}\bigl[P^{-1}\mathcal{T}[UXU^{\top}]\bigr]U$

Appendix

Table 3: Inputs to the Quadratic Program. The rate of change of the critical block of the slack,

U^{\top}\dot{\Lambda}U=\eta(\mathsf{A}+\mathsf{B}[X])

, for each optimizer.

\eta\mathsf{A}

is the rate of change under the free flow.

\eta\mathsf{B}[X]

is the rate of change due to the oscillations. For RMSProp, Adam, and AdamW,

\Lambda

denotes that optimizer’s slack.

Optimizer	$\mathcal{L}$	$\Omega$	$W_{\nu}$
Gradient descent	$L(\bar{w})$	$\{\Lambda(\bar{w})\succeq 0\}$	—
Heavy ball	$L(\bar{w})-\tfrac{\eta\beta}{2(1-\beta)}\\|\bar{m}\\|^{2}$	$\{\Lambda_{\mathrm{HB}}(\bar{w})\succeq 0\}$	—
Nesterov	$L(\bar{\theta})-\tfrac{\eta\beta^{2}}{2(1-\beta)}\\|\bar{m}\\|^{2}$	$\{\Lambda_{\mathrm{NAG}}(\bar{\theta})\succeq 0\}$	—
RMSProp	$L(\bar{w})$	$\{\Lambda_{\mathrm{RMSProp}}(\bar{w},\bar{\nu})\succeq 0\}$	$\tfrac{8(1-\beta_{2})}{\eta}\,P^{2}\operatorname{Diag}\bigl(\sqrt{\bar{\nu}}\bigr)$
Adam	$L(\bar{w})-\tfrac{\eta\beta_{1}}{2(1-\beta_{1})}\bar{m}^{\top}P^{-1}\bar{m}$	$\{\Lambda_{\mathrm{Adam}}(\bar{w},\bar{\nu})\succeq 0\}$	$\tfrac{8(1-\beta_{2})(1+\beta_{1})}{\eta(1-\beta_{1})}\,P^{2}\operatorname{Diag}\bigl(\sqrt{\bar{\nu}}\bigr)$
AdamW	$L(\bar{w})-\tfrac{\eta\beta_{1}}{2(1-\beta_{1})}\bar{m}^{\top}P^{-1}\bar{m}$	$\{\Lambda_{\mathrm{AdamW}}(\bar{w},\bar{\nu})\succeq 0\}$	$\tfrac{4(1-\beta_{2})(1+\kappa)(1+\beta_{1})}{\eta(1-\beta_{1})}\,P^{2}\operatorname{Diag}\bigl(\sqrt{\bar{\nu}}\bigr)$

Appendix

Table 4: Central Flow Objectives. For each optimizer, we give the generating function

\mathcal{L}

and the feasible region

\Omega

of the objective. For the preconditioned methods, we give the metric

W_{\nu}

through which the normal cone acts on the second moment.

	MLP			CNN			ViT
Optimizer	$\eta$	$t_{h}$	$t_{c}$	$\eta$	$t_{h}$	$t_{c}$	$\eta$	$t_{h}$	$t_{c}$
Gradient descent	$0.05$	$8\,700$	$9\,700$	$0.05$	$19\,600$	$20\,600$	$0.1$	$4\,600$	$5\,100$
RMSProp	$2\times 10^{-6}$	$1\,500$	$2\,500$	$2\times 10^{-6}$	$3\,700$	$4\,800$	$2\times 10^{-6}$	$3\,700$	$4\,200$
Heavy ball	$0.3$	$3\,800$	$4\,800$	$0.7$	$1\,600$	$2\,800$	$0.2$	$5\,700$	$6\,200$
Nesterov	$0.2$	$3\,100$	$4\,100$	$0.4$	$1\,200$	$2\,200$	$0.2$	$3\,700$	$4\,300$
Adam	$10^{-5}$	$3\,900$	$4\,900$	$5\times 10^{-5}$	$1\,000$	—	$10^{-5}$	$4\,000$	$4\,500$

Appendix

Table 5: Runs. The step size

\eta

, handoff

t_{h}

, and crossing

t_{c}

of each run.

Optimizer	Effective sharpness	Threshold $S^{\star}$
Gradient descent	$\lambda_{1}(H)$	$\dfrac{2}{\eta}$
RMSProp	$\lambda_{1}(P^{-1/2}HP^{-1/2})$	$\dfrac{2}{\eta}$
Heavy ball	$\lambda_{1}(H)$	$\dfrac{2(1+\beta)}{\eta(1-\beta)}$
Nesterov	$\lambda_{1}(H)$	$\dfrac{2(1+\beta)}{\eta(1-\beta)(1+2\beta)}$
Adam	$\lambda_{1}(\hat{P}^{-1/2}H\hat{P}^{-1/2})$	$\dfrac{2(1+\beta_{1})}{\eta(1-\beta_{1})}$
AdamW	$\lambda_{1}(\hat{P}^{-1/2}H\hat{P}^{-1/2})$	$\Bigl(1-\dfrac{\eta\zeta}{2}\Bigr)\dfrac{2(1+\beta_{1})}{\eta(1-\beta_{1})}$

Appendix

Table 6: Effective Sharpness and Thresholds. The effective sharpness of each optimizer and its threshold

S^{\star}

Figure 15: Double-Precision Flow State. Distance between the center of midpoint rod flow and the discrete midpoint, with the flow’s state stored in single and double precision (gradient descent, MLP,

\eta=0.05

). With single precision, the distance grows as rounding error accumulates, reaching several times that of the double-precision flow.

Figure 16: Flow Comparison for Gradient Descent. MLP:

\eta=0.05

. CNN:

\eta=0.05

. ViT:

\eta=0.1

. The rows show, from top to bottom: the loss at each center; the sharpness at each center; the top six eigenvalues at the discrete optimizer’s midpoint; the oscillation amplitude; the alignment of each flow’s oscillation with the discrete-time oscillation; and the distance from each flow’s center to the discrete optimizer’s midpoint. Dashed vertical lines mark the handoff, and dotted vertical lines mark the crossing into the edge of stability.

Figure 17: Flow Comparison for RMSProp (no bias correction). MLP:

\eta=2\times 10^{-6}

. CNN:

\eta=2\times 10^{-6}

. ViT:

\eta=2\times 10^{-6}

. All with

\beta_{2}=0.99

and

\epsilon=10^{-8}

. The rows show, from top to bottom: the loss at each center; the preconditioned sharpness at each center; the raw sharpness at each center; the top six preconditioned eigenvalues at the discrete optimizer’s midpoint; the oscillation amplitude; the alignment of each flow’s oscillation with the discrete-time oscillation; the distance from each flow’s center to the discrete optimizer’s midpoint; and the norm of the second-moment estimate. Dashed vertical lines mark the handoff, and dotted vertical lines mark the crossing into the edge of stability.

Figure 18: Flow Comparison for Heavy Ball Momentum. MLP:

\eta=0.3

. CNN:

\eta=0.7

. ViT:

\eta=0.2

. All with

\beta=0.9

Figure 19: Flow Comparison for Nesterov Momentum. MLP:

\eta=0.2

. CNN:

\eta=0.4

. ViT:

\eta=0.2

. All with

\beta=0.9

Figure 20: Flow Comparison for Adam (with bias correction). MLP:

\eta=10^{-5}

. CNN:

\eta=5\times 10^{-5}

. ViT:

\eta=10^{-5}

. All with

\beta_{1}=0.9

\beta_{2}=0.999

, and

\epsilon=10^{-7}

Figure 21: Flow Comparison for AdamW (with bias correction). MLP:

\eta=10^{-5}

\zeta=10

. CNN:

\eta=5\times 10^{-5}

\zeta=40

. ViT:

\eta=10^{-5}

\zeta=10

. All with

\beta_{1}=0.9

\beta_{2}=0.999

, and

\epsilon=10^{-7}

Figure 22: Flow Comparison for Gradient Descent Under Cross-Entropy Loss. MLP:

\eta=0.015

Figure 23: Flow Comparison for RMSProp Under Cross-Entropy Loss (no bias correction). MLP:

\eta=10^{-5}

\beta_{2}=0.99

\epsilon=10^{-8}

Figure 24: Rod Turning for Gradient Descent. MLP. Top: Two critical eigenvalues,

\eta=0.05

. Bottom: Three critical eigenvalues,

\eta=0.075

. In each half, the first row compares the time-averaged extent matrix with

\Sigma_{\mathrm{central}}

. The second row compares the turning rate with the linear law, as in Figure 7 .

Figure 25: Rod Turning for RMSProp (no bias correction). MLP,

\eta=2\times 10^{-6}

\beta_{2}=0.99

. Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with

\Sigma_{\mathrm{central}}

. The second row compares the turning rate with the linear law, as in Figure 7 .

Figure 26: Rod Turning for Heavy Ball Momentum. MLP,

\eta=0.3

\beta=0.9

. Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with

\Sigma_{\mathrm{central}}

. The second row compares the turning rate with the linear law, as in Figure 7 .

Figure 27: Rod Turning for Nesterov Momentum. MLP,

\eta=0.2

\beta=0.9

. Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with

\Sigma_{\mathrm{central}}

. The second row compares the turning rate with the linear law, as in Figure 7 .

Figure 28: Rod Turning for Adam (with bias correction). MLP,

\eta=10^{-5}

\beta_{1}=0.9

\beta_{2}=0.999

. Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with

\Sigma_{\mathrm{central}}

. The second row compares the turning rate with the linear law, as in Figure 7 .

Figure 29: Rod Turning for AdamW (with bias correction). MLP,

\eta=10^{-5}

\zeta=10

\beta_{1}=0.9

\beta_{2}=0.999

. Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with

\Sigma_{\mathrm{central}}

. The second row compares the turning rate with the linear law, as in Figure 7 .

Figure 30: Position and Momentum Extents of the Momentum Optimizers. MLP, with two critical eigenvalues. Heavy ball:

\eta=0.3

\beta=0.9

. Nesterov:

\eta=0.2

\beta=0.9

. Adam (with bias correction):

\eta=10^{-5}

\beta_{1}=0.9

\beta_{2}=0.999

. AdamW: as Adam, with

\zeta=10

. The columns show, from left to right: the ratio of the momentum amplitude to the position amplitude, with its value on a period-two orbit (dashed); the alignment between the position and momentum extents; and the distance between their time averages, against the width of the averaging window. For Adam and AdamW, the momentum extent is preconditioned.

Figure 31: Rate of Change of the Relative Momentum. Runs as in Figure 30 . The relative momentum $\tilde{\gamma}$ is the momentum extent divided by the position amplitude. Left: The relative change of $\tilde{\gamma}$ per iteration, for the discrete oscillation and for rod flow, with the linear law for rod flow (dashed). Right: Measured against the linear law.

Figure 32: Power Spectrum of the Discrete Oscillation. MLP with two critical eigenvalues. Runs as in the top halves of Figures 24 , 25 , 26 , 27 , 28 and 29 . Each panel shows the power spectrum of the oscillation of the iterates, normalized to unit area, so that the area under the curve over any range of frequencies is the share of the power in that range. The gray band marks the frequencies within

0.03

\tfrac{1}{2}

, and the number in each panel gives the share of the power in this band (dark area). Dashed lines mark where period doubling to periods four, eight, and sixteen would add power.

Figure 33: The Self-Stabilization Loop. Gradient descent and rod flow on an MLP. Top:

\eta=0.05

. Bottom:

\eta=0.075

. In each half, the left panel shows the phase portrait of the loop in the log-amplitude and the excess sharpness. The right panels show the log-amplitude and the progressive sharpening constant over time. The dashed curve is the progressive sharpening constant of the stable flow plus the response of the center to the oscillation, with no free parameters.

Figure 34: The Kicked Ring. Gradient descent and rod flow on an MLP, kicked by multiplying the amplitude of the oscillation by

1.2

. Top:

\eta=0.05

. Bottom:

\eta=0.075

. In each half, the left panel shows the phase portrait, with dots marking the state at the kick and arrows marking the kick. The right panels show the log-amplitude and the progressive sharpening constant over time. The dashed curves are predictions. For the log-amplitude, only the frequency and the decay rate are predicted; the amplitude, phase, and trend are fitted.

Gradient Descent Dynamics

Gradient Descent Dynamics studies how iterative gradient-based updates shape optimization trajectories, stability, convergence, and the solutions models learn.

See all 166 papers

Jun 16, 2026

Edge Flow: A Tractable and Predictive Continuous-Time Model for Gradient Descent at the Edge of Stability

Gradient descent in deep learning may operate at the edge of stability (EoS), a regime in which the largest eigenvalue of the loss Hessian hovers near the stability threshold

2/η

, where

η

is the learning rate. Classical analysis tools such as gradient flow and the descent lemma do not apply here, motivating the search for a continuous-time model valid at EoS. We propose Edge Flow, a system of three coupled ordinary differential equations that provides a tractable, faithful, and predictive model of gradient descent dynamics at EoS. Edge Flow decomposes the dynamics into a center, an oscillation direction, and an oscillation magnitude. The center follows a modified gradient flow on a symmetrized loss; the direction tracks a top eigenvector of the Hessian via Rayleigh quotient dynamics; and the magnitude grows or decays exponentially depending on whether the sharpness exceeds or falls below the threshold

2/η

. Crucially, sharpness stabilization emerges from the coupled dynamics via a self-stabilization feedback loop. Discretizing Edge Flow only requires two gradient evaluations and one Hessian--vector product at each iteration. We demonstrate empirically that Edge Flow tracks the dynamics of gradient descent at least as faithfully as previously proposed continuous-time EoS models, while in addition resolving the oscillation of the sharpness at the onset of EoS, and that it provides a principled framework for understanding and mitigating instabilities in this regime.

Pierre Marion

Inria, École Normale Supérieure, PSL Research University

Edge of Stability Neural Network Training Dynamics Gradient Descent Dynamics

Ask PDF

May 7, 2026

A Rod Flow Model for Adam at the Edge of Stability

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.

Eric Regis, Sinho Chewi

Yale University

Edge of Stability Momentum Methods Adaptive Gradient Methods Gradient Descent Dynamics

Ask PDF

Sep 1, 2026

The Multiple Timescales of Gradient Descent on the Edge of Stability: A Perturbative Derivation of the Central Flow

The central flow of Cohen et al. (2025) is an empirically accurate continuous-time model of gradient descent at the edge of stability in deep learning, However, its derivation is heuristic. We propose a perturbative regime in which the central flow is the limit of gradient descent: we assume that the loss decomposes as

f = g + \varepsilon h

; in the limit

\varepsilon \to 0

, the dynamics of gradient descent with learning rate

η

converge to the gradient flow of

h

constrained to the minimizers of

g

of sharpness at most

2/η

. Our approach is formal rather than rigorous; it treats gradient descent as a singularly perturbed dynamical system in

\varepsilon

. Three timescales emerge: a fast timescale of oscillations along the sharpest direction, an intermediate timescale of the self-stabilization mechanism, and a slow timescale of the dynamics along the minimizers of

g

-the central flow. Using the method of multiple scales, a classical formal method from singular perturbation theory, we derive the expansion of the dynamics in

\varepsilon

: the central flow emerges as the leading-order term in the expansion, while the self-stabilization mechanism appears in the next-order term. We study this mechanism beyond previous analyses: with a single eigenvalue at the edge of stability, we compute the slow drift of the energy of the fluctuations; with several eigenvalues at the edge of stability, we derive the self-stabilization system and explain why fluctuations persist.

Raphaël Berthier

Sorbonne Université, Inria, Centre Inria de Sorbonne Université, Paris, France

Dynamical Systems Edge of Stability Gradient Descent Dynamics

Ask PDF

Abstract

Gradient Descent Dynamics

Explore similar work

Edge Flow: A Tractable and Predictive Continuous-Time Model for Gradient Descent at the Edge of Stability

A Rod Flow Model for Adam at the Edge of Stability

The Multiple Timescales of Gradient Descent on the Edge of Stability: A Perturbative Derivation of the Central Flow