On a continuous time model of gradient descent dynamics and instability in deep learning
Authors: Mihaela Rosca
Yan Wu
Chongli Qin
Benoit Dherin
Venue: Accepted by TMLR ()

Abstract: The recipe behind the success of deep learning has been the combination of neural networks and gradient-based optimization. Understanding the behavior of gradient descent however, and particularly its instability, has lagged behind its empirical success. To add to the theoretical tools available to study gradient descent we propose the principal flow (PF), a continuous time flow that approximates gradient descent dynamics. To our knowledge, the PF is the only continuous flow that captures the divergent and oscillatory behaviors of gradient descent, including escaping local minima and saddle points. Through its dependence on the eigendecomposition of the Hessian the PF sheds light on the recently observed edge of stability phenomena in deep learning. Using our new understanding of instability we propose a learning rate adaptation method which enables us to control the trade-off between training stability and test set evaluation performance.

Published in Transactions on Machine Learning Research (01/2023)
On a continuous time model of gradient descent dynamics and instability in deep learning
Mihacla Rosca
mihaelacr Sedtulaaarrdeeemminddd com com
DeepMind, University College London
Yan Wu Deep
yaruuudeeepmiinn com com
Chongli Qin
Atonnliggi@ddeepmiiinn
Deep
Benoit Dherin
dherin encerin@@googe.coom com
Google
Reviewed on OpenReview:
https: // openreview. net/ forum? id= EdEERPRKiPnn
Abstract
The recipe behind the success of deep learning has been the combination of neural net- works and gradient-based optimization. Understanding the behavior of gradient descent however, and particularly its instability, has lagged behind its empirical success. To add to the thcoretical tools available to study gradient descent we propose the principal flow (PF), continuous time flow that approximates gradient descent dynamics. To our knowledge, the PF is the only continuous flow that captures the divergent and oscillatory behaviors of gradient descent, including escaping local minima and saddle points. Through its depen- dence on the eigendecomposition of the Hessian the PF sheds light on the recently observed edge of stability phenomena in deep lcarning. Using our new understanding of instability we propose a learning rate adaptation method which enables us to control the trade-off between training stability and test set evaluation performance.
1 Introduction
Our goal is to use continuous time models to understand the behavior of gradient descent. Using continuous
dynamics to understand discrete time systems opens up tools from dynamical systems such as stability
analysis, and has long history in optimization and machine learning (Glendinning 1994 Saxe el al. 2013
Nagarajan and Kolter 2017: Lampinen and Ganguli, 2018 Arora et al. 2018] 8 Advani et al. 2020 El Elkabetz
and Cohen 2201;; Vardi and Shamir 2021 Franca et al. 2020 Barrett and Dherin 2021 Smith et al. 2021). Most theoretical analysis of gradient descent using continuous time systems uses the negative gradient flow, but this has well known limitations such as not being able to explain any behavior contingent on the learning rate. To mitigate these limitations we find a new continuous time flow which reveals important new roles of the Hessian in gradient descent training. To do so, we use backward error analysis (BEA), method with a long history in the numerical integration community (Hairer et al.] 2006) that has only recently been used in the deep learning context [Barrett and Dherin 2021 Smith et al. 2021
We find that the proposed flow sheds new light on gradient descent stability, including but not limited to divergent and oscillatory behavior around fixed point. Instability areas of training where the loss consistently increases and edge of stability behaviors (Cohen et al. 2021 arras of training where the loss does not behave monotonically but decreases over long time periods are pervasive in deep learning and occur for all lcarning rates and architectures Cohen et al. (2021) ; Gur-Ari et al. (2018); Gilmer et al. b2 Lewkowyez et al. 2020). We use our novel insights to understand and mitigate these instabilities.
The structure of the presented work is follows:

---

Published in Transactions on Machine Learning Research (01/2023)
We discuss the advantages of continuous time approach in Section where also highlight the limitations of existing continuous lime flows.
We introduce the principal flow (the PF), flow in complex space defined by the cigendecompo- sition of the Hessian (Section Bp. To our knowledge the PF the first continuous time flow that captures that gradient descent can diverge around local minima and saddle points. We show that using complex flow is crucial in understanding instabilities in gradient descent.
We show the PF is better than existing flows modelling neural network training dynamics in Section | In Section [ we use the PF to shed new light on edge of stability behaviors in deep learning. We do so by connecting changes the loss and Hessian eigenvalues with core quantities exposed by the PF and neural network landscapes explored through the behavior of gradient flows.
Through a continuous time perspective we demonstrate empirically how to control the trade-off between stability and performance in deep learning in Section 6 We do using DAL (Drift Adjusted Learning rate), an approach to sctting the learning rate dynamically based on insights on instability derived from the PF.
We end by showcasing the potential of integrating our continuous time approach with other opti- mization schemes and highlighting how the PF can be used as tool for existing continuous time analyses in Section
Notation: We denote as E the loss function, 0 the parameter vector of dimension D, V2 E the loss Hessian and X, the Hessian's i'th largest eigenvalue with u, the corresponding eigenvector. Since if u, is an eigenvector of V so -u, we always use u, such that Re E u,] 2 O; this has no effect on our results and is only used for convenience. For continuous time flow (hh) refers to the solution of the flow at time h.
Experiments: A list of figures with details on how to reproduce each of them is provided in the Appendix. Code available at https://g/t http://gtttttbbtee www www.mm/eeemindddiii..e dideemiinnndisceetisatiioooo drift
2 Continuous time models of gradient descent
The aim of this work is to understand the dynamics of gradient descent updates with learning rate h
0. - -V
(1)
-
from the perspective of continuous dynamies. When using continuous time dynamics to understand gradient descent it is most common to use the negative gradient flow (NGF)
0- -
(2)
Gradient descent can be obtained from the NGF through Euler numerical integration, with an error of O(h2) after one gradient descent step. Studying gradient descent and its behavior around equilibria and beyond has thus taken two main approaches: directly studying the discrete updates of Eq (Bartlett al.
2018a by Mescheder et al.] 2017 Gunasekar et al.] 2018 Du et al.) 2019
Allen-Zhu et al. 2019 Du and Hup
2019
Ziyin ot al. 2021a] Liu et al.
2021), or the continuous time NGF of Eq 2 (Glendinning
1994 Saxc
et al. 2013 Nagarajan and Kolter
2017 Lampinen and Ganguli 2018
Arora al.) 2018 Advani et al.
2020
kkabetz and Cohen 2021
Vardi and Shamir 2021 Franca el. al. 2020: Balduzzi et. al. 2018). The
appeal of continuous time systems lies in their connection with dynamical systems and the plethora of tools that thus become available, such as stability analysis; the simplicity by which conserved quantities can be obtained (Du et al.) 2018 Franca et al. 2020); and analogies that can be constructed through similarities with physical systems (Franca et al.] 2020). Because of the availability tools for the analysis of continuous time systems, it has been previously noted that discrete time approaches are often more challenging and discrete time proofs are often inspired from contimuous time ones May 1976 Elkabetz and Cohen 2021) ). We use an example to showcase the case of contimuuu time analyses: when following the NGF the loss E
decreases since dE dET 20
F112 Showing that and when following the discrete time gradient
2

---

