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 ( η=10−5 , β1=0.9 , β2=0.999 , ζ=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 ( η=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 gˉt and with the gradient half-difference [[g]]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 ( η=0.05 ). Top: Norms of the mean gˉt and the half-difference [[g]]t of consecutive gradients. Before the edge of stability, ∥gˉt∥≫∥[[g]]t∥ . In the steady-state edge-of-stability regime, ∥gˉt∥≪∥[[g]]t∥ . Bottom: Norm of the alternating term gˉt⊗[[g]]t+[[g]]t⊗gˉt in δ^t⊗δ^t , relative to that of the non-alternating terms gˉt⊗2+[[g]]t⊗2 . The alternating term is negligible before the crossing but peaks just after the onset, when gˉt and [[g]]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 ( η=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 Σcentral . Right: Their normalized Bures–Wasserstein distance to Σcentral falls toward zero.
Figure 7: Turning Rate of the Rod. Gradient descent on an MLP ( η=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 ( η=0.3 , β=0.9 ). There are two critical eigenvalues. Left: The amplitude ratio ∥γ∥/∥δ∥ over time. For period-two orbits, the amplitude ratio is 2/η (dashed line). The ratio hovers around this value but shows non-trivial fluctuations. Middle: The alignment cos(−γ,δ) between the position and momentum extents. Right: The normalized Bures–Wasserstein distance between the time averages ⟨4η2γγ⊤⟩T and ⟨δδ⊤⟩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 ( η=0.05 ). Center: CNN ( η=0.05 ). Right: ViT ( η=0.1 ). The flows are handed off from the discrete iterates 1000 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 ( η=0.05 ). Top: The top six eigenvalues of the Hessian at the discrete midpoint. λ1 , λ2 , and λ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 ( η=10−5 ). Center: CNN ( η=5×10−5 ). Right: ViT ( η=10−5 ). All three with β1=0.9 , β2=0.999 . All flows except the stable flow self-stabilize the preconditioned sharpness at its threshold 38/η , 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
η
Step size
β
Momentum coefficient
β1
First-moment decay rate
β2
Second-moment decay rate
ϵ
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 η=0.05 ( left ) and η=0.075 ( right ), run in single and double precision from the same initialization. Shown is the sharp component ∣⟨at,u1⟩∣ 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.
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, η=0.05 ). Top: The magnitude of the oscillation along the top eigenvector 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 η=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 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
A
B[X]
Gradient descent
U⊤T∗[∇L]U
21U⊤T∗T[UXU⊤]U
Heavy ball
U⊤T∗[βmˉ+(1−β)∇L]U
21−βU⊤T∗T[UXU⊤]U
Nesterov
U⊤T∗[β2mˉ+(1−β2)∇L(θˉ)]U
21−β2U⊤T∗T[UXU⊤]U
RMSProp
U⊤T∗[P−1∇L]U
21U⊤T∗[P−1T[UXU⊤]]U
+η1−β2U⊤∇νˉΛ[(∇L)⊙2−νˉ]U
+η1−β2U⊤∇νˉΛ[diag(HUXU⊤H)]U
Adam
U⊤T∗[P−1(β1mˉ+(1−β1)∇L)]U
21−β1U⊤T∗[P−1T[UXU⊤]]U
Appendix
Table 3: Inputs to the Quadratic Program. The rate of change of the critical block of the slack, U⊤Λ˙U=η(A+B[X]) , for each optimizer. ηA is the rate of change under the free flow. ηB[X] is the rate of change due to the oscillations. For RMSProp, Adam, and AdamW, Λ denotes that optimizer’s slack.
Optimizer
L
Ω
Wν
Gradient descent
L(wˉ)
{Λ(wˉ)⪰0}
—
Heavy ball
L(wˉ)−2(1−β)ηβ∥mˉ∥2
{ΛHB(wˉ)⪰0}
—
Nesterov
L(θˉ)−2(1−β)ηβ2∥mˉ∥2
{ΛNAG(θˉ)⪰0}
—
RMSProp
L(wˉ)
{ΛRMSProp(wˉ,νˉ)⪰0}
η8(1−β2)P2Diag(νˉ)
Adam
L(wˉ)−2(1−β1)ηβ1mˉ⊤P−1mˉ
{ΛAdam(wˉ,νˉ)⪰0}
η(1−β1)8(1−β2)(1+β1)P2Diag(νˉ)
AdamW
L(wˉ)−2(1−β1)ηβ1mˉ⊤P−1mˉ
{ΛAdamW(wˉ,νˉ)⪰0}
η(1−β1)4(1−β2)(1+κ)(1+β1)P2Diag(νˉ)
Appendix
Table 4: Central Flow Objectives. For each optimizer, we give the generating function L and the feasible region Ω of the objective. For the preconditioned methods, we give the metric Wν through which the normal cone acts on the second moment.
MLP
CNN
ViT
Optimizer
η
th
tc
η
th
tc
η
th
tc
Gradient descent
0.05
8700
9700
0.05
19600
20600
0.1
4600
5100
RMSProp
2×10−6
1500
2500
2×10−6
3700
4800
2×10−6
3700
4200
Heavy ball
0.3
3800
4800
0.7
1600
2800
0.2
5700
6200
Nesterov
0.2
3100
4100
0.4
1200
2200
0.2
3700
4300
Adam
10−5
3900
4900
5×10−5
1000
—
10−5
4000
4500
Appendix
Table 5: Runs. The step size η , handoff th , and crossing tc of each run.
Optimizer
Effective sharpness
Threshold S⋆
Gradient descent
λ1(H)
η2
RMSProp
λ1(P−1/2HP−1/2)
η2
Heavy ball
λ1(H)
η(1−β)2(1+β)
Nesterov
λ1(H)
η(1−β)(1+2β)2(1+β)
Adam
λ1(P^−1/2HP^−1/2)
η(1−β1)2(1+β1)
AdamW
λ1(P^−1/2HP^−1/2)
(1−2ηζ)η(1−β1)2(1+β1)
Appendix
Table 6: Effective Sharpness and Thresholds. The effective sharpness of each optimizer and its threshold S⋆ .
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, η=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: η=0.05 . CNN: η=0.05 . ViT: η=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: η=2×10−6 . CNN: η=2×10−6 . ViT: η=2×10−6 . All with β2=0.99 and ϵ=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: η=0.3 . CNN: η=0.7 . ViT: η=0.2 . All with β=0.9 . 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 19: Flow Comparison for Nesterov Momentum. MLP: η=0.2 . CNN: η=0.4 . ViT: η=0.2 . All with β=0.9 . 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 20: Flow Comparison for Adam (with bias correction). MLP: η=10−5 . CNN: η=5×10−5 . ViT: η=10−5 . All with β1=0.9 , β2=0.999 , and ϵ=10−7 . 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 (none on the CNN, which never crosses).
Figure 21: Flow Comparison for AdamW (with bias correction). MLP: η=10−5 , ζ=10 . CNN: η=5×10−5 , ζ=40 . ViT: η=10−5 , ζ=10 . All with β1=0.9 , β2=0.999 , and ϵ=10−7 . 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 22: Flow Comparison for Gradient Descent Under Cross-Entropy Loss. MLP: η=0.015 . 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 23: Flow Comparison for RMSProp Under Cross-Entropy Loss (no bias correction). MLP: η=10−5 , β2=0.99 , ϵ=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 24: Rod Turning for Gradient Descent. MLP. Top: Two critical eigenvalues, η=0.05 . Bottom: Three critical eigenvalues, η=0.075 . In each half, the first row compares the time-averaged extent matrix with Σ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, η=2×10−6 , β2=0.99 . Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with Σ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, η=0.3 , β=0.9 . Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with Σcentral . The second row compares the turning rate with the linear law, as in Figure 7 .
Figure 27: Rod Turning for Nesterov Momentum. MLP, η=0.2 , β=0.9 . Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with Σ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, η=10−5 , β1=0.9 , β2=0.999 . Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with Σ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, η=10−5 , ζ=10 , β1=0.9 , β2=0.999 . Top: Two critical eigenvalues. Bottom: Three critical eigenvalues. In each half, the first row compares the time-averaged extent matrix with Σ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: η=0.3 , β=0.9 . Nesterov: η=0.2 , β=0.9 . Adam (with bias correction): η=10−5 , β1=0.9 , β2=0.999 . AdamW: as Adam, with ζ=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 γ~ is the momentum extent divided by the position amplitude. Left: The relative change of γ~ 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 of 21 , 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: η=0.05 . Bottom: η=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: η=0.05 . Bottom: η=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.