Understanding Optimization in Deep Learning with Central Flows

## Abstract

Abstract Optimization in deep learning remains poorly understood, even in the simple setting of deterministic (i.e. full-batch) training. A key difficulty is that much of an optimizer’s behavior is implicitly determined by complex oscillatory dynamics, referred to as the “edge of stability.”
The main contribution of this paper is to show that an optimizer’s implicit behavior can be explicitly captured by a central flow : a differential equation which models the time-averaged optimization trajectory. We show that these flows can empirically predict long-term optimization trajectories of generic neural networks with a high degree of numerical accuracy. By interpreting these flows, we reveal for the first time 1) the precise sense in which RMSProp adapts to the local loss landscape, and 2) an acceleration via regularization mechanism, wherein adaptive optimizers implicitly navigate towards low-curvature regions in which they can take larger steps. This mechanism is key to the efficacy of these adaptive optimizers. Overall, we believe that central flows constitute a promising tool for reasoning about optimization in deep learning.

## 1 Introduction

Figure: Figure 1: The central flow models the time-averaged (i.e. smoothed) trajectory of the oscillatory optimizer. In this representative cartoon, gradient descent takes an oscillatory path through weight space. The central flow (black, dashed) is a smooth curve that characterizes this trajectory, whereas gradient flow (gray, dashed) takes a different path.
Refer to caption: /html/2410.24206/assets/x1.png

Optimization in deep learning remains poorly understood, even in the simple setting of deterministic (i.e. full-batch) training.
A key difficulty is that much of an optimizer’s behavior is determined *implicitly* by complex oscillatory dynamics (Xing et al., ; Wu et al., ; Jastrzębski et al., , ; Cohen et al., ). As a result, an optimizer’s update rule often sheds little light on its actual behavior.