Published in Transactions on Machine Learning Research (01/2023)
33
0 = E
gradient descerd
NGF
global min
Init point
0.5
discretization drift
0.5
-14
gridlent descert
IGR flos
0111
-15
InE soint
0,
-1554
-3
-21
0.0
0. -2.
global man
-20
-10
-03
20
(a) Discretisation drift.
(b) 2D convex
Banana function.
Figure 1: Motivation. Using continuous time flows to understand gradient descent is limited by the gap between the discrete and continuous dynamics. In the case of the negative gradient flow, we call this gap discretization drift. Other flows have been introduced to capture part of the drift, but they also fail to capture the oscillatory unstable behavior of gradient descent.
descent update in Eq is more challenging and requires adapting the analysis on the form of the loss function E. Classical convergence guarantoes associated with other optimization approaches such as natural gradient are also derived in cont time (Amari 1998 Ollivier 2015a b). By analyzing the properties of continuous time systems one can also determine whet optimizers should more closely follow the underlying continuous time flow (Song et al. 2018 Qin et al.) 2020 what regularizers should be constructed
ensure
convergence stability (Nagarajan and Kolter 2017 Balduzzi et al. 2018 Rosca et al. 2021). construct
converge guarantees in functional space for infinitely wide networks (Jacol et al. 2018
Lee el al. 2019
2.1 Limitations of existing continuous time flows
The well-known discrepancy between Euler integration and the NGF, often called discretization error or discretization drift (Figure [(a]] leads certain limitations when using the NGF describe gradient descent, namely: the NGF cannot explain divergence around local minima for high learning rates convergence to flat minima as often seen in the training of neural networks. Critically, since the NGF does not depend on the learning rate, cannot explain any learning rate dependent behavior.
The appeal of continuous time methods together with the limitations of the NGF have inspired the machine learning community to look for other contimuus time systems which may better approximate the gradient descent. trajectory. One approach to constructing continuous time flows approximating gradient descent that takes into account the learning rate is backward error analysis (BEA). Using this approach, Barrett and Dherin 2021 introduce the Implicit Gradient Regularization flow (IGR flow):
h = - VVEVEE
(3)
which tracks the dynamics of the gradient descent step O, 0,-1 -Voo with an error of O thus reducing the order of the error compared to the NGF. Unlike the NGF flow, the IGR flow depends on the learning rate h. This dependence explains certain properties of gradient descent, such as avoiding trajectories with high gradient norm; the authors connect this behavior to convergence flat minima.
Like the NGF flow however, the IGR flow docs not explain the instabilities of gradient descent, as we illustrate in Figure [ Indeed, [Barrett and Dherin (2021) (their Remark 3.4) show that performing stability analysis around local minima using the IGR flow does not. lead to qualitatively different. conclusions from those using the NGF: both NGF and the IGR flow predict gradient descent to be always locally attractive around a local minimum (proofs Section A.4.1). contradicting the empirically observed behavior of gradient descent. To understand why both the NFG and the IGR flow cannot capture oscillations and divergence around local minimum, we note that stationary points E 0 are fixed points for both flows. We visualize an example in Figure((a)) since to go from the initial point to the gradient descent iterates requires passing through the local minimum, both flows would stop at the local minimum and never reach the following gradient descent
3

---

Published in Transactions on Machine Learning Research (01/2023)
3.0
loss
loss
2.5
NGF
2.5
principal (ours)
IGR flow
gradient descent
2.0
gradient descent
starting point
Re(c(())
starting point
Re(L(()]
1.5
global minimum
1.5
global minimum
1.0
0.5
0.5
0.0
-0.5
-1.0
-0.5
0.0
0.5
1.0
1.5
2.0
2.5
3.0
-1.0
-0.5
0.0
0.5
1.0
1.5
2.0
2.5
3.0
0
Re(0]
(a) Real flows.
(b) Complex flow.
Complex flows capture oscillations and divergence around local minima. In the real space, the trajectory going from the starting point to the second gradient descent iterate goes through the global minima, and real flows stop there, In complex space however, that need not be the case.
Figure 2:
iterates. In the case of neural networks we show in Figure 31 in the Appendix that while the IGR flow is better than the NGF at describing gradient descent, & substantial gap remains.
The lack of ability of existing continuous time flows to model instabilities empirically observed in gradient descent such those shown in Figure has been used motivation to use discrete-time methods instead (Yaida 2018 Liu et al. 2021). The goal our work is to overcome this issue by introducing a novel continuous time flow which captures instabilitics obscrved in gradient descent. To do so, we follow the footsteps of Barrett and Dherin (2021) and use Backward Error Analysis. By using a continuous time flow we can leverage the tools and ad vanatages of continuous time methods discussed earlier in this section; by incorporating discretization drift into our model of gradient descent we can increase their applicability explain unstable training behavior. Indeed, training instabilities; a key reason why is that, unlike cxisting flows, it operates in complex space. In Scction we show the importance of operating in complex space in order understand oscillatory and instability behaviors of gradient descent.
show in Figure [(b] that the flow propose captures the
2.2 Backward error analysis
Backward error analysis (BEA) is a tool in numerical analysis developed to understand the discretization error numerical integrators. We now present an overview of how use il in the context of gradient descent; for general overview see Hairer et al. (2006). BEA provides modified vector field:
F0) VoE h()
ha ((),
(4)
by finding functions Ji, fn such that the solution of the modified ODE at order n, that is,
-
(z)
(5)
follows the discrete dynamics of the gradient descent update with an error |0, - ((h)] of order (h*+2), where (hh) the solution of the modified equation truncated order at time h, with O(0) 0,-1- The full modified vector field with all orders (n oo)
f(0) = -V (((()) + + (n() +...,
(6)
is usually divergent and only forms an asymptotic expansion. What BEA provides is the Taylor expansion in of an unknown I-dependent, vector field (0) developed at h 0:
i(0) Taylor
(7)
4

---

Published in Transactions on Machine Learning Research (01/2023)
Thus strategy for finding fn is to find series of the form in Eq 6 via BEA and then find the function fr such that its Taylor expansion in h at O results in the found series. Using this approach we can find the flow ă = (0) which exactly describes the gradient descent step 0: 0 - (0,)))
While flows obtained using BEA are constructed to approximate one gradient descent step, the same flows can be used over multiple gradient descent steps as shown in Section A.9 in the Appendix.
BEA proofs. The general structure of BEA proofs is as follows: start with a Taylor expansion in h of the modified flow in Eq [] write each term in the Taylor expansion as a function of V and the desired fo (this often requires applying the chain rule repeatedly); group together terms of the same order in h in the expansion; and identify fa such that all terms of ((h)) are 0 for p 2 2, as is the case in the gradient descent update. A formal overview of BEA proofs can be found in Section A.1 in the Appendix.
Since we are only looking for the first correction term, we only need to lind Ji. We perform Taylor expansion to find the value of (h) up to order O(h)) and then identify fi from that expression such that the
We now exemplify how to use BEA to find the IGR flow (Eq 3 Barrett and Dherin 2021
error |00, - Oh))ll is of order O(h²). We have: O(h) 0:-1 + 11011))) (0.1))
(0,.)) 0(1"). We know by the definition of the modified vector field (Eq 5) that Ø(1) = =Vel + (0). We can then use the chain rule to obtain ((2) (0))+ -Vol + O(h) = -VgE do
+0(h) V + O(h). +0(h)) We can then write -0 dt
Thus (h) = 0 - (V + 12 (0. +
02 - (1) = 0,1 - VV - -001) ( (0000
After simplifying we obtain 0: - (()= (0. 1 VVF((-)))) 0 For the error to be of order O1h)) the terms of order (hh)) have to be 0. This entails f1 = - V leading to Eq 3
3 The principal flow
In the previous section we have seen how BEA can be used to define continuous time flows which capture the dynamics of gradient descent up to a certain order in learning rale. We have also explored the limitations of these flows, including the lack of ability to explain oscillations observed empirically when using gradient descent. To further expand our understanding of gradient descent via continuous time methods, we would like to get an intuition for the structure of higher order modified vector fields provided by BEA, We start with the following modified vector field, which we will call the third order flow (proof in Section A.2
0 -- -- - 1 - -VEEOEE 112 (LVDD)VO +
I V (VEE)V 12
The third order flow tracks the dynamics of the gradient descent step 02 = -V of Ohh)), thus further reducing the order of the error compared to the IGR flow. Like the IGR flow and the NGF, the third order flow has the property that 0 0 V = 0 and thus will exhibit the same limitations observed in Figure 2 The third order flow allows us to spot a pattern: the correction term of order O(h)) in the BEA modified flw describing gradient descent contains the term (V) V)" F and terms which contain higher order derivatives with respect to parameters, terms which we will denote as C(V))))
with an error
Our approach. We will use the terms of the form VV to construct new continuous time flow. We will take a three-step approach. First, for an arbitrary order O(h") we will find the terms containing only first and second order derivatives the modified vector field given by BEA and show they are of the form (V3 we will use the series to find the modified flow given by BEA (Theorem 3.2 All proofs are provided in Section A of the Appendix.
(Theorem 3.1.. Second, we will use all orders to create series (Corollary (.1p. Third,
Thcorem 3.1 The modified vector field with an error of order (hh+2)) to the gradient descent update 0- 11-0----- 1V has the form:
n
12 (V
(9)
++1
5

---

Published in Transactions on Machine Learning Research (01/2023)
where E(V denotes the family of functions which can be written sum of terms, each term containing a derivative of higher order than 3 with respect lo paramelers.
The result is proven by induction. The base cases for n 1,2 and 3 follow from the NGF, IGR and third order flows. For higher order terms, the proof uses induction to find the term fi depending on Vg and VoE only and follows the BEA proof structure highlighted in Section 2.2 but Step 3 is modified to not account for terms in E). From the above, we can obtain the following corollary by using all orders n and the eigen decomposition of Vy
Corollary 3.1 The full order modified flow obtained by performing BEA on gradient descent updates is of the form:
1P
V + +C(VE) = Σ
1P
Σ
VoE+
(10)
p+1
p=0
p=0
1=0
D
D
Σ
Σ
(VoE= u.ju; + ((V)))
(11)
p
1=0
where X; and u, are the respective eigenvalues and eigenvectors of the Hessian VaE.
If Ao > 1/h the BEA scries above diverges. Generally BEA scries are not convergent and approx- imate the discrete schome only by truncation (Hairer et al.) 2006). When the series in Eq [1 di- verges, truncating it up to any order n however will result in flow which will not be able to caplure instabilities, even in the quadratic case. ways predict the loss function will decrease for a quadratic loss where a minimum exists, since: Ep + h2 Eb (X)((o)"" which is never positive for any quadratic loss where minimum exists (i.e. when A; 2 O, Vi). The above also entails that the flows always predict convergence around a local minimum, which is not the case for gradient descent which can diverge for large learning rates.
Such flows (including the IGR flow) will al-
dE VoET 1 de
To further track instabilities can use the BEA series to formulate the following flow: Definition 3.1 We define the principal flow (PF) as
D-1 log (1 hdd))
V u u
(12)
hd:
1=0
We note that limx-so log(1-h2) -1 and thus the PF is well defined when the Hessian V is not invertible, hd Unlike the NGF and the IGR flow, the modified vector field of the PF cannot be always written the gradient of a loss function in R, and can be complex valued.
Theorem 3.2 The Taylor expansion in h at h = 0 of the PF vector field coincides with the series coming from the BEA of gradient descent (Eq 11.
Proof:
Using the Taylor expansion Taylor
log(1=2)
Ep=0 p=0 p4T
2
we obtain:
2
Taylor
log(1 hdde)
(VoET
u
D-1
Σ
Σ
E--11xx] (VoET
(13)
hd4
We have used BEA to find the flow that when Taylor expanded at h = 0 leads to the series in Eq 11 When the BEA series in Eq [I converges, namely Ag < 1/h, the PF and the flow given by the BEA series are the same. When Ao > 1/h however, the PF is complex and the BEA scries diverges. While in this case any BEA truncated flow will not be able to track gradient descent closcly, we show that for quadratic losses the PF will track gradient descent exactly, and that is a good model of gradient descent around fixed points. We show examples of the PF tracking gradient descent exactly in the quadratic case in Figures 2(b) and 5

---

Published in Transactions on Machine Learning Research (01/2023)
Negative Gradient Flow
IGR Flow
Principal Flow
EPo -(VLL" aver (hAs) --1
2im - -(1+ (V uu)uu oIGn(h).) = (1 >x)
0 2
-h-.)
(V
u,)
hd.
(pr(hh)) (g(()))
Ad.
Table 1: Undorstanding the differences between the flows discussed in terms of the cigendocomposition of the Hessian. All flows have the form 0 = ED- (V (V)b)u,)u with different a summarized here.
Remark 3.1 For quadratic losses of the form 0 AO + the PF captures gradient descent exactly.
This case has been proven in Hairer et al. 22066 The solution of the PF can also be computed evactly in terms of the eigenvalues of E:E: (tt)= Ing(1-Ax, 0 u + + D-
log(1-hhi)
hd:
Remark 3.2 In small enough neighborhood around critical point (where higher order derivatives can be ignored) the PF can be used describe gradient descent dynamics closely. We show this also using a linearization argument in Section A.5 in the Appendiz.
Definition 3.2 The terms C(V)E) are called non-principal terms.
The term 12 V E (V)E) in Eq is non-principal term (we will call this term non-principal third order term).
Definition 3.3 We define the principal flow with third order non principal term as
D-1
log(1 1 -h2,)
0
=
(V FT u.),
V
V
(14)
1=0 hds
12
third order non principal
General theoretical bounds on the error between contimuous time flows and gradient descent are challenging construct in the case of a general parametrised E(O) as the error will be determined by the shape of E. We know the conditions which determine when certain flows follow gradient descent exactly. The NGF and gradient descent will follow the same trajectory in areas where vgEV&E = 0 (see Theorem 6.1) and thus E' has a constant gradient in time, since dV&E = VoEVe under the NGF. The PF generalises the NGF, dt
= in that it follows the same trajectory as gradient descent not only for trajectories where V 0, but also when E is quadratic. Informally, we can state that the closer we are to these exact conditions, the more likely the flows are capture the dynamics of gradient. deecentt. Formally, bounds the error between GD and NGF can be provided by the Fundamental Theorem (Theorem 10.6 in Wanner and Hairer 1996)) which has recently been adapted to a neural network parametrisation by Elkabetz and Cohen (0221) ; this bound depends on the mdgnitude of the smallest Hessian eigenvalue along the NGF trajectory. We hope that future work can expand the Fundamental Theorem such that error bounds between the PF and gradient descent can be constructed for deep neural networks. Here we take an empirical approach and show that although not exact outside the quadratic case the PF captures key features of the gradient descent dynamics in stable unstable regions of training, around and outside critical points, for small examples or large neural networks.
3.1 The principal flow and the eigen decomposition of the Hessian
All flows considered here have the form form Ø = EP-- (Vo u,),, where a function computing the corresponding coellicient; we will denote the one associaled with each flow as aNGF, 0IGR and opF respectively. For a side-by-side comparison between the NGF, IGR flow and the PF as functions of the Hessian eigendecomposition see Table | Since VoETu, u, 2 0, the function determines the sign of a modified vector field in the direction u, For brevity it will be useful to define the coefficient of u, in the vector field of the PF:
Definition 3.4 We call se; = log(1-hAd hd, (V u;) = a u, the stability coefficient for eigendirection i. sign(sc;) = sign(arr(hx))).
In order to understand the PF and how it is different from the NGF we explore the change in cach cigendi- rection u, and we perform case analysis on the relative value of the cigenvalues A; and the learning rate h. To do so, we will compare (NGr( and apr (ht;) since the sign of (N determines the direction which minimises E given by u, Since our goal to understand the behavior of gradient descent, we perform
7

---

Published in Transactions on Machine Learning Research (01/2023)
( (Mgg(
Re(aPF(hA)] Re(aNGF(IAi]
10
8
2
10
-10
8
2
0
6
10
Au h
Ai h
(a) Real part.
(b) Imaginary part.
Figure 3: Comparing the coefficients aNGF and app across the training landscape.
E(2) -3
starting point solution trajectory
0.75
0.04
0.50
0.02
Imagiz)
0.1
Imagiz!
0.25
-0.02
0.0
starting point souttoon
-0.25
-0.50
-0.04
-0.1
trajectory
-0.75
0.10
0.15
0.20
0.00 Retz)
0.75
Relz)
Relz)
(a) λy 1/h, (Po() <
(b) 1/h 2/h, Re(apF(Aph))
2/h Ao, Re(apr(Aahh) >
Figure The behavior of PF on E(z) 1 with solution x(t) c When do < 1/h,
a(t) = (1 (0)) which is real space and converges to the equilibrium. When An > 1/h, =(t) = (h 1)th (cos( + isin sit(///))) 2(0). This cxhibits oscillatory behavior, and when Au 2/h, diverges.
the case by case analysis of what happens at the start of a gradient descent iteration and thus use real values for A; and u, even when the PF complex valued. We visualize aNGF and app in Figure 3 and we use Figure 4 to show examples of each case using simple function.
Real stable case: A; < 1/h. (ign(aNcr(hh)) = sign(apr(h>.)) -1.
hy. < 0. The coefficients of both the NGF and PF in cigendiroction u, are both negative and real. The case is exemplified in Figure [(a)
(hA,) - and apr (hA;) log(1-AX.)
Complex stable case: 1/h < A < 2/h. sign(aner (hA.)) = sign( Relarr = 1. apr(hAs) E C.
(NGF(hA) = -1 and (Pp((h.) log(1-AAi
log(- - C and Relapp(hA.)] logí log(-1+hA;) < 0. by
The real part of the coefficient of the NGF and PF in cigendirection u, are both negative. The imaginary part of app can still introduce instability and oscillations, as we show in Figure 4(b)
Unstable complex case: 2/h < Ai. sign(a (hA:)) sign( (e(app(())))) app(hdd) EC
UNGF(h)) -1 and apr (hA:) loR(1-AAd l C and Relapp(hAd)] loR(-14AA:) 0. hds
Aλ: The real part of the coefficient of the NGF in eigendirection u; is negative, while the real part of the coefficient
AA:
of the PF is positive. The PF goes in the opposite direction the NGF which minimises E; this change in sign can cause instabilities. The imaginary component can still introduce oscillations, however the larger Xyh, the smaller the imaginary part of app. We visualize this case in Figure 4(c)
The importance of the largest eigenvalue Ag- The largest eigenvalue Xn plays an important part in the PF. Since hAg hλg Vi, A determines where in the above cases the PF is situated and thus whether there
8

---

Published in Transactions on Machine Learning Research (01/2023)
are oscillations and unstable behavior in training. For all flows of the form we consider write:
dF((0) dFF(0)² do
D-1 dt d0 dt = VoET
D-1
V i.u = Σ a(hh))(VoE u.)
(15)
and thus if (hh))) E R and ahhx)) < 0 Vi then delt) di S 0 and following the corresponding flow minimises E. In the case of the PF this gets determined by λo If d * then app(hA.) < 0 Vi (real stable case above) and the PF minimises E. If 1/h An å then Re(avr((h)]]] 0 Vi (complex stable case above) close to gradient descent iteration A, E R we can write that dRe]E(0)] E Re(app(hAi)] Vo u.)) and thus the real part of the loss function decreases. If do 2 then Re(aps (hAo)] (unstable complex case above) and if (V un)² sufficiently large we can no longer ascertain the behavior of E. We present a discrete time argument for this observation in Section A.8.1
Building intuition. For quadratic objective E(() _ AO the PF describes gradient descent. exactly. We show examples Figuros _ and 5 Unlike the NGF the IGR flow, the PF captures the oscillatory and divergent behavior of gradient decent. Importantly, to capture the unstable behavior which occurs when expand intuition outside the quadratic case, we
Ao > 1/h the imaginary part the PF is needed. show the PF for the banana function (Rosenbrock) 1960) in Figure D and additional example in 1D with uon-quadratic function (Figure [30 in the Appendix). In this case, the PF no longer follows the gradient descent trajectory exactly, but we still observe the importance the PF capturing instabilities of gradient descent; we also obscrve that adding non-principal terms rostabilize the trajectory.
Remark 3.3 For the banana function, the principal terms have a destabilizing effect when h > 2/20 while the non principal terms can have a stabilizing effect.
0.5
gradient descont
gracient descert
gradiient descent
Pf
PF global min init point
PF
giobal min point
global min Init point
01
-11
-04
02
03
41
0.1
22
0.
8:
0g
02
(a) 1/h (stability)
(b) 1/h 2/h (oscillations)
(c) Ao 2/h (divergence)
Figure 5: Quadratic losses in 2 dimensions. The PF captures the behavior of gradient descent exactly for quadratic losses, including oscillatory behavior and divergence.
3.2 The stability analysis of the principal flow
We now perform stability analysis on the PF, to understand how it can be used to predict certain behaviors of gradient descent around critical points of the loss function E. Consider 0* such critical point, i.c veE(O*)= = 0. For a critical point 0* to be exponentially asymptotically attractive, all eigenvalues of the Jacobian evaluated at 0* need to have strictly negative real part.
The PF has the following Jacobian critical points (proof in Section A.4 the Appendix):
D-1 log(1 hA²) JpF(0")
u,
(16)
h
where X, up are the eigenvalues and eigenvectors of the Hessian V E(0*). We thus have that the eigenvalues of the Jacobian (rr(0") at the critical point 0r are I/r log(1 hAT) for 1 1,..., D.

---

Published in Transactions on Machine Learning Research (01/2023)
10
200
201
93
344
0.0
gradient descent
gradient descent
gradient descent
ICR flow
-05
IGR flow
IGR flow
init point
init point
order non ren-priridppe torm
-100
global min
-1.
global min
-10
global min
Inst point
-1.a
655
-12
30
-10
1.5
(a)
Ao < 1/h
(b) Ao 2/h(Ao 1.9/h)
(c) Ao >> 2/b(Ao 5/h)
Figure 6: Banana function. The PF can capture instability and the gradient descent trajectory over many iterations when A is close to 2/h. When Ao >> 2/h (right) the PF does not track the GD trajectory over many gradient descent steps, but when including non-principal term the flow is able to caplure the general trajectory of gradient descent and unstable behavior of gradient descent.
Local minima. Suppose that 0 a local minimum. Then all Hessian eigenvalues are non-negative X; We perform the stability analysis in cases given by the value of X; corresponding to the cases in Section 3.1
h < 1/2;- The corresponding eigenvalue of the Jacobian 1/ log(1 - hxz) is negative, since < 1 - hi2 < The principal vector field attractive in the corresponding eigenvector direction.
h E 11/X, 2/1)). The corresponding eigenvalue of the Jacobian 1 log(1 h\ī) т log(h>; 1) + i is complex, with negative real part since since hΛ; - 1 < 1. The principal vector field is attractive in the corresponding eigenvector direction.
h 2 2/ The corresponding cigenvalue of the Jacobian log(1 - hAT) = log(h>; - + is complex, with non-negative real part, since since hλ; - 2 1. The principal vector field is not attractive in the corresponding eigenvector direction, and if h > 2/ it is repelled in the corresponding eigenvector direction.
The last case tells us that the PF is not always attracted to local minima, as it is not attractive in eigendrec- tions where h 2 2/ Thus like gradient descent, the PF can be repelled around local minima for large learning rates. This is in contrast the NGF and the IGR flow, which always predict convergence around a local minimum: the cigenvalues of the NGF Jacobian are - and for the IGR flow the eigenvalues are A; //2 A;² both are negative when X; is positive. For derivations see Section A.4.1 in the Appendix.
A²
Remark 3.4 For quadratic losses, where the PF is exact, the results above recover the classical gradient descent result for quadratic losses namely that gradient descent convergences if do < 2/h, otherwise diverges.
Saddle points. Suppose that 0 is a strict saddle point. In this case there exists X; such that X; < O. We want to analyse the behavior of the PF in the direction of the corresponding eigenvector u, In that case, log(1 hAg) > 0 which entails that the PF is repelled in the eigendirections of strict saddle points. Note that this is also the case for the NGF since the corresponding eigenvalues of the Jacobian of the NGF would be -A3, also positive. Unlike the NGF however, the subspace of eigendirections that the PF is repelled by can be larger since it includes also eigendirections where X 2/h
4 Predicting neural network gradient descent dynamics with the principal flow
Computing the PF on large neural networks during training is computationally prohibitive, as it requires finding all eigenvalues of the Hessian matrix once for each step of the flow simulation, corresponding to many eigen-decompositions per gradient descent step. To build intuition about the PF for neural networks, we start with a small MLP for a 2 dimensional input regression problem, with random inputs and labels. Here we can understand the behavior of the PF since we can compute its modified vector field exactly and compare it with the behavior of gradient descent. We show rosults in Figure 7 where visualize the norm of the difference between gradient descent parameters at each iteration and the parameters produced by the continuous time flows we compare with. We observe that short lerm the principal flow is better than all other flows at bracking the behavior of gradient descent. As the number of iterations increases however,
10

---

Published in Transactions on Machine Learning Research (01/2023)
the PF accumulates error in the case of X 2/h; this is likely due to the fact that while gradient descent parameters are real, this is not the case for the PF, as discussed in Remark 4.1 Since we are primarily concerned with using the PF to understand gradient descent for small number iterations this will be less of concern in our experimental settings. Additional results which confirm the PF is better than the other flows at tracking gradient descent on a bigger network trained the UCI breast cancer dataset (Asuncion and Newman (007)) are shown in Figure 32 in the Appendix.
Remark 4.1 On the multiple iteration behavior of the PF. We note that while gradient descent paramelers are real for any ileralion 0,, Oz+1, Octn when approzimale the behavior of gradienl descent by initializing (()) = 0, and ranning the PF for time nh, there is nothing enforcing that ((h), ((nh) will be real when the PF is complex valued (Ao 1/h). We also note that in that case the symmetric Hessian is not Hermitian and the eigenvalues and eigenvector of the Hession will not be real; furthermore, the eigenvectors need not form a bosi For long term trajectories (larger n),, this can have an effect on long term error between gradient descent and PF trajectories, through an accumulating effect of the imaginary part in the PF. This can be mitigated by using the PF to understand the short term behavior of gradient descent (small n).
1 {
2 <
0.014
NGE
NSS
dl wt 29
0.02
8 comerll with
0.12
KGR Now
Third order flow
0.8
FP
0.015-
0.6-
8.88-
trajectoey
0.028
0.06
cada 1r
0.034
[
0.2
0.032
0.22
0.004
0.0
2
(a) Ag 1/h
(b) 1/h < do < 2/h
(c) Ao >> 2/h
Figure 7: Error between gradient descent paramcters and parameters obtained following continuous time flows for multiple iterations: | (nh)]l with 0(0) = 0y. For small n, the PF is better at capturing the behavior of gradient descent across all cases.
4.1 Predicting V un using the principal flow
For large neural networks, instead of simulating the PF describing how the entire parameter vector changes in time we can use the PF to approximate changes in a scalar quantity only. This will allow us to compare the predictions of the PF against the predictions of the NGF and IGR flow on realistic settings. To do so, we first have to compute how the gradient changes in time:
Corollary 4.1 If 0 follows the PF, then: (VoE) = E.- log(1 hAd)
(VoET ET u u.),:
h
This follows from applying the chain rule and using the definition the PF. We contrast this with how the gradient evolves the parameters follow the NGF:
Corollary 4.2 If 0 follows the NGF, then: (voE)= = DD- - -V,
Corollary 4.3 Tj 0 follows the IGR flow, then: (VoF) = Σ (A: + 2X2 (VoETu.)u
We would like to use the above to assess how V u; changes in time under the above flows and check their predictions empirically against results obtained when training neural networks with gradient descent. Since
"To avoid the concern around the eigenvectors of the Hessian no longer forming a basis, one can use the Jordan normal form instead, as we show in Section A. We don't take this approach here as most of our following analysis is not affected, and is concerned with the behaviour of the PF around one gradient descent iteration. Furthermore, support of the Jordan normal form code libraries is limited (especially for complex matrices), and we did not find this experiments where we simulate the PF outside the quadratic case for few iterations. We note, however, that mathematical analysis of long-term PF trajectories for general functions should use the Jordan normal form.
a significant issue in the
11

---

Published in Transactions on Machine Learning Research (01/2023)
CIFAR-10, Learning
CIFAR-10, Learning rate
CIFAR-10, Learning
perfect prediction
perfect prediction
perfect prodiction
3-
prediction
prodiction
prediction
NGF prediction
NGF prediction
NGF prediction
IGR flow prediction
IGR flow prediction
IGR flow prediction
V,ETuo prediction
V.E uo pred prediction citiiiion
V.ETun prediction
0
-2-
-3
-2
2 3
-3
-2
2
-3
RyE Us
7vE Un
TyE uu
Figure 8: Predictions of VoET Up according to the NGF, IGR flow and the PF. On the a axis we plot the value of VeL un as measured empirically in training, and on the y axis we plot the correspouding prediction according to the flows from the value of the dot product at the previous itcration. The 'exact match' line indicates a perfect prediction, the upper bound of performance. The PF performs best from all the compared flows, however for higher learning rates its performance degrades when VeE" un is large; this is due to the fact that the higher the learning rate and the higher the gradient norm, the more likely it is that the additional assumption we used that X; do not change does not hold.
is eigenvector of the Hessian it also changes in time according to the changes given by the corresponding flow, making (Vol?'u,) difficult to calculate. Even when if we wrote an exact flow for (VoETu,) u.), it would be computationally challenging to simulate since finding the new values of u, would depend on the full Hessian and would lead the same computational issues we are trying avoid in the case of large neural networks. In order to mitigale these concerns, we will make the additional approximation that X and u, do not change inside an iteration which will allow us to approximate changes to V F u, and compare them against empirical obscrvations. We note that we will not use this approximation for any other results.
Remark 4.2 If we assume that X; u, do not change between iterations, if 0 follows the PF then (VoETu)) = log(1- 1-;]1 nd:) V u,
Remark 4.3 If we assume that Ai, u, do not change between iterations, if follows the NGF we can write (VoETu) u;) = o u,
Remark 4.4 If we assume that do u, do not change between iterations, if 0 follows the IGR flow we can write (VoETu,) = (>: + x2) V u,
The above flows have the form - cx, with solution xxt)= x0))e We can thus test those solutions empirically by training neural networks with gradient descent with learning rate h and at cach step compute Eo (u,)-1 and compare it with the prediction zhh) obtained from the solution from each flow initial- ized at the previous iteration, i.e. x0))= Vo ((
We show results with a VGG model trained on CIFAR-10 in Figure 8 The results show that the PF is substantially better than the NGF and IGR flow at predicting the behavior of V ug. Since the NGF and the IGR flow solutions scale the initial value by the inverse of an exponential of magnitude given by do for large Ao this leads to small prediction, which is not aligned with what is observed empirically. We also note that the higher the value of VoET ug the worse the prediction of the PF; these are the areas where the approximations made in the above remarks are likely not to hold due to large gradient. norms.
4.2 Around critical points: escaping sharp local minima and saddles
The stability analysis we performed in Section 3.2 showed the PF is repelled by local minima where Xj > 2/h: that is, even if the model is close to sharp local minima (with \i > 2/h), that local minima. will not be attractive and training will continue until a shallow minima is reached. We provide experimental evidence to support that hypothesis in the context of neural networks in Figure 33 in the Appendix; these results are consistent with observations in the deep learning literature (Jastrzebski al. 2018
Cohen et al. 2021
12

---

Published in Transactions on Machine Learning Research (01/202)
CIFAR-
VSG CIFAR 10
Aq whin the less increases between Iar saratios
Learning rate
Learning rate
1.0
0.005
0.005
0.05
05
Ap 27
0.01
0.1
0.5
0.05
Training toos
0.4
100
0.2
0-
0.50
1775
0.00
0.25
2.00
1:25
1.50
1:75
2.02
0.0
0.89
0.24
0,35
0.26
0.577
0.28
0.29
01
Iteration
Iteration
Figure 9: Edge of stability in neural networks (Cohen et al. 20210: instability occurs when Ao > 2/h.
Furthermore, while saddle points have long been considered challenge with high dimensional optimisa- tion (Dauphin et al., 2014) in practice gradient descent has not been observed to converge to saddles (Lee et al. (0166). Our analysis suggests that saddles will be repelled not only in the direction of strictly negative eigenvalues, but also in the eigendirections with large posítive eigenvalues when large learning rates are used; this can explain why neural networks do not converge to non-strict saddles which exist in deep neural landscapes (Kawaguchi 2016) but need not be repelling for the NGF and IGR flow (existing analyses of
escaping saddle points by gradient descent apply only to striet saddles (Du et al. 2017 Lee al. 2016)
5
The principal flow, stability coefficients and edge of stability results
Edge of stability results. Cohen et al. (2021) did a thorough empirical study show that when training deep neural networks with full batch gradient descent the largest eigenvalue of the Hessian, Åo keeps growing until reaching approximately 2/h (a phase of training they call progressive sharpening), after which it remains in that area; for mean squared losses this continues indefinitely while for cross entropy losses they show it further decreases later in training. They also show that instabilities in training occur when λo > 2/h. Their empirical study spans neural architectures, data modalitios and loss functions. We visualize the edge of stability behavior they observe in Figure [ since we use a cross entropy loss A decreases later in training. We also visualize that iterations where the loss increases compared to the previous ileralion overwhelmingly occur when Ag > 2/h. Cohen et al./ (2021) also empirically observe that H un has oscillatory behavior in the edge of stability area but is 0 or small outside it.
Continuous-time models of gradient descent at edge of stability, To investigate if existing continuous time flows and the PF capture gradient descent behavior at the edge of stability we train a 5 layer MLP on the toy UCI Tris dataset (Asuncion and Newman 2007); this simple setting allows for the computation of the full eigenspectrum of the Hessian. We show results in Figure 10 the NGF and IGR flow have a larger error compared to the PF when predicting the parameters at the next gradient descent iteration in the edge of stability regime; the NGF and IGR flow predict the loss will decrease, while the PF captures the loss increase observed when following gradient descent. As we remarked in Section [ the NGF and the IGR flow do not capture instabilities when the eigenvalues of the Hessian are positive, which has been remarked to be largely the case for neural network training through empirical studies (Sagun et al., 2017]; Ghorbani et al,, 2019 Papyan 2018) and we observe here (Figure 34 in the Appendix). We spend the rest of the section using the PF to understand and model edge of stability phenomena using a continuous time approach.
Connection with the principal flow: stability coefficients. The PF captures the key quantities observed in the edge of stability phenomenon: the eigenvalues of the Hessian A; and the threshold 2/h. These quantities appear in the PF via the stability coefficient sc; log( IoR(1- V ET u, app(>x)) E u: hA of eigendirection u, Through the PF, by connecting the case analysis in Section 3.1 with existing and new empirical observations, we can shed light on the edge of stability behavior in deep learning.
First phase of training (progressive sharpening): Ao < 2/hh. This entails Re[scy] = Relapr(h>))) < 0, Vi (Real stable and complex stable cases of the analysis in Section 3.1 sign(aNar) = sign(apF) = -1 and following the PF minimises E or its real part (Eq 15). To understand the behavior of Ao; we now have to make use of empirical observations about the behavior of the NGF early in the training of neural networks. Tt has been
13

---

Published in Transactions on Machine Learning Research (01/2023)
bradient DOCDRY
0.033
6.002
Prinibal
5.000
-0.002
remm
0.035
5.0015
5.0000
21-
-2.008
NGP
I2R
0.0000
0,012
225
220
Iberation
250
260
Figure 10: Comparing different continuous time models of gradient descent at the edge of stability area small 5 layer MLP, with 10 units per layer. We show the local parameter prediction error ||0, 0h:; 0,-1)l for the NGF, IGR and PF flows (middle), as well as EhO 0.- 0.- E((,) (right).
-P,E
GD
40
2 35
100
30
25
-V6F
GD
2/h
50
GD
2.0
0,5
L055
0.0
Iteration
Reisig -aa
0.10
-0.2
0.05
-033
0.00
1.0
15 Rteraion
20
25
60
61 62
63
64 65 66 67 68 69 70 72
73
76
78
Iteration
(a) Early training
(b) Edge stability.
Figure 11: Understanding the edge of slability results using the PF on a layer MLP: we plot the behavior of the NGF A = F. and the positive gradient flow 0 = Vo initialized at each gradient descent iteration parameters, and see that the behavior of gradient descent is connected to the behavior of the respective flow through the stability coefficient. Figure [I) shows that even when λo > 2/h, if the real port of the slabilily coefficient sen is negalive close there no instabililies in the loss and the eigenvalue Xa keeps increasing, as it does when following the NGF in that region.
empirically observed that in early areas of training, Åo increases here when following the NGF (Cohen et al. 2021); we further show this in Figure 55 the Appendix. Since this part of training gradient descent follows closely the NGF, it exhibits similar behavior and An increases. We show this case in Figure [1
Second phose of training (edge of stability) An 2 2/h. This entails Re(scu = Re(app(hA)) 2 0. (Unstable complex case of the analysis Section 3.1 We can no longer say that following the PF minimizes E. sign Re(eReaaapp Rerrr (hAp)) since aNGr(h)))
sa (h)o)) that in that direction gradicnt descent resembles the positive gradient flow Ô E rather than the NGF. The positive gradient flow component can cause instabilities, and the strength of the instabilities depends on the stability coefficient sco = Vo un- We show in Figures [1(b) and 13 how the behavior of the loss and An affected by the behavior of the positive gradient flow when Ag 22h h.
-1 and sign Re (((((((())))))
0 meaning
14

---

Published in Transactions on Machine Learning Research (01/2023)
MNIST layers MLP
CIFAR-10 VGG
1.3
loss increasing
loss incre accressing
1.2
Loss
1.0
1.9
0.9
0.4
Re(scal
1.0
Retscal
0.5
0.0
1/
WV
W1
2.0
50
100
150
200
250
225
250
275
300 Iteration
325
350
375
400
Iteration
Figure 12: The loss function and stability coefficients: areas where the loss increases correspond to areas where the sep is large. The highlighted areas correspond to regions where the loss increases.
CIFAR-10, Resnet-18
200
-P
150
2 100
50
0.5
0.0
0.0
-0.2
-0.6
100
400
500
Iteration
Figure 13: Loss instabilities, As and stabilitiy coefficients for CIFAR-10. Together with the behavior of gradient descent, we plot the behavior of the NGF and positive gradient. flow initialized at O, and simulated for time h for each iteration t. The analysis performed based the PF suggests that when Re > 0 and large we should expected gradient descent to exhibit behaviors close to those of the positive gradient flow. What we observe cmpirically that increases loss value of gradient descent are proportional to the increase of the positive gradient flow in that area (can be seen best between iterations 200 and 350); the same behavior can be seen in relation to the eigenvalue Ag-
More than Ag: the importance of stability coefficients. While the sign of the real part of the stability coefficient sco is determined by Ao its magnitude is modulated by the dot product V Vo ug, since sCo (hpr(h)))V V E up. The magnitude of VoET un plays an important role, since if An is the only eigenvalue greater than 2/h training is stable VoE un = 0, as we observe in Figure 11 To undersland inslabililies, we have to look al slabilily coesficients, not only eigenvalues. We show in Figure 12 how the instabilities in training can be related with the stability coefficient scy: the increases in loss occur when the corresponding Re[sca] is positive and large. In Figure 13 we show results with the behavior of Ag: Xg increases or decreases based the behavior of the corresponding flow and the strength of the stability cocfficient and that gets reflected in instabilities in the loss function; specifically when Xn > 2/h, we use the posilive gradient flow and see how the strength its fluctnations alTect the changes both the loss value and do of gradient descent. We show additional results in Figures 38 and 39 in the Appendix.
15

---

Published in Transactions on Machine Learning Research (01/2023)
VGG CIFAR-10 Full batch
VGG CIFAR-10 Full batch
GD (LR 0.05)
GD (LR 0.05)
NGF
10*
NGF
Training loss
2.5
NGF with UyETup sign swap
NGF with VyETu, sign swap
10a
1.5
107
0.5
10
12
2
12 14.
Physical Time
Physical Time
Figure 14: One cigondirection is sufficient to lead to instabilitics. To create a situation similar to that of the PF, we construct a flow given by the NGF in all eigendirections but un: in the direction of ug we change the sign of the flow. This leads to the flow i = Fo unll + EP-1 - We show this flow can be very unstable when initialised in an odge of stability arca.
Is one eigendirection enough to cause instability? One question that arises from the PF is whether the leading eigendirection un can be sufficient to cause instabilities, especially in the context of deep networks with millions of parameters. To assess this we train a model with gradient descent until it reaches the edge of un + + stability (Ao ss 2/h), after which we simulate the continuous flow 0 = uouu coefficients of the modified vector field of this flow are negative for all cigendirections except from up which is positive; this is also the case for the PF when λ is the only cigenvalue greater than 2/h. In Figure 14 we empirically show that a positive coefficient for ug can be responsible for an increase in loss value and a significant. change in Ag in neural network training.
-D- -Ve L u,u, The
Decreasing the learning rate. Cohen et al. 2021 show that if the edge of stability behavior is reached and the learning rate is decreased, the training stabilizes and Ag keeps increasing (Figure 35 in the Appendix). The PF tells us that decreasing the learning rate entails going from Re[scy] 2 Re(scu] S since Àg < 2/h after the learning rate change. Since all stability coefficients are now negative, this reduces instability. The increase in An is likely due to the behavior of the NGF in that area (as can be seen in Figure 14 when changing from gradient descent training to the NGF in an edge of stability area leads to an increase of Aa).
The behavior of VoE un. The PF also allows us to explain the unstable behavior of VoE un around
edge of stability areas. As done in Section 4.1 we assume that X; u, do not change substantially be-
tween itcrations and write VoËTu, = log(1 1gg(1-RX.. RA.) VoETu, u; under the PF, with solution (VoETu,)() (;)()) =
A
(VoETu, inel,, This solution has different behavior depending on the value of do relative to 2/h: decrcasing below 2/h and increasing above 2/h. We show this thcoretically prodicted behavior in Figure 15 alongside empirical behavior showcasing the fluctuation of Va ug in the edge of stability area, which confirms the theoretical prediction. We also compute the prediction error of the proposed flow and show it can capture the dynamics of V un closely in this setting. We present a discrete time argument for this observation in Section A.8.2 We note that the stable behavior early in training together with the oscillatory behavior of V un in the edge of stability area which we predict and observe can explain the results of Cohen et al. 2021 on the behavior of 8T ug, since O accumulates changes given by gradient updates.
Why not more instability? To determine why there isn't more instability in the edge of stability area we have to consider that neural networks are not quadratic, which has two effects. Firstly, when following the PF the landscape changes slightly locally; this leads to changes in stability coefficients and thus the behavior of gradient descent as we have consistently seen in the experiments in this section. Secondly, non-principal terms can have effect; while we do not know all non-principal terms in Section B in the Appendix we provide a justification for why the non-principal term we do know (Eq 14 can have a stabilizing effect by inducing a regularisation pressure to minimise (V )2 in certain parts of the training landscape.
In this section we have shown the PF closely predicts the behavior of gradient descent in neural network training. This has led to additional insights, including the importance of stability coefficients in determining causally showing one eigendirection is sufficient to cause
instabilities in gradient descent. (Figures 11 12
13
16

---

Published in Transactions on Machine Learning Research (01/2023)
0.6
0,030
> 2/h
A
2/h
VGG CIFAR-1 Full batch, Learning 0.21
0.025
Integration ertor
imag/V 5i
0.2
0,029
0.0
0.5
0.015
-0.2
0.010
-0.4
0.005
-0.6
-0.6
-0.4
-0.2
0.2 Re(V,ETul
0.0
0.4
0.6
0.00
0.5
1.0 1.5
2.0
25
0.04 Leaamnng raite
0.10
Doratiion
Figure 15: Predicting the unstable dynamics of VoLTu u in the edge of stability area. (A R 2/h) using the PF. Left: the predicted behavior of Vo Vo under Vo u, log(1-A>)
h V E² u, with an inflection point at À = 2/h. Middle: empirical behavior of Ve u for a model shows instabilities in the edge of stability area (highlighted). Right: the approximation made to derive the flow is suitable around Rz 2/h.
instability (Figure 14h and change and being able to closely predict the behavior the dot product between the gradient and the largest eigenvector (Figure 15). This evidence suggests that the PF captures significant aspects of the behavior of gradient descent in deep learning; this is likely due to the specific structure of neural network models. While we take a continuous time approach, a discrete time approach can be used to motivate some of our observations (Section A.8); this is complementary to our approach but nonetheless related, since it also does not account for higher order derivatives of the loss and further suggests the strength of a quadratic approximation of the loss in the case of neural networks, as observed by Cohen et al. [022
6 Stabilizing training by adjusting discretization drift
The PF allows us to understand not only how gradient descent differs from the trajectory given by the NGF, but also when they follow each other very closely. Understanding when gradient descent behaves like the NGF flow reveals when the existing analyses of gradient descent using the NGF discussed Section 2 are valid. It also has practical implications, since in areas where gradient descent follows the NGF closely training can be sped up by increasing the learning rate. Prior works have empirically observed that gradient descent follows the NGF early in neural network training (Cohen et al.] (021) and this observation can be used to explain why decaying learning rates (Loshchilov and Hutter 2016) learning rate warm up (He ltt al. 2019) are successful when training neural networks: having a high learning rate in areas where the drift is small will not cause instabilities and can speed up training while decaying the learning rate avoids instabilities later in training when the drift is larger.
6.1 ViEVE determines discretization drift
In previous sections we have seen that the Hessian plays an important role in defining the PF and in training instabilities. We now want to quantify the difference between the NGF and the PF in order to understand when the NGF can be used as model of gradient descent. We find that: Remark 6.1 In region of the space where ViEVoL = the PF is the same as the NGF.
To see why, we can expand
D-1
VIEVOF= = Σ 1=0
X V. u.u.
(17)
II VEEVEE we have that X, V V = 0,Vi E {1,.., D}, thus either X, =0 leading to aNGF = (pF()))) = -1 or VoET 0. Then 6- ΣP- (pF(h)) (VeTu,
u.)u; = -EP-1 aNcr(hA.) u;),; Thus comparing the PF with the NGF reveals an important quantity: VEEVE.. Further investigating this quantity reveals it has connection with the total drift, since:
17

---

Published in Transactions on Machine Learning Research (01/2023)
MLP batch MNIST
CIFAR-10
VGG Batch 1024 CIFAR-10
0.5
drift
0.10
drift
drift
11 VEEV
115 ViEVEEII
0.15
11 ViEV,
0.08
0.3
0.06
0.10
0.04
0.05
0.02
0.0
0.00
0.00
200
400
600
800
1000
0
500
1000
1500
2000
0
500
1000
1500
2000
lteration
Iteration
Iteration
Figure 16: Connection between
V E and the per iteration drift mcasured during training.
MLP Full batch MNIST
VGG Full batch CIFAR-10
VGG Batch size 1024 CIFAR-10
Spearman Cornelation
Spearman Correlation
Spearman Correlation
1.0
0.8
0.8
0.8
0.6
0.6
0.6
CC(drift ([]EVEEII)
CCidrif, [[R,]])
0.
0.2
0.2
CCidrit, [|CEVEE))
0.00
0.02
0.04
0.06
0.08
0.10
0.00
0.02
0.04
0.06
0.08
0.10
0.00
0.02
0.04
0.06
0.08
0.10
Learning
Learning
Learning
Figure 17: Correlation between (IVEE Vel we plot the correlation with the individual terms as well.
V and the per iteration drift. Since
||V
Theorem 6.1 The discretization drift (error between gradient descent and the NGF) after 1 iileration 0, = -Vo- is VVE(0')V for set of parameters 0' in the neighborhood of 0.-1-
This follows from the Taylor reminder theorem in mean value form (proof in Section A.10 This leads Corollary 6.1 In region of space where oVE 0 gradient descent follows the NGF.
Thus the PF revealed ViEVEL as a core quantity in the discretisation drift of gradient descent. To further see the connection between with the PF consider that ||V = - V.V
EPJ I uull;; the higher each term in the sum, the higher the difference between the NGF and the PF. To mcasure the connoction between per iteration drift and oV in neural network training we approximate it via 0. - NGP h)] where NGF is the numerical approximation to the NGF initialised
at 0--1. Results in Figures 16 and 17 show the strong correlation between per iteration drift and V
throughout training and across learning rates. Since Theorem 6.1 tells us the form of the drift but not the exact value of 0', we have used 0,-1 instead to evaluate VoLVEL and thus some error exists.
Understanding this connection advantageous since computing discrotization drift computationally ex- pensive as requires simulating the continuous time NGF but computing V via Hessian-vectoo products is cheaper and approximations are available, such VEEVEE Rs VoE(O+.V.E) (.E(O) which only requires an idnilional backward pass Geiping el al. (2021.
6.2 Drift adjusted learning rate (DAL)
A natural question to ask is how use the correlation between V and the itcration drift improve training stability; V captures all the quantities we have shown to be relevant to instability highlighted by the PF: A and Vo u, (Eq. 17). One way to use this information is to adapt the learning rate of the gradient descent update, such as using [ as the learning rate. This learning rate slows
down training when the drift is large areas where instabilities are likely to occur and speeds up training in regions of low drift areas where instabilities unlikely to occur. Computing the norm of
18

---

Published in Transactions on Machine Learning Research (01/2023)
Resnet-50 50, Batch size 2048, Imagenet
VGG Full batch CIFAR-10
Resnet-1B 18 Full batch CIFAR, 10
7.0
2.5
Learning rate
Learning rate
Learning rale
7.5
0:01
0.00s
2.5-
0.005
0.5
6.1
Tra loss.
0.05
0.1
0.5
DAL
0.05
DAL
2.5
3.5
3.0
1.0
6.0
250
500
750
1003 Beewiton
1253
1500
1753
2000
0.2
200
435
660
100
1002
10000
20203
30000
40000
50000
Ihenatio
Iteration
Figure 18: DAL: using the learning rate
[ results in improved stability without requiring a hyper-
parameter sweep.
1e Full batch CIFAR-10
Resnet batch CIFAR-10
Resnet-18 18 Full batch CIFAR 10
Learning rate
A.005
0.5
107
102
10"
Lanihrrtt
Upddte 92m
2014
30
10--
Learning rate
0.005
k0-
Learning rate
10.
10-
0.005
202
425
800
1000
202
402
620
820
1000
Itaration
Itrationn
Figure 19: Key quantities in DAL versus fixed learning rate training: learning rate, and update norms.
the update provided by this learning rate shows a challenge however since 2/ ViEVo 2 to & this implies that when using this learning rate the norm of the gradient descent update will never be 0 and thus training will not result in convergence. Furthermore, the magnitude of the parameter update will be independent of the gradient norm. To reinstate the gradient norm, we propose using the learning rate
2
2
h())
(18)
V Ell E
where g(0) is the unit normalised gradient Fg F|. We will call this learning rate DAL (Drift Adjusted Learning rate). As shown in Figure 16 [g)) has a strong correlation with the per iteration drift. Another interpretation of DAL can be provided through signal to noise perspective: the size of the learning signal obtained by minimising E is that of the update h VoEll, while the norm of the noise coming from the drift be approximated as 1²2 oV thus the 'signal to noise ratio' can be approximated A VVl /(4 VVEVe 2/ ()E)(0)))) [], which when using DAL (Eq 18 is 1; thus DAL can be seen as balancing the gradient signal and the regularising drift noise in gradient descent training.
We use DAL to set the learning rate and show results across architectures, models and datasets in Figures 18 (with additional results in Figure 42 in the Appendix). Despite not requiring a learning rale sweep, DAL is stable compared to using fixed learning rates. To provide intuition about DAL, we show the learning rate and the update norm in Figure [9]] for DAL the learning rate decreases in training after which it slowly increases when reaching areas with low drift. Compared to larger learning static learning rates where the update norm can increase in the edge of stability area with DAL the update norm steadily decreases in training.
6.3 The off between stability and performance
Since are interested in understanding the optimisation dynamics of gradient descent, we have so far locused on training performance. We now try to move our atlention to Lest performance and generalization. Previous works [T et al., 2019 Barrett and Dherin 2021; Jastrzebski et al., 2019) have shown that higher learning rates lead to better generalization performance. We now try to further connect this information with the per iteration drift and the PF. To do so, we use learning rates with various degrees of sensitivity to itcration drift using DAL-p:
2
hp) (IV3E
(19)
19

---

Published in Transactions on Machine Learning Research (01/2023)
VGG Full Batch CIFAR-10
Resnet-18 Full batch CIFAR-10
Imagenet Resnet-50 Batch Size 2048
Train loss.
20
0.05
0.01
Learning
0.5
3.5
5.0
0.01
0.005
Leamiing rote
0.5
0.1
0.08
Learnnng
DAL-0.5
0.1
3.5
0.05
DA1-0.5
3.0
63-
00
Test accurary
60
64-
Tep
13050
23083
3058s
40000
Iteration
Beration
Iteraiion
Figure 20: DAT.0.5: increased training speed and generalization compared to a sweep of fixed learning rates.
1.2
VGG Full batch CIFAR-10
1.0
Resnet-18 Full batch CIFAR-10
10
Resnet-50 Batch size 8192 Imagenet
p DAL-p
0.2
0.0
0.25
0.75
Teat arcumccy
0.6
0.6
0.4
P DAL-P
in DAL-p
0.125
0.75
1.5
Test
0.125
1.0
3.0
0.2
0.0
500
750
1900
1250
1580
1750
2002
0.0
250
1050
1250
3500
1750
2003
10000
200000
30000
40000
Itoration
Ioortion
Iteration
Figure 21: DAL-P sweep: discretization drift helps test performance at the cost of stability. Corresponding training curves and loss functions are present in the Figure 44 in the Appendix; results showing the same trends across various batch sizes are shown in Figure 41
The higher p, the slower the training and less drift there is; the lower p, there is more drift. We start with extensive experiments with p 0.5, which we show in Figure and show more results in Figure B Compared to p 1 (DAL), there is faster training but at times also more instability. Performance on the test set shows that DAL-0.5 performs as well or better than when using fixed learning rates.
Remark 6.2 We find that across datasets and batch sizes, DAL-0.5 performs best in terms of the stability generalization trade-off and in these settings can be used as a drop in replacement for a learning rute sweep.
To further investigate the connection between drift and test set performance, we perform a set of sweeps over the power p and show results in Figure 21 These results show that the higher the drift (the smaller p), the more gencralization; additional results across batch sizes showing the same trend are shown in Figure 41 in the Appendix. We also show in Figure 22 the correlation between mean per iteration drift and test accuracy both for learning rate and DAL-P sweeps. The results consistently show that the higher the mean iteration drift, the higher the test accuracy. We also show that the mean iteration drift has connection to the largest eigenvalue Ag: the higher the drift, the smaller Ag. These results add further evidence to the idea that discretization drift is beneficial for generalization performance in the deep learning setting. We also notice that DAL-p with smaller values of leads to a small λn compared to vanilla gradient descent even when large learning rates are used for the latter; this could explain its generalisation capabilities as lower sharpness has been connected to generalisation in previous works (Keskar et al. 2016; Jastrzebbkki et al., 2018 Foret et al. 2020). To consolidate these results, we use the method of Li et al. (2018) to visualise the loss landscape learned by DAL-p compared to that learned using gradient descent, and observe that even when reporting similarr accuracies, DAL-p converges to flatter landscape; this is observed even when small
batch sizes are used. Results are shown in Figures 46 48 50 in the Appendix.
Inspired by understanding when the PF is close to the NGF, in this section we investigated the total discretisation drift of gradient descent. This led us to DAL-p, method to automatically set the lcarning rate based on approximation to the per iteration drift of gradient descent; we have seen that DAL produces stable training and further connected discretisation drift, generalisation and flat landscapes as measured by leading Hessian eigenvalues.
20

---

Published in Transactions on Machine Learning Research (01/2023)
140
Lea ming
Leerning rate
DAL-p
1.20-
0.05
Mean An before convergene
500
0:125
0.68
0.07
0.23
Test acou/acy
0.01
1.0
0.02
0.04
0.09
0.08
2.75
0.5
0.75
0.
0.03
54
0.05
L.
Moan A. before
1.5
2.0
0.44
0.25
0.83
0.52
0.43
Mean iteration drift (approx]
Mean iteration drift (approx!
Mean iteration drift (approx)
Mean iteration drift (appros)
(a) Fixed learning rate sweep.
(b) DAL-p sweep.
Figure 22: The correlation between drift, test set performance and Ap full batch training on CIFAR-10. The same pattern can be scen in SGD results in Figure 53
7 Future work
Beyond gradient descent. this work we focused on understanding vanilla gradient descent. Under- standing discretization drift via the PF can be beneficial for improving other gradient based optimization algorithms as well, as we briefly illustrate for momentum updates with decay m and learning rate h:
V = mV -
1 NV
0, = 0. +
(20)
We can scale Veb(0,-1) in the above not by a fixed learning rate h, but by adjusting the learning rate according to the approximation to the drift. This has two advantages: it removes the need for lcarning rate sweep and it uses local landscape information in adapting the moving average, such that in areas of large drift the contribution decreased, while it increased areas where the drift is small (a more formal justification provided in Section A.10 This leads to the following updates:
V: = mV 2
((F(
- 0.
(21)
As with DAL-P, we can use powers to control the stability performance trade-off: the lower p, the more the current updale contribution is reduced in high drift (instability) areas. We tested this approach Imagenet and show results in Figure 23 The results show that integrating drift information improves the speed of convergence compared to standard gradient descent (Figure 21), and leads to more stable training compared to using a fixed learning rate. We present additional experimental results in the Appendix.
Imagenet Batch 8192 with momentum 0.9
Imagenet Batch 8192 with momentum 0.9
Learning rate 0.05
Learning rate 0.1
Learning
0.8
Training less
3-
DAL- 0.5 with momentum momentum
(Tap 1)
2000
8000
14000
0.0
2000
12000
14000
atioo
Rtertion
Figure 23: DAL with momentum: integrating drift. information results in faster and more stable training compared to fixed learning rate sweep. Compared to vanilla gradient descent there is also significant performance and convergence speed boost.
Just as momentum is a common staple of optimization algorithms, so are adaptive schemnes such as Adam (Kingma and Bal 2015) and Adagrad (Duchi et al. 2011). which adjust the step taken for each
21

---

Published in Transactions on Machine Learning Research (01/2023)
parameter independently. We can also use the knowledge from the PF to set per parameter learning rate: instead of using V (o) to set a global learning rale, we can use the per parameter informa- to adapt the learning rate of each parameter. We present preliminary
tion provided by V= results in the Appendix (Figures 51 and 52 The above two approaches (momentun and per-parameter learning rate adaptation) can be combined, bringing us closer to the most commonly used deep learning optimization algorithms. While we do not explore this avenue here, we are hopeful that this understanding of discretization drift can be leveraged further to stabilize and improve deep learning optimization.
Non- con-principal terms. This work focuses understanding the effects of the PF on the behavior of gradient descent. The principal terms however are not the only terms in the discretization drift: we have found one non-principal lerm (Eq [ and have seen that it can have a stabilising effect (Figure G We provide a preliminary explanation for the stabilising effect of this non-principal term together with results measuring its value in neural network training in Section [ in the Appendix. One promising avenue of non -oriciici terms is thcorctically modelling the change of the eigenvalues X, in time; another promising direction is that of implicit regularisation: while existing work which uses BEA in deep learning has found important implicit regularisation effects (Barrett and Dherin 2021: Smith et al. 2021 Rosca et al. 2021), we have shown here that considering only eflects of O(h))) is not sufficient to capture the intricacies of gradient descent, which suggests that other implicit regularisation effects could be uncovered using the non-principal terms.
Neural network theory. Many theoretical works studying at gradient descent in the neural network context use the NGF (Du et al.) 2018: Elkabetz and Cohen| 2021 Kunin et al.) 2021 Jacot et al. 2018). We posit that replacing NGF in these theoretical contexts with PF may yield interesting results. In contrast to the NGF, the PF allows the incorporation of the learning rate into the analysis, and unlike existing continuous time models of gradient descent, it can model unstable behaviors obscrved in the discrete case. An example can be seen using the Neural Tangent Kernel: Jacot et al. (2018) model gradient descent using the NGF to show that in the infinite wide limit gradient descent for neural norksrks follows kernel gradient descent. The PF can be incorporated in this analysis either by replacing the NGF with the PF as a model of gradient descent or by studying the difference in the PF for infinitely wide and finite width networks, since discretisation drift could be responsible for the observed gap between finite and infinite networks in the large learning rate case (Loe et al. 2020).
8 Related work
Modified flows for deep learning optimization. Barrett and Dherin 20021 found the first order correction modified flow for gradient descent using BEA and uncovered its regularization effects; they were the first to show the power of BEA in the deep learning context. [mith et al. (2021) find the first order error correction term in expectation during one epoch of stochastic gradient descent. Modified flows have also been used for other optimizers than vanilla gradient descent: Franca et al. (2020): Shi al. (2021) compare momentum and Nesterov accelerated momentum; Kunin et al. (2021) study the symmetrics of deep neural networks and use modified vector fields to show commonly used discrete updates break conservation laws present when using the NGF (for gradient descent they use the IGR flow while for momentum and weight decay they introduce different flows); Kovachki and Stuart, (2021) use modified flows to understand the behavior of momentum by approximating Hamiltonian systems; França et al. (2021) construct optimizers controlling their stability and convergence rates while Li et al. (2017) construct optimizers with adaptive learning rates in the context of stochastic differential equations. In the context of two-player games, Rosca
ct al. 2021 compute the first order BEA correction terms while while/ Chavdarova et al. (2021) use high-resolution differential equations to shed light on the properties of different saddle point optimizers.
In concurrent work Miyagawa (022)) use BEA to find modified flow coined 'Equations of Motion' (EOM) to describe gradient descent and find higher order terms, including non-principal terms; their focus is however on (1M(1), which is the IGR flow, which they use to understand scale and translation invariant layers. Their approach does not expand to complex space and does not capture the instabilities studied here (see also the discussion on the difference between the full modified flow provided by BEA and the PF in Section 3
Edge of stability and the importance of the Hessian. There have been number of empirical studies on the Hessian in gradient descent. Cohen et al. (2021) observed the edge of stability behavior and performed an extensive study which led to many empirical observations used in this work. Jastrzebski el al.] (2018) performed a similar study in the context of stochastic gradient descent. Sagun et al. (2017 Ghorbani et al.
22

---

Published in Transactions on Machine Learning Research (01/2023)
2019: Papyan (2018) approximate the entire spectrum of the Hessian, and show that there are only lew negalive Tigenvalues, plenty of eigenvalues centered around 0 and a lew posilive eigenvalue, with large magnitude. Similarly, Gur-Ari et al. (2018) discuss how gradient descent operates in a small subspace. Lewkowyez et al. (2020) discuss the large learning rate catapult in deep learning when the largest eigenvalue exceeds 2/h. Gilmer et al. (2021) assess the effects of the largest Hessian eigenvalue in a large number of empirical settings.
There have been series of concurrent works aimed at theoretically explaining the empirical re- sults above. Ahn et al. (2022) connect the edge of stability behavior with what they coin as which they empirically show is 0 in stable areas of training and in the edge of stability arcas. To see the connection between the relative progress ratio and the quantities discussed in this paper, one can perform Taylor expansion on -1 + h/2 VE While this ratio is related to the quantities we discuss, we also note significant differences: it is scalar, and not a parameter length vector and thus does not capture per cigendirection behavior as we see with the stability coefficients (Section 5 Arora et al. 2022) prove the edge of stability result occurs under certain conditions either on the learning rale or on the loss function. Ma al.] (2022) empirically observe the multi-scale structure of the loss land- scape in neural networks and use it to theoretically explain the edge of stability behavior of gradient descent.
the 'relative progress ratio': (00-AV,E--ER
EL ~ -AVEE" "VAE"VvEEt?" "4E++2/2VE"" V3VVE
E10
A Vobill
Chen and Bruna (2022) use low dimensional theoretical insights around a local minima to understand the edge of stability behavior. Damian et al. (2022) use cubic Taylor expansion to show that gradient descent follows the trajectory of a projected method which ensures that Àu < 2/h and V ET
u = O; their work is what inspired us to write the third order nou-principal term in the form of Eq 182 in the Appendix, after we had previously noted its stabilizing properties. These important works are complementary to our own work; they do not use continuous time approaches and tackle primarily the edge of stability problem or its subcases, while we focus on understanding gradient descent and applying that understanding broadly, including but not limited to the edge of stability phenomenon.
Discrete models of gradient descent. The desire to understand learning rate specific behavior in gradient descent. has been a motivation in the construction of discrete time analyses. These analyses have provided great insights, from studying noise in the stochastic gradient descent setting (Liu et al.] 2021 Ziyin et al. 2021b), the study of overparametrized neural models and their convergence (Gumasekar et al. 2018 Du et al. 2019 Allen-Zhu et al. 2019), providing examples when gradient descent can converge to local maxima (Ziyin ee al. 202 the importance of width for proving convergence in deep linear networks (Du and Hu 2019) We differ from these studies both in motivation and execution: we are looking for a continuous time flow which will increase the applicability of continuous time analysis of gradient descent. We do so by incorporating discretisation drift using BEA and showing that the resulting flow is useful model of gradient descent, which captures instabilities and escape of local minima and saddle points.
Understanding the difference between the negative gradient flow and gradient descent. Elkabetz and Cohen (021)) recently examined the differences between gradient descent and the NGF in the deep learning context; their work examines the importance of the Hessian in determining when gradient descent follows the NGF. Their theoretical results show that neural networks are roughly convex and thus for reasonably sized learning rates one can expect that gradient descent follows the NGF flow closely. Their results complement ours and their approach might be extended to help us understand why the PF is sufficient to shed light on many instability observations in the neural network training.
Second-order optimization. By using second order information (or approximations thereof) to set the learning rate, DAL is related to second-order approaches used in deep learning. Many second-order hoss can be seen as approximates of Newton's method 0, = 0, inverse of Hessian can be prohibitively expensive for large models, many practical methods approximate it with tractable alternatives (Martens and Grosse 2015 Foret et al. (2020) propose an optimisation scheme dircetly aimed at minimising sharpness, and show this can improve generalisation.
(0,-0 (o Since computing the
Connection between drift and generalization. We have made the connection between increased drift and increased generalization. This connection was first made by BBrrett and Dherin (2021) through the IGR flow. Generalization has also been connected to the largest eigenvalue Xo (Hochreiter and Schmidluber
1997 Keskar et al. 2016 Jastrzębski et al.) 2018 Lewkowycz et al. 2020); recently Kaur et al. (0022
however showed a more complex picture, primarily in the context of stochastic gradient descent. The largest
23

---

Published in Transactions on Machine Learning Research (01/2023)
eigenvalue could be confounder to the drift we have observed in Section 6.3 we hope that future work can deepen these connections.
9 Conclusion
We have expanded on previous works which used Backward Error Analysis in deep learning find new continuous time flow, called the Principal Flow, to analyze the behavior of gradient descent. Unlike ex- isting flows, the principal flow operates in complex space which enables to better capture the behavior of gradient descent compared cxisting flows, including but not limited instability and oscillatory behavior, We use the form of the Principal Flow to find new quantities relevant to the stability of gradient descent, and shed light on newly observed empirical phenomena, such as the edge of stability resulls. After under- standing the core quantities connected to instabilities deep learning we devised an automatic learning rate schedule, DAL, which exhibits stable training. We concluded by cementing the connection between large discretization drift and increased generalization performance. We ended by highlighting future work avenues including incorporating the principal flow in existing thoorctical analyses of gradient descent which use the negative gradient flow, incorporating understanding of the drift of gradient descent in other optimization approaches and specializing the PF for neural network function approximators.
Acknowledgments. We would like to thank the TMLR anonymous reviewers and the TMLR Action Editor for their useful feedback and comments. We would also like to thank Soham De and Michael Munn for discussions and feedback; and Frederic Besse, Mare Deisenroth, Patrick Cole, Shakir Mohamed and Timothy Lillicrap for their support.
24

---

Published in Transactions on Machine Learning Research (01/2023)
References
Paul Glendinning. Slability, instability and chaos: on introduction to the theory of nonlinear differential equations. Cambridge university press, 1994.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXio:1312.6122 2013.
Vaishnavh Nagarajan and Zico Kolter. Gradient descent gan optimization locally stable. In Advances in neural informalion processing systems, pages 5585-5595, 2017.
Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprinl ar Xiv: 809. 10374, 2018.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameeteriaation. International Conference on Machine Learning, pages 244 253. PMLR, 2018
Madlu S Advani, Andrew M Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks, 132:428 446, 2020.
Omer Elkabetz and Nadav Cohen. Continuous vs. discrete optimization of deep neural networks. Advances in Neural Information Processing Systems, 34:4947 4960, 2021.
Gal Vardi and Ohad Shamir. Implicit regularization in relu networks with the square loss. In Conference Learning Theory, pages 4224 4258. PMLR, 2021.
Guilherme Franca, Jeremias Sulam, Daniel Robinson, and Rene Vidal. Conformal symplectic and relativistic optimization. Conference on Neural Information Processing Systems (NeurIPS 2020). 2020.
David GT Barrett and Benoit Dherin. Implicit gradient regularization. In International Conference on Learning Represendations, 2021.
Samuel L Smith, Benoit Dherin, David GT Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. In International Conference on Learning Regresentations, 2021.
Ernst Hairer, Marlis Hochbruck, Arieh Iserles, and Christian Lubich. Geometric numerical integration. Oberwolfach Reports, 3(1):005 882, 2006.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent neural networks typically occurs at the edge of stability. arXiv preprint W.Xi::210300006. 2021.
Guy Gur-Art Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint (rXf::1812.0454, 2018
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instability in deep learning. arXiv preprint 0at:22100044000 2021.
Aitor Lewkowyez, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.0222 2020.
Peter Bartlett, Dave Helmbold, and Philip Long. Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks. International conference on machine learning, pages 521-530. PMLR, 2018a.
Peter Bartlett, Steven N Evans, and Philip M Long. Representing smooth functions as compositions of near-identity functions with implications for deep network optimization. orXiv preprint orXiu: 180, 05012, 2018b.
Lars Mescheder, Sebastian Nowozin, and Andreas Geiger. The numerics of gans. In Advances in Neural Information Processing Systems, pages 1825-1835, 2017.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems, 31, 2018
25

---

Published in Transactions on Machine Learning Research (01/2023)
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine leurning, pages 1675-1682 1685. PMLR, 2019.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pages 242-252, PMLR, 2019.
Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In Inter- national Conference on Machine Learning, pages 1655-1664. PMTR, 2019.
Liu Ziyin, Botao Li, James B Simon, and Masahito Ueda. Sgd can converge to local maxima. In International Conference Learning Representations, 2021a.
Kangqiao Liu, Liu Ziyin, and Masahito Ueda. Noise and fluctuation of finite learning rate stochastic gradient descent. In International Conference Machine Learning, pages 7045-7056. PMLR, 2021.
David Balduzzi, Sebastien Racaniere, James Martens, Jakob Foerster, Karl Tuyls, and Thore Graepel. The mechanics of n-player differentiable games. In International Conference on Machine Learning, pages 354- 363. PMLR, 2018.
Simon Du, Wei Hu, and Jason Lee. Algorithmic regularization in learning deep homogeneous models: Layers automatically balanced. Advances in Newral Information Processing Systems, 31, 2018.
Robert M May. Simple mathematical models with very complicated dynamics. Nature, 261:459, 1976.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251 276, 1998.
Yann Ollivier. Riemannian metries for neural networks i: leedforward networks. Information and Inference: A Journal of the IMA, 4(2):108-153, 2015a.
Yann Ollivier. Riemannian metrics for neural networks ii: recurrent networks and learning symbolic data sequences. Information and Inference: A Journal of the IMA, 4(2):154 193, 2015b.
Yang Song, Jiaming Song, and Stefano Ermon. Accelerating natural gradient with higher-order invariance. In International Conference on Machine Learning, pages 4713-4722. PMLR, 2018
Chongli Qin, Yan Wu, Jost Tobias Springenberg, Andrew Brock, Jeff Donahue, Timothy P Lillicrap, and Pushmeet Kohli. Training generative adversarial networks by solving ordinary differential equations. 2020.
Mihaela C Rosca, Yan Wu, Benoit Dherin, and David Barrett. Discretization drift in two-player games. In International Conference on Machine Learning, pages 9064-9074. PMLR, 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systemss, 31, 2018.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as lincar models under gradient descent. Advances in neural information processing systems, 32, 2019.
Sho Yaida. Pluctuation-ddissipatoon relations for stochastic gradient descent. In International Conference Learning Represeniations, 2018.
Gerhard Wanner and Ernst Hairer. Solving ordinary differential equations II, volume 375. Springer Berlin Heidelberg New York, 1996.
HoHo Rosenbrock. An automatic method for finding the greatest least value of function. The compuler journal, 3(3):175-18) 184, 1960.
Arthur Asuncion and David Newman. Uci machine learning repository, 2007.
Stanislaw Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey, On the relation between the sharpest directions of dnn loss and the sgd step length. ar Xiu preprinl arXiv:1807.0555 807.00031, 2018.
26

---

Published in Transactions on Machine Learning Research (01/2023)
Yann N Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. Identifying and attacking the saddle point problem in bigh-dinensional non-convex optimization. Advances neural information processing systems, 27, 2014.
Jason D Lee, Max Simchowitz, Michael I Jordan, and Benjamin Recht. Gradient descent only converges to minimizers. In Conference on learning theory, pages 1246 1257. PMLR, 2016.
Kenji Kawaguchi. Deep learning without poor local minima. Advances in neural information processing systems, 29, 2016.
Simon S Du, Chi Jin, Jason D Lec, Michael I Jordan, Aarti Singh, and Barnabas Poczos. Gradient descent can take exponential time to escape saddle points. Advances in neural information processing systems, 30, 2017.
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of paraaaaametri neural networks. arXi preprint arXiv:1706.044: 04454, 2017.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-22222 2232-2241. PMLR, 2019.
Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint wrXiw:1811.07062 2018.
llya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arxiv:1608. 0,9333 2016.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. Bag of tricks for image clas- sification with convolutional neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 558 567, 2019.
Jonas Geiping, Micah Goldblum, Phillip E Pope, Michael Moeller, and Tom Goldstein. Stochastic training is not necessary for generalization. arXiv preprint Xrxiv:2109. 14119, 2021.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training noural networks. Advances in Neural Information Processing Systems, 32, 2019.
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Kraysztof Geras. The break-even point on optimization trajectorics of deep neural networks. In International Conference on Learning Representations, 2019.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Gencralization gap and sharp minima. arXie preprint arXiv:1609. 161.00669004896, 2016.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur, Sharpness aware minimization for efficiently improving gencralization. In International Conference on Learning Representations, 2020.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 2015.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Daniel Kunin, Javier Oogastuy-Brena and Hidenori Tanaka Ganguli, Surya Daniel L.K. Yamins. Symmetry, conservation laws, and learning dynamics in neural networks. In International Conference on Learning Representations, 2021.
Jachoon Lee, Samuel Schocnholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems, 33:15156-1511 2020.
27

---

Published in Transactions on Machine Learning Research (01/2023)
Bin Shi, Simon S Du, Michael I Jordan, and Weijie J Su. Understanding the acceleration phenomenon via high-resolution differential equations. Mathemalical Programming, pages 1 70, 2021.
Nikola B Kovachki and Andrew M Stuart. Continuous time analysis of momentum methods. Journal of Machine Learning Research, 22(17):1 40, 2021.
Guilherme França, Michael 1 Jordan, and René Vidal. On dissipative symplectic integration with applications gradient-based optimization. Journal of Statistical Mechanics: Theory and Experiment, 2021 2021 (4):043402, 2021.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, volume 70, pages 2101-2110, 2017.
Tatjana Chavdarova, Michael Jordan, and Manolis Zampetakis. Last-iterate convergence of saddle point optimizers via high-rosolution differential equations. arXiv preprint arXiv:21 13826, 2021.
Taiki Miyagawa. Toward equation of motion for deep neural networks: Continuous time gradient descent and discretization error analysis. In Advances in Neural Information Processing Systems, 2022.
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient
descent. arXie preprint arXiv:220400105 010500 2022.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arxiv preprint orXiv:2205.097455 2022.
Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure of neural network loss functions: The effect on optimization and origin. arxiv preprinl urXiv:2204.110 11326, 2022.
Lei Chen and Joan Bruna. On gradient descent convergence beyond the edge of stability. arXie preprint arXiv:2206.00 2.Ni::220600112.. 04172, 2022.
Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent the edge of stability, Xiviv preprint arXiv:2209.158: 2022.
Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. Strength of minibatch noise in sgd. anXiv preprint arXiv:2102.059 2021b.
James Martens and Roger Grosse. Optimizing neural networks with frontceer-factored approximate curva- ture. International conference on machine learning, pages 2408-2417 PMLRR 2015.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Simran Kaur, Jeremy Cohen, and Zachary C Lipton. On the maximum hossian cigenvaluc and genoralization. arXiv preprint ar (xXi::220611655, 2022.
David M Grobman. Homcomorphism of systems of differential equations. Doklady Akademūi Nauk SSSR, 128(5):880 881, 1959.
Philip Hartman. A lemma in the theory of structural stability of differential equations. Proceedings of the American Mathematical Society, 11(4):610 620, 1960.
Yann LeCun, Lawrence D Jackel, Léon Bottou, Corinna Cortes, John S Denker, Harris Drucker, Isabelle Guyon, Urs Muller, Eduard Sackinger, Patrice Simard, et al. Learning algorithms for classification: A comparison on handwritten digit recognition. Neural networks: the statistical mechanics perspective, 261 (276):2, 1995.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pollern recognilion, pages 248 255. Teee, 2009.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint or Xie: 1409 1556, 2014.
28

---

Published in Transactions on Machine Learning Research (01/2023)
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on compuler vision and pallern recognilion, pages 770-778, 2016.
Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:5511.0788. 2015.
Abien Fred Agarap. Deep learning using rectified lincar units (relu). arXiv preprint 10xx:1808.08975 2018.
29


Reviews:
[{"review_id": "PXfZHQvz09", "reviewer": "Reviewer_Psao", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: the strength of the paper is the novelty of the proposed approximation and the extensive experiments used to validate its empirical relevance.\n\nWeakness: \n1. I feel that there is not any new significant insight the new theory provides. The theory mostly confirms previous experimental discoveries and the insights the previous works have provided. If theory is really novel and meaningful, it should also provide some new insight that no previous work has discovered.\n\n2. The proposed DAL algorithm also does not feel very novel. It can also be obtained by a simple application of Newton's method.\n\n3. This is my major criticism. I do not feel that the paper is sufficiently motivated. The paper has glossed over the discussion of a very important motivation: why is a continuous-time approximation needed in the first place? \n- For example, the instability threshold 2/h can be directly understood from a discrete-time treatment (in fact, even more precise treatments have been established). The DAL method also does not require a continuous-time model to justify. See the requested changes section for more detail\n\n4. A different question -- gradient flow is known to work very well for NTKs, does the present work offer a better understanding for NTKs when trained at a larger learning rate?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a continuous-time flow approximation to the gradient descent algorithm.\n\nThe proposed approximation is found to approximate the actual gradient descent much better than the existing flow methods.\n\nThe paper then proposes the DAL, a second-order method for setting the learning rate to improve the stability of training.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "While I feel that the scientific content of the paper is decent, the discussion of the motivation and the related works needs extensive updates for me to recommend acceptance.\n\n1. Extensive discussion of why continuous-time approximation is needed, especially given many works exist to study the original discretized version. Especially, the authors need to point out clearly, in light of the previous works, what the previous works on discrete-time gradient descent cannot offer and the continuous-time theories excel at. I do personally believe that continuous-time theories are important and have a lot of merits -- but the authors are responsible for explaining why. To my knowledge, the following works are worth discussing (saddle points, stability, etc).\n- https://proceedings.mlr.press/v139/liu21ad.html\n- https://openreview.net/forum?id=SkNksoRctQ\n- https://openreview.net/forum?id=uorVGbWV5sw\n- https://openreview.net/forum?id=9XhPLAjjRB\n\n2. In fact, the authors did not mention any work that studies gradient descent in discrete-time, as if they do not exist. Outside the references I point out above, the authors should look for more works in this direction and discuss\n\n", "broader_impact_concerns": "Fine.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "dsxCVvh4Fd", "reviewer": "Reviewer_uCdx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "trengths: \n\tThe idea using BEA to provide a more precise continuous flow for (S)GD is interesting.\n\tThe paper is well-written.\n\nWeakness:\nI find a technical flaw which may severely deteriorate the contribution of this paper.\nFrom Eq. (11) to Eq. (12) (or the proof of Theorem 13), the authors replace \\sum_{p=0}^\\infty {- 1/(p+1) * z^p}  with log(1-z)/z, which is not reasonable to me. A simple example is that the former term diverges to -\\infty when z>1 while the latter is a finite complex number (corresponding to the case λ_i>1/h in page 7.  All of the divergent and oscillatory analyses are based on this case). However, the latter term yields a complex number when z>1 (which is the focus of this paper). In a word, log⁡(1-z)/z is only a proper approximation for the case |z|≤1, but the authors use the approximation for all the cases, and erroneously arrived at the conclusion “for the case λ_i>1/h, PF is a complex flow” (to me, it seems PF even cannot be defined in this case).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new continuous time flow to approximate gradient descent. Specifically, based on the Backward Error Analysis (BEA) technique in numerical analysis, the authors propose the principle flow (PF) with higher order error terms. Based on the principle flow, the authors are able to explain divergent and oscillatory behaviors in deep learning training, and propose a learning rate adaptation method to reduce instabilities. ", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Please refer to the part “Weakness”. Since the conclusions of this paper heavily based on the above analysis, I hope the authors can comment on this.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Vms096vLLA", "reviewer": "Reviewer_hnSw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The PF can well characterize the stability of GD. Compared with NGF and IGR, its extension to the complex space allows different stability behaviors of the dynamics around stationary points. This is very interesting and may potentially introduce more tools from continuous dynamics to study GD. \n\n2. The paper provides extensive experimental evidence to justify the arguments made.\n\nWeaknesses:\n\n1. From the reviewer's perspective, the explanation to the edge of stability phenomenon in Section 5 is not enough. The explanations focus on the stability of GD on a quadratic landscape. However, on a quadratic landscape the GD cannot stay at the edge of stability---it either converges or diverges. Can the PF be used to explain why GD can stay at the edge of stability?\n\n2. The contents in Section 6 seems not related with PF. Instead, the relation between the GD and NGF is studied. Could the authors provide some motivations from the PF for this section? \n\n3. It is shown clearly that the PF characterizes the same stability of GD around a minimum. However, the paper does not provide a case in which the PF can help analyze the the behavior of GD that cannot be done by directly studying GD---it seems all results for PF in the paper can be obtained by simply look at GD. Can the PF extend our understanding for GD? ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors first proposes the Principal Flow (PF) as a continuous-time approximation for the discrete gradient descent method. The PF is derived from the approach of modified equation in numerical analysis, which is prevalent in the analysis of numerical methods for differential equations. The PF dynamics depends on the spectrum of the Hessian of the loss function. Compared with the traditional gradient flow (NGF), the PF changes the magnitude the update (right hand side of the dynamics) along each eigenvector direction of the Hessian by a logarithm function. The logarithm function makes the dynamics complex, and thus allows richer behaviors around stationary points. \n\nAfter deriving the PF, the authors first studied the stability of the PF with respect to the eigenvalues of the Hessian. The results are parallel with that for GD. Showing the advantage of PF over NGF and IGR on the approximation of GD behaviors. Then, the PF is used to explain the behaviors of GD in different regimes observed in the edge of stability phenomenon. Extensive experiments show that the PF can indeed characterize GD behaviors in both the stable and edge of stability regimes. \n\nFinally, the authors studied the closeness of GD and NGF, showing that the \"discretization drift\" is the crucial quantity that controls this closeness. From here, an algorithm which scales the learning rate using a quantity related with the discretization drift is proposed. Numerical results show that the proposed algorithm (DAL) can improve the stability of training (over GD). However, a tradeoff of stability and generalization on the other hand shows algorithms that depend less on the discretization drift generalize better. ", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Questions are posted in the \"strengths and weaknesses\" part. ", "broader_impact_concerns": "The reviewer does not have concerns on the broader impact of the paper. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The article proposes a continuous time model to understand the behavior of gradient descent. 
There was a healthy amount of discussion following the initial reviews. During the discussion, several perceived shortcomings and technical aspects could be clarified. The authors made various updates to the manuscript in response to the discussion with the reviewers. 

Overall, I find the article discusses a topic that is of interest to the TMLR audience. While the final reviewer recommendations are mixed and there are still some aspects that could be further improved, after evaluating the discussion, rebuttal, article updates, and discussing with the reviewers, I find that the merits of this article outweigh the weaknesses and that the paper meets the bar for acceptance. 

I am recommending accept with minor revision and ask the authors to take one final pass over the article taking the final reviewer comments into account. In particular, uCdx indicated that the technical flaws they perceived initially have been resolved for the case of quadratic functions, but that for general cases they were not fully addressed. Reviewer hnSw indicated that the article proposes a good continuous approximation for GD, but that they would expect to see new insights. The authors gave a few examples in the forum later. Reviewer Psao was most favorable after the discussion, but also had minor final suggestions. 

Decision: Accept (poster)