Understanding Optimization in Deep Learning with Central Flows
Authors: Jeremy Cohen
Alex Damian
Ameet Talwalkar
J Zico Kolter
Jason D. Lee
Venue: ICLR 2025 Poster (2025)
Keywords: Edge of Stability
Optimization Dynamics
Adaptive Optimizers
RMSProp
Primary area: optimization
TL;DR: We show how to understand the dynamics of oscillatory optimizers in deep learning

Abstract: Optimization in deep learning remains poorly understood.  A key difficulty is that optimizers exhibit complex oscillatory dynamics, referred to as "edge of stability," which cannot be captured by traditional optimization theory.  In this paper, we show that the path taken by an oscillatory optimizer can often be captured by a  _central flow_: a differential equation which directly models the time-averaged (i.e. smoothed) optimization trajectory. We empirically show that these central flows can predict long-term optimization trajectories for generic neural networks with a high degree of numerical accuracy.  By interpreting these flows, we are able to understand  how gradient descent makes progress even as the loss sometimes goes up; how adaptive optimizers ``adapt'' to the local loss landscape; and how adaptive optimizers implicitly seek out regions of weight space where they can take larger steps.  These insights (and others) are not apparent from the optimizers' update rules, but are revealed by the central flows.  Therefore, we believe that central flows constitute a promising tool for reasoning about optimization in deep learning.

Published as a conference paper at ICLR 2025
UNDERSTANDING OPTIMIZATION IN DEEP LEARNING WITH CENTRAL FLOWS
Jeremy M. Cohen & Alex Damian & Ameet Talwalkar' & J. Zico Kolter' & Jason D. Lee²
Carnegie Mellon University, Princeton University. Flatiron Institute
* *Alex and Jeremy contributed equally; author order determined by coin flip Corespondence to: (ad278 perr7rrriintooooom edu, jcohen& flat www.troniinnitteeoor.o
ABSTRACT
Optimization in deep learning remains poorly understood. A key difficulty is that optimizers exhibit complex oscillatory dynamics, referred to "edge of stability," which cannot be captured by traditional optimization theory. In this paper, we show that the path taken by an oscillatory optimizer can often be captured by a central flow: a differential equation which directly models the time-averaged (i.e. smoothed) optimization trajectory. empirically show that these central flows can predict long-term optimization trajectories for generie neural networks with a high degree of numerical accuracy. By interpreting these flows, we are able to understand how gradient descent makes progress up: how adaptive optimizers "adapt" to the local loss landscape; and how adaptive optimizers implicitly seek regions of weight space where they can take larger steps. These insights (and others) are not apparent from the optimizers' update rules, but revealed by the central flows. Therefore, we believe that central flows constitute a promising tool for reasoning about optimization in deep learning.
the loss sometimes goes
INTRODUCTION
While there is a rich body of work the theory of optimization, few works attempt to analyze optimization in "real" deep learning settings. Instead, even works motivated by deep learning often rely unrealistic assumptions such convexity, restrict their analyses simplified architectures. Practitioners cannot use such theories to reason directly about their optimization problems. Our goal in this paper is to develop optimization theory that applies directly deep learning problems. This is a difficult task: prior research has shown that, even in the seemingly simple setting of deterministic (i.e. full-batch) training, optimization typically operates
complex, oscillatory regime called the edge of stability (Xing et al. 2018 Wu et al. 2018) Jastrzebski et al. 2019 20200 Cohen et al. 2021 2022). The dynamies in this regime cannot be captured by traditional optimization theory.
In this paper, devise a methodology for analyzing these complex deep learning dynamics. We argue that an oscillatory optimizer should be conceived of as "oscillating around"' a particular path through weight space, which we call a central flow. The central flow can be interpreted as the "time-averaged" (i.e. smoothed) optimization trajectory (Figure 2).
We show that this central flow
Gradient Descent on ResNet
Gradient Descent on Transformer
RMSProp on Mamba
1.4
n-
n 2/150
n= 2/15
n= 1e-05
loss
1.0
2/100
2/100
1.0
2w-05
0.8
1.5
0.0
0.6-
1.0
0.5
0.0
0.2
1000
2000
3000
3000
5000
step
step
step
Figure 1: Our theory can deliver accurate numerical predictions on actual neural networks. For example, our central flows can accurately predict (dashed black) the time-averaged (i.c. smoothed) loss curves of gradient descent and RMSProp three architectures different learning rates.

---

Published as a conference paper at ICLR 2025
Figure The central flow models the time-averaged (smoothed) trajectory of the oscillatory optimizer. In this representative cartoon, gradient descent takes an oscillatory path through weight space. The central flow a smooth curve that characterizes this trajectory, whereas gradient flow takes different path.
Gradient Descent
Gradient Flow
Central Flow
can be captured by a differential equation which explicitly characterizes how the optimizer's path through weight space depends the local loss landscape and the optimizer's hyperparameters. As a smooth curve, the central flow is easier to understand than the original oscillatory trajectory. Hence, by reasoning about the central flow, we can reason more easily about the original optimizer.
We start in Section 3 by analyzing gradient descent, the simplest optimizer. The dynamics of gradient descent in deep learning are complex and potentially intimidating; for example, the training loss curve behaves non-monotonically. Yet we show that gradient descent is merely "oscillating around" a well- behaved central flow that does decrease the loss monotonically. We then examine a simple adaptive optimizer in Section 4, before turning to RMSProp (i.e. Adam without momentum) in Section 5 We show that much of the behavior of these optimizers is actually implicil their oscillatory dynamics, and we render such behaviors explicit via central flow analysis. Our analysis reveals precisely how these adaptive optimizers implicitly: (1) adapt their step size(s) to the local curvature, and (2) steer towards lower-curvature regions where they can take larger steps.
While we derive each central flow using informal mathematical reasoning, we show that these flows accurately predict long-term optimization trajectories variety of deep learning settings high standard of empirical proof. Thus, we believe that central flows hold promise as a framework for analyzing, reasoning about, and perhaps even inventing, deep leamnigg optimizers.
Our code can be found at:
httpsaubb.com//cccsaaaa//eteeeee flows
2 RELATED WORK
Edge of Stability The dynamics of optimization in deep learning remain poorly understood, even in the scemingly simple setting of deterministic (i.e. full-batch) training. Indeed, recent research showed that gradient descent on neural networks typically operates in a regime termed the "edge of stability" (EOS) in which (1) the largest Hessian eigenvalue equillibrates around the eritical threshold 2/17, and (2) the algorithm oscillates along high-curvature directions without diverging (Xing et al. 2018: Wu et al. 2018 Jastrzębski al.) 2019 2020 Cohen et al.! (02T). These dynamics could be explained by existing optimization theory, which led Cohen et al. (202T) to observe that there was no explanation for how or why gradient descent can function properly deep learning.
Subsequently, several studies sought to theoretically explain EOS dynamics. Some works analyzed EOS dynamics on specific objectives (Agarwala et al. (2023: Ahhn et al. [024] [Cen & Bruna (023)
Even et al. 2024] Kreisler et al. [2023) Song & Yun |2233; Li et al. 2022 Wu et al. 2024 Zhu et al. 2023), while other works (Arora et al. [0222]00 Lyu et al. 2022 Damian et al. 2023), gave generic analyses based on a local third-order Taylor expansion of the loss, which is one order higher than is normally used to analyze gradient descent. Similar arguments were first used by Blane et al. (0019) to study implicit regularization in SGD, Our work is most directly inspired by Daamian al. 2023 as their analysis applies to generic objective functions, and holds throughout training, not just near convergence. However, whereas they analyze the fine-grained oscillatory dynamics, argue that analyzing the time-averaged dynamics is simpler, and is sufficient for most purposes.
Prior works have proposed continuous time processes to approximate the trajectory of gradient descent (Barrett & Dherin 2021 Rosca et al. (023). SGD (Li et al. 2017)]
Continuous-time models
Smith al. [202T: Li et al.) (021). )) and adaptive optimizers (Malladi al. 2022] Cattanco al. 2024] Compagnoni et al. 2024). However, in the deterministic setting, these processes cannot model optimization trajectory over long timescales in the EOS regime, which our goal.
Understanding adaptive optimizers Ma et al. (2022) observed that RMSProp and Adam oscillate, and Cohen et al. (202)) showed that such dynamics can be viewed as an adaptive version of the edge of stability a finding which will leverage. KKhaled et al. (2023) and Mishkin et al. (2024) observed that on quadratic functions, certain adaptive optimizers adapt their effective step size to the maximum stable step size; we show this holds more generally, beyond quadratics. Experiments in Roulet
2

---

Published as a conference paper at ICLR 2025
as aililaation variance along each eigenvector
train loss
top Hessian eigenvalues
distance to GD trajectory
1780
undbable
undah
500
Lnddatee
incaate
eatel
Lnceatee
Zuresttle
<tabl
Luntbbbe
0.25
gradient
30
empiriall vartance predicted verience
central flow
gredient flom
0.26
grod ent flow
0.3-
0.22
(averaged)
200
0.2
central Nlo
100
2 30
0.3
gradient flom
0.18
0.0
20.00
3000
4000
5000
2000
3000
1000
3000
2000
3000
4000
5000
20.00
3000
4000
5000
step
stop
stop
stop
Figure 3: Central flow for gradient descent. A ViT is trained on CIFAR- using gradient descent with 7 2/200 (blue). Gradient descent enters EOS at step 3000 and after step 4500, multiple cigenvalues are unstable (dark gray). The central flow (black) accurately models gradient descent even at EOS, whereas gradient flow (red) follows different path. Our analysis also accurately predicts the covariance with which gradient descent oscillates around the central flow (third panel).
Note that the first panel plots the central flow's prediction for the time-averaged train loss (eq. (6)
al. (2024) and Waan et al. (2024d) can be explained by the phenomenon we call "acceleration via regularization." Many works have also conducted rigorous convergence analyses of adaptive optimizers, under various regularity assumptions on the objective (Duchi et al.) a [201T] Reddi et al.
2018:
Chen et al.,
2019a|b Zaheer et al. 2018;; Zou et al. 22099; Défossez et al. 2022) Li & Lin
2024| Chen et al.
2022, Wang et al. 2024a Yang et al. 2024] Guo et al. 2021 Shi et al. 2021
Zhang et al. 2022
Crawshaw et al.
2022. Li et al. 2024 Wang et al. 2024b
Hong & Lin 2024
Zhang et al.
2024
Wang et al.
2024c
Hübler et 2024
GRADIENT DESCENT
The simplest optimizer is deterministic gradient descent with fixed learning rate ry: Wh++1 = Wt - v/V ((w)).
(GD)
Cohen et al. (222)) showed that traditional optimization analyses cannot capture the typical dynamics of gradient descent deep learning. We now present an analysis that does capture these dynamics.
3.1 THE DYNAMICS OF GRADIENT DESCENT
To understand the dynamics of gradient descent in deep learning. it is instructive to first consider the case of quadratic objectives. On quadratics, gradient descent oscillates if the curvature (i.e. Hessian) is too large relative to the learning rate. For example, consider one-dimensional quadratic L(z)
Szz,, with global curvature S. The gradient descent iterates { evolve via Tt+1 = (1- (1-11)) )
If S exceeds the critical threshold 2/1, then (1 n)S) < -1, the iterate I, flips signs and grows in magnitude at each step, i.e. gradient descent oscillates with exponentially growing magnitude. More generally, on a quadratic in multiple dimensions, gradient descent oscillates with exponentially growing magnitude along any Hessian eigenvector with an eigenvalue exceeding 2/n-
While deep learning objectives are not globally quadratic, a local quadratic Taylor approximation suggests that if the Hessian H (w) has any cigenvalues exceeding 2/7, gradient descent will oscillate along the corresponding eigenvectors. This suggests that gradient descent cannot function properly if the sharpness S(w)= == x₁ (H (w)), defined as the largest Hessian eigenvalue, exceeds 2/n-
Why, then, does gradient descent converge in deep learning? The natural explanation is that the sharpness S(w) remains below 2/1 throughout training. Yet, Cohen et al. (202 empirically observed a very different reality. When training neural networks. the sharpness tends to rise (a still- still-unexplained phenomenon termed progressive sharpening), and usually reaches the critical threshold 2/1 during training. When this occurs, gradient descent oscillates along the top Hessian eigenvector(s), as expected from quadratic Taylor approximation. Yet, perhaps surprisingly, these oscillations do not cause divergence. Instead, gradient descent enters a regime termed the edge of stability (EOS) in which the sharpness equilibrates around the critical threshold 2/1,. the algorithm oscillates without diverging along the top Hessian eigenvector(s)) and the loss behaves non- monomoonttniaally. Cohen et al.) (2021) noted that these dynamics could not be explained by traditional optimization theory.
Damian et al. (203) showed that the key for understanding these EOS dynamics is to Taylor-expand the objective to third order, which is one order higher than traditionally used in analyses of gradient
3

---

Published as a conference paper at ICLR 2025
descent. Let us informally sketch this argument. Suppose that gradient descent is oscillating around reference point w. along the top Hessian eigenvector u, with current magnitude z, so that w = w + zu. Due to the oscillation, the optimizer follows the gradient at rather than the gradient at How do the two relate? A Taylor expansion of VI. around w yields: (see Lemma
V VL(w)= (1) gradient reference point
VL(U)) w)
r:Sww)e
2 VS
o(x²)
(1)
(2) oscillation
(3) sharpness reduction
The third term, which arises from the cubic term in the Taylor expansion of the loss, reveals that a gradient step on the loss with step size 7) automatically includes gradient step on the sharpness of the loss with step size 2 2 Thus, oscillations automatically trigger reduction of sharpness. This the crucial ingredient that is missing from traditional optimization theory.
In the special case where only the largest Hessian eigenvalue crosses 2/n (e.g. steps 3000-4500 in Figure 33, which was analyzed in Damian et al. (2023), the EOS dynamics consist of repeated cycles in which: (a) the sharpness rises above 2/ 2/17: (b) this triggers growing oscillations along the top Hessian eigenvector; (c) such oscillations reduce sharpness via eq. (T). pushing it below 2/ n/ (d) the oscillations consequently shrink in magnitude. In the more typical case when multiple Hessian eigenvalues have reached 2/1 (e.g. steps 4500-5500 in Figure B gradient descent oscillates along all the corresponding eigenvectors, causing all such eigenvalues stay regulated around 2/n-
While the fine-grained EOS dynamics are difficult to analyze, we will now show that a simple time- averaging argument allows us tractably characterize the macroscopic path taken by gradient descent. This analysis will recover the main result of Damian et al. (2023) (albeit non-rigorously) while also generalizing to the more realistic and challenging setting of multiple oscillating directions.
3.2 DERIVING THE GRADIENT DESCENT CENTRAL FLOW
So long as the sharpness S(w,) remains below the critical threshold 2/n,, gradient descent said be stable. While gradient descent is stable, it does not exhibit sustained oscillations, and its trajectory is empirically -ell-approximateed by that of gradient flow! du = V (w)); however, once gradient di descent enters the EOS regime, its trajectory rapidly departs from that of gradient flow (Cohen et al., 2021). We will now derive a more general ODE, which we call a central flow, which models the time-averaged trajectory of gradient descent in both the stable and EOS regimes. We will derive the central flow using a heuristic time-averaging argument, and we will empirically show that it can accurately predict gradient descent trajectories in variety of neural network settings.
We abuse notation and use E to denote "local time-averages" of deterministic quantities see Appendix A.2 The central flow is intended to model the time-averaged trajectory B[wr]. To simplify notation, we will also use Wú == E[w: to denote this time-averaged trajectory.
3.2. THE SPECIAL CASE OF ONE UNSTABLE EIGENVALUE
We will introduce time-averaging methodology by analyzing the special case when only the largest Hessian eigenvalue has crossed the critical threshold 2/1,. this setting, gradient descent oscillates along a single direction the top Hessian eigenvector. We will therefore model the GD trajectory by wt = w; + Uy where w; is the gradient descent iterate, W; is the time-averaged iterate, 4e is the top Hessian eigenvector We, and x, denotes the displacement between W, and W, along the de direction. Note that by definition, E[z,] = 0, i.c. the time-averaged displacement is zero. To track the evolution of w,, we time-average both sides of the gradient descent update and Taylor expand the time-averaged gradien using eq. (1):
Wr++1 We E. L(w)) RJ Tt - VV + S(m,, Biuz) ur + E[a ((TT))
Note that the second term in the time-averaged gradient is 0 because E[r,]] 0. Therefore, we model the time-averaged iterates W. by the sharpness-penalizee gradient flow wt) defined by:
dw dt
7 VL(w) + (t)
implicit sharpocss penalty
Here, (t) is a still-unkoown quantity intended to model E[z²], the instantaneous variance of the oscillations time f. This quantity also controls the strength of the implicit sharpness penalty. To
4

---