To address this challenge, we develop a new methodology for analyzing optimization in deep learning. To analyze an optimizer, we derive a *central flow*: a differential equation which directly models the time-averaged (i.e. smoothed) trajectory of the oscillatory optimizer.
We use informal mathematical reasoning to derive central flows for gradient descent ([Section 3](#S3)), a simple adaptive optimizer ([Section 4](#S4)), and RMSProp, i.e. Adam with no momentum ([Section 5](#S5)), and we empirically demonstrate that these flows can predict long-term optimization trajectories of neural networks with a high degree of numerical accuracy.
We are unaware of any other theoretical analyses of deep learning optimization with a similar degree of predictive power.

Unlike an optimizer’s update rule, the central flow directly exposes the optimizer’s behavior, thereby making it possible to reason about optimization.
For example, our analysis reveals, for the first time, the precise sense in which RMSProp adapts its effective step sizes to the local loss landscape.
It also shows that the two adaptive optimizers implicitly regularize the curvature along their trajectories, giving rise to a mechanism we call *acceleration via regularization*, whereby implicitly regularizing curvature drives the trajectory towards low-curvature regions where the optimizer can take larger steps. This mechanism is key to the efficacy of these optimizers.

We are optimistic that our methodology holds promise as a framework for analyzing, and perhaps even inventing, deep learning optimization algorithms beyond the ones studied here.

## 2 Related Work

##### Edge of Stability

The dynamics of optimization in deep learning remain poorly understood, even in the seemingly 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 *critical threshold* $2/\eta$, and (2) the algorithm oscillates along high-curvature directions without diverging (Xing et al., ; Wu et al., ; Jastrzębski et al., , ; Cohen et al., ). These dynamics could not be explained by existing optimization theory, which led Cohen et al. () to observe that there was no explanation for how or why gradient descent can function properly in deep learning.

Subsequently, several studies sought to theoretically explain EOS dynamics. Some works rigorously analyzed EOS dynamics on specific objective functions
(Agarwala et al., ; Ahn et al., ; Chen and Bruna, ; Even et al., ; Kreisler et al., ; Song and Yun, ; Li et al., ; Wu et al., ; Zhu et al., ), while other works (Arora et al., ; Lyu et al., ; Damian et al., ), gave generic analyses based on a local *third-order* Taylor expansion of the loss, which is one order higher than is normally used in the theoretical analysis of gradient descent.
Similar arguments were first used by Blanc et al. () to study implicit regularization in SGD.

Our analysis is most directly inspired by Damian et al. (), which rigorously analyzed EOS in the special case where gradient descent oscillates along a single direction. Whereas they analyze the *fine-grained* oscillatory dynamics, we argue that analyzing the *time-averaged* dynamics is simpler, and is sufficient for many purposes. We first reproduce their main result using a simple, albeit non-rigorous, time-averaging argument.
We then show that this time-averaging methodology easily extends to the more realistic and challenging setting where gradient descent oscillates along *multiple* directions simultaneously, as well as to the analysis of two adaptive optimizers.

##### Understanding Adaptive Optimizers

Ma et al. () observed that RMSProp and Adam oscillate, and Cohen et al. () showed that such dynamics can be viewed as an adaptive version of the edge of stability, a finding which we will leverage.
Khaled et al. () and Mishkin et al. () observed that on quadratic functions, certain adaptive optimizers implicitly adapt their effective step size to the maximum stable step size; we show this holds more generally, beyond quadratics.
Experiments in Roulet et al. () and Wang et al. () are explained by the phenomenon we call “acceleration via regularization.”
Many works have also conducted rigorous convergence analyses of adaptive optimizers, generally focused on deriving rates of convergence to a global minimizer or stationary point (Duchi et al., ; Reddi et al., ; Chen et al., , ; Zaheer et al., ; Zou et al., ; Défossez et al., ; Li and Lin, ; Chen et al., ; Wang et al., ; Yang et al., ; Guo et al., ; Shi et al., ; Zhang et al., ; Crawshaw et al., ; Li et al., ; Wang et al., ; Hong and Lin, ; Zhang et al., ; Wang et al., ; Hübler et al., ).

## 3 Gradient Descent

We begin by studying the simplest first-order optimizer: gradient descent with a fixed step size $\eta$.

$$ $\displaystyle w_{t+1}=w_{t}-\eta\nabla L(w_{t}).$ (1) $$

Traditional analyses of gradient descent do not apply in deep learning settings (Cohen et al., ).
In this section, we develop a new analysis technique which is capable of analyzing gradient descent in deep learning.

- •
In [Section 3.1](#S3.SS1), we describe the dynamics of gradient descent in deep learning. Gradient descent typically operates in a regime termed the *edge of stability* (EOS) characterized by complex oscillatory dynamics (Cohen et al., ).
The oscillations dominate the trajectory over short timescales, causing the loss to behave non-monotonically, as shown in [Figure 2](#S3.F2). More importantly, they also affect the *long-term* trajectory by implicitly steering gradient descent away from high-curvature regions of the loss landscape. As a result, the raw update rule [eq. 1](#S3.E1) is largely uninformative as to actual behavior of the algorithm.
- •
In [Section 3.2](#S3.SS2), generalizing a result by Damian et al. (), we show that while the gradient descent trajectory is highly oscillatory, its *time-averaged* (i.e. smoothed) trajectory can be captured by a differential equation called a *central flow*.
We derive this central flow using informal mathematical reasoning, and empirically demonstrate that it can successfully predict long-term optimization trajectories of neural networks with a high degree of numerical accuracy. For example, [Figure 2](#S3.F2) shows that the central flow can predict training loss curves on a variety of neural network architectures. Because the central flow makes *explicit* the path that is *implicitly* taken by gradient descent, it makes reasoning about the optimizer’s behavior easier.
- •
In [Section 3.3](#S3.SS3) we use this central flow to understand the behavior of gradient descent. For example, we show that gradient descent’s non-monotonic loss curve can be viewed as the superposition of the loss along the central flow plus a contribution from the oscillations. The loss along the central flow is a smoothly varying quantity that monotonically decreases, and therefore constitutes a hidden progress metric for gradient descent.

Our analysis of gradient descent will set the stage for our subsequent analyses of more complex optimizers.

Figure: Figure 2: The central flow accurately predicts the loss curve of gradient descent. We train three architectures at three different learning rates on a subset of CIFAR-10. For each gradient descent trajectory, we also run the central flow starting from the same initialization. The central flow (black) accurately predicts the time-averaged, i.e. smoothed, loss curves of gradient descent (colors, solid). Raw loss curves are plotted in semi-transparent colors.
Refer to caption: /html/2410.24206/assets/x2.png

### 3.1 The Dynamics of Gradient Descent

The *sharpness* $S(w):=\lambda_{1}(H(w))$, defined as the largest Hessian eigenvalue, plays a key role in the dynamics of gradient descent.
In particular, gradient descent oscillates whenever the sharpness is “too large” relative to the learning rate $\eta$.
For example, consider optimizing a one-dimensional quadratic objective $L(x)=\tfrac{1}{2}Sx^{2}$, which has global sharpness $S$. The gradient descent iterates $\{x_{t}\}$ evolve via $x_{t+1}=(1-\eta S)x_{t}$. If $S$ exceeds the *critical threshold* $2/\eta$, then $(1-\eta S)<-1$, so the iterate $x_{t}$ flips signs and grows in magnitude at each step, i.e. gradient descent oscillates with exponentially growing magnitude. More generally, for a quadratic objective in multiple dimensions, gradient descent oscillates with exponentially growing magnitude along all Hessian eigenvectors with eigenvalues exceeding $2/\eta$.
While deep learning objectives are not globally quadratic, a *local* quadratic Taylor approximation suggests that if the sharpness exceeds $2/\eta$, gradient descent will oscillate along the top Hessian eigenvector(s). This suggests that gradient descent cannot function properly when the sharpness exceeds $2/\eta$.

Figure: Figure 3: Gradient descent on the quadratic $\frac{1}{2}Sx^{2}$ with different values of $S$ but the same learning rate $\eta$.
Refer to caption: /html/2410.24206/assets/x3.png

In light of this discussion, why does gradient descent converge in deep learning? The natural explanation is that the sharpness remains below $2/\eta$ throughout training. Yet, to the contrary, Cohen et al. () observed that gradient descent tends to *continually increase the sharpness*, an empirical phenomenon dubbed “progressive sharpening.” As a result, the sharpness generally rises above $2/\eta$ during training.
When this occurs, the optimizer indeed oscillates with growing magnitude along the sharpest direction(s), just as a quadratic Taylor approximation would predict.
Yet, gradient descent does not subsequently diverge. Instead, as shown in [Figure 4](#S3.F4), it enters an operating regime called the *edge of stability* (EOS), in which: (1) gradient descent oscillates along sharpest directions without diverging, and (2) the sharpness equilibrates around the critical threshold $2/\eta$. When only the top eigenvalue has reached $2/\eta$ (steps 600–1450 in [Figure 4](#S3.F4)), gradient descent oscillates along the top Hessian eigenvector, and this eigenvalue equilibrates around $2/\eta$. When *multiple* Hessian eigenvalues have reached the critical threshold $2/\eta$ (steps 1450–2250 in [Figure 4](#S3.F4)), gradient descent oscillates simultaneously along all the corresponding eigenvectors, and all such eigenvalues equilibrate around $2/\eta$.
These dynamics cannot be explained by traditional optimization theory, as the sharpness frequently exceeds the critical threshold $2/\eta$.

Figure: Figure 4: A typical gradient descent trajectory. We train a ViT on a subset of CIFAR-10. The sharpness (top Hessian eigenvalue) rises until reaching the critical threshold $2/\eta$, at which point the dynamics enter the edge of stability (EOS) regime. At EOS, the sharpness stays dynamically regulated around the critical threshold $2/\eta$ (center), the optimizer oscillates along high-curvature directions (right), and the loss decreases non-monotonically (left).
Refer to caption: /html/2410.24206/assets/x4.png

Damian et al. () showed that to understand EOS dynamics, it suffices to consider a *third-order* Taylor expansion of the loss function, which is one order higher than traditionally used in analyses of gradient descent. Their analysis focused on the special case of a single oscillating direction. Let us informally sketch this argument. Suppose that gradient descent is oscillating around a reference point $\overline{w}$, along the top Hessian eigenvector $u$, with current magnitude $x$, so that $w=\overline{w}+xu$. Then Taylor-expanding the gradient $\nabla L(w)$ around the reference point $\overline{w}$ gives (see [Lemma 1](#Thmlemma1) for the full derivation):

$$ $\displaystyle\nabla L(w)\quad=\quad\quad\underset{\mathclap{\text{(1) gradient at reference point}}}{\boxed{\nabla L(\overline{w})}}\quad+\quad\underset{\text{(2) oscillation}}{\boxed{xS(\overline{w})u}}\quad+\quad\underset{\text{(3) sharpness reduction}}{\boxed{\tfrac{1}{2}x^{2}\nabla S(\overline{w})}}\quad+\quad\mathcal{O}(x^{3})$ (2) $$

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 $\eta$ automatically includes a gradient step on the *sharpness* of the loss with step size $\tfrac{1}{2}\eta x^{2}$.
Thus, oscillations automatically trigger reduction of sharpness.
This acts as a form of negative feedback which automatically stabilizes training.
Damian et al. () conducted a fine-grained analysis of EOS dynamics in the setting of a single oscillating direction (e.g. steps 600–1450 in [Figure 4](#S3.F4)).
In this setting, the EOS dynamics consist of repeated cycles in which: (a) progressive sharpening drives the sharpness above $2/\eta$; (b) this triggers growing oscillations along the top Hessian eigenvector; (c) such oscillations force the sharpness back below $2/\eta$; (d) the oscillations consequently shrink in magnitude. When *multiple* Hessian eigenvalues have reached the critical threshold $2/\eta$ (e.g. steps 1450–2250 in [Figure 4](#S3.F4)), causing gradient descent to oscillate simultaneously along all the corresponding eigenvectors, a Taylor expansion analogous to [eq. 2](#S3.E2) shows that these oscillations also act to reduce the curvature in these directions. However, the dynamics in this setting are more complex, and may be chaotic.

At the edge of stability, oscillations dominate the short-term dynamics of optimization, accounting for much of the loss and gradient norm. Perhaps more importantly, they also affect the long-term optimization trajectory, by steering the optimizer away from high-curvature regions of weight space.
Thus, the raw update rule [eq. 1](#S3.E1) is largely uninformative as to the behavior of the optimizer; instead, much of the optimizer’s behavior is implicitly determined by these complex oscillatory dynamics.

In the next section, we will aim to understand the overall path that gradient descent takes through weight space.
While the *fine-grained* dynamics of gradient descent at the edge of stability are difficult to analyze, we will see that a heuristic *time-averaging* argument suffices to characterize this overall path.
Our time-averaging analysis in the next section will not only recover (albeit non-rigorously) the analysis of Damian et al. () for a single oscillating direction, it will also easily generalize to the more realistic and challenging setting of multiple oscillating directions.

### 3.2 Deriving the Gradient Descent Central Flow

In this section, we will derive a differential equation which models the trajectory of gradient descent. The standard continuous-time approximation to gradient descent is the gradient flow:(^1^11We fold $\eta$ into the definition of gradient flow so that there is a correspondence between step $t$ of gradient descent and time $t$ of gradient flow. This will especially be useful when analyzing adaptive optimizers where the effective step size is a dynamic quantity.)

$$ $\displaystyle\frac{dw}{dt}=-\eta\nabla L(w).$ (3) $$

Cohen et al. () observed that trajectory of gradient descent is well-approximated(^2^22 Barrett and Dherin () argued that the accuracy of the gradient flow approximation can be improved by adding a penalty on the squared gradient norm. However, their modified flow does not hold in the EOS regime, and in the stable regime, we found that the accuracy improvement it brings is relatively small. Therefore, for simplicity, we leave out any such term from our flows.)(^3^33It remains theoretically unexplained why gradient flow is such a good fit to gradient descent. Existing bounds for the distance between gradient descent and gradient flow increases exponentially with time, with an exponent determined by the most negative Hessian eigenvalues (Elkabetz and Cohen, ). Empirically, such bounds are overly conservative.)
by that of gradient flow so long as training is *stable*, i.e. so long as the sharpness $S(w)$ remains below $2/\eta$. However, once the sharpness reaches $2/\eta$ and the optimizer enters the EOS regime, gradient descent departs from the gradient flow trajectory and takes a different path.
For example, the gray line in [Figure 5](#S3.F5)(d) plots the weight-space distance between the gradient descent iterate at step $t$ and the gradient flow solution at time $t$. This distance remains small so long as training is stable, but starts to grow large once training enters EOS, indicating that these trajectories diverge.

We now derive a more general differential equation, which we call a *central flow*, that approximates the trajectory of gradient descent even at the edge of stability.
The central flow directly models the *time-averaged* (i.e. smoothed) trajectory of the oscillatory optimizer. In other words, the central flow averages out the oscillations while retaining their lasting effect on the optimization trajectory.

We will derive the central flow using informal mathematical reasoning, and we will empirically demonstrate that it can accurately predict long-term optimization trajectories of neural networks with a high degree of numerical accuracy.
For example, the black line in [Figure 5](#S3.F5)(d) plots the weight-space distance between gradient descent at step $t$ and the central flow at time $t$. Observe that this distance stays small throughout training, indicating that while the gradient descent and gradient flow trajectories diverge, the gradient descent and *central* flow trajectories remain close.

Figure: Figure 5: Central flow for gradient descent. A ViT is trained on CIFAR-10 using gradient descent with $\eta=2/100$ (blue). Gradient descent enters EOS at step 600 and after step 1450 multiple eigenvalues are unstable (dotted red). The central flow (black) accurately models gradient descent even at EOS, whereas gradient flow (gray) follows a different path. The distance between the central flow and gradient descent (black) even remains smaller than the distance between consecutive iterates of gradient descent (blue) throughout the trajectory.
Refer to caption: /html/2410.24206/assets/x5.png

In [Section 3.3](#S3.SS3), we will use this central flow to interpret the behavior of gradient descent. Because the central flow is a smooth curve, it satisfies the chain rule. This will allow us to easily reason about properties of the gradient descent trajectory, including its non-monotonic training loss curve.

#### 3.2.1 The Special Case of One Unstable Eigenvalue

We will introduce our time-averaging methodology by analyzing the special case when only the largest Hessian eigenvalue has crossed the critical threshold $2/\eta$, and gradient descent oscillates along a single direction — the corresponding eigenvector. The general case of multiple unstable eigenvalues will then be tackled in [Section 3.2.2](#S3.SS2.SSS2).

We start our analysis at the instant when the sharpness $S(w)$ reaches $2/\eta$. From this point onward, we will model the GD trajectory by $w_{t}=\overline{w}_{t}+x_{t}u_{t}$ where $w_{t}$ is the gradient descent iterate, $\overline{w}_{t}$ is the time-averaged iterate, $u_{t}$ is the top Hessian eigenvector at $\overline{w}_{t}$, and $x_{t}$ denotes the displacement between $w_{t}$ and $\overline{w}_{t}$ along the $u_{t}$ direction.(^4^44In reality, gradient descent oscillates around $\overline{w}$ in *two* directions: the $u$ direction and the $\nabla S(\overline{w})$ direction. The latter is responsible for the fluctuations in the sharpness. However, when modeling the time-averaged gradient, we will ignore the effects of oscillations along $\nabla S(\overline{w})$, which corresponds to Damian et al. (, Assumption 5). The success of our experiments validates that these effects can be safely neglected.)
Note that by definition, $\operatorname{\mathbb{E}}[x_{t}]=0$, i.e. the time-averaged displacement is zero. To track the evolution of $\overline{w}_{t}$, we begin by time-averaging both sides of the gradient descent update:

$$ $\displaystyle\overline{w}_{t+1}\enskip=\enskip\operatorname{\mathbb{E}}[w_{t+1}]\enskip=\enskip\operatorname{\mathbb{E}}[w_{t}-\eta\nabla L(w_{t})]\enskip=\enskip\overline{w}_{t}-\eta\underbrace{\operatorname{\mathbb{E}}[\nabla L(w_{t})]}_{\mathclap{\text{time-averaged gradient}}}.$ (4) $$

Thus, the time-averaged iterates follow the (negative) time-averaged gradient.
To approximate the time-averaged gradient, we time-average the Taylor expansion of the gradient given in [eq. 2](#S3.E2):(^5^55When time-averaging [eq. 2](#S3.E2), we assume that the eigenvector $u$ changes slowly relative to the displacement $x$ so that $\operatorname{\mathbb{E}}[x_{t}u_{t}]\approx\operatorname{\mathbb{E}}[x_{t}]u_{t}$.)

$$ $\displaystyle\operatorname{\mathbb{E}}[\nabla L(w_{t})]\enskip\approx\enskip\nabla L(\overline{w}_{t})+\underbrace{\cancel{S(\overline{w}_{t})\operatorname{\mathbb{E}}[x_{t}]u_{t}}}_{0\text{ because }\operatorname{\mathbb{E}}[x_{t}]=0}+\tfrac{1}{2}\operatorname{\mathbb{E}}[x_{t}^{2}]\nabla S(\overline{w}_{t})\enskip=\enskip\nabla L(\overline{w}_{t})+\underbrace{\tfrac{1}{2}\operatorname{\mathbb{E}}[x_{t}^{2}]\nabla S(\overline{w}_{t})}_{\mathclap{\text{implicit sharpness penalty}}}.$ (5) $$

Thus, the time-averaged gradient $\operatorname{\mathbb{E}}[\nabla L(w_{t})]$ is equal to the gradient at the time-averaged iterate $\nabla L(\overline{w}_{t})$, plus an implicit sharpness penalty whose strength is proportional to $\operatorname{\mathbb{E}}[x_{t}^{2}]$, the variance of the oscillations at step $t$. Substituting [eq. 5](#S3.E5) into [eq. 4](#S3.E4) and switching to continuous time, we therefore model the time-averaged iterates $\overline{w}_{t}$ by the sharpness-penalized gradient flow $w(t)$ defined by:

$$ $\displaystyle\frac{dw}{dt}=-\eta\quantity\Big[\;\nabla L(w)+\underbrace{\tfrac{1}{2}\sigma^{2}(t)\nabla S(w)}_{\mathclap{\text{implicit sharpness penalty}}}\;].$ (6) $$

Here, $\sigma^{2}(t)$ is a still-unknown quantity intended to model $\operatorname{\mathbb{E}}[x_{t}^{2}]$, the instantaneous variance of the oscillations at time $t$. This quantity also controls the strength of the implicit sharpness penalty.
To determine $\sigma^{2}(t)$, we argue that only one value is consistent with empirical observation.
Empirically, once the sharpness reaches the critical threshold $2/\eta$, it does not continue to rise indefinitely; rather, it remains dynamically regulated around $2/\eta$. Thus, we will enforce that the central flow never increases the sharpness $S(w(t))$ past $2/\eta$. The time derivative of the sharpness under a flow of the form [eq. 6](#S3.E6) can be easily computed using the chain rule:

$$ $\displaystyle\frac{dS(w)}{dt}=\expectationvalue{\nabla S(w),\frac{dw}{dt}}=\underbrace{\vphantom{\tfrac{1}{2}}\eta\expectationvalue{\nabla S(w),-\nabla L(w)}}_{\mathclap{\begin{subarray}{c}\text{change in sharpness}\\ \text{under gradient flow}\end{subarray}}}\;-\;\underbrace{\tfrac{1}{2}\eta\sigma^{2}(t)\norm{\nabla S(w)}^{2}}_{\mathclap{\begin{subarray}{c}\text{sharpness reduction}\\ \text{from oscillations}\end{subarray}}}\quad\quad$ (7) $$

When the first term, the change in sharpness under the gradient flow, is *negative*, gradient descent will leave the edge of stability and will once again follow gradient flow — this is made precise in [Section 3.2.2](#S3.SS2.SSS2). Therefore, we focus on the case where this first term is positive, i.e. where progressive sharpening holds. As the sharpness is currently at $2/\eta$ and must remain at $2/\eta$, we must have that $\tfrac{dS(w)}{dt}=0$. Since $\frac{dS(w)}{dt}$ is linear in $\sigma^{2}(t)$, we can easily solve for the unique value of $\sigma^{2}(t)$ that ensures $\tfrac{dS(w)}{dt}=0$:

$$ $\displaystyle\sigma^{2}(t)=\frac{2\expectationvalue{\nabla S(w),-\nabla L(w)}}{\norm{\nabla S(w)}^{2}}.$ (8) $$

Intuitively, this is the unique $\sigma^{2}(t)$ for which the downward force of oscillation-induced sharpness reduction “cancels out” the upwards force of progressive sharpening so the sharpness remains locked at $2/\eta$. The central flow for a single unstable eigenvalue is given by substituting this $\sigma^{2}(t)$ into [eq. 6](#S3.E6).

##### The Projection Interpretation

While we have derived the central flow as a sharpness-penalized gradient flow, it can be equivalently interpreted as a *projected* gradient flow. In particular, substituting [eq. 8](#S3.E8) into [eq. 6](#S3.E6) gives:

$$ $\displaystyle\frac{dw}{dt}$ $\displaystyle=-\eta\quantity[\nabla L(w)-\frac{\nabla S(w)\nabla S(w)^{\top}}{\norm{\nabla S(w)}^{2}}\nabla L(w)]=-\eta\Pi_{\nabla S(w)}^{\perp}\nabla L(w)$ $$

where $\Pi_{v}^{\perp}:=I-\frac{vv^{T}}{\norm{v}^{2}}$ denotes the projection onto the orthogonal complement of $v$. Intuitively, this flow projects out the $\nabla S$ direction from the gradient $\nabla L$ to keep the sharpness fixed at $2/\eta$. This re-derives a main result of Damian et al. () in a simpler, but less rigorous way.
(^6^66This flow actually differs slightly from the constrained trajectory in Damian et al. (). In Damian et al. (), the gradients along the constrained trajectory are forced to be orthogonal to the top eigenvector of the Hessian, which prevents this theory from applying to certain models (Kreisler et al., , Appendix A). Our time-averaging analysis does not have this restriction.)

#### 3.2.2 The General Case (Multiple Unstable Eigenvalues)

When multiple eigenvalues have reached the critical threshold $2/\eta$, gradient descent oscillates in the span of the corresponding eigenvectors. We assume the displacement $\delta_{t}:=w_{t}-\overline{w}_{t}$ between the true process and the time-averaged process lies in the span of these eigenvectors. For example, when only one eigenvalue has reached $2/\eta$, taking $\delta_{t}=x_{t}u_{t}$ recovers the analysis in [Section 3.2.1](#S3.SS2.SSS1). We model the iterate $w_{t}$ as $w_{t}=\overline{w}_{t}+\delta_{t}$ where, by definition, $\operatorname{\mathbb{E}}[\delta_{t}]=0$. As above, the time-averaged iterates $\overline{w}_{t}$ should follow the time-averaged gradient $\operatorname{\mathbb{E}}[\nabla L(w_{t})]$. To compute the time-averaged gradient, we first Taylor-expand the gradient around $\overline{w}_{t}$:

$$ $\displaystyle\nabla L(w_{t})\enskip\approx\enskip\nabla L(\overline{w}_{t})\;+\;H(\overline{w}_{t})\delta_{t}\;+\;\tfrac{1}{2}\nabla_{\overline{w}_{t}}\expectationvalue{H(\overline{w}_{t}),\delta_{t}\delta_{t}^{T}}\;+\;\mathcal{O}(\|\delta_{t}\|^{3}).$ (9) $$

The third term in this Taylor expansion can be interpreted as the gradient of the curvature in the $\delta_{t}$ direction.(^7^77Indeed, notice that $\expectationvalue{H(\overline{w}_{t}),\delta_{t}\delta_{t}^{T}}=\delta_{t}^{T}H(\overline{w}_{t})\delta_{t}$, the curvature in the $\delta_{t}$ direction.) Time-averaging both sides of [eq. 9](#S3.E9) gives that the time-averaged gradient is equal to:

$$ $\displaystyle\operatorname{\mathbb{E}}[\nabla L(w_{t})]\enskip\approx\enskip\nabla L(\overline{w}_{t})\;+\;\underbrace{\cancel{H(\overline{w}_{t})\operatorname{\mathbb{E}}[\delta_{t}]}}_{0\text{ because }\operatorname{\mathbb{E}}[\delta_{t}]=0}\;+\;\underbrace{\tfrac{1}{2}\nabla_{\overline{w}_{t}}\expectationvalue{H(\overline{w}_{t}),\operatorname{\mathbb{E}}[\delta_{t}\delta_{t}^{T}]}}_{\text{implicit curvature penalty}}$ (10) $$

As above, the third term shows that the time-averaged gradient $\operatorname{\mathbb{E}}[\nabla L(w_{t})]$ contains an implicit curvature penalty whose strength and direction are determined by the *covariance* of the oscillations $\operatorname{\mathbb{E}}[\delta_{t}\delta_{t}^{T}]$. Substituting [eq. 10](#S3.E10) into the time-averaged gradient descent update (eq. [4](#S3.E4)) and switching to continuous time, we therefore model the time-averaged iterates $\overline{w}_{t}$ by the curvature-penalized gradient flow $w(t)$ defined by:

$$ $\displaystyle\frac{dw}{dt}=-\eta\quantity\Big[\;\nabla L(w)+\underbrace{\tfrac{1}{2}\nabla_{w}\expectationvalue{H(w),\Sigma(t)}}_{\mathclap{\text{implicit curvature penalty}}}\;].$ (11) $$

Here, $\Sigma(t)$ is a still-unknown quantity intended to model $\operatorname{\mathbb{E}}[\delta_{t}\delta_{t}^{T}]$, the instantaneous covariance of the oscillations at time $t$.
This matrix also controls the implicit curvature penalty.
As above, we will argue that only one value of $\Sigma(t)$ is consistent with empirical observation:

- 1.
Empirically, Hessian eigenvalues which reach the critical threshold $2/\eta$ do not continue to rise further; rather, each eigenvalue either remains dynamically regulated around $2/\eta$, i.e. it stays at the edge of stability, or it drops below $2/\eta$, i.e. it leaves the edge of stability (Cohen et al., ). Thus, we impose the condition that $\Sigma(t)$ should not allow any Hessian eigenvalues to rise beyond the critical threshold $2/\eta$.
- 2.
Since gradient descent oscillates within the span of the unstable eigenvectors, we impose the condition that $\Sigma(t)$, which models the covariance of these oscillations, should be supported(^8^88We say that a symmetric matrix $\Sigma$ is supported within a subspace $\mathcal{U}$ if $\operatorname{span}[\Sigma]\subseteq\mathcal{U}$.) within the span of the Hessian eigenvectors whose eigenvalue is equal to $2/\eta$.(^9^99For gradient descent, the unstable eigenvectors have eigenvalues which fluctuate around $2/\eta$. However, for the central flow (which prevents any eigenvalues from rising above $2/\eta$), the unstable eigenvectors have eigenvalues which are exactly equal to $2/\eta$.)(^10^1010Similar to footnote [4](#footnote4), this neglects the motion in the top Hessian eigenvalues. The success of our experiments justifies this simplification.)
- 3.
Since $\Sigma(t)$ models a covariance matrix, we impose the condition that $\Sigma(t)$ should be positive semidefinite.

These three conditions turn out to imply a unique value of $\Sigma(t)$. In particular, we detail in [Section C.3](#A3.SS3) that $\Sigma(t)$ is the unique solution to a type of convex program known as a *cone complementarity problem* (CCP). The central flow for gradient descent is defined as [eq. 11](#S3.E11) with this value of $\Sigma(t)$. We note that $\Sigma(t)$ can be efficiently represented numerically as it is a low rank matrix, with rank at most the number of unstable eigenvalues.

A formal definition for the gradient descent central flow based on this CCP formulation is given in [Definition 2](#Thmdefinition2).

##### Understanding the Central Flow

At *almost all* times $t$, the $\Sigma(t)$ that solves the CCP is equal to the unique $\Sigma$ which causes all Hessian eigenvalues currently at $2/\eta$ to remain fixed in place at that value.(^11^1111This $\Sigma$ can be computed by solving a linear inverse, generalizing [eq. 8](#S3.E8).) However, [eq. 11](#S3.E11) with this $\Sigma$ would never allow an eigenvalue to leave the edge of stability, and hence would not properly model the trajectory of gradient descent.(^12^1212Further, this $\Sigma$ can have negative eigenvalues, which is unphysical since $\Sigma$ models a covariance matrix.) The CCP formulation fixes this issue by allowing eigenvalues to enter and leave EOS when appropriate.

[Figure 5](#S3.F5) depicts a central flow trajectory (black line).
Initially, all Hessian eigenvalues are below $2/\eta$, so $\Sigma(t)=0$ and the central flow reduces to the gradient flow.(^13^1313If all Hessian eigenvalues are below $2/\eta$, condition #2 forces $\Sigma(t)=0$, and the central flow reduces to the gradient flow.)
Once the top Hessian eigenvalue reaches $2/\eta$ around step 600, $\Sigma(t)$ becomes a rank-one matrix, and the central flow keeps the top Hessian eigenvalue locked at $2/\eta$, as it mimics the effects of oscillating along the the top eigenvector direction.(^14^1414When one Hessian eigenvalue is at $2/\eta$, the CCP returns the rank-one matrix $\Sigma(t)=\sigma^{2}(t)\,u(t)\,u(t)^{T}$, where $u(t)$ is the top Hessian eigenvector and $\sigma^{2}$ is given by [eq. 8](#S3.E8); thus, the central flow reduces to the flow in [Section 3.2.1](#S3.SS2.SSS1).)
Once the second Hessian eigenvalue also reaches $2/\eta$ around step 1400, $\Sigma(t)$ becomes a rank-two matrix, and the central flow keeps the top two eigenvalues both locked at $2/\eta$, as it mimics the effects of oscillating simultaneously along the top two eigenvector directions.

##### The Projection Interpretation

The projection interpretation in [Section 3.2.1](#S3.SS2.SSS1) generalizes to the case of multiple oscillating directions. In particular, the central flow can be written as *projected gradient flow* on the constrained problem: $\min_{w}L(w)$ such that $S(w)\leq 2/\eta$, which projects out the components of the gradient that would increase sharpness above $2/\eta$. The precise formulation is given in [Definition 3](#Thmdefinition3).
This projection interpretation will be used in [Section 3.3](#S3.SS3) to show that the loss along the central flow decreases monotonically.

##### Empirical verification

While our derivation employs informal mathematical reasoning, our extensive experiments described in [Section 6](#S6) show that the central flow can accurately predict long-term optimization trajectories of neural networks. For example, in [Figure 5](#S3.F5), the central flow is an excellent approximation to gradient descent. We discuss the factors that affect the quality of this approximation in [Section 6](#S6). A promising direction for future research is to identify conditions under which the central flow can be proved to approximate the real optimization trajectory.

### 3.3 Interpreting Gradient Descent via its Central Flow

Figure: Figure 6: Reasoning about the loss curve using the central flow. While the gradient descent loss (blue) behaves non-monotonically, the central flow loss (solid black) monotonically decreases. The central flow can also predict the *time-averaged* loss of gradient descent (dashed black). Due to the implicit curvature regularization induced by oscillations, the slope of the central flow loss curve is less steep than that of gradient flow (orange lines).
Refer to caption: /html/2410.24206/assets/x6.png

We now use the central flow to understand the behavior of gradient descent. As a smooth flow(^15^1515Formally, the central flow is not smooth at *breakpoints*, i.e. times when an eigenvalue enters or leaves EOS. However, the central flow is always smooth *from the right*, which allows us to use the chain rule to reason about its behavior in the immediate *future*.), the central flow is an simpler object to reason about than the oscillatory gradient descent trajectory. In particular, we can use the chain rule to quantify the evolution of any differentiable metric.

For example, the rate of loss decrease under the central flow is simply given by $\frac{dL(w)}{dt}=\expectationvalue{\nabla L(w),\frac{dw}{dt}}.$ Combining this with the projection interpretation ([Definition 3](#Thmdefinition3)), we prove in [Proposition 1](#Thmproposition1) that $\frac{dL(w)}{dt}\leq 0$, i.e. that the loss along the central flow $L(w(t))$ is monotonically decreasing (black line in [Figure 6](#S3.F6)). This stands in stark contrast to the loss along the gradient descent trajectory, which behaves non-monotonically (blue line in [Figure 6](#S3.F6)). Thus, the central flow loss is a hidden progress metric for the optimization process.

As is visible in [Figure 6](#S3.F6), the gradient descent train loss $L(w_{t})$ is generally higher than the central flow train loss $L(w(t))$, due to the oscillations. However, since the central flow also models the oscillation covariance $\Sigma(t)$, it can render predictions for the *time-averaged* train loss along the gradient descent trajectory:

$$ $\displaystyle\operatorname{\mathbb{E}}[L(w_{t})]\;\approx\;L(w(t))+\tfrac{1}{2}\expectationvalue{H(w(t)),\Sigma(t)}\;=\;\underbrace{L(w(t))}_{\mathclap{\text{central flow loss}}}\;+\;\underbrace{\tfrac{1}{2}S(w(t))\tr(\Sigma(t))}_{\text{effect of oscillations}}$ (12) $$

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 a contribution from the oscillations.

Finally, we also show in [Proposition 1](#Thmproposition1) that the central flow decreases the loss at a *slower rate* than would the gradient flow: $\frac{dL(w)}{dt}>-\eta\|\nabla L(w)\|^{2}$. Therefore, while the oscillations do not cause gradient descent to diverge, the implicit curvature regularization they induce does slow down optimization.

Having introduced the central flows methodology with an analysis of gradient descent, we will now use this methodology to analyze two adaptive optimizers.

## 4 Scalar RMSProp

As a stepping stone to the analysis of RMSProp, we now study “Scalar RMSProp,” a simple adaptive optimizer which uses one global adaptive step size, rather than separate adaptive step sizes for each coordinate:(^16^1616Note that we have re-indexed $\nu$ compared to the standard definition of RMSProp (i.e. $\nu_{t+1}\to\nu_{t}$). This does not affect the trajectory and just ensures the effective learning rate at step $t$ is determined by $\nu_{t}$, rather than $\nu_{t+1}$, which simplifies the notation.)(^17^1717This algorithm was also studied by Lyu et al. (). However, their analysis only applies along a manifold of global minima, as $\eta\to 0$.)

$$ $\displaystyle\nu_{t}$ $\displaystyle=\beta_{2}\nu_{t-1}+(1-\beta_{2})\|\nabla L(w_{t})\|^{2},\quad w_{t+1}=w_{t}-\frac{\eta}{\sqrt{\nu_{t}}}\nabla L(w_{t}).\qquad\qquad\qquad$ (13) $$

The algorithm maintains an exponential moving average (EMA), $\nu$, of the squared gradient norm, and takes gradient steps of size $\eta/\sqrt{\nu}$, which we call the *effective step size*.(^18^1818The terms “learning rate” and “step size” are usually interchangeable. In this paper, to avoid ambiguity, we will use the phrase “learning rate” to denote the hyperparameter, and “step size” or “effective step size” to denote the actual step sizes that are taken.)
The EMA hyperparameter $\beta_{2}$ is a knob that interpolates the algorithm between gradient descent when $\beta_{2}=1$ and normalized gradient descent (NGD) when $\beta_{2}=0$.(^19^1919When $\beta_{2}=1$, Scalar RMSProp reduces to gradient descent with learning rate $\eta/\sqrt{\nu_{0}}$. Conversely, when $\beta_{2}=0$, it reduces to NGD with learning rate $\eta$: $w_{t+1}=w_{t}-\eta\cdot\frac{\nabla L(w_{t})}{\|\nabla L(w_{t})\|}$.)

In this section, we will use the central flow framework to understand the behavior of this simple adaptive optimizer. We first describe the dynamics of Scalar RMSProp in [Section 4.1](#S4.SS1). We then leverage this to derive a central flow in [Section 4.2](#S4.SS2).
Finally, in [Section 4.3](#S4.SS3), we interpret this flow to understand the optimizer’s behavior. In particular:

- •
In [Section 4.3.1](#S4.SS3.SSS1), we make precise how Scalar RMSProp adapts its step size to the local loss landscape. Specifically, we show that oscillatory dynamics implicitly set the effective step size to the value $2/S(w)$, where $S(w)$ is the current sharpness; this value is the *largest stable step size* at the current weights $w$.
Thus, our analysis reveals how oscillations can enable a *gradient*-based optimizer to adapt to the local *curvature*.
- •
In [Section 4.3.2](#S4.SS3.SSS2), we show that step size adaptation is not the full story: Scalar RMSProp also implicitly regularizes curvature throughout training, and in fact, at EOS, the hyperparameters $\eta,\beta_{2}$ only affect the time-averaged trajectory by modulating the strength of this curvature regularization.
- •
Bringing it all together, in [Section 4.3.3](#S4.SS3.SSS3) we describe how the interplay between step size adaptation and curvature regularization gives rise to a mechanism we call *acceleration via regularization*, whereby the optimizer implicitly steers itself towards low-curvature regions where it can take larger steps. We show that this mechanism is key to the efficacy of Scalar RMSProp and to the function of its hyperparameters.

These points will generalize to RMSProp in [Section 5](#S5), but are simpler to understand for Scalar RMSProp.

Figure: Figure 7: A typical Scalar RMSProp trajectory. We train a ViT on a subset of CIFAR-10 using Scalar RMSProp with $\eta=0.02$. While the top eigenvalues of the Hessian $H(w)$ evolve freely (right), the top eigenvalue of the *effective* Hessian $H(w)/\sqrt{\nu}$ equilibrates at the critical threshold $2/\eta$ (center).
Refer to caption: /html/2410.24206/assets/x7.png

### 4.1 The Dynamics of Scalar RMSProp

The dynamics of Scalar RMSProp revolve around the *effective sharpness*, defined as $S^{\text{eff}}:=S(w)/\sqrt{\nu}$.
First, the effective sharpness controls the oscillations: when $S^{\text{eff}}>2/\eta$, Scalar RMSProp oscillates with growing magnitude 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 2](#S3.E2), thereby decreasing the effective sharpness via its *numerator*.
The other mechanism, new to Scalar RMSProp, is that oscillations increase the gradient norm and hence $\nu$, 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/\eta$, as depicted in the bottom left plot in [Figure 8](#S4.F8).
The fine-grained dynamics are complex and challenging to analyze, even in the case of a single oscillatory direction.
Fortunately, we will see in the next section that analyzing the *time-averaged* dynamics is much simpler.

### 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., , cf.):(^20^2020 The $1-\beta_{2}\to\frac{1-\beta_{2}}{\beta_{2}}$ correction is necessary for small values of $\beta_{2}$. For example, when $\beta_{2}=0$ (i.e. NGD), $\nu_{t}=\norm{\nabla L(w_{t})}^{2}$ so in the continuous time ODE, $\nu(t)$ needs to adapt “instantly” to $\norm{\nabla L(w(t))}^{2}$. See [Lemma 4](#Thmlemma4) for additional justification for this correction term.)

$$ $\displaystyle\frac{dw}{dt}=-\frac{\eta}{\sqrt{\nu}}\nabla L(w),\quad\quad\frac{d\nu}{dt}=\frac{1-\beta_{2}}{\beta_{2}}\quantity\big[\norm{\nabla L(w)}^{2}-\nu].$ (14) $$

However, at the edge of stability, the trajectory of Scalar RMSProp deviates from [eq. 14](#S4.E14). 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 (and remains at) the edge of stability. See [Section C.4](#A3.SS4) for our full derivation which accounts for multiple eigenvalues at EOS and for eigenvalues entering and leaving EOS.

In [section 3.2.1](#S3.SS2.SSS1), we derived an approximation for the time-averaged gradient, $\operatorname{\mathbb{E}}[\nabla L(w)]$. Using the first two terms of [eq. 2](#S3.E2), we can also derive a time-averaged approximation for the squared gradient norm $\operatorname{\mathbb{E}}[\norm{\nabla L(w)}^{2}]$:

$$ $\displaystyle\operatorname{\mathbb{E}}[\norm{\nabla L(w)}^{2}]\approx\norm{\nabla L(\overline{w})}^{2}+\cancel{2\expectationvalue{\nabla L(\overline{w}),u}S(\overline{w})\operatorname{\mathbb{E}}[x]}+S(\overline{w})^{2}\operatorname{\mathbb{E}}[x^{2}]$ $$

where we again used $\operatorname{\mathbb{E}}[x]=0$ to ignore the middle term. Based on these time averages, we make the ansatz that the joint dynamics of $(w_{t},\nu_{t})$ follow a central flow $(w(t),\nu(t))$ of the form:

$$ $\displaystyle\frac{dw}{dt}=-\frac{\eta}{\sqrt{\nu}}\quantity\Big[\underbrace{\nabla L(w)+\tfrac{1}{2}\sigma^{2}(t)\nabla S(w)}_{\operatorname{\mathbb{E}}[\nabla L(w_{t})]}],\quad\quad\frac{d\nu}{dt}=\frac{1-\beta_{2}}{\beta_{2}}\quantity\Big[\underbrace{\norm{\nabla L(w)}^{2}+S(w)^{2}\sigma^{2}(t)}_{\operatorname{\mathbb{E}}[\norm{\nabla L(w_{t})}^{2}]}-\nu]$ (15) $$

where $\sigma^{2}(t)$ is a still-unknown quantity intended to model $\operatorname{\mathbb{E}}[x_{t}^{2}]$, the instantaneous variance of the oscillations.
As in our analysis of gradient descent, there is a unique value of $\sigma^{2}(t)$ that maintains $S^{\text{eff}}(w,\nu)=2/\eta$. To compute it, we expand $\frac{dS^{\text{eff}}}{dt}$ using the chain rule:
$\frac{dS^{\text{eff}}}{dt}=\langle\frac{\partial S^{\text{eff}}}{\partial w},\frac{dw}{dt}\rangle+\frac{\partial S^{\text{eff}}}{\partial\nu}\cdot\frac{d\nu}{dt}$.
Plugging in $\frac{dw}{dt},\frac{d\nu}{dt}$ from [eq. 15](#S4.E15) shows that $\frac{dS^{\text{eff}}}{dt}$ is linear in $\sigma^{2}$. Thus, there is a unique value of $\sigma^{2}$ that will ensure $\frac{dS^{\text{eff}}}{dt}=0$, which is given by:

$$ $\displaystyle\sigma^{2}(w;\eta,\beta_{2})=\frac{\beta_{2}\overbrace{\expectationvalue{-\nabla L(w),\nabla S(w)}}^{\mathclap{\text{progressive sharpening}}}+(1{-}\beta_{2})\overbrace{\quantity[S(w)^{2}/4-\|\nabla L(w)\|^{2}/\eta^{2}]}^{\mathclap{\text{effect of mean reversion on $\nu$}}}}{\beta_{2}\underbrace{\tfrac{1}{2}\norm{\nabla S(w)}^{2}}_{\mathclap{\text{sharpness reduction}}}+(1{-}\beta_{2})\underbrace{S(w)^{2}/\eta^{2}}_{\mathclap{\text{effect of oscillation on $\nu$}}}}.$ (16) $$

The central flow for Scalar RMSProp with a single unstable eigenvalue is given by [eq. 15](#S4.E15) with this value of $\sigma^{2}$.(^21^2121The Scalar RMSProp central flow can be interpreted as a projected flow in the augmented space $(w,\nu)$ under a certain non-Euclidean norm. However, because this flow is not a gradient flow, it does not immediately suggest a decreasing potential function for Scalar RMSProp.)
The full Scalar RMSProp central flow, derived in [Section C.4](#A3.SS4), is given in [Definition 5](#Thmdefinition5).
In [Figure 8](#S4.F8) and [Appendix D](#A4) we empirically verify that this flow accurately predicts long-term optimization trajectories of neural networks.

Figure: Figure 8: Central flow for Scalar RMSProp. A ViT is trained on a subset of CIFAR-10 using Scalar RMSProp with $\eta=2/100$ and $\beta_{2}=0.995$ (blue). It enters EOS around step 300, and around step 1300 multiple eigenvalues go unstable (dotted red). The central flow (black) accurately models the time-averaged trajectory of Scalar RMSProp even at the edge of stability, whereas the naive stable flow (gray) follows a different path. The effective step size (ESS) plot shows that after reaching EOS, the ESS equilibrates around the *maximal locally stable step size* $2/S$.
Refer to caption: /html/2410.24206/assets/x8.png

The analysis in this section highlights the potential of our time-averaging methodology. With just a single invocation of the chain rule, we have characterized the long-term trajectory of a complex dynamical system involving mutual interactions between the oscillations, the sharpness, and the adaptive learning rate.

### 4.3 Interpreting Scalar RMSProp via its Central Flow

We now interpret the Scalar RMSProp central flow to shed light on the behavior of the algorithm and the function of its hyperparameters $\eta$ and $\beta_{2}$.
Because the dynamics usually transition from stable to EOS quite early in training, we focus on interpreting the central flow in the EOS regime.(^22^2222In the stable regime ($S^{\text{eff}}<2/\eta$), the central flow is given by the stable flow [eq. 14](#S4.E14). For this flow, $\frac{dw}{dt}$ is directly proportional to that of gradient flow, implying these flows traverse the same trajectory, just at a different speed (i.e. with a nonlinear time-rescaling). In this regime, the effective step size generally increases monotonically, so Scalar RMSProp follows gradient flow with a learning rate warmup.) We discuss step size adaptation in [Section 4.3.1](#S4.SS3.SSS1), implicit curvature reduction in [Section 4.3.2](#S4.SS3.SSS2), and finally in [Section 4.3.3](#S4.SS3.SSS3), we show that the interplay between these two mechanisms is key to the success of Scalar RMSProp.

#### 4.3.1 Implicit step size selection

[Uncaptioned image]: /html/2410.24206/assets/x9.png

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 $S^{\text{eff}}:=S(w)/\sqrt{\nu}$ is fixed at $2/\eta$. This condition can be rearranged into a statement about the effective step size: $\eta/\sqrt{\nu}=2/S(w)$. Notably, the value $2/S(w)$ is the *largest stable step size* for the current location $w$ in weight space. In other words, while the algorithm is at EOS, the oscillatory dynamics continually adapt the effective step size to the current largest stable step size. This is the precise sense in which Scalar RMSProp “adapts” to the local loss landscape.

#### 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 a joint flow over $(w,\nu)$. However, at EOS, because $\eta/\sqrt{\nu}=2/S(w)$, we can eliminate $\nu$ from the expression for $\frac{dw}{dt}$, and write the central flow in terms of $w$ alone:(^23^2323Note that at EOS we can rearrange the EOS condition as $\nu=\eta^{2}S(w)^{2}/4$, which lets us write $\nu$ as a function of $w$ and eliminate $\nu$ everywhere in the central flow. This was already used to derive the expression for $\sigma^{2}$ in [eq. 16](#S4.E16).)

$$ $\displaystyle\frac{dw}{dt}=-\underbrace{\frac{2}{S(w)}}_{\mathclap{\text{effective step size}}}\bigl{[}\nabla L(w)+\underbrace{\frac{1}{2}\sigma^{2}(w;\eta,\beta_{2})\nabla S(w)}_{\mathclap{\text{implicit sharpness penalty}}}\bigr{]}$ (17) $$

where $\sigma^{2}(w;\eta,\beta_{2})$ is given by [eq. 16](#S4.E16). In other words, the time-averaged trajectory of Scalar RMSProp at EOS is essentially equivalent to that of the following simpler-to-understand 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. 16](#S4.E16). This characterization, which directly exposes the optimizer’s implicit behaviors, is far from obvious from the raw update rule (eq. [13](#S4.E13)).

Interestingly, the hyperparameters $\eta,\beta_{2}$ are not used to determine the effective step size $2/S(w)$. Instead, their only role is to modulate $\sigma^{2}$, which controls the strength of the implicit sharpness penalty. The effect of the learning rate hyperparameter $\eta$ is to *monotonically increase* $\sigma^{2}$ — indeed, the numerator of [eq. 16](#S4.E16) is increasing in $\eta$ while the denominator is decreasing in $\eta$, which implies the overall expression for $\sigma^{2}$ is increasing in $\eta$. The simplest case is that of NGD, i.e. when $\beta_{2}=0$, for which [eq. 16](#S4.E16) reduces to $\sigma^{2}\approx\frac{\eta^{2}}{4}$. Meanwhile, the effect of the hyperparameter $\beta_{2}$ is to monotonically interpolate $\sigma^{2}$ between that of NGD when $\beta_{2}=0$ and that of gradient descent when $\beta_{2}=1$.(^24^2424We note that which of these is larger is situation dependent, so $\sigma^{2}$ can be either monotonically increasing or monotonically decreasing in $\beta_{2}$. That said, because when $\beta_{2}=0$, $\sigma^{2}(w;\eta,0)\approx\eta^{2}/4$ and when $\beta_{2}=1$, $\sigma^{2}(w;\eta,1)$ is independent of $\eta$, a general rule is that for small learning rates, $\sigma^{2}$ is monotonically increasing in $\beta_{2}$, while for large learning rates, $\sigma^{2}$ is monotonically decreasing in $\beta_{2}$.) The interpretations of $\eta,\beta_{2}$ generalize to the setting of multiple oscillating directions, as detailed in [Lemma 3](#Thmlemma3).

Figure: Figure 9: “Acceleration via regularization” for RMSProp-Norm. Starting from the same initialization, we run the RMSProp-Norm central flow at two different learning rates (in blue and orange), as well as an ablated flow $\frac{dw}{dt}=-\frac{2}{S(w)}\nabla L(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 inset), the flows with higher curvature regularization optimize slower; however, over the longer run, they take larger steps and optimize faster (CNN / CIFAR-10 / MSE).
Refer to caption: /html/2410.24206/assets/x10.png

#### 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 $S(w)$, 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 effectiveness.
In [Figures 9](#S4.F9) and [15](#A1.F15), we compare the Scalar RMSProp central flow to an ablated version which adapts the step size to $2/S(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 $\eta$.
We have seen that at EOS, the only direct effect of $\eta$ on the central flow is to modulate the strength of sharpness regularization, with higher $\eta$ inducing stronger sharpness regularization.
Thus, counterintuitively, the *instantaneous* effect of a higher $\eta$ is often to *slow down* optimization.
However, as we illustrate in [Figures 9](#S4.F9) and [15](#A1.F15), over longer timescales, higher $\eta$ steers the trajectory into lower-sharpness regions, in which Scalar RMSProp’s effective step size will be larger, thereby tending to *speed up* optimization. Thus, as one would expect of a learning rate hyperparameter, larger $\eta$ can accelerate optimization; however they do so through this *indirect* mechanism.

## 5 RMSProp

We now study RMSProp (Tieleman and Hinton, ), which is equivalent to Adam (Kingma and Ba, ) without momentum. RMSProp maintains an EMA $\nu$ of the elementwise squared gradients $\nabla L(w)^{\odot 2}$, and uses *per-coordinate* effective step sizes of $\eta/\sqrt{\nu}$:(^25^2525Our analysis can accommodate both bias correction and an $\epsilon$-dampening (dividing by $\sqrt{\nu+\epsilon}$ rather than $\sqrt{\nu}$) which are used by Adam (see [Section C.7](#A3.SS7)). However, to simplify exposition, the main text focuses on this simpler version of RMSProp.)

$$ $\displaystyle\nu_{t}=\beta_{2}\nu_{t-1}+(1-\beta_{2})\nabla L(w_{t})^{\odot 2},\quad w_{t+1}=w_{t}-\frac{\eta}{\sqrt{\nu_{t}}}\odot\nabla L(w_{t}),\qquad\qquad$ (18) $$

where $\odot$ represents the entrywise product.
RMSProp can also be viewed as preconditioned gradient descent $w_{t+1}=w_{t}-P_{t}^{-1}\nabla L(w_{t})$ with the dynamic preconditioner $P_{t}:=\mathrm{diag}(\sqrt{\nu_{t}}/\eta)$.(^26^2626Folding $\eta$ into the preconditioner is unconventional, but will make the analysis clearer.)
While Adam employs the same dynamic preconditioner and has achieved widespread success in deep learning, it has remained unclear why this specific preconditioning strategy is so effective (Kunstner et al., ; Orabona, ; Martens, ).

Figure: Figure 10: Central Flow for RMSProp. A ViT is trained on a subset of CIFAR-10 using RMSProp with $\eta=4{\times}10^{-5}$ and $\beta_{2}=0.995$ (blue). It enters EOS at step 230 and multiple eigenvalues become unstable at step 360 (dotted red). The central flow (black) accurately models the time-averaged trajectory of RMSProp even at EOS, whereas the naive stable flow (gray) follows a different path.
Refer to caption: /html/2410.24206/assets/x11.png

In this section, we use the central flows framework to understand the behavior of RMSProp.
We will show that the success of RMSProp is intrinsically tied to its oscillatory dynamics. Our analysis of these dynamics will allow us to understand basic aspects of RMSProp’s behavior for the first time.

We start by describing the dynamics of RMSProp in [Section 5.1](#S5.SS1).
We then derive a central flow in [Section 5.2](#S5.SS2). Finally, in [Section 5.3](#S5.SS3), we interpret this flow to understand the optimizer’s behavior. In particular:

- •
In [Section 5.3.1](#S5.SS3.SSS1), we show that RMSProp’s preconditioner is *implicitly* determined by the algorithm’s oscillatory dynamics, and we make this preconditioner *explicit* for the first time. Specifically, we show that RMSProp computes its preconditioner by solving a convex program (eq. [19](#S5.E19)) involving the Hessian. This reveals that RMSProp is implicitly a second-order optimizer, despite only accessing the loss through first-order gradients.
- •
In [Section 5.3.2](#S5.SS3.SSS2), we show that, like Scalar RMSProp, the success of RMSProp relies not only on this preconditioning strategy, but also on an *acceleration via regularization* mechanism whereby implicitly regularizing curvature allows the optimizer to take larger steps later in training.

### 5.1 The Dynamics of RMSProp

The dynamics of RMSProp revolve around the *effective sharpness* $S^{\text{eff}}(w;\nu)$, defined as the largest eigenvalue of the preconditioned Hessian $P^{-1}H(w)$ where $P=\operatorname{diag}(\sqrt{\nu}/\eta)$.(^27^2727Note that $P^{-1}H(w)$ is similar to $P^{-1/2}H(w)P^{-1/2}$ which is symmetric and therefore diagonalizable with real eigenvalues. Therefore $P^{-1}H(w)$ also has real eigenvalues.)
When $S^{\text{eff}}>2$, the iterates oscillate along the top right eigenvector of the preconditioned Hessian.
In turn, such oscillations reduce $S^{\text{eff}}$, via a combination of two mechanisms: (1) they implicitly reduce curvature, and (2) they increase the gradient, growing $\nu$ and therefore $P$.
The net effect is that the effective sharpness $S^{\text{eff}}$ stays regulated around $2$ throughout training (Cohen et al., ).(^28^2828Note that in this section we have absorbed the learning rate $\eta$ into the definition of $P$ so the critical threshold is $2$ rather than $2/\eta$. This simplifies the analysis.)

### 5.2 Deriving the RMSProp Central Flow

In [Section C.5](#A3.SS5) we derive a central flow $(w(t),\nu(t))$, in the same way as above, which models the time-averaged trajectory of RMSProp.
We verify its accuracy in [Figure 10](#S5.F10) and [Appendix D](#A4).

### 5.3 Interpreting RMSProp via its Central Flow

We now interpret the RMSProp central flow to understand the behavior of RMSProp.
Because the dynamics usually transition from stable to EOS early in training, we focus on the EOS regime.

Figure: Figure 11: Understanding RMSProp using its central flow. Left: the EMA $\nu(t)$ converges to its stationary value $\bar{\nu}({w(t)})$ early in training. While running the RMSProp central flow, we plot the cosine similarity between the EMA $\nu(t)$ and its stationary value $\bar{\nu}(w(t))$. This cosine similarity approaches $1$ during training, indicating that $\nu$ approaches stationary. Thus, we can reason about RMSProp using the stationary preconditioner $\bar{P}(w)=\eta/\sqrt{\bar{\nu}(w)}$. Middle: this preconditioner is not optimal. While running the central flow, we compare the unregularized optimization speed $\|\nabla L(w(t))\|_{P^{-1}}^{2}$ of the stationary preconditioner (black) to that of an alternative preconditioner (red) which leaves out the second term in the objective [eq. 19](#S5.E19). The latter quantity is larger, indicating that the alternative preconditioner would optimize faster. Right: stationary flow accurately predicts optimization speed. The stationary flow [eq. 20](#S5.E20), which incorporates implicit curvature regularization, predicts (black) the rate of loss decrease (blue) more accurately than a naive estimate (red) which uses the stationary preconditioner but does not incorporate curvature regularization.
Refer to caption: /html/2410.24206/assets/x12.png

#### 5.3.1 The stationary preconditioner

The central flow for RMSProp is harder to interpret than that for Scalar RMSProp, because even at EOS, $\nu$ cannot be expressed as a closed-form function of $w$, and instead remains an independent variable. This reflects the fact that for any $w$, there are potentially many values for $\nu$ that could stabilize optimization, and the actual value used by RMSProp depends on the historical trajectory.
Nevertheless, it turns out that, in some circumstances, $\nu$ implicitly converges under the dynamics of RMSProp to a value that depends on the current $w$ alone.
In particular, imagine holding the weights fixed at some value $w$, and letting $\nu$ “catch up.”
What value would $\nu$ converge to?
We show in [Proposition 3](#Thmproposition3) that for any $w$, there is a unique $\nu$ that satisfies the stationarity condition $\frac{d\nu}{dt}=0$.
We call this unique $\nu$ the *stationary $\nu$* for the weights $w$, denoted as $\overline{\nu}(w)$.
Empirically, we observe that $\nu(t)$ usually converges to the stationary value $\bar{\nu}(w(t))$ during training ([Figure 16](#A1.F16)).

In [Proposition 2](#Thmproposition2), we show that the corresponding *stationary preconditioner* $\overline{P}(w)=\operatorname{diag}(\sqrt{\overline{\nu}(w)}/\eta)$ is, remarkably, the optimal solution to a convex optimization problem over preconditioners:

$$ $\displaystyle\overline{P}(w)\quad:=\quad\operatorname*{arg\,min}_{\mathclap{\text{$P$ diagonal, $P{\succeq}0$}}}\qquad\tr(P)+\tfrac{1}{\eta^{2}}\underbrace{\norm{\nabla L(w)}^{2}_{P^{-1}}}_{\mathclap{\text{optimization speed}}}\mbox{\quad such that\quad}\underbrace{H(w)\preceq 2P}_{\text{local stability}}.\quad$ (19) $$

That is, RMSProp implicitly solves the convex program [eq. 19](#S5.E19) to compute its preconditioner. This is the precise sense in which RMSProp “adapts” its preconditioner to the local loss landscape.

We can now understand RMSProp’s preconditioning strategy by interpreting the optimization problem [eq. 19](#S5.E19).
The constraint $H(w)\preceq 2P$ is equivalent to $S^{\text{eff}}\leq 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 preconditioning strategy could be simply summarized as maximizing the *harmonic mean* of the effective step sizes while maintaining local stability — a sensible preconditioning strategy.(^29^2929Interestingly, this SDP is the dual to the max-cut SDP: $\max_{\Sigma\succeq 0}\expectationvalue{\Sigma,H}$ such that $\Sigma_{ii}=1$ for all $i$. Thus, this preconditioning strategy could be described as solving the max-cut SDP with the Hessian as the weight matrix, and using the resulting dual variable as its preconditioner.)

However, matters are complicated by the presence of the second term in the [eq. 19](#S5.E19) objective.
The quantity $\norm{\nabla L(w)}^{2}_{P^{-1}}$ is the instantaneous rate of loss decrease under preconditioned gradient flow with preconditioner $P$. Minimizing this term actually *slows down* optimization ([Figure 17](#A1.F17)).
Therefore, the presence of this term implies that RMSProp’s preconditioning strategy is not optimal.

Figure: Figure 12: Acceleration via regularization for RMSProp. We compare the RMSProp stationary flow (blue) to an ablated flow (orange) which uses the same preconditioning strategy but leaves out the implicit curvature regularization. Initially, the stationary flow optimizes slower (left, initial), due to the implicit curvature regularization. But over time, it navigates to lower-curvature regions (right), where it takes larger steps (middle), and optimizes faster (left, final).
Refer to caption: /html/2410.24206/assets/x13.png

#### 5.3.2 Implicit curvature reduction & acceleration via regularization

Understanding the preconditioning strategy implicitly employed by 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 implicit curvature reduction induced by oscillations.
Substituting $\bar{P}$ into the RMSProp central flow, we can obtain a *stationary flow* over $w$ alone, which assumes that the preconditioner $P$ is always fixed at its stationary value $\overline{P}$ (eq. [19](#S5.E19)):

$$ $\displaystyle\frac{dw}{dt}=\underbrace{\overline{P}(w)^{-1}}_{\mathclap{\begin{subarray}{c}\text{stationary}\\ \text{ preconditioner}\end{subarray}}}\quantity\Big[\;\nabla L(w)+\underbrace{\tfrac{1}{2}\nabla_{w}\expectationvalue{\Sigma,H(w)}}_{\text{implicit curvature penalty}}\;].$ (20) $$

where $\Sigma=\Sigma(w;\eta;\beta_{2})$ is defined as the solution to a cone complementarity problem (eq. [32](#A3.E32)). This characterization suggests that the time-averaged trajectory of RMSProp is essentially equivalent to that of the following simpler-to-understand algorithm: at each iteration, compute the preconditioner $\bar{P}(w)$ using [eq. 19](#S5.E19) and then take a preconditioned gradient step on a curvature-penalized objective.
While our experiments indicate that the stationary flow [eq. 20](#S5.E20) is less accurate than the central flow, it can still accurately predict many behaviors of RMSProp. For example, [Figure 18](#A1.F18) shows that the stationary flow [eq. 20](#S5.E20) can accurately predict the instantaneous optimization speed of the central flow, even though it is only given access to $w(t)$ and not $\nu(t)$.

Because the stationary flow makes RMSProp’s implicit behaviors *explicit*, it allows us to reason about these behaviors directly, and disentangle them from one another.
In [Figure 19](#A1.F19), we compare the stationary flow against an ablated flow which uses the same preconditioning strategy, but leaves out the curvature penalty. We find that in the long run, this ablated flow navigates into sharper regions in which it takes smaller steps, and optimizes slower.
This implies that RMSProp’s implicit curvature reduction effect is crucial for its efficacy as an optimizer.

## 6 Experimental Results

We have used informal mathematical reasoning to derive central flows for three optimizers. In this section, we empirically validate the accuracy of these flows.
To assess how well a central flow approximates its corresponding optimizer, we run both the flow and the optimizer simultaneously, starting from the same initialization.
For a baseline, we also run the corresponding *stable flow* (i.e. for gradient descent, the gradient flow), which does not account for the effect of the oscillations and hence should fail to properly model the trajectory in the EOS regime.(^30^3030We run the stable flow until its effective sharpness reaches $20/\eta$, at which point it becomes computationally intractable to discretize.)

Our experiments span a variety of neural network architectures: a CNN, a ResNet (He et al., ), a vision transformer (ViT) (Dosovitskiy et al., ), a GPT-2 (Radford et al., ) style transformer, and a recurrent neural network (RNN).
We evaluate the CNN, ResNet, and ViT on an image classification task, and we evaluate the GPT-style transformer and RNN on a sequence prediction task, using both mean-squared-error (MSE) and cross entropy losses.
Since discretizing the flows is computationally expensive, we are restricted to small scale datasets: our image dataset is a subset of CIFAR-10 with 1,000 examples and either 4 or 10 classes, while our sequence dataset is a sorting task with 1,000 sequences of length 8 and alphabet size 4. Even at these modest scales, our full set of 115 experiments required 2000 GPU-hours.

Our full set of experiments, which can be found in [Appendix D](#A4), demonstrate that the central flow accurately approximates the true optimization trajectory in a variety of deep learning settings.
In many experiments, the weight-space distance between the central flow and the real optimizer trajectory stays at essentially zero throughout training (e.g. [Figure 5](#S3.F5)).
Even in cases where this distance grows non-negligibly during training, the central flow’s predictions for derived quantities such as the train loss, or individual network outputs, often still closely match those of the real optimizer trajectory (e.g. [Footnote 33](#footnote33)).
We are unaware of any other theoretical framework for reasoning about optimization in deep learning which has a comparable level of predictive power.

We now discuss some factors which affect the accuracy of the central flow approximation.
First, the quality of the central flow approximation often degrades as the learning rate grows increasingly large (see [Footnote 33](#footnote33)).(^31^3131It is unclear whether there exists a correction to the central flow that could improve its accuracy in these settings, or whether the underlying trajectory is simply too chaotic to be captured by any flow.)
Second, the central flow approximation is usually more accurate for MSE loss than for cross-entropy loss. The issue may be that higher-order terms cause the sharpness equilibrium point to differ slightly from $2/\eta$ for cross-entropy loss, as previously noted by Damian et al. () in their Appendix F. An interesting direction for future work is to design corrections to the central flow that improve its predictions in this situation. Third, the regularity of the loss landscape also seems to play a role (Wang et al., ); in [Figure 14](#A1.F14), we consider a family of activations which interpolate between GeLU (smooth) and ReLU (not smooth).
We observe that the weight-space discrepancy between the flow and the optimizer trajectory grows larger as the activation function becomes less smooth.

Figure: Figure 13: 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).(^33^3333While the rate of error accumulation in this figure can be partially explained by the fact that larger learning rates travel farther through weight space over a fixed number of steps $t$, this effect persists even after rescaling time and plotting $\eta\times t$, rather than $t$, on the $x$ axis.)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).
Refer to caption: /html/2410.24206/assets/x14.png

## 7 Discussion

### 7.1 Modeling decisions

##### Deterministic setting

Our analysis is restricted to the simple setting of deterministic (i.e. full-batch) training, whereas practical deep learning generally involves minibatch training. We study the full-batch setting because, as the simplest special case, understanding full-batch training is a necessary prerequisite for understanding minibatch training. However, we also believe that understanding full-batch training might suffice for some practical purposes, such as designing optimizers. For example, Kunstner et al. () showed that the advantage of adaptive methods over SGD grows larger with larger batch sizes, suggesting that the relevant algorithmic principles can be best understood in the deterministic setting.

An interesting direction for future research is to try to extend our central flows methodology to the stochastic setting.
Like deterministic optimizers, stochastic optimizers are known to implicitly regularize the curvature along their trajectories, and in fact this effect is *stronger* in the stochastic setting (Keskar et al., ; Jastrzębski et al., , ).
However, extending the central flows methodology to the stochastic setting may be nontrivial; due to the randomness, it is not clear whether there exists a deterministic differential equation around which SGD oscillates.
An interesting question is whether there exists a differential equation that can predict derived metrics such as network predictions or training loss curves, even if it cannot model the weight-space trajectory of SGD.

Finally, while our analysis in this paper sheds light on adaptive optimizers in the deterministic setting, it remains to be seen if these optimizers exhibit similar preconditioning strategies in the stochastic setting.

##### Black-box model of the loss

Our analysis treats the loss function as a black box, and never uses that the optimization problem at hand involves training a neural network.
The advantage of this approach is its generality: we expect our analysis to apply to generic deep learning architectures and learning problems, including those that do not yet exist. The disadvantage, however, is that the predictions made by our theory are at the abstraction level of the *loss landscape*, and would need to be further translated in order to make concrete claims about the network architecture or learning problem. For example, our theory tells us that the learning rate hyperparameter modulates the strength of an implicit sharpness penalty, but does not tell us how this sharpness penalty affects learning. Nor does our theory shed light on how different layers of the neural network are mechanistically implicated in progressive sharpening or sharpness reduction.

On the one hand, the loss landscape level of abstraction is in some sense “natural” — the overall path that optimizers follow really does intrinsically depend on the (effective) sharpness. But on the other hand, understanding many important aspects of optimization in deep learning will likely require cracking open the black box a bit more.

### 7.2 Takeaways from our analysis

##### The unreasonable effectiveness of time-averaging

As demonstrated by prior works on EOS, it is challenging to analyze the oscillatory EOS dynamics in fine-grained detail. Our work shows that, perhaps surprisingly, simple heuristics allow us to analyze the *time-averaged* dynamics with excellent numerical accuracy. Interestingly, the success of this time-averaging approach seems to imply that the oscillations only affect the macroscopic trajectory in an ergodic sense, i.e. via their *covariance* rather than via their fine-grained details. An promising direction for future work is to identify realistic conditions under which our heuristic time-averaging arguments can be made rigorous.

##### Necessity of third-order Taylor expansions

While optimization theory generally relies on second-order Taylor expansions of the loss, Damian et al. () showed that a *third-order* Taylor expansion is key for understanding the convergence of gradient descent; such a Taylor expansion reveals that oscillations implicitly trigger curvature reduction, a form of negative feedback which stabilizes optimization.
In this work, we have shown that a third-order Taylor expansion is similarly necessary for understanding the *acceleration via mechanism* which underlies the success of adaptive optimizers.
Thus, our work further drives home the necessity of a third-order Taylor expansion when analyzing optimization in deep learning.

##### Oscillatory first-order methods are implicitly second-order methods

A key surprise over the last decade is that there are no second-order optimizers(^34^3434We use “second-order” to refer to optimizers that explicitly use Hessian-vector products to adapt to second-order information. We therefore regard RMSProp, Adam, and Shampoo as first-order optimizers, as they only access the loss through gradients.) which consistently outperform first-order adaptive optimizers. Our work partially demystifies this observation. We show that when first-order optimizers oscillate, they implicitly pick up second order information. Thus, even though RMSProp is a first-order optimizer, it implicitly employs a sophisticated second-order preconditioning strategy, detailed in [Section 5.3.1](#S5.SS3.SSS1).
Further, this preconditioning strategy is efficient, requiring no more gradient queries than gradient descent does.
An exciting direction for future work is to *intentionally* design first-order adaptive methods with such implicit preconditioners in mind.

##### Adapting to curvature is not enough

Traditional optimization theory views the curvature of the loss as a pre-existing feature of the optimization problem, and views the job of an optimizer as *adapting* to this pre-existing curvature. We have shown that RMSProp does not merely passively adapt to the curvature; it also actively *shapes* the curvature along its trajectory, by steering away from high-curvature regions where it would need to take small steps.
Further, we have shown that this effect is crucial for the optimization efficacy of RMSProp: a variant that only adapts to curvature (but does not shape it) ends up optimizing slowly. Thus, our work suggests that *acceleration via regularization* is a vital design principle for adaptive optimizers.

## 8 Conclusion

In this paper, we have developed a framework for analyzing deep learning optimization algorithms.
To analyze an algorithm, we derive a *central flow* which directly models the time-averaged trajectory of the oscillatory optimizer.
We have empirically demonstrated 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 dynamics), and they are empirically consistent across different deep learning settings, which suggests that 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.

## Appendix A Additional Figures

Figure: Figure 14: Accuracy of central flow degrades as activation function becomes less smooth. We consider networks with the $\beta$-GeLU activation function from Dauphin et al. (), defined as $x\mapsto\Phi(\beta x)$ where $\Phi$ is the standard Gaussian CDF. This activation interpolates between GeLU when $\beta=1$ and ReLU when $\beta=\infty$. Subfigure (a) plots this activation function with varying $\beta$. Subfigure (b) shows that when $\beta$ is larger (i.e. when the activation is less smooth), the approximation error between the central flow $w(t)$ and the optimizer trajectory $w_{t}$ grows faster. Subfigure (c) plots the loss curve, and the network’s output on a test example, for both the optimizer trajectory and the central flow. Fortunately, even when $\beta=20$, at which point $\beta$-GeLU is a very close approximation to ReLU, the central flow accurately predicts the overall training loss curve.
Refer to caption: /html/2410.24206/assets/x15.png

Figure: Figure 15: “Acceleration via regularization” for RMSProp-Norm. Starting from the same initialization, we run the RMSProp-Norm central flow at various learning rates, as well as an ablated flow $\frac{dw}{dt}=-\tfrac{2}{S(w)}\nabla L(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 are able to take larger steps and optimize faster.
Refer to caption: /html/2410.24206/assets/x18.png

Figure: Figure 16: The EMA $\boldsymbol{\nu}$ reaches stationarity during training. While running the RMSProp central flow, we monitor the cosine similarity between $\nu(t)$, the real EMA, and $\bar{\nu}(w(t))$, the stationary EMA. This cosine similarity rises to high values during training, implying that $\nu(t)$ reaches stationarity. Thus, we can reason about the stationary $\nu$ (and in particular, the corresponding stationary preconditioner [eq. 19](#S5.E19)) in order to reason about RMSProp’s preconditioning strategy. Note that we compute $\bar{\nu}(w)$ using the Burer-Monteiro factorization, as described in [appendix B](#A2).
Refer to caption: /html/2410.24206/assets/x22.png

Figure: Figure 17: RMSProp’s implicit preconditioner is suboptimal. We compare RMSProp’s stationary preconditioner, defined as the solution to the optimization problem [eq. 19](#S5.E19), to an alternative preconditioner defined as the solution to an analogous optimization problem without the second term in the objective. We assess the efficacy of each preconditioner $P$ by computing $\|\nabla L(w(t))\|^{2}_{P^{-1}}$, the instantaneous rate of loss decrease under the preconditioned gradient flow with preconditioner $P$. We observe that the rate of loss decrease for the alternative preconditioner (in green) is higher than the rate of loss decrease for the RMSProp stationary preconditioner (in orange), indicating that the alternative preconditioner would be better. Both preconditioners are computed using the Burer-Monteiro factorization, as described in [Appendix B](#A2).
Refer to caption: /html/2410.24206/assets/x23.png

Figure: Figure 18: Stationary flow predicts the rate of loss decrease. The stationary flow [eq. 20](#S5.E20), which incorporates implicit curvature regularization, predicts (green) the rate of loss decrease (blue) more accurately than a naive estimate (orange) which uses the stationary preconditioner but does not incorporate curvature regularization. Note that all the estimates are off early in training, as the preconditioner has not yet reached stationarity (see [Figure 16](#A1.F16)).
Refer to caption: /html/2410.24206/assets/x24.png

Figure: Figure 19: “Acceleration via regularization” for RMSProp stationary flow. We compare the RMSProp stationary flow (blue) to an ablated version (orange) which uses the same preconditioning strategy but leaves out the implicit curvature regularization. Initially, the stationary flow optimizes slower (left, initial), 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 a 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.
Refer to caption: /html/2410.24206/assets/x25.png

## Appendix B Additional Experimental Details

##### Architectures

Our CNN is comprised of 4 layers with GeLU activations and uses average pooling before a linear readout layer. We use a modified ResNet without BatchNorm which has 20 layers and uses GeLU activations. Our vision transformer (ViT, (Dosovitskiy et al., )) is based on the default PyTorch (Paszke et al., ) vision transformer and uses the GeLU activation as well. Our recurrent neural network (RNN) has a single layer and uses $\tanh$ activations. Our GPT2-style (Radford et al., ) transformer is based on the “nano” version of minGPT ([github.com/karpathy/minGPT](github.com/karpathy/minGPT)).

##### Discretizing the flows

We discretize each flow using Euler steps. We found that Euler steps of size $1/4$ sufficed to approximate the central flows, we use Euler steps of size $1/10$ to approximate the stable flows. For each trajectory: the true trajectory, the central flow, and the stable flow, we recompute the top eigenvectors of the Hessian from scratch every 5 steps. Eigenvalues and eigenvectors are computed using a sparse eigenvalue solver and hessian-vector products. For the intermediate steps, we reuse the eigenvectors and approximate the eigenvalues by the quadratic forms $u^{T}\nabla^{2}L(w)u$. We compute the gradients of the eigenvalues by computing the contraction of the third derivative $\nabla^{3}L(w)$ with the eigenvectors, i.e. $\nabla^{3}L(w)[u,u]$ which can be implemented in standard PyTorch.

##### Computational Resources

In total, our 115 experiments required about $2000$ GPU-hours on a mix of $A6000$ and $A100$ GPUs.

##### Computing the RMSProp Stationary Preconditioner

To verify that the RMSProp preconditioner indeed solves the convex program [eq. 19](#S5.E19), we explicitly solve this convex program throughout the training trajectory and compare it with the true RMSProp preconditioner. To solve [eq. 19](#S5.E19), we apply a fixed point iteration based on the fixed point equations for $\Sigma,P$. To derive our fixed point iteration, we factorize $\Sigma=DD^{T}$ where $D\in\mathbb{R}^{d\times r}$ (as in the Burer-Monteiro factorization (Burer and Monteiro, )) and $r\ll d$ is intended to upper bound the number of unstable eigenvalues. Then the formula for $\frac{d\nu}{dt}$ reduces to:

$$ $\displaystyle\nu=\nabla L(w)^{\odot 2}+\operatorname{diag}[H\Sigma H]=\nabla L(w)^{\odot 2}+(HD)^{\odot 2}\mathbf{1}.$ $$

In addition, we will use that the span of $\Sigma$ lies in the critical subspace, so

$$ $\displaystyle\operatorname{diag}[\nu^{-1/2}]HD=\frac{\eta}{2}D.$ $$

We begin with a random initial guess for $D$ and iteratively update $D,\nu$ by:

$$ $\displaystyle\nu$ $\displaystyle\leftarrow\nabla L(w)^{\odot 2}+(HD)^{\odot 2}\mathbf{1}$ $\displaystyle D$ $\displaystyle\leftarrow\frac{\eta}{2}\operatorname{diag}[\nu^{-1/2}]HD.$ $$

If this fixed point iteration converges, we have

$$ $\displaystyle\nu=\nabla L(w)^{\odot 2}+(HD)^{\odot 2}\mathbf{1}\mbox{\quad and\quad}\operatorname{diag}[\nu^{-1/2}]HD=\frac{2}{\eta}HD$ $$

so that $\frac{d\nu}{dt}=0$ and the span of $\Sigma=DD^{T}$ 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. 19](#S5.E19) is:

$$ $\displaystyle\max_{\Sigma\succeq 0}\expectationvalue{\Sigma,H(w)}+2\absolutevalue{\nabla L(w)}\cdot\sqrt{1-2\operatorname{diag}\Sigma}\mbox{\quad such that\quad}\operatorname{diag}\Sigma\preceq 1/2.$ $$

## Appendix C 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.

### C.1 Tensor Notation

For a $k$-tensor $T$ and a $s$-tensor $A$ with $s<k$ we define $T[A]$ to be the contraction of $T$ with $A$ along its last $s$ indices, i.e.

$$ $\displaystyle(T[A])_{i_{1},\ldots,i_{k-s}}=\sum_{j_{1},\ldots,j_{s}}T_{i_{1},\ldots,i_{k-s},j_{1},\ldots,j_{s}}A_{j_{1},\ldots,j_{s}}.$ (21) $$

We also define $\expectationvalue{T,T^{\prime}}$ for two $k$-tensors $T,T^{\prime}$ of the same shape by:

$$ $\displaystyle\expectationvalue{T,T^{\prime}}=\sum_{i_{1},\ldots,i_{k}}T_{i_{1},\ldots,i_{k}}T^{\prime}_{i_{1},\ldots,i_{k}}.$ (22) $$

We will also use $T[u_{1},\ldots,u_{j}]$ to denote $T[u_{1}\otimes\cdots\otimes u_{j}]$.

### C.2 On Local Time Averaging

We intentionally do not specialize to a specific notion of “local time-average”. The only properties of the local time-averaging operator $\operatorname{\mathbb{E}}$ that we use are:

- 1.
linearity, i.e. $\operatorname{\mathbb{E}}[f+g]=\operatorname{\mathbb{E}}[f]+\operatorname{\mathbb{E}}[g]$ and $\operatorname{\mathbb{E}}[cf]=c\operatorname{\mathbb{E}}[f]$ for any constant $c$
- 2.
the local time average of a constant $c$ is itself: $\operatorname{\mathbb{E}}[c]=c$
- 3.
in the EOS regime when the sharpness oscillates around $2/\eta$, the time-average is coarse enough to smooth out these oscillations so that $S(\operatorname{\mathbb{E}}[w_{t}])=2/\eta$

When we empirically verify the accuracy of our central flows, we use Gaussian smoothing with a standard deviation of $50$ steps, i.e. we define

$$ $\displaystyle\operatorname{\mathbb{E}}[f_{t}]:=\frac{\sum_{s}f_{s}c_{t-s}}{\sum_{s}c_{t-s}}\mbox{\quad where\quad}c_{j}:=\exp(\frac{-j^{2}}{2\times 50^{2}}).$ $$

### C.3 Gradient Descent

We begin by proving the Taylor expansion used in [Section 3.1](#S3.SS1):

###### Lemma 1 .

Assume that $\overline{w}$ is such that the top eigenvalue of $H(\overline{w})$ has multiplicity $1$ and let $w=\overline{w}+xu$ where $u$ is the top eigenvector of $H(\overline{w})$. Furthermore, assume that $\sup_{w^{\prime}\in U}\norm{\nabla^{4}L(w^{\prime})}<\infty$ for some open neighborhood $U$ of $\overline{w}$. Then,

$$ $\displaystyle\nabla L(w)=\nabla L(\overline{w})+S(\overline{w})xu+\tfrac{x^{2}}{2}\nabla S(\overline{w})+\mathcal{O}(x^{3}).$ $$

###### Proof.

By Taylor’s theorem,

$$ $\displaystyle\nabla L(w)=\nabla L(\overline{w})+H(\overline{w})xu+\tfrac{x^{2}}{2}\nabla^{3}L(\overline{w})[u,u]+\mathcal{O}(x^{3}).$ $$

Because $u$ is an eigenvector of $H(\overline{w})$ with eigenvalue $S(\overline{w})$, the second term can be simplified to $S(\overline{w})xu$. Finally, by Danskin’s theorem (or equivalently the standard formula for the derivative of an eigenvalue):

$$ $\displaystyle\nabla_{\overline{w}}S(\overline{w})=\nabla_{\overline{w}}\quantity[\max_{\norm{u}=1}u^{T}H(\overline{w})u]=\nabla_{\overline{w}}\quantity[u^{T}H(\overline{w})u]=\nabla^{3}L(\overline{w})[u,u]$ $$

where $u$ is the argmax of the second expression, i.e. the top eigenvector of the Hessian at $\overline{w}$.
∎

We now define some notation that will be necessary for the remainder of the section. First, let $U(w):=\ker[\eta H(w)-2]$ denote the *critical subspace*, i.e. the directions in weight space that are at the edge of stability. Let $\operatorname{Sym}(U(w))\subseteq\mathbb{R}^{d\times d}$ denote the subspace of symmetric matrices whose span is contained in $U(w)$.

###### Lemma 1 .

###### Proof.

#### C.3.1 The Differential Varational Inequality Formulation

As in the main text, we will assume that GD oscillates around the central flow $w(t)$ with covariance $\Sigma(t)\in\operatorname{Sym}(U(w))$ so that it follows:

$$ $\displaystyle\frac{dw}{dt}=-\eta\quantity[\nabla L(w)+\tfrac{1}{2}\nabla^{3}L(w)[\Sigma(t)]].$ (23) $$

To determine $\Sigma$, we impose three conditions for all times $t$:

- •
non-negativity: As a covariance matrix, $\Sigma(t)$ is positive semidefinite, i.e. $\Sigma(t)\succeq 0$
- •
stability: The sharpness remains bounded by $2/\eta$, i.e. $H(w(t))\preceq(2/\eta)I$
- •
complementarity: The span of $\Sigma(t)$ (i.e. the span of the oscillations) is contained within the critical subspace $U(w(t))$. Equivalently, $\Sigma(t)\in\operatorname{Sym}(U(w(t)))$.

We say that $(w(t),\Sigma(t))$ follow the GD central flow if they follow [eq. 23](#A3.E23) along with these conditions:

###### Definition 1 (GD Central Flow, DVI Formulation) .

We say that $\{w(t),\Sigma(t)\}_{t\geq 0}$ follow the GD central flow if for almost all $t$ they satisfy [eq. 23](#A3.E23) along with the conditions: $\Sigma(t)\succeq 0$, $H(w(t))\preceq(2/\eta)I$, and $\Sigma(t)\in\operatorname{Sym}(U(w))$.

[Definition 1](#Thmdefinition1) is an example of a *differential variational inequality* (DVI). It provides conditions that the central flow must satisfy but it is not a-priori clear that such a flow exists or is unique. To prove this, and to turn [Definition 1](#Thmdefinition1) into a form that we can simulate numerically, we will show that for almost all times $t$, $\Sigma(t)$ solves a low dimensional convex *cone complementarity program* (CCP) This CCP will be guaranteed to have a unique solution, which can be easily computed numerically.

###### Definition 1 (GD Central Flow, DVI Formulation) .

#### C.3.2 The Cone Complementarity Problem Formulation

We will start with the DVI formulation, [Definition 1](#Thmdefinition1), and prove two additional conditions on $\Sigma(t)$ which will allow us to solve for $\Sigma(t)$. For simplicity, we will look for a choice of $\Sigma(t)$ is right-continuous so that we can reason about the immediate future using the chain rule. First, we consider the change in the Hessian along the central flow ansatz [eq. 23](#A3.E23), restricted to the critical subspace. We will denote this quantity by $\dot{H}(t)\in\operatorname{Sym}(U(w(t)))$. This generalizes the computation of $\frac{dS}{dt}$ in [Section 3](#S3). To compute $\dot{H}$, we define the linear operator $\mathcal{T}_{w}:\operatorname{Sym}(U(w))\to\mathbb{R}^{d}$ by

$$ $\displaystyle\mathcal{T}_{w}[\Sigma]:=\nabla_{w}\expectationvalue{H(w),\Sigma}\quad\forall\;\Sigma\in\operatorname{Sym}(U(w)).$ $$

Intuitively, $\mathcal{T}_{w}$ will play a role analogous to that of $\nabla S(w)$ in [Section 3](#S3). The operator $\mathcal{T}_{w}:\operatorname{Sym}(U(w))\to\mathbb{R}^{d}$ takes as input a $d{\times}d$ matrix $\Sigma\in\operatorname{Sym}(U(w))$ and returns the gradient of $\expectationvalue{H(w),\Sigma}$, the $\Sigma$-weighted Hessian.
Meanwhile, its transpose $\mathcal{T}^{\top}_{w}:\mathbb{R}^{d}\to\operatorname{Sym}(U(w))$ takes as input a direction, and returns the directional derivative of the Hessian, restricted to the critical subspace. Therefore, $\frac{dw}{dt}=-\eta\quantity[\nabla L(w)+\tfrac{1}{2}\mathcal{T}_{w}[\Sigma]]$, and $\dot{H}$ is given by:

$$ $\displaystyle\dot{H}(t):=\frac{dH}{dt}\evaluated{}_{U(w)}=\mathcal{T}_{w}^{T}\quantity[\frac{dw}{dt}]=\underbrace{\mathcal{T}_{w}^{T}[-\nabla L(w)]}_{=:\alpha(w)}-\underbrace{\tfrac{1}{2}\mathcal{T}_{w}^{T}\mathcal{T}_{w}}_{=:\beta(w)}[\Sigma].$ $$

To simplify notation, we define $\alpha(w),\beta(w)$ by:

$$ $\displaystyle\alpha(w):=\mathcal{T}_{w}^{T}[-\nabla L(w)]\in\operatorname{Sym}(U(w)),\quad\beta(w):=\tfrac{1}{2}\mathcal{T}_{w}^{T}\mathcal{T}_{w}\in\operatorname{Sym}(U(w))\otimes\operatorname{Sym}(U(w)),$ $$

so that $\dot{H}=\alpha-\beta[\Sigma]$. In order to maintain stability, i.e. $H\preceq(2/\eta)I$, we need to ensure $\dot{H}\preceq 0$, i.e. $\alpha\preceq\beta[\Sigma]$. However, this condition only lower bounds $\Sigma$ and does not yet uniquely determine it. To do so, we will need to differentiate the complementarity condition. Recall that $\dot{H}$ measures the change in the Hessian, restricted to the critical subspace. The directions in the kernel of $\dot{H}$ will remain at EOS. However, the curvature in the directions where $\dot{H}\prec 0$ will drop below $2/\eta$, so these directions will leave the critical subspace. In other words, the critical subspace at time $t+\epsilon$ will be contained in $\ker[\dot{H}]$, which is a subspace of the critical subspace. Because we assumed $\Sigma$ is right-continuous, this implies that $\Sigma(t)$ must lie in this subspace, i.e. $\operatorname{span}[\Sigma]\in\ker[\dot{H}]$. Equivalently, because $\Sigma\succeq 0$ and $\dot{H}\preceq 0$, we can rewrite this condition as $\expectationvalue*{\Sigma,\dot{H}}=0$. Together, these constraints:

$$ $\displaystyle\Sigma\succeq 0,\quad\dot{H}\preceq 0,\quad\expectationvalue{\Sigma,\dot{H}}=0,\quad\dot{H}=\alpha-\beta[\Sigma]$ $$

define a *cone complementarity problem* for $\Sigma$ ([Definition 8](#Thmdefinition8)). We show in [Lemma 5](#Thmlemma5) that if $\beta$ is symmetric and has full rank, the solution to this CCP is unique, and we denote it by $\mathrm{CCP}(\alpha,\beta)$. 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 $\{w(t)\}_{t\geq 0}$ follows the GD central flow if for almost all $t\geq 0$, $w$ satisfies [eq. 23](#A3.E23) with $\Sigma\in\mathrm{CCP}(\alpha(w),\beta(w))$.

[Definition 2](#Thmdefinition2) is the simplest version of the central flow to simulate numerically. By picking a basis for $\operatorname{Sym}(U(w))$, which has dimension $\frac{k(k+1)}{2}$, we can materialize the linear operator $\mathcal{T}_{w}$ as a $\frac{k(k+1)}{2}\times d$ dimensional matrix, compute $\alpha(w),\beta(w)$, solve the low-dimensional CCP to compute $\Sigma$, and take a small Euler step on the flow using this $\Sigma$.

In the next section, we will see that this CCP formulation can be reinterpreted as a projected gradient flow, which will help to interpret the behavior of the gradient descent central flow.

###### Definition 2 (GD Central Flow, CCP Formulation) .

#### C.3.3 The Projection Formulation

Let $\mathbb{S}_{\eta}:=\{w~{}:~{}S(w)\leq 2/\eta\}$ 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)\}_{t\geq 0}$ follows the gradient descent central flow if for almost all $t$,

$$ $\displaystyle\frac{dw}{dt}=\mathrm{proj}_{T_{\mathbb{S}_{\eta}}(w)}[-\eta\nabla L(w)]\mbox{\quad where\quad}\mathbb{S}_{\eta}:=\{w~{}:~{}S(w)\leq 2/\eta\},$ (24) $$

where $\mathrm{proj}_{T_{\mathbb{S}_{\eta}}(w)}$ denotes the orthogonal projection onto the *tangent cone* of $\mathbb{S}_{\eta}$ at iterate $w$, i.e. the set of directions that, if followed, would not immediately increase the curvature in the critical subspace: $T_{\mathbb{S}_{\eta}(w)}=\{v\in\mathbb{R}^{d}:\mathcal{T}_{w}^{T}[v]\preceq 0\}$.

The equivalence between [Definition 3](#Thmdefinition3) and the CCP formulation [Definition 2](#Thmdefinition2) ia a consequence of this simple lemma:

###### Lemma 2 .

Let $w\in\mathbb{S}_{\eta}$. Then

$$ $\displaystyle\mathrm{proj}_{T_{\mathbb{S}_{\eta}}(w)}[v]=v-\tfrac{1}{2}\mathcal{T}_{w}[\Sigma]\mbox{\quad where\quad}\Sigma\in CCP(\mathcal{T}_{w}^{T}[v],\tfrac{1}{2}\mathcal{T}_{w}^{T}\mathcal{T}_{w}).$ $$

###### Proof.

Note that the tangent space of $\mathbb{S}_{\eta}$ is given by the set: $\{v~{}:~{}\mathcal{T}_{w}^{T}[v]\preceq 0\}$. Therefore the projection is given by solving the quadratic program:

$$ $\displaystyle\min_{\delta}\|\delta\|^{2}\mbox{\quad such that\quad}\mathcal{T}_{w}^{T}[v+\delta]\preceq 0.$ $$

The KKT conditions imply that there exists a $\Sigma$ such that:

$$ $\displaystyle\delta=-\tfrac{1}{2}\mathcal{T}_{w}[\Sigma],\quad\expectationvalue{\Sigma,\mathcal{T}_{w}^{T}[v+\delta]}=0,\quad\Sigma\succeq 0,$ $$

which implies that $\Sigma\in CCP(\mathcal{T}_{w}^{T}[v],\tfrac{1}{2}\mathcal{T}_{w}^{T}\mathcal{T}_{w})$.
∎

We now intuitively reconcile [eq. 24](#A3.E24) with the CCP definition:

- •
When $S(w)<2/\eta$, $w$ is in the interior of $\mathbb{S}_{\eta}$ so $U(w)=\emptyset$, the tangent cone is the entire space, and the projection is the identity map. Therefore [eq. 24](#A3.E24) reduces to gradient flow.
- •
When there is a single eigenvalue at $2/\eta$, $w$ is on the boundary of $\mathbb{S}_{\eta}$ and the tangent cone is given by the halfspace: $T_{\mathbb{S}_{\eta}}(w)=\{v:\expectationvalue{\nabla S(w),v}\leq 0\}$. If the negative gradient lies outside this halfspace (i.e. if gradient flow threatens to increase the sharpness above $2/\eta$), then the projection onto the halfspace is given by the projection onto the hyperplane: $-\eta\Pi_{\nabla S(w)}^{\perp}\nabla L(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 $T_{\mathbb{S}_{\eta}(w)}=\{v:\mathcal{T}_{w}^{\top}[v]\preceq 0\}$ requires solving a semidefinite quadratic program for which $\Sigma$ is the Lagrangian dual variable. The KKT conditions of this quadratic program are equivalent to the CCP that defines $\Sigma$ above.

###### Definition 3 (GD Central Flow, Projection Formulation) .

###### Lemma 2 .

###### Proof.

#### C.3.4 The Rate of Loss Decrease

###### Proposition 1 .

Under the GD central flow ([definition 3](#Thmdefinition3)), for almost all $t$ we have

$$ $\displaystyle\frac{dL(w)}{dt}=-\eta\norm{\mathrm{proj}_{T_{\mathbb{S}_{\eta}}(w)}[-\nabla L(w)]}^{2}.$ $$

This implies $\frac{dL}{dt}\leq 0$ (i.e. the loss monotonically decreases) and $\frac{dL}{dt}\geq-\eta\norm{\nabla L(w)}^{2}$ (i.e. the loss decreases at a slower rate than under the gradient flow).

###### Proof.

By the chain rule we have

$$ $\displaystyle\frac{dL(w)}{dt}$ $\displaystyle=\expectationvalue{\nabla L(w),\frac{dw}{dt}}$ $\displaystyle=\expectationvalue{\nabla L(w),\mathrm{proj}_{T_{\mathbb{S}_{\eta}}(w)}[-\eta\nabla L(w)]}$ $\displaystyle=-\eta\expectationvalue{-\nabla L(w),\mathrm{proj}_{T_{\mathbb{S}_{\eta}}(w)}[-\nabla L(w)]}$ $\displaystyle=-\eta\norm{\mathrm{proj}_{T_{\mathbb{S}_{\eta}}(w)}[-\nabla L(w)]}^{2}$ $$

where we used that for any orthogonal projection $\mathrm{proj}[cv]=\mathrm{proj}[c]$ when $c\geq 0$ to pull out the $\eta$, and that $\expectationvalue{v,\mathrm{proj}[v]}=\norm{\mathrm{proj}[v]}^{2}$ to get the last equality. Finally, the comparison with gradient flow follows from the inequality $\norm{\mathrm{proj}[v]}\leq\norm{v}$ for any orthogonal projection.
∎

###### Proposition 1 .

###### Proof.

### C.4 Scalar RMSProp

We follow a similar derivation to gradient descent. We define $U(w,\nu):=\ker[\nu^{-1/2}H(w)-(2/\eta)I]$, and let $\operatorname{Sym}(U(w,\nu))$ denote the subspace of symmetric matrices with span in $U(w,\nu)$. We now proceed with the time-averaging derivation.

If $w=\overline{w}+\delta$ with $\operatorname{\mathbb{E}}[\delta]=0$, we have that

$$ $\displaystyle\operatorname{\mathbb{E}}\norm{\nabla L(w)}^{2}\approx\operatorname{\mathbb{E}}\norm{\nabla L(\overline{w})}^{2}+\operatorname{\mathbb{E}}\norm{\nabla^{2}L(\overline{w})\delta}^{2}=\norm{\nabla L(\overline{w})}^{2}+\expectationvalue{\nabla^{2}L(\overline{w})^{2},\Sigma}$ (25) $$

where $\Sigma=\operatorname{\mathbb{E}}[\delta\delta^{T}]$. Because we assume that $\Sigma\in\mathcal{S}_{\text{crit}}$, $\nabla^{2}L(\overline{w})\Sigma=S(\overline{w})\Sigma$, so this expression is equal to

$$ $\displaystyle\operatorname{\mathbb{E}}\norm{\nabla L(w)}^{2}\approx\norm{\nabla L(\overline{w})}^{2}+S(\overline{w})^{2}\tr(\Sigma).$ $$

This suggests the central flow ansatz:

$$ $\displaystyle\begin{split}\frac{dw}{dt}&=-\frac{\eta}{\sqrt{\nu}}\quantity[\nabla L(w)+\tfrac{1}{2}\nabla^{3}L(w)[\Sigma(t)]]\\ \frac{d\nu}{dt}&=\tfrac{1-\beta_{2}}{\beta_{2}}\quantity[\norm{\nabla L(w)}^{2}+S(w)^{2}\mathrm{tr}(\Sigma(t))-\nu].\end{split}$ (26) $$

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),\nu(t),\Sigma(t))\}_{t\geq 0}$ satisfy the Scalar RMSProp central flow if they satisfy [eq. 26](#A3.E26) along with the conditions: $\Sigma(t)\succeq 0$, $\nu(t)^{-1/2}H(w(t))\preceq 2/\eta$, and $\Sigma(t)\in\mathcal{S}_{\text{crit}}(w(t),\nu(t))$ for almost all $t$.

As for gradient descent, this definition is fairly opaque and does not give a way of actually computing $\Sigma$. To do so, we will again derive a cone-complementarity problem formulation of the Scalar RMSProp central flow. The first step in this derivation is to differentiate the stability constraint: $\nu^{-1/2}H(w)\preceq 2/\eta$. We will fix a time $t$ and use $U$ to denote $U(w,\nu)$ to simplify notation. We will also define the linear operator $\mathcal{T}:\operatorname{Sym}(U(w,\nu))\to\mathbb{R}^{d}$ by

$$ $\displaystyle\mathcal{T}[\Sigma]=\nabla_{w}\expectationvalue{H(w),\Sigma}\mbox{\quad for\quad}\Sigma\in\operatorname{Sym}(U(w,nu)).$ $$

Let $\dot{H}:=\frac{d}{dt}\nu^{-1/2}H(w)\evaluated{}_{U}$ 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 $\dot{H}\preceq 0$. We can compute $\dot{H}$ using the chain rule.

$$ $\displaystyle\dot{H}$ $\displaystyle:=\frac{d}{dt}\nu^{-1/2}H(w)\evaluated{}_{U}$ $\displaystyle=\nu^{-1/2}\quantity(\frac{d}{dt}H(w))\evaluated{}_{U}+\quantity(\frac{d}{dt}\nu^{-1/2})H(w)\evaluated{}_{U}$ $\displaystyle=\nu^{-1/2}\mathcal{T}^{T}\quantity[\frac{dw}{dt}]-\frac{1}{2\nu^{3/2}}\frac{d\nu}{dt}\frac{2\sqrt{\nu}}{\eta}I$ $\displaystyle=\frac{\eta}{\nu}\mathcal{T}^{T}\quantity[-\nabla L(w)-\tfrac{1}{2}\mathcal{T}[\Sigma]]+\frac{1-\beta_{2}}{\beta_{2}}\frac{1}{\eta\nu}\quantity[\nu-\norm{\nabla L(w)}^{2}-S(w)^{2}\tr(\Sigma)]I.$ $$

We can now assume that $(w,\nu)$ are at EOS so that $\nu=\frac{\eta^{2}S(w)^{2}}{4}$. Otherwise, complementarity forces $\Sigma=0$. Substituting this gives:

$$ $\displaystyle\dot{H}$ $\displaystyle=\frac{4}{\eta S(w)^{2}}\mathcal{T}^{T}\quantity[-\nabla L(w)-\tfrac{1}{2}\mathcal{T}[\Sigma]]$ $\displaystyle\qquad+\frac{1-\beta_{2}}{\beta_{2}}\frac{4}{\eta^{3}S(w)^{2}}\quantity[\frac{\eta^{2}S(w)^{2}}{4}-\norm{\nabla L(w)}^{2}-S(w)^{2}\tr(\Sigma)]I.$ $$

We will now group the constant terms and the terms linear in $\Sigma$. Define:

$$ $\displaystyle\alpha(w)$ $\displaystyle=\frac{4}{\eta S(w)^{2}}\mathcal{T}^{T}\quantity[-\nabla L(w)]+\frac{1-\beta_{2}}{\beta_{2}}\frac{4}{\eta^{3}S(w)^{2}}\quantity[\frac{\eta^{2}S(w)^{2}}{4}-\norm{\nabla L(w)}^{2}]I$ $\displaystyle\beta(w)$ $\displaystyle=\frac{2}{\eta S(w)^{2}}\mathcal{T}^{T}\mathcal{T}+\frac{1-\beta_{2}}{\beta_{2}}\frac{4}{\eta^{3}}I\otimes I.$ $$

Then $\dot{H}=\alpha-\beta[\Sigma]$. As in the derivation for GD, if $\dot{H}$ is strictly negative definite in some direction, these directions will drop from the critical subspace so right continuity forces $\Sigma\in\ker[\dot{H}]$. This gives us the cone complementarity definition of the Scalar RMSProp central flow:

###### Definition 5 (Scalar RMSProp Central Flow, CCP Formulation) .

We say that $\{(w(t),\nu(t)\}_{t\geq 0}$ follow the Scalar RMSProp central flow if they satisfy [eq. 26](#A3.E26) where $\Sigma\in CCP(\alpha(w),\beta(w))$.

We can now use this formulation to prove [eq. 16](#S4.E16) when $U(w,\nu)$ is one dimensional. In the one dimensional case, $\alpha,\beta$ are both scalars and assuming that $\alpha(w)>0$ the solution to the CCP is simply $\Sigma=\alpha(w)/\beta(w)$. In addition, in this one dimensional case we simply have $\mathcal{T}=\nabla S(w)$. Therefore,

$$ $\displaystyle\sigma^{2}$ $\displaystyle=\frac{\frac{4}{\eta S(w)^{2}}\expectationvalue{-\nabla L(w),\nabla S(w)}+\frac{1-\beta_{2}}{\beta_{2}}\frac{4}{\eta^{3}S(w)^{2}}\quantity[\frac{\eta^{2}S(w)^{2}}{4}-\norm{\nabla L(w)}^{2}]}{\frac{2}{\eta S(w)^{2}}\norm{\nabla S(w)}^{2}+\frac{1-\beta_{2}}{\beta_{2}}\frac{4}{\eta^{3}}}$ (27) $\displaystyle=\frac{\beta_{2}\expectationvalue{-\nabla L(w),\nabla S(w)}+(1-\beta_{2})\quantity[\frac{S(w)^{2}}{4}-\frac{\norm{\nabla L(w)}^{2}}{\eta^{2}}]}{\beta_{2}\tfrac{1}{2}\norm{\nabla S(w)}^{2}+(1-\beta_{2})S(w)^{2}/\eta^{2}}.$ (28) $$

###### Definition 4 (Scalar RMSProp Central Flow, Differential Variational Inequality Formulation) .

###### Definition 5 (Scalar RMSProp Central Flow, CCP Formulation) .

#### C.4.1 The effect of the hyperparameters η , β 2 𝜂 subscript 𝛽 2 \eta,\beta_{2}

###### Lemma 3 .

Fix an iterate $w=w(t)$ and define $\nu=\nu(w;\eta):=\frac{\eta^{2}S(w)^{2}}{4}$ so that $S^{\text{eff}}(w,\nu)=2/\eta$. Let $U$ be the top eigenspace of $H(w)$. Then if $\Sigma(t)\succ 0$ under the Scalar RMSProp central flow ([Definition 5](#Thmdefinition5)), we have that

$$ $\displaystyle\frac{\partial}{\partial\eta}\frac{dH}{dt}\evaluated{}_{U}=-cI_{U}\mbox{\quad where\quad}C\geq 0$ $$

and where $I_{U}$ denotes the identity matrix on the subspace $U$. In other words, larger learning rates $\eta$ more aggressively decrease the curvature and they do so uniformly across all eigenvalues in the critical subspace $U$. In addition,

$$ $\displaystyle\frac{\partial}{\partial\beta_{2}}\frac{dH}{dt}\evaluated{}_{U}=-CI_{U}$ $$

where $C\geq 0$ when $\tr\Sigma\evaluated{}_{\beta_{2}=1}\leq\tr\Sigma\evaluated{}_{\beta_{2}=0}$ and $C\leq 0$ otherwise. Therefore $\beta_{2}$ can either monotonically increase or decrease the curvature regularization depending on whether gradient descent ($\beta_{2}=1$) or normalized gradient descent ($\beta_{2}=0$) would have a larger oscillation variance.

###### Proof.

The condition $\Sigma(t)\succ 0$ implies that $\dot{H}(t)=0$ and $\Sigma=\beta^{-1}\alpha$. We will rescale $\alpha,\beta$ by $\eta S(w)^{2}/4$:

$$ $\displaystyle\hat{\alpha}$ $\displaystyle=\mathcal{T}^{T}\quantity[-\nabla L(w)]+\frac{1-\beta_{2}}{\beta_{2}}\quantity[\frac{S(w)^{2}}{4}-\frac{\norm{\nabla L(w)}^{2}}{\eta^{2}}]I$ $\displaystyle\hat{\beta}$ $\displaystyle=\tfrac{1}{2}\mathcal{T}^{T}\mathcal{T}+\frac{1-\beta_{2}}{\beta_{2}}\frac{S(w)^{2}}{\eta^{2}}I\otimes I.$ $$

We will also use $\hat{\alpha}_{w},\hat{\beta}_{w}$ to refer to just the first terms in $\hat{\alpha},\hat{\beta}$:

$$ $\displaystyle\hat{\alpha}_{w}=\mathcal{T}^{T}\quantity[-\nabla L(w)]\mbox{\quad and\quad}\hat{\beta}_{w}=\tfrac{1}{2}\mathcal{T}^{T}\mathcal{T}.$ $$

Then $\Sigma=\hat{\beta}^{-1}\hat{\alpha}$. Differentiating this with respect to $\eta,\beta_{2}$ gives:

$$ $\displaystyle\frac{\partial}{\partial\eta}\Sigma$ $\displaystyle=\frac{1-\beta_{2}}{2\eta^{3}\beta_{2}}\hat{\beta}^{-1}[I]\quantity[\norm{\nabla L(w)}^{2}+S(w)^{2}\tr\Sigma]$ $\displaystyle\frac{\partial}{\partial\beta_{2}}\Sigma$ $\displaystyle=-\frac{1}{\beta_{2}^{2}}\hat{\beta}^{-1}[I]\quantity[\frac{S(w)^{2}}{4}-\frac{\norm{\nabla L(w)}^{2}}{\eta^{2}}-\frac{S(w)^{2}}{\eta^{2}}\tr\Sigma]$ $$

Then:

$$ $\displaystyle\frac{\partial}{\partial\eta}\frac{dH}{dt}\evaluated{}_{U}=-\hat{\beta}_{w}\quantity[\frac{\partial}{\partial\eta}\Sigma],\quad\frac{\partial}{\partial\beta_{2}}\frac{dH}{dt}\evaluated{}_{U}=-\hat{\beta}_{w}\quantity[\frac{\partial}{\partial\beta_{2}}\Sigma]$ $$

Therefore it suffices to compute $\hat{\beta}_{w}\hat{\beta}^{-1}[I]$. To compute this, we use the Sherman–Morrison formula. Let $c=\frac{1-\beta_{2}}{\beta_{2}}\frac{S(w)^{2}}{\eta^{2}}$. Then,

$$ $\displaystyle\hat{\beta}_{w}\hat{\beta}^{-1}[I]=I\quantity[1-\frac{c\beta_{w}^{-1}[I,I]}{1+c\beta_{w}^{-1}[I,I]}]=I\quantity[\frac{1}{1+c\beta_{w}^{-1}[I,I]}].$ $$

Finally, note that $\quantity[\frac{1}{1+c\beta_{w}^{-1}[I,I]}]\geq 0$, which immediately implies the result for $\eta$. For $\beta_{2}$, we have shown that

$$ $\displaystyle\frac{\partial}{\partial\beta_{2}}\frac{dH}{dt}\evaluated{}_{U}=-CI_{U}$ $$

where $C\geq 0$ if and only if

$$ $\displaystyle\tr\Sigma\leq\frac{\eta^{2}}{4}-\frac{\norm{\nabla L(w)}^{2}}{S(w)^{2}}.$ $$

To rewrite this in a form that is independent of $\Sigma$, note that

$$ $\displaystyle\tr\Sigma$ $\displaystyle=\tr[\hat{\beta}^{-1}\hat{\alpha}]$ $\displaystyle=\expectationvalue{\hat{\beta}^{-1}[I],\alpha}$ $\displaystyle=\expectationvalue{\alpha,\hat{\beta}_{w}^{-1}[I]}\quantity[\frac{1}{1+c\beta_{w}^{-1}[I,I]}]$ $\displaystyle=\frac{\tr[\hat{\beta}_{w}^{-1}\hat{\alpha}_{w}]+c\beta_{w}^{-1}[I,I]\quantity[\frac{\eta^{2}}{4}-\frac{\nabla L(w)^{2}}{S(w)^{2}}]}{1+c\beta_{w}^{-1}[I,I]}.$ $$

Therefore $\tr\Sigma$ is a weighted average between $\tr[\hat{\beta}_{w}^{-1}\hat{\alpha}_{w}]$ and $\frac{\eta^{2}}{4}-\frac{\nabla L(w)^{2}}{S(w)^{2}}$, so $\tr[\Sigma]\leq\frac{\eta^{2}}{4}-\frac{\nabla L(w)^{2}}{S(w)^{2}}$ if and only if $\tr[\hat{\beta}_{w}^{-1}\hat{\alpha}_{w}]\leq\frac{\eta^{2}}{4}-\frac{\nabla L(w)^{2}}{S(w)^{2}}$. This first expression is just $\tr\Sigma\evaluated{}_{\beta_{2}=1}$ and the second is just $\tr\Sigma\evaluated{}_{\beta_{2}=0}$.
∎

The condition $\Sigma(t)\succ 0$ is equivalent to requiring that we are not currently at a “breakpoint,” i.e. a time $t$ at which an eigenvalue drops from EOS. The set of such $t$ constitute a measure zero set, so the above lemma holds for almost all $t$.

###### Lemma 3 .

###### Proof.

#### C.4.2 The small β 2 subscript 𝛽 2 \beta_{2} correction

The following lemma shows that the $\beta_{2}\to\frac{1-\beta_{2}}{\beta_{2}}$ correction for the Scalar RMSProp and RMSProp central flows allows the continuous EMA to match the discrete EMA for approximately linear targets $f(t)$:

###### Lemma 4 .

Let $f:\mathbb{R}\to\mathbb{R}$ be a continuous time process with $|f^{(2)}(t)|\leq\Delta$ for all $t$. Then if

$$ $\displaystyle\nu_{t}$ $\displaystyle=\beta_{2}\nu_{t-1}+(1-\beta_{2})f(t)$ $\displaystyle\nu^{\prime}(t)$ $\displaystyle=\frac{1-\beta_{2}}{\beta_{2}}\quantity[f(t)-\nu(t)]$ $$

we have that for all integers $t\geq\tilde{O}((1-\beta_{2})^{-1})$,

$$ $\displaystyle\absolutevalue{\nu(t)-\nu_{t}}\leq O\quantity(\frac{\beta_{2}\Delta}{(1-\beta_{2})^{2}}).$ $$

###### Proof.

First define

$$ $\displaystyle\delta_{t}:=\nu_{t}-f(t)-\frac{\beta_{2}}{1-\beta_{2}}f^{\prime}(t)$ $\displaystyle\delta(t):=\nu(t)-f(t)-\frac{\beta_{2}}{1-\beta_{2}}f^{\prime}(t).$ $$

Then we have that:

$$ $\displaystyle\delta_{t}$ $\displaystyle=\beta_{2}\nu_{t-1}+(1-\beta_{2})f(t)-f(t)-\frac{1-\beta_{2}}{\beta_{2}}f^{\prime}(t)$ $\displaystyle=\beta_{2}\nu_{t-1}-\beta_{2}f(t)-\frac{\beta_{2}}{1-\beta_{2}}f^{\prime}(t)$ $\displaystyle=\beta_{2}\delta_{t-1}-\beta_{2}[f(t)-f(t-1)]-\frac{\beta_{2}}{1-\beta_{2}}f^{\prime}(t)+\frac{\beta_{2}^{2}}{1-\beta_{2}}f^{\prime}(t-1)$ $\displaystyle=\beta_{2}\delta_{t-1}-\beta_{2}[f(t)-f^{\prime}(t)-f(t-1)]+O\quantity(\frac{\beta_{2}^{2}\Delta}{1-\beta_{2}})$ $\displaystyle=\beta_{2}\delta_{t-1}+O\quantity(\frac{\beta_{2}\Delta}{1-\beta_{2}}).$ $$

Therefore,

$$ $\displaystyle\delta_{t}=\beta_{2}^{t}\delta_{0}+O\quantity(\frac{\beta_{2}\Delta}{(1-\beta_{2})^{2}}).$ $$

Similarly,

$$ $\displaystyle\delta^{\prime}(t)$ $\displaystyle=\frac{1-\beta_{2}}{\beta_{2}}\quantity[f(t)-\nu(t)]-f^{\prime}(t)-\frac{\beta_{2}}{1-\beta_{2}}f^{\prime\prime}(t)$ $\displaystyle=-\frac{1-\beta_{2}}{\beta_{2}}\delta(t)+O\quantity(\frac{\beta_{2}\Delta}{1-\beta_{2}})$ $$

so

$$ $\displaystyle\delta(t)=e^{-\frac{1-\beta_{2}}{\beta_{2}}t}\delta(0)+O\quantity(\frac{\beta_{2}^{2}\Delta}{(1-\beta_{2})^{2}}).$ $$

Therefore subtracting the bounds on $\delta_{t}$ and $\delta(t)$ gives:

$$ $\displaystyle\absolutevalue{\nu_{t}-\nu(t)}\lesssim\frac{\beta_{2}\Delta}{(1-\beta_{2})^{2}}$ $$

as desired.
∎

###### Lemma 4 .

###### Proof.

### C.5 RMSProp

We first motivate the stability condition $\lambda_{max}(P^{-1}H)\leq 2$ on a quadratic. First, note that $P^{-1}H$ is similar to $P^{-1/2}HP^{-1/2}$ which is symmetric so $P^{-1}H$ has real eigenvalues and real eigenvectors. Next, consider gradient descent on the quadratic $\tfrac{1}{2}w^{T}Hw$ with preconditioner $P$. The update is:

$$ $\displaystyle w\leftarrow w-P^{-1}Hw=(I-P^{-1}H)w.$ $$

Therefore, $w_{t}=(I-P^{-1}H)^{t}w_{0}$. If the eigenvalues of $P^{-1}H$ are in the range $(0,2)$, then $w\to 0$. Otherwise, it diverges along the corresponding *right* eigenvectors of $P^{-1}H$. We can reinterpret this stability condition using the equivalent condition $H\preceq 2P$. Note that the top right eigenvectors of $P^{-1}H$ correspond to the top eigenvectors for the generalized eigenvalue problem $Hv=\lambda Pv$.

We again follow a similar derivation to gradient descent. We define the critical subspace by $U(w,\nu):=\ker[H-2P]$ where $P=\text{diag}(\sqrt{\nu}/\eta)$, and use $\operatorname{Sym}(U(w,\nu))$ to denote the set of symmetric matrices on this subspace. We can now proceed with the time-averaging argument.

As before, we assume $w=\overline{w}+\delta$ with $\operatorname{\mathbb{E}}[\delta]=0$. Rather than expanding the squared gradient norm, we expand the element wise squared gradient:

$$ $\displaystyle\operatorname{\mathbb{E}}[\nabla L(w)^{\odot 2}]$ $\displaystyle\approx\nabla L(\overline{w})^{\odot 2}+\operatorname{\mathbb{E}}[(H(\overline{w})\delta)^{\odot 2}]$ (29) $\displaystyle=\nabla L(\overline{w})^{\odot 2}+\operatorname{diag}\quantity[H(\overline{w})\Sigma H(\overline{w})]$ (30) $$

where $\Sigma:=\operatorname{\mathbb{E}}[\delta\delta^{T}]$. This can be further simplified using the fact that $\Sigma$ is in the top eigenspace, so $H(\overline{w})\Sigma=2P\Sigma$. Therefore,

$$ $\displaystyle\operatorname{diag}\quantity[H(\overline{w})\Sigma H(\overline{w})]=4P\operatorname{diag}[\Sigma]P=\frac{4\nu}{\eta^{2}}\odot\operatorname{diag}[\Sigma].$ $$

This suggests the central flow ansatz:

$$ $\displaystyle\begin{split}\frac{dw}{dt}&=-\frac{\eta}{\sqrt{\nu}}\odot\quantity[\nabla L(w)+\tfrac{1}{2}\nabla^{3}L(w)[\Sigma(t)]]\\ \frac{d\nu}{dt}&=\frac{1-\beta_{2}}{\beta_{2}}\quantity[\nabla L(w)^{\odot 2}+\frac{4\nu}{\eta^{2}}\odot\operatorname{diag}[\Sigma(t)]-\nu].\end{split}$ (31) $$

As for gradient descent and Scalar RMSProp, this immediately implies a differential variational inequality definition:

###### Definition 6 (RMSProp Central Flow, Differential Variational Inequality Formulation) .

We say that $\{(w(t),\nu(t)\}_{t\geq 0}$ follow the RMSProp central flow if for almost all $t\geq 0$, they satisfy [eq. 31](#A3.E31) along with the conditions: $\Sigma(t)\succeq 0$, $H(w(t))\preceq 2P(t)$, and $\expectationvalue{H(w(t))-2P(t),\Sigma(t)}=0$ where $P(t)=\sqrt{\nu(t)}/\eta$.

We will now show that $\Sigma$ can be computed as the solution to a cone complementarity problem. Fix some time $t$, let $U=U(w,\nu)$, and define the linear operator $\mathcal{T}:\operatorname{Sym}(U)\to\mathbb{R}^{d}$ by $\mathcal{T}[\Sigma]=\nabla_{w}\expectationvalue{H(w),\Sigma}$. We can differentiate the stability condition in the critical subspace as:

$$ $\displaystyle\dot{H}$ $\displaystyle:=\frac{d}{dt}(H(w)-2P)\evaluated{}_{U}$ $\displaystyle=\mathcal{T}^{T}\quantity[\frac{dw}{dt}]-2\frac{\partial P}{\partial\nu}\frac{d\nu}{dt}\evaluated{}_{U}$ $\displaystyle=\mathcal{T}^{T}P^{-1}\quantity[-\nabla L(w)-\tfrac{1}{2}\mathcal{T}[\Sigma]]$ $\displaystyle\qquad+\frac{1-\beta_{2}}{\beta_{2}}\operatorname{diag}\quantity[\frac{1}{\eta\nu^{1/2}}\odot\quantity[\nu-\nabla L(w)^{\odot 2}-\frac{4\nu}{\eta^{2}}\odot\operatorname{diag}[\Sigma]]]\evaluated{}_{U}.$ $$

We again group the constant terms and the terms linear in $\Sigma$. Define:

$$ $\displaystyle\alpha(w,P)$ $\displaystyle:=\mathcal{T}^{T}[-P^{-1}\nabla L(w)]+\frac{1-\beta_{2}}{\beta_{2}}P^{-1}\quantity[P^{2}-\frac{\operatorname{diag}[\nabla L(w)^{\odot 2}]}{\eta^{2}}]\evaluated{}_{U}$ $\displaystyle\beta(w,P)[\Sigma]$ $\displaystyle:=\tfrac{1}{2}\mathcal{T}^{T}P^{-1}\mathcal{T}[\Sigma]+\frac{1-\beta_{2}}{\beta_{2}}\frac{4}{\eta^{2}}P\operatorname{diag}[\Sigma]\evaluated{}_{U}.$ $$

Note that we chose to define $\beta$ through its action on $\Sigma$ to avoid unnecessarily complicated tensor notation. Then $\dot{H}=\alpha-\beta[\Sigma]$. As in the derivations of gradient descent and Scalar RMSProp, we require that $\Sigma$ solves the CCP:

$$ $\displaystyle\Sigma\succeq 0,\quad\dot{H}\preceq 0,\quad\expectationvalue{\Sigma,\dot{H}}=0.$ (32) $$

Together these define the CCP formulation, which can be efficiently simulated:

###### Definition 7 (RMSProp Central Flow, CCP Formulation) .

We say that $\{(w(t),\nu(t)\}_{t\geq 0}$ follow the RMSProp central flow if for almost all $t\geq 0$, they satisfy [eq. 31](#A3.E31) with $\Sigma(t)\in CCP(\alpha(w(t),P(t)),\beta(w(t),P(t))$ where $P(t)=\operatorname{diag}[\sqrt{\nu(t)}/\eta]$.

###### Definition 6 (RMSProp Central Flow, Differential Variational Inequality Formulation) .

###### Definition 7 (RMSProp Central Flow, CCP Formulation) .

### C.6 Solving for the stationary ν 𝜈 \nu in RMSProp

Setting $\frac{d\nu}{dt}=0$ in [eq. 31](#A3.E31) gives:

$$ $\displaystyle\nu=\nabla L(w)^{\odot 2}+(4/\eta^{2})\nu\odot\operatorname{diag}[\Sigma].$ (33) $$

In addition, by complementarity we know that

$$ $\displaystyle\Sigma\quantity(H(w)-2\operatorname{diag}[\sqrt{\nu}/\eta])=0.$ (34) $$

###### Proposition 2 .

If $w,\nu$ satisfy [eqs. 33](#A3.E33) and [34](#A3.E34), then $P:=\operatorname{diag}(\sqrt{\nu}/\eta)$ minimizes the convex program [eq. 19](#S5.E19).

###### Proof.

Let $p=\sqrt{\nu}/\eta$. The convex program [eq. 19](#S5.E19) can be written as:

$$ $\displaystyle\min_{p}\sum_{i}p_{i}+\frac{\nabla L(w)_{i}^{2}}{p_{i}}\mbox{\quad such that\quad}H(w)\preceq 2\operatorname{diag}(p).$ (35) $$

If $\hat{\Sigma}$ is the dual variable for the semidefinite constraint, the KKT conditions for this program are:

$$ $\displaystyle 1-\frac{\nabla L(w)_{i}^{2}}{p_{i}^{2}}-2\hat{\Sigma}_{ii}=0\quad\forall i,\quad\hat{\Sigma}\quantity(H(w)-2\operatorname{diag}(p))=0.$ (36) $$

We will prove that $(p,\hat{\Sigma})=(\frac{\sqrt{\nu}}{\eta},\frac{2}{\eta^{2}}\Sigma)$ solve the KKT conditions [eq. 36](#A3.E36). Note that the complementary slackness KKT condition is equivalent to the complementarity condition on $\Sigma$. In addition, dividing the stationarity condition for $\nu$ ([eq. 33](#A3.E33)) by $\nu$ gives

$$ $\displaystyle 1-\frac{\nabla L(w)^{\odot 2}}{\nu_{i}}-(4/\eta^{2})\Sigma_{ii}=0\quad\forall i$ (37) $$

which is equivalent to the first condition in [eq. 36](#A3.E36) after substituting $\hat{\Sigma}=\frac{2}{\eta^{2}}\Sigma(t)$.
∎

###### Proposition 3 .

For any $g,H$, the solution to [eq. 19](#S5.E19) is unique.

###### Proof.

Assume there are two minimizers $P,P^{\prime}$ and let $p:=\operatorname{diag}(P),\delta:=\operatorname{diag}(P^{\prime}-P)$. Then by convexity, $\operatorname{diag}[p+\epsilon\delta]$ also minimizes [eq. 19](#S5.E19) for any $\epsilon\leq 1$. Therefore, differentiating the objective function in this direction gives:

$$ $\displaystyle\sum_{i}\delta_{i}\quantity[1-\frac{1}{\eta^{2}}\frac{g_{i}^{2}}{p_{i}^{2}}]=0.$ $$

Taking another derivative implies that:

$$ $\displaystyle\sum_{i}\frac{g_{i}^{2}}{p_{i}^{3}}\delta_{i}^{2}=0.$ $$

This implies that $\delta_{i}=0$ in any direction where $g_{i}\neq 0$. Let $I$ be the set of indices for which $g_{i}\neq 0$, and for any vector $p$, let $p$ denote the vector $p$ restricted to the indices in $I$. Define the linear map $g$ by

$$ $\displaystyle g[v_{I}]_{i}:=\begin{cases}v_{i}&i\in I\\ p_{i}&i\notin I\end{cases}.$ $$

In other words, $g$ takes a reduced vector $v_{I}$ and fills in the missing entries with $p$. Next, define the operator $\mathcal{A}$ by

$$ $\displaystyle\mathcal{A}^{T}[v_{I}]=\operatorname{diag}[g[v_{I}]]\oplus\operatorname{diag}[v_{I}].$ $$

Then both $p_{I},p_{I}^{\prime}$ minimize the following reduced SDP:

$$ $\displaystyle\min_{p_{I}}\sum_{i\in I}p_{i}\mbox{\quad such that\quad}\tfrac{1}{2}H(w)\oplus 0_{|I|\times|I|}\preceq\mathcal{A}^{T}(p).$ $$

Now we apply (de Carli Silva and
Tunçel, , Proposition 1) with $(\mathcal{A},\mathbf{1}_{|I|})$. First, note that $\mathcal{A}[I_{d+|I|}]=2\mathbf{1}_{|I|}$ which satisfies the first condition. Next, for any $y\neq 0$, we can take $z=|y|$ to satisfy the second condition, as in the proof of (de Carli Silva and
Tunçel, , Corollary 2) Therefore $p_{I}=p^{\prime}_{I}$, and as we have already shown equality on $I^{c}$, we must have $p=p^{\prime}$.
∎

###### Proposition 2 .

###### Proof.

###### Proposition 3 .

###### Proof.

### C.7 Arbitrary Preconditioned Methods

In this section, we derive a 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 $\nu$ be the “state” of the preconditioner, and let $P$ be a function that maps the state $\nu$ to a symmetric, positive definite matrix. Specifically, we consider the discrete update:

$$ $\displaystyle\nu_{t}=\nu_{t-1}+f(\nu_{t-1},w_{t}),\quad w_{t+1}=w_{t}-P(\nu_{t})^{-1}\nabla L(w_{t}).$ $$

This formulation is very general and includes a wide variety of updates including:

- •
*GD with a learning rate schedule $\eta(t)$:* Set $f(\nu,w)=1$ so that $\nu(t)=t$, and $P(t)=\eta(t)^{-1}I$
- •
*Vanilla RMSProp:* Set $f(\nu,w)=(1-\beta_{2})[\nabla L(w)^{\odot 2}-\nu]$ and $P(\nu)=\operatorname{diag}[\sqrt{\nu}/\eta]$
- •
*RMSProp with $\epsilon$, bias correction, and learning rate schedule $\eta(t)$:* Set $\nu=[v,t]$, $f([v,t],w)=[(1-\beta_{2})[\nabla L(w)^{\odot 2}-v],1]$ and define
$P([v,t])=\operatorname{diag}\quantity[\frac{1}{\eta(t)}\cdot\sqrt{\frac{v}{1-\beta_{2}^{t}}+\epsilon}].$

Note that this trick of embedding $t$ into the state variable $\nu$ allows us to automatically derive central flows for any smooth hyperparameter schedules (e.g. $\eta(t),\beta_{2}(t),\epsilon(t)$) as a simple corollary.

As for RMSProp, the stability of this algorithm requires $\lambda_{max}(P^{-1}H)\leq 2$ or equivalently $H\preceq 2P$. To derive the central flow, we assume that $w=\overline{w}+\delta$ with $\operatorname{\mathbb{E}}[\delta]=0$ and $\operatorname{\mathbb{E}}[\delta\delta^{T}]=\Sigma$. Taylor expanding and time-averaging the update for $\nu$ gives:

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\nu,w)]=\operatorname{\mathbb{E}}[f(\nu,\nabla L(\overline{w}+\delta))]\approx f(\nu,\nabla L(\overline{w}))+\tfrac{1}{2}\nabla^{2}_{\overline{w}}f(\nu,\overline{w})[\Sigma]$ $$

These motivate the central flow ansatz:

$$ $\displaystyle\begin{split}\frac{dw}{dt}&=-P(\nu)^{-1}\quantity[\nabla L(w)+\tfrac{1}{2}\nabla^{3}L(w)[\Sigma]]\\ \frac{d\nu}{dt}&=f(\nu,w)+\tfrac{1}{2}\nabla^{2}_{w}f(\nu,w)[\Sigma].\end{split}$ (38) $$

We say that $\{w(t),\nu(t),\Sigma(t)\}_{t\geq 0}$ satisfy the DVI formulation of the central flow if for almost all $t\geq 0$ they satisfy [eq. 38](#A3.E38), $\Sigma(t)\succeq 0$, $H(w(t))\preceq 2P(\nu(t))$, and $\expectationvalue{\Sigma(t),H(w(t))-2P(\nu(t))}$.

To derive the CCP formulation which is efficiently computable, we differentiate the stability condition. Fix an iterate $t$, let $U:=\ker[H-2P(\nu(t))]$ be the critical subspace at time $t$, and define $\dot{H}:=\frac{dH}{dt}\evaluated{}_{U}$ under [eq. 38](#A3.E38). Then for the stability condition to remain true, we need $\dot{H}\preceq 0$. To compute it, define $\mathcal{T}:\operatorname{Sym}(U)\to\mathbb{R}^{d}$ by $\mathcal{T}[\Sigma]:=\nabla_{w}\expectationvalue{H(w),\Sigma}$. Then,

$$ $\displaystyle\dot{H}$ $\displaystyle:=\frac{dH}{dt}\evaluated{}_{U}$ $\displaystyle=\mathcal{T}^{T}\quantity[\frac{dw}{dt}]-2\nabla_{\nu}P(\nu)\quantity[\frac{d\nu}{dt}]\evaluated{}_{U}$ $\displaystyle=\mathcal{T}^{T}P(\nu)^{-1}\quantity[-\nabla L(w)-\tfrac{1}{2}\mathcal{T}[\Sigma]]-2\nabla_{\nu}P(\nu)\quantity[f(\nu,w)+\tfrac{1}{2}\nabla^{2}_{w}f(\nu,w)[\Sigma]]\evaluated{}_{U}.$ $$

Splitting up the constant terms and the terms linear in $\Sigma$, we define:

$$ $\displaystyle\alpha(w,\nu)$ $\displaystyle:=\mathcal{T}^{T}[-P(\nu)^{-1}\nabla L(w)]-2\nabla_{\nu}P(\nu)|_{U}f(\nu,w)\in\operatorname{Sym}(U)$ $\displaystyle\beta(w,\nu)$ $\displaystyle:=\tfrac{1}{2}\mathcal{T}^{T}P(\nu)^{-1}\mathcal{T}+\nabla_{\nu}P(\nu)|_{U}\nabla^{2}_{w}f(\nu,w)\in\operatorname{Sym}(U)\otimes\operatorname{Sym}(U).$ $$

Then $\dot{H}=\alpha-\beta[\Sigma]$, so by the same arguments as for gradient descent,Scalar RMSProp,RMSProp, $\Sigma$ needs to solve $CCP(\alpha(w,\nu),\beta(w,\nu))$, and plugging this $\Sigma$ into [eq. 38](#A3.E38) gives the CCP formulation for the central flow for this preconditioned method.

### C.8 Cone Complementarity Problems

In this section we will let $k\geq 1$ be a positive integer, and we will use $\operatorname{Sym}_{k}(\mathbb{R})$ to denote the set of $k\times k$ symmetric matrices.

###### Definition 8 (Cone Complementarity Problem) .

Let $\alpha\in\operatorname{Sym}_{k}(\mathbb{R})$, and let $\beta\in\operatorname{Sym}_{k}(\mathbb{R})\otimes\operatorname{Sym}_{k}(\mathbb{R})$ be a linear operator on symmetric matrices. We say that the matrix $X\in\operatorname{Sym}_{k}(\mathbb{R})$ solves the cone complementarity problem $\mathrm{CCP}(\alpha,\beta)$ if:

$$ $\displaystyle X\succeq 0,\quad\alpha-\beta[X]\preceq 0,\quad\expectationvalue{X,\alpha-\beta[X]}=0.$ $$

###### Lemma 5 .

Let $\beta\in\operatorname{Sym}_{k}(\mathbb{R})\otimes\operatorname{Sym}_{k}(\mathbb{R})$ be a symmetric linear operator on symmetric matrices.

- 1.
If $\beta\succeq 0$ as an operator on $\operatorname{Sym}_{k}(\mathbb{R})$, then $CCP(\alpha,\beta)$ has a solution for all $\alpha$.
- 2.
If $\beta\succ 0$, this solution is unique. Otherwise, all solutions $X,X^{\prime}$ differ by a matrix in the kernel of $\beta$.

###### Proof.

Consider the quadratic program

$$ $\displaystyle\min_{X}\tfrac{1}{2}\beta[X,X]-\expectationvalue{X,\alpha}\mbox{\quad such that\quad}X\succeq 0.$ $$

Because $\beta\succeq 0$, this program is convex. The KKT conditions for this program are $\beta[X]-\alpha\succeq 0$ and the complementary slackness condition $X(\beta[X]-\alpha)=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,X^{\prime}$ be solutions to the conic complementarity program. Then taking the trace of the complementarity relations we have:

$$ $\displaystyle\expectationvalue{\alpha,X}-\beta[X,X]=0,\quad\expectationvalue{\alpha,X^{\prime}}-\beta[X^{\prime},X^{\prime}]=0.$ $$

In addition, because both $\beta[X]-\alpha$ and $X^{\prime}$ are PSD and vice-versa, we must have

$$ $\displaystyle\expectationvalue{X^{\prime},\beta[X]-\alpha}\leq 0,\quad\expectationvalue{X,\beta[X^{\prime}]-\alpha}\leq 0.$ $$

Adding these four conditions gives:

$$ $\displaystyle\beta[X,X]+\beta[X^{\prime},X^{\prime}]-\beta[X,X^{\prime}]-\beta[X^{\prime},X]\leq 0\iff\beta[X-X^{\prime},X-X^{\prime}]\leq 0.$ $$

However, since $\beta\succeq 0$ this implies that $\beta[X-X^{\prime},X-X^{\prime}]=0$. When $\beta\succ 0$ this implies that $X=X^{\prime}$, and when $\beta\succeq 0$ this implies that $X-X^{\prime}$ is in the kernel of $\beta$.
∎

###### Definition 8 (Cone Complementarity Problem) .

###### Lemma 5 .

###### Proof.

## Appendix D Full Experiments

### D.1 Gradient Descent

#### D.1.1 CIFAR 10

Figure: (a) $\eta=0.01$
Refer to caption: /html/2410.24206/assets/x31.png

Figure: (a) $\eta=0.01$
Refer to caption: /html/2410.24206/assets/x34.png

Figure: (a) $\eta=0.01$
Refer to caption: /html/2410.24206/assets/x37.png

#### D.1.2 CIFAR 10 (4 class)

Figure: (a) $\eta=0.01$
Refer to caption: /html/2410.24206/assets/x40.png

Figure: (a) $\eta=0.005$
Refer to caption: /html/2410.24206/assets/x43.png

Figure: (a) $\eta=0.01$
Refer to caption: /html/2410.24206/assets/x46.png

Figure: (a) $\eta=0.005$
Refer to caption: /html/2410.24206/assets/x49.png

Figure: (a) $\eta=0.01$
Refer to caption: /html/2410.24206/assets/x52.png

Figure: (a) $\eta=0.005$
Refer to caption: /html/2410.24206/assets/x54.png

#### D.1.3 Sorting

Figure: (a) $\eta=0.00125$
Refer to caption: /html/2410.24206/assets/x56.png

### D.2 RMSProp-Norm

#### D.2.1 CIFAR 10 (4 class)

Figure: (a) $\eta=0.005$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x59.png

Figure: (a) $\eta=0.005$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x62.png

Figure: (a) $\eta=0.005$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x65.png

Figure: (a) $\eta=0.005$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x68.png

Figure: (a) $\eta=0.005$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x71.png

Figure: (a) $\eta=0.005$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x74.png

### D.3 RMSProp

#### D.3.1 CIFAR 10 (4 class)

Figure: (a) $\eta=1\times 10^{-5}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x77.png

Figure: (a) $\eta=5\times 10^{-6}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x80.png

Figure: (a) $\eta=1\times 10^{-5}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x83.png

Figure: (a) $\eta=5\times 10^{-6}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x86.png

Figure: (a) $\eta=1\times 10^{-5}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x89.png

Figure: (a) $\eta=5\times 10^{-6}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x92.png

#### D.3.2 Sorting

Figure: (a) $\eta=1\times 10^{-5}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x95.png

Figure: (a) $\eta=1\times 10^{-5}$, $\beta_{2}=0.995$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x98.png

#### D.3.3 The effect of β 2 subscript 𝛽 2 \beta_{2}

Figure: (a) $\eta=2\times 10^{-5}$, $\beta_{2}=0.95$, $\epsilon=10^{-7}$, bias correction
Refer to caption: /html/2410.24206/assets/x101.png

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
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.
soundness: 3
strengths: 1. The article is clearly written and reader-friendly, with insightful discussions on the oscillatory behavior of optimization methods.

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

3. Extensive empirical results are presented, supporting the validity of the central flow.

Overall, I think this is a nice paper.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. Why is there a mismatch between Fig. 4 and Fig. 5, in terms of sharpness and $2/S$? 

2. How does the sharpness evolve (for a numerical example) when optimizing using full RMSProp? 

3. Do the authors think this approach can be used when there may not be an obvious choice of $\sigma$? 

4. Why are $x_t u_t$ not present in Eq. 2?
soundness: 3
strengths: 1. This work was well written and well structured, making it easy to follow along. 

2. The authors analyzed a wide range of neural networks, including modern architectures such as ViT.

3. The authors developed their approach in settings where the interpretation was straightforward to make their point more clear. 

4. The figures were well done and easy to read. 

5. The authors convincingly show that their central flow approach accurately captures the general behavior of neural network model training. 

6. 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.
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.
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? 

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

3. 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? 

4. 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? 

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

MINOR POINTS: 
1. PSD should be defined. 

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

3. Figure 1 should be referenced in the text. 

4. The figure in Sec. 4.3.1 is not referenced anywhere nor does it have a label.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: What are the geometric implications of the curvature penalty?
soundness: 3
strengths: The work presents novel contributions. Additionally, it represents an advancement in the study of locally analyzed optimizers and suggests progress for other optimizers.
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.
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.'
2. 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.
3. 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). 
4. 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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
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. 
2. 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?
3.  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?
soundness: 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.
2. 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. 
3. 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. 
4. 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. 
5. 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.
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.
weaknesses: Some comments
1. 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.
2. 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. 
3. 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.

Minor weakness
1. 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. 
2. Figure 3, the legend is difficult to read might be possible to increase the font a bit.  
3. 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.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
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?
soundness: 2
strengths: - The work is important.  Understanding the behaviors of optimization algorithms is very important, and this paper provides new insights through a novel approach.
- The paper is engaging.  It is well written and easy to read.
- Conclusions and key takeaways are clearly stated.
- The work is wide-ranging, and makes a number of important conclusions in a single paper.
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.

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

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

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

I have skimmed these appendices, but not closely checked the mathematics.

Meta-review: The paper introduces a novel framework for understanding the long-term dynamics of optimization in deep learning, particularly in the "edge of stability" regime. By modeling the time-averaged behavior of GD and RMSProp, the paper reveals some new insights, including implicit regularization mechanisms that steer optimization trajectories toward low-curvature regions. The claims are supported by empirical results and detailed derivations.

### **Strengths**
- Introduces a new tool to analyze optimization trajectories
- Offers insights into RMSProp's implicit behavior, such as regularization against sharpness and "acceleration via regularization."

### **Weaknesses**
- A significant portion of the theoretical derivations and experimental results are relegated to the appendices, making the main paper feel incomplete.
 The framework relies on informal reasoning (e.g., time-averaging steps and Taylor expansions), and some equations (e.g., Eq. 7) are not fully derived in the main text.

The reviewers are overall positive and I therefore recommend acceptance but I strongly encourage the authors to better integrate key derivations into the main text.

Decision: Accept (Poster)