Published as a conference paper at ICLR 2025
determine (t), argue that only one value is consistent with empirical observation. Empirically, the sharpness stays dynamically regulated around 2/1.. Therefore, we will enforce that the central flow keeps the sharpness fixed at S(w(tt) = 2/n-
The time derivative of the sharpness under eq. can be easily computed using the chain rule:
(S(w)
dw
(V 2 (0)
(3)
dt
dt
progreasive shappning
sharpness reduction from oscillations
Solving for dS(w) de = 0 yields the unique o (t) that keeps the sharpness fixed in place:
o² (²(() (VS(ww) ( -V(W))
(4)
|]S(w)]]
Intuitively, this is the unique (1) for which the downward force of oscillation -nnduced sharpness reduction "cancels out" the upwards force of progressive sharpening so the sharpness is locked at //y, The central flow for single unstable eigenvalue is given by substituting this o (t( into eq. 2
3.2.2 THE GENERAL CASE (MULTIPLE UNSTABLE EIGENVALUES)
We now generalize this analysis to the setting where multiple eigenvalues have reached the critical threshold 2/ and gradient descent oscillates in the span of the corresponding eigenvectors. We assume the displacement ó, == Wt wú between the true process and the time-averaged process lies in the span of these eigenvectors. Note that by definition, E[,,] = because Wi == E. By taking we obtain a similar ansatz for the central flow:
a cubic Taylor expansion similar to Lemma [
dw dt
VL(w) + V (HT)),,(())
implicit curvature pemalty
Here, E(t) is still-unknown quantity intended model EE (8,66)), the instantaneous covariance of the oscillations at time t. This matrix also controls the strength and direction of the implicit curvature penalty, As above, we will argue that only one value of (t) is consistent with empirical observation. First, since Hessian eigenvalues that have reached 2/1 empirically do not rise further, we impose that the flow eq. (5) should not increase any Hessian eigenvalues beyond 2/n. Second, since gradient descent oscillates within the span of the unstable eigenvectors, we impose that Σ(t), which models the covariance of these oscillations, should be supported within the span of the Hessian eigenvectors with eigenvalue 2/n. Third, since (() models covariance matrix, we impose that (() should be PSD. These three conditions turn out to imply a unique value of (t). In particular, we detail in Appendix A..e that that (t) must be the unique solution
cone type of convex program known
complementarity problem (CCP). The central flow for gradient descent is defined as eq. (5) with this value of E(t) (Definition 2). This central flow can also be interpreted
projected gradient flow on the constrained problem: min,, L(w) such that Sww) S 2/1,, as we detail in Appendix A.3.3 Note that the central flow automatically reduces to the gradient flow when S(w) 2/1/
3.2.3 EXPERIMENTAL VERIFICATION
Although this derivation employs informal mathematical reasoning, our experiments demonstrate that this central flow can successfully predict long-term optimization trajectories on a variety of neural networks with a high degree of numerical accuracy. For example, in Figure we run gradient descent side by side with the central flow, and monitor the distance between the two trajectories in weight space. Observe that the distance between these two trajectories stays small, whereas the distance between the gradient descent and gradient flow trajectories grows large. Moreover, observe that the central flow's Et) accurately predicts the covariance with which gradient descent oscillates around the central flow; in particular, each eigenvalue of (t) accurately predict the instantancous variance of oscillations along the corresponding eigenvector of (t) (see Figure [] for more explanation).
Appendices B and C.I show our full set of experiments, and Appendix B.4 discusses the factors that affect the quality of the central flow approximation. A promising direction for future research is to rigorously identify the conditions under which gradient descent follows the central flow.
3.3 INTERPRETING GRADIENT DESCENT VIA ITS CENTRAL FLOW
As smooth flow, the central flow an simpler object to reason about than the oscillatory gradient descent trajectory. For example, the training loss L (w,) along the gradient descent trajectory behaves non-monotonically, as shown in Figure
5

---

Published as a conference paper at ICLR 2025
train loss
This led Cohen et al. (0221) to observe that there was no explanation
for why the loss still decreases over longer timescales. The central flow
0.25
L(w))
Etl (wi] L(w(t))
provides simple way to reason about this phenomenon Because the central flow is a smooth curve, the chain rule quantifies the rate of loss
predittion
decrease: di
dL(w)
(V)) duy dr Using this, we prove in Proposition
0.23
that dL(ww) dt I O, i.e that the central flow loss [(w(t)) is monotonically decreasing. Thus, the central flow loss is hidden progress metric for the optimization process.
0.22
0.21
As visible in Figure 4 the gradient descent train loss L(wx) is generally higher than the central flow train loss [(w(,)), due to the oscillations. However, since the central flow also models the oscillation covariance
5000
5500
6000
step
Figure 4: Train loss curve
(t), it can render predictions for the time-averaged train loss (Figure 4
B[L(w))] v L(w(t)) + (W(())), 20()) = -(w(tt) + ((((()))(((()))
(6)
central flow loss
effect of oscillations
Thus, the central flow allows us to decompose the time-averaged training loss curve into the loss along the central flow, which decreases monotonically, plus contribution from the oscillations.
SCALAR R.SPPROP
As stepping stone RMSProp, now study "Scalar RMSProp." simple adaptive optimizer which uses one global adaptive step size, rather than separate adaptive step sizes for each coordinate:
3 1 (1 32) || [L(w))] W+++
Voc VL(w).
(Scalar RMSProp)
The algorithm maintains an exponential moving average (EMA), v, of the squared gradient norm, and takes gradient steps of size VV, which we call the effective step size! Sii The hyperparameter |3
interpolates the algorithm from GD when B2 = 1 to normalized GD (NGD) when B3 = 0.0
Our analysis in this section will make precise how Scalar RMSProp adapts its effective step size to the local curvature, and will reveal that Scalar RMSProp's 's efficacy as an optimizer relies on hidden ability to implicitly steer towards low-curvature regions in which it can take larger steps.
4.1 THE DYNAMICS OF SCALAR RMSPROP
The dynamics of Scalar RMSProp revolve around the effective sharpness, defined as Sell == (w)) /VV. First, the effective sharpness controls the oscillations: when S > 2, Scalar RM- SProp oscillates with growing magniude along high curvature direction(s). Second, such oscillations in turn trigger a reduction of effective sharpness. This occurs via a combination of two distinct mechanisms. One mechanism, shared with gradient descent, is that oscillations implicitly reduce sharpness due to Equation (T), thereby decreasing the effective sharpness via its numerator. The other mechanism, new Scalar RMSProp. is that oscillations increase the gradient norm and hence v, thereby decreasing effective sharpness via its denominator. These dynamics give rise to a negative feedback loop that keeps the effective sharpness automatically regulated around 2, depicted in F The fine- grained dynamics are challenging to analyze, even in the case of a single oscillatory direction. Fortunately, we will see the time-averaged dynamics more tractable.
4.2 DERIVING THE CENTRAL FLOW
Recall that while gradient descent trains stably, it is well-approximated by gradient flow. One can derive an analogous "stable flow" for Scalar RMSProp (Ma et al. 2022 )
d w() - nh V (VL(((())
1v(1)
1-02 [[(((( v()].
32
However, at the edge of stability, the trajectory of Scalar RMSProp deviates from eq. 77 We will now derive a more general central flow that characterizes the time-averaged trajectory even at EOS. In the main text, we will focus on the case where one eigenvalue is unstable.
In section 1 we derived an approximation for the time-averaged gradient, [ Using the first two terms of eq. we can also derive time averagee approximation for E [
E [ [V 2 +
S ((w) 2122

---

Published as a conference paper at ICLR 2025
cillatioo variance along each eigenvector
train loss
top effective Hessian eigs
distance Scalar RMSProp
Lettabbe
cittee incable
cabe indstte Lnvtsbie
thee
Limbbbee
2upeaatee
0.26
RMSProp
pirpicall predicted verience
flows stable flo
centrall flo stable flow
statle flow
10
0.22
5.0x10
0.05
0.29
0.00
1500 2000 2500
1000 1500
2000
500
1000
2500
500
10000
2500
step
stop
stop
stop
Figure 5: Central flow for Scalar RMSProp. Mamba is trained on a synthetic sequence prediction task using Scalar RMSProp with 7) 2/ and 82 = 0.99 (blue). The central flow (black) accurately models the time-averaged trajectory of Scalar RMSProp even at the edge of stability, whereas the naive stable flow (red) follows different path. As with gradient descent, our analysis can predict the exact covariance E with which Scalar RMSProp oscillates around the central flow (third panel).
where we again used E[r] 0 to ignore the middle term. Based on these time averages, we make the ansatz that the joint dynamics of (w,, 14) follow central flow (w(t), v())) of the form:
duc dt
Vl(w(t)) + + (t)V((()))
dv
1-82 D2
|VL(w) S(w) -y] (8)
E[VL(w)]]
where (t) is still-unknown quantity intended to model E E[r,], the instantaneous variance of the oscillations. As in our analysis of gradient descent, there is a unique value of o²()) that maintains Sett(w v) 2. To compute it, we expand ds=
di using the chain rule: ds dt
a
du
25 de
du
df
de
Plugging in du: du from eq. ( shows that dt di dste that will ensure dt
d is linear in o2. Thus, there is unique value of o² de
= 0, which is given by:
progressive sharpening
effect of mean reversion on
32 LWw), VS(w)) +(1-32) ((()) ( [VL(w)]
d² (w; n, n.)
B2 |[V5(w)] -1-11,1 B2) S(w)²/n²
sharpness reduction
effect oscillation
The central flow for Scalar RMSProp with a single unstable eigenvalue is given by eq. (8) with this value of 0². In Appendix A.4 4 we extend this flow to the case of multiple oscillating directions. In
Figure 5) and Appendix C.2
we validate this flow empirically.
4.3 INTERPRETING SCALAR RMSPROP VIA ITS CENTRAL FLOW
We now interpret the Scalar RMSProp central flow to shed light the behavior of the algorithm and the function of its hyperparameters and 32. Because the dynamics usually transition from stable to EOS quite early in training, focus interpreting the central flow in the EOS regime ?
4.3.1 IMPLICIT STEP SIZE SELECTION
The central flow renders explicit the step size strategy that is implicit in the oscillatory dynamics of Scalar RMSProp. Recall that while the central flow is at EOS, the effective sharpness Seft := n)(w))/ //v/ fixed This condition can rearranged into statement about the effective step size: // VD 2ww( Notably, the value 2/5(w) the maximal locally stable step size for the current location w weight space. In other words, while the algorithm is EOS, the oscillatory dynamics continually adapt the effective step size to the current maximal stable step size (figure right). This is the precise sense which Scalar RMSProp "adapts" to the local loss landscape.
Effective step size
optimiee)
5.0044-
Stw00 (aking flow
5.0036
9.003
step
Notice that while Scalar RMSProp only accesses the loss function via first-order gradients, its oscillatory dynamics implicitly adapt the effective step size to the curvature.
7

---

Published as a conference paper at ICLR 2025
train loss L(wit))
sharpness S(wit))
effective step size: nlv 2//((w()))
200
150
3.04
ablation w/o
curvature reg
030
100
0,02
0.68
3.2
0.04
0.06
3.01
500
1600
1500
2002
1003 1500
2003
503
1500 2000
Figure 6: "Acceleration via regularization" for Scalar RMSProp. Starting from the same initial- ization, we run the Scalar RMSProp central flow two different learning rates (in blue and orange), as well as an ablated flow die we)
2 V (w) (in black) with curvature regularization removed.
These three flows all use the same step size strategy but differ in the strength of implicit curvature regularization. Initially (see insct), the flows with higher curvature regularization optimize slower;
however, over the longer run, they take larger steps and optimize faster (Mamba/ / MSE).
4.3.2 IMPLICIT CURVATURE REDUCTION
Understanding the implicit step size strategy employed by Scalar RMSProp is not sufficient to fully characterize the behavior of the algorithm. To do so, we need to return to the central flow, which additionally accounts for the curvature regularization induced by oscillations. In general, the Scalar RMSProp central flow is joint flow over (w, v). However, at EOS, because 17/ VV 2/S(w), eliminate from the expression for dw and write the central flow in terms of alone
dw dt etfective step
2
[L(w) + 00 (w: 7, S(w))
(10)
S(w)
implicit sharpness penalty
where (² 17, 32) is given by eq. (9) In other words, the time-averaged trajectory of Scalar RMSProp at EOS essentially equivalent that of the following sumpler-o-uunderstand algorithm: at each iteration, compute the sharpness S(w), and take a gradient step of size 2/S(w) on a sharpness- regularized objective, where the strength of the sharpness regularizer is given by eq. (9)
Notably, the hyperparameters 77, 32 are not used to determine the effective step size. Instead, their only role is modulate o², which controls the strength of the implicit sharpness penalty. The effect of the learning rate hyperparameter 7) is to monotonically increase o² indeed, the numerator of eq. is increasing 17 while the denominator decreasing in 1,, which implies the overall expression for is increasing in 7) " Meanwhile, the effect of the hyperparameter 32 is to monotonically interpolate o² between that of NGD when 82- 0 and that of gradient descent when 3g 100 The interpretations
of 1, 32 generalize to the setting of multiple oscillating directions, as detailed in Lemma
4.3.3 ACCELERATION VIA REGULARIZATION
To fully grasp the modus operandi of Scalar RMSProp. it is necessary to consider the link between step size adaptation and curvature regularization. By regularizing sharpness Sww), Scalar RMSProp is able to steer itself towards regions where the maximal locally stable step size of 2/S(w) is larger. In such regions, Scalar RMSProp can and does take larger steps. Thus, by regularizing sharpness, Scalar RMSProp enables larger steps later in training. We call this mechanism acceleration via regularization. Our experiments suggest that this mechanism is a critical component of the algorithm's 's effectiveness. In Figures 6 and 15 we compare the Scalar R.M.Prop central flow an ablated version which adapts the step size to (/5(w) but does not regularize sharpness. Over the long term, this ablated flow optimizes slower than the Scalar RMSProp central flow, because it traverses very sharp regions of weight space in which it is forced to take small steps.
The mechanism of "acceleration via regularization" is also key for understanding the function of the learning rate hyperparameter 7). We have seen that at BOS, the only direct effect of r on the central flow is to modulate the strength of sharpness regularization, with higher 77 inducing stronger sharpness regularization. Thus, counterintuitively. the instantaneous effect a higher r) is often to slow down optimization. However, 7 steers the trajectory into lower-sharpness. regions, in which Scalar RMSProp's 's effective step size will be larger, thereby tending to speed up optimization.
illustrate in Figures 6 and 15] over longer timescales, higher
8

---

Published as a conference paper at ICLR 2025
oscillation variance along each eigenvector
train loss
top effective Hessian eigs stab
distance to RMSProp
edge crstatiity
edit of eaaiity
edge stabilly
0.30
RMSProp flow sta ble flew
central
predicted variance
sbable flow
2
6.15
2.0x
6.10
500
750
1900
1250
1500
500
1000
1500
1000 1250
1300
750
1000 1250
1500
stop
stop
stop
stop
Figure Central Flow for RMSProp. ResNet is trained on a subset of CIFAR-10 using RMSProp with 77 2x 1x100-5 and B2 0.99 (bluc). The central flow (black) accurately models the aver-aee trajectory of RMSProp even at EOS, whereas the naive stable flow (red) follows different path. As for gradient descent and Scalar RMSProp, we are able to predict the exact covariance E with which RMSProp oscillates around the central flow (third panel).
RMSPROP
We now study RMSProp TTieleman Hinton 2012), which maintains an EMA v of the elementwise squared gradients V (w)) and uses per-coonlinate effective step sizes of n// VV E
Vt 32 32) VLW,)
W++1 = We
VL(w))
(RMSProp)
where @ represents the entrywise product. It is useful to view RMSProp as preconditioned gradient descent 11+1 w; P VV(ww, with the dynamic preconditioner P. diag( (Vi/)))
In this section, will show that RMSProp's preconditioner is implicitly determined by the al- gorithm'ss 's oscillatory dynamics, and we will make this preconditioner explicit for the first time. Interpreting this preconditioner sheds light on its efficacy, while also revealing potential direction for future improvement. Finally, we will show that adaptive preconditioning is not the full story: RMSProp also implicitly regularizes curvature, and this behavior important for its success.
5.1 THE DYNAMICS OF RMSPROP; DERIVING A CENTRAL FLOW
The dynamics of RMSProp revolve around the effective sharpness S (w: v/. defined as the largest eigenvalue of the preconditioned Hessian p-1 where P diag( vv/n) [4] When Sell the iterates oscillate along the top right eigenvector of the preconditioned Hessian. In turn, such oscillations reduce Seef via a combination of two mechanisms: (1) they implicitly reduce curvature, and (2) they increase the gradient, growing v and therefore P. The net effect is that the effective sharpness Sreff stays regulated around 2 throughout training (Cohen al. 2022
In Appendix A.5
we derive a central flow (wr(), v()), in the same way above, which models the time-averaged trajectory of RMSProp. We verify its accuracy in Figure and Appendix C.3
5.2
INTERPRETING RMSPROP VIA ITS CENTRAL FLOW
We now interpret the RMSProp central flow to understand the behavior of R.MSProp.. Because the dynamics usually transition from stable to EOS early in training, we focus on the EOS regime.
5.2. THE STATIONARY PRECONDITIONER
The central flow for RMSProp is harder to interpret than that for Scalar RMSProp. because even at EOS, v cannot be expressed as a closed-form function of and instead remains an independent variable. Nevertheless, it turns out that, in some circumstances, v implicitly converges under the dynamics of RMSProp to value that depends on the current alone. Namely, imagine holding the weights fixed at some value w, and letting v "catch up." What value would z converge to? We show in Proposition] that for any w, there is unique / that satisfies the stationarity condition du 0. We call this the stationary y/ for the weights w, denoted as v(w). Empirically. we observe that vtt) usually converges its stationary value i/ (()( during training (Figure Appendix B.3)

---

Published as a conference paper at ICLR 2025
In Proposition 2 we show that the corresponding stationary preconditioner P(w) diag( (v is, remarkably, the optimal solution to a convex optimization problem over preconditioners:
Ff(w/m)
P(w)
arg min dingonal, P
tx(P) + +
(L(w))
such that
(ww) 2P
(11)
optimization speed
local stability
That is, RMSProp implicitly solves the convex program eq. 11 to compute its preconditioner. This is the precise sense which RMSProp "adapts" to the local loss landscape.
We can now understand RMSProp's 's preconditioning strat- egy by interpreting the optimization problem eq. (TT). The constraint H (w) 2P is equivalent to Self < 2 and hence stipulates that the preconditioner P should keep RMSProp locally stable. The first term of the objective, tr(P), is the sum of the inverse effective step sizes. If this were the only term in the objective, RMSProp's 's preconditioning strategy could be summarized as maximizing the harmonic mean of the effective step sizes while maintaining local stabil- ity sensible preconditioning strategy. In fact eq. (II) would reduce to the dual of the max-cut SDP (Goemans
Coordinates of
10-1
10-4
10:
10 10
10-12
1000
1500
2000
2500
step
& Williamson 1995).
However, matters are complicated by the presence of the second term in the eq. [TI objective. The quantity (VL(w)ll )| 1 is the instantaneous rate of loss decrease under preconditioned gradient flow with preconditioner P. Minimizing this term actually slows down optimization (Figures 24 and 25 Therefore, the presence of this term implies that RMSProp's 's preconditioning strategy is not optimal.
Figure 8: Across 10 coordinates (col- ors), the EMA v(t;, (dots) matches its stationary value ((w(()); (dashed).
5.2.2 IMPLICIT CURVATURE REDUCTION AND ACCELERATION VIA REGULARIZATION
Understanding the preconditioning strategy implicitly employed by RMSProp is not sufficient fully characterize the behavior of the algorithm. To do so, we need to return to the central flow, which additionally accounts for the implicit curvature reduction induced by the oscillations. Substituting P into the RMSProp central flow, we can obtain a stationary flow in terms of alone, which assumes that the preconditioner P is always fixed at its stationary value P (eq. dw di
- -P(w) V((w)
V (2, (WW)) }
(12)
stationary preconditioner
implicit curvature penalty
where Σ E(w; 77; 82) is defined as the solution to a cone complementarity problem (eq. 24p. Figure 26 and Figure 27 show that this stationary flow can accurately predict the instantancous optimization speed of the central flow, given only access to w and not vtt).
Empirically, we find that the implicit curvature reduction effect is beneficial for the efficacy of the optimizer. In Figure 19, we compare the central flow against an ablated flow that leaves out the implicit curvature penalty. We find that in the long run, this ablated flow navigates into sharper regions in which takes smaller steps, and optimizes slower.
CONCLUSION
In this paper, we have developed a methodology for analyzing deep learning optimizers. To analyze an optimization algorithm, we derive central flow which models the optimizer's 's time-averaged tra- jectory, rendering explicit what was implicit in the oscillatory dynamics. We have empirically shown that these central flows can accurately predict long-term optimization trajectories of neural networks, and by interpreting these flows we have obtained new insights about optimizers' behavior.
These advances are made possible by the fact that we adopt different goals from most works in optimization. Rather than try to characterize global convergence rates, we set ourselves the more modest goal of characterizing the local optimization dynamics throughout training. The local dynamics are important, they are more interesting than may have been assumed (even vanilla gradient descent gives rise to rich, complex dynamics), and they empirically consistent across different deep learning settings, which suggests that a general theory is feasible. We believe that similar analyses can be fruitfully conducted for other optimizers, and we hope to inspire work in that direction.
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGEMENTS
AD acknowledges support from an NSF Graduate Research Fellowship and a Jane Street Graduate Research Fellowship. JDL acknowledges support of Open Philanthropy, NSF IIS 2107304, NSF CCF 2212262, NSF CAREER Award 21 2144994, and NSF CCF 2019844. AT acknowledges support from National Science Foundation grants 11S1705 121, IIS 838017, IIS2046613, IIS21 and funding from Meta, Morgan Stanley, Amazon, Google, and Scribe. JC would like to thank Jim and Marilyn Simons for their support of basic research via the Flatiron Institute. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of any of these funding agencies.
The authors are grateful for extensive feedback from Zhili Feng, Nikhil Ghosh, Yiding Jiang, and Sam Sokota.
Reproducibility statement All code is available at central flows
Ethics statement This paper presents work whose goal is to advance the field of machine learning. As such, there are many potential societal consequences of our work.
BIBLIOGRAPHY
Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. In Proceedings of the 40th International Conference on Machine Learning. [CML'23, 2023.
K Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via edge of stability. Advances in Neural Information Processing Systems, 36, 2024.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In International Conference on Machine Learning, pp. 948-1024. PMLR, 2022.
David Barrett and Benoit Dherin. Implicit gradient regularization. In International Confer- ence Learning Representations, 2021. URL /ttpttss////opeeeiieee www.neee/eeee www.eneee//orimeeemm 3q5IqUrkeF
Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov. Better plain vit baselines for imagenet- 11, 2022.
Guy Blane, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural net- works driven by an omsteein-uhlenbeck like process. In Annual Conference Computationnl Learning Theory, 2019. URL /ttps://ap..semart http://cchc.aaaoaoooooooooaaaa//
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miine and JAX: composable transformations of Python+NumPy programs, 2018. URL
Qiao Zhang.
http:// /ttp:///gtthb..... com/ com/jax m1/jax
Samuel Burer and Renato DC Monteiro. Local minima and convergence low-rank semidefinite programming. Mathematical programming, 103(3):427-444, 2005.
Matias D. Cattaneo, Jason Matthew Klusowski, and Boris Shigida. On the implicit bias of Adam. In Proceedings of the 41st International Conference on Machine Learning, 2024.
Congliang Chen, Li Shen, Fangyu Zou, and Wei Liu. Towards practical adam: Non-convexity. convergence theory, and mini-batch acceleration. Journal of Machine Learning Research, 23(229); 1-47, 2022.
Lei Chen and Joan Bruna. Beyond the edge stability via two-step gradient updates. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
11

---

Published as a conference paper at ICLR 2025
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. In Internationall Conference on Learning Representations, 2019a. /tLpsttss://pp/eee//// her.rreeee...ttttoooreooo-l- x309tm
Zaiyi Chen, Zhuoning Yuan, Jinfeng Yi, Bowen Zhou, Enhong Chen, and Tianbao Yang. Universal
stagewise learning for non-convex problems with convergence averaged solutions. In Interna- tional Conference on Learning Representations, 2019b. htR.sttts: /oppn//e////eevvv ew net/
Forum?id= 2orumm?.d==Sx552cccF
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representations, 2021. URL Stt.httts:/// httllhttps://peneeeeiee hhe
Jeremy M. Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E. Dahl, and Justin Gilmer. Adaptive gradient methods the edge of stability, 2022.
Enea Monzio Compagnoni, Tianlin Liu, Rustem Islamov, Frank Norbert Proske, Antonio Orvieto, and Aurelien Lucchi. Adaptive methods through the lens of sdes: Theoretical insights the role of noise. arxiv preprint arXiv:2411.1595. 2024.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems, 35:9955-9968, 2022.
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gra- dient descent the edge of stability. In The Eleventh International Conference Learning Representations, 2023, URL htt.tts://peeeevieew.nn net wtwp://www.n....oooooooooottttte
Yann Dauphin, Atish Agarwala, and Hossein Mobahi. How hessian structure explains mysteries in sharpness regularization. Advances in neural information processing systems, 37, 2024.
Marcel K. de Carli Silva and Levent Tunçel. Strict complementarity in maxcut sdp, 2018. URL wttp:////////./ooe.oo.oo
Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL httt::tt/tponrreeeeeeee mert/ooum?? wttoooommmmiiiioooaoao
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s)gd over diagonal linear networks: implicit bias, large stepsizes and edge of stability. In Proceedings of the 37th International Conference Neural Information Processing Systems, NIPS '23, 2024.
Michel X Goemans and David P Williamson. Improved approximation algorithms for maximum cut and satisfiability problems using semidefinite programming. Journal of the ACM (JACM), 42(6): 1115-1145, 1995.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024. URL htt.h::://pppeeeeeeee.e ew. net. forum? d= EEYskw VY2
Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianba Yang. A novel convergence analysis for algorithms of the adam family. arXiv preprint urXiv:2//2.0344 2021.
12

---

Published as a conference paper at ICLR 2025
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 1997.
Yusu Hong and Junhong Lin. On convergence of adam for stochastic optimization under relaxed assumptions. arXiv preprint arXiv:2402.03962 2024.
Florian Hübler, Junchi Yang, Xiang Li, and Niao He. Parameter-agnostic optimization under relaxed
smoothness. In International Conference on Artificial Intelligence and Statistics, pp. 4861-4864 PMLR, 2024.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. pmlr, 2015.
Stanislaw Jastrzębski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho and Krzysztof Geras The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2020. URL https:/7
4C44WBB
Stanistaw Jastrzebski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey, On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019. URL /ttp::://///oeerreeeeeee
het//oozun.1d--SkgE..
Andrej Karpathy. mingpt- demo.ipynb paster/demo... ipynb 2020.
Ahmed Khaled, Konstantin Mishchenko, and Chi Jin. Dowg unleashed: An efficient universal parameter-free gradient descent method. Advances in Neural Information Processing Systems, 36: 6748-6769, 2023.
Andrew V Knyazev. Toward the optimal preconditioned eigensolver: Locally optimal block pre- conditioned conjugate gradient method. SIAM Journal Scientific Computing, 23(2)5517-541, 2001.
Itai Kreisler, Mor Shpigel Nacson, Daniel Soudry, and Yair Carmon. Gradient descent monotonically decreases the sharpness of gradient flow solutions in scalar networks and beyond. In International Conference on Machine Learning, pp. 17684-17744. PMLR, 2023.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of adam under relaxed assumptions. Advances in Neural Information Processing Systems, 36, 2024.
Huan Li and Zhouchen Lin. On the O(vd/t!/4) convergence rate of RMSProp and its momen- tum extension measured by 4 norm: Better dependence on the dimension. arXiv preprint 0rXiv:2402.00389 2024.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference Machine Learning, 2017.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic dif- ferential equations (sdes). In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 12712-12725. Curran Associates, Inc., 2021. URL
Siles/paper/20221
pdf
Zhouzi Li. Zixuan Wang, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability. In Neural Information Processing Systems, 2022.
13

---

Published as a conference paper at ICLR 2025
LucidRains. Vision transformer pytorch. vit-pytores 2024.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35: 34689-34708, 2022.
Chao Ma, Lei Wu, and E Weinan. A qualitative study of the dynamic behavior for adaptive gradient algorithms. In Mathematiaal and Scientific Machine Learning, pp. 671-692. PMLR, 2022.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. In S. Koyejo, hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35. pp. 7697-7711. Curran Associates, Inc., 2022. URL https::/pproeeednnggg neurips www.eeeerrreerrprreeoe wc/paper Files//aaper20222 fi F11e2222c21000222222000626666663 orr.://00.....a0.oo44oooooeeerrr
Mo-
Aaron Mishkin, Ahmed Khaled, Yuanhao Wang, Aaron Defazio, and Robert M Gower. Directional smoothness and gradient methods: Convergence and adaptivity. Advances in neural information processing systems, 2024.
Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th International Conference on Machine Learning (ICML), 2010.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018. URL httl:its:/opperrviiieee
Mihaela Rosca, Yan Wu, Chongli Qin, and Benoit Dherin. On continuous time model of gradient descent dynamics and instability in deep learning, 2023.
Vincent Roulet, Atish Agarwala, Jean-Bastien Grill, Grzegorz Swirszez, Mathieu Blondel, and Fabian Pedregosa. Stepping on the edge: Curvature aware learning rate tuners. arXiv preprint arXiv:2407.06183, 2024.
Naichen Shi, Dawei Li, Mingyi Hong, and Ruoyu Sun. RMSprop converges with proper hyper- parameter. In International Conference Learning Representations, 2021. URL https
//openreview het//oorumidd=-DDddyy
Samuel L. Smith, Benoit Dherin, David G. T. Barrett, and Soham De. On the origin of implicit regularization stochastic gradient descent, 2021.
Minhak Song and Chulhee Yun. Trajectory alignment: Understanding the edge of stability phc-
nomenon via bifurcation theory. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL http:/::////eeeeeeeee.eeeee/ httpapaeeesssa..e..e..ooooeeeeaa httt://oorumm10ooPJJAAA000..
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26, 2012.
Alexandre Torres-Leguuet. mamba.py. httpp:////tttthoo.o///////a/////
2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen. Closing the gap between the upper bound and lower bound of adam's iteration complexity. Advances in Neural Information Processing Systems. 36, 2024a.
Bohan Wang, Huishuai Zhang, Qi Meng, Ruoyu Sun, Zhi-Ming Ma. and Wei Chen. On the convergence of adam under non-uniform smoothness: Separability from sgdm and beyond. arXiv preprint arXiv:2403.1 /5/46, 2024b.
14

---

Published as a conference paper at ICLR 2025
Bohan Wang. Yushun Zhang. Huishuai Zhang, Qi Meng. Ruoyu Sun, Zhi-Ming Ma, Tie- Yan Liu, Zhi-Quan Luo, and Wei Chen. Provable adaptivity of adam under non-unifom smoothness. In Proceedings of the 30th ACM SIGKDD Conference Knowledge Discovery and Data Mining. Association for Computing Machinery, 2024c. ISBN 9798400704901.
Mingze Wang, Jinbo Wang, Haotian He, Zilin Wang, Guanhua Huang, Feiyu Xiong, Zhiyu Li, Weinan E. and Lei Wu. Improving generalization and convergence by enhancing implicit regularization, 2024d. URL http:htsss//aaii..o.////o//////p http://rrg..../......../ 2405. 20763
Jingfeng Wu, Vladimir Braverman, and Jason D Lee. Implicit bias of gradient descent for logistic regression at the edge of stability, Advances in Neural Information Processing Systems, 36, 2024.
Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in prratttraeteri learning: A dynamical stability perspective. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL /ttpps:/prrceeedingg http://owwoeeii....e.reeerereeee neur ccpppppppppcccppprpppppppppp
-aieeeee---ppprrr-ee pdf
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018.
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd, 2018.
Junchi Yang, Xiang Li, Ilyas Fatkhullin, and Niao He. Two sides of one coin: the limits of untuned sgd and the power of adaptive methods. Advances in Neural Information Processing Systems, 36, 2024.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiy Kumar. Adaptive methods for nonconvex optimization. Advances neural information processing systems, 31, 2018.
Qi Zhang, Yi Zhou, and Shaofeng Zou. Convergence guarantees for rmsprop and adam in generalized- smooth non-convex optimization with affine noise variance. arXiv preprint MrXiv:2404.01436. 2024.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification on update rules. Advances in neural information processing systems, 35: 28389, 2022.
Xingyu Zhu Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of- stability training dynamics with a minimalist example. In The Eleventh International Confer- ence on Learning Representations, 2023. URL btttp:://ppenreviieee 77EagBsMAEO,
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for conver- gences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 11127-11111 2019.
NOTES
Barrett & Dherin 2021 argued that the accuracy of the gradient flow approximation can be improved by adding penalty on the squared gradient norm account for discretization error. However, their modified flow does not the stable regime, we found that the improvement
work (and is not intended work) in EOS regime, in accuracy it brings is small. Therefore, for simplicity, we leave out any such term from our flows.
2. We fold 1y into the definition of gradient flow that there correspondence between step of gradient descent and time of gradient flow. This will especially be useful when analyzing adaptive optimizers where the effective learning rate is changing.
3. In taking the aime-average of eq. assume that the eigenvector changes slowly relative the displacement x so that E[zru.] B[xx) [u..
4. Our explanation generalizes the explanation in Damian al. 2023). which applied only to the special case when only the top Hessian eigenvalue has reached 2/n.
15

---

Published as a conference paper at ICLR 2025
5. The terms "learning rate" and "step size" usually interchangeable. this paper, to avoid ambiguity, will use the phrase "learning rate" to denote the hyperparameter, and "step size" "effective step size" denote the actual step sizes that are taken.
Note that have re- re-indexed compared the standard definition of RMSPPrp (i.e. M441 vt). This does affect the trajectory and just ensures the effective learning step is determined by ve, rather than Pt+1> which simplifies the analysis.
7. When S2 Scalar RMSProp reduces gradient descent with learning rate n/ /Pu. Conversely, when v.((w)) 82 it reduces NGD with learning vy: wc+1 " [VL(W))]
The B2
correction is necessary for small values of Sa. For example, when ,32 (i.c. NGD), D2
| (wwe) Il so in the continuous time ODE, v(t) needs to adapt "instantly" to [V See I Lemma for additional justification for this correction term.
the stable regime (S* 2//n),, the central flow is given by the stable flow eq. for this flow is directly proportional to that of gradient flow, implying that this flow traverses the gradient flow trajectory, just at different speed (i.e. with nonlincar time-rescaling).
7 Note that the value of du
de
10. Note that at EOS we can rearrange EOS condition as v = n/ (wr)²/ /4, which lets us write v as a function of w and eliminate v everywhere in the central flow. This was already used to derive the expression for o² in eq.
11. The formula for o² is somewhat opaque for general 82. However, is relatively simple in the special case of NGD (3g 0). In this case, eq. (9) reduces to oo o²(w;;) z 4
12. note that which these is larger situation dependent, so a² can either monotonically increasing or monotonically decreasing in 92. That said, because o (w: v,, 0) Rs // and o"(w; v,. 1) is independent of 1/ a general rule is that for small learning rates, o² is monotonically increasing in 82, while for large learning rates, o² is monotonically decreasing in 32.
13. Our analysis can accommodate hoth bias correction and an e-dampening (dividing by vv + rather than vD) which are used by Adam (see Appendix A.6). However, simplify exposition, the main text focuses on this simpler version of RMSProp.
14. Note that p- ( (w) is similar to p-1/2 (w)) which symmetric and therefore diagonalizable with real eigenvalues. Therefore P p- (ww) also has real eigenvalues.
15. Note that in this section we have absorbed the learning rate yj into the definition of P so the critical threshold
2 rather than 2/n. This simplifies the analysis.
16. This reflects the the actual value used by RMSProp depends on the historical trajectory.
for any w, there are potentially many values v could stabilize optimization,
16

---

Published as a conference paper at ICLR 2025
TABLE OF CONTENTS
Introduction
2
Related Work
3
Gradient Descent
3.1
The Dynamics Gradient Descent
3.2
Deriving the Gradient Descent Central Flow
3.2.1
The Special Case of One Unstable Eigenvalue
3.2.2
The General Case (Multiple Unstable Eigenvalues)
3.2.3
Experimental verification]
3.3
Interpreting Gradient Descent via its Central Flow
4
Scalar RMSProp
4.1
The Dynamics Scalar RMSProp
4.2
Deriving the Central Flow
4.3
Interpreting Scalar RMSProp via its Central Flow
4.3.1
Implicit step size selection
4.3.2
Implicit curvature reduction
4.3.3
Acceleration via regularization
5
RMSProp
5.1
The Dynamics of RMSProp: Deriving a Central Flow
5.2
Interpreting RMSProp its Central Flow
5.2.1
The stationary preconditioner
5.2.2
Implicit Curvature Reduction and Acceleration via Regularization
10
6
Conclusion
10
7
Acknowledgements
11
Bibliography
Table of Contents
17
Central Flow Derivations
19
A.
Tensor Notation
19
A.2
On Local Time Averaging
19
................................
A.3
Gradient Descent
19
A.3.1
The Differential Variational Inequality Formulation
20
A.3.2
The Cone Complementarity Problem Formulation)
20
A.3.3
The Projection Formulation
21
A.3.4
The Rate of Loss Decrease
22
A.4
Scalar RMSProp
22
A.4.1
The effect of the hyperparameters 52
23
A.4.2
The small 32 correction
25
A.5
RMSProp/
26
A.5.1
The RMSPrrp stationary
27
A.6
Arbitrary Preconditioned Methods
28
A.7
Cone Complementarity Problems
29
A.8
Discretizing the central flows]
30
................................
A.8.1
Discretizing the Gradient Descent Central Flow
30
A.8.2
Discretizing Preconditioned Optimizers
31
B
Experiments
33
B.1
Gradient descent experiments
34
B.2
Scalar RMSProp experiments
37
B.2.1 Acceleration via regulaniztion
39
B.3
RMSProp experiments
41
17

---

Published as a conference paper at ICLR 2025
B.3.1
Acceleration via regularization
43
B.3.2
Stationary preconditioner,
43
B44
Known failure modes for the central flows
53
B.5
Implementation details
55
8.6
Architecture details)
57
B 7
Dataset details
57
Bulk experimental data
57
Gradient Descent
58
Scalar RMSProp
78
RMSProp
97
18

---

Published as a conference paper at ICLR 2025
A CENTRAL FLOW DERIVATIONS
We will now derive the general central flows for gradient descent, Scalar RMSProp, and RMSProp. We begin by defining the necessary tensor notation.
A.1 TENSOR NOTATION
For a k-tensor T and a s-tensor A with s < k we define T[A] to the contraction of T with along its last indices, i.e.
(T(A))
Σ Th
(13)
We also define (T,T") for two k-tensors T, T' of the same shape by:
(T,T")
Σ Th Trm..d
(14)
We will also use Tuu,,... My] to denote T(u 2 60 uz].
A.2 ON LOCAL TIME AVERAGING
We intentionally not specialize a specific notion of "local time-average". The only properties of the local time-averaging operator E that we use are:
1. linearity, i.e. [ + 9) = [E[] + E[g] and E[cs] = eRfs )
for any constant c
2. the local time average constant e itself: EE[c]
3. in the EOS regime when the sharpness oscillates around 2/1, the time-average is coarse enough to smooth out these oscillations so that ((E(w,) = 2/11
When we empirically verify the accuracy of our central flows, we use Gaussian smoothing with standard deviation of 7 steps, i.e. we define
Elfi = Eosores
-
where
Cj == exp 2 2x
E.ar. Ct-s
A.3 GRADIENT DESCENT
We begin by proving the Taylor expansion used in Section 3.1 Lemma 1. Assume that w is such that the top eigenvalue of H (w)) has multiplicity 1 and let wu where u is the top eigenvector of H (w). Furthermore, assume that sup [VL((")]] < co for some open neighborhood B of Then,
V( = V V S(Tjru +
Proof. By Taylor's theorem,
V((w) ( + H(
[u, (zx²)
Because is eigenvector of H(w) with eigenvalue S(w). the second term can be simplified to S(w)ru. Finally, by Danskin's 's theorem (or cquivalently the standard formula for the derivative of an eigenvalue):
max "7 H()e V (n("HH(())) [||||==1
V
L
where u is the argmax of the second expression, i.e. the top eigenvector of the Hessian w
We now define some notation that will be necessary for the remainder the section. First, let (ww kerln () 2/] denote the critical subspace, i.e. the directions in weight space that at the edge stability. Let Sym(U(()) S Rdxe denote the subspace of symmetric matrices whose span is contained in Uww).
19

---

Published as a conference paper at ICLR 2025
A.3.1 THE DIFFERENTIAL VARIATIONAL INEQUALITY FORMULATION
As in the main text, we will assume that GD oscillates around the central flow w(t) with covariance ((t E Sy)(U(w)) so that it follows:
dw
di - -
+ (
(15)
To determine E. we impose three conditions for all times t:
non-negativity: As a covariance matrix, (t) is positive semidefinite, i.c. E(t) = 0
stability: The sharpness remains bounded by 2/n, i.c. ((((t)) = (2/n)4
complementarity: The span of E(t (i.e. the span of the oscillations) is contained within the critical subspace U' (w(()). Equivalently, E() E Sym(U (wr()))).
We say that (w((), E()) follow the GD central flow if they follow eq. (5) along with these condi- tions:
Definition 1 (GD Central Flow, DVI Formulation). We say that (t), [(t) follow the GD central flow if for almost all t they satisfy eq. (15) along with the conditions: ( 2 0, ni (w((t) I
21, and (tt E Sym(U(w))
Definition I is an example of differential variational inequality (DVI). It provides conditions that the central flow must satisfy but is not a-priori clear that such a flow exists or unique. To prove this, and to turn Definition [ into a form that we can simulate numerically, we will show that for almost all times I, E() solves low dimensional convex cone complementarity program (CCP) This CCP will be guaranteed to have a unique solution, which can be casily computed numerically.
A.3.2 THE CONE COMPLEMENTARIT PROBLEM FORMULATION
We will start with the DVI formulation, Definition] 1 and prove two additional conditions on E(() which will allow us to solve for Et).. For simplicity, we will look for a choice of E(t) right- continuous so that the change in the Hessian along the central flow ansatz eq. (15), restricted to the critical subspace. We will denote this quantity by (I) E (ym(U(w(()) This generalizes the computation of di in dt
can reason about the immediate future using the chain rule. First, we consider
Section
To compute H, we define the linear operator Ta Sym(U(w)) + Rd by
[ = V. (Hww,, Σ) A E Sy)(UUw))..
Intuitively, Ta will play a role analogous to that of ( in Section 3 The operator Ten Sym(U(w:) Rdd takes input dxd matrix Σ Sym(U(w)) and returns the gradient of (H (w), Σ). the W-weighted Hessian. Mcanwhile, its transpose T Hd + Sym(U(w)) takes as input direction, and returns the directional derivative of the Hessian, restricted to the critical subspace. Therefore, du =
-- V [ (ww) + mTT [-]] and i is given by:
de
dH
dw
i1(1) di
T
- [L(W)
TT.
dt
(8w))
=:a(w)
=;p(w)
To simplify notation, we define a(w), ((w) by:
a(w) : T [-VL E Sy)) B(w) = 1T E Sym(((w)) & Sym(U(w)),
so that A == {[=]. In order to maintain stability, i.e. H (2/n)T, we need ensure A < 0, i.e. a 3 (]. However, this condition only lower bounds Σ and does not yet uniquely determine it. To do so, we will need to differentiate the complementarity condition. Recall that H measures the change in the Hessian, restricted to the critical subspace. The directions in the kernel of H will remain at EOS. However, the curvature in the directions where 1t < O will drop below 2/n,, so these directions will leave the critical subspace. In other words, the critical subspace at time t e will be contained in [f) which is subspace of the critical subspace. Because we assumed Σ right-continuous, this implies that (t) must lie in this subspace, i.e. span [2]] ker kerl.)]. Equivalently, because Σ ž 0 and
i O, we can rewrite this condition as (E, H) = O. Together, these constraints:
E Σz 0, H 3,
(I, H 0. A [|
20

---

Published as a conference paper at ICLR 2025
define cone complementarity problem for Σ (Definition 8 We show in Lemma 5 that if / is symmetric and has full rank, the solution to this CCP is unique, and we denote it by CCP(a. 8). Otherwise it denotes the set of all solutions. We can now define the central flow in terms of this CCP:
Definition 2 (GD Central Flow, CCP Formulation). We say (120 follows the GD central flow if for almost all f 2 0, w satisfies eq. 15 with 2 E CCP(a(w), (())).
Definition 2 is the simplest version of the central flow to simulate numerically. By picking a basis for Sym(U(w)) which has dimension k(k+1) we can materialize the linear operator Tu as k(k+1) x d dimensional matrix, compute a ((w). (,,,. (w)), solve the low-dimensional CCP to compute Σ, and take a small Euler step the flow using this Σ.
In the next section, we will see that this CCP formulation can be reinterpreted as a projected gradient flow, which will help interpret the behavior of the gradient descent central flow.
A.3.3 THE PROJECTION FORMULATION
Let 8,, :- {w Sww) < 2/1) denote the stable set. In this section we will show that the gradient descent central flow can be reinterpreted as projected gradient flow constrained to this set. Definition 3 (GD Central Flow, Projection Formulation). We say that (w(t))22 follows the gradient descent central flow if for almost all 1,
dw
di = projss,, (w)l where =- S(w) S 2/n},
(16)
where projt:, (w) denotes the orthogonal projection onto the tangent cone of Sy at iterate wr i.e. the set of directions that, if followed, would not immediately increase the curvature in the critical subspace: (w)(w) {u E Rdd T < 0}.
The equivalence between Definition 3 and the CCP formulation Definition ia a consequence of this simple lemma:
Lemma 2. Let w E S.p Then
proin,,(w)ll = v- Tm where ΣECCP(TT [o], T Tm).
Tel [u] 3 0}. Therefore the Proof. Note that the tangent space of S,, is given by the set: {o=
projection is given by solving the quadratic program:
min |0 such that [e + 0] 30
20
The KKT conditions imply that there exists a such that:
8 == T
(2,
0]) =0, Eこ
which implies that Σ e CCP(TT 1 TT.
We now intuitively reconcile eq. (16) with the CCP definition:
When S(w) < 2/1,, w is in the interior of S., so (w) = 0, the tangent cone is the entire space, and the projection is the identity map. Therefore eq. (16) reduces to gradient flow.
When there is single eigenvalue 2/1,, w is the boundary of S. and the tangent cone is given by the halfspace: Ts, (w) = {E (VS(w),)) S 0}. If the negative gradient lies outside this halfspace (i.e. if gradient flow threatens increase the sharpness above 2/m). then the projection onto the halfspace is given by the projection onto the hyperplane: -n S(a)) VL(w). Otherwise, if the negative gradient already lies in the halfspace, the projection is the identity map, so the central flow follows gradient flow and leaves EOS.
In general, computing the orthogonal projection onto Ts,(w) = fo T [v] < 0} requires solving a semidefinite quadratic program for which Σ is the Lagrangian dual variable. The KKT conditions of this quadratic program are equivalent to the CCP that defines above.
21

---

Published as a conference paper at ICLR 2025
A.3.4 THE RATE OF Loss DECREASE
Proposition Under the GD central flow (definition By for almost all we have
dL(w) dt = -
dL S 0 (i.e. the loss monotonically decreases) and dL This implies at decreases at a slower rate than under the gradient flow).
dt 2 -nl VL(W)]l (i.e. the loss
Proof. By the chain rule we have
dL(w) dl
dw
((w)), di
VL(w), (w)) (wol
proim, (w)) VL(w)))
- [[] -
where we used that for any orthogonal projection proj[cu] = proj[c] when 2 to pull out the n and that (v, proj[v]> [[pro][u]]| to get the last equality. Finally, the comparison with gradient flow follows from the inequality [[pro][e]] S ||v|]] for any orthogonal projection.
A.4 SCALAR RMSPROP
follow similar derivation to gradient descent. We define U(w, v/) := and let Sym(U(( v)) denote the subspace of symmetric matrices with span in U@w, v). We now proceed with the time-averaging derivation.
- 2/],
If uw 8 with [0]
0, we have that
B [ L (T)
((v Σ)
(17)
where E[ expression is equal to
Because we assume that Σ Sym(U(T, v)), VLL(T)E S(T) so this
E VL ~ [V +
This suggests the central flow ansatz:
dw 17
dt Vv VL(w) + + L
(18)
dy 1-3 S(w)) =
dt 3,
We can now give the differential variational inequality definition for the Scalar RMSProp central flow:
Definition 4 (Scalar RMSProp Central Flow, Differential Variational Inequality Formulation). We say that (w((t, vt),, h(t)) satisfy the Scalar RMSProp central flow if they satisfy eq. (18) along with the conditions: E(t) z 0, (p(t))1/22 (1/2 = and Σ(t) E (ym(( v(t)) for almost all t.
As for gradient descent, this definition is fairly opaque and does not give a way of actually computing will again derive come-complementarrit problem formulation of the Scalar
E. To do RMSProp central flow. The first step in this derivation is to differentiate the stability constraint: mp 1/221 (w) x 2. We will fix a time and use U to denote U U(w, v) to simplify notation. We will also define the linear operator T SymnU v)) Rd by
T V.. (H(w), Σ} for EE Sym(U(w,
22

---

Published as a conference paper at ICLR 2025
(w) lo be the time-derivative of the preconditioned Hessian restricted to the critical subspace. In order to avoid violating the stability condition, we need to enforce H A 0. We can compute H using the chain rule.
Let H := == (0011/(H(w)))))))))
H = dt
1/2 H(v))
17
np 1/2
H(w)
v-12
+ dt
dt
1
1
ddw
dv
-12-1/27
2/D1
dt
2v3/2 dt
1 B2 1
7
[VL(w)]] -VW
32
We can now assume that (w, v) are at EOS so that == -²0((w)) Otherwise, complementarity forces O. Substituting this gives:
# 4 5(w)²
T
[T[VVI())
2 7
1 1 82
4
IVL(w)] ((w))
82
Sw)
We will now group the constant terms and the terms linear in E. Define:
B₂
1y S(w)
IVL(())))
alw)= = S(w)2 VL(w)] +
32
ST)
4
32 82
4 1 2 I.
B(w) = S(w)2
T
Then H = B|Σ]. As in the derivation for GD, if H is strictly negative definite in some direction, these directions will drop from the critical subspace right continuity forces E kerl [ This gives us the cone complementarity definition of the Scalar RMSProp central flow: Definition 5 (Scalar RMSProp Central Flow, CCP Formulation). We say that {ww(t), v(t) follow the Scalar RMSProp central flow if they satisfy eq.
(88) where 2 E CCP(((w), ((w)).
We can now use this formulation to prove eq. ( when U(w, v) is one dimensional. In the one dimensional case, a. 3 are both scalars and assuming that a(w) > O the solution to the CCP is simply = a (ww / In addition, in this one dimensional case we simply have VS(w). Therefore,
0² S(w)2 < VL(),, VS(w)) + B2
1-82
4
S(ww)
[V[(w)1
8(0)2
(19)
2
[V5(w)
+ 1-32
3(w)
82 (VVL(W) VS(w))
3z)
S(w)² [
4
(20)
+ V
(w))/m
A.4.1 THE EFFECT OF THE HYPERPARAMETERS TJ 8,3
wi() and define v(w; 11) nJ² S(w)² that Sel v) = 2. Lemma 3. Fix iterate Let U be the top eigenspace of H(w). Then if E(t) > O under the Scalar RMSProp central flow (Definition 5), we have that
0 dH =-clu where C
On dt lu and where Iy denotes the identity matrix the subspace U. In other words, larger learning rates n more aggressively decrease the curvature and they do so uniformly across all eigenvalues in the critical subspace U. In addition.
d.
-CTv
03, dt
1
23

---

Published as a conference paper at ICLR 2025
tr and C' S 0 otherwise. Therefore B2 can either monotonically increase or decrease the curvature regularization depending on whether gradient descent (32 = 1) or normalized gradient descent (32 0) would have a larger oscillation variance.
where C 2 0 when E
8==1
Proof. The condition {t) > implies that H(t) = O and Σ = a. We will rescale a, 3 by S(w)2/44
B2
(ww)
|V
â = 1- +
$2
4
r72
B TT + 1 - 32 S(w)²
8
82 1/2
We will also use aw Bn to refer to just the first terms in å, B:
āw V[ and
=== = = TT
Then Σ == 3- a. Differentiating this with respect to 17, 82 gives:
B2 1-1
Σ
[7]
[ + 5(w) tr
On
21ya B2
0 Σ
S(w)2
S(w)²
3 |M
tr Σ
032
32
172
Then:
On -
i dH
0
0
dH
8m
032 dt
032
Therefore suffices to compute 3 [ To compute this, we use the Sherman-Morrison formula.
1-0, S(m)²
Let
Then,
B2 7
1T, 1 =1 1
[
B [1]
1 cB. [1,1]
Finally, note that shown that
2 0, which immediately implies the result for 77. For 82, we have
1-c3 [I
dH -CIv
032 dt lu
where C' if and only if
tr M 2
VL(w)]I
S(w)²
To rewrite this in form that is independent of Σ, note that
tr L = tr 3-1a
(3 3 [1],
(o,
+ -
tr 3-
c [7,
n²
V
S(w)
1 cB [1,
Therefore tr Σ a weighted average between 3 aw
and x22
VL(w)²
so tr[2]] <
S(w)=
S(w)²
and only tr 3 áw just Lr Σ 0==0"
s n²
VL(W)² This first expression is just tr 1,=1 and the second is
S(w)²
The condition E(t) 0 is equivalent to requiring that not currently a "breakpoint," i... time at which an eigenvalue drops from EOS. The set of such constitute a measure zero set, so the above lemma holds for almost all t.
24

---

Published as a conference paper at ICLR 2025
A.4.2 THE SMALL B2 CORRECTION
The following lemma shows that the /3 + 32 correction for the Scalar RMSProp and RMSProp
central flows allows the continuous EMA to match the discrete EMA for linear targets f(t): Lemma 4. Let f R + R be continuous time process with | (1) S for all L. Then if
Vt- -2 + (1)
v1(1 1-32 [5(1) - v(O]
A3
we have that for all integers 1 2 O(11 - 12)
(v(() u1 0
A2A
(1-12))
Proof. First define
82 ('(1)
;= = - f(1)
1 82
32
8(1) = v(1) - 1(1)
(1).
1 B2
Then we have that:
1 32 (t())
6. 82 + (1 - (2) f(t)
82
824-11 - (2f(1)
1 - B2 (t)
B2
B2
828:-1 - Belf(t) f(t - 1)] 1 - B3
(t( + 1- B2
f((t 1)
133
328,-1 - 32(f(t) f(( - f(t 1)] 0
1 1 82
:28:-11 + ( 324
1 -2
Therefore,
8, 3300 +0 ( A.2
(1- B2)
((1-22)
Similarly,
1 82 (f(t) - v(t)] -f(())
82
8'(1) =
(*)
B2
1- B,
1 B2 ((t) + +0 32
B2
I- Br
SO
1-82 (2()
834
8(1) = -
8(0) +
(1-
Therefore subtracting the bounds on de and 8(t) gives:
321
1/ v()) S
(1 B2)
desired.
25

---

Published as a conference paper at ICLR 2025
A.5 RMSPROP
We first motivate the stability condition Amaz (P-1H)
quadratic. First, that P-1 is similar p-1/2Hp-1/2 which is symmetric so P-1 H has real eigenvalues and real eigenvectors. Next, consider gradient descent on the quadratic _ w Hw with preconditioner P. The update is:
w f- w w- W
w.
Therefore, = (-))-- If the eigenvalues of P-1H are in the range (0, 2), then - 0. Otherwise, it diverges along the corresponding right eigenvectors of P-1H. We can reinterpret this stability condition using the equivalent condition H = 2P. Note that the top right eigenvectors of P- correspond to the top eigenvectors for the generalized eigenvalue problem Hv APv.
We again follow a similar derivation to gradient descent. We define the critical subspace by U(w, v) == ker 2P] where
diag(v and use Sym(U (w, v)) to denote the set of symmetric matrices this subspace. We can now proceed with the time-averaging argument.
As before, we assume w = w + with E[6] = O. Rather than expanding the squared gradient norm, expand the element wise squared gradient:
[[V [I(w)02 ([(((W)))) 02
(21)
V 02 diag H (T) (T)
(22)
where Σ = ].] This can be further simplified using the fact that Σ is in the top eigenspace, so H(w)> = 2PE. Therefore,
4v diag [( 4P P 1
2 diag(E].
This suggests the central flow ansatz:
dw 7
VL(w)) [w)
dt VV
(23)
dy 1- $2
4v 17
82 VL(w)
+ diag(=((t] ]
dl B2
As for gradient descent and Scalar RMSProp, this immediately implies a differential variational inequality definition:
Definition (RMSProp Central Flow, Differential Variational Inequality Formulation). We say that (wt(t, v()) follow the RMSProp central flow if for almost all 2 O. they satisfy eq. (33
along with the conditions: (t) > 0, (w(()) = (t), and (H (w()) - 2P()), E{t) = where
P(t) = Vu(t)/n.
We will now show that Σ can be computed as the solution to cone complementarity problem. Fix some time f, let U = Uww v), and define the linear operator T Sym(U) - Rdd by T[=] = V.. (Hw),, Σ). We can differentiate the stability condition in the critical subspace as:
d (H(w) - 2P)
A = di
U
dw
2 a du
TT
dt
Ov dt 1s
T p 1 1 ((w)
T
1 - 82
1
4v c diag Σ VL(w)02 _
diag
82
nu
1/2
We again group the constant terms and the terms linear in E. Define:
1 1 82 p-
p2 diag V (w) 02
a(w, P) = :) VLw)] +
B2
112
:= 1 + 1 B2 4
B(w, pJ[E] =
P diagg(?)
32 n
26

---

Published as a conference paper at ICLR 2025
Note that chose to define 3 through its action Σ avoid unnecessarily complicated tensor notation. Then H 3 3/2.]. As in the derivations of gradient descent and Scalar RMSProp, require that 2 solves the CCP:
10 H =
E, ( =0
(24)
Together these define the CCP formulation, which can efficiently simulated: Definition 7 (RMSProp Central Flow, CCP Formulation). We say that {(n(tt) (())(())) 120 fol- low the RMSProp central flow if for almost all 2 0. they satisfy eq. (33 with {() E
CCP(a(w(t)) P(t)), ((w(t), P(t)) where P(t) = diagí /(t///l
A.5.1 THE RMSPROP STATIONARY
Setting dt
0 in eq. (3
gives:
((w)) 02 (4/n²)v @ E..
(25)
In addition, by complementarity we know that
[(H(w)- 2 diag, VD/nl) O. Proposition 2. If w.v satisfy eqs. 25 and (6). then P = diag( vv/n) minimizes the convex program eq.
(26)
Proof. Let vv/n. The convex program eq. 11) can be written as:
VL(w); such that H(w) = diag(p).
min Σ
P
(27)
P.
E the dual variable for the semidefinite constraint, the KKT conditions for this program are:
1 VL(W);
25 =0 Vi, ((H(w) - 2 diag(p)) 0.
(28)
1/2 p
We will prove that (p. 5) = (2 1 Σ) solve the KKT conditions eq. 28 Note that the comple-
mentary slackness KKT condition is equivalent to the complementarity condition on Σ. In addition,
dividing the stationarity condition for (eq. (25)) by v gives
VL(w) 32
(4/1))2 Vi
(29)
which is equivalent to the first condition in eq. (28) after substituting Σ -
= [() and v = 1rp
Proposition 3. For any L(w),, (w), the solution to eq.
111 is unique.
Proof. For notational simplicity, let g = V..(w) and let H Hw))
Assume there are two minimizers P, P' and let == diag( P), = diagg(P'-P Then by convexity,
diag/p + co] also minimizes eq. [11 for any e < 1. Therefore, differentiating the objective function in this direction gives:
Σ9,
1
112
Taking another derivative implies that:
Σ 9
62 =
This implies that 6; 0 in any direction where 9i + O. Let / be the set of indices for which gi * and for any vector p. let p denote the vector restricted to the indices in L. Define the linear map by
[2z iEl
gfvr]s =
pj i
27

---

Published as a conference paper at ICLR 2025
In other words, g takes reduced vector by and fills in the missing entries with p. Next, define the operator A by
d diagig(un) [/v]]] diag(er].
Then both pr, p/J minimize the following reduced SDP:
min Σ
such that 1 (Ww) @ 11 3 A²)
Pr
iE!
Now we apply de Carli Silva & Tunçel 20018 Proposition 1) with (A, 1n)) First, note that A(Ta+ (+ z|] = 21 [r] which satisfies the first condition. Next, for any y of O. we can take = = ly to satisfy the second condition, as in the proof of (de Carli Silva & Tunçel 2018
Corollary 2) Therefore O
pI p/. and we have already shown equality on 1 we must have p = p..
Numerically computing the RMSProp stationary v and preconditioner To numerically compute v(w) and/or P(w) at some weights we use a fixed point iteration to solve eq. 11 for P(w). Then, one can easily recover v(w) from P(w).
To solve eq- [II we apply fixed point iteration based on the fixed point equations for Σ, P. To derive our fixed point iteration, we parameterize E in the factorized form Σ = DD where D E Rdxrr (as in the Burer-Monteiro factorization (Burer Monteiro 2005 and r < d intended to upper bound the number of unstable eigenvalues. Then the formula for du
at reduces to:
(w + diag]//EI1) L(ww) + (H
where define H = (w) for notational simplicity. In addition, we will use that the span of Σ (and of D) lies in the critical subspace,
1/2 HD 2D. "
diag(v-1/21HD
We begin with random initial guess for D and iteratively update D, by:
V - VL(w) 62 + (HD)
7 diagly HDD
2
If this fixed point iteration converges, we have
VL(w) + 621 and diag(u- 10g]u-1/2||DD 2 1)
du so that dt 0 and the span of Σ = DD is in the critical subspace.
To verify the results of this fixed point iteration, we compare the objective values of the primal and dual programs. The dual program to eq. (IT) is:
max (E, H(w)) 2 [ww) VI- 2 diag Σ E2
such that
diag Σ = 1/2.
A.6 ARBITRARY PRECONDITIONED METHOMS
In this section, we derive central flow for an abstract preconditioned method. This general central flow will reduce to our central flows for gradient descent, Scalar RMSProp, and RMSProp. However, we emphasize that this paper does not claim that the central flow derived in this section will be correct for any preconditioned method. We include this section both because it allows us to easily generalize our central flows to minor variants of the same algorithms (e.g. gradient descent with a learning rate schedule, RMSProp with bias correction), and because we hope it can be a starting point for others to derive central flows for other optimizers.
We will let v be the "state" of the preconditioner, and let P be function that maps the state // to symmetric, positive definite matrix. Specifically, we consider the discrete update:
= Pt-1 + f wx), Wi+1 = wt - VL(w).
This formulation is very general and includes wide variety of updates including:
28

---

Published as a conference paper at ICLR 2025
GD with learning rate schedule n(t): Set f(v, w) 1 that v/t) and P(t) n(()-11
Vanilla RMSProp: Set f(v, w) (2-32)(VL(u))
v) and P(v) diagivp/m)
RMSProp with 6, bias correction, and learning rate schedule ry(t): Set v = [u, f([v, t), w) [[1 B2) VL))22 v). 1] and define
P([u, {) = diag
U
(t()) V
1 32
Note that this trick of embedding into the state variable v allows us to automatically derive central flows for any smooth hyperparameter schedules (e.g. r7(t,, (t),), e(t) as simple corollary.
As for RMSProp, the stability of this algorithm requires Amaa P H) S 2 or equivalently // - 2P. To derive the central flow, assume that w + å with [[ō] 0 and E[65"1 Σ. Taylor expanding and time averainggg the update for // gives:
EiS(U, w)] = [ VL(w + 5)) ~ f(U, VL(w)) + V (V,,
These motivate the central flow ansatz:
dw dt
P(()) (L(w) + (P()) [
(30)
dv f(u,w) + f [NB]
dt
We say that (wt,), (t), {2) satisty the DVI formulation of the central flow if for almost all t 2 0 they satisfy eq. (30). (1) 0, (w())) = 2P(x(t)). and (E(t), H (w=(t) 2P(((t))).
To derive the CCP formulation which is efficiently computable, we differentiate the stability condition. dH
Fix an iterate f, let U == kerlH (P(((t))) be the critical subspace at time t, and define H == at lo
under eq. (30). Then for the stability condition remain true, we need A = O. To compute it, define T Sym((U) + Rd by T [ = V. (Hw),, Σ). Then,
dH
A = dl
U
=TT dw
dv
22 P(U)
dt
dt
1
P('
- - + + w) w>|E|
Splitting up the constant terms and the terms linear in E, we define:
a(w,v) = [(][
27 2V, f(v, w) E Sym(U)
= P
w) Sym(U) & Sym(U).
Then H a a- 3[.) so by the same arguments as for gradient Scalla,,,Salar needs to solve CCP(a v), ((w. v)). and plugging this Σ into eq. (30) gives the CCP formulation for the central flow for this preconditioned method.
RMMProp.
A.7 CONE COMPLEMENTARITY PROBLEMS
In this section we will let be a positive integer, and we will use Sym, to denote the set of k x & symmetric matrices.
Definition 8 (Cone Complementarity Problem). Let E Sym,(R) and let 3 E Sym, (R)@Sym, be a linear operator on symmetric matrices. We say that the matrix X E Sym, (R) solves the cone complementarity problem CCP B)
X 7 0,
3[x (X,a B[X)) 0.
Lemma S. Let 3 E Sym, (R) @ Sym, (IR) be a symmetric linear operator on symmetric matrices.
1. If 3 Z 0 as an operator on Sym, (R), then CCP(a,, has solution for all a.
29

---

Published as a conference paper at ICLR 2025
2. If 3 > 0, this solution is unique. Otherwise, all solutions X, X' differ by matrix in the kernel of B.
Proof. Consider the quadratic program
min 3[ X] (X,a) X
such that X ビ 0.
(31)
Because 3 > O, this program is convex. The KKT conditions for this program are 3| a t 0 and the complementary slackness condition X [3[X] a) 0. Therefore any KKT point to the original quadratic program, including its global optimum, solve the conic complementarity problem which proves (1). To prove (2), let X' be solutions to the conic complementarity program. Then taking the trace of the complementarity relations we have:
(a, X) 3|X, X] O, (a.
0.
(32)
In addition, because both ][X] and X' are PSD and vice-versa, must have
(X' (X', [[x] a) 2 0, (x,
(33)
Adding these four conditions gives:
x|X, x 3| - 3 2 - B - "'] 5. (34)
However, since 3 > 0 this implies that 3[X - - = O. When B > 0 this implies that X X', and when this implies that X X' is in the kernel of
A.8 DISCRETIZING CENTRAL FLOWS
It is not immediately trivial discretize the central flows (Definitions and 7 in terms of the critical subpace, i.e. the subspace of (effective) Hessian eigenvectors with eigenvalues the critical threshold. However, when simulating these flows numerically, there never any eigenvalues "exactly" equal to the critical threshold, so a naive discretization is not possible.
These defined
Instead, we will describe an alternate Euler discretization scheme that recovers these central flows as the size of each discretization time step goes O. We will first discuss the of gradient descent as the intuition is clearest (Appendix A.8.I). D), and then we we will discuss the preconditioned methods (Appendix A.8.2).
A.8.1 DISCRETIZING THE GRADIENT DESCENT CENTRAL FLOW
Recall that the gradient descent central flow can be written as the unique trajectory (w(t), E(t) that satisfies: (Definition
1. dur 1 (VL(W + V (H(w), Σ)
oi
2. (ww)= 2/n1
3. Σ 0
4. Σ E kcrl (w)) 2/nT]
where we write and Σ rather than w(() and Σ((), for notational simplicity.
As discussed above, the last constraint creates challenges when simulating these flows in discrete time, as (ww will never have eigenvalues "exactly" equal 2/n on a computer. As a result, we will slightly relax these constraints to allow them to be practically simulated. We will fix discretization step size e and we will require:
1. Wh+t + Wi en VL(w.) + V (HW)) Σx)
2. H(w++)) 2/n
3. Σ, 0
4. Σ, (er[H((wtte) 2/nl].
30

---

Published as a conference paper at ICLR 2025
This is again not possible to directly discretize, since is challenging reason analytically about H (wi+e) where is the Hessian of the deep learning objective. Therefore, we linearize H (wr+ee) around Wi, and approximate H (We+e) as:
H (wi++c) Re (we)
en (WN + V mn (H(w),, Ex) [T
Hessian
change
Fuler step - t
We will refer to this approximation H' [E,, and crucially, this approximation is linear in E. We can therefore ask for trajectory {we, } that satisfies:
1. Whte < W/ en) VL(w) + 2 Vw, (H))) ((()))E)))))) Ex)
2. H [ 2/n1T
3. E =
4. Σg kcr(H" -
We can now use these constraints to derive CCP for E that can be solved numerically on a computer. First, note that Σ, is supported the span of the cigenvectors that will be the edge of stability in the next Fuler step. To reduce the search space, we will track all eigenvectors of (w;) above some fixed threshold (e.g. 1.95/1)). as these have the potential become unstable after the next Buler step. We will store these cigenvectors in a matrix V Ra* where k denotes the number of such eigenvectors. We will then parameterize Σ VXVT where X E g PSD. The above four conditions now imply that X solves a A x k dimensional CCP:
x a =0 (X, - ([x]] =
where
= V (
][X] = V en (
after sub-step regularization
effect of curvature rogularization
This gives a k 1s CCP for X where is is the number of eigenvectors with eigenvalues stability threshold (i.e. above 1.95/n). When simulating these flows numerically, we first compute then solve this CCP to get X & and we then update
the
E k and E Rexxxxxk using condition (1) above. Ase + 0, this process will converge to the continuous time central flow Definition 2
Interpretation as discretization of projected gradient flow We note that while this discretization scheme may seem bit strange, actually matches the natural discretization scheme for the projection interpretation for the gradient descent central flow (Definition 3). Recall that the gradient descent central flow can be written as gradient flow on the constrained optimization problem: min., L(w) such that w E S, := (w H(w) = 2/nl}. The natural idea is to discretize this using projected Euler method with discretization step-size
Wh+t proje, [w. - En LWWV)
However, this requires computing a projection onto S,,, which is intractable since H(w) is the Hessian of neural network. Therefore, we approximate this projection by its linearization around Specifically, rather than project onto the set where Hww) z 2/nI, we will project onto the linearization of this constraint around we:
H(w)) VH(w)) w.] 2/1/T.
Computing the projection onto this set requires solving the exact same CCP as the one described above, and the proof is exactly the same as the proof for Lemma [ This therefore the "natural" discretization for these types of differential variational inequalities and their related forms.
A.8.2 DISCRETIZING PRECONDITIONED OPTIMIZERS
In our codebase, we view GD, Scalar RMSProp. and RMSProp as special cases of the more general framework described Appendix A.6 this section we assumed that v, w follow the discrete time updates:
= Pt-1 + f w.), Wt+1 = Wt P(r)) (ww.).
Recall that GD, Scalar RMSProp. and RMSProp can all be written as special cases of this framework. To derive central flow for updates of this form, imposed the conditions:
31

---

Published as a conference paper at ICLR 2025
1. de P(v) P(v)-'(V((uu + 3 [
dt
2. f(v,w)+ (V
3. (w) 3 2P())
4. 0
5. Σ E ker(H (w) 2P())
As for GD, the key to numerically simulating such process is to linearize both H and P. Given oscillation covariance E, that we will describe how to compute, we will take sub-step of size and update w, by:
Wi++ wt c(P(v) (VL( + + [ww))[
Vi+e + 14 + + w.) + V. (Ve, w [-]]]
We will also define H', P' be the linearizations of H, around wy:
H' = (w,) (w:))lutte wal, P' P(v)) VP(4) vil-
Note that because we+ Mete are linear in Ex, so are H' P.. We will therefore enforce that:
H' 2P'
. Σ 00
E, ker[H' 2P').
This will again give us a CCP for E. As for GD, note that the last condition implies that E is supported on the generalized eigenvectors the edge of stability the next step. We therefore track all generalized eigenvectors which solve H(w))u )u XP(w))u with eigenvalue À above some fixed threshold (e.g. 1.95), as these have the potential to become unstable at the next substep. Let V E Id store these vectors in the GD section. Then we decompose Σ VXV" where X
IR kxk Plugging this into the above conditions gives a CCP for X:
x a 3[
(X.a-3(X]]
where
VT H(w)) + (VH(w Awy 2 P(w)) + +cV((r)) DU
8 'VT V (x)) ( P(v)) - 'VH 2VP(()) [(r, wx) [E]]]
2 Σ VXVT.
Aw, -
Au f(r, w.).
Again as for GD, computing 3 only require metrox-vector product access to the Hessian and consor-vector-vecte product access the third derivatives. When simulating the central flows, we first construct the basis V by solving the generalized eigenvalue problem on H H(w)), P(v.) and keeping the eigenvaluesleieenveetoss near 2, and we then compute these values of a E ge and * We then solve the CCP to get X E * and then use this to take Euler step of size e on the central flow. As c - 0, this process reduces to the continuous time central flow described in Appendix A.6 and in the special case of GD where 0 and P(v) 1/1,, this reduces to the discrete time update for gradient descent described above.
32

---

Published as a conference paper at ICLR 2025
B EXPERIMENTS
This section contains an overview of our experimental results. More experimental data can be found in Appendix
which contains "bulk" data from the full experiments. C
Our code can be found at: /ttpp:///iiee http://www...ooooooooccaaacaaaaa flows
Architectures We experiment six architectures:
1. a convolutional neural network (CNN);
2. ResNet (He al. 2016);
3. a Vision Transformer (ViT) (Dosovisskiy et al. 2021
4. LSTM (Hochreiter & Schmidhuber 1997);
5. a (sequence) Transformer (Vaswani et al. 2017):
6. and Mamba sequence model (Gu & Dao 2024).
Due to the computational expense of discretizing the central flows (see Appendix B.5). we ex- periment small-scale instances of these architectures. Architectural details can be found in Appendix B.6.
Datasets We test the vision architectures (CNN, ResNet, ViT) subset of CIFAR-10 (Krizhevsky) 2009). We test the sequence architectures (LSTM, Transformer, Mamba) on synthetic sorting task (Karpathy, 2020). Further details these datasets can be found in Appendix B.7 For each architecture and each dataset, we test both cross-entropy loss and MSE loss. As discussed in Appendix B.4, the central flow tends be somewhat more accurate with MSE loss.
We discretize the central flows using the variant Euler's method outlined in Appendix A.8 Implementation details can be found in Appendix B.5
Implementation
The remainder of this section is structured as follows:
1. In Appendix B.I
we present experiments for gradient descent.
2. In Appendix B.2 we present experiments for Scalar RMSProp, and we experimentally demonstrate the accellrattinn via regularization" effect that was described in Section 4.3
3. In Appendix B.3 we present experiments for RMSProp, and we experimentally support the stationarity analysis that was described in Section 5.2 well the "acceleration via regularization" phenomenon that was described in Section 5.2
4. In Appendix B.4
discuss known failure modes for the central flows, i.e. situations where the central flow poorly approximates the optimizer trajectory.
33

---

Published as a conference paper at ICLR 2025
B.1 GRADIENT DESCENT EXPERIMENTS
We tested the gradient descent central flow in the six deep learning settings described above. We experimented with both mean squared error (MSE) loss and cross-entropy loss. Figure Figure show that, in each of these six settings, the central flow can accurately predict the time-averaged (smoothed) loss curve of gradient descent across several different learning rates.
CNN
ResNet
VIT
0.5
n00.01
0.35
0.01
0.35
B.01
0,0133
0.4
n=0.00 022
0.30
0.02
n D.C2
flws
0.30
0.3
0.20
0.25
0.15
0.20
0.10
0.15
1000
2000
1000
4000
1000
2000
5000
40000
1000
2008
1000
4000
snde ccon
LSTM
Transformer
Mamba
0.4
n 0.6267
G 0.02
11 D.005
n-0.04
0.0267
0.30
0.01
0
1
0.3-
cenitral Bows
Lentinal flows
0.25
cenirral flows
0.2
0.20
0.15
0.0 1000
0.10
10000
4000
5.000
1000
2000
3000
5000
Figure 9: Central flow predictions for gradient descent loss curves (MSE). For six architectures, and three learning rates each, we compare the actual loss curve (colors) to the central flow's 's prediction for the time-averaged loss curve eq. (6) (black dashed). The faint colored lines the raw loss curves; the thick colored lines are the time-averaged (smoothed) versions. Observe that the central flow accurately predicts the time-averaged loss curves.
CNN
ResNet
VIT
1.4
0m0.0
1.4
n= 9.01
1.4
1.2
0=0.00657
0.0133
n-0011
1.0
1.0
contral floss
contral flos
6.8
0.6
0.6
0.6
0.01
0.4
0.0333
contral flows
0.0
0.0
0.0
1000
2030
3000
5000
6000
500
1000
1500
2000
2500
3000
1000
2000
3000
LSTM
Transformer
Mamba
1.4
2.5
n=
1.4
n=0.002
1.2
n- 0.02
2- 0.0133
1.2
7- 0.002267
0-04
1.0
n=0.004
contral Tows
central flows
n-0.01
3.5
0.8
0,00.3
0.6
0.6
centrol Soos
0:60
0.4
0.4
0.5
6.0
0.0
0.0
1000
2000
3000
5000
1000
1000
2000
Figure 10: Central flow predictions for gradient descent loss curves (CE). Similar figure to Figure 9 but for cross-entropy loss. Observe that the prediction is sometimes a bit off, especially at the end, and especially with larger learning rates.
34

---

Published as a conference paper at ICLR 2025
real vs. predicted cs cillatiion covariance 1w
top hessian cigenwalues
400
gradiert ark descent dentraa fow
gradiort Nlw [tpp
200
flow
100
winn Eift genvectors
0.15
flow prodiction
220 eipeniouus
2000
4000
6000
60.00
2000
4103
6002
8000
2000
4363
6000
5000
step
step
step
gradientt norm?
best ccuuaacy (8)
dista nce to gradient descent
1.25
gradient descent
contral
grudient descen: (averge)
grudient flow
contrul flow
contral flow rroiction
0.75
50
gradient descert Indipoinss
0.25
45-
-.. -ortral flow
2000
6000
8000
2000
6050
8000
0.03
2000
4300
6:00
network outputs on test example
network outputs on test mple
network outputs on example
gradient descont
grasiert descert
1.00
gradlert descent
central Tow
0.75
central faa
central Tow
0.6-
0.50
0:5
0.00
0.0-
0.03
-0.5
0
2003
4000
6000
8000
2:00
-1033
5000
8000
2000
4200
5000
D020
stop
Stop
stop
Figure 11: A close look at the gradient descent central flow. Using gradient descent with v) = = we train a VIT on subset of CIFAR-10. The central flow (black) accurately models the time-averaged trajectory of gradient descent (blue), whereas the gradient flow (red) takes a different path. As described in Appendix [B.5] we stop running the gradient flow once the sharpness gets too high, which is why the red lines stop early.
Top left: The loss along the central flow (solid black) decreases monotonically, whereas the loss along the gradient descent trajectory (light blue) behaves non-mononn non-monotonically once the dynamics enter EOS. While the loss is higher along the gradient descent trajectory than along the central flow, the central flow can accurately predict the fime-averaged loss along the gradient descent trajectory, using eq. (6) (dashed black); this be seen accurately match the time average of the actual gradient descent loss curve (dark blue). Finally, the train loss along the gradient flow (in red) decreases faster, because it follows a different, unregularized path.
Top center: we plot the top three Hessian eigenvalues under gradient descent (colors) and under the central flow (black). Under GD, the top Hessian eigenvalues equilibrate around 2/1; under the central flow they are fixed exactly at 2/1,. In red, we plot the top Hessian eigenvalue under the gradient flow, which rises beyond 2/1,. Note that for GD, we report the Hessian eigenvalues at the second-order midpoints (see Appendix B.5)) rather than at the iterates themselves, as this makes for clearer plots. Top right: we show that the central flow's (t) accurately predicts the covariance of the oscillations. In black, we plot the nonzero eigenvalues of E(t); the number is always the same as the number of Hessian eigenvalues at 2/nz. In faint colors, we plot the squared magnitude of the displacement between gradient descent and the central flow along each eigenvector of E(t). In thick colors, plot the time-averages of these displacements, i.e. the empirical variance of the oscillations along each eigenvector of E(t.. Observe that the eigenvalues of Et) accurately predict the instantancous variance of the oscillations along the corresponding eigenvector. More details in Appendix B.5 Middle left: in solid black, we plot the squared gradient norm along the central flow; dashed black, we plot the central flow's prediction for the time averraee squared gradient norm along the trajectory (similar to eq. (6). This accurately predicts the time average (dark blue) of the squared gradient norm at the iterates (light blue). Notice that most of the gradient norm comes from the oscillations.
Middle center: we plot the test accuracy under gradient descent (bluc) and the central flow (black). For gradient descent, we report the test accuracy second-order midpoint, this removes much of the oscillations. Because the central flow matches the gradient descent trajectory, the test accuracy is nearly the same across both trajectories.
Middle right: The distance in weight space between gradient descent and the central flow (black) stays small over time, indicating that these two trajectories stay close. By contrast, the distance between gradient descent and the gradiens flow (red) grows rapidly once the dynamics enter EOS. Bottom row: we show the network's 's final-layer predictions on three arbitrary examples. Under gradient descent (colors) these predictions oscillate due to the oscillations in weight space. Under the central flow (black), the predictions evolve smoothly while following the same overall path.
35

---

Published as a conference paper at ICLR 2025
For one of these runs, Figure I1 depicts more granular set of metrics. (Analogous figures for all of the other runs are given in Appendix ( Note the following:
The distance in weight space between the central flow and the real optimizer trajectory stays small over time. In comparison, the distance between the gradient flow and the real optimizer trajectory grows much faster over time.
The central flow accurately predicts the evolution of individual network outputs over time. This can be viewed as visualization of the central flow's accuracy in function space.
The central flow's ' prediction [() for the instantancous covariance matrix of the oscillations accurately predicts the empirical variance of oscillations along each eigenvector of E(())
The central flow can accurately predict the (smoothed) training loss curve, as well as the (smoothed) squared gradient norm curve.
Nevertheless, the central flow is just an approximation to the gradient descent trajectory, and this approximation break down under several different circumstances which we detail in Ap- pendix B.4
36

---

Published as a conference paper at ICLR 2025
B.2 SCALAR RMSPRRP EXPERIMENTS
We tested the Scalar RMSProp central flow in the six deep learning settings described above. We experimented with both mean squared error (MSE) loss and cross-entropy loss. Figure T2 and Figure 13 show that, in each of these six settings, the central flow can accurately predict the time- averaged (smoothed) loss curve of Scalar RMSProp across several different learning rates.
For one of these runs, Figure 14 depicts more granular set of metrics. (Analogous figures for all of the other runs are given in Appendix
ResNet
0.5
7- 0.0
0.007
0,007
0.35
- 0.02
7-
0.02
0.35
f 0.02
0.3
central Dows
0.30
flows
ceritnal flows
0.30
0.2
0.20
0.25
0.15
0.20
0.0
500
1500
3000
1000
20000
stop
stop
LSTM
Transformer
Mamba
0.4
0.5
0.005
7-0.00
0.000
0.35
q 0.01
0.3
7- 0.04
0.4
4 0.02
[ 0.82
0.30
0.3
0.2
0.2
0.20
0.1
0.0
0.10
2000 ROCO
step 8080
2000
4000
6C00
8000
1000
2000
4000
Mtep
step
Figure 12: Central flow predictions for Scalar RMSProp loss curves (MSE). For six architectures, and three learning rates each, we compare the actual loss curve (colors) to the central flow's prediction for the time- averaged loss curve eq. (6) (black dashed). The faint colored lines the raw loss curves; the thick colored lines are the time-averaged (smoothed) versions. Observe that the central flow accurately predicts the time-averaged loss curves.
1.4
1.25
g=16
76
1.25
1e 05
1.00
n- 2e
1.00
0- 4e 05
contral lows
central flows
0.75
0.75
05
0.50
0.50
0.8
- 76-06
0.25
0:25
0.6
05
0.00
0.00
2500
LSTM
Transformer
Mamba
1w05
2.5
1.4
7e-06
2w-05
2.0
1.0
2e-05
15
71 1e-25
0.50
d- Ze 05
0.6
0.25
contral flows
0.2
0.00
0.5
2000
1900
2000
1000
2000
4000
ztop
stop
Stap
Figure 13: Central flow predictions for Scalar RMSProp loss curves (CE). Similar to Figure 12 but for cross entropy.
37

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigervaluus
real Vs. predicted oscillation covariance
0.40
stable fow
Scalar RMSPhop
0.0012
0.35
Scalar AMSProp laverage)
nlo (top)
central flow
0.0008
empirical voriance BI'
--- carital al flow prediction
2-
0,0006
wigenval Laws
0.25
0.0004
0.0052
0.0000
0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4.00
5000
step
step
step
gr gradient norm?
top hessian eigeniva lues
effective step size ANV
0.085
Scaler RMSPhpp
Scwler RM SPrp flow
Scaler RMSShop
central laverage) contral flow pred ition
800
0.030
(top)
0.025
600
0.020
400
0.015
200
0.030
0.925
0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
200p
stop
network outputs on test example
test accuracy (3)
distance to Scalar RMSProp
1.0
0.25
central
0.20
0.6
Scaler RMSProp
0.35
contral flw
62-
60-
Stabo RMSPrup (m (midponts)
0.0
flow
58
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
stop
stop
Figure 14: A close look at the Scalar RMSProp central flow. Using Scalar RMSProp with ? 2/400, 3, 0.99, and bias correction, we train Mamba on a synthetic sequence prediction task with MSE loss. The central flow (black) accurately models the time-averaged trajectory of Scalar RMSProp (blue), whereas the stable flow (red) takes a different path. As described in Appendix [B.5] we stop running the stable flow once the effective sharpness gets too high, which is why the red lines stop early.
Top left: see Figure [ caption.
Top center: see Figure caption.
Top right: see Figure caption.
Middle left: Figure caption.
Middle center: while the top eigenvalues of the effective Hessian equilibrate at a fixed value, the top eigenvalues of the "raw" Hessian keep evolving. Once the dynamics enter EOS, both Scalar RMSProp (blue) and the central flow (black) start to regularize these eigenvalues relative to the stable flow (red).
Middle right: while the effective step size n/ v along the Scalar RMSProp trajectory is oscillatory, ESS along the central flow varies smoothly. At EOS, this quantity stays locked 2/S(w). Bottom left: we show the network's final-layer predictions on an arbitrary example. Under Scalar RMSProp (colors) these predictions oscillate due to the oscillations in weight space. Under the central flow (black), the predictions evolve smoothly while following the same overall path.
Bottom center: see Figure [I caption.
Bottom right: see Figure 11 caption.
38

---

Published as a conference paper at ICLR 2025
B.2.1 ACCELERATION VIA REGULARIZATION
Here, we provide experimental support for the acceleration via regularization phenomenon that was described in the main text. We conduct the following experiment: starting from some initial point we in weight space, we run the Scalar RMSProp central flow multiple learning rates 1). For each that the dynamics start EOS, i.e S(wo)/ /vro 2/n or equivalently
we set the initial /// yPa = 2/5(wo). Thus, each flow has form:
dw dt
VL(W)+ + (u( I(t:n))
(35)
S(w) adapt stop size
regularize curvature
These flows all adapt the current step size to (ww( but differ in the curvature regularization term, with larger 77 translating stronger curvature regularization, per Lemma We also run ablated flow which adapts the step size to 2/St 2/5(w) (w) but does not regularize curvature at all:
dw dt
2 S(w)
VL(w)
(36)
As shown in Figure stronger the curvature regularizer, the lower the sharpness S(w) along the trajectory. Since the Scalar RMSProp central flow keeps the effective learning rate 77/ / /v(() locked at 2w5( stronger curvature regularization thus leads to larger effective learning rates later in training. This, in turn, empirically causes the train loss to decrease faster over the long run (even though at the outset, stronger curvature regularization often causes the optimization be slower).
39

---

Published as a conference paper at ICLR 2025
effective
train 1/w/1)
sharpness S(wiit)
zV vit 2:S(vid))
003
0.20
0.01
0.15
400
0.10
2.01
1800
2000
15:0
offective
train loss L(wic))
sharpness S(wit))
//// - 2/S(witt)
actation are
9.025
0.22
0.18
502
1000
820
effective
train loss L(MEE))
sharpness S(wttt)
g/v'stt) 2/50wit))
0.26
0.23
4.0
0.23
SUTwture mng
0.033
n: 0.51
0.92
3.005
500
1000
1500
2000
510
1000
15000
2000
1000
1500
2000
stop
step size
train loss ((w20))
sharpness S(w1tt)
r/'y UTO w/S1
0.26
200
0.04
150
0.03
0.10
attetton web
100
0.08
0.02
0.06
0.01
5:0
1000
2500
2000
503
1000
1592
2000
0
302
1000
1500
2000
step
dtep
effective step size
train loss 1/wto)
sharpnesss S(w1t))
N0V 2/S1
0,045
300
0,040
0.28
0.035
0.24
Attriion wit
0,030
0.51
0,255
0.000
200
400
620
800
1000
230
400
603
100
1000
0
203
400
500
800
3000
S10p
stop
5009
effective step size
train losk ((w20))
sharpniess S(w1ti)
r/vivt)) w/tl
202
0.05
15.
0.12
0.10
Jdtation WNo
302
0.08
0.01
500
1000
1500
2003
503
1000
1505
2000
0
503
1000
3500
2000
2tpp
stop
J0ep
Figure 15: "Acceleration via regularization" for Scalar RMSProp. Starting from the same initialization, we run the Scalar RMSProp central flow various learning rates, well as an ablated flow dar di - w(w) V ((w) with curvature regularization removed. These three flows all use the same step size strategy but differ in the strength of implicit curvature regularization. Initially (see inset), the flows with higher curvature regularization often optimize slower; however, over the longer run, they able to take larger steps and optimize faster.
40

---

Published as a conference paper at ICLR 2025
B.3 RMSPROP EXPERIMENTS
We tested the RMSProp central flow in the six deep learning settings described above. We experi- mented with both mean squared error (MSE) loss and cross-entropy loss. Figure and Figure 17 show that, in each of these six settings, the central flow can accurately predict the time-averaged (smoothed) loss curve of RMSProp across several different learning rates.
For one of these runs, Figure 18 depicts more granular set of metries. (Analogous figures for all of the other runs are given in Appendix
CNN
ResNet
VIT
0.40
0-0.007
0.4
0-0.01
0.35
0-0.00
q 0.01
0,00
0.3
nows
flows
0.25
0.20
6.1
0.10
0.0
2000
4090
6000
5000
500
1000
1500
2000
2500
3000
0
1000
2000
3000
4000
step
Stap
step
LS1M
Transformer
Mamba
0.4
7-
0.01
0.5
7
0.60
1 0.005
0.02
0.01
0.35
0.51
central Rows
central flows
0.25
0.10
0
2000 4000
6000
3000
2000
4000
6000
8000
1000
2000
3000
4000
5000
S189
stop
stop
Figure 16: Central flow predictions for RMSProp loss curves (MSE). For six architectures, and three learning rates each, we compare the actual loss curve (colors) to the central flow's 's prediction for the time-avenn averaged loss curve eq. (6) (black dashed). The faint colored lines are the raw loss curves; the thick colored lines are the time-averaged (smoothed) versions. Observe that the central flow accurately predicts the time-averaged loss curves.
CNN
ResNet
76
0- 1c 05
1.25
16:5
1.25
2u-05
1.00
1.00
n
0.75
contral lows
0.75
central flows
1.0
Se-05
0.50
0.8
n 76-05
n 1e-05
0.00
500
1000
1500
2000
2500
2000
1000
2000
2000
4000
1000
2000
3000
4000
1.25
n 1w05
Ze
2.5
1.2
1 7e-06
1:
1.00
7 66-05
n 26-05
contral
central flows
0.75
1.5
0.8
2=
0.50
2- 2c 05
0.25
46-05
0.4
0.00
2000
3000
4000
5000
1000
1500
2000
1000
2000
3000
4000
5000
Figure 17: Central flow predictions for RMSProp loss curves (CE). Similar to Figure 16 but for cross-entropy loss.
41

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigeniaalues
real VS. predicted escillaion tion covariance
stable flow
RNSPrpp
0,00330
emprical lence
0.00925
along Zio cigenvectors
central (average)
flww (top)
eigenvalues
0.3
Central flow prediction
0,00015
6.2
0.00010
1-
0.00005
0.1
1000
1500
2000
2500
3000
5c0
1000
1500
200n
2500
3000
500
1000
1500
2008
2500
3000
gradient norm?
top hessian eigenvalues
test accuracy (%)
10
RMSProp
1000
RNSPmp
70-
central laverage)
800
flow (top)
contral nlow prediction
400
RNSPup (midgeints)
central
1000
1500
2000
2500
3000
5C0
1000
1500
2000
2500
3000
500
1000
1500
2000 2500 3000
network cutputs on test example
coordinates of
distance to RMSProp
contral
0.4
6.3
10-9
10:11
RNSProp
10:47
RNSProp
1000 1500 2000
2500
3000
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Figure 18: A close look at the RMSProp central flow. Using RMSProp with 7) 2 x 10 32 = 0.99, e = 10-8 and bias correction, we train a ResNet on a subset CIFAR-IO with MSE loss. The central flow (black) accurately models the time-averaged trajectory of RMSProp (blue), whereas the stable flow (red) takes a different path. As described Appendix B.5 we stop running the stable flow once the effective sharpness gets too high, which is why the red lines stop early.
Top left: see Figure caption.
Top center: see Figure caption.
Top right: we show that the central flow's E(t accurately predicts the covariance of the oscillations. In black, we plot the nonzero generalized eigenvalues of ((t with respect to the preconditioner P(t): the number always the same as the number of effective Hessian eigenvalues
In faint colors, we plot the squared magnitude of the displacement between gradient descent and the central flow along each generalized eigenvector of (t) with respect to the preconditioner P(t). thick colors, we plot the avmeraverages of these displacements, i.e. the empirical variance of the oscillations along each generalized eigenvector of (t). Observe that the generalized eigenvalues of Et) accurately predict the instantaneous variance of the oscillations along the corresponding generalized eigenvectors. More details in Appendix [B.5]
Middle left: Figure 11 caption.
Middle center: while the top cigenvalues of the effective Hessian equilibrate at fixed value, the top eigenvalues of the "raw" Hessian keep evolving.
Middle right: see Figure Ii captioo.
Bottom left: we show the network's final-laye final-layer predictions on an arbitrary example. Under RMSProp (colors) these predictions oscillate due to the oscillations in weight space. Under the central flow (black), the predictions evolve smoothly while following the same overall path. Bottom center: while the entries of oscillatory along the RMSProp trajectory, they vary smoothly along the central flow.
Bottom right: see Figure II caption.
As with the other optimizers, we find that the central flow approximation becomes less accurate as 77 becomes larger. We also expect the RMSProp central flow to break down when Bg becomes very small, as then RMSProp would not resemble preconditioned gradient descent with a slowly varying preconditioner.
42

---

Published as a conference paper at ICLR 2025
B.3.1 ACCELERATION VIA REGULARIZATION
Here, we verify the acceleration via regularization phenomenon for RMSProp. Our main claim is that RMSProp's s implicit curvature reduction behavior which is made precise by its central flow beneficial for the RMSProp's efficacy as an optimizer.
Starting from some initial point (Wo, 26) from an RMSProp central flow trajectory, we run both the RMSProp central flow, as well as an ablated flow which has the implicit curvature regularization disabled. (That is, this flow sets {((() so as to stay at the edge of stability purely via the effect of oscillations on v.) Figure 19 shows that the original RMSProp central flow: (1) traverses a lower- sharpness trajectory than the ablated flow; (2) is able to use higher effective learning rates in the long run; (3) optimizes faster in the long run.
B.3.2 STATIONARY PRECONDITIONER
Here, we empirically support our stationarity analysis from Section 5.2 which aims to understand the local dynamics of RMSProp based on the current alone, eliminating from the picture.
Convergence to stationary v We first show empirically that under the RMSProp central flow (w((), (()), the EMA v(() approximately converges to its stationary value (w(()). To numerically estimate ((w). we use the fixed point scheme outlined in Appendix A.5.1 to obtain P(w), and then we obtain (ww) from P(w).
Figures 20 and 21 plot the cosine similarity between v(() and ((w(()) over time. We can see that this cosine similarity goes to nearly 1, indicating that the ((w(()) is indeed a good approximation to the R.MSProp EMA vI)).
Suboptimality of stationary preconditioner We have shown that the stationary preconditioner P(w) at weights w can be characterized as the optimal solution to eq. (11. a convex optimization problem over preconditioners:
P(w)
arg min P dingonal, P tr(P) |V(W))|
such that
optimization speed
local stahility
It is reasonable why minimizing the first term in eq. 11 would be desirable in preconditioner: the diagonal elements of the preconditioner are the inverse effective learning rates (P; = vui/n). so the trace (r(P) is the sum of the inverse effective learning rates. Thus, minimizing tr(P) is equivalent to maximizing the harmonic mean of the effective learning rates. This is an intuitively reasonable criterion we want the effective learning rates to be large while still maintaining local stability, and one natural-enough way to measure the "largeness" of a vector is via the harmonic mean.
By contrast, the second term in eq. (1I appears to be undesirable. Prepending minus sign, the (WW)
quantity - L L(w) under the preconditioned gradient flow rate of change in the loss under the central flow, as the central flow also posesses an implicit curvature regularization term, but it is reasonable proxy. This suggests that for training to be fast, we want to be large and negative, or equivalently, we want || [[W(())] p-- to be large and positive. Yet the optimization problem eq. (11 tries to decrease [[ p- i in other words, it
is the instantaneous rate of change in the loss P-'VL(w) weights w. This is not quite the
tries to make training slower. Thus, its presence is undesirable.
One can consider an alternative optimization problem over preconditioners that omits the problematic second term:
P(w)
arg min P diagonal. P
tr (P(P) such that
H (ww) 2P.
(37)
local stability
From the definitions of eq. I1] and eq. 37 as optimization problems, it is guaranteed that precondi- tioned gradient flow using this alternative preconditioner P(w) would decrease the loss at least as fast as preconditioned gradient flow using the RMSProp preconditioner P(w):
(V) < | P(w)
(38)
To illustrate this point empirically, and see how large the gap is, in Figure 24 and Figure 25 we compute both || [VL(() (?? and ||V [())] along the trajectory of the RMSProp central flow.
43

---

Published as a conference paper at ICLR 2025
fixed point iteration similar that described in Appendix A.5.1 but To compute P(w),
which leaves out the gradient term. We observe that [VL(w)] P(w) is often substantially bigger than [[V[(w)
P(w) p(w)" indicating that alternative preconditioner would have been non-negligibly better. The degree to which the alternative preconditioner is better empirically depends the RMSProp learning rate This makes sense, because as grows, the effect of the second term in Equation [11 diminishes.
Thus, our analysis suggests natural direction for improvement for RMSProp's ' implicit precondition- ing scheme. It is interesting open question whether there exists optimizer update rule that has eq. (37) as its implicit preconditioner.
Stationary flow
We have shown that RMSProp's preconditioner converges the stationary value eq. (II). However, this statement does not fully capture the local dynamics of the optimizer, as the RMSProp central flow is not merely preconditioned gradient flow it also includes curvature penalty which refeett the effect of the oscillations on the curvature.
To approximate the local dynamics of RMSProp closed system in w alone, we propose the following stationary flow:
dw dt [(w) VL(w))
m ((r)) implicit curvature penalty
stationary
preconditioner
where = [(w; my: 32) the unique value of which satisfies the RMSProp CCP eq. (44.. but using the stationary preconditioner P(w) rather than the RMSProp preconditioner derived from Equation (I2) models RMSProp implicitly performing preconditioned gradient descent survature-penalizee objective, using the current stationary preconditioner P(w).
Figures and show that the stationary flow eq. 122 can indeed accurately predict the instantaneous
rate Loss decrease along the central flow.
44

---

Published as a conference paper at ICLR 2025
ha narnonic mean of effective stop siboed
train loss
sharpness
certral 100iinn
343
central 3012ion
123
CNN
200
central abrasion
83-
0.10
0.92
0.05
500
820
harmonk mean of effective step sizees
train loss
sharpness
centra
0.4
centrl
10 ldti
600
st4vdon
straiin
ResNet
0.3
400
9.2
200
5.3
8:0
600
harmonic effective step sizees
train loss
sharpness
cortra abilton
2000
cential diii
0.30
contul
abioion
LSTM
1500
2.00
0.68
500
300
1500
2000
1300
2500
2000
2300
1300
2500
harmoric
train loss
sharpness
effective
Transformer
0.25
ablatio
253
central
200
9.35
0.13
0:10
0.02
252
750
1250
1505
250
750
1250
1503
250
350
1250
harmanic
effective
0.225
5000
5.30
Mamba
0.200
0.00
0.355
1000
0.152
2000
5.02
0.102
0.00
530
1000
15:00
2000
500
1000
1500
2000
0
500
1000
1502
2000
Figure 19: "Acceleration via regularization" for RMSProp. We compare the the RMSProp central flow (blue) an ablated version (orange) which leaves the implicit curvature regularization (maintaining stability purely by the effect of oscillations on v)). At first, the central flow often optimizes slower than the ablated flow, due to the presence of implicit curvature regularization. But over time, it navigates to lower-curvature regions (right), where it takes larger steps (middle). and optimizes faster (left). Each row is different DL setting. The left column plots the train loss, the middle column plots the harmonic mean of the effective learning rates, the right column plots the sharpness. These experiments all use MSE loss.
45

---

Published as a conference paper at ICLR 2025
CNN
1o:05
(-- 20.05
G- 40:05
03306 0166001
440335
0.29
0.95
0.95550
2.50
3.590
0.82
0.90
0.75
300
10000
1520
2000
25:00
500
1000
15:00
2003
2500
500
1030
1593
2020
2500
xtep
18ep
ResNet
26-05
4e-05
0000000
000001 3100400
888800
2.25
1.0000
0,990
0.90
0.50
02275
0.90
0.997
0.97
0.85
0.35
0.06
2.50
0.75
0.80
10:00
1500
1000
1500
2500
1000
1502
2000
2500
VT
1.000
1.000
1,000
1 1.00000
0.939350
73
3.998
0 99995
0.900
5.991
0.097
2000
1000
1500
2500
3000
1500
1000
2000
stap
500e
LSTM
n= 15:05
n= 2e-05
q= 4e 05
1.000
1.50
1.00
024140 200000
0.015
0.98
0.900
3.406525
(19097
0.95
940
0.985
0.250
0.34
3000
4003
3000
6030
7000
1000
2000
3000
4000
2003
4003
5200
0020
s7ep
step
stpp
Transformer
n 2e-05
4e-05
0 0.0001
1.00
0. 0.9995
0.59500
0.99
0.0900
0.956
0.9955
0.904
0.58
0.900
0.32
500
1003
1500
3000
2500
550
2000
1500
2000
2500
500
1000
1500
2000
2500
<tep
step
Marbba
1e-05
2e-05
46-05
0.2090
0.900
0.0996
0.50
0.955
0.800
0375
2000
2000 30000
4000
5000
1000
2:000
20000
4000
5020
step
Step
Figure 20: The EMA v reaches stationarity during training (MSE). While running the RMSProp central flow, starting at the time when training enters EOS, monitor the cosine similarity between the EMA v(t) and the stationary EMA (w()).. This cosine similarity rises to high values (nearly 1) during training. implying that vt)) reaches stationarity.
46

---

Published as a conference paper at ICLR 2025
CNN
76:06
1e:05
20-05
c32e66 9160000
0.9494
0.09300
0 1:998
0.92035
0. 0.99965
0.:950
0.99775
0.9555
1000
1500
2000
2500
500
1200
1500
2000
2500
502
1000
1500
2000
2500
step
step
ResNet
05
Ze 05
05
1,0000
1,00000
0.9999
3,9995
0:39916
339109
0.998
9.9050
0.9997-
5.9235
9.596
VT
7 Se-06
2 76-06
7 1e-05
1,0000
1,0000
1,000
3,9999
0.999
0.99975
9.9958
9.598
3.9997
0.0997
9.597
3,9996
1500
2000
2500
3000
3500
1500
1500
2000
25:00
3000 3500
1000
2030
3000
St0p
2209
Stop
LSTM
o= 1e-05
2e 05
6e 05
1.00000
0.9 0.0000
2.59293
0.9093732
0.9998
0.9996
1:200:00
0.359:9
0.99950
3,54590
0.00977
0,9900
0.80000
3300
3730
4200
4250
4300
4750
2300
2500
3000
2500 4200 4500
1002
2010
3000
4002
5000
step
step
step
Transformer
1e-05
2e-05
4e-05
1,0000
0.9000
5.5986
0.99956
0,0994
0.9598
3,9992
0.3597
3,5000
1520
2003
2530
Zup
Mamba
n- 76-06
1e-05
1,00000
1,0000
0
0.999922
0.9959
0.90900
0.99285
0.9957
0.99280
0.9919
0.99975
0.4455
Figure 21: The EMA a reaches stationarity during training (cross-entropy). This figure is analogous Figure 20 but for cross-entropy loss.
47

---

Published as a conference paper at ICLR 2025
CNN
1e-05
2e-05
46-05
10
10-
10
10a4
10.1
10-
10
10
30-1
10-15
10-31
300
1200
1590
2000
2300
300
1000
1500
2000
25.0
540
1000
15:00
2000
2500
st6p
step
step
ResNet
05
05
n 000000
20
10.
10
30-4
10
10-1
xith Rlwe
10
10
10-11
10:12
10:12
10-19
1000
1500
2000
2500
500
1000
1500
2000
2520
5:0
1000
15:00
2000
2500
slap
stap
sEap
VIT
7e:06 06
1e:05
2o:05
10
10.
19--
10:44
30-9
10-1
10-41
1500
2009
2500
3003
3500
1000
1500
2000
2500
3000
3500
1000
2009
3003
step
step
sEap
LSTM
2e:05
10
10-5
30-4
10-7
10-a
3000
4000
5000
6000
7030
8000
2000
3002
4003
2000
4000
6000
8000
stap
stap
step
Transformer
11
n 0.0001
10
10-5
10
10-1
10-1
10-18
10-19
500
1000
1500
2000
2500
500
1000
1500
2000
25:00
510
1000
1500
2000
2500
step
step
step
Mamba
1e-05
n
10
1.00
10
10
10
10
10-14
10
13-19
10-31
12-14
1000
2000
3030
4000
1000
2000
3003
4000
5000
1030
2000
3000
4000
5000
step
Figure 22: Stationary EMA is accurate at a coordinate-wise level (MSE). While running RMSProp central flow, starting at the time when training enters FOS, we plot the evolution of ten coordinates of the actual EMA v(t) (dots) and the stationary EMA (ww()) (half-black dashed lines). Each color is different coordinate, and the ten coordinates network. We can see that, starting soon after training reaches EOS, the stationary EMA ((w(t)) becomes excellent approximation to the real EMA v(t), coordinatewise level. We can also see that both the real EMA and the stationary EMA evolve significantly (in tandem) during this time.
uniformly spaced throughout the
48

---

Published as a conference paper at ICLR 2025
CNN
76-06
1e-05
2e-05
10
10
10"1
10-1
10
10.19
10
1000
1500
2000
2500
500
1020
15:00
2000
2103
500
1000
15:00
2000
2500
st6p
step
step
ResNet
05
2c:05
46.05
10'
10
10
10"
30-4
xith
10
10-11
10-12
10-55
30-19
1500
2000
25:00
3000
3500
1000
2000
5003
1200
2000
3003
slap
stap
sEap
VIT
06
7c:06
1e:05
10
10.1
10
10-9
16-9
10-1
10:12
10-11
1500
2009
2500
3003
3500
1000
1500
2030
2500
3000
3500
1003
2000
3000
stap
step
sEap
LSTM
2e:05
6e
10
10
10-1
13-11
10
3500
3759
4000
4250
4500
4750
2000
2500
3030
3500 stap
4000
4500
1000
2000
30.00
4000
5000
38.9
Transformer
16-05
1 2e:05
n
10
10
15-4
10
13-39
10"0
1900
1509
2000
509
1000
1500
2009
2500
5:0
1000
1500
2009
2500
step
step
step
Mamba
n
201
10
10
10
10
10-1
10:10
13-43
13-3
10-20
13-11
1050
2000
3000
4000
5000
1000
2000
3003
4000
5200
1000
2000
3000
4300
5000
step
step
Figure 23: Stationary EMA accurate at a coordinate-wise level (CE). Analogous Figure 22 but for cross-entropy loss.
49

---

Published as a conference paper at ICLR 2025
CNN
1e-05
2e-05
46-05
10-3
10-1
10-a
procond Boner
10-1
10-1
Nloo second berm
10
10
10
302
1203
1500
2003
2500
500
1000
1500
2000
2500
500
1000
1500
2000
2500
stpps
stapa
stapa
ResNet
f) 05
2
0.0001
IVL(
10
10
prwcond foone
Meo socond temm
10-4
10-4
10-
1000
1500
2000
2500
500
1000
15:00
2000
2500
505
1000
1500
2000
2500
steps
stapa
atapa
ViT
76-06
f) 1e-05
1 2e-05
[VL(WI)
procand Boner Wefe Socond temm
10-4
10
1300
2000
2303
3000
3500
1000
1500
2003
23:00 steps
3000 2500
1000
2000
3005
Steps
steps
LSTM
n
1
n
10-
|VL/W/l
10-3
10-1
procand Boner
10-1
10
second
10a1
10
3020
4020
5020
6020
7000
8000
2002
3000
4200
2003
4003
5000
8000
staps
stspa
Transformer
n 2c:05
n 0.0001
10-1
10-3
10-2
statanary second bern
10
10
floner
10'
1000
1500
2000
2500
500
1000
1500
2500
500
1000
35:00
2000
2500
stapa
M Namba
1)= 1e-05
1 2e-05
f) 4e-05
10-1
10-2
10m
|VL(WII
10--
13-2
statiaary flnner
2011
13a4
1000
20000
3000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Figure 24: RMSProp stationary preconditioner is suboptimal. We compare the RMSProp stationary preconditioner, defined the solution to the optimization problem eq. (17). an alternative preconditioner defined as the solution to eq. (37). similar optimization problem but without the second term in the objective. We assess each preconditioner P by reporting || [() the instantaneous rate of decrease in the loss under the preconditioned gradient flow with preconditioner P. Observe that this value is higher under the alternative preconditioner (orange) than under the RMSProp stationary preconditioner (blue), meaning that the alternative preconditioner would decrease the loss faster. The gap between the two preconditioners tends to be smaller when v) is larger, which makes sense, because the second term in eq. (11) is proportional to 77
50

---

Published as a conference paper at ICLR 2025
CNN
7e-06
1e-05
2e-05
10-1
procond Boner
Nlo second term
10-1
10
10a
10-1
10-4
10-4
1000
1550
2000
2500
500
1000
15:00
2000
2500
300
1000
1500
2000
2500
Stops
stapa
stppa
ResNet
f) 05
45:05
1.0
10-
10
IVL(
precand honer
Meo socond temm
10-1
10-1
10-4
10-1
1590
2000
2500
5000
3500
1000
2000
3000
1000
2000
3000
staps
slapa
stapa
VIT
Se-06
f
10
10
10
procand Boner
Socond term
10-1
10-1
10-1
1300
2000
2303
3000
3500
1000
1500
2003
23:00
3200
3500
1000
2000
3000
Steps
Steps
steps
LSTM
1e-05
1
10-1
10-0
10-1
10-7
10
10:
10-1
10-1
10-1
procand boner
second
10
10-4
10-0
10
10a4
3500
3750
4000
4250
4300
4750
2000
2500
5000
3500
4203
4300
1000
2000
3000
4000
5000
Steps
steps
Transformer
n- 15:05
n
n
10
10:1
10t1
statonary
10-1
10-2
10-2
foner fioner
second
10-2
1000
1500
1000
1500
2000
2500
500
1000
3500
2000
2500
stsps
M Mamba
n-
1)
1)
2e
10-1
10-1
10-1
statianry flnne
1000
2000
3000
4000
5000
1000
2000
3000
4000
1000
2000
3000
4000
5000
Figure 25:
RMSProp stationary preconditioner suboptimal. This figure is analogous to but for cross-entropy loss.
Figure 24
51

---

Published as a conference paper at ICLR 2025
CNN
1e-05
2e-05
7 4e-05
decrease .dL/dt
rate of loss
10"
10
10
curveture
CUrVeSUre heg
20:
10
30
503
1200
1500
2000
2500
500
1000
1500
2000 2520
320
1000
1500
2000
2500
steps
Steps
ResNet
n- Ze
1) 05
0.0001
rate of loss decrease -didt
10-
10-1
10
cst male CUEVeSSTH heg
10-1
10
10
1000
1500
2000
25:00
500
3000
1500
2000
2500
500
1000
1500
2000
2500
slaps
stapa
3Eapa
VIT
n- 76-06
1 1e 05
1 2e:05
rate of loss decrease
real
mati mate Curvo l0g
1500
2003
2500
3000
3500
1000
2003
3000
1000
2000
3000
staps
stapa
sdppa
LSTM
20.00 05
Q 2e:05 05
4e 05
rate of loss dd//tt
10
10-
and mate w
10
10-b
JurV l0g
3000
4003
5000
6000
7000
8000
2000
3000
4000
2000
4.000 steps
6000
8000
stape
stapa
Transformer
q- 2c:05
q 4e-05
1 0,0001
rate of loss decrease dd//tt
real
20--
16-1
10--
and made w
SuFV l0g
20-1
16-3
10-4
500
1000
1500
2000
2500
520
1000
1500
2000
2500
500
1000
15:0
2000
2500
stags
stapa
staps
LSTM
n= 1e-05
1
n
decrease dL/do.
rate of loss
20
10-1
10
real
and w
10-1
10
Surc 10g
3000
4003
5000
5000
7000
8000
2000
3000
4000
2030
4.000
6000
8000
stage
Steps
steps
Figure 26: Stationary flow accurately predicts the instantancous speed of optimization (MSE). The stationary flow eq. (12). which incorporates an implicit curvature regularizer, predicts (black) the rate of loss decrease di (bluc) more accurately than a naive estimate [W))") (|)p--(()))
(w) (in red)
which uses the stationary preconditioner but does not incorporate curvature regularization. Observe that the gap between the two estimates is larger when 7) is larger, suggesting that, like Scalar RMSProp. the implicit regularization of RMSProp increases strength with 1).
52

---

Published as a conference paper at ICLR 2025
CNN
7e-06
1e-05
2e-05
rate of loss
10-4
10-4
10-4
curveture
CUrVeOUTe teg
10:4
10
10th
1000
1300
2000
2500
300
1000
13000
2000
2500
320
1000
1500
2000
2500
S1606
steps
Steps
ResNet
n- Ie
1) 05
?)
05
10.1
30
-dd/at
rate of loss
10-4
10-4
10-1
mate CUEVeOUTH teg
10-4
10
10:15
3500
2000
2503
3000
3500
1000
2000
3000
1000
2000
3000
stapa
3Eppu
VIT
06
f 05
1 2e:05
decrease -didtt
rate of loss
matimats male CurvoRurC l0g
10-1
1000
2000
3000
1020
2000
3000
1000
2000
3000
staps
atapa
3dppa
LSTM
n 26:05
Q 2e-
f)
dd/att
10-1
rate of loss
end made w
10--
10-1
10-4
JurV l0g
10
3500
4000
4500
2000
3003
40000
3000
2000
30.00
4030
5000
stpss
atapa
xtapa
Transfa meer
q-
1)
1 4e
decrease dd/dtt
10-1
10-1
10-3
rate of loss
real
and mate w
4
4x.10
4
CurV l0g
10-1
3x
3
1.000
15:00
2000
500
1000
1500
2000
2500
500
1000
1500
2000
2500
stacs
steps
LSTM
n
f
1
decrease dL/dot
10m
10
10-1
rate of loss
w
10--
10
10:4.
SurV 10g
10
10-
10a
3500
4000
4500
2000
3003
4000
3030
2000
3000
4030
5000
steps
steps
STeps
Figure 27: Stationary flow accurately predicts the instantaneous speed of optimization (CE). Same as Figure 26 but with cross-entropy loss.
B.4 KNOWN FAILURE MODES FOR THE CENTRAL FLOWS
So long as an optimizer is initialized stably (i.e. with effective sharpness below the critical threshold), the central flow is intended to approximate the real optimizer trajectory the long term. However,
53

---

Published as a conference paper at ICLR 2025
this approximation can fail. Below, we describe some common failure modes which, in experience, consistently cause the real optimization trajectory to deviate from the central flow. A promising direction for future work is to precisely identify conditions under which the central flow does does not approximate the real optimizer trajectory.
Sufficiently large learning rates For all three optimizers studied in this paper, we reliably observe that the learning rate hyperparameter is made increasingly large, the real optimization trajectory tends to deviate more from the central flow, as illustrated in Figure 28 (As an extreme example, even if the optimizer is initialized stably, very large learning rates sometimes cause the real optimizer to explosively diverge in the middle of training, whereas this never happens under the central flow.) We do not know whether some corrected version of the central flow would be more successful at capturing the real optimization trajectory in these scenarios, or whether the real trajectory is simply too chaotic to be captured by any flow.
training loss
network output test example
weight space distance between with and
1- 0.81
0.3
n 02
0.0
central flows
0.2
-0.2
0.1
0.0
0.0
1000
2000
3000
4000
2000
step
step
step
Figure 28: Central flow approximation is less accurate at larger learning rates. We run both gradient descent and its central flow at three learning rates (colors). The larger the learning rate, the faster the growth in the accumulated approximation error (right). Indeed, at larger learning rates, the network's output on an arbitrary test example can be visually seen to be slightly different between the central flow and gradient descent (middle). Nevertheless, the central flow approximation is still accurate enough here to accurately capture the train loss curves (left). (CNN/ / CIFAR-10/ / MSE)
Sub/super-quaaraatiitty Sometimes, at the edge of stability, the sharpness does not equilibrate exactly 2/1; rather, it equilibrates slightly above or below 2/1 (e.g. 2.1/n 1.9/1)). Damian et al. (2023) showed that this is due to contributions from higher-order derivatives, beyond the third-order derivatives we study. For example, when the fourth derivative of the loss along the top eigenvector direction is positive, the local loss landscape is super-quadratic, and the sharpness at the iterates will be above 2/1; when this fourth derivative is negative, the loss loss landscape is sub-quadratic, and the sharpness at the iterates will be below 2/17. Since our derivation for the central flow assumes that the sharpness along the timeraeeraeed trajectory is located exactly at 2/1z. this leads to a discrepancy between the central flow and the real optimizer which causes error to gradually accumulate over the long term. For the special case of one unstable eigenvalue, [Damian et al. (2023) Appendix F derived corrected version of their constrained trajectory (analogous to our central flow) which they empirically showed to match the real gradient descent trajectory even in the sub/super-quadratee setting. It would be interesting to try to re-derive this correction under the central flow framework, which would require extending this correction to the setting of multiple unstable directions.
Smoothness of architecture The smoothness of the architecture (e.g. of the activation function) seems to affect the accuracy of the central flow approximation, although it is not clear how to precisely quantify "smoothness." " In Figure 29) we show that as a network's 's activation function is interpolated from smooth to non-smooth, the accuracy the central flow approximation degrades, both in weight space and function space.
Large spikes When the EOS dynamics lead to extremely large spikes (c.g. in the gradient norm), we have found such spikes can cause the real trajectory to deviate noticeably from the central flow. This may be related to the large learning rate issue described above.
Interactions with loss criterion We have empirically found that the large spike issue are more common with cross- loss than with mean squared error loss (although we do not have a satisfactory explanation for these observations). Consequently. the central flow approximation is often more accurate under the MSE loss than the cross-entropy loss.
issue and
54

---

Published as a conference paper at ICLR 2025
(a) B-GeLU activation function
(b) weight-space distance
between w(t) and W,
14:
622
8.0
3-50
1-100
2-220
54-
830
25-
665
0.00
000
1990
zice
2508
3500
vooo
(c) loss and network output on test example
[- 1.0
5.0
f- 10.0
B- 20,0
0.4
0.4
0.4
0.4
12 toss
train lose
0.3
0.3
train
0.3
0.2
0.1
0.1
0.3
networ ourput example
1000
2000
3000
4000
1000
2000
3000
4000
netarkk cupuu on test esamde
1000
2000
3000
4000
network kutoorkkouruu.. crample
1000
2000
30000
4000
0.4
0.3
0.4
0.2
0.0
0.0-
0.1
0.3
-0.2
0.0
1000
2000
3000
4000
3000
2000
3000
4000
1000
2000
3000
4.000
1000
2000
3000
4000
step
Figure 29: Accuracy of central flow degrades as activation function becomes less smooth. We consider networks with the B-GeLU activation function from Dauphin et al. (2024). defined as x (Bxr where is the standard Gaussian CDF. This activation interpolates between (smooth) GeLU when 3 and (non-smooth) ReLU when @ co. Subfigure (a) plots this activation function with varying B. Subfigure (b) shows that when 3 is larger (i.e. when the activation is less smooth), the approximation error between the central flow w1) and the optimizer trajectory wy grows faster. Subfigure (c) plots the loss curve, and the network's 's output on a test example, for both the optimizer trajectory and the central flow. Fortunately, even when 20, which point 3-GeLU is a very close approximation to ReLU,, the central flow accurately predicts the overall training loss curve.
B.5 IMPLEMENTATION DETAILS Our code can be found at:
flowss
In order to reuse code between the central flows for gradient descent, Scalar RMSProp, and RM- SProp, we view all three optimizers as instances of preconditioned gradient descent W/+1 = Wt - (v.,), as described in Appendix A.6 For example, for gradient descent, we have Pt = -I. For all three optimizers, we define the effective Hessian as P.1 (wtt), and the EOS condition is that the largest eigenvalue of this matrix is 2.
Eigenvalue computation To regularly recompute the top eigenvalues and eigenvectors of the effective Hessian, use the LOBPCG algorithm (Knyazev 2001), which only requires access to the Hessian via Hessian-vecton products, and which allows us to warm-start using the previously computed eigenvectors. We were originally inspired by the LOBPCG implementation in Jax's jax.exxe popperimmnnn rimmetaalll...m.mmmm............ spa rse. spa.se. linalg (Bradbury et al. 2018).
Stable flows To discretize the stable flows (e.g. gradient flow), we use Euler's method. We dynamically adapt the discretization step size based on the current effective sharpness. Discretizing the stable flow would take prohibitively long time in regions of weight space where the effective sharpness is too high, so we automatically terminate the stable flow if the effective sharpness exceeds certain threshold.
55

---

Published as a conference paper at ICLR 2025
Central flows To discretize the central flows, we use the scheme described in Appendix A.8 At each discretization time step, we re-estimate the top eigenvectors and eigenvalues of the effective Hessian, we compute the necessary third derivatives (quamtities of the form V V H (w)), where W4 Uy are critical eigenvectors of the effective Hessian), we solve the CCP to compute (t) and to determine the update for w (and, for adaptive optimizers, v), and we apply that update.
Discretizing the central flows is computationally expensive, due to the need to continually re- estimate the top sigenvector/eeigenvalees of the effective Hessian, and to compute the gradients of these eigenvalues (third derivatives). The time complexity of each discretization step grows quadratically with the number of eigenvalues that are at the edge of stability. In particular, if there are k digenvalues/eigenveetoo at the edge of stability, then we need to compute a third derivative Vulu? H [w)u,] for each distinct pair uj Uj of these cigenvectors, including the case where i = j. which comes out to k + (+) (122) third derivative computations.
The quality of the central flow approximation is enhanced by starting it 10-15 steps into training. This likely due to the size of the gradients during these first few steps, and the central flow could possibly be improved during this phase by incorporating the implicit gradient norm penalty from Barrett & Dherin (202T). Thus, most of our experiments first run the original optimizer for 10-15 steps and then use this as starting point from which to run both the original optimizer and the central/stable flows.
Burn-in
How many eigenvalues to track? For all processes (i.e. discrete optimizers, central flows, stable flows), we track all eigenvalues of the effective Hessian that are above the threshold 1.5. We then track the same number of eigenvalues of the Hessian. Note that for gradient descent and Sealar RMSPPrp the Hessian eigenvalues are trivially related to the effective Hessian eigenvalues, whereas for RMSProp we need to do an extra eigenvalue solve to obtain the Hessian eigenvalues.
Second-order midpoints When reporting metrics from the discrete optimizers (gradient descent, Scalar RMSProp, and RMSProp, as opposed to their central flows), we usually report the top Hessian eigenvalues measured not at the optimizer iterates {wy} themselves, but at the (second order) midpoints { defined wh := [2w, W wt+1] (so named because it is the midpoint of the midpoints 1 [w- wt-1] and 1 [u+++1 we)). This empirically makes the Hessian measurements crisper (less "noisy"), while not altering the fundamental patterns.
Assessing the central flow's prediction for the oscillation covariance The central flow's Σ(t) is intended to model the "instantaneous covariance matrix" of the oscillations at time Since E(() is continually evolving, it is not immediately obvious how to numerically verify this prediction. We now describe our solution, first for gradient descent. At time if the Hessian has & eigenvalues at criticality, then (t) will have k nonzero eigenvalues, which are intended to predict the instantancous variance of the oscillations along the corresponding eigenvectors. Thus, as we simultaneously run both the optimizer and the central flow, at each time/step t we measure the squared magnitude of the displacement between the discrete optimizer and the central flow along each eigenvector of E(!), i.e. (vf (w: - w (=))) where vj is the j-th eigenvector of E(). These quantities are plotted in faint lines colored lines in plots such as Figure 3 (third plot), and each eigenvallu/ccinnvectr j is given a different color. Note that the eigenvectors and cigenvalues of (t) are gradually evolving (except at breakpoints where a new Hessian eigenvalue enters or leaves EOS, causing Et) to be discontinuous). We then smooth these time series by convolving them with Gaussian kernel, and plot the resulting smoothed time series as thick colored lines in plots such as Figure 3 (third plot). These are our empirical measurements of the variance along the different eignevector directions. We then compare these to the eigenvalues of [(), which are our predictions for these variances.
For the two adaptive optimizers, we assess whether the generalized eigenvalues of E(t with respect to the preconditioner P(t) accurately predict the instantaneous variance of the oscillations along the corresponding generalized eigenvectors. That is, if Uj is a generalized eigenvector of (t) so that H (wo(t)) = 2P(()), then we measure (v/f P(w. - w())))² and compare the Gaussian smoothing of that quantity to the j-th generalized eigenvalue of (t).
For the experiments in Appendix B and Appendix C. for each optimization algorithm (GD, Scalar RMSProp, RMSProp), we used fixed Gaussian kernel bandwidth for all experiments and for all time within each experiment. However, we believe that it may be more appropriate to use different bandwidths at different times within single run, and we used this in Figures 3 5 and
56

---

Published as a conference paper at ICLR 2025
Numerical issues in the central flow In some runs (e.g. Figure 34.17 the central flow's Σ(t)
predictions are erratic for stretches. This phenomenon seems to be a numerical issue that arises when discretizing the flow. We are optimistic that the issue is fixable.
B.6 ARCHITECTURE DETAILS
Here we describe our architectures. Note that our code for all architectures can be found at: http::/gittuu..commoooom www.sm/mooosalll///oooo cen. flows
Both our derivations and the analytic formulas for the central flows rely on higher-order information about the loss function (e.g. Hessians and third derivatives). As result, require that all of the architectures are smooth. This rules out the commonly used ReLU activation (Nair & Hinton 2010). Note that it might be possible to central flows ReLU nets by replacing ReI.U with close approximation that is smooth, as in Figure 29
CNN Our CNN has four layers, an initial channel width of 32, and 3x3 convolutional kernels. uses the GeLU activation function, average pooling, and linear readout layer.
ResNet We use a ResNet (He et al. 2016) with 20 layers and GeLU activations. We use GroupNorm
(Wu & He| 2018) in place of BatchNorm (loffe & Szegedy 2013), we empirically find that BatchNorm often leads to sub/super-quaazaticit (see the discussion in Appendix B.4
Vision Transformer We use the Vision Transformer (ViT) (Dosovitskiy et al.) 2021 implementa- tion from LucidRains (2024). Our ViT has depth 3, embedding dimension 64, number of heads 8, MLP dimension 256, and patch size 4. We use the ViT modifications proposed in |Beyer et al. (002)) We initialize the weights and biases of the final linear layer to be zero, as this makes the curvature low initialization. For unknown reasons, we found that the core PyTorch LayerNorm implementation
(written in C++) leads to third derivatives being computed incorrectly; thus, we substituted in an alternative implementation written in vanilla PyTorch, which empirically fixed the issue.
LSTM Our LSTM (Hochreiter & Schmidhuber 1997 has 2 layers, an embedding dimension of 48, and a hidden dimension of 48.
(Sequence) Transformer Our sequence transformer has 4 layers, an embedding dimension of 32, an MLP dimension of 128, and attention heads. We disabled dropout, to make the network deterministic. As with the VIT (see above), initialize the weights and biases of the final linear layer to be zero, and we substitute a vanilla PyTorch LayerNorm implementation in place of the default C++ LayerNorm implementation.
Mamba We use the Mamba (Gu & Dao 2024 implementation from Torres-Leguee 2024). Our
Mamba has 2 layers and model dimension of 64.
B.7 DATASET DETAILS
Here we describe our datasets. The code can be found at: http://github www.///owwlll//eenoooooaooo
flows
We test the vision architectures on subset of CIFAR-I0 that contains 1000 training examples, all from the first 4 CIFAR- classes. We use the standard preprocessing of subtracting the lataset-wide channel-wist mean, and dividing by the dataset-widh channel-wise standard deviation. When training using MSE loss, we encode the ground truth class as and the others as 0.
CIFAR-10
Sorting We test the sequence architectures on the synthetic sorting task described in Karpathy (002)). The network is fed a sequence of numbers and is then tasked (via a language modeling loss) with returning these numbers in sorted order. We used numbers through 4, and sequences of length 8. The size of the training datraset was usually 1,000 (except for Mamba, where it was 250).
C BULK EXPERIMENTAL DATA
This section contains the "bulk" experimental data from all of our experiments. Appendix C.1 has GD experiments, Appendix C.2 has Scalar RMSProp experiments, and Appendix has RMSProp experiments. See Figure 11 for fully annotated example of a gradient descent trajectory, Figure [4] for a fully annotated example of a Scalar RMSProp trajectory, and Figure 18 for a fully annotated example of RMSProp trajectory.
57

---

Published as a conference paper at ICLR 2025
C.I
GRADIENT DESCENT
real predicted osc ation covariance 1c-5
train loss
top hessian eigenvalues
gradient flow
descent gradient (tup)
empittccll varfance
0.4
slong 2ihh elgenvectors sigenvalues
grachent descent (average)
central flow
central fow preciction
0.2
0.1
3000
1000
20:00
3000
2000 step
3,000
stop
stop
gradient norm?
test accuracy
distance gradient descent
gradicnt descont
0.8
contral flow
gradient descent loverage)
gradient flow
central flow
50
0.4
40
gradient descent (midponnts)
0.2
contral flow
0
1000
2000
3000
4000
1000
2000
3000
4000
10:00
2000
3000
4000
step
step
step
network outputs on test example
network outputs on test example?
network outputs on best example 3
gradient descent
1.0
0.8
gradient descent control
flow
0.50
0.5
0.4
0.25
0.00
0.0
-0.25
0..5
20000
3000
2000 step
500n
step
step
Figure 30.1: Gradient descent central flow for CNN with MSE loss, 17 = 0.01.
real vs. predicted oscill ation covariance 1e-5
train
top hessian eigenvalues
300
gradient flow
gradient descent
gradient (average)
0.3
coww
200
0.2
central now prediction
100-
empirical varfance along Eth elgenvectors 210 digenvalues
0.1
0.0
1000
2000
3000
4000
1000
2000
3000
4500
1000
2000
3000
4000
step
gradient norm!
test accuracy (%)
distance to gradient descent
1.25
control
gradient flow
50
0.75
gradient descent
0.50
40
cenital flom central flow prediction
graddent descent (midpuints)
--. central flow
0.00
1008
2000
3000
4000
1000
2000
3000
4000
10:00
2000
3000
stop
stop
step
network outputs on test exampte
network outputs on test example
network outputs on test example
0.75
gradient descent
central
gradient descent
0.6
0.50
0.5
0.4
gradient descent
0.2
central
0.00
0.0
0.9
-6.25
-0.5
-0.2
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
stop
step
Figure 30.2: Gradient descent central flow for CNN with MSE loss, 17 0.013333.
58

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
200
0.4
gredient flw
gradient descert
0.00012
descent
gradient descent (average)
gradient (top)
central
0.00008
0.2
central flow prediction
100-
0.00006
0.00004
emstrical variance ce
0.1
50-
0,00002
zicl eigervalues
0.0
0.00000
1000
2000
3000
4000
1000
2000
3000
4.000
2000
2000
3000
4000
step
step
step
gradient norm?
test accuracy (%)
distance gradient descent
central flow gradient
50-
2.0
gradient oescent (iverage)
Central flow
gradient descent (midpontss
0.5-
central flow prediction
central flow
30
0.0
1000
2000
3000
4000
1000
2000
2000
4000
1000
2000
2000
4090
step
step
step
natwork outputs on test example
network outputs on test example 2
netwark autputs on test example
1.0
grodiont descent
0.8
0.5
0.5
0.4
0.0
0.2-
central
0.0
-0.5
centrol flom
gradd and descent
-0.5
-0.2
1000
2000
3000
4000
1000
2000
3000
4000
0
1000
2000
3000
4000
stop
stop
step
Figure 30.3: Gradient descent central flow for CNN with MSE loss, 17 = 0.02.
real vs. predicted escillation covariance le-5
train
top hesslan aigenvalues
6.35
gradient flw
400
empiniaal varance
gradient (nverage)
300
along E11 eigerivectors
500 eigenvollass
0.25
flw central flo prediction
2
200
giadient descert central flow gradient flow (top)
0.15
0.19
1000
2000
3000
4000
5000
8000
1000
2000
3000
4000
5000
5000
1000
2000
3000 step
4000
5000
6000
stop
gradient norm?
test accuracy (3)
distance to gradient de scent
2.0
grodient descent
0.4
contral
gradient descent (average)
gradient flw
1.5
central flw
0.3
central flw presiction
gradient descont (mridpointa)
--- central flw
0.0
0.0
1080
2000
3000 stop
4000
5000
8000
1000
2000
3010
4000
5000
6000
10000
2000
3000
4000
stop
stop
network outputs on test example
network outputs on test example 2
network outputs on test example
0.5
gradient ouscent
1.00
gradient descert
0.6
centrol flow
central flom
0.5
0.50
0.4
gradient descent central flow
0.00
0.1
-0.25
1000
2000
4000 step
3000
5000
6009
1000
2000
400c step
3000
5000
6000
1000
2000
3000
4000
5000
6000
step
Figure 30.4: Gradient descent central flow for ResNet with MSE loss, 77 0.01.
59

---

Published as a conference paper at ICLR 2025
real vs. predicted ed asion cov covariance 1c
train loss
hessian eigenvalues
200
gradiont flw
gradicnt descent centtal flw
ompitical variance
0.3
giadient descent
tignnvetooss
gradieent descent loverage)
200
gradiien flow (tup)
Xii sigunvalass.
central
prediction
0.1
0
1000
2000
3000
4000
5000
0
1000
2003
Sc00
4000
5000
0
1000
2000
3000
4000
5000
stop
stop
step
gradient norm?
tast accuracy (56)
distance to gradient descent
contral grodient
gradient descent laverage!
1.5
...
central now prediction
contral fow
0.6
55-
50-
(midppints)
0.2
0.0
45
0.0
1000
5000
10000
3000
step
step
step
outputs example
outputs example
outputs example
gradiett descent central fow
gradett descent contralfiow
0.6
0.5
0.2
0.0
grodient descent
central flow
-0.2
1000
2000
3000
4000
5000
5000
1000
2:00
3000
4000
5000
6000
1600
2000
3000
4000
5000
6000
Figure 30.5: Gradient descent central flow for a ResNet with MSE loss, 0.013333.
real vs. predicted oscillation covariance
200
tap hessian eigenvalues
1e-5
graolent flw
graolent sescont
empirical varlance
descent
150
central
G-
aking 21t) eigeriectioss
210 sigansalues
0.2
central flow prediction
100
1
2000
3000
4000
1000
2000
3000
4000
10000
2000
3000
4000
gradiont norm
test accuracy (%6)
distance to gradient descent
flow gradicnt flow
65
1.0
0.75
gracient descect
(average)
55-
0.50
0.5
Bom
gradieet descent (midpainta)
central flom prediction
50
central flow
0.0
0.00
1000
2000
4000
1000
2000
4000
2000
3000
4000
stop
stop
step
network outputs on test example
network outputs on test example 2
network outputs on test example
gradient descent
gradient descunc flo
0.8
gradient descent
central
0.6
0.4
0.2
0.2
0.0
6.2
-0.5
0,2
1000
2000
3000
4.00
1000
2000
3000
4000
1000
2000
3000
4000
step
step
Figure 30.6: Gradient descent central flow for ResNet with MSE loss, 7) 0.02.
60

---

Published as a conference paper at ICLR 2025
real vs. predicted escilaaion covariance 1e: 5
train loss
top hessian eigenivaluss lues
0.35
400
gredientt descent
0.30
300
conital flw
gradiant now (top)
0.25
gradient How
200
0.20
central flww
gradient descert (uverage)
100
2
empirical vorance
sigerivectons
6.15
centrol flow prosiction
200 elgeniaalass
0
2000
4000
6000
8000
2000
40000
50000
8000
2000
4000
6000
2000
step
step
step
gra dient norm?
test accuracy (1)
distance to gradient descent
1:25
gradient descent
central
cential (overage)
60
1.00
central flow prediction
0.75
50-
0.50
gredient (midpoints)
0.25
central flw
0.00
6000
stop
stop
stop
network outputs test example
network outputs test example
network outputs test example
0.5
gradient descent
grostent descent
1.00
gradient descent
central nlw
central flw
0.75
central flow
0.4
0.50
0.25
0.25
0.00
0.00
-0.25
2000
4000
6000
2000
0
2000
4000
60000
8000
2000
4000
6000
8000
stop
stop
step
Figure 30.7: Gradient descent central flow for VIT with MSE loss, = 0.01.
train loss
top hessian eigenvallues
real vs. predicted osdilaaion on covariance
0.35
0.00012
empirical
gradient descant
250
contral
0.00010
along 201 wigerowctors.
gradient descent laverage)
grodient flow (top)
0.00008
EID eigervatues
0.25
central flow
central flow prediction
150
0.00006
6.20
3,00004
50-
0.00002
2000
4000
6000
8000
2000
4000
6000
3000
2000
4000
6000
8000
gra gradient norm?
test accuracy (3)
distance gradient descent
gradient descent
1.2
central flo
gradient descent lnveragel
gradient flow
central flo
60-
1.0
control flow prostction
0.5
0.5
50-
0.4
gredient descent (midppinis)
centrol
0.0
2000
4000
6000
8000
2000
4000
6000
5000
2000
4000
6000
8000
step
step
network outputs test example 1
network outputs on test example 2
network outputs on test example
0.5
grodient descent
1.00
gradientt descont
groslent descent
central flow
central flow
0.8
central fow
0.75
0.6
0.2
0.50
0.4
0.25
0.2
0.00
0.0-
-0.25
2000
4000
600D
2000
2000
4000
600D
8000
2000
4030
6000
8000
step
step
stop
Figure 30.8: Gradient descent central flow for VIT with MSE loss, 17 0.0133333.
61

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenvalues
real Vs. predicted oscillation covariance
200
0.35
gradient flww
gradient
gradiet descent central gradient (top)
0.00015
emplical along zici eigervectors Eith eigenvaluts
0.30
gradlent laverage)
0.25
cerdial fow prediction
central flow
100
0.00010
6.20
50
0.60005
0.15
0.00000
0
1000
2000
3000 4000 5000 stop
6000
1000
2000
3000
4000
5000 6000
1000
2000
3000
4000
5000
G00D
step
stop
gradient nomm
test accuracy (%)
distance gradient descent
gradieen sscent
65-
1.2
contral flow
gradient dieseet (average)
60-
1.0
gradient flom
Cow presictio
55
0.6
50-
0.6
gradiant descent imiápoinss
--- centrol flow
0.2
0.0
1000
3000
4000
5000
5000
2000
4000
5000
2000
3000
soco 6000
step
step
step
network outputs on test example
notwork outputs on example
network outputs on example
1.0
gradient a sscent
gradlent descent
graolent descont
0.5
central flow
1.0
central faw
0.8
cond Bow
0.6
0.4-
0.5
0.0
0.2
0.0
0.0
-0.2
2000
4000
5000
2000
4000
5000
4000
step
step
step
Figure 30.9: Gradient descent central flow for a VIT with MSE loss, 1) = 0.02.
train less
top hessian eigenvalues
real vs. predicted oscillation covariance
gradiett
0.00030
gradient descant
125
central
along 301) wigenivectors
0.3
gradiant descent (average)
grodient flow Htpp
eigenvalaes
contral flow
0.00020
central nlw prediction
75
0.00015
0.2
0.00010
0.1
25-
0.00005
1000
2000
3000
4000
5000
6000
1000
2000
30000
4000
5000
6000
1000
2000
2000
4000
5000
6000
step
gradient norm?
test accuracy (%)
distance to gradient descent
gradicnt descont
now gradient frw
2.0
gradient descent lovenge)
central flow
80
1.5
centrol flow prediction
1.0
0.5
60-
gradient descent (miqpoinhs) control
1000
2000
3000
4000
5oco
6000
1000 2000 3050 4000 5000 6000
1000
2000
2000
4000
5000 6000
stop
step
network outputs on test example
network outputs on test example
network outputs on test example 3
0.8
gradicnt descont
grodient descont
gradient descent
0.75
... central flow
0.5
-.-
central Sow
0.75
centrol flow
0.50
0.50
0.25
0.0
0.00
0.00
0.2
0.25
0.25
10000
2000
3000 stop
4500
5000
6000
1000
2000
3000
4000
5000
8000
10000 2000
3000
5000 6000 step
4000
stop
Figure 30.10: Gradient descent central flow for LSTM with MSE loss, 17 0.02666.
62

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
1.00
gradivnt flow
gradient descent
0.0006
empirica verience
0.5
gradient
gradient descent (average)
gradiant flow (Cop)
central
0.0905
slong 3i5 cigenvecters
eigenvalues
central
0.0004
0.2
central flow prediction
40
0,0002
20-
0.0001
0.0000
1000
2000
3000 step
4000
5000
6000
1000
2000
3000
4000
5000
5000
1000
2000
3000
4009
5000
6000
step
step
gradient norm?
test accuracy (%)
distance to gradient descent
2.0
100
central flw gracieet
1.5
gradient
15-
gradient laveraee)
1.0
central flow
60-
centtal Row prediction
0.5
40
gradient descant (miopoints)
0.5
control
0.0
0.0
1000
2000
3000
40000
5000
6000
1000 2000 3000 4000 5000 5000
1000
2000
3000
4009
Se00
6090
step
step
step
network outputs an test example
network outputs on test example
natwork putputs on test example 3
gradient d:scant central
0.75
0.75
0.50
gracient descore
0.2
g gradintt descent contral/ow
0.25
contral flow
0.0
0.25
0.00
-0.2
0.00
-0.4
1000
2000
2000
4000 stop
soco
80c0
1000
2000
3000
4000
5000
5000
6
1000
2000
2000 4600 sed0 6000 step
step
Figure 31.11: Gradient descent central flow for LSTM with MSE loss, 17 = 0.04.
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
gradiett flom
desceet
0.00125
gradient discant
40-
gradient descont (average)
gradient flow (tp)
0.00100
contral flow
50
empiriaal variance eigenvectors -.. 50tt eigencalues
0.2
central flow presiction
0.00075
20
0.00050
0.1
10
0.00025
1000
2000
3000
4000
5000
6000
1000
2000
30000
4000
5000
6000
1000
2000 3000
4000
5900 5000
gradient norm?
test accuracy (%)
di distance to gradient descent
1.2
3-
1.0
80-
0.8-
contral flow
0.6-
gradient sescent
60
gradient flow
0.4
gradient descent (avenge)
0.2
central flow central presictioon
40
50
central
gracient descent (midgeints)
1000
2000
3000 step
4000 5000 6000
1000 2000 3000 4000 56.0 6000
1000
2000
3000
4000
5900 6000
stop
network outputs on test example
network outputs on test example
network outputs on test example
1.00
1.00
gradent central Howm
0.50
0.4
0.50
graciant descent
0.2
gradient descant
contral flow
0.0
contral flow
0.00
0.00
-0.25
-0.25
1.000
2000
3000 step
4000
5000
3000
2000
4000 stop
3000
5000
6000
3000 2000
s000 4000 5000 6000 stop
Figure 30.12: Gradient descent central flow for LSTM with MSE loss, 17 0.08.
63

---

Published as a conference paper at ICLR 2025
real vs. predicted oscillation covariance 1e
train less
hessi an eigenvalues
200
gradient
empéical veriance
0.4
giadient descard
gradient descent laverage
150
--- 2tft eigenvalaes
along eigenvectors
0.3
central prediction
100
0.2
gradient descertt
central flow
0.1
gradient flow (top)
0000
2000
5000 step
4000
5000
5000
1000
2000
5000
4000
5000
6000
1000
2000
3000 stop
4000
5000
6000
stop
gradient norm?
test accuracy (%)
distance to gradient descent
100
gradient (mverage)
0.5
central Flom
0.8
central flow
0.4
... central flow prodiction
0.2
(midpsints)
0.1
flow
0.0
0.0
3000
5000
4000
2000
stop
stop
stop
network outputs test example
network outputs test example
network outputs test example
gradient descert
0.8
central flom
0.4
0.4
gradient descent
contral
0.2
0.2
0.0
gradint descent
-0.4
central flw
-0.2
-0.2
1900
2000
2000
4600
500e
6009
0
1000
2000
3000
4.000
5000
6000
1000
2000
2000
4000
50co
6000
stop
stop
Figure 30.13: Gradient descent central flow for a Transformer with MSE loss, = 0.02.
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
gradient descant
125
central
0.000125
gradient descent (average)
gradient fow (top)
0.000100
0.3
contral flow
empirical vatiance along eigenvertoes Zicl eigevialues
central prediction
75
0.000075
0.2
6,000050
0.1
25
0.000025
1000
2000
3000 step
4000
5000 6000
1000
2000
3000
4000
5000
6000
1000
2000
2000
4000
5000
6000
gradient norm?
test accuracy (%)
distance gradient descent
1.0
gradient oescent
gradent destent laverago)
0.8
gradient flow
0.8
central flow
central flow prediction
0.6
0.4
0.4
0.2
gradlntt descert (midpoinss central
1000
2000
3000
4000
5000
6000
1000 2000 3000 4000 5000 6000 step
1000
2000
2000
4000
Soce 6000
step
natwork outputs on test example
network outputs test example 2
network autputs on test example
0.4
0.6
0.2
6.6.
gradient descerk
gradient descent contralfiom
0.0
contral flow
0.2
-0.
0.0
gradiett delcert
-0.4
-02
-62
1000
2002
3000
4000
5000
6000
1000 2000
3000
4n3o 5000 6oco
1000
2000 3000
4000
5000 6000
step
step
Figure 30.14: Gradient descent central flow for Transformer with MSE loss, 77 0.02666.
64

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenvalues
real Vs. predicted oscillation covariance
0.4
gredient flow
100
gredient descent
0.00925
emplical var varince atong Hig digerpcctors
0.3
gradient
gradiant descent laverage)
gradient flow (top)
flow
0.00920
eigehvatues
flow centrall flow prediction
0.00015-
0.2
40
0.00011
0.1
20
0.00005
0.0
0.00000
1000
2000
3000 step
4000
5900
6000
1000
2000
3000
4000
5000
6000
1000
2000
2000
4000
5000
6000
step
step
gradient norm?
test accuracy (%)
distance to gradient descent
1.0
gradent
100
central flow
gradhent
gradient
0.8
central flow
1.25
central flow prediction
1.00
0.6
0.75
0.4
0.50
0.2
gradiant descent imlopoints)
contral flow
0.5
0.0
0.00
1000
2000
3000
4000
5900
5000
0
1000
2000 3000 4000 5000 5000
0
1000
2000
3000
4000
5000
6090
stop
stop
network autputs on test example
network outputs on test example 2
network outputs on test example 3
0.6
gradient descont central
0.6
0.8
0.4
0.4
0.2
gracient descent
0.4
gr gradint descent contral
central flow
0.0
0.2-
-0.2
0.0-
-0.4
-02
-0.2
1000
2000
3000
4000
5000
0000
0
1000 2000
3000
4000 5000 6000
1000 2000
3000 4000 5000 6000 stop
stop
Figure 30.15: Gradient descent central flow for a Transformer with MSE loss, = 0.04.
real vs. predicted oscillation covariance 1w-6
train
top hessian cigenvalues
800
0.30
grodient flow
gradlent descont central flow gracient (top)
gradient descent
600
lnverage)
6.25
flow central flw prediction
emprical variance along 2/9/ eire
400
516) wigervakess
6.20
2
1000
2000
3000
4000
5000
5000
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
500c
6000
gradient norm?
test accuracy (9)
distance to gradient descent
68
grodient descent
contral flow
gradient descent (average)
66-
0.4
graolent flww
cential flw
centrol flw presiction
64
0.3
62-
60-
gradieet descent (midpointa)
0.1
central flow
58
0.0
1080
2000
4000
5000
5000
1900
2000
3010
4000
scon
8000
1000
2000
3000
4000
5000
step
stop
stop
network outputs on test exampte
network outputs on test example 2
network outputs best example 3
1.00
0.6
gracient descent central
0.6
gradient descent
gracient t descent
0.50
contral flow
control flow
0.2
0.25
0.0
6.00
0.0
1000
2000
4000 step
3000
5000
6000
1000
2000
3000
4000
5000
6000
1000
2000
3000 step
4000
5000
6000
step
Figure 30.16: Gradient descent central flow for Mamba with MSE loss, 17 0.005.
65

---

Published as a conference paper at ICLR 2025
real vs. predicted oscilaaion covariance
0.275
gredient flow
train loss
400
top hessian eigenvaaues
gradkent descent
20
1c 5
emplicaal variance eigeniecctos Sitt eipencalaes
6.250
gradient descent (average)
gradienn descant
300
cantral flow
gradieet flow (top)
1.5
control
0.200
prediction
200
1.00
0.150
1000
2000
3000
4000
5000
6000
1000
2000
3000
4600
5000
6000
1.000
2000
3000 step
4000
s000
6000
step
step
gradient norm?
test accuracy (%)
distance to gradient descent
68
0.5
central flow gradient
1.5
66-
gradient descent
gradient desrent (everage)
62
gradient descenti (midpoints)
flow prediction
0.0
0.0
3000
4000
5000
1000
2000
3000
4000
5000
6000
4000
5000
step
stop
step
network outputs test example
network outputs test example
network outputs test example
1:00
0.8
grsdlent descent central flom
gradlent descent
central flow
0.75
0.6
grad lent descent
0.50
contral flow
02-
0.25
-0.2
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
6000
1000
2000
2000 step
4000
50co
6000
stop
step
Figure 30.17: Gradient descent central flow for a Mamba with MSE loss, = 0.01.
real vs. prodicted oscillation covariance 1--5
train
top hessian eigenvalues
200
gradient flow
IM gradient descont
empirical voriance
0.25
gradient descent
150
contral flow
along digenvecters 2 sigenccllass
gradient descent
gradient flow (top?
0.20
central flow prediction
100
6.15
0.10
1000
2000
t0n0
4000
sonD
60c0
1000
2000
3000
4000
5000
6con
1000
2000
1000
4000
5000
6000
step
step
step
gradient norm?
test accuracy (36)
distance to gradient descent
70.0
central gradiant
1.00
67.
0.75
65.0
gradent descent
62.5
gradient
0.25
60.0
gradient descont (midpoints)
central flow prediction
central Flow
0.00
0.0
1000
1000
4000 5000
3000
4.000
5000
step
step
step
network outputs on test example
retwork outputs on example 2
network outputs test example
1.00
1.00
gracient descert
1.00
gratient descent
6.75
0.75
central Row
0.75
central flow
0.50
central flow
gradient descent
0.50
0.50
6.25
0.25
0.25
0.50
1000
2000
4000 step
3000
5000
6000
0
1000
2000
3000
4000
5000
6000
1000
2000
3000 step
4000
5000 6000
Figure 30.18: Gradient descent central flow for a Mamba with MSE loss, 7) 0.02.
66

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenivalues
real Vs. predicted oscillation cova covariance
800
1.25
gradieett flow
gradiiett
gradient descent
flow
0,000
1,00
gradient descent laverage)
gradient flow (top)
flow
0.75
contral flw precktion
400
0,00030
6.50
200
0.00005
empirial variance sigenvecters 2141 eigenvvaluss
6.25
6.00
0,00000
1000
2000
3000 step
4000
5000
6000
1000
2000
3000 step
4000
5000
6000
1000 2000
3000
4000
5000
6000
step
gradient norm?
test accuracy (86
di stance to gradient descent
70
central flow gradient
50
30
gradient descent (unerage)
40-
10
control flow presiction
contral flow
30
centrol flow
gradient descort (midponts)
0.0
1000
2000
3000
4000
5000
6000
1000 2000 3000 4000
5000 6000
1000
2000
3000
4000
5000
6090
step
step
step
network outputs on test example
network outputs on test example 2
network putputs on test example 3
20
gradient descent
central
-5
10-
gradient descere
gradient descent
control flow
-10
contral flw
10
1000
2000
3000 step
4000
5000
6000
1000 2000
3000 4000 step
5000 8000
1000
2000
4000 stop
3000
Soco 6000
Figure 31.1: Gradient descent central flow for a CNN with CE loss, 77 = 0.005.
train loss
top hessi an eigerivalues
real vs. predicted oscillation covariance
gradient
gradient descent
0.00025
gradient descent
500
flw
0,00020
gradkent descent laveragel
gradient flw (top)
central flow
flow prediction
300
empitical varlence. along 209 elgenvectors ZIO digenssiues
100-
1000
2000
3000 step
4000
5000
5000
1000
2000
3000
4000 5000
6000
1900
2000 3000 4000
5000
6000
gr gradient nor??
test accuracy (%)
di stance to gradi ent descent
70
1.50
control
40
1.25
gradient flow
60
30
1.00
50
20
gradient descent gradient déscent (nverage)
0.50
flw prediction
flow
central fow
gradient descert (midpeints)
0.25
1000
2000 3000
4000
5000 6000
1000 2000 3000 4000 5000 6000 stop
1900
2000
3000
4000
5000 6000
stop
stop
network outputs an test example
network outputs on test example
network outputs test example 3
10-
5-
0-
-5
-30-
-5-
-20
gradlent descent
gradlent descent
-10
gradient descent
flow
-20
1.000
2000
3000
4000
5000
6000
1000
2000
3000 4000 stop
5000 6000
1000 2000 3000
5000 0000 stop
4000
stop
Figure 31.2: Gradient descent central flow for CNN with CE loss, 17 0.006666.
67

---

Published as a conference paper at ICLR 2025
train loss
top hessi an eigenvalues
real vs. predicted oscillation covariance
gradient few
400
gredert descent
0,0094
ampiicaal verience slong Z10) digenwecters Zift eigenvatues
gradient
cendral flow
gradlent descent (average)
gradient (top)
0.0003
central
canitrll flow prediction
200
0,0002
100
0,0091
0.0000
1000
2000
2000
4000
5000
5000
0
1000
2000
3000 step
4000
5000
5000
1000
2000
3000 step
4000
5000
5000
step
gradient norm?
test accuracy (%)
distance to gra viadient descent
70
central flox
30
60-
1.5
gradlent descent
50
1.0
central flow
gradient descant (average)
40-
gradient descent Imioponns)
0.5
contral flow prediction
central flow
0.0
1000 2000 2000 4000 5000
5000
1000
2000 3000
4000 5000 5000
0
1000
2000
3000
4000
5000 6090
stop
netwark outputs on test example
network putputs on test example 2
network putputs on test example 3
15
grodient descont control
5
10
-5
-5
-15
gradient descent
gracient descen
-20
centr flow
-15
-15
central flox
1000
2000
2000 stop
4009
Scou
5000
1000
2000
3000
4000
5000
5000
1000 2000
3000
4000
50c0 6000
stop
Figure 31.3: Gradient descent central flow for a CNN with CE loss, 17 = 0.01.
train
hessian cigervalues
predicted cscillation covariance
400
0.00010
1.25
gradlent flow
gradient descent
1.00
gradient descent (average)
gradlent descent
300
central finw
grediett flow Hop
0.00008
centrol
0.00006
empliical variance along I66 wigervectors
prediction
0.50
0.00004
E2t eigenvolucs
0.25
0.00002
0.00000
1000
2000
3000
1000
2000
4000
1000
3000
4000
step
step
step
gradient norm?
test accuracy
distance gradient descent
70
centrol flom
65-
2.0
gradient flow
gradiemt descun:
55
1.0
gradient descent laverage) central
50-
prodllnt imidpoints)
central preoiction
45
0.0
1000
2000
3000
4000
10000
2000
3000
4000
10000
2000
3000
4000
step
step
step
network outputs test example 1
network outputs on test example
network outputs on test example
7.5
10
5.0
2.5
-25-
gradient decent
-10
gradlent descent
-5.0
gredient descent
flow
1000
2000
3000
4000
10000
2000
3000
4000
1000
2000
3000
4000
stop
stop
stop
Figure 31.4: Gradient descent central flow for ResNet with CE loss, 1 0.01.
68

---

Published as a conference paper at ICLR 2025
train loss
top hessian elgenvalues
real vs. predicted oscillation covariance
300
1:2
gredient flow
gradient
250
gredient descort
flow
0.000156
1.0
gradient descant (average)
200
gradient nlw (top
0.000125
0.8
6.6
central flow
central flo prediction
150
0.000100
0.c00075
2011 cigervsiues
along eigenvectors
0.4
100
n.c00050
6.-
50-
0.000025
6.0
0.000000
500
1000
1500
2000
2500 3000
50e
1000
1500
2000
2500 3000
500 1000 1500
2000
2500 3000
stop
stop
stop
gradlent norm?
test accuracy (%)
distance gradient descent
central
gradient now
descent
55
gradient descant (averoge)
central flow flo prediction
gradient descont (midpaints)
centrol
500
1000
1500 2000 2500 2000
1000
2000
2000
500
1000 1500 2000 2500 3000
step
network outputs on test example
nctwork outputs test example 2
network outputs on test example
5.0
5-
-25
gradient
gracient desce
gradeet descent
-10
1000
2000
3000
10000
2000
3000
1000
2000
3000
step
step
step
Figure 31.5: Gradient descent central flow for ResNet with CE loss, 77 0.013333.
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
1.2
gradient
gradient flow gradient flo (top)
empinical eigenvectors SIC) olg cigorialucs
1.0
gradkent descant (average)
gradient descent
150
0.8
central flow
0.6
prodiction
100
0.4
0.0003
0.2
0.0
0.0000
1000
1500
2500
500
1000
1500 step
2000
2500
1500 step
2000
25:00
step
gradient norm?
test accuracy
distance to gradient descent
3.0
central flow
gradianttoo
85
2.5
80-
2.0
gradient gradient (overgg) central flow central flow predktton
55
1.5
1.0
gradient descent (midpoints)
central flow
0.5
45
0.0
500
1.000
1500
2000
2500
500
1900
1500
2000
2500
3000
500
1000
1500
2000
25:00
step
step
step
network outputs on test example 1
network outputs on test example 2
network outputs on test example 3
10
oradiant descort
5.0
2.5
2.5
-4-
-5
gradient descent
-5.0
gradiart descont
zentral
-10
flo
500
Lo00
1560
2000
2500
3000
500
1600
1500
2000
2500
3000
0
500
1000
1500
2000
2500
3000
step
step
step
Figure 31.6: Gradient descent central flow for ResNet with CE loss, n 0.02.
69

---

Published as a conference paper at ICLR 2025
train loss
top hessian elgenvvauus
real VS. predicted escillation tion covariance
400
gradient descent
emprical lence
1.2
flow
0.0005
along Zio cigenvectors
gradknn flw (top)
0.0004
eigenvalues
gradient flw
200
0.0003
0.8
gradient
gradiant descent laverage)
0,0002
0.6
contral flo prediction
contral nlw
100
0.0001
0,0000
0
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
step
step
step
gradient
test accuracy (%6
distance to gradient descent
gradient descent
central flow
20
1.0
wcisent
contral
contral flo prediction
0.6
10
grad lent descort (midponns)
0:2
control flom
1
20000
3000
4000
0
1000
2000
5000
4000
1000
2000
3000
4090
stop
stop
stop
network outputs on test example 1
network outputs on test example 2
network outputs an test example 3
gradient descent
gradient descent central
2-
gradient desceet
-3
centiaal Blw
1000
2:000
2000
4000
1000
2000
3000
4000
1000
200
3000
4000
step
step
stop
Figure 31.7: Gradient descent central flow for VIT with CE loss, = 0.01.
train loss
hessian oigenvalues
predicted oscillation covariance
300
0,0008
1.2
gradient descent
omglrical varlance eignnvertors
contral
1.0
200
gradient flow (top)
x00 sigenvaluus
graclent flo
0.0004
gradlent descent
gradient (average)
100
0.4
fll
0.0002
6.2
contral flo prediction
0
0.0000
2000
3000
10000
2000
4000
3000
step
step
step
gradient
best accurary 196)
distance gradient descent
30
graolent descont
centrol flow
gradlent (overgoy
ent
60-
20
central fow prddction
55-
0.8
15
50
45
gradiant descent Imidpoints)
0.2
contral flow
0
1000
20000
3000
4000
0
1000
2000
3000
4000
1000
2000
3000
4000
step
step
step
network outputs on test example 1
network outputs on test example
network putputs an test example 3
central
grasient descont
descont contral flow
descont
control
1.000
2000 stop
3000
4000
30.0
2000
3000
4000
1,000
2000 step
3000
4000
step
Figure 31.8: Gradient descent central flow for VIT with CE loss, 77 0.013333.
70

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
200
groolent Row
grodient Sescort
cmp empiriical veriance
gradlent
flow
slong Zith cli envectors
1.5
gradient laverage)
ttep)
0.0008 200 wigenvaluuss
flow
1.0
central flow prodiction
0.0004
50-
0.0
0-
0.0000
3000
3000
3000
Stop
step
step
gradient norm?
acuraccy (%)
to gradient descent
20
gradienn ocseont (ivergge)
gradiirtt ocscent
1.0
central flow
graolent flow
15
centrol
piediction
0.8
10
55-
0.6
50
0.4
contral Imioponis)
45-
0.0
20000
3000
1000
2000
3000
4000
2000
3000
step
step
step
network cutputs test example
network outputs test example
network test example
q cescent
gracient descent
gracient descent
5.0
control flow
5.0
control flow
5.0
central flow
2.5
0.3
0.0
0.0
2.5
2.5
-5.0
-7.5
-7.5
1000
2000
2000
4000
1000
2000
3000
4000
1000
2000
3000
4000
stop
Figure 31.9: Gradient descent central flow for VIT with CE loss, 77 0.02.
train
hessian aigenvalues
predicted oscillation covariance
300
g cradient flow
gradient descent
omgirical varlance eigenvecttes
1.2
gradient descent
contral flo
1.0
gradient descent
200
gradient now thop!
200 siganvalues
0.00020
flow predictio lion
0.00015
0.6
100
0.00010
0.4
0.00005
0
0.00000
3000
5000
1000
2000
6000
2000
4000
step
step
step
gradient
best accurary 196)
distance gradient descent
graalent descont
central flow
(overgo))
central flo prediction
60-
0.6
50
0.4
40
gradiant descent Imidpoints)
0.2
50
contral flow
0
1000
2000
3000
4000
S000
6000
0
1000 2000 3000 4000 5000 6000
1000
2000
3000
4000
5000
6000
step
step
step
network outputs on test example 1
network outputs on test example
network putputs an test example 3
2
central fom
central
1
1000 2000
3000 step
4000
5000
6000
1000
2000
3000
4000
5000
6000
1000 2500 3000 4000 5000 6000 step
step
Figure 31.10: Gradient descent central flow for LSTM with CE loss, 17 0.01333.
71

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenivalues
real vs. predicted oscillation covariance
2000
1.25
gradentt now
gradant descent
gracient descers
0.0006 Z19 eigenvalues
emptiiall veriance along 2081 eigervecters
gradient descent (averape)
gracknt Sow (top)
central
0.75
central flow prediction
100
0.50
0.25
50
0.0002
0.00
0.0000
1000
20000
30000 step
4000
5000
6000
0
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
6000
step
step
gradient norm?
test accuracy (%)
distance to gr gradient descent
central Blo
gracient
90
2.0
descent (average)
central
00-
central flww prediction
1.0
gradiant descent Imiopointi
contral flow
0
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000 6000
0
1000
2000
3000
4000
5000
6000
step
step
step
network outputs on test example 1
network outputs on test example 2
network putputs on test example 3
pradient descent
descent
graclent Sescont
centrol
central
-2-
2-
1000 2009
3000
4000
5000
6000
1000 2000
3000
4000
5000
5000
1000
2000
3:000
4000
5000 6000
step
step
Figure 31.11: Gradient descent central flow for a LSTM with CE loss, = 0.02.
train loss
top hessian eigenvalues
real vs. predicted oscilation covariance
2.5
destent nlw
gradeet descent
0.0015
gradient descent (ovarags)
gradient flow (oo)
centrol flo
80
empaiic verience eigerwectors Elcl eigenvalles
prediction
0.0010
1.0
40
0.5
0.0005
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
6000
step
gradient norm?
test accuracy (3)
distance to gradient descent
1.0-
gradient now
0.8
gradient descemt
gradient descent laverage)
0.4
central flow central prediction
40-
gradient dincert (m (milpaints)
flom
1000
2000 3000 4000 5000
6000
1000 2000 3000 4000 5000 60C0 step
1000
2000
3000
4000
5000
6000
stop
network outputs test example
network outputs on test example
network outputs on test example
groslent centri flm
gr gradient descent control flo
gracient descont contral flo
-2
1000 2000
3000 step
4000
5000
6000
1000 2000
3000 step
4000
5000
6000
1000 2000 3000 4000 step
5000 6000
Figure 31.12: Gradient descent central flow for LSTM with CE loss, "7 0.04.
72

---

Published as a conference paper at ICLR 2025
train loss
top hessian elgenivalues
real vs. predicted oscilationn covariance
2,5
gredientt flow
400
gradientt descent
engilicall verieece sling sid dgghvectats
2.0
descent
0.00015
gradient descont (average)
gradient now (top)
eignvvales
1.5
central
centtall flow prediction
200
0.00010-
1.0
0.5
100
0.00005
0.00000
1000
2000 step
3000
4000
1000
2000
2000
4000
1000
2000
3000
4000
step
step
diadient norm?
test accuracy (16)
distance to gradient descent
gradient descent
300
0.8
central flow
gtadient
gradient
contral flow
0.6-
contral nlw prediction
60
0.6-
gradient descent (middoinnss)
20
--- control flw
0.0
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
stop
stop
stop
network outputs on test example
network outputs on test example 2
network outputs on test example 3
gradiont descent
flw
2-
-1-
gradient descent central flww
gradient descent central flow
1000
2000
3000
4000
10G0
2000
3000
4000
0
1000
2000
3000
4000
stop
step
stop
Figure 31.13: Gradient descent central flow for a Transformer with CE loss, 0.01.
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
2.0
gradient descant
250
central
0.68225
aling sigettoectors
gradient descent laverage)
gradient flow (top)
0.00020
201 eigenvalues
1.5
central flow
central flow prediction
150
0.00015
1.0
0.00610
0.5
50
0.00005
1000
2000
3000
4000
1000
2000
2000
4000
1000
2000
2000
4000
step
step
gra dient norm?
test accuracy (36)
distance to gradient descent
100
10-
1.0
flow gradient flow
80-
0.6
0.6
gradient descent
gradient descent (nverage)
40
cenital flow centrol prediction
20
centrol
gredient descent (midpsinis)
0.2
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
step
step
network outputs on test example
network outputs on test example
network outputs test example 3
gradient descert
gradient descent
central flom
contral
-2-
-2-
gradient desiwor
10000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
step
step
Figure 31.14 Gradient descent central flow for Transformer with CE loss, 7 0.013333.
73

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenivllus
real vs. predicted escilation covariance
2,5
gradiect flow
200
gradient descent
0.0005
2.0
gredient
0.0004
gradiant descent laverage)
gradknnt flo (top)
central flow preckcion
flow
100
0.0003
atong aigenvectors
0:0002
eigenvalues
50-
0.000
0,0000
1
2.000 step
3000
4000
1000
2000
300
4000
1000
200 0
2000
4000
step
step
gradient
test accuracy (%)
distance to gradient descent
100
60-
grediant
gradient
gradient descent laverage)
contral flow contral nlo prodiction
grad lent descort (' (midponts) contral flom
20
1000
2000
3000
4000
1000
2000
3000
4000
1000
200
3000
4090
stop
stop
stop
network outputs on test example
network outputs on test example 2
network outputs an test example
grasient descont
or adadentt descent Gentral now
adleet discent
contralloow
-3
1000
20000
2000
4000
1000
2000
3000
4000
1000
200
3000
4000
step
step
stop
Figure 31.15: Gradient descent central flow for a Transformer with CE loss, 0.02.
real vs. predicted oscillation covariance
2000
top elgenvalues
14-5
gradient flow
gradient descent flow
gradient
1500
1,00
central flow prediction
1000
0.75
6.8
0.50
emplitaa voriance
ator g Eicl eigenwectors
0.25
5039 eigenvalues
1000
2000
3000
4000
5000
5000
1000
2000
3000
4000
s000
6000
1000
2000
3000
4000
5000
6000
gradient norm?
test accurary (36)
distance to gradient descent
68
descent
contral
gradient descent laverago)
66
grad lent flow
central now
15
contral nlo prediction
64
10
62-
60-
gr ucdlent descent (midpannss
--- central
0.0
2000
4000
1000
2000
3000
3000
stop
stop
stop
network outputs on test example
network outputs on test example
network outputs on test example
grodiant descent contral
-2
gradient descect
gradiert descent
1.000 2000
3000 step
4000
s000
0000
1000
2000
3000
4000
5000
6000
1000
2000
3000 step
4000
5069
6000
step
Figure 31.16: Gradient descent central flow for Mamba with CE loss, 0.002.
74

---

Published as a conference paper at ICLR 2025
real vs. predicted osdilation covariance
train less
1500
top hessian elgenivalues
1c
grecient flw
gradent descent
gradient descert
ceréral flow
1.0
gradient descent laveragel
1000
gradient flow (tas)
15
flow
0.8
flo prediction
1.-
emplrtica variance
wight 3010 eigenvalass
0.4
0
1000
2000
5000 step
4000
5000
5000
0
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
6000
step
step
gradient norm?
test accuracy (16)
distance to gradient descent
20
gradient descent
68
central gradient
grad (average)
15
central flw
65-
cantral now prediction
64-
0.6
10
62-
0.4
60
(micpontss)
0.2
contral flow
0.0
2000
2000 step
4000
5000
6000
4000
5050
step
step
network outputs test example
network outputs test example
network outputs test example
grodient descent
gradient descent
central flw
gradient descect
cential few
1000 2000
3000 step
4000
5000
5000
1000
2000
3000
4000
5000
6000
1000
2000
2000
4000
5oco
6000
stop
step
Figure 31.17: Gradient descent central flow for a Mamba with CE loss, 17 = 0.002666.
real vs. predicted oscillation covariance 14-5
top hessian eigenvalues
1000
1.2
graolent flow
empirical variance aling 110 eigetieectors
800
sott eigenvalass
contral flow
central flow prediction
2
0.6
gradlent descert central flom gradient
1000
2000
3000
4000
5000
sec0
1000
2000
3000
4000
5000
6000
1000
2000
3000 step
4000
S000
6000
gradient norm?
test accurary (36)
distance to gradient descent
68
contral
66
0.3
gradient flw
10-
graclent descent
(nverage)
62-
cenit flow
60-
gr grudient descent (midpainis)
cantral al flow prediction
---
central flow
58
0.0
1000
2000
Soon stop
4000
5000
8000
1000
2000
3000
600e
1000
2000
3000
stop
stop
network outputs on test example
network outputs on test example
network outputs on test example
gradient descent
graciennt descant central
contral
gradient descont central
-2
-2
1.000
2000
3000 step
4000
5000
0000
1000
2000
3000
4000
5000
6000
1000
2000
3000 step
4000
5060 6000
step
Figure 31.18: Gradient descent central flow for Mamba with CE loss, 17 0.004.
75

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
400
greddent flow
emppriicl variance
gradient descert
along 2it) sigervectors
gradient descart (average)
0,00010
Eiti eigenvalues
control flow
0.00008
centrol now prediction
200
c.n0co6
100
central flom
or gradient descert
0. 0.09004
0.2
gradient flow (top)
0.00002
0,00000
0
1000
2000
3000 stop
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
stop
stop
gradient narm?
test accuracy (95)
distance to gradient descent
12
gradient descent
75
central gradlent fow
gradient laverage!
10
central Now
central nlw predktioo
0.4
65
gradient descent (midpeints)
0.2
central flo
0
1000
2000
3000
4000
3000
1000 2000 3000
4000 5000
1000
2000
3000
4000
5000
step
step
network outputs on test example 1
network outputs on test example 2
network outputs on test example
gradient descent
grsdlent sescent
7.5
gradlent descont
Central now
central
5.0
cantrall flow
2.5
2.5
-2.
0.0
5.00
25-
1000
2000
3000 step
4000
5000
1000
20000
3000
4000
5000
0
1000
2000
3000
40000
5000
step
Figure 31.19: Gradient descent central flow for Mamba with CE loss, 17 0.01.
train loss
top hessian eigenvalues
real vs. predicted oscillation covariance
300
fow
gradlent descent
empirical varience alena 7i5t eigenvectess Eldl eigonvalues
250
0.00015
grad descent (average)
200
grostent
0.5-
centrall flow predictioe
central flow
150
0.00030
0.4-
100
5.00005
0.2
0.0
0
0.00000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
10:00
20:00
40.00
5000
stop
stop
stop
gradient
test accuracy (%)
distance gradient descent
10-
1.-
central flow
gradient flow
70
ad descent
descent
central flom
gradicnt descont (midpoirtsi
02-
cantral Bow prediction
caniaal flow
0.0
1000
2000
3000
20000
stop
stop
stop
network outputs example
network test example
network outputs test example
7.5
7.5
gradiont descent
central flow
5.0
central Flw
gradient descent
central flow
gradient descent
5.0
2.5
0.0
0.9
-2.
5.0
-8
-7.5
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000 stop
4000
5000
stop
stop
Figure 31.20: Gradient descent central flow for a Mamba with CE loss, 17 0.013333.
76

---

Published as a conference paper at ICLR 2025
train loss
top hessian eigenivaluee
real vS. predicted oscillation covariance
gradied flw
gradiect Cescent
emptical lence along 2in digenvectors
flw
gradient descont (average)
gradiart
eigenvalues
central flo precitton
1.0
0.00015
50-
0.0001 0
0.00005
0,00000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
step
step
step
gradient norm?
test accuracy (%)
distance to gradient descent
gradient descent
80
1550
central flow
contral
1.25
contral now prediction
100-
0.75
0.5
gradian descert imidponts)
80
contral flom
0.55
0.00
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4600
5000
step
step
step
networkk outputs on test example 1
network outputs on test example
network outputs an test example
10
descent
5.0
grostent descort
2.5
0.0
2.5
-5
5.0
-5
gracieet descent
contral flow
-7.5
L000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
stop
step
Figure 31.21: Gradient descent central flow for a Mamba with CE loss, 0.02.
77

---

Published as a conference paper at ICLR 2025
C.2 SCALAR RMSPROP
top hessian eigenvalues
vs. predicted oscillation covariance
stable flo
Scalor RM SProp
eme rical vorfance
0.4
Scalar RMSProp
central flom
along 2ith elg envectoors
(average)
(top)
210 eigenvaluez
centrol how prediction
2
0.2
0.0010
0.1
1-
0.0005
0,0000
2000
40:00 steg
6000
8000
2000
4000 step
6000
8000
2000
4:000
6000
8000
step
gradient
tap sigenvalues
effective stop n/vo
Scalar RMSFrop
Scalar RMSProp
0.04
Scalor RMESPro
central
flow prediction
0.02
Scalar RM SKrpP
central flow
stable flow (tap)
0.00
2000
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
8000
step
step
step
network outputs on test example
test accuracy (%)
distance Scalar RMSProp
0.6
60
central fow
0.4
1.00
50
0.2
0.0
40
0.50
-0.2
Scallar RMMSPrPp
central Bow
30
Scaler RMSPtop (midgeints)
central flow
0.25
2000
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
8000
step
step
step
Figure 32.1: Scalar RMSProp central flow for CNN with MSE loss, 77 =0.003, 32 =0.99, and bias correction.
78

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted escillation covariance
Mtable fow
Scalar RMRSProp flom
RMSProp (average)
(tog)
0.3
central flow predicion
0.002
0.2
0.1
0.001
smpiicaal variance along sigervecters 20/1 eigencell.ss
0.0
2003
4000 step
6000
8000
2000
4000 step
6000
8000
2000
4000
6000
8000
step
gradient norm?
top hessian eigenvalues
offective step size NVV
12
Scaler RMS5hop
10
0.06
0.04
RMS/vop
400
central flow
Scalar RMSProp (average)
200
central flom
Scalar RMMSProp
0.02
ceniral flow predictio
slable flow (op))
0.00
2000
4000
6000
8000
2000
4000
6COD
8000
2000
4000
6000
8000
stop
network outputs on test example
test accuracy (%)
distance to Scalar RMSProp
1.2
central
0.6
69-
1.0
50
0.0
0.4
Scalar RMSProP
Scalar RMSRrop (ntlopoints)
-0.2
central flow
centrol
2000
4000
60.00
8000
2000
4000
5000
8000
2000
4000
6000
8000
stop
Figure 32.2: Scalar RMSProp central flow for CNN with MSE loss, 1/ =0.066, 92 0099 and bias correction.
loss
effective hessian eigenva
predicted ascillation covariance
0.005
stable
RMShrop
0.5-
Scalar RMSProp
Scalor RMSPhop taveroge)
3
stable flow (top)
central flo
0.004
flow
0.003
0.2
flow prediction
2-
0.1
along IDD eigeneectors
0.0
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
gradient norm?
top hessian eigenvalues
effective step size r/v
300
RMSProp
4-
250
central
3-
200
0.100
150
0.075
Scalar RMSProp
Scatar RMSProp (averoge)
100
Scalor RMSProp
0.050
centrrll flo
central flow
2000
4000
800D
2000
4000
6000
2000
4080
5000
8000
network outputs on test ex example
test accuracy (%6)
distance to Scalar RMSProp
central
60
0.50
0.00
0.25
Scator RMSProp
40
Scalar RMSPrag (midgolntsi
centrall flo
centiral flow
2006
4000
6000
8006
0
2000
4000
6000
8000
2000
4000
6000
8000
stop
stop
Figure 32.3: Scalar RMSProp central flow for CNN with MSE loss, 77 =0.01, 1, 82 =0.99, and bias correction.
79

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted oscillation covariance
stabic flow
Scalar RMSProp
0. 0.00125
RMSProp
central now
0.30
RMSProp (everage)
(top)
centrel fow
0.25
central flo prediction
0.00075
along 200 eigenvectors
dgenwalues
0.00050
0.15
1-
0.00025
0.00000
500
1000
15500 st0p
2000
2500
3000
500
1000
1500
2000
2500
3000
500
1000
1500 stup
2000
2500 3000
step
gradient norm?
top eigenvalues
effective step size WYV
1.50
Scaler RMMPPmp
0.025
Sealer RMSSProp
Scaler IMSPrup (everage)
central flow
1.25
Flow
1.00
200
0.75
150
0.015
0.50
100
RMSProp
central nlow
0.25
50
stable flow (top)
0.000
0.00
1000
1500
2000
2500
500
1000
1500
2000
3000
500
1000
1500
2500 3000 step
2000
step
stap
network outputs tast example
accuracy cy
distance Scalar RMSProp
70
central flow
0.4
65
1.5
stable flow
1.0
0.2
0.5
RMSProp
RMSPhop Imiopoints)
contral
nlow
2000
2500
suic
2000
4000
6000
0.0
500
1966
1500
2500
stop
stop
Figure 32.4: Scalar RMSProp central flow for a ResNet with MSE loss, 17 =0.01, 82 0099, and bias correction.
loss
top affective hessian cigenvalues
vs. prodicted oscillation covariance
6.35
stablc flow
Scalar RNSProp central fow
RMSPmp
0.30
(average)
(top)
0.0035
cential
6.25
central fow prediction
0.0010
atong Zid sigenvecters eigenvalues
6.20
0.0005
0.10
5.0000
500
1000
1500 step
2000
2500 3000
500
1900
1500 step
2000
2500
3000
500
1000
1500 step
2000
2500
2000
gradiont
top hesslan cigenvalles
effective stop n/To
1.50
Scalor RMSProp
0.030
Scalor RMSProp central
Scalar RMSProp laverage)
1.00
central pred ition
150
6.75
100
Scalar RMSProp
0.25
50
central flow
stable flow Hep!
0.015
6.00
0.010
500
1000
1500
2000
2500
2000
500
1900
1500
2000
2500
3000
500
1000
1500
2000
25.00 2000
step
step
network outputs on test example
Lest accuracy (36
distance to Scalar RM5Prop
70
1.50
central flw
65
1.25
stable
0.3
60
0.2
55
0.75
0.50
Staler AM4Shop
50
Scalir RMSShop Imispnints]
cuntral flow
45
cwntrr flow
0.25
500
1000
1500 step
2000
2500
3000
1000
2000
3000
4000
5000
500
1000
1500
2000
2500
3000
step
Figure 32.5: Scalar RMSProp central flow for a ResNet with MSE loss, 17 =0.02, 83 =0.99, and bias correction.
80

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigervaluus
real vs. predicted oscillation covariance
0.35
stable flww
Scaler RMSProp flow ltop
emptrical verience
6.30
slong 3121 digenvectors
Scalar RMSProp (average)
flow
0.003
centir flow predk tin
2-
0.20
0.002
6.15
0.001
0
500
1000
1500 step
2000
2500
3000
500
1000
1500
2000
2500
3030
500
1000
1500 step
2000
2500
3000
gradient norm?
top hessian eigeniallues
effectivve step size GVV
1.5
Scaler RMSPhrp
150
0.05
Scaler RMS5hop
central
centrol flow pred ction
100
0.04
Scalar RMSProp
0.03
0.5
central flowr
Mtuble flow [top!
0.62
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3030
0
500
1000
1500
2000
2500
3000
Stap
natwork outputs on test example
test accuracy (36)
distance to Scalar RMSProp
0.6
1.25
central
1.00
60
0.50
Scalor RMSProp
50
Scalar RMSProp (midpointsi
0.25
central
central
45
50c
1000
1500
2000
2500
3000
1000
2000
3000
4000
5000
500
1000
1500
2000
2500
2000
Figure 32.6: Scalar RMSProp central flow for a ResNet with MSE loss, =0.03, 82 0099 and bias correction.
train
effective hessian eigenvalues
real predicted oscillation coveriance
RMEPPro
Scalar RMMPrrp
central flow
0.0025
Stalar RMSPrrp loverage)
stablc nlo (top)
flow
prediction
0.0015
6.250
along B(1) eigeniectios
6.225
0.0005
20st
0.0000
1000
2000 step
4000
1000
2000
3000
1000
2000
3000
gradient norm?
top hessian eigenvalues
effective step size nivi
RMSPhop
RMSProp
contral flow contral flo prediction
RMShhop (mverage)
500
stable flow Itop
central flow
0.05
--. central
400
0.06
300
200
0.5
1.00
2000
3000
4000
1.000
2000
3030
4000
1000
2000
3000
4000
network outputs on test example
test accuracy (%)
distance to Scalar RMSProp
1.0
RMSProp
central
60-
6.8
0.6
0.6
0.1
Scolar RMSSPop Imidpaints)
central flow
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
stop
stop
St0p
Figure 32.7: Scalar RMSProp central flow for a ViT with MSE loss, 77 =0.01, 83 =0.99, and bias correction.
81

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted esciliation covariance
6.350
siabli flow
Scalar RMSProp
0.325
Scalor RMSProp
central
RMSProp laverage)
ttop)
flow
central flo prediction
2-
0,250
1-
emptrical vartance
0,225
Zist dienvvauuss
0.000
1000
2000 step
3000
4000
1000
2000
3000
4000
1000
2000 stap
3000
4000
step
top hessian eigenvalues
affective step WVV
Scaler RMShhpp
Sc wier RMSPPrp
0.08
Sraler RMtSProp
Scaler RMSihop (sowrage)
600
centraal flow
central flow
stable
0.06
300
0.04
200
100
1000
3000
10000
2000
4000
10000
3000
4000
stap
step
step
outputs example
test accuracy
distance RMSProp
Scalar RMSPmp
central
central Flow
60
stable flw
0.6
0.-
50
0.1
(mkoppinto)
flow
0.0
2000
5000
3000
stop
Figure 32.8: Scalar RMSProp central flow for ViT with MSE loss, 17 00002 82 00.99, and bias correction.
top effective hessian cigenvalues
vs. predicted oscillation
6.35
stable fow
Scater RMSProp
Scalor RMSProp
flw
0.005
emplrical variance along 205) digenvecctrs
0.30
(sverage)
0.064
261 egenvvlless
central flw prediction
0.25
0.002
0.001
0.000
1,00
2000 step
3000
4000
1000
2000 step
3000
4000
1000
2000
2000
4000
step
gradi ent
top hosslan cigenvalles
effective stop ntV
300
Sca RMSProp
Sca RMSPmop laverage)
250
0.10
Scalor RMSProp
central
200
0.68
centtel pred ition
150
0.06
100
Scalar RMSProp
central flowe
stuble flww [top!
1000
2000
3000
4000
1000
2000
3000
4000
0
10000
2000
3000
4000
step
step
network oulputs on test example
test accuracy (%6
di stannce to Scalar RM4Prop
0.6
Scaler RMSShop
0.8
central fow
60
stable
0.6
0
0.2
Scaler RM5Prop (middcintt)
0.2
central flow
10000
2000
3000
4000
3000
2000
3000
4000
3000
2000
3000
4000
step
step
Figure 32.9: Scalar RMSProp central flow for a ViT with MSE loss, 77 =0.03, 83 =0.99, and bias correction.
82

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigervaal
real vs. predicted escillation covariance
0.4
Mtable few
Scaler rURSihopp flom
Scalor
0.003
0.3
Scalar RMEPrrp (avorage)
(top)
central flom
centrol flom prediction
0.002
0.2
0.001
empitical variacee aiong 219 eigenvectors
0.1
210 sigarsaatess
2000
4D00
6000
8000
2000
60000
6000
8000
2000
4000
6000
8000
gradient norm?
top hessi an eigenvaalues
effective Step size N/V
2.5
Scalor
400
0.25
RMSProp
Scalar laverage!
2.0
contral now
1.5
contral now prodiction
1.0
0.10
Scaler RMSSrrp flom
0.5
100
stable (top)
0.0
0.00
20:00
8000
stop
stop
network test example
test accuracy (%)
distance Scalar RMSProp
0.8
central flow
0.6
80-
2.0
stable flom
0.4
Scaler RMSPPop
1.5
0.2
flow
60
0.0
-0.2
--- central flow
Scaler RMS//op (midscints)
0.5
-0.4
0.0
4000
5000
2000
4000
8000
2000
stop
5top
stop
Figure 32.10: Scalar RMSProp central flow for LSTM with MSE loss, 17 =0.01, (3 -0.99, and bias correction.
top effective hossian eigenvalues
vs. prodicted ascillation covariance
stabic flow
Scalor RMSProp
Scallrr RMSPmp
central flw (top)
0.3
(average)
flow
0.2
central rlo prediction
2
0.004
atong Zid aigenvecters
eigenvalues
0.1
1
6.002
0,000
2000
4000 step
6000
8000
2000
4000
6000
80000
2:000
4000
6000
8000
step
step
gradient
top hessian eigenvalues
effective stop nVv
2.5
200
Scalor RMSProp central
2.0
150
1.5
1.0
Scalor RMSProp
100
Scalar RM5Pmp (sverage)
Scaler RMSProp
0.5
central flow
central flow
0.1
conttal flow preciction
stable flow (tap)
0.0
0.0
2000
4000
6000
8000
2000
4000
600D
8000
2000
4000
6000
8000
step
step
network outputs on test exa mple
test accuracy (%)
distance to Scalar RMSProp
0.75
100
1.50
central flw
stable
1.25
RMSProp
1.00
centraa flow
60
0.75
0.00
Staler #MtShop (mispoints)
0.25
central flow
0.25
2000
4000
6000
8006
0
2000
4000
6000
8000
2000
4000
6000
8000
step
step
step
Figure 32.11: Scalar RMSProp central flow for a LSTM with MSE loss, 77 =0.02, .3 =0.99, and bias correction.
83

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian elgenv alues
real vs. predicted oscillation covariance
stable Row
Scaler RMSProp nlow
0.3
RMShhop
0.010
Scalar RM5Prop laverage!
stable itpp!
central flow
0.008
empiricca variance EIlt
0.2
central flew prediction
21th eigervalues
0.004
0.1
0.002
0.000
2000
40000
6000
8000
2000
60000
60.00
8000
2000
4000
60.00
8000
gradienk norm?
top hessian eigenvallus
effective step size nNV
0.6
RMSProp
1.5
0.5
1.0
Stalor RMSProp
40
Scwier RMShrop (nverage]
Sealer RMSProp flow
0.2
prediction
stable Htp!
0.1
0.0
0.0
6000
4000
st3p
stop
stop
network test example
accuracy (%)
cistance Scalar RMSProp
1.00
150
central flow
90
stable flom
1.00
Staler AMSirop
0.2
0.0
0.50
-0.4
0.2
50
40
Scoler IMMSirop Imidpointa)
central flww
0.00
2000
4000
2000
4000
6000
stop
5top
stop
Figure 32.12: Scalar RMSProp central flow for LSTM with MSE loss, 17 =0.03, (32 -0.99, and bias correction.
top effective hessian eigenvalues
vs. predicted osciliation covariance
stable flow
Scalor RM SProp central flom (tap)
0.4
Scalar RMSProp
(average)
0.3
cen tral how prediction
2
0.2
empirical variance eigenvectors 501 sigenculless
0.1
1-
0.0002
0,0000
2000
40:00 step
6000
8000
2000
4000 step
6000
8000
2000
4000 step
6000
8000
gr
tap hessian eigenvalues
effective stop n/vo
Scalar RMSProp
Scalor RMSSPro central
Scalar RMSProp
300
central flow prediction
200
Scalar RM SPrpp
0.02
central flow
stable flow (tap)
2000
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
8000
step
step
step
network outputs on test example
test accuracy (%)
distance Scalar RMSProp
0.4
central flw
0.2
80
60
-0.
Scaliar RMMShhp
20
Scaler RM5Phop (midgeints)
-0.4
central Bow
central flow
2000
4000
6000
8000
2000
4000
6000
8000
2000
4000 step
6000
8000
step
step
Figure 33:13: Scalar RMSProp central flow for a Transformer with MSE loss, 77 =0.01, 32 =0.99, and bias correction.
84

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted oscillation covariance
Vladle flow
Scaler RM SPrup
0.0020
0.4
Scaler INSPPrpp
flow
Scalar RMSProp (average)
stable itop)
0,0015
0.3
central
---
central flow prediction
2-
0,0010
0.2
ompirical variance along EXt eigenveactors 2/1 wigervaluse
1-
0.0005
0.1
2000
40000 step
6000
8000
2000
4000 step
60D0
8000
2000
40000 step
6000
8000
Or radiennt norm?
top hessian eigenvalues
effective stop NVV
0.8
Sealer RMSSPrpp
IMS5Prp laverage)
200
central
Scaler RMSSPop
flow
0.6
central flow prediction
150
Scalor RM SProp
0.2
50
central from
viable flow (tpp)
2000
40000
6000
8000
2000
4000
6000
8000
2000
40000
6000
8000
stop
network outputs on test example
test accuracy (%)
distance to Scalar RMSProp
100
0.4
central
0.2
0.0
60
Scalar RMMPPop
Scalar RMSProp (midooints)
central flow
20
central
2:000
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
8000
stop
Figure 32:14: Scalar RMSProp central flow for a Transformer with MSE loss, 17 0000 83 -0.99, and bias correction.
train
top effective hessian eigenvalues
vs. predicted escillation covariance
stable flaw
Sca br RMS Prop
0.004
0.4
Scalor AMSProp
central
0.3
central (average)
0.003
central flo prediction
stong zidi sigenvecters eigenvalaes
0.2
0.1
0.001
0.0
2000
4000
6000
8000
2000
4000
5000
8000
2000
4000
6000
8000
step
step
step
gradient norm?
tap hessian sigenvalues
affective step size niV'
0.8
120
Scalor RMSProp
0.08
central flww
0.6
80
60
0.05
Scaler RMSProp (nverage)
Scalor RMSProp
0.04
ceniral flow
central flow prediction
20
central flow slable flow [tpp)
0.0
0.02
2003
4300
abco
8000
2000
4000 step
600D
B00D
0
2000
4000
6DOD
B000
step
step
network outputs on test example
test accuracy (%)
distance to Scalar RMSProp
0.6
Scaler 194/hop
100
3.0
central now
flow
stable
0.0
60
1.5
-0.2
-0.4
central Blw
Scelar RMMSProp (ntippanth))
0.5
2000
4000
60:00
8000
2000
4000
80C0
RCOD
2000
4000
6000
8000
step
step
step
Figure 32:15: Scalar RMSProp central flow for a Transformer with MSE loss, 77 =0.03, 32 =0.99, and bias correction.
85

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigervalluss
real vs. predicted oscillation covariance
0.40
stable fow
Scalar RMSPhop
0.0012
0.35
Scalar RM5itop
Scalar AMSProp laverage)
noo (top)
central flow
0.0008
empirical voriance B11
--- carital flow prediction
2-
0,0006
wigenvall Laws
0.25
0.0004
0,0002
0.0000
0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4.000
5000
step
step
step
gradient norm"
top hessian eigeniva lues
effective step size NNV
0.085
Scaler RMSPhup
Scwler RM SPrp flow
Scaler RMSShop
IMSihup laverage)
800
0.030
central
central
stable nlow (top)
0.025
contral flow pred ition
600
0.020
400
0.015
200
0.030
0.905
0
1000
2000
3000
4000
5000
10000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
200p
Stop
stop
network outputs on test example
test accuracy (36
distance to Scalar RMSProp
1.0
0.25
central
0.20
0.6
Scalur RMSProp
0.35
central flow
62-
0.30
0.2
60-
Stalo or RMSPrup (m (midponts)
0.05
0.0
flow
58
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
step
Stop
stap
Figure 36: Scalar RMSProp central flow for a Mamba with MSE loss, 7) =0.007, B2 =0.99, and bias correction.
loss
top affective hessian cigenvalues
real vs. predicted esciliation covariance
stablc flow
Scalar RNSProp
empirical variance slorg itd aigetwecters
0.30
Scalor RMSPmp
central now
0.0020
cential (average)
ttpp
2001 eigenvalass
0.25
0.0015
central fow prediction
3.0005
0.15
5.0000
1000
2000
3000 step
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4:000
5000
step
step
gradiont
top hessian cigenvalues
affective step n///
2.0
Scalor RMSProo
0.04
central
0.33
1.0
RMSProp
Scalar RMSimm (average)
Scalor RMSProp
0.02
centiaa flow
100
central flow
central few pretiction
stable flow Hop!
0.91
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
step
step
network outputs on test example
test accuracy (86
distance to Scalar RM5Prop
68
central flow
stable
0.5
0.15
central
62
0.2
60
Scalior RMSShrp Imiipnints]
0.05
cwntir flow
58
1000
2000
3000
4000
5000
2000
4000
GOOD
8000
1000
2000
3000
4300
5000
step
step
step
Figure 32.17: 17: Scalar RMSProp central flow for a Mamba with MSE loss, 17 =0.01, B2 =0.99, and bias correction.
86

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicteed oscillation covariance
stable fow
empricaal verience slong IJ0 aigenwecters
0.30
0.004
Scalar RMSProp (average)
Zift eigenvatues
0.25
0.003
canti flow prediction
Scalar RMSProp
6.15
central flinr
0.001
stuble flow Hoop!
6.10
0.000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5030
1000
2000
3000
4000
5000
step
step
gradient norm
top hessian eigenivalues
effectivve step size aVV
Scaler RMSShop
125
0.06
100
0.75
75
6.50
Scalie RMShrop
Scalar RMSProp (average)
Scalar RMSProp
0.04
0.25
central flow
central flow prediction
25
central flowr
stuble flow [top!
0.03
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
stop
natwork outputs on tast example
test accuracy (36)
di stance to Scalar RMSProp
1.0
70
0.5
stable
0.4
Scala RMSPrup
0.4
contral flow
0.0
Scalar RMSProp (midpointsi
contral
0.1
1000
20000
3000
4000
5000
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
stop
Figure 32:18: bias correction.
18: Scalar RMSProp central flow for a Mamba with MSE loss, 1) 00.02, 82 =0.99, and
87

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted oscillation covariance
1.25
slable flw
Scale RMS/hop flow
centiall (weragee
ttop
0.0004
0.75
centtaal few pred alion
0.0033
0.0002
Ampircaa variance slong eigervectors
6.25
2019 eigensoliues
8.00
0.0000
10:0
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
step
step
step
gradient norm?
top hessian eigenvvalus lues
offective stop size W/V
0.6-
6:00
Scaler RMSivop
nlww
Itop!
0.05
0.06
Scaler RMSSPoop
(wergg)
0.34
central
0.33
central fow pred shion
2.00
0.32
0.01
0.00
30:0
2000
3000
4000
5000
0
10000
2000
3000
4000
5000
10000
2000
3000
4000
5000
step
step
network outputs on test example
test accuracy (%)
distance to Scalar RMSProp
60-
central flow
55
50
central flow
45
40
Seater RMSPhop Imiopnints)
-10
35
central flowe
1000
2000
3000
4000
5000
2000
4000
600D
8000
0
1000
2000
3000
4000
5000
step
step
Figure 33.1: Scalar RMSProp central flow for a CNN with CE loss, 77 =0.003, B3 =0.99, and bias correction.
train loss
top effective hessian eigervalues
real vs. predicted oscillation covariance
1.25-
stable flow
Scala RMSShop
RMSProp
nlww stable Itep!
Scalor RMSProp (average)
0.0035
6.75
.-
central few pred ition
central flw
0.0010
empintcal variance 210 eigensectoon
eipensalues
6.25
5.0005
10:0
2000
3000
6000
5000
1,000
2000
3000
4000
5000
10000
2000
3000
40000
5000
step
step
step
gradientt norm?
top hessian eigeniva lues
effective stop size N/VVV
Scalia RMS/hop
0.125
Scalar RMSSPop
6:00
now
central
stable ttep
0.100
1.5
Scalar RMMSPrED (average)
cencial
1.0
centnal few pred kiion
0,050
0.5
0.025
0.000
10:00
2000
3000
6000
5000
0
1.000
2000
3000
4000
5000
1.000
2000
3000
4000
5000
step
step
step
network outputs on test example
Lest accuracy (96
distance to Scalar RMSProp
60-
central Blw stable
55
RMS Prop
50
central flow
45
-5
40
Sealer RMS/hhp Imispnints
10-
35
central al flow
1000
2000
3000
4000
500D
2000
4000
600D
8008
10000
2000
3000
4000
5000
step
step
Figure 33.2: Scalar RMSProp central flow for a CNN with CE loss, 77 =0.006, (3 =0.99, and bias correction.
88

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
reall vs. predictted oscill ation cov covariance
sladle flow
Scaler
Scalie RMSprop flow stable
0.005
empircaal verience along Zid aigenwecters Zift eigencstaes
RMSProp (average)
2.004
central flow
centiral Nlo precittion
2-
0.002
0.001
3000
2000
3000
4000
5000
10000
2000
3000
4000
5000
1000
2000
3000
4000
5000
step
step
gredient norm?
top hessi an eigenvaauus
step size gNVVV
600
Scalier RMSProp
0.50
Scaler RMSShrpp
flow
0.25
flow
stable
0.20
Scalar (average]
3-
central flom
100
concrol flom prediction
200
1.00
3000
2000
3000
4000
5000
1.000
2000
3000
4000
50000
0
1000
2000
3000
4000
5000
Stop
Stap
network putputs an test example
test accuracyy (%)
distance to Scalar RMSProp
flow stable
2-
Scalor RMSPTOP
Scalor RMEprop imispointti
1-
central
1000
2000
3000
400
5000
2000
4000
6000
8200
1000
20000
3000
4000
5000
Figure 33.3: Scalar RMSProp central flow for a CNN with CE loss, 00.01, correction.
2,001,,220 f2 =0.99, and bias
top eigenvalues
prodicted oscillation covariance
stablc flow
Scalar RNSProp
RMSPmp
central flw
cential (average)
0.004
1.0
central fow prediction
0.003
0.002
empitical variance ainng eigenvectors
1-
210
digenvalues
0.4
5.000
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
500
1000
1500 step
2000
2500
2000
step
step
gradient norm?
top elgenvalues
step size niv
800
Scalor RMEPPro
500
central
RMSProp
Scalar RMS3mm (average)
200
Scalar RMSProp
centiaal flow
central few pretiction
100
stable flow Hep!
central flow
500
1000
1500 step
2000
2500
3000
0
500
1000
1500
2000
2500
3000
500
1000
1500 step
2000
2500
3000
step
network outputs test example
test accuracy (36)
distance to Scalar RMSProp
70-
central fow
65
central
-1-
-2
50-
45
Scaler RMSShop Imispnints]
cwntr flow
0.5
500
1000
1500
2000
2500
3000
500
1000
1500 step
2000
2500
3000
0
500
1630
1500
2000
2500
3000
step
step
Figure 33.4: Scalar RMSProp central flow for a ResNet with CE loss, 17 =0.01, B2 =0.99, and bias correction.
89

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenivalues
real vs. predicted escillation covariance
viable flow
Scaler RMSivop flow stable tto!
Scular
0.008
1.0
Scalor RMSProp (average)
central
empinical verbance E00) sigencectors
centiaal few preciction
2-
0.6
0.004
110 eigurssluss
0.002
0.000
500
1000
1500 step
2000
2500
3000
0
500
1000
1500
2000
2500
3000
500
10:00
1.500 step
2000
2500
3000
step
gradient norm?
top hessian eigenvalues
effective step size ANV
Scaler RMS5hop
0.08
RMMPPrp
Scalar RMSProp (average)
Scalar RM5Prop
central flow
central flow
0.02
cential few preciction
stable flow Heep!
500
1000
1500 stap
2000 2500 3000
0
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500 3000
2300
stop
network outputs on test example
test accuracy (36)
distance to Scalar RMSProp
Scabor RM1PPop
70-
central
65
1.25
1.00
55-
0.50
50-
Scalar RMSProp Imiopoints)
0.25
-3
45
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
500
1690
1500
2000
2500
2000
Stap
2300
st4p
Figure 33.5: Scalar RMSProp central flow for a ResNet with CE loss, 17 =0.02, B2 =0.99, and bias correction.
train loss
top effective hessian el genvalues
real vs. predicted ascillation atiion covariance
1.2
(MMPPrpp
flow
along 310 eigenvectors
Scalar IMEPPrpp (averagu)
stable flow [op!
0,0100
sigervalues
central flow
0.6
predtctionn
2
0.0075
0.4
0.0025
0.0
0
0.0000
1000
1500
2500
500
1000
1500
2000
2500
3000
1000
2000
25:00
step
gradient nom??
hessian eigenvalues
affective stop size gVviv
200
0.150
Scaler RMSProp
central flow
0.100
Scaler HMSPPop
Scalar AMSProp
0.050
contral
centrol flow prediction
stable flow (top)
0.025
0-
1500
2500
500
1500
step
step
step
outputs example
test accuracy
distance RMSProp
central flw
sable fow
65
60
0.5
0.6
55
0.4
Scator RMSRroo contral flom
Scalar RM5Prop imidpoints)
central flow
500
1000
1500
2000
2500
3000
0
500
1000
1500 step
2000
2500
3000
500
1000
1500 stpp
2000
25000
3000
step
Figure 33.6: Scalar RMSProp central flow for a ResNet with CE loss, 17 =0003, B2 =0.99, and bias correction.
90

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvallues
real vs. predicted esciliation covariance
stabic flow
Scalor RMSProp
1.2
RMSPPrp
central
0.015
(weragee)
--. contral fow predicttion
0.5
empirical varfance
EVI eigenvalies
0.6
0.000
1000
2000 step
3000
4000
1000
2000 step
3000
4000
1000
2000 step
3000
4000
gradient norm"
hessian aigenvalues
affectiva step
0.08
Scalor RMSProp
2000
Scalar RMSProp control nlow
Scalar #MSProp
80
Scalar RMSProp laveroge)
centrol
1500
flom (top)
0.06
control flww
prediction
1000
0.04
20
500
0.22
0.00
2000
step
step
network outputs test example
test accuracy (%)
distance to Scalar RMSProp
cennral
1.0
60
0.8
stable flow
0.5-
Scaler RMSPPop
55
0.3
centrral flow
Scalor RMSProp imidpointsi
1000
2000
50.00
4000
1000
2008
3,00
4000
1000
2000
3000
4000
Figure 33.7: Scalar RMSProp central flow for ViT with CE loss, =0.01, 82 =0.99, and bias correction.
top eigenvalues
real prodicted escillation covariance
stable flow
Scalar RMSProp
emplttal voriance
Scalar RMSProp
RMSShop (nverage)
flom
0.025
along Eidl eigenwectors
500 eigenvaluus
flow central flo prediction
0.020
0.8
0.010
6.6
0.005
0.000
0
1000
2000 step
3000
4000
10000
2000
3000
4000
1000
2000
3000
4000
step
step
gradient norm?
top hessian elgenvalues
stop size 1//1
80
Scalor RM15Prop
Scalor RMSSPro
Scalor RM15hop laveragel
1000
central
60
central
800
0.100
0.075
40
Scalar RMSSPro
200
central flow
stable flow (to))
0.025
1000
20000
3000
4000
1000
2000
3000
4000
1.000
2000
3000
4000
step
step
step
network outputs on test example
test accuracy (36)
distance to Scalar RMSProp
central fow
stable
55
0.5
0.6
-2
50
Scaler RMSShhp
Scaler RMMSPrp (midponts)
0.2
central flw
central flw
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
step
step
step
Figure 33.8: Scalar RMSProp central flow for a ViT with CE loss, 77 =0.02, /8 =0.99, and bias correction.
91

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenivalues
real vs. predicted oscillation covariance
1.2
stable feww
Scaler RMSProp
0.04
Scaler
flow
Scalar RMSProp (average)
ltop!
0.03
flo
empirical variacce
0.8
centi flow prediction
2
0.52
vigencoll
IIO
6.6
0.4
1000
2000 step
3000
4000
1000
2000 step
3000
4030
1000
2000
3000
4000
step
gradient norm?
top hessian eigenvaaues
offective stop size g/V
Scaler RMShhop
0.20
Scaler RMS5Pop
60
1015/hp
600
contral
0.15
contral now prediction
0.30
200
Scalar RMSProp
central flow
0.05
siuble flow [top!
0.00
10000
2000
3000
4000
1000
2000
3000
4030
0
1000
2000
3000
4000
stop
Stop
Stop
network outputs on test example
test accuracy (36)
distance to Scalar RMSProp
65-
central
0.6
0.4
Scalor RM15Prop
50-
Scalar RMSProp (midpoinnsi contral
contral
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
stop
St0p
st6p
Figure 33.9: Scalar RMSProp central flow for a ViT with CE loss, 1) =0.03, 82 =0.99, and bias correction.
train loss
top affective hessian el genvaluus
real vs. predicted oscillation covariance
1.4
0.012
RMMPProp
Scalar RMM5Prop (average)
3-
central flow
0.008
flom prediction
Scaler RMSirrp
0.004
empiiicall
0.4
along 211 sigervectors
6.2
sigenvalez
0.000
6000
2000
4000
8000
2000
4000
6000
step
gradient norm?
hessian eigenvalues
offective step size N
25
RMSSProp
1250
0.4
Scalee RMSProp
Scaler HMMS/rop (wverage)
central flow
20
central flom
1000
0.3
15
contral prediction
750
10
250
stable flow (top)
0-
4000
8000
step
step
step
outputs example
test accuracy
RM5Prop
100
1.50
central al flw
1.25
stabe flw
central
0.75
-1-
40
Scalar RM5Prop imidpoints)
0.25
central flow
0.00
2006
4000
6000
8000
0
2000
4000
6000
8000
2000
4300
6000
8000
step
step
step
Figure 33.10: Scalar RMSProp central flow for LSTM with CE loss, =0.01, 82 =0.99, and bias correction.
92

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvaaluss
real vs. predictted oscillation covariance
1.25
Mable flw
0,520
1.00
PMM javoragel
central
0.75
centraal flow predicioon
0.010
0.50
Scaler RMSProp
empliicall veriance
6.25
centrall flinr
0.005
along sigervecters
stubll flow ltop!
21/1 eigeticll.sss
0.000
2000
4000 step
6000
8000
2000
4030
6000
8030
2000
4000 step
6000
8000
gradient norm?
top hessian elgenivalues
effectivv step size VVV
20-
600
0.8
Scaler RMSShop
0.6
15-
10-
RMShhop
300
0.4
Scalar RMSProp (average)
200
Scalor RMSProp
0.2
central flow
central flowr
central flow prediction
100
siuble flow [top!
2000
4000
6000
8000
2000
4030
6000
8030
0
2000
4000
6000
8000
stop
Stap
natwork outputs on test example
test accuracy (36)
di stance to Scalar RMSProp
100
1.25
central
1.00
Scaler RMSimup central flow
0.50
2-
Scalar RMSProp (midpoinnsi
contral
0.25
2000
4000
6000
8000
2000
4000
6009
8000
2000
4000
6000
8000
stop
Stap
Figure 33.11: Scalar RMSProp central flow for LSTM with CE loss, 0000 32 00.99, and bias correction.
top effeetive hessian cigenvalues
vs. predicted oscillation covariance
1.25
stable flo
Scalor AM5Phop
Scalar RMSPhop
central
0.025
1.00
lnverage)
0.020
6.75
central flow prediction
along 310 digenvectors
eigenvalues
0.010
0.25
0.005
0.000
2000
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
8000
step
step
step
gradi ent
top hessian cigenvalues
effective stop gvv
1.2
Scalor RMSProp central
1.0
30.0
200
7.5
5.0
Scaler RMSProp (sverage)
Scaler RM Shop
0.4
central flow
centnal flow
0.2
central flow prediction
stable flow (top)
2000
4000
5000
8000
2000
4000
6000
8000
0
2000
4000
6000
8000
step
step
network outputs on test example
Lest accuracy (%)
di distance to Scalar RM4Prop
Scaler RMS2mop
100
1.2
stable
central fow
80
0.8
-1-
Scaler RMSProp (middnintt)
50
central flow
2000
4000
00:00
80C0
2000
4000
6000
BD00
2000
4000
6000
8000
step
step
Figure 33.12: Scalar RMSProp central flow for a LSTM with CE loss, =0.03, 8g =0.99, and bias correction.
93

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenivalues
real vs. predicted escillation covariance
2.5
slable flow
Scaler RMS Prop
0.006
2.0
0.005
Scalor RMSProp (average)
central
3.006
centiaal few preciction
2-
ampírical verfance along sigenvectors
0.001
210 eigeticalaes
0.000
10000
2000
3000 step
4000
5000
6000
0
1000
2000
3000 step
4000
5000
6000
1000
20000
3000 step
4000
5000
6000
gradient norm?
top hessian eigenvalues
effective step size ANT
Scaler RMS5hop
600
400-
RMSShop
Scalor RMSProp (average)
Scalar RMSProp
0.04
central flow central few preciction
central flo
auble flow Hepp
0.00
0
10000
2000
3000 stop
4000 5000 6000
0
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
6000
Stop
stop
netwark outputs on test example
test accurary (36)
distance to Scalar RMSProp
100
80-
4-
60-
40-
2-
Scalor RMSPPrp
Scalar RMSProp Imispointt)
contral flow
contral
central
1000
2000
3000 stop
4000
5000
6000
2000
4000
5000
8000
1000
2000
3000
4000
5000 6000
St0p
st4p
Figure 33:13: Scalar RMSProp central flow for a Transformer with CE loss, =0.01, 32 -0.99, and bias correction.
top eigenvaluus
vs. predicted covariance
stable flow
Scalar RMSProp
Scalar RMSProp
central flom
(overage)
0.008
central flom prediction
aton 2/0 aigenvecters
0.002
0.000
1000
2000
3.000 steg
4000
5000
6000
10000
2000
2000 step
4000
5900
6000
1000
2000
300
4000
5000
6000
step
gradient
top elgenvalues
effective step nVV
Scalor RMSProp
0.20
Scalor RMEPPro
central
central
0.15
Scalor RMSProp
laverage)
rentral flow Lunti flo prediction
9.10
100
0.05
9.00
0
1000
2000
3000
4000
5000
6000
1.000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
6000
step
step
step
network outputs on test example
test accuracy (36)
distance to Scalar RMSProp
100
80
central
Flow
2-
Scaler RMSiPPop (m diponts))
40
cential flw
1000
2000
3000
4000
5000
6000
2000
4000
6000
8000
1000
2000
3005
4000
5000
6000
step
step
step
Figure 33.14: Scalar RMSProp central flow for a Transformer with CE loss, =0.02, B3 =0999, and bias correction.
94

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted escillation covariance
3.0
Mabie Row
RMt5Prap
Scale RMSivop
flww
0.0150
empiiical verience
along 20d digenwecters
central laverage)
ttpp
0.0125
ZTD eigencstaes
0.0100
... central Now prediciion
2-
0.0075
1.0
0.0050
0.0025
10000
2000
3000 step
4000
3000
6000
0
1000
2000
3000
4000
5000
6000
1000
20:00
3000
4.00
5000
6000
step
step
gradient norm?
top hessian eigenvalues
step size WNV
250
Scalue RMSivop now
60
Scwler RMSPhop
20.0
Itpp
150
central
cential New pred etion
50-
0
1,000
2000
3000
4000
5000
6000
0
1000
2000
3000
4000
5000
0000
1000
2000
3000
4000
5000 6000
step
step
step
network outputs test example
test accuracy (%6
distarce to Scalar RMSProp
100
2.5
central flo
80-
1.5
central fiow
60-
1.8-
Sealer RMShhop Imiopaints)
-3
40
central flow
1000
2000
3000
4000
5000
6000
2000
4000
6OOD
0
1000
20:00
3000
4000
5000
6000
step
step
step
Figure 33.15: Scalar RMSProp central flow for a Transformer with CE loss, =0.03, 82 =0.99, and bias correction.
hessian eigenvalues
prodicted
1.2
stablc flow
0.006
RMSPmp
9.005
cential (average)
5.004
central flw prediction
0.003
along 509 eigenvectors
eigenwalues
Sealer RMSProp
0.6
stable row Itop!
contral flow
0.001
5.000
1000
2000
3000
4.000
5000
1000
2000
300
4000
5000
1000
2000
3000
40000
5000
step
step
step
gradient norm?
top elgenvalues
step size niVD
25
Scabr RM15Prop
Scalar RM5Prop central flow
Scalor RMSSPro
20
Scalor RM55hop lnveragel
0.015
central
(top)
15
central prediction
1000
0.010
1.0
500
1000
2000
3000
4000
5000
0
1.000
20000
3000
4000
5000
1000
2000
3000
4000
5000
step
step
step
network outputs best example
test accuracy (36)
distance to Scalar RMSProp
central fow
stable
RMSProp central
64-
0.15
62-
0.10
60-
Scalirr RMSShop Imiopnints]
0.05
cuntrr flow
58
1000
2000
3000
4000
5000
2000
4000
500D
8008
1000
2000
3003
4000
5000
sep
step
step
Figure 33.16: Scalar RMSProp central flow for a Mamba with CE loss, 17 =0.007, 83 =0.99, and bias correction.
95

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenivalues
real vs. predicted escillation covariance
1.2
saable flww
Scaler RMSivop flow stable ttp!
emprical verience slong ZiO digenwectors
1.0
Scalor RMSProp (average)
Zift eigenvalars
central
centtaal few preciction
0.004
0.8
0.6
0.002
0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3003
4000
5000
step
step
step
gradient norm?
top hessian eigenvaauee
step size N/VVV
Scaler IM5hhop
0.025
Scaler RMSSPop
1015/hp
1000
0.029
contral
15
contral now prodiction
0.015
000-
10-
Scalar RM5Prop
0.100
200
central flw
Jble flow Hepp!
0.005
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
stop
S10p
stop
network outputs on test example
test accurary (36)
distance to Scalar RMSProp
0.25
Scabor RM15Prp
68.
central
66-
0.15
62-
0.10
60
Scalar RMSProp Imiopoints)
0.05
-2-
58
1000
2000
3000
4000
5000
2000
4000
6000
3000
1000
2000
3000
4000
5000
stop
Stop
step
Figure 33.17: Scalar RMSProp central flow for a Mamba with CE loss, =0.01, .3 00.99, and bias correction.
effective hessian eigenvalues
predicted oscill ation covariance
emglr
1.0
Scalor RorSProp
Scalar RMISPro loverage)
0.009
EVT aigenvatees
flow
prediction
0.006
0.6
Scalar RMtSPrp
0.002
(top)
1000
3000
4000
5000
1000
2000
3000
4000
0.000
2000
2000
3000
4000
gradient norm?
top hessian eigenvalues
effective step size n/V
0.05
RMSPhop
500
RMSProp
RMSProp
Scater RMSProp (averaee)
central flow
0.04
--. central
central flow central now prediction
stob it flow Itop!
300
0.63
200
0.62
100
1000
2000
1000
AD00
5000
1000
2000
3000
4000
5000
1.000
2000
3000
40on
5000
network outputs on test example
test accuracy (%)
distance to Scalar RMSProp
RMSPPrp
flw
68
65
-2
Scolor RMSSPro (nridpoints)
central flow
56
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
6000
5000
stop
2t8p
Figure 33.18: Scalar RMSProp central flow for a Mamba with CE loss, =0.02, 32 =0.99, and bias correction.
96

---

Published as a conference paper at ICLR 2025
C.3 RMSPROP
top effective eigenvalues
oscillation
siaiie flom
RM5/hop
0.0004
empiriaal verience along cigenvectors
0.4
AMStrop lnverge)
RMSProp
central flo
stable flow
Iidl eigenwalues
central preciction
0.0002
0.2
0.1
0.8001
500
1000
1500
2000
2500
2000
500
1900
1500
2000
2500
2000
500
1000
1500
2000
2500
2000
nom
eigenvalues
accuracy
PM5pPrpp
RM5hhop
0.6
RMSPrrp (average)
500
stable
central Row
60
prediction
0.4
50
300
0.2
100
RMSPup (midgeinis)
central flow
500
1000
1500
2000
2500
3000
500
10000
1500
2000
2500
3000
500
1008
1500
2000 2500 3000
outputs on example
RMSProp
10-5
2.5
contral flo
0.4
1.0
RMSProp central flow
10-11
RMSProp
0.5
central flow
-0.2
500
1000
1500
2000
2500
300D
500 1000 1500 2030
2500
3000
0
500
1000
1500
2000
2500
3000
Figure 34.1: RMSProp central flow for a CNN with MSE loss, 77 =7c-06, 83 =0.95, e =cc-08, and bias correction.
97

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real Vs. predicted oscillation covariance
slabie flow
RMSProp
emplical varience along Ho digetwactors
0.00125
0.3
central flow
RMEPrro (average)
flow (top)
0.00100
eigenvalues
coniral flow prediction
0.00075
0.2
0.00059
0.1
0.00025
0.0
0.00000
500
1000
1500
2000
2500
3000
500
1000
1500
2008
2500
3000
500
1000
1,500
20.00
2500
3000
gradient norm?
top hessian eigenvalues
test accuracy (%6
RM5Prop
RMShmp
500
central flow prediction
400
stable (top)
60-
300
50-
100
RNSPup (midgeints)
central
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500 3000
notwork outputs on test example
coordinates of
distance to RMSPrap
10-1
central
0.4
2.0
0.2
10:
10
0.0
10 x
1.0
RNSProp
10-4
RM
0.0
1000
2500
500 1000 1500 2000
2500
3000
500
1000
1500
2000
2500
Figure 34.2: RMSProp central flow for a CNN with MSE loss, 17 = 1e-05, B3 0095, =1c-08 and bias correction.
train loss
top effective hessian eigenvalues
real vs. predicted oscillation covariance
0.4
stable flw
RMSProp
RMSPrup (average)
3-
stable Hap!
0.0020
0.2
0.3
central flow
control flow prediction
2-
0.9015
eigetivecturs 210 digenvakes
empitiical varance
0.1
0.0005
0.0
0.0000
2500
25.00
2000
2500
gradient norm?
top hessian sigenvalues
test accuracy (36)
2.0
250
60
1.0
AMSProp
100
RMSProp
50
central flow
40
RMSinop (midpaints) central flow
contrll flow prediction
stable flo Hep!
0.0
1000
1500
2000
2500
3000
500
1000
15:00
2000
2500
3000
500
1000
1500 2000 2300 3000
network outputs test example
coordinates of
distanca RMSProp
0.6
central flow
0.4
10-1
2.0
stable fow
0.2
10-2
1.0
0.2
10-9
RMSFFrp
IMM SPMpp
0.5
-0.4
central now
10
central flow
500
1000
1500
2000
2500
3000
500
1000 1520 2030
2500 3000
500
1000
1500
2000
2500
3000
Figure 34.3: RMSProp central flow for a CNN with MSE loss, 77 =2c-05, /3 =0.95, =cc-08 and bias correction.
98

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigeniaalues
real VS. predicted escillaion tion covariance
stable flow
RNSPrpp
0,00330
emprical lence
0.00925
along Zio cigenvectors
central (average)
flww (top)
eigenvalues
0.3
Central flo prediction
0,00015
6.2
0.00010
1-
0.00005
0.1
1000
1500
2000
2500
3000
5c0
1000
1500
200n
2500
3000
500
1000
1500
2008
2500
3000
gradient norm?
top hessian eigenvalues
test accuracy (%)
10
RMSPhop
1000
RNSPmp
70-
central
800
flow (top)
contral nlow prediction
400
RNSPup (midgeints)
central
1000
1500
2000
2500
3000
5C0
1800
1500
2000
2500
3000
500
1000
1500
2000 2500 3000
network cutputs on test example
coordinates of
distance to RMSProp
contral
0.4
6.3
10-9
10:11
10-14
RNSProp
RNSProp
1000
1500
2000
2500
3000
500
1900
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Figure 34.4: RMSProp central flow for a ResNet with MSE loss, = -c-00 B2 =0.95, = and bias correction.
loss
eigenvalues
vs. predicted oscillation
stuble flww
RMSProp
empricaa werience along cigenwectors
0.4
RMSPhop (nverge))
RMSPPop
central nlw
stable flow (top)
0.00125
Ztl eigenvalues
0.3
0.00300
central prediction
0.2
3.00050
0,00000
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
2000
500
1000
1500
2000
2500
2000
norm?
eigenvalues
accuracy (%)
10
RMt/hop
RM5hhop
70-
RMSProp (average)
central flw
stable
65-
contral prediction
600
400
RMSPup (midgsinis)
central flww
0
500
1000
1500
2000
2500
1000
500
10000
1500
2000
2500
3000
50n
1565 2010 2500 3000 distance RMSProp
1000
network outputs on example
coordinates
0.5
RMSProp
contral foo
central
10-
10:13
10:1s
RMSPro
central fow
500
1000
3500
2000
2500
3000
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Figure 34.5: RMSProp central flow for a ResNet with MSE loss, =2c-05, /2 =0.95, =1c-08, and bias correction.
99

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalluss
real vs. predicted escillation covariance
0.4
slablle flw
RMSSPrpp nlww stable ttp!
emppiiaal yarfence along 15f elgenvectors Tin eigenwatues
0.003
0.3
RMSProp (average)
central luo prediction
0.002
0.001
0.000
500
1000
1500
2000
2500
300n
500
1000
1500
2000
2500
1000
500
10.00
15:00
20000
2500
3000
gradient norm?
top hessian eigenvalues
test accur acy (96)
RMSSPopp
RMSSPrp flow
70-
contral flo prediction
contral
300
sable Htp!
200
55-
50-
RNSPup (midgeints)
45
central
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
1000
500
1000
1500
2000
2500
3000
network outputs on best axample
coordinates of
distance to RMSProp
10-2
RMSProo
central
6-
10
0.2
10-11
RMSProp
10-14
500
10000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
500
1600
15:00
2000
2500
3000
Figure 34.6: RMSProp central flow for a ResNet with MSE loss, 4c-05 B2 =0.55 clc-08 and bias correction.
eigenvaaluss
vs. predicted oscillation covariance
0.5
stable flw
empiiea vet letc along 11D ciganvectors
RMSPPop
0.4
RMS5Prp (nvergee)
3-
0,0006
Bti eigenvalues
central now prediction
0,0004
RMSPmop central flow stable flow (top)
0,0002
6.1
0,0000
1000
20000
2000
4000
10000
2000
300
4000
1000
2000
3000
4000
gradient
eigenvalues
accuracy (%)
RMS/Prop
RMSProp central nlw
RMSProp laverage)
60-
2000
nlw
contral prediction
1500
50-
40-
500
RMSPup (midpsinis)
30-
central flow
0
1000
200
3000
4000
1000
2000
1000
4000
1000
2000
5000
4000
network outputs on example
coordinates
distance RMSProp
10
6.6
RMSProp central
RMSProp
contral flow
10 6.
0.4
10-0
10-30
10
10-34
6.0
1000
2.000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
Figure 34.7: RMSProp central flow for a VIT with MSE loss, 55-006 B3 =0.95, =le-08, and bias correction.
100

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenivalues
real Vs. predicted escilation covariance
0.5
stable flow
RNSPip flow (top)
0,0020
emprical lence
along Zin cigenvectors
0.4
central (áverage)
0.0015
eigenvalues
0.3
Central flo prediction
6.001
1-
0.0005
0.0000
1000
2000
1000
4000
1000
2000
300
4000
1000
200.
3000
4000
gradient
top hessian eigenvalues
test accuracy (%)
10
RMSProp contral (average)
RNSPmpp
60-
2500
flow (top)
contral nlw prediction
50-
1500
40-
500
RNSPup (midgeints)
central
1000
2000
1000
4000
1.000
2000
300.
4000
1000
200
1000
4000
network outputs on test example
coordinates o
distance to RMSProp
RMSProp
10-4
central
10.
1.0
10
10:30
10-9
0.4
RNSProp
10
1000
4000
2000
3000
1000
3000
Figure 34.8: RMSProp central flow for a VIT with MSE loss, =7-06 32 =0.95, =1c-08 and bias correction.
train less
top effective hessian eigenvalues
real vs. predicted oscillaiion covariance
0.5
RMSProp flow stable fow (top)
RMSSPop
RMMSPopp laverage) central flow
0.003
empirical veriince along sigenvectors
central flow prediction
0.002
Z1O eigeneslaas
0.1
0.001
0.000
500
1000
1500
2000
3000 3500
500
1900
1500
2500 3000 3500
500
1000
1500
2000
2500 3000
gradient norm?
top hessian eigenvaauues
test accuracy (%)
15.9
(sverage)
60
12.5
(hop!
central flw prediction
1000
7.5
2.5
20
RMSPrrp (midpolnts)
500
1000
1500
2010
2500 3000 3500
500
1000 1500 2000 2500 3000 3500
1980
2000
3000
network outputs on test example
coordinates of
distance to RMSProp
0.8
RMSihop
centiall flow
stable
1.5
0.4
10
1.0
10-11
0.5
RMSPPr
10
central
1900
2000
3000
1000
2000
3000
500
1000 1500 2000 2500 2000 3500
Figure 34.9: RMSProp central flow for a VIT with MSE loss, -e-05 B2 =0.95, = and bias correction.
101

---

Published as a conference paper at ICLR 2025
real vs. predicted oscillation covariance
train loss
top effective hessian eigenvalues
1e 6
RMShrap
empéical veriance
central fllw
along 100 eigenvectors
stable flow (tos)
2t1t elgenvalaes
0.3
stable flow
2-
3-
0.2
RMSFrop
RMSProp (averaee)
1
0.1
rentral flow
prediction
1
1000
2000
3000
4000
10:00
2000
3000
4000
10:00
2000
3000
4000
gradient norm?
top hessi an eigenvalues
test accuracy (%)
0.30
RMSProp
200
RMSProp
100
RMS/top (average)
centtal flow
stable flow (top)
80
0.20
central flww prediction
150
0.15
100
0.05
40-
RMSPrp (midpoints)
central flow
0.00
1000
2000
3000
4000
1000
2000
3000
4000
2000
4000
6000
8000
network outputs test excmple
coordinates of
distance RMSPPrP
0.3
RMSProp central flw
10-4
RMSSrop central Flow
1.2
central flw
stable flom
0.2
10-1
0.8
0.1
0.6
0.0
10-a
10-3
0.2
10000
20000
3000
4000
1000
2000
3000
4003
1000
2000
3000
4000
Figure 30.10: RMSProp central flow for a LSTM with MSE loss, 17 = 11-00, B2 =0.95, 1c-088 =1c-88 and bias correction.
top effective hessian eigenvalues
real predicted oscillation govariance
stable flow
RMSPhop central flww stable fow (top)
empirical vet verience along digenwectors
iMSPtrp
0.00029
RMSPPrp (avergee)
3
2011 eigenvalaes
0.3
0.00015
central flom prediction
2-
0.2
0.1
0.00005
0.00000
1000
2000
2000
4000
1000
2000
2000
4000
1000
2000
3000
4000
gradient norm?
top hessian eigenvalues
lest accuracy
0.8
PMop
1500
RMSShrp central flow
100
PMSP (average)
0.6
central
stable flow Itep)
80
flow prediction
1000
0.4
750
500
0.2
250
40
IMSPup (midpsinit)
central flow
0.0
1900
2000
3000
4000
1.000
2000
3000
4000
0
1850
2006
3000
4000
network outputs on example
coordinates
distance to RMSProp
0.3
contral flow
1.5
flow
15-2
0.1
RMSProp central flom
10r4
0.0
10-1
0.5
0.1
10-41
SPro central flom
1000
2000
3000
4000
0
1.000
2000
300D
4000
1000
2000
3000
4000
Figure 34.11: RMSProp central flow for a LSTM with MSE loss, 17 =2e-05, B3 =0.95, e = 1e-08 and bias correction.
102

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian elgenvalues
real vs. predicted oscillation covariance
vlatie flom
RMSivop
flo
0.0010
central
RMSProp laverage)
3-
stable (top)
0.0008
central flow prediction
2-
0.2
0.1
1-
alie g 210 eigetveecton
empirical veriance
2001 eiguncal.ass
1000
2000
1000
4000
1000
2000
5000
4000
1:000
2000
3000
4000
OP r'adient norm?
top hessian eigenvalues
test accuracy (%)
PM5P
RKSivop
IMS (averigge)
flo
flow central flow prediction
stable (top)
20-
1500
3-
60-
1000
900
RNSPup (midgeints)
central
1000
2000
3000
4000
1c0n
2000
3000
4000
2000
4000
6000
8000
network outputs on test example
coordinates of
distance to RM5Prop
10"
0.3
central
10
0.2
10:4
1.5
0.1
RM15Prop
10
1.0
contral flow:
0.0
10-1
0.5
RMSProp central
-0.1
10-9
0.0
1.000
2000
30000
4000
1000
2000
300.
4000
1.000
2000
3000
4000
Figure 32.12: RMSProp central flow for a LSTM with MSE loss, 17 =6-055 B2 =0.55, 1cc-08 and bias correction.
top effective hessian eigenvalues
real predicted oscillation covariance
0.6
stable flow
RMSProp central flow stable fow (top)
0.00010
empircaal verience along aigenwactors
IMSPPrp
RMSSPop (avenge))
2011 eigenvataes
0.4
central now prediction
2.
0.00005
0.3
0.00004
0.2
0.00002
0.00000
500
1000
1500
2000
2500
500
1000
1500
2000
2500
500
1000
1500
2000
2500
norm?
eigenvalues
lest accuracy
300
1MMS/ropp
RMSShrp
iMSSPrpp (average)
250
central flww
flw
stable flow Itep)
now prediction
200
150
4-
100
IMSPup (midpsinis)
central flow
500
1000
1500
2000
2500
500
1000
1500
2000
2530
500
1000
1500
2000
2500
network outputs on example
coordinates of
distance to RMSProp
RMSProp central flow
RMSSPrp centrol
contral flo
0.4
10"4
10-7
0.2
0.75
10-*
0.0
10-11
10-42
500
1000
1500
2000
2500
0
500
1000
1500
2000
2500
500
1000
1500
2000
2500
Figure 33: RMSProp central flow for Transformer with MSE loss, 7) = =1c-05, B2 =0.95, e = -c- 08, and bias correction.
103

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted oscillation covariance
0.6
vlatie flom
RMSivop
empiniall verience
0.5
RMSProp laverage)
3-
flw stable (top)
0.0004
along sift eigenvecters
Fid
0.3
central
central flow prediction
2-
0.0003
0.2
0.0001
0.0
500
1000
1500
2000
2500
5.0
1000
1500
200n
25000
50.0
1000
1500
20000
2500
Of radiennt nomm?
top bessian eigenvalues
test accuracy (%)
AMShrop
RMSProp
100
central
300
centr
stable flow (top)
80-
central flow prediction
200
40
20-
INSPup (midgeints)
central
500
1000
1500
2000
2500
5.00
1000
1500
200n
2500
500
1000
1500
2000
2500
network outputs on test example
coordinates of
distance to RMSProp
10-1
1.50
0.4
10-1
0.2
RM15Prop
1.00
central flow
0.50
0.2
10-1
0.25
central
1000
1500
2000
2500
10000
1500
2000
2500
500
1500
2:000
Figure 14: RMSProp central flow for Transformer with MSE loss, 17 =2c-05, B2 =0.95, e =le- 08, and bias correction.
train loss
top effective hessian eigenvalues
real vs. predicted escillation covariance
0.0075
stable flow
RMSProp nlw Rop!
0.5
RMSirop (average)
central flow
control flow prodiction
0.0015
emptrical variance
SI0 eigervalues
0.2
0.1
0.0005
0.0
0.0000
1500
2000
500
1000
1500
1000
1500
2000
gradient norm?
top hessian eigenvalues
test accuracy (%)
RMS Prop
250
RMEProp central flow
Pro laveragel
central flow
stable sow (top)
central flow precistion
150
60
100
20
RMS5rrp [mispoints]
central flow
1000
1500
2000
500
1000
1500
2000
500
1,000
1500
2000
2500
network outputs example
coordinates of
distance RMSProp
RM5 central fow
10
centrol flw
0.6
10-4
stable flow
0.4
0.2
10-2
4-
0.0
10-4
10
RMMP
10-1c
contral flow
-0.4
0
5.00
10:00
15:20
2000
2500
500
1000
1500
2000
2500
500
1000
1500
2000
Figure 35: RMSProp central flow for Transformer with MSE loss, 7) =4c-05, B3 =0.95, =le- 08, and bias correction.
104

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian ei genvalues
real vs. predicted oscillation covariance
stable flom
RMShhrp
emprical veriance
RMSPrup
0.00025
along 2id eigetieccoos
RMSProp (average)
(top)
eigenvaluks
contral flow
0.00020
0.3
central flow prediction
2-
0.00015
0.00010
0.00005
6.1
0.00000
1000
2000
3000
5000
1000
2000
4000
5000
1000
2000
3000
4000
gradient norm?
top hessian eigenvalues
lest accuracy (%)
70
RMSProp
8000
RMSProp (everage)
central flo
central flow
stable flow (hop)
cundial fow prodiction
2000
RMSSrop (midpoinss)
contral flow
2000
5000
2000
4000
1000
3000
network outputs test example
coordinates of
distance RMSProp
RMSProp central flom
10-2
central flow
0.5
stable from
6.8
10-1
0.4
0.3
0.4
10
0.2
10
0.2
RMSProp
10-42
contral flow
0.0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Figure 34.16 16: RMSProp central flow for Mamba with MSE loss, 17 =7e-06, 32 =0.95, 1e-008 and bias correction.
hessian eigenvalues
predicted oscillation covariance
stable flow
RM SPrup
RMSProp central flom
0.5038
empirnc veriance along 315 cigonvectors
0.4
RM
stable
20th eigenvalues
0.0006
0.3
central flom prediction
0.0004
0.2
0.3062
0.1
0.0000
0
10000
2008
3000
4000
5000
1000
2000
5000
4000
5c0n
1000
2000
3:000
4000
5000
eigenvalues
accuracy
RMS/hop
RMS/hhp
RMSPtop (average)
5000
cuntral flow
12.5
4000
prestction
3000
7.5
30-
2000
2.5
1000
50
IMSPup (midpeinis)
45
central flow
0.0
1.000
2008
3000
4000
5000
1000
2000
30000
4000
5000
0
3000
2000
3000
4000
5000
network outputs on example
coordinates
distance RMSProp
1
0.6
contral foo
0.75
10
0.50
10
04
central flw
10
RMSProp
0.1
10-12
central flom
1005
2000
1000
40000
5000
1000
2000
5000
4300
5000
1000
2000
3000
4000
5000
Figure 34.17: RMSProp central flow for Mamba with MSE loss, = 1c-05, B2 =0.95, llc-08 and bias correction.
105

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvaalues
real vs. predicted oscilaation covariance
stabic Tlow
RMSProp
emplrical varance along lidl aganvectars 321) sigenvalass
0.4
RMSProp
Oiverage)
3
0.5015
central flow prodiction
0.3019
0.3005
0.1
0.0000
1000
2000
2000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3.00
4000
5000
norm?
top eigenvalues
accuracy
12-
RMSPPro
RMSProp
RM1SProp (pverage)
central flow
10
central
2000
stable
1500
6-
50-
RMSihhpp (miopaints)
control flow
2000
5000
network example
coordinates of
distance RMSProp
1.00
10-1
0.6
10
MMSPPrp central flow
0.25
10--
0.00
10
-0.25
30-12
0.0
4000
Figure 34.18: RMSProp central flow for Mamba with MSE loss, =2--55, 82 =095, =1e-08 and bias correction.
106

---

Published as a conference paper at ICLR 2025
real vs. predicted escillaioo covariance 1c: 5
in loss
effective hessian eigenvalues
1.25
stable flow
RMSSPrp
RMSPmp centr flw stable flw (top)
emper vetterce
1,00
RMSProp caveragel
digenvalues
central flow
0.75
prwd
0.50
500
1000
1500
2000
2500
3000
500
1009
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
gradient norm?
top eigenvalues
test accuracy
RMSProp
250
flow
6.100
50
0.075
RMSProp
0.050
RMSirop (mvenage)
100
40-
central
50
flow prediction
1500
2000
2500
1500
2000
2500
1000
1500
2000
2500
2000
network putputs on test example
coor diorten of
distance to RMSProp
10
1.5
10-1
centrol flow
0.0
1000 1500
3000
1.000
1500
3000
1500
Figure 35.1: RMSProp central flow for CNN with CE loss, 7) 7e-06, ,3 =0.5, 1e-08 and bias correction.
train
top effective hessian eigenvalues
real predicted
1.25
viable fow RMSSrop
0.000 125
emtiriccal varlience along cigenvecters
1.00
RMSProp
0.000 100
Σirl eigervallee
0.75
prediction
0,000
c.00c.50
RMSProp
centi flow
0,000
stable flow (top)
0.000000
500
1000
1500
2000
2500
2000
500
1000
1500
2000
2500 3000
500
1000
1500
2000
2500 2000
gradient norm?
eigenvalues
accuracy
300
RMShhop
0.15
250
central flow
60
200
50-
RMSProp
0.05
RMSPtop (everage)
central flow
RMSPup (midpeinis) central flow
central flow prediction
500
1090 1900 200n 2500 1000
500
1000
1500
2000
2500
3000
0
500 1800 1500 2000 2500 3000
networkk outputs test example
distance RMSProp
10-4
RMSProp
6-
contral flow stable
5.0
central
10-4
5
2.5
0.0
3
-25
-5.0
central flow
RMSProp
1
500
1000
1500
2008
2500
3000
500 1.00 1500
2000
2500 3000
500
1000
1500 2000
2500
3000
Figure 35.2: RMSProp central flow for CNN with CE loss, 77 1e-05 Bg =0.95, e =le-08, and bias correction.
107

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted oscillation covariance
1.25
slable now
RMMShhpp
amppptcaal verience along 310 digenwecters Zit etgencstaes
1.00
central laversge)
(top)
0.75
centtal flow prediction
0.0004
0.50
0.0032
0.25
500
1.000
15:00
2000
2500
3000
0
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
2000
gradient norm?
top hessian eigenivalue lues
test accuracy (%)
0.5
RMSSPrpp
0.4
central nlww
stable itep!
60-
2.00
RMSthap (averagH)
50
central flow
0.2
central rio predictio
0.1
50-
RMSPrup (médpeints))
central
500
1000
15:00
2000
2500
3000
5:00
1000
1500
2000
2500
3000
500
1000
1500
2000
25.00 2000
network outputs on test example
coordinates
distance RMSProp
10-1
RMSivop
central flow
10
10
-5
10 15
AMShtop central
500
1000
15:00
2000
2500
3000
500
1000 1500 2000 2500 3006
500
1000
1500
2000
2500
2000
Figure 35.3: RMSProp central flow for a CNN with CE loss, 1) =2e-05, .3 =0955, c 1e-08 and bias correction.
real vs. predicted oscillation covariance 1--5
train
effective hessian genvalues
stable flom
RMS hoop contral now stable
empiniccal variance
RMSProp
along 329 eigenvecters
RMSProp laverage)
I1O eigenwalues
1.0
central prediction
0.8
0.6
1000
2000
3000
2000
3000
2000
3000
gradient n
top hessian eigenvalues
test accuracy (%)
RMS/rop
1000
RM5
70
RMSProp laverage)
contral flo
1.0
contral flow
800
stable flow (top)
0.8
centrai prediction
60
400
55-
0.4
200
50-
RMSincp (midpointa) central
1000
2000
3000
4000
1000
2000
3000
4000
10000
2000
3000
4000
network outputs on test example
coordinates of
distance to RMSProp
10-1
flow stable flow
34
-2
RMM5Phop
10-37
RMM5Prpp central flow
central fumm
1800
2000
1000
4000
1000
2000
3000
4000
1000
2000
3000
4000
Figure 35.4: RMSProp central flow for a ResNet with CE loss, 77 = 1c-05, /32 =0.95, e =1c-08, and bias correction.
108

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian di genvalues
real vs. predicted escillation covariance
1.4
viadle flow
0.0006
emplicaal varience
along 201 eigennections
(average)
3
0.0005
sigenvaluus
1.0
central few
central flow redictiion
2-
0.0003
0.4
RMSPPOP
0.0000
0.0000
Itop)
0.0000
1000
20000
3000
1000
2000
3000
4000
1000
2000
2000
gradient norm?
top hessian eigenvalues
test accuracy (86
3000
Prop Prop liverage)
2500
central flo
central from central flow prediction
viable fiow [hop!
05-
1000
55-
500
(midpoinss)
contral flow
0
45
3000
10000
2000
4000
1.000
netwerk outputs example
coordinates of
distance RMSProp
PMMP central from
10-4
2.0
stable frw
central flow
10-1
1.5
10:30
1.0
-2-
RMSProp contral now
1000
20000
3000
4000
1000
2000
5000
4000
1000
2000
3000
4000
Figure 35.5: RMSProp central flow for a ResNet with CE loss, 1) 2e-05, 33 =0.95, 6 1e-08, and bias correction.
top effective hessian eigenvalues
predicted oscillation coveriance
1.25
vtatie flom RMS/rop
central al flow stable flor stopl
Prop
0.0020
empiriaal verience along cigenvectors
1.00
AMStrop lnvernge)
3
Zidl eigenwalues
central preciction
0.0005
1000
2000
3000
4000
1000
2000
3000
4000
10000
2000
3000
4000
eigenvalues
accuracy
RM5P
2000
centr al nlo
stable flow (top)
65-
RMSProp
1500
60-
flow central flow prediction
1000
50
IMSPup (midpsinis)
central flow
1000
2:000
1000
4000
1000
2000
3000
4000
1000
2000
3000
4000
outputs on example
distance RMSProp
RMSProp central
10-9
10-9
10-12
1-
10-35
RMSProp
contral flow
central flow
stable flom
-7.5
1.000
7000
3000
4000
1000
2000
3000
4030
1.000
2000
3000
4000
Figure 35.6: RMSProp central flow for a ResNet with CE loss, 17 4cc-0 /3 =0.95, 6 =1c-08, and bias correction.
109

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real vs. predicted oscillation covariance
1.4
viatie flom
RMSivop
0.0005
empiriaal verience
1.3
flom stable (top)
along sift eigenvectors
laverage)
3-
Zid
1.2
1.1
central flow prediction
2-
0.0003
0.0002
1-
0.9
0.0001
0.8
1000
2000
1000
4000
1000
2000
3000
4000
1.000
2000
3000
4000
radient norm?
top hessian eigenvalues
test accuracy (%)
PM5P
RKSivop
IMS
4000
flow central flow prodiction
stable (top)
3000
55
2000
1000
RNSPup (midgeints)
central
1000
2000
1000
4000
1c0n
2000
3000
4000
1000
2000
3000
4000
network outputs on test example
coordinates v
distance to RMSProp
RNSProp
10-4
RMSProp
1.0
centrol flow
flww
10-4
Now
0.8
stable
0.5
0.6
0.0
10 -0
0.4
10-12
0.2
1.0
10-11
0.0
1.000
2000
30000
4000
1000
2000
3000
4000
1.000
2000
30000
4000
Figure 35.7: RMSProp central flow for ViT with CE loss, 55e-66 32 =0.95, e == =1c-08, and bias correction.
eigenvalues
vs. predicted oscillation covariance
stable flow
RMSProp central flow
emppi vet lntce along SID cigenvectors
RMS5Prp (nvergee)
RMSPPrpp
3-
stable flow (top)
0.0020
Bti eigenvalues
central now prediction
0.0015
2-
0.001
0.0005
0.0000
1000
20000
2000
4000
1000
2000
3000
4000
1000
2000
3000
4000
norm?
top eigenvalues
accuracy (%)
30
RMS/hop
RM5hhop
25
RMSProp (average)
6000
central flo
60-
stable
20
contral prediction
15
55-
10
2000
RMSPup (midpoinis)
central flow
0
1000
200
1000
4000
200
2000
1000
4000
1000
2000
5000
4000
network outputs on test example
coordinates
distance RMSProp
RMMSPro central
10-4
contral foo
10:4
10-4
10
-1
RMSPrp
10:34
central fow
1000
20000
3000
4000
0
1006
2000
3000
4000
1000
2000
3000
4000
Figure 35.8: RMSProp central flow for ViT with CE loss, =7c-06, 3g =0.95, = 1c-088 and bias correction.
110

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalles
real vs. predicted oscilation covariance
1.4
stable flow
RK5Prop
empiical variance
along zio aigervectors
1.2
central (averaga)
(tag)
6.0125
Bitt
centtal flow predition
6,0075
2-
6,0025
1000
2000
300n
4000
1.000
2000
300n
4con
1,000
2000
3000
4000
gradient norm?
top hessian elgenvallues
test accuracy (96
125
RMS/rop
MMS/hop
6000
RK5Prop
100
flow contral flow prediction
stable (tag)
4000
50
RNSPup (midgeints)
central
1000
2000
3000
40.0
1000
2000
300n
4000
1000
2000
3000
4000
notwork putputs on test example
coordinates of
distance to RMSProp
10-4
2-
RMSRrap
10-1
1.25
contral flo
10
1.00
0.75
10-10
0.50
10:
RNSProp
0.25
10:14
1000
3000
4000
1.00
2000
3000
4000
1.000
2000
3000
4000
Figure 35.9: RMSProp central flow for ViT with CE loss, = =c-00, 82 =0.95, =0,95, = and bias correction.
real vs. precicted oscillation covariance
1.4
top effective eigenvalues
RMSSProp
3.0
1e-6
empiitcaal variance
contral flow
alogg eipenvecceor
1.2
I1O sigenwaluss
1.0
stable flow
RMSProp
RMSProp
0.6
0.5
flow prediction
1000
2000
3000
0.0
gradient norm?
top hessian eigenvalues
test accuracy (%)
RM5hrup central flow 2tble flow itop)
(milpoonnta)
0.10
300
contral al flow
0.08
200
60
RMSProp jUMShop (sverage)
100
0.02
cantral flow
flow
1000
2000
3000
4000
5000
1000
4000
5000
10000
2000
3000
5000
network outputs on test example
coordinates of
distance to RMSProp
0.6
0.4
contralfiiw
10-1
0.56
6.2
10-
0.04
0.0
0.2
10-2
RMS Prop
central flow
0.00
5
10000
2000
1200
4000
5000
10000
2000
3000
4000
50C5
10000
2000
3000
4000
5000
Figure 10: RMSProp central flow for LSTM with CE loss, 77 1c-05, B3 =0.95, =1c-08, and bias correction.
111

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian elgenvalues
real vs. pre oscillation covariance
slatie flow
RM5ivop flo stable (top)
empiriaal verience
RM5imop
aleng sift eigenvectors
RMSProp (average)
3-
Fidi eigenwalues
1.00
central
0.75
central ('ow prediction
2-
0.50
1-
0.25
0.0000
1.000
2000
1000
4000
5000
10000
2000
3000
40on
5000
1000
2000
3000
4000
5000
OF d'adientnonn norm?
top hessian eigenvalues
test accuracy (%)
2.0
PM5Prop
RKSivop
100
AMSPrrp (averigge)
3000
flo
1.5
central flow
stable (top)
80
central flow prediction
2000
1.0
60
0.5
1000
RNSPup (midgeints)
central
1.000
2000
1000
4000
5000
1000
2000
3000
400n
5000
1000
2000
3000
4000
5000
network outputs on test example
coordinates of
distance to RMSProp
1.0
10a
RNSProp
RMSProp
centrol flow
New
10-3
Nlw
stable
0.5
30
0.6
0.4
10-
10-
0.2
10
0.0
1000
20000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4.00
5000
Figure 31: RMSProp central flow for LSTM with CE loss, 77 =cc-55 B2 =0.95, e =cc-08, and bias correction.
train
effective hessian eigenvalues
predicted escillation covariance
1.25
stable flww
RMSProp
RMSProp central flow stable flow (top)
0.003
1.00
RMSProp laverage)
central flow prediction
0.002
0.001
emplrtcal variance sigenvectioss
0.25
+lgencaluss
0.00
1000
5000
1000
3000
4000
5000
0.000
2000
3000
gradient norm?
top hessian eigenvalues
test accuracy (36
20-
RMSProp
100
3000
central flow
sable fow (bop)
15
RMSProp
RMMPPrpp laverags) central flo
2000
10
60
central flow preaiction
1000
RMSProp Imiopoints)
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
network outputs on test example
coordinates of
distance to RMSProp
10-2
1.5
RMSProp
2.0
fow stable
1.0
10
0.5
0.3
10
3.0
-10-
10-
RM5Prop
0.5
central
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Figure 12: RMSProp central flow for LSTM with CE loss, 7) =6c-05, B3 =0.95, 6 = -1-008, and bias correction.
112

---

Published as a conference paper at ICLR 2025
real vs. predicted oscillation covariance
train less
top effective hessian eigenvalues
1e
RMStrop
1.
empéical veriance
2.4
centaal flow
along eigenvectors
stable flow (top)
0.8
St1i eigenvalaes
stable
2-
2.1
RMSFrop
RMSProp (average)
central flo
500
1000
1500
2000
2500
500
10:00
1500
2000
2500
1000
2000
2500
gradient norm?
top hessi an eigenvalues
test accuracy (%)
RMSProp
RMSProp imidpoints))
1.5
100
50
stable (tup)
80
RMSPrap laverage)
0.5
20
central flow
central flow prediction
0.0
500
1000
1500
2000
2500
500
1000
1500
2000
2500
5.00
1000
1500
2000
2500
network example
coordinates
distance RMSProp
centrol flom
10
stable flom
10
-0.2
10-19
0.4
-0.4
RNSProp
10:11
RMSProp
0.2
32
500
1000
1500
2000
2500
5.00
1.000
1,500
2000
2500
1000
1500
2000
2500
Figure 35.13: RMSProp central flow for a Transformer with CE loss, =1e-05, 82 =0.95, 6055,==le00 1c-088 and bias correction.
eigenvalues
oscillation covariance
2.4
siadil flow
RM5hhop central flw
empirical verience slong cipcrivectors
RMSirop
stable
Zib eigenvatues
2.0
central preciction
0.00015
1.8
0.00010
1.4
0.00000
1000
1500
2000
2500
1000
1500
2000
2500
500
1000
1500
2000
2590
gradient nom?
eigenvalues
accuracy
1.5
RMSihop
central flom
80-
stable
300
40-
200
0.5
RMShrop (mverage)
central flow
100
20
RMSPup (midpsinis) central flow
central flow presiction
500
1000
1500
200n
2500
520
10000
1500
2000
2500
500
1900
1900
2000
2500
network outputs test example
coordinates
distance to RMSProp
10.4
1.2
contral foo
10-1
-0.25
0.8
-0.50
10-4
-0.55
10-33
0.4
-1.00
RMSProc contral flom
11
RMSProp
central flww
0.2
10
500
1000
1500
2030
2500
500
1000
1500
2000
2500
0
520
1000
1500
2000
2500
Figure 35.14: RMSProp central flow for a Transformer with CE loss, 77 =2c-05, Bg =0.95, -c-088 and bias correction.
113

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigenvalues
real Vs. predicted oscillation covariance
2.5
slabie flow
RMSProp
0.00325
along Hio digetwactors
empiical berience
2.0
RMESrro (average)
flow (top)
0,00100
eigenvalues
central flow
1.5
coniral flow prediction
2
0.00075
0.0005
1.0
1-
500
1000
1500
2000
2500
500
1000
1500
2000
25000
5.00
1,000
15000
2000
2500
gradient norm?
top hessian eigenvalues
test accuracy (%)
100
2.0
RMShmp
500
(top)
1.5
60-
400
RMSPI0P
40-
RM5Phop (average)
200
central flow
20-
RNSPup (midgeints)
central
500
1000
1500
2000
2500
500
1000
1500
2000
25.00
500
1000
1500
2000
2500
notwork outputs on test example
coordinates of
distance to RMSPrap
10
central
10a7
10-2
RNSProp
10-11
RM SPr0
1000
1500
2000
2500
500
1,000
1500
2000
2500
5.0
1000
1500
2000
2500
Figure 35.15: RMSProp central flow for a Transformer with CE loss, =4c-05, 82 =0.95, =c--0 and bias correction.
train loss
effective hessian eigenvalues
predicted oscillation covariance
stable flom
RMSSPra
empirical varlance
RM43hmp
central flos
0.0005
ctgeneecttra
1.2
RMSProp (awerage)
stable foo (top)
B15 clgenvalues
central
0.0004
1.0
3.0003
0.8-
0.0002
0.6-
3.0001
1000
3000
40000
3000
3.0000
3000
gradient norm?
top hessian cigenvalues
test accuracy (%)
70
20
RMSProp RMSProp (average)
RMERrrp
floe
control flow centrol flo prediction
stable fow (top)
60
15000
50-
10000
40
5000
RMSProp Imi
30
1000
3000
4000
1000
3000
4000
1000
2000
3000
4000
5000
network outputs on test example
coordinates
distance RMSProp
1.0
RMSPPrp
0.10
centrol
flow
0.5
10
0.06
10-4
-0.
10-a
-1.5
RMSProp central
10-19
0.02
1000
2000
3000
4000
5000
2000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Figure 36: RMSProp central flow for a Mamba with CE loss, 7c--6 (3 =0.95, e = 1c-08 and bias correction.
114

---

Published as a conference paper at ICLR 2025
train loss
top effective hessian eigervallus
real vs. predicted oscillation covariance
1.4
wable flow
amptiiaal varience slong Rist elgenvectors
1.2
RM5Prop laverage)
1.0
central flow
central flom prediction
0,6010
0.6
RMShhpp
0.0005
central flow
0.4
stable flow
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1500
2000
3000
4000
5000
gradient norm?
top hessian eigenivalues
test accuracy (16
RMS/Pop
RMSShrp
125
MM5/top
65-
100
central flow prediction
flo
30000
(top)
75
20000
55
50
50
25
10000
RNSPup (midgeints)
40
central
1000
2000
9000
4000
5000
1000
2000
3000
4000
5000
1000
2000
1000
4000
5000
network outputs on test example
coordinates of
distance RMSProp
0.20
central
16-
10
0.15
10
0.10
2-
RMSRrap
10:43
RMSProp
1000
2000
3000
4000
5000
1000
2000
3000
40.00
5000
1300
2000
3500
4000
5000
Figure 35.17: RMSProp central flow for a Mamba with CE loss, = 1e-05, B2 =0.95, =1c-08 and bias correction.
train loss
top effective hessian elgenvalues
real vs. predicted escillation covariance
1.4
0.005
empirical varlance
1.2
RMSProp
RMSProp laveragel
0.004
sigunvulues
central flow
0.003
0.8
0.002
0.4
flow
0.2
stable flow (top!
0.000
1000
2000
5con
2000
3000
5o0n
1800
2000
3000
gracieett norm?
top hessian eigenvalues
test accuracy (%)
RAtEProp
30000
RMSProp
PMSPrup (overage)
central flom sable flow (top)
300
central flow
25000
flow prodiction
20000
200
15000
60
10000
RMSProp imidpoints) central now
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
10:00
2000
3000
4000
5000
outputs example
coordinates of
RMSProp
RoMiProp
central flow
central flow
10-2
stable flw
10
10
0-
10-4
10-10
IIMPProp central flow
-2
0
1000
2000
1000
4000
5000
1000
2000
3000
4000
5000
1000
7000
3000
4000
5000
Figure 38;18; RMSProp central flow for a Mamba with CE loss, =2c-05, =0.95, =cc-08, and bias correction.
115


Reviews:
[{"review_id": "d5JziU6soC", "reviewer": "Reviewer_DEu8", "summary": "The paper investigates optimization methods in deep learning, focusing on the \"edge of stability\" regime. Since the gradient flow differential equation fails to capture both the oscillatory behaviors and the averaged trends of gradient descent in many neural network training scenarios, the authors introduce central flows, a new differential equation that approximates the \"time-averaged\" optimization dynamics up to high accuracy. The paper also provides valuable insights into adaptive optimizers, such as RMSProp, and introduces the concept of acceleration via regularization, where adaptive optimizers naturally steer trajectories toward low-curvature regions, enabling larger steps and improving optimization efficiency.", "questions": "The authors present extensive numerical experiments that demonstrate the central flow as a strong approximation to average optimization dynamics. While the numerical results are impressive, could the authors discuss any potential limitations or scenarios where the central flow may not perform as effectively? Understanding the boundaries of the method’s applicability would provide helpful context for further investigation.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The article is clearly written and reader-friendly, with insightful discussions on the oscillatory behavior of optimization methods.\n\n2. The time-averaging and Taylor expansion framework is intuitively sound and appears to be broadly applicable, capable of handling various adaptive optimizers, and could potentially be extended to more general optimization methods.\n\n3. Extensive empirical results are presented, supporting the validity of the central flow.\n\nOverall, I think this is a nice paper.", "weaknesses": "The central flow derivations are grounded in empirical observations, and several mathematical steps rely on informal reasoning. For example, it is not clear how to rigorously justify the time averaging step and quantify the high-order error in the Taylor expansion. Rigorous proofs or formal analysis would enhance the robustness of the claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uLvVC9ZKKF", "reviewer": "Reviewer_mqGu", "summary": "This work develops a central flow approach with which to study the training dynamics of neural networks. The authors use this approach to investigate the edge of stability and optimization via RMSProp. The authors argue that their analysis can shed new light on how RMSProp adapts to a local loss landscapes.", "questions": "1. Why is there a mismatch between Fig. 4 and Fig. 5, in terms of sharpness and $2/S$? \n\n2. How does the sharpness evolve (for a numerical example) when optimizing using full RMSProp? \n\n3. Do the authors think this approach can be used when there may not be an obvious choice of $\\sigma$? \n\n4. Why are $x_t u_t$ not present in Eq. 2?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This work was well written and well structured, making it easy to follow along. \n\n2. The authors analyzed a wide range of neural networks, including modern architectures such as ViT.\n\n3. The authors developed their approach in settings where the interpretation was straightforward to make their point more clear. \n\n4. The figures were well done and easy to read. \n\n5. The authors convincingly show that their central flow approach accurately captures the general behavior of neural network model training. \n\n6. The identification of a possible weakness of the pre-conditioning of RMSProp is interesting and suggestive of a way in which to use the authors' work to improve DNN optimization.", "weaknesses": "1. My main reservation of this work as it stands (which I hope the authors can easily clarify for me), is the claim that the central flow analysis of RMSProp-Norm demonstrates how the optimizer regularizes against sharpness and moves towards regions of parameter space with lower sharpness. While I understand the argument mathematically, and the analysis of the central flow (Fig. 5) illustrates this point, it does not seem consistent with the results of training ViT (Fig. 4). In particular, in Fig. 4 the Hessian eigenvalues increase with training and $2/S$ decreases (which more closely mirrors the ablation in Fig. 5). This seems to suggest that training is moving towards a region with greater sharpness. How do the authors reconcile this? Are the authors claiming that the network would be headed towards an even more sharp region if it was not for RMSProp-Norm, and that the optimizer is reducing sharpness, albeit not as completely as in the case of just analyzing the central flow? \n\n2. I think this work would also benefit from showing the sharpness of the full RMSProp. Providing numerical evidence that use of RMSProp steers the network away from sharp areas of parameter would make this point more salient. \n\n3. While the goal of understanding edge of stability was motivated, I think this work could also benefit from more discussion on how the central flow approach could be used for other training dynamics phenomena. In particular, as I understood this work, one of the keys to utilizing the method successfully was the ability to find $\\sigma$. This was only possible since the training limited to some fixed value. How general do the authors expect this kind of solving for $\\sigma$ to be? \n\n4. The authors show that the distance between the central flow weights and the GD weights are small (Fig. 3). Is this because the distance is an average across many weights? Or is the difference between oscillations and no oscillations a matter of small differences in the weights? \n\n5. One last point that was unclear to me is why $x_t u_t$ are not present in Eq. 2, since they are time varying and I would expect $d w(t) / dt$ to depend on their value. \n\nMINOR POINTS: \n1. PSD should be defined. \n\n2. The authors should make a note somewhere that the footnotes are at the end of the main text. It was confusing not to see them at the bottom. \n\n3. Figure 1 should be referenced in the text. \n\n4. The figure in Sec. 4.3.1 is not referenced anywhere nor does it have a label.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cA5otIsh3f", "reviewer": "Reviewer_X29f", "summary": "The authors developed a methodology for creating optimizers in the context of deep learning. They empirically demonstrated that central flows govern long-term optimization trajectories, providing insight into the behavior of optimizers.", "questions": "What are the geometric implications of the curvature penalty?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The work presents novel contributions. Additionally, it represents an advancement in the study of locally analyzed optimizers and suggests progress for other optimizers.", "weaknesses": "1. In the abstract, the introduction and related work the authors repeats the same phrase 'Optimization in deep learning remains poorly understood' (lines 11, 26 and 68).  Maybe the authors can phrases like: 'How optimization works in deep learning continues to be largely unclear' or 'The optimization process in deep learning is still not well understood.'\n2. The contributions of the work are unclear in the introduction, which is somewhat difficult to follow in that sense. Maybe the authors can add a explicit numbered list with the contributions.\n3. In general, it would have been valuable to provide a bit more discussion or formal mathematical argumentation on how some equations were derived (e.g., equations 1 or 7). \n4. The proofs are generally well-written and clear, but since they are in the appendix, the flow of the paper's main narrative can sometimes be interrupted. Would be helpful to add some discussion of the results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7Bka2PkgGa", "reviewer": "Reviewer_HuNC", "summary": "The paper shows that an optimizer's implicit behavior comprising complex oscillatory dynamics (EOS) can be captured explicitly by ODE (termed as \"central flow\" by authors) modeling a time-averaged optimization trajectory. This a very interesting idea and probably the first explicit & numerically precise characterization that predicts a long-term optimization path well. Using central flow they analyze the local adaptivity of RMSProp and provide insight into an interesting phenomenon termed \"acceleration via regularization\" which attempts to reason the efficacy of adaptive optimizers in general.", "questions": "1. The ODE derivation is quite general. However, I wonder if it's possible to analyze non-convex objective functions with central flow and if they still track real optimization trajectories. \n2. Can the authors comment on why they chose to analyze RMSProp only as opposed to using ADAM? What sort of issues did they face of trying with ADAM, if any?\n3.  Central flow clearly shows how adaptive optimizers like RMSProp additionally account for the curvature regularization induced by oscillations. While it can be seen that they bias trajectory towards lower curvature, I wonder if central flow helps understand what happens in the presence of pathological curvatures (specifically low-curvature critical points such as saddle points). Does central flow demonstrate whether these points are circumvented when using adaptive optimizers?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. This is a novel study to deeply understand an optimizer's implicit dynamics and attempt to explicitly characterize it. The authors exploit central flow as a tool to analyze adaptive optimizers (RMSProp) and consequently draw interesting observations regarding their efficacy. They provide a strong theoretically backed argument as to how the EOS regime leads to implicit step size selection and curvature reduction.\n2. A new characterization of phenomena \"acceleration via regularization\" (observed using central flow) where they show how adaptive optimizers (RMSProp) navigate to a low generalization landscape where they can take larger steps. This claim is backed by a comprehensive set of experiments that showcase their claim precisely. \n3. The development and intuition of central flow and how it tracks time-averaged optimization trajectory to account for precise implicit behavior is robust. While the derivation might have started from informal reasoning (as accounted in the paper) I find the subsequent logic intuitive and reasonable. Each claim is backed by theoretical analysis in the Appendix which also looks correct. \n4. Detailed explanation of the experimental setup and good accounting of the limitations. They make it extremely clear what is the purpose of their work while also providing a comparison against other ODE methods like stable flow. Each claim is backed by a relevant ablation or experimental evaluation. \n5. There are comprehensive notes for the reader to provide quick additional reference or context to a claim being made in the paper which is very useful.", "weaknesses": "Some comments\n1. In the large-scale evaluation reported, I majorly see MSE loss being used. How does the ODE analysis deviate or negatively impacted by commonly used loss functions for classification tasks? While this might be part of future work I think addressing such limitations in the main body would improve the paper.\n2. The analysis relies heavily on the largest eigenvalue estimation for ODE analysis. I think it might be helpful to account for how precise/reliable such estimations are when empirically implemented. \n3. I might have missed this but I do not see any ablations on the behavior of $\\sigma^2$, that models $\\mathbb E[x_t^2]$ (equation 2). Is it because it has a fixed value due to the 3 desiderata described in line 249? In either case, I think it might be good to show its evolution across training to support its expected behavior.\n\nMinor weakness\n1. Figure 1: x-axis labels missing. While I can infer what is being plotted I think the figures should be precise on labels for faster understanding. \n2. Figure 3, the legend is difficult to read might be possible to increase the font a bit.  \n3. Figure 4: The gray curves are mostly hard to see. I think it might be helpful to use a color palette that makes the curves more discernable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WP19QoqIGV", "reviewer": "Reviewer_VSTS", "summary": "The paper uses mathematical reasoning to derive central flows of optimization functions. Central flows are functions which predict the averaged behavior of oscillating optimization functions when applied to neural networks.  These functions are empirically accurate, and reveal several interesting insights into the behavior of different optimization functions.\n\nIn gradient descent, central flows are used to describe the oscillatory effects of sharpness, and the sharpness-decreasing effects of oscillation, resulting in a steady-state (on average) of the sharpness of the points of the path taken by gradient descent.  These effects have been described previously, but not so generally.\n\nIn RMSProp, central flows are used to reason about what the paper calls effective sharpness.  They show that RMSProp, despite being a first-order algorithm, nonetheless adapt to the second-order curvature.  The path of the optimizer is driven to low-curvature spaces, where it can safely take large steps, improving its performance.", "questions": "Would the creation of central flows be easily created for other optimization algorithms, or do they depend upon specific characteristics of these two algorithms?  If they were second-order algorithms, would the derivation break down?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 4, "strengths": "- The work is important.  Understanding the behaviors of optimization algorithms is very important, and this paper provides new insights through a novel approach.\n- The paper is engaging.  It is well written and easy to read.\n- Conclusions and key takeaways are clearly stated.\n- The work is wide-ranging, and makes a number of important conclusions in a single paper.", "weaknesses": "The main weakness is that this is not really a conference paper.  Essential parts that justify the work (the mathematical derivations and experimental results that support the conclusions) are included in the 41-page appendix.  There is some point where the appendix is so essential, this cannot really be considered a 10-page paper.  The work is interesting, important, and well laid out, but the ten pages are not, on their own, an acceptable paper, as the appendices are so voluminous and essential that the reader would have insufficient reason to believe the claims of the authors without them. I cannot find anything in the reviewer guidelines to help me understand the expectations in this area, so I leave it to the AC for guidance and discussion.\n\nIf this is a problem for the conference, the most useful thing to promote to the core paper would be some of the experiments which demonstrate the fidelity of the predictions from the central flows to the true observed optimizations.  If it's not a problem, I believe this otherwise-excellent paper should be highlighted by the conference.\n\nI have skimmed these appendices, but not closely checked the mathematics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)