Rethinking the limiting dynamics of SGD: modified loss, phase space oscillations, and anomalous diffusion
Authors: Daniel Kunin
Javier Sagastuy-Brena
Lauren Gillespie
Eshed Margalit
Hidenori Tanaka
Surya Ganguli
Daniel LK Yamins
Venue: ICLR 2022 Submitted (2022)
Keywords: learning dynamics
loss landscape
stochastic differential equation
modified equation analysis
hessian
geometry
physics
fokker-plank
modified loss
probability currents
diffusion

Abstract: In this work we explore the limiting dynamics of deep neural networks trained with stochastic gradient descent (SGD). We find empirically that long after performance has converged, networks continue to move through parameter space by a process of anomalous diffusion in which distance travelled grows as a power law in the number of gradient updates with a nontrivial exponent. We reveal an intricate interaction between the hyperparameters of optimization, the structure in the gradient noise, and the Hessian matrix at the end of training that explains this anomalous diffusion. To build this understanding, we first derive a continuous-time model for SGD with finite learning rates and batch sizes as an underdamped Langevin equation. We study this equation in the setting of linear regression, where we can derive exact, analytic expressions for the phase space dynamics of the parameters and their instantaneous velocities from initialization to stationarity. Using the Fokker-Planck equation, we show that the key ingredient driving these dynamics is not the original training loss, but rather the combination of a modified loss, which implicitly regularizes the velocity, and probability currents, which cause oscillations in phase space. We identify qualitative and quantitative predictions of this theory in the dynamics of a ResNet-18 model trained on ImageNet. Through the lens of statistical physics, we uncover a mechanistic origin for the anomalous limiting dynamics of deep neural networks trained with SGD.

Under review as a conference paper at ICLR 2022
THE LIMITING DYNAMICS OF SGD: MODIFIED LOSS, PHASE SPACE OSCILLATIONS, ANOMALOUS DIFFUSION
Anonymous authors Paper under double-blind review
ABSTRACT
In this work we explore the limiting dynamics of deep neural networks trained with stochastic gradient descent (SGD). As observed previously, long after per- formance has converged, networks continue to move through parameter space by a process of anomalous diffusion in which distance travelled grows as a power law in the number of gradient updates with nontrivial exponent. We reveal an intricate interaction between the hyperparameters of optimization, the structure in the gradient noise, and the Hessian matrix at the end of training that explains this anomalous diffusion. To build this understanding, we first derive continuous- time model for SGD with finite learning rates and batch sizes as an underdamped Langevin equation. We study this equation in the setting of linear regression, where we can derive exact, analytic expressions for the phase space dynamics of the parameters and their instantaneous velocities from initialization stationar- ity. Using the Fokker-Planck equation, we show that the key ingredient driving these dynamics is not the original training loss, but rather the combination of a modified loss, which implicitly regularizes the velocity, and probability currents, which cause oscillations in phase space. We identify qualitative and quantitative predictions of this theory in the dynamics of ResNet-18 model trained Im- ageNet. Through the lens of statistical physics, we uncover a mechanistic origin for the anomalous limiting dynamics of deep neural networks trained with SGD,
INTRODUCTION
Deep neural networks have demonstrated remarkable generalization across a variety of datasets and tasks. Essential their success has been a collection of good practices how to train these models with stochastic gradient descent (SGD). Yet, despite their importance, these practices are mainly based heuristic arguments and trial and error search. Without general theory connecting the hyperparameters of optimization, the architecture of the network, and the geometry of the dataset, theory-driven design of deep learning systems is impossible. Existing theoretical works studying this interaction have leveraged the random structure of neural networks at initialization [1] 2 B and their infinite width limits in order to study their dynamics 4 S different approach and study the training dynamics of pre-trained networks that are ready to be used for inference. By leveraging the mathematical structures found at the end of training, we uncover an intricate interaction between the hyperparameters of optimization, the structure in the gradient noise, and the Hessian matrix that corroborates previously identified empirical behavior such as anomalous limiting dynamics. Not only is understanding the limiting dynamics of SGD a critical stepping stone to building a complete theory for the learning dynamics of neural networks, but recently there have been a series of works demonstrating that the performance of pre-trained networks be improved through averaging and ensembling (9] 10, IT. Combining empirical exploration and theoretical tools from statistical physics, we identify and uncover a mechanistic explanation for the limiting dynamics of neural networks trained with SGD.
X. Here we take
2 DIFFUSIVE BEHAVIOR IN THE LIMITING DYNAMICS OF SGD
A network that has converged in performance will continue to move through parameter space 12. 13 [5). To demonstrate this behavior, resume training of pre-trained convolutional networks while tracking the network trajectory through parameter space. Let 0, e R" be the parameter vector for a pre-trained network and 0g IR be the parameter vector after k steps of resumed training.

---

Under review as a conference paper at ICLR 2022
We track two metries of the training trajectory, namely the local parameter displacement òg between consecutive steps, and the global displacement A5 after & steps from the pre-trained initialization:
6% --=4 02 - 0k-
4 02 O2. (I)
VGG-16
3x
As shown in Fig. [ neither of these differences converge to zero across a variety of architectures, indicating that despite performance convergence, the networks continue to move through parameter space, both locally and glob- ally. The squared norm of the local displacement [[ remains near a constant value, indicating the network essentially moving at a constant instantaneous speed. This observation is quite similar to the "equilibrium" phe- nomenon or "constant angular update" observed in Li et al. [17] and Wan et al. [3] respectively. However, these works only studied the displacement for parameters im- mediately preceding normalization layer. The constant instantaneous speed behavior we observe is for all param- eters in the model and is even present in models without normalization layers.
VGG-16 with BN
10
||8112
VGG-11 with BN
R4
10
12x10
RedNet-18
15000
30000
45000
Stop
109
c=1.085
== 1.075
10
1.370
||
<= 1.401
1.616
10-1
While the squared norm of the local displacement is es- sentially constant, the squared norm of the global dis- placement ||A,112 2/2 is monotonically growing for all net- works, implying even once trained, the network continues to diverge from where has been. Indeed Fig. innicatts power law relationship between global displacement and number of steps, given by || x k. As we'll see in section [8] this relationship is indicative of anoma- lous diffusion where c corresponds to the anomalous dif- fusion exponent. Standard Brownian motion corresponds to c = 1. Similar observation were made by Baity-Jesi al. 14 who noticed distinct phases of the training tra- jectory evident in the dynamics of the global displace- ment and Chen et al. [15] who found that the exponent of diffusion changes through the course of training. A parallel observation is given by Hoffer et al. [18] for the beginning of training, where they measure the global dis- placement from the initialization of an untrained network and observe a rate xx log(k), a form of ultra-slow diffusion. These empirical observations raise the natural questions, where is the network moving to and why? To answer these questions we will build a diffusion based theory of SGD, study these dynamics in the setting of linear regression, and use lessons Iearned in this fundamental setting to understand the limiting dynamics of neural networks.
0.5
Step
Figure 1: Despite performance con- vergence, the network continues to move through parameter space. We plot the squared Euclidean norm for the local and global displacement (5, and Ax) of five classic convolutional neu- ral network architectures. The networks are standard Pytorch models pre-trained on ImageNet [16). Their training is re- sumed for 10 additional epochs. We show the global displacement on log- log plot where the slope of the least squares line is the exponent of the power law || o ke. See ap- pendix H for experimental details.
3 RELATED WORK
There is long line of literature studying both theoretically and empirically the learning dynamics of deep neural networks trained with SGD. Our analysis and experiments build upon this literature.
Continuous models for SGD. Many works consider how improve the classic gradient flow model for SGD to more realistically reflect momentum [19] discretization due to finite learning rates [10) and stochasticity due to random batches [22 223. One line of work has studied the dynam- ics of networks in their infinite width limits through dynamical mean field theory (24] 25 26 271, while a different approach has used stochastic differential equations (SDEs) to model SGD directly, the approach we take in this work. However, recently, the validity of this approach has been ques- tioned. The main argument, as nicely explained Yaida (28). is that most SDE approximations simultaneously assume that At - 0+, while maintaining that the learning rate y At is finite. The works Simsekli et al. (2 and Li et al. [30] have questioned the correctness of the using the central limit theorem (CLT) to model the gradient noise Gaussian, arguing respectively that the heavy-tailed structure in the gradient noise and the weak dependence between batches leads the CLT to break down. In work, we maintain the CLT assumption holds, which discuss fur-
2

---

Under review as a conference paper at ICLR 2022
ther in appendix A but importantly we avoid the pitfalls of many previous SDE approximations by simultaneously modeling the effect of finite learning rates and stochasticity.
Limiting dynamics. A series of works have applied SDE models of SGD study the limiting dynamics of neural networks. In the seminal work by Mandt et al. [31], the limiting dynamics were modeled with multivariate SGD with assumptions on the geometry of the loss and covariance matrix for the gradient noise. This analysis was extended by Jastrzębski et al. [12] through additional assumptions on the covari- ance matrix to gain tractable insights and applied by Ali et al. (2) to the simpler setting of linear regression, which has a quadratic loss. A different approach was taken by Chaudhari and Soatto [33, which did not formulate the dynamics as an OU process, nor assume directly a structure on the loss or gradient noise. Rather, this analysis studied the same first-order SDE via the Fokker-Planck equation to propose the existence of a modified loss and probability currents driving the limiting dynamics, but did not provide explicit expressions. Our analysis deepens and combines ideas from all these works, where our key insight is to lift the dynamics into phase space. By studying the dynamics of the parameters and their velocities, and by applying the analysis first in the setting of linear regression where assumptions are provably true, we are able to identify analytic expressions and explicit insights which lead to concrete predictions and testable hypothesis.
-hiteinn-hhleeeckk process by combining first-order SDE model for
Stationary dynamics. A different line of work avoids modeling the limiting dynamics of SGD with an SDE and instead chooses to leverage the property of stationarity. These works 28 34
35 36
assume that eventually the probability distribution governing the model parameters reaches stationarity such that the discrete SGD process is simply sampling from this distribution. Yaida [28] used this approach to derive cectuaationddisipatiin parameters and hyperparameters of SGD. Liu et al. (35] used this approach to derive properties for the stationary distribution of SGD with a quadratic loss. Similar to our analysis, this work identifies that the stationary distribution for the parameters reflects a modified loss function dependent on the relationship between the covariance matrix of the gradient noise and the Hessian matrix for the original loss.
relations that link measurable quantities of the
Empirical exploration. Another set of works analyzing the limiting dynamics of SGD has taken a purely empirical approach. Building the intuition that flat minima generalize better than sharp minima, Keskar et al. (37 demonstrated empirically that the hyperparameters of optimization in- fluence the eigenvalue spectrum of the Hessian matrix at the end of training. Many subsequent works have studied the Hessian eigenspectrum during and at the end of training. Jastrzębski et al.
38 Cohen al. [39] studied the dynamics of the top eigenvalues during training. Sagun et al. 40 Papyan [41) Ghorbani et al. (42) demonstrated the spectrum has a bulk of values near zero 433 demonstrated that the learning dynamics plus small number of larger outliers. Gur-Ari et al.
are constrained to the subspace spanned by the top eigenvectors, but found no special properties of the dynamics within this subspace. In our work we also determine that the top eigensubspace of the Hessian plays a crucial role in the limiting dynamics and by projecting the dynamics into this sub- space in phase space, we see that the motion is not random, but consists of incoherent oscillations leading to anomalous diffusion.
4 MODDLLIG SGD AS AN UNDERDAMPED LANGEVIN EQUATION
Following the route of previous works BI I2 B3 studying the limiting dynamics of neural networks, we first seek to model SGD as continuous stochastic process. We consider a net- work parameterized by 0 E Rm. training dataset {xI, .N .NN of size N, and training loss [()) - + اا {(0,x;) with corresponding gradient g ac The state of the network at the t.th step of training is defined by the position vector Og and velocity vector is of the same dimension. The gradient descent update with learning rate 1), momentum B, and weight decay is given by
1 Sug g(Ox) A02,
0k41 0x + 17
(2)
where we initialize the network such that ng and 6g the parameter initialization. In order to understand the dynamics of the network through position and velocity space, which we will refer to as phase space, we express these discrete recursive equations as the discretization of some unknown ordinary differential equation (ODE), sometimes referred to as a modified equation as in 44 20) While this ODE models the gradient descent process even at finite learning rates, it fails to account for the stochasticity introduced by choosing a random batch B of size S drawn uniformly from the set of N training points. This sampling yields the stochastic gradient (I) = Y/4 Etes Ve(0, x.). To model this effect, we make the following assumption:
3

---

Under review as a conference paper at ICLR 2022
Assumption 1 (CLT). We assume the batch gradient is noisy version of the true gradient such that (5( - ( is a Gaussian random variable with mean O and covariance (Z))).
Incorporating this model of stochastic gradients into the previous finite difference equation and applying the stochastic counterparts to Euler discretizations, results in the standard faft-diffusion stochastic differential equation (SDE), referred to as an underdamped Lngevin equation,
0 =
2 (I++) (g(0) + 10 (1 B)u)
di +
dW.
(3)
VnS(I+B)
where W, is a standard Wiener process. This is the continuous model we will study in this work: Assumption (SDE). We assume the underdamped Langevin equation (3) accurately models the trajectory of the network driven by SGD through phase space such that (()l) = 8y, and vnn)) ok-
See appendix for further discussion on the nuances of modeling SGD with an SDE.
5 LINEAR REGRESSION WITH SGD IS AN ORNSTEIN- -ULEEIIE-UULEBEECK PROCESS
Equipped with a model for SGD, we seek to understand its dynamics in the fundamental setting of linear regression, one of the few cases where we have a complete model for the interaction of the dataset, architecture, and optimizer. Let X E R be the input data, Y E RN be the output labels, and 0 E Rd be our vector of regression coefficients. The least squares loss is the convex quadratic loss L(0) 2N ||Y -x0] with gradient g(0) H0-l, bb, where H = x and b XTY, Plugging this expression for the gradient into the underdamped Langevin equation 3 results in the multivariate Censtein-Uhleneee (OU) process,
xd
N
and rearranging terms,
0 1/
[. U/
d
(H XI)
201 2(1-8)
dt + 22
2( 0 (() (1+
dW..
(1++))
D
where A and D are the drift and diffusion matrices re-
spectively, K = S(1 - 32) ( an inverse temperature con-
8 0.9
$q.07
stant, and (H + x is the ridge regression solu-
tion. The solution to an OU process a Gaussian process. By solving for the temporal dynamics of the first and sec- ond moments of the process, we can obtain an analytic expression for the trajectory any time t. In particular, we can decompose the trajectory as the sum of a deter- ministic and stochastic component defined by the first and second moments respectively.
2oo
60 120
B0R 00 oo0
Pashion
Step
B = 0.99
36 6g
Velccty
240
o.0ggng 0.0
120 88 Step
Deterministic component.
Using the form of A we can decompose the expectation as sum of harmonic oscilla- tors in the eigenbasis {q1, ... 4m of the Hessian,
60
Postion
Oscillatory dynamics in linear regression. We train a linear network to perform regression on CIFAR-10 dataset by using MSE loss on the one-hot encoding of the la- bels. We compute the hessian of the loss, as well as its top cigenvectors. The position and velocity trajectories are projected onto the first eigenvector of the hessian and visualized in black. The theoretically derived mean, equa- tion (5), is shown in red. The top and bottom panels demonstrate the effect of varying momentum on the oscillation mode.
Figure 2:
0: U1
q 0
E
(1)
+ bc(
(5)
Here the coefficients (t) and (t) depend on the op- timization hyperparameters 7, X, and the respective eigenvalue of the Hessian P; further explained in appendix FF We verify this expression nearly perfectly matches empirics complex datasets under various hy- perparameter settings as shown in Fig. [
Stochastic component. The cross-covariance of the pro- cess between two points in time 1. 5 &,
Cov Ut
0. V.
=
(B B-c Al Be --) (T--))
(6)
4

---

Under review as a conference paper at ICLR 2022
where B solves the Lyapиnov equation AB + BAT = 2D. In order to gain analytic expressions for B in terms of the optimization hyperparameters, eigendecomposition of the Hessian, and covariance of the gradient noise, must introduce the following assumption:
Assumption (Simultaneously Diagonalizable). We assume the covariance of the gradient noise is spatially independent Σ(0) Σ and commutes with the Hessian HE = EH, therefore sharing common eigenbasis.
6 UNDERSTANDING STATIONARITY VIA THE
-AANCKK EQUATION
The OU process is unique in that it is one the few SDEs which we can solve exactly. As shown in section[5] we were able to derive exact expressions for the dynamics of linear regression trained with SGD from initialization to stationarity by simply solving for the first and second moments. While the expression for the first moment provides an understanding of the intricate oscillatory relationship in the deterministic component of the process, the second moment, driving the stochastic component, much more opaque. An alternative route to solving the OU process that potentially provides more insight is the Fokker-Planc' equation.
The Fokker-Planck (FP) equation is PDE describing the time evolution for the probability distribu- tion of particle governed by Langevin dynamics. For an arbitrary potential o and diffusion matrix D, the Pokker-Plancer equation (under an Itō integration prescription)
0.p V (Vop + V (is (k-1Dp)) Dp));
where p represents the time-dependent probability distribution, and J is a vector field commonly referred to as the probability current. The FP equation is especially useful for explicitly solving for the stationary solution, assuming one exists, of the Langevin dynamics. The stationary solution Pss equivalently V -as- = O. From this second definition we see that there are two distinct settings of stationarity: detailed balance when Jos = 0 everywhere, or broken detailed balance when V Jas = 0 and Jas + 0.
by definition obeys a.p..
For general OU process, the potential is convex quadratic function (xx) = xT Ar defined by the drift matrix A. When the diffusion matrix is isotropic (D x I) and spatially independent (V Gibbs distribution Pas (x) x e --x((z) determined
D the resulting stationary solution by the original loss (xx) and is in detailed balance. Lesser known properties of the OU process arise when the diffusion matrix is anisotropic or spatially dependent 45 46). In this setting the solution is still Gaussian process, but the stationary solution, if it exists, is no longer defined by the Gibbs distribution of the original loss (zx), but actually a modified loss (xx) Furthermore, the stationary solution may be in broken detailed balance leading to non-zero probability current (x). Depending on the relationship between the drift matrix A and the diffusion matrix D the resulting dynamics of the OU process can have very nontrivial behavior.
In the setting of linear regression, anisotropy in the data distribution will lead to anisotropy in the gradient noise and thus an anisotropic diffusion matrix. This implies that for most datasets we should expect that the SGD trajectory is not driven by the original least squares loss, but by a modified loss and converges to stationary solution with broken detailed balance, predicted by Chaudhari and Soatto (33.. Using the explicit expressions for the drift A and diffusion D matrices we can compute analytically the modified loss and stationary probability current,
( v) = 0
0
(
((e)
0
(
0
(s, v) = -QU
Pas.
(8)
where Q is skew-symmetri matrix and U is positive definite matrix defined as,
0 E())
-E()]
n(1 m(+3) E(0)-1 (H + AT)
0
U
(9)
=(0)-
These new fundamental matrices, Q and U, relate to the original drift A and diffusion D matrices through the unique decomposition = (D+ QUU, introduced by Ao [47] and Kwon et al. [48] Using this decomposition we can easily show that B U-1 solves the Lyapunov equation and indeed the stationary solution p.s is the Gibbs distribution defined by the modified loss (0, v) in equation (8). Further, the stationary cross-covariance solved in section 5 reflects the oscillatory dynamics introduced by the stationary probability currents (f,(0, u) in equation (8). Taken together, gain the intuition that the limiting dynamics of SGD in linear regression driven by modified loss subject to oscillatory probability currents.
5

---

Under review as a conference paper at ICLR 2022
7 EVIDENCE OF A MODIFIED LOSS AND OSCILLATIONS IN DEEP LEARNING
Does the theory derived in the linear regression setting (sections 6 help explain the empirical phenomena observed in the non-linear setting of deep neural networks (section ? In order for the theory built in the previous sections to apply to the limiting dynamics of neural networks, must introduce simplifying assumptions on the loss landscape and gradient noise at the end of training:
Assumption 4 (Quadratic Loss). We assume that at the end of training the loss for neural network can be approximated by the quadratic loss L(0) = (0 - 1²) ( training loss Hessian and μ. is some unknown mean vector, corresponding to local minimum.
p)), where H ž is the
Assumption 5 (Covariance Structure). We assume the covariance of the gradient noise is propor- tional to the Hessian of the quadratic loss I(0) = =² H where E R* is some unknown scalar.
Under these simplifications, then the expressions derived in the linear regression setting would apply to the limiting dynamics of deep neural networks and depend only on quantities that we can easily estimate empirically. Of course, these simplifications are quite strong, but without arguing their the- oretical validity, we can empirically test their qualitative implications: (1) a modified isotropic loss driving the limiting dynamics through parameter space, (2) implicit regularization of the velocity trajectory, and (3) oscillatory phase space dynamics determined by the Hessian eigen-structure.
Modified loss. As discussed in section 66] due to the anisotropy of the diffusion matrix, the loss landscape driving the dynamics at the end of training is not the original training loss [()), but a modified loss (0,, v) in phase space. As shown in equation (8). the modified loss decouples into term a that only depends on the parameters 0 and a term V. that only depends on the velocities v. Under assumption 5 the parameter dependent component is proportional the convex quadratic,
Ve x (0 - --) 11 (M AD) + 7(1 BB)
(0 10-10 p).
(10)
This quadratic function has the same mean 1l as the training loss, but a different curvature. Using this expression, notice that when X = O, the modified loss is isotropic in the column space of H, regardless of what the nonzero eigenspectrum of H is. This striking prediction suggests that no matter how anisotropic the original training loss reflected by poor conditioning of the Hessian eigenspectrum - the training trajectory of the network will behave isotropically, since it is driven not by the original anisotropic loss, but modified isotropic loss.
We test this prediction by studying the limiting dynamics of a pre-trained ResNet-18 model with batch normalization that we continue to train on ImageNet according to the last setting of its hyper- parameters [49]. Let 0, represent the initial pre-trained parameters of the network, depicted with the white dot figures and
0.012
0.008
0.004
0.000
0.004
430
-0.012
-.0008
0.004
0.000
0.004
-0.012
-0.008
-0.004
0.000
0.004
-0012
-0.008
0.004
0.000
0.004
q1
-
Training loss fraain
Test lass fret
Modified loss Vp
Figure 3: The training trajectory behaves isotropically, regardless of the training loss. We resume training of pre-trained ResNet-18 model on ImageNet and project its parameter trajec- tory (black line) onto the space spanned by the eigenvectors of its pre-trained Hessian Q1, 930 (with eigenvalue ratio pt/ /P3 ~ 6). We sample the training and test loss within the same 2D subspace and visualize them as a heatmap in the left and center panels respectively. We visualize the mod- ified loss computed from the eigenvalues (p1, Psu) and optimization hyperparameters according to equation (IO) in the right plot. Note the projected trajectory is isotropic, despite the anisotropy of the training and test loss.

---

Under review as a conference paper at ICLR 2022
We estimate the top thirty cigenvectors 91..... 930 of the Hessian matrix H. evaluated at 9, and project the limit- ing trajectory for the parameters onto the plane spanned by the top @i and bottom (30 eigenvectors to maximize the illustrated anisotropy with estimates. We sample the train and test loss in this subspace for a region around the projected trajectory. Additionally, using the hyperparam- eters of the optimization, the eigenvalues P1 and P3o. and the estimate for the mean μ 6, H g+ (g. is the gra- dient evaluated at 0.), we also sample from the modified loss equation (10) the same region. FigureB3 shows the projected parameter trajectory on the sampled train, test and modified losses. Contour lines both the train and test loss exhibit anisotropic structure, with sharper curva- ture along eigenvector qi compared to eigenvector gao. as expected. However, as predicted, the trajectory appears to cover both directions equally. This striking isotropy of the trajectory within highly anisotropic slice of the loss landscape indicates qualitatively that the trajectory evolves in a modified isotropic loss landscape.
3.0
1.5
0.0
-1.5
330
930
5
Q1
Figure 4: Implicit velocity regulariza- tion defined by the inverse Hessian. The shape of the projected velocity tra- jectory closely resembles the contours of the modified loss L
Q1
Velocity
1.5
Implicit velocity regularization. A second qualitative prediction of the theory is that the velocity is regulated by the inverse Hessian of the training loss. Of course there are no explicit terms in either the train or test losses that depend on the velocity. Yet, the modified loss con- tains a component, V, X v H- that only depends on the velocities This additional term can be understood form of implicit regularization on the velocity trajectory. Indeed, when we project the velocity trajectory onto the plane spanned by the and Q30 eigenvectors, as shown in Fig. 4 we see that the trajectory closely resembles the curvature of the inverse Hessian H- The modified loss effectively penalizing SGD for moving in eigenvectors of the Hessian with small eigenvalues. A similar quali- tative effect was recently proposed by Barrett and Dherin 21 as consequence of the discretization error due to finite learning rates.
0.0
-1.5
300 600 Poo 900 208
Step
Position
430
Volacity
0.0
0.8
300 600 900 208508800
0.00 Boogoo44
Step
Position
Figure 5: Phase space oscillations are determined by the eigendecomposi tion of the Hessian. We visualize the projected position and velocity trajecto- ries in phase space. The top and bottom panels show the projections onto @i and 430 respectively. Oscillations at differ- ent rates are distinguishable for the dif- ferent eigenvectors and were verified by comparing the dominant frequencies in the fast Fourier transform of the trajec- tories.
Phase space oscillations. A final implication of the the- ory that at stationarity the network is in broken detailed balance leading to non-zero probability currents flowing through phase space:
U v0) = |- .(- (H + A1) (0 - (0-1))0 Pss. (11)
These probability currents encourage oscillatory dynamics in the phase space planes characterized by the eigenvectors of the Hessian, at rates proportional to their eigenvalues. We consider the same projected trajectory of the ResNet-18 model visualized in figures 3 and 4 but plot the trajectory in phase space for the two eigenvectors Φi and 430 separately. Shown in Fig. 5, we see that both trajec- tories look like noisy clockwise rotations. Qualitatively, the trajectories for the different eigenvectors appear to be rotating at different rates.
The integral curves of the stationary probability current are one-dimensional paths confined to level sets of the modified loss. These paths might cross themselves, in which case they are limit cycles, or they could cover the entire surface the level sets, in which case they are space-filling curves. This distinction depends the relative frequencies of the oscillations, as determined by the pairwise
'To estimate the cigenvectors of H. we subspace iteration, and limit ourselves to 30 eigenvectors to constrain computation time. See appendix H for details.
7

---

Under review as a conference paper at ICLR 2022
ratios of the cigenvalues of the Hessian. For real-world datasets, with large spectrum of incom- mensurate frequencies, we expect to be in the latter setting, thus contradicting the suggestion that SGD in deep networks converges to limit cycles, claimed in Chaudhari and Soatto [3].
UNDERSTANDING THE DIFFUSIVE BEHAVIOUR OF THE LIMITING DYNAMICS
Taken together the empirical results shown in section 7 indicate that many of the same qualitative behaviors of SGD identified theoretically for linear regression are evident in the limiting dynamics of neural networks. Can this theory quantitatively explain the results we identified in section 2
Constant instantaneous speed. As noted in section 2 we observed that at the end of training, across various architectures, the squared norm of the local displacement ||6: remains essentially constant. Assuming the limiting dynamics are described by the stationary solution the expectation of the local displacement
172
Egs [1] = S(1 - B²)
tr (H)
(12)
as derived in appendix G We cannot test this prediction directly as we do not know o² and comput- ing tr(H) is computationally prohibitive. However, we can estimate o²ir(H) by resuming training for a model, measuring the average ||ō, 112,, and then inverting equation 12 Using this single es- timate, we find that for a sweep of models with varying hyperparameters, equation (12) accurately predicts their instantaneous speed. Indeed, Fig. shows an exact match between the empiries and theory, which strongly suggests that despite changing hyperparameters at the end of training, the model remains in the same quadratic basin.
1.00
Eq. (q.(11)
0.75
0.50
0.6-
0.5
0.25
0.00
0.25
3.50
5.75
1.00
250
500
750
1000
0.80
0.85
0.60
0.85
Difusion
Learning
7]
Batch size S
Momentum
|
Figure 6: Understanding how the hyperparameters of optimization influence the diffusion. We resume training of pre-trained ResNet- models on ImageNet using a range of learning rates, batch sizes, and momentum coefficients, tracking 6 and | Starting from the default hyperparam- eters, namely 7 = 1e 4, S = 256, and = 0.9, we vary each one while keeping the others fixed. The top shows the measured || in color, with the default hyperparameter setting highlighted in black. The dotted line depicts the predicted value from equation (12). The bottom row shows the estimated exponent found by fitting line to the || trajectories on log-log plot. The dotted line shows e 1, corresponding to standard diffusion.
Exponent of anomalous diffusion. The expected value for the global displacement under the sta- tionary solution also be analytically expressed in terms of the optimization hyperparameters and the eigendecomposstion of the Hessian as,
n Egs [|7 = S(1 B2)
tr (r( 1+ 21
Σ
Σ (.C.(())
(13)
where (A:): is trigonometic function describing the velocity of harmonic oscillator with damp- ing ratio = (1 -1-30 / + B) (p, + A). see appendix G for details. As shown empirically in section 2 the squared norm || | monotonically increases as power law in the number of steps, suggesting its expectation is proportional to t" for some unknown, constant c. The exponent c de- termines the regime of diffusion for the process. When Brownian diffusion. For e > 1 or < the process corresponds to anomalous super-diffusion or sub-diffusion respectively. Unfortunately, it is not immediately clear how to extract the explicit
1. the process corresponds standard
8

---

Under review as a conference paper at ICLR 2022
exponent from equation
However, by exploring the functional form of (k) and its relation- ship to the hyperparameters of optimization through the damping ratio Çi, we can determine overall trends in the diffusion exponent c.
13
Akin to how the exponent c determines the regime of diffusion, the damping ratio & determines the regime for the harmonic oscillator describing the stationary velocity-velocin correlation in the 7th eigenvector of the Hessian. When & = 1, the oscillator is critically damped implying the velocity correlations converge zero as quickly as possible. In the extreme setting of Ci(k) - for all Seipssppliii
1, k. then equation (3) simplifies to standard Brownian diffusion, Ee. [[[Ael]]] x t. When & 1,
the oscillator is overdamped implying the velocity correlations dampen slowly and remain positive even over long temporal lags. Such long lasting temporal correlations in velocity lead to faster global displacement. Indeed, in the extreme setting of ()))) = 1 for all 1, k, then equation (1) simplifies to a form of anomalous super-diffusion, E.. ||| x f.. When & < 1. the oscillator is underdamped implying the velocity correlations will oscillate quickly between positive and negative values. Indeed, the only way equation (13) could describe anomalous sub-diffusion is if (k) took
on negative values for certain 1, k.
Using the same sweep of models described previously, we can empirically confirm that the opti- mization hyperparameters each influence the diffusion exponent c. As shown in Fig. 6 the learning rate, batch size, and momentum can each independently drive the exponent e into different regimes of anomalous diffusion. Notice how the influence of the learning rate and momentum on the dif- fusion exponent e closely resembles their respective influences on the damping ratio G. Interest- ingly, a larger learning rate leads to underdamped oscillations, and the resultant temporal velocities' anti-correlations reduce the exponent of anomalous diffusion. Thus contrary to intuition, a larger learning rate actually leads to slower global transport in parameter space. The batch size on the other hand, has no influence on the damping ratio, but leads to an interesting, non-monotonic influence on the diffusion exponent. Overall, the hyperparameters of optimization and eigenspectrum of the Hessian all conspire to govern the degree of anomalous diffusion at the end of training.
9 DISCUSSION
Through combined empirics and theory based on statistical physics, we uncovered an intricate in- terplay between the optimization hyperparameters, structure in the gradient noise, and the Hessian matrix at the end of training.
Significance. The significance of our work lies in (1) the identification/verifcati of multiple empirical phenomena (constant instantaneous speed, anomalous diffusion in global displacement, isotropic parameter exploration despite anisotopic loss, velocity regularization, and slower global parameter exploration with faster learning rates) present in the limiting dynamics of deep neural networks, (2) the emphasis on studying the dynamics in velocity space in addition to parameter space, and (3) concrete quantitative as well as qualitative predictions of an SDE based theory that we empirically verified in deep networks trained on large scale datasets (indeed some of the above nontrivial phenomena were predictions of this theory). Of course, these contributions directly build upon series of related works studying the immensely complex process of deep learning. To this end, we further clarify the originality of our contributions with respect to some relevant works.
Originality. The empirical phenomena we present provide novel insight with respect to the works of Wan et al. [13], Hoffer et al. [18), and Chen et al. 15] We observe that all parameters in the network (not just those with scale symmetry) move at a constant instantaneous speed at the end of training and diffuse anomalously at rates determined by the hyperparameters of optimization. In contrast to the work by Liu et al. 35), we modeled the entire SGD) process as an OU process which allows to provide insight into the transient dynamics and identify oscillations in parameter and velocity space. We build on the theoretical framework used by Chaudhari and Soatto [33] and provide explicit expressions for the limiting dynamics in the simplified linear regression setting and conclude that the oscillations present in the limiting dynamics are more likely to be space-filling curves (and not limit cycles) in deep learning due to many incommensurate oscillations.
Overall, by identifying key phenomena, explaining them in simpler setting, deriving predictions of new phenomena, and providing evidence for these predictions at scale, we are furthering the scientific study of deep learning. We hope our newly derived understanding of the limiting dynamics of SGD, and its dependence on various important hyperparameters like batch size, learning rate, and momentum, can serve as a basis for future work that can turn these insights into algorithmic gains.

---

Under review as a conference paper at ICLR 2022
REFERENCES
[1] Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pages 29-53. Springer, 1996.
[2] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli,, and Jascha Sohl-Dickstein. Deep infor- mation propagation. arXiv preprint 161 2016.
[3] Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.0016. 2017.
[4] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint
2018.
[5] Jaehoon Lee, Lechao Xiao, Samuel S Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. arXiv preprint arXiv: 1902.06720, 2019.
[6] Mei Song, Andrea Montanari, and P Nguyen. A mean field view of the landscape of two-layers neural networks. Proceedings of the National Academy of Sciences, 155667665-7767, 2018.
[7] Grant M Rotskoff and Eric Vanden-Eijnden. Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 7 7146-7155, 2018.
[8] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. arXiv preprint arXiv: 18050058055 2018.
[9] Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wil- son. Loss surfaces, mode connectivity, and fast ensembling of dnns. In Proceedings the the 32nd International Conference on Neural Information Processing Systems, pages 8803-8812, 2018.
[10] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv: 2018.
[11] Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. simple baseline for bayesian uncertainty in deep learning. arXiv preprint arXiv:1902.0244 1902.02476, 2019.
[12] Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint 17X1v:1711.04623, 2017.
[13] Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics of deep neural networks with batch normalization and weight decay. arXiv preprint 9884499006.08419 2020.
[14] Marco Baity-Jesi, Levent Sagun, Mario Geiger, Stefano Spigler, Gérard Ben Arous, Chiara Cammarota, Yann LeCun, Matthieu Wyart, and Giulio Biroli. Comparing dynamics: Deep neural networks versus glassy systems. In International Conference on Machine Learning, pages 314 323. PMLR, 2018.
[15] Guozhang Chen, Cheng Kevin Qu, and Pulin Gong, Anomalous diffusion dynamics of learning in deep neural networks. arXiv preprint arXiv:2009.10588 2020.
[16] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. Neural Information Processing Systems Workshop, 2017.
[17] Zhiyuan Li. Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with tradi- tional optimization analyses: The intrinsic learning rate. In Advances in Neural Information Processing Systems, volume 33, pages 14544 14555. Curran Associates, Inc., 2020.
10

---

Under review as a conference paper at ICLR 2022
[18] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the gen- eralization gap in large batch training of neural networks. Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[19] Ning Qian. On the momentum term gradient descent learning algorithms. Neural networks, 12(1):145-151, 1999.
[20] Daniel Kunin, Javier Sagastuy-Brens Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dy- namies. arXiv preprint urXiv:2012.04728. 2020.
[21] David GT Barrett and Benoit Dherin. Implicit gradient regularization. arXiv preprint
DXXix2009.11162 2020.
[22] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pages 2101-2110 10. PMLR, 2017.
[23] Samuel L Smith, Benoit Dherin, David GT Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. arXiv preprint arXiv:2101.12176, 2021.
[24] Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynam- ical mean-field theory for stochastic gradient descent in gaussian mixture classification. arXiv preprint arXiv:2006.06099 2020.
[25] Stefano Sarao Mannelli, Eric Vanden-Eijnden, and Lenka Zdeborová. Optimization and gen- eralization of shallow neural networks with quadratic activation functions. arXiv preprint arXiv:2006.15459 2020.
[26] Francesca Mignacco. Pierfrancesco Urbani, and Lenka Zdeborová. Stochasticity helps to nav- igate rough landscapes: comparing -eeeenttddeceet-baseee algorithms in the phase retrieval problem. Machine Learning: Science and Technology, 2021.
[27] Stefano Sarao Mannelli and Pierfrancesco Urbani. Just momentum: Analytical study of momentum-based acceleration methods in paradigmatic high-dimensional non-convex prob- lems. arXiv preprint 102. 2021.
[28] Sho Yaida. Fluctuation Fluctuation-dissipation relations for stochastic gradient descent. arXiv preprint arXiv:1810.000044 2018.
[29] Unut Simsekli, Levent Sagun, and Mert Gurbuzzaalbban A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827-5837. PMLR, 2019.
[30] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochas- tic differential equations (sdes). arXiv preprint arXiv:2102.124700 2021.
[31] Stephan Mandt, Matthew Hoffman, and David Blei. A variational analysis of stochastic gra- dient algorithms. In International conference machine learning, pages 2016.
PMLR,
[32] Alnur Ali. Edgar Dobriban, and Ryan Tibshirani. The implicit regularization of stochastic gradient flow for least squares. In International Conference on Machine Learning, pages 233- 244. PMLR, 2020.
[33] Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational infer- ence, converges to limit cycles for deep networks. In 2018 Information Theory and Applica- tions Workshop (ITA), pages 1-10. IEEE, 2018.
[34] Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. In Advances in Neural Information Processing Systems, volume 32, Curran Associates, Inc., 2019.

---

Under review as a conference paper at ICLR 2022
[35] Kangqiao Liu, Liu Ziyin, and Masahito Ueda. Noise and fluctuation of finite learning rate stochastic gradient descent. In International Conference on Machine Learning, pages 7045- 7056. PMLR, 2021.
[36] Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. On minibatch noise: Discrete- time sgd, overparamettizattoo, and bayes. arXiv preprint arXiv:2102.05375, 2021.
[37] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp min- ima. arXiv preprint arXiv: 00866,990466,... 2016.
[38] Stanistaw Jastrzebski, 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 2018.
length. arXiv preprint arXivv/807..00
[39] Jeremy M Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gra- dient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00m 00065, 2021.
[40] Levent Sagun, Utku Evei, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint 0rXiv:1706.04454 2017.
[41] Vardan Papyan. The full spectrum of deepnet hessians scale: Dynamics with sgd training and sample size. arXiv preprint arXiv:
2018.
[42] Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net opti- mization via hessian eigenvalue density. In International Conference on Machine Learning. pages 2232-2241. PMLR, 2019.
[43] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprint arXiv: /8/2.04754 2018.
[44] Nikola B Kovachki and Andrew M Stuart. Analysis of momentum methods. arXiv preprint 2019.
[45] Crispin W Gardiner et al. Handbook of stochastic methods, volume 3. springer Berlin, 1985.
[46] Hannes Risken. Fokker-planck equation. In The Fokker-Planck Equation, pages 63-95. Springer, 1996.
[47] Ping Ao. Potential in stochastic differential equations: novel construction. Journal of physics
A: mathematical and general, 37(3):255, 2004.
[48] Chulan Kwon, Ping Ao, and David Thouless. Structure of stochastic dynamics near fixed points. Proceedings of the National Academy of Sciences, (02(37): 2005.
[49] Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recogni- tion, pages 770-778, 2016.
[50] Tomaso Poggio, Kenji Kawaguchi, Qianli Liao, Brando Miranda, Lorenzo Rosasco, Xavier Boix, Jack Hidary, and Hrushikesh Mhaskar. Theory of deep learning ii:: explaining the non- overfitting puzzle. arXiv preprint 1=0ivv1801.001 2017.
[51] Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07476 2016.
[52] Valentin Thomas, Fabian Pedregosa, Bart Merriënboer, Pierre-Antoine Manzagol, Yoshua Bengio, and Nicolas Le Roux. On the interplay between noise and curvature and its effect on optimization and generalization. In International Conference on Artificial Intelligecce and Statistics, pages 3503-3513, PMLR, 2020.
[53] Richard Jordan, David Kinderlchrer, and Felix Otto. The variational formulation of the fokker- planck equation. SIAM journal on mathematical analysis, 29(1): 1-17, 1998.
[54] Richard Jordan, David Kinderlehrer, and Felix Otto. Free energy and the fokker-planch equa- tion. Physica D: Nonlinear Phenomena, 107(2-4):26--271, 1997.
12

---

Under review as a conference paper at ICLR 2022
A
MODELING SGD WITH AN SDE
As explained in section in order to understand the dynamics of stochastic gradient descent we build continuous Langevin equation phase space modeling the effect of discrete updates and stochastic batches simultaneously.
A.1 MODELING DISCRETIZATION
To model the discretization effect we assume that the system of update equations (2) is actually discretization of some unknown ordinary differential equation. To uncover this ODE, we combine the two update equations in (2) by incorporating previous time step 0x-1, and rearrange into the form of finite difference discretization, as shown in equation (??). Like all discretizations, the Euler discretizations introduce error terms proportional the step size, which in this case is the learning rate 7). Taylor expanding 0x+1 and f around 02. its casy to show that both Euler discretizations introduce second- -oder error term proportional to 2 700
0k+1 O2 = n O(1²²,
O2 02 = 2
" 0
7
1
Notice how the momentum coefficient / E [0, 1] regulates the amount of backward Euler incorpo- rated into the discretization. When 3 = 0. we remove all backward Euler discretization leaving just the forward Euler discretization. When = 1, we have equal amounts of backward Euler as forward Euler resulting a central second-order discretization as noticed [T9].
A.2 MODELING STOCHASTICITY
order to model effect of stochastic batches, we first model a batch gradient with the following assumption:
Assumption (CLT). We assume the batch gradient a noisy version of the true gradient such that (o) g(0) is Gaussian random variable with mean 0 and covariance (0)).
The two conditions needed for the CLT to hold are not exactly met in the setting of SGD. Indepen- dent and identically distributed. Generally perform SGD by making complete pass through the entire dataset before using sample again which introduces a weak dependence between samples. accurately models the dependence between
While the covariance matrix without replacement samples within a batch, it fails to account for the dependence between batches. Finite variance. A different line of work has questioned the Gaussian assumption entirely because of the need for finite variance random variables. This work instead suggests using the generalized central limit the- orem implying the noise would be a heavy-tailed a-stable random variable (29] Thus, the previous assumption is implicitly assuming the i.idd and finite variance conditions apply for large enough datasets and small enough batches.
Under the CLT assumption, we must also replace the Euler discretizations with Suler-Maruyee discretizations. For a general stochastic process, d.x, = pull adW,, the Maruuaaaa method extends the Huler method for ODEs to SDEs, resulting in the update equation X = X + Alp
/ Vttag, where & (0 1). Notice, the key difference is that if the temporal step size is At n). then the noise is scaled by the square root /T.. fact, the main argument against modeling SGD with SDE, nicely explained in Yaida (28), is that most SDE approximations simultaneously assume that At + 0+, while maintaining that the square root of the learning rate v is finite. However, by modeling the discretization and stochastic effect simultaneously can avoid this argument, bringing us to our second assumption:
Assumption 2 (SDE). We assume the underdamped Langevin equation (3) accurately models the trajectory of the network driven by SGD through phase space such that ((nk) Re 02 u(jk) se un.
This approach of modeling discretization and stochasticity simultaneousty is called stochastic mod- ified equations, as further explained in Li et al. 22
"The difference between forward Euler and backward Euler discretization is second-order central dis- cretization, 6k-1 -0
0g
Pp-0 0
00.+1 +0--
n7 O
1
13

---

Under review as a conference paper at ICLR 2022
B STRUCTURE IN THE COVARIANCE OF THE GRADIENT NOISE
As we've mentioned before, SGD introduces highly structured noise into an optimization process, often assumed to be essential ingredient for its ability to avoid local minima.
Assumption (Covariance Structure). We assume the covariance of the gradient noise is propor- tional to the Hessian of the quadratic loss Σ(0) o² H where R* is some unknown scalar.
In the setting of linear regression, this a very natural assumption. If we assume the classic genera- tive model for linear regression data 1/+ 2: 1+++0 are where, 0 Rd is the true model and N(,, 1), then provably E()) = x² H.
Proof. We estimate the covariance Σ(0) N E 9ig! ggt. Near stationarity gg* ΣN
N 2i=1 9:9 and thus,
(0) gig
Under the generative model y; x:T0 + where (0, 1) and R+, then the gradient gi is
gi (x. (0 0) -
and the matrix 9:9 is
9.9. (x] - 0)
Assuming 0 R f at stationarity, then (x: 0) oc)² ~ 0². Thus,
N प.o.
d² x =
(0)
~
N
Also notice that weight decay is independent of the data batch and thus simply shifts the gradient distribution, but leaves the covariance of the gradient noise unchanged.
While the above analysis is in the linear regression setting, for deep neural networks it is reasonable to make the same assumption. See the appendix of Jastrzębski et al. [12 for discussion on this assumption in the non-linear setting.
Recent work by Ali et al. [32 also studies the dynamics of SGD (without momentum) in the setting of linear regression. This work, while studying the classic first-order stochastic differential equation, made point to not introduce point that even in the setting of linear regression, a constant covariance matrix will fail to capture the actual dynamics. To illustrate this point they consider the univariate responseless least squares problem.
assumption
the diffusion matrix. In particular, they make the
minimize BER
(aa0)²
2n
1=1
As they explain, the SGD update for this problem would be
k II ( Ex:))) SEB
Ox+1 0s 17
Σ
T 02
from which they conclude for small enough learning rate then with probability one 9% + They contrast this with the Amstein-Uhlenbeck process given by constant covariance matrix where while the mean for 0g converges to zero its variance converges to positive constant. So process with a constant covariance matrix fails to capture the updates of SGD? In many ways this problem is not simple example, rather pathological edge case. Consider the generative model that would give rise to this problem,
this discrepancy evidence that an
Oz + 0g 0.
In otherwords, the true model Ø 0 and the standard deviation for the noise = 0. This would imply by the assumption used in our paper that there would be zero diffusion and the resulting SDE would simplify to deterministic ODE that exponentially converges to zero.
14

---

Under review as a conference paper at ICLR 2022
C A QUADRATIC LOSS AT THE END OF TRAINING
Assumption (Quadratic Loss). We assume that at the end of training the loss for a neural network
can be approximated by the quadratic loss [(0) (0 1T) [H) (0 p.), where 0 is the
training loss Hessian and μ is some unknown mean vector, corresponding to local minimum.
This assumption has been amply used in previous works such Mandt al. (31, Jastrzębski al.
[12]. and Poggio et al. [50]. Particularly, Mandt et al. [3 discuss how this assumption makes sense for smooth loss functions for which the stationary solution the stochastic process reaches deep
local minimum from which it is difficult to escape.
well-studied fact, both empirically and theoretically, that the Hessian is low-rank near local minima as noted by Sagun et al. [51]. and Kunin et al. (20]. This degeneracy results in flat directions of equal loss. Kumin al. [0) discuss how differentiable symmetries, architectural features that keep the loss constant under certain weight transformations, give rise these flat directions. Im- portantly, the Hessian and the covariance matrix share the same null space, and thus we can always restrict ourselves to the image space of the Hessian, where the drift and diffusion matrix will be full rank. Further discussion on the relationship between the Hessian and the covariance matrix can be found in Thomas al. [52
is also a well known empirical fact that even at the end of training the Hessian can have negative eigenvalues 4 This empirical observation is at odds with assumption that the Hessian is positive semi-definite ž 0. Further analysis is needed to alleviate this inconsistency.
15

---

Under review as a conference paper at ICLR 2022
D SOLVING AN ORNSTEIN-UHLENBEG PROCESS WITH ANISOTROPIC NOISE
We will study the multivariate OrnsteinUUhlenbeck process described by the stochastic differential equation
dX; = A(u Xx)dtt + V2w-IDdW,
Xo = Io.
(14)
where E S is a positive definite drift matrix, μ E RT a mean vector, x E RR is some positive constant, and D) STT is positive definite diffusion matrix. This OU process is unique in that it one of the few SDES we can solve explicitly. We can derive an expression for X7 as,
XT =r=c e - xo + (1 (I- e--e-A) a [ 2(0
(15)
Proof. Consider the function f(t, = EA where e a matrix exponential. Then by Itô's
Lemma we can evaluate the derivative of Xt) as
Xt) (Ac* ex A(μ xx)) dt 2x
= AA + e / 2
Integrating this expression from t = 0 10 T gives
{- AcA pdt +
T
((T,x))-(00.
eat 2VVi--DDVV
0
e xx x0 (T - 1) a (' 22
which rearranged gives the expression for Xy-
From this expression it is clear that Xy is Gaussian process. The mean of the process is
E [ = e AT'
To + (I - e e μ.
(16)
and the covariance and cross-covariance of the process are
Var(Xr) = =
Alt De AT(-T)
(17)
Cov(Xr, Xx) 11
A 2ee 1T (Tt-5))
(18)
These last two expressions are derived by Itô Isometry
D.1 THE LYAPUNOV EQUATION
We can explicitly solve the integral expressions for the covariance and cross-covariance exactly by solving for the unique matrix B E ST that solves the Lyapumov equation,
AB BAT 2D.
(19)
If B solves the Lyapunov equation, notice d
e A(t- T) T) ABe AT +
- A(t-T)2De (T=-S)
Using this derivative, the integral expressions for the covariance and cross-covariance simplify as,
Var(Xr) = k -1 e AT Be Be Be=A")
(20)
Cov(XT,X) (x) = -1 B - AT Be (a-ATT) e (T
(21)
where we implicitly assume T S.
"It's Lemma states that for any Itō drift- difft-diffuson process dX, = midt + o.dWs and twice differentiable
scalar function ft, x), then df tt X.) (f: utfa of faz + oof.
Itō Isometry states for any standard Itô process X., then E (IC x(dw:))) = [ x2dt]
16

---

Under review as a conference paper at ICLR 2022
D.2 DECOMPOSING THE DRIFT MATRIX
While the Lyapunov equation simplifies the expressions for the covariance and cross-covariance, does explain how actually solve for the unknown matrix B. Following method proposed by Kwon et al. [48) we will show how to solve for B explicitly in terms of the drift A and diffusion D.
The drift matrix A can be uniquely decomposed as,
(22) where our symmetric diffusion matrix, Q is skew-symeetric matrix (i.e. Q and U is positive definite matrix. Using this decomposition, then B U-1, solves the Lyapunov equation.
A= (
Proof. Plug B = U-1 into the left-hand side of equation 19
AU + --- AT (D
(D + +
= 2D
Here used the symmetry of A, D,U and the skew-symmetry of Q
All that is left is to do is solve for the unknown matrices Q and U. First notice the following identity, AD DA = QA + AQ
(23)
Proof. Multiplying A (D Q)U on the right by (D-Q) gives,
A(D - Q) (D + QD (D + Q) which rearranged and using A AT gives the desired equation.
Let VAVT be the eigendecomposition of A and define the matrices D = VT DV and @ = VTQV. These matrices observe the following relationship.
Q:j
X3 Dig.
(24)
P: + X3
Proof. Replace A in the previous equality with its eigendecompsoition. VAVTD DVAVT QVAVT + VAVTQ.
Multiply this equation on the right by V and on the left by VT,
Aδ DA = QA + AQ.
Looking at this equality element-wise and using the fact that diagonal gives the scalar equality for any j.
(A: 1,) (A+ 2
which rearranged gives the desired expression.
Thus, Q and are given by,
Q VQVT
U UD++Q))~14
(25)
This decomposition always holds uniquely when A, > 0, as Y= exists and (D + Q) is invert- ible. See [48] for a discussion on the singularities of this decomposition.
A,+A;
D.3 STATIONARY SOLUTION
Using the Lyapunov equation and the drift decomposition, then X, pr, where Pr N ( AAT To + (1 e-T) u, K
(26)
In the limit as T + oc, then Similarly, the cross-covariance converges to the stationary cross- covariace,,, CXv, x- (T-S),
and pr - Pas where
Pas = N ( -U-1).
(27)
(28)
17

---

Under review as a conference paper at ICLR 2022
E A VARIATIONAL FORMULATION OF THE OU PROCESS WITH ANISOTROPIC NOISE
In this section we will describe alternative, variational, route towards solving the dynamics of the OU process studied in appendix [
Let + R R be arbitrary, non-negative potential and consider the stochastic differential equation describing the Langevin dynamics of particle in this potential field.
dx:- (xt + 22 D(,))
Xo 20,
(29)
where D(X;) is an arbitrary, To E RT is the particle's 's initial position. The Fokker-Planck equation describes the time evolution for the probability distribution of the particle's s position such that p(x, 1) P(X, = x). The FP equation the partial differential equatior
diffusion matrix, λ is temperature constant, and
0,p v (V
(D(X.)p)),
p(z, 0) 8(1o),
(30)
where V- denotes the divergence and 6(x0) is a dirac delta distribution centered the initialization To. To assist in the exploration of the FP equation we define the vector field,
[æx t) = - vo(X))p - V (D(X.)p),
(31)
which is commonly referred as the probability current. Notice, that this gives an alternative ex- pression for the FP equation, 0,p- = VJ, demonstrating that (zz t) defines the flow of probability mass through space and time. This interpretation is especially useful for solving for the stationary solution Pas. which is the unique distribution that satisfies,
O.Pss =
JAs = O,
(32)
where J,s is the probability current for Pss- The stationary condition can be obtained in two distinct ways:
1. Detailed balance. This is when (x) = for all x E Q. This is analogous to reversibility for discrete Markov chains, which implies that the probability mass flowing from a state to any state j is the same as the probability mass flowing from state j to state i.
2. Broken detailed balance. This is when V [z.((()) = but J.s(x) (x) of 0 for all Ω. This is analogous to inrevrriiilliyy for discrete Markov chains, which only implies that the total probability mass flowing out of state equals to the total probability mass flowing into state i.
The distinction between these two cases is critical for understanding the limiting dynamics of the process.
E.1 THE VARIATIONAL FORMULATION OF THE FOKKER- PLANEK EQUATION WITH ISOTROPIC DIFFUSION
We will now consider the restricted setting of standard, isotropic diffusion (D = 7). It is easy enough to check that in this setting the stationary solution is
(x4(2)
Pas (x)
7. = Jn
xx) da,
(33)
Z
where P.* is called Gibbs distribution and Z is the partition function. Under this distribution, the stationary probability current is zero (Jaa (x) = 0) and thus the process is in detailed balance. Interestingly, the Gibbs distribution pas has another interpretation the unique minimizer of the the Gibbs free energy functional,
= [ - HH())
(34)
where E [ is the expectation of the potential & under the distribution and H())
- SS (orr))og(p(x)))) the Shannon entropy of p.
"This PDE is also known the Forward Kolmogorov equation.
18

---

Under review as a conference paper at ICLR 2022
Proof. To prove that indeed Pas is the unique minimizer of the Gibbs free energy functional, consider the following equivalent expression
/ p(a (x))((z)d + -1 / p(z))oo(((x))ddz
F(p) = 22
/ p(a) (log(p(1;) Log (ps (x))) dr 151 1og(2)
=
= (p Pss) x log((Z)
From this expressions, it is clear that the Kullback-Leibler divergence is uniquely minimized when p Pss-
In other words, with isotropic diffusion the stationary solution Pea can be thought of as the limiting distribution given by the Fokker-Planck equation the unique minimizer of energetic-entropic functional.
Seminal work by Jordan al. [53] deepened this connection between the Fokker-Planck equation and the Gibbs free energy functional. In particular, their work demonstrates that the solution p(z, t) the Fokker-Plance equation is the Wasserstein gradien! flow trajectory the Gibbs free energy functional.
Steepest descent is always defined with respect distance metric. For example, the update equa- tion, 2k+1 = Tg 17 (xg.). for classic gradient descent on a potential ((x). can be formu- lated as the solution to the minimization problem Fg++1 argmin, // (x) fd(2, 2x)² where d(x, 2x) = ||æ - Ix || is the Euclidean distance metric. Gradient flow is the continuous-time limit of gradient descent where we take 77 + 0+ Similarly. Wasserstein gradient flow is the continuous time limit of steepest descent optimization defined by the Wasserstein metric. The Wasserstein metric is a distance metrie between probability measures defined as,
W P2) inf pel()). J x"*R"
|a yl y1²r(dx, dy).
(35)
where li and #2 are two probability measures on R" with finite second moments and I( H2) defines the set of joint probability measures with marginals μi and μ12. Thus, given an initial distri- bution and learning rate y/. we can use the Wasserstein metric to derive sequence of distributions minimizing some functional in the sense of steepest descent. the continuous limit as 7) - 0+ this sequence defines continuous trajectory of probability distributions minimizing the functional. Jordan et al. [54] proved, through the following theorem, that this process applied to the Gibbs free energy functional converges to the solution the Fokker-Planck equation with the same initializa- tion:
Theorem (JKO). Given an initial condition po with finite second moment and an nj >> define the iterative scheme Py with iterates defined by
Ps = argmin,," (E [E) K- H()) + W (p.p"=))
As 17 + 0+, then Py + weakly in L where p is the solution to the Fokker-Planck equation with the same initial condition.
See 54 for further explanation and 53 for a complete derivation.
E.2 EXTENDING THE VARIATIONAL FORMULATION THE SETTING OF ANISOTROPIC DIFFUSION
While the JKO theorem provides very powerful lens through which to view solutions to the Fokker- Planck equation, and thus distributions for particles governed by Langevin dynamics, it only applies in the very restricted setting of isotropic diffuuion.. In this section we will review work by Chaudhari and Soatto (33 extending the variational interpretation to the setting of anisotropic diffusion.
Consider when D(X;) is an anisotropic, sratially-depende original Gibbs distribution given in equation (33) does not necessarily satisfy the stationarity dition equation (32). In fact, it is not immediately clear what the stationary solution is or the dynamics have one. Thus, Chaudhari and Soatto 3 make the following assumption:
diffusion matrix. In this setting, the
19

---

Under review as a conference paper at ICLR 2022
Stationary Assumption. Assume there exists unique distribution Pas that is the stationary solution to the Bokker-PPanncee equation irregardless of initial conditions.
Under this assumption we can implicitly define the potential - Using this modified potential we can express the stationary solution as a Gibbs distribution,
Pas (1) x e -V()))
(36)
Under this implicit definition we can define the stationary probability current as Jes(a) (x) where
j(x)= = -v -
D(x) + (z))
(37)
The vector field j(x) reflects the discrepancy between the original potential & and the modified
potential v according to the diffusion (xx). Notice that in the isotropic case, when D(x) = I,
then 1 V and j(x) O. Chaudhari and Soatto (33 introduce another property j(x) through assumption,
Conservative Assumption. Assume that the force j(x) is conservative (i.e. V fj((z) 0).
Using this assumption, Chaudhari and Soatto 33 extends the variational formulation provided by the JKO theorem the anisotropic setting,
Theorem 2 (CS), Given an initial condition Po with finite second moment, then the energetic- entropic functional
F(p) [w))) K
monotonically decreases throughout the trajectory given by the solution to the Fokker-Plance equa- tion with the given initial condition.
In other words, the Fokker-Plant -lankk equation (30) with anisotropic diffusion can be interpreted as mini- mizing the expectation of modified loss V while being implicitly regularized towards distributions that maximize entropy. The derivation requires we assume a stationary solution Pss exists and that the force j(x) implicitly defined by Pas is conservative. However, rather than implicitly define (x) and j(x) through assumption, we can explicitly construct a modified loss & (x) such that the resulting j(r) satisfies certain conditions, then the stationary solution exists and the variational formulation will apply as well. We formalize this statement with the following theorem,
Theorem 3 (Explicit Construction). If there exists potential (x)) such that either j(x) 0 or V (z)(=) and (xx 1 j(x), then Pas is the Gibbs distribution & e (x)(()) and the variational formulation given in Theorem 2 @apllies.
E.3 APPLYING THE VARIATIONAL FORMULATION TO THE OU PROCESS
Through explicit construction we now seek to find analytic expressions for the modified loss (x)
and force j(z) hypothesised by Chaudhari and Soatto 33 in the fundamental setting of an OU
process with anisotropic diffusion, as described in section D We assume the diffusion matrix is anisotropic, but spatially independent, V D(x) = 0. For the OU process the original potential generating the drift is
(xx) = (x- 4x
(38)
Recall, that in order to extend the variational formulation we must construct some potential (x) such that V j(x)= 0 and Vy 1 j(r).. is possible to construct (xx) using the unique decompo- sition of the drift matrix = (D + Q)U discussed in appendix D Define the modified potential,
(x( (z 1) 4 14).
(39)
Using this potential, the force j(x)) is
j(x) - p)) + (z))) - -H) H).
(40)
Notice that (xx) is conservative, V V -QU (z- -= because Q skew-symmetric.
Additionally, (zx) is orthogonal, (V = (x)) (z: μ4) = 0, again because Q is skew-symmetrie. Thus, we have determined a modified potential (xx that results in a conservative
orthogonal force jxx) satisfying the conditions for Theorem 3 Indeed the stationary Gibbs distri-
bution given by Theorem agrees with equation (77) derived via the first and second moments in
appendix[D
(x) x N
20

---

Under review as a conference paper at ICLR 2022
In addition to the variational formulation, this interpretation further details explicitly the stationary probability current, Jea(2) = (ss) and whether or not the the stationary solution is in broken detailed balance.
F EXPLICIT EXPRESSIONS FOR THE OU PROCESS GENERATED BY SGD
We will now consider the specific OU process generated by SGD with linear regression. Here we repeat the setup as explained section 5
Let X E Rxx Y E RN be the input data, output labels respectively and 0 E Rd be our vector of regression coefficients. The least squares loss the convex quadratic loss [ 2N
X0I22 | with gradient g = HО- b. where H = x and b = x Plugging this expression for the gradient into the underdamped Langevin equation (3), and rearranging terms, results in the multivariate -hieebe-khleebeck (OU) process,
1 ||
N
0c Ut
d
A
dt + VDW, DdW,,
(41)
where A and D are the drift and diffusion matrices respectively,
A =
X))
2(1-8) +(++3)
D
0
2(-9))
(42)
(1+3)
11+
R. = 5(1 - B2) is a temperature constant, and 1 HH + AT) -b is the ridge regression solution.
F.1 SOLVING FOR THE MODIFIED LOss AND CONSERVATIVE FORCE
In order to apply the expressions derived for a general OU process in appendix D and we must first decompose the drift as A = (D + Q)U. Under the simplification _ = o² H discussed in appendix B then the matrices Q and U, as defined below, achieve this,
0 O²H
o²H
2 al H
AT)
0
Q
U=
(43)
0
0
1 H
Using these matrices we can now derive explicit expressions for the modified loss (0, v) and con- servative force j(0, v). First notice that the least squares loss with L2 regularization is proportional to the convex quadratic,
(0) (0 T + (0) p).
(44)
The modified loss is composed of two terms, one that only depends on the position,
H H AT)
) (0 14)
V v = (0 1
(45)
1)(1 + 3)
and another that only depends on the velocity,
H o²
(e(u) TT
v.
(46)
The conservative force j(8, v) is
w m(++8] HH + Al) (0 m
j(0 v)
(47)
and thus the stationary probability current is (0, = j(0, v u)p
F.2 DECOMPOSING THE TRAJECTORY INTO THE EIGENBASIS OF THE HESSIAN
As shown in appendix D the temporal distribution for the OU process at some time T 2 is,
Pr (1= . 1e-AT AT
0g
+ (1-e- A7)
p
K (U-I- (0-1-e-ATU--e--
vo
Here we will now use the eigenbasis {q..... 4m of the Hessian with eigenvalues {pI...., Pm] to derive explicit expressions for the mean and covariance of the process through time.
21

---

Under review as a conference paper at ICLR 2022
Deterministic component. We can rearrange the expectation
!
100
E
+
AT
0
vo
Notice that the second, time-dependent term is actually the solution to the system of ODEs
0 1
0
with initial condition [00 - vo]". This system of ODEs can be block diagonalized by factorizing OSOT where O is orthogonal and S is block diagonal defined
Q1
0
Zm
0
((T++335 (p1 + A)
2(1-8) +[1++8]
O
S
Q1
4m
1
m[++3] Pp A)
201 11-91 9) 7 [++0 D)
In otherwords in the plane spanned by [q: 0)* and
q:]
the system of ODEs decouples into the
2D system
0
a
6.
o[I+B) (p= +
201
201-8)
8)
b.
(113)
This system has simple physical interpretation as a damped harmonic oscillator, we let b, = di. then unravel this system into the second order ODE
0 2 1-
2
B) à; +
(p: + AJa; 0
71+
711 + B)
which is in standard form (i.e. F + 2 + w²r for 7
*1++a and wi V MI+カ (pi X).
1-8
2
Let (0) (8. n. 4.) and (0) (vo, q:), then the solution in terms of y and w; is
rh (a: cosh (v>² wit) + (a,(0)+t,,0) sinh (V>2- w
V---w3
7 wi
(1())= (a,(0) + (Ta; (0) + (0))t)
7 Wi
e- (c((0) (Vuž - 721)
(o. sin (Vw? 721) ww²-z²
7 < wi
Differentiating these equations gives solutions for (t)
rr (0( cosh (v7² - -
sinh Vクtt
? >> wi
b₂()) ve (0) (wfa, + (b)(
? wi
e- (b((0) cos (WwF 726)
202 a. (0) (0)
sin (Važ - 721)
7 < wi
Combining all these results, we can analytically decompose the expectation as the sum,
( + (t)))
0 qi
E
as (t)
Intuitively, this equation describes a damped rotation (spiral) around the OLS solution in the planes defined by the the eigenvectors of the Hessian rate proportional to the respective eigenvalue.
22

---

Under review as a conference paper at ICLR 2022
Stochastic component. Using the previous block diagonal decomposition A OSOT we can simplify the variance as
Var (0)) x- -
= =x- (U-1
=x-1 - (0T0-00 oO
Interestingly, the matrix OTU-1O is also block diagonal,
bl) Pr
0
P1
o²p
0 0
21++))²² 2
(H+
H
0=
0
o²H H
m1
Pm
0
o² Pm
Thus, similar to the mean, we can simply consider the variance in each of the planes spanned by If we define the block matrices,
las 0]" and [0 q:} T
n0
0 m(1+0) (ps + A)
D
25 B) Pi+A
S
211 8
S(1- Pi
(L++
then the projected variance matrix in this plane simplifies as
c Var
= D. - S. E.T 8:TT
Using the solution to damped harmonic osccilator discussed previously, we can express the matrix exponential e- S.T explicitly in terms of Vh² w/f |. then the matrix exponential is
1-3 and
n(+ (pi + A). If we let a,
M(++8]
cosh (ael) + I sinh (oil)
sinh (axl)
e 7
wi sinh (ast) a;
a
cosh (ast) as
sinh (a,t)
wi
1 + 7t
e S1t
e nt
w/
wit 1 1-
cos (a,t) + 2
sin (ast)
a; sin (a,t) cos (axl) بي sin (art)
e
a: 2 sin (ail)
< w
23

---

Under review as a conference paper at ICLR 2022
G ANALYZING PROPERTIES OF THE STATIONARY SOLUTION
Assuming the stationary solution given by equation (??) we can solve for the expected value of the norm of the local displacement and gain some intuition for the expected value of the norm of global displacement.
G.1 INSTANTANHOUS SPEED
ss [1 = Ess [10 0=]
Egs [||u++1|(]]
n/2rr (Ess [UR+1Ub++1)
112tr (Var,s (Ux+1) Eso |Kt1] Ess (vx+1l")
11irr (N 2)-
n/²
S(1 B2) 1r (a²
Note that this follows directly from the definition of 8ş in equation ① and the mean and variance of the stationary solution in equation ( ???), as well as the follow-up derivation in appendix F
G.2 ANOMALOUS DIFFUSION
Notice, that the global movement A, 0.-00.0. On can be broken up into the of the local movements
A. Ein o, where & 8, : Applying this decomposition,
Ess [IA, [1 = E
6
Ess
[1] + Σ
Eas K6 5,31
i=1
i
As we solved for previously,
1/2
E,s [[6611]]] = = [||= = 1 (v:)) = S(1 - 32) tr (o² H).
By similar simplification, we express the second term in terms of the stationary cross- covariance,
Ege [65, [6), - Ess [(v,, Uy)l 1/²tr (Cov.o(Vh.
Thus, to simplify this expression we just need to consider the velocity-velocity covariance Co, Uz). At stationarity, the cross-covariance for the system in phase space,
[0 oj]
Cov.s = -U-le AT
where x = S(1 - B²), and
U n(1 |)+ H-1 -H
I 20 8) v(++3)
AT)
T
[[[+8)
As discussed when solving for the mean of the OU trajectory, the drift matrix A be block diagonalized as A = OSOT where is orthogonal and S is block diagonal defined as
0
q1
0
mm
m1+#3 (PI + A)
0[I+サ]
2(1-8)
O
0
In
2
(pm
+ 2)
2(1-8) m(++9)
M1+
24

---

Under review as a conference paper at ICLR 2022
Notice also that O diagonalizes U-1 such that,
(11+9)00 P1
2
P1+
o²
A = 00
Pm X
0
0² pm
Applying these decompositions, properties of matrix exponentials, and the cyclic invariance of the trace, allows us to express the trace of the cross-covaitee cross-covariance as
tr (Cov = k 'tr (U--1-11)--)))
= -tr'tr U 0e -05
= tr Ae S
(Ace-
where A& and S₁ the blocks associated with each eigenvector of H. As solved for previously in the variance of the OU process, we can express the matrix exponential e SSlli--l explicitly in terms
1-8 and W/ of = x(+3)
V / (px + A). If we let r = /i - jl and ag = v/7² - wwtl then the
matrix exponential
cosh (axT) + 2 sinh (axT) ok " sinh (ax7)
sinh (auT) 2 sinh (ayeT)
e 7-
cosh (a47) as
Selii-ll
e -
W YT
T
? wk
cos (a + I sin (axT) we sin (OUTT)
1 sin (a++) cos (akT)
an
? We
I sin (OxT)
Plugging in these expressions into previous expression and restricting to just the 1 velocity com- ponent, we see
1 02 pke (cosh (AET) - 7 sinh (AET)
7 > wk
og 1 p2pppe-rr (1 -T)
Cov.s(Mi, (y.k) = x-1
[Age- -[(--51
7 Wk
1,1
k: l
-17 - (cos (O4E)) 2 sin (OLT)
y < Wk
au
Pulling it all together,
11²0
E.. | []
Ir (H)))++ 21
Σ c((())
5(1 - B22
k=1
where (k)) is defined as
e 7h cosh (ark) 2 sinh (ark))
? > wl
(k))
e 7 (1 11-7th yk)
= wl
e 7k cos (azk) 2 sin (a,k)
? < wl
1-8 for 7 = 7(1+A)* Wi =
2 ((1..)) (pi + A), and αq = vh² - wfl
25

---

Under review as a conference paper at ICLR 2022
G.3 TRAINING LOSS AND EQUIPARTITION THEOREM
In addition to solving for the expected values of the local and global displacements, we can consider the expected training loss and find interesting relationship the equipartition theorem from classical statistical mechanics.
The regularized training loss is (0) (0 T (1) \
11011 where is the Hessian matrix and y is the true mean. Taking the expectation with respect to the stationary distribution,
1 2 ((H + 2I) [00T]) n H [0] + 24
Ess (x)
5 Hp
The first and second moments of the stationary solution are
o² E,, [00T] S(1 3) 2
(H XI) H up
Plugging these expressions in and canceling terms get
E.s [Cx(0]] 4S(1 tr (o² +
77
A マ |111 2
Define the kinetic energy of the network K(v)) ml|ll1 where 31 8) is the per- parameter "mass" of the network according to our previously derived Langevin dynamics. At sta- tionarity,
E.. [K(()] n11+ 88 tr (L.ss [voT]) = 4S(1 - tr (o²
4
where we used the fact that [xss [vvT] 301 o2 In otherwords, at stationarity,
E. [Cx( = E.. [K(e]] +
2 |11
This relationship between the expected potential and kinetic energy can be understood as a form of the equipartition theorem.
26

---

Under review as a conference paper at ICLR 2022
H EXPERIMENTAL DETAILS
H.1 COMPUTING THE HESSIAN EIGENDECOMPOSITION
Computing the full Hessian of the loss with respect the parameters is computationally intractable for large models. However, equipped with an autograd engine, we can compute Hessian-vector prod- ucts. We use the subspace iteration on Hessian-vector products computed on variety of datasets. For Cifar-10 10 we use the entire train dataset to compute the Hessian-vector products. For Imagenet, we use a subset 40,000 images sampled from the train dataset to keep the computation within rea- sonable limits.
For experiments on linear regression, the Hessian is independent of the model (it only depends on the data) and can be computed using any model checkpoint. For all other experiments, the Hessian eigenvectors were computed using the model its initial pre-trained state.
H.2 FIGURE
We resumed training for variety of ImageNet pre-trained models from Torchvision 16 for 10 epochs, with the hyperparameters used at the end of training, shown in table
We kept track of the norms of the local and global displacement, ||6 112 and || every 250 steps in the training process, to keep the length of the trajectories within reasonable limits. || is visualized directly, along with its 15 step moving average. We then fit power law of the form ak to the | 1 trajectories for each model, using the last 2/3 of the saved trajectories. We visualize the || | trajectories along with their fits log-log plot.
Model
Dataset
Opt.
Epochs
Batch size
LR r
Mom, $
WD
VGG-16 VGG-11 w/BN
ImageNet
SGDM
10
256
10-5
0.9
5 x 10
ImageNet
SGDM
10
256
10-5
0.9
5 x 10
VGG-16 w/BN
ImageNet
SGDM
10
256
10-5
0.9
x 10
ResNet-18
ImageNet
SGDM
10
256
10 10-4
0.9
10
ResNet-34
ImageNet
SGDM
10
256
10
0.9
10-4
Table 1: Figure experiments training hyperparameters.
H.3 FIGURE 2
For this figure trained a linear regression model Cifar-10, using MSE loss on the one-hot encoded labels. The hyperparameters used during training are shown in table
At every step training the full of model weights and velocities were stored. The top 30 eigen- vectors of the hessian were computed as described in appendix H.]l using 10 subspace iterations.
The saved weight and velocity trajectories were then projected onto the top eigenvector the hessian and were visualized in black. Using the initial weights and velocities, the red trajectories were computed according to equation (5).
Model
Dataset
Opt.
Epochs
Batch size
LR rj
Mom. 3
WD
Linear Regression
Cifar- 10
SGDM
512
10-5
0.9
0
Linear Regression
Cifar-10
SGDM
4
512
10-5
0.99
0
Table Figure experiments training hyperparameters.
H.4 FIGURE 3
For this figure constructed arbitrary -haiteen--hllbecckk process with anisotropic noise which would help contrast the original and modified potentials. We sampled from 2 dimensional OU process of the form
27

---

Under review as a conference paper at ICLR 2022
d
21 2/2
A
x [t2
dt + 10 1
where set |-0.1, 0.05JT and arbitrarily construct A such that it's eigenvectors are aligned with Q1 = [-1, and 42 = [1,1]] and it's eigenvalues are and as follows:
D 9). v= ( 1). A=V-
The background for the left panel was computed from the convex quadratic potential Φ(x) f Az bz. The background for the right panel was computed from the modified quadratic
(x) = uUx uz, with U = (D + +Q ' and u = UA- (see equation 25 Both
were sampled in a regular 40 40 grid in [-0.1, [.0.1
H.S FIGURES 4, 5, AND
Starting with the ImageNet pre-trained ResNet-18 from Torchvision [16. we resumed training for epochs with the hyperparameters used at the end training, shown table The top 30 Hessian eigenvectors were computed as described in appendix H.] using 10 subspace iterations. During training, we tracked the projection of the weights and velocities onto eigenvectors (1 and 4su-
For Figure black. The background for the left and center panels was computed taking the ImageNet pre- trained ResNet- 18 from Torchvision [16] and perturbing its weights in the Qı and 93o directions in region close to the projected trajectory. The training and test loss were computed for a grid of 20 x 20 perturbed models. The background for the right panel was computed according equation (10)
show the projection of the position trajectory onto cigenvectors Q1 and tso in 2D in
Model
Dataset
Opt.
Epochs
Batch size
LR
Mom. 3
WD
ResNet-18
ImageNet
SGDM
5
256
10-4
10
0.9
10
Table 3: Figures 4,5,6 experiments training hyperparameters.
H.6 FIGURE
We resumed training for an ImageNet pre-trained ResNet- 18 from Torchvision 16 for 2 epochs, using the sweeps of hyperparameters shown in table We indicate a sweep a particular hyperpa- rameter by [A, B), which denotes 20 evenly spaced numbers between A and B, inclusive.
We kept track of the norms of the local and global displacement, || and || 122. every step in the training process. The value for ||6 1 the end of the two epochs is shown the top row of the figure. We fitted power law of the form ak to the | 1l trajectories for each model on the full trajectories. The fitted exponent e for each model is plotted the bottom row the figure.
Model
Dataset
Opt.
Epochs
Batch size S
LR rj
Mom. B
WD
ResNet-18
ImageNet
SGDM
2
[32, 1024]
10-4
0.9
10-4
R8sNet-18
ImageNet
SGDM
2
256
[10 10 10-21
0.9
10
ResNet-18
ImageNet
SGDM
2
256
10
[0.8, 0.99]
10
Table 4: Figure 7 experiments training hyperparameters.
H.7 INCREASING RATE OF ANOMALOUS DIFFUSION
Upon further experimentation with the fitting procedure for the rate of anomalous diffusion plained in Figure [ we observed an interesting phenomenon. The fitted exponent e for the power law relationship A x ke increases function of the length of the trajectory we fit to. As can
28

---

Under review as a conference paper at ICLR 2022
be seen in Figure increases at diminishing rate with the length of the trajectory. This could be indicative A 1 being governed by a sum power laws where the leading term becomes dominant for longer trajectories.
101 spochs
spocte
spoche
1.069
1,117 2>
1,157 >P
|112
1.28
Diffusion
1.20
1,190
1,219
1.245
spadhs
apocha
100 spocha
501
1.04
20
50
80
100
Step
Epochs
Figure 7: The fitted rate of anomalous difflusion increases with the length of trajectory fitted. The left panel shows the fitted power law on training trajectories of increasing length from the pre- trained ResNet-1e 18 model. The right panel shows the fitted exponent function of the length of the trajectory.
29


Reviews:
[{"review_id": "oxvil8ov4Xi", "reviewer": "Reviewer_pbf9", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written; a very pleasant reading. I find interesting the characterization of diffusion and displacement as a function of hyperparameters. Although it is not new that SGD obeys a modified loss, I find it insightful to see that this modified loss can be isotropic despite the loss being anisotropy, and to learn under which conditions this occurs. I also appreciate the clarity of the assumptions that are needed to obtain every single result, while maintaining the discussion at an intuitive level.\n\n\nMajor Criticisms.\n\n- Result (a) is not new. The anomalous diffusion at the end of training was already shown in arXiv:1803.06969, exactly in the terms defined in this manuscript (\"distance travelled grows as a power law in the number of gradient updates with a nontrivial exponent\"). \n\n- Result (c). That SGD obeys a modified loss was also stated in previous literature, such as Ref.[21] or, in a different framework, arxiv:1803.01927.\n\n- The title is very bold, inviting a rethinking of the limiting dynamics. I would say that this is because the authors assume that the deep learning community thinks that [statement (g) up here] “the network converges in parameter space”, or “the network stays within a local region”. These are however very naive intuitions that nobody with sufficient experience in deep learning theory would believe. Also, one of their main claims is [statement (f)] that detailed balance does not hold, implicitly suggesting that the results of e.g. Refs.[16,23] could be flawed. Again, nobody believes that detailed balance holds for SGD dynamics -not even at late stages of training-, since detailed balance is a sufficient condition (not even necessary, see e.g. https://aip.scitation.org/doi/full/10.1063/1.4863991) for equiibrium, but steady state distributions generally do not satisfy it. There is a whole section in Ref.[16] devoted to showing that SGD is out of equilibrium. Another example is arXiv:1803.06969, where they state that at the end of learning the system diffuses at the bottom of the landscape with a time-dependent diffusion constant (so this is clearly not an equilibrium process, and the network does not stay confined within a local region). In other words, I don't think that statements (f) and (g) add anything new to the current knowledge.\n\n- The analysis of the limiting cycles is nice, but I do not see what we learn more than was already presented in Ref.[21]. The only novelty would be that the paths at the end of learning are not limit cycles, but rather space filling curves. However, in the way as it is presented, this looks more like a conjecture than an actual result. As for the description in terms of velocities, this is nice, but I don't find it surprising that the velocities oscillate, especially in the directions of the top 30 eigenvectors. A system confined by quadratic walls will have an oscillating velocity, as any harmonic oscillator. I would find it more surprising if the authors measured the same kind of oscillation in the direction of the bottom eigenvectors.\n\n\nMinor comments:\n\n- The authors project the limiting trajectory for the parameters onto the plane spanned by the top q 1 and bottom q 30 eigenvectors. I assume that  q1 and q30 were chosen in order to maximize the anisotropy on the related plane, in order to show that the trajectory is instead isotropic. I suggest to explain that this is the reason of the choice, since at first it wasn't clear to me why this was being done.\n\n- Fig.3. It's nice to see that the trajectory is isotropic when no regularization is used. But why is the trajectory not centered around the minimum of the modified loss? Also, it looks like the minimum of test and modified losses coincide. Why is this?\n\n- typo in section 8: trigonomentric\n\n- Implicit regularization of the velocity trajectory. I assume that the authors mean that the velocity appears in the modified loss as an L2 regularization term. Is this right? Is there any further consequence to this?\n\n- In figure 6 - bottom, the dots are measured from fits of the trajectory. The dashed line is c=1 diffusion, but this is a little confusing, because the dashed line in the top three plots is the theoretical prediction of Eq11. ", "summary_of_the_paper": "The paper studies the long-time dynamics of deep neural networks. The author(s) (1) show some empirical findings related to the mean square displacement, (2) model SGD as an underdamped Langevin Equation, relate it to an Ornstein Uhlenbeck process in a linear regression setting, and use it to study the limiting dynamics of SGD, (3) use the Fokker-Planck formalism to show that the steady state weight distribution obeys a modified loss which is isotropic in the absence of L2-like regularization, (4) provide empirical evidence that their findings are relevant also beyond the context of linear regression.\nAccording to the authors, these are the novelties provided by the paper:\n\n(a) \"We find empirically that long after performance has converged, networks continue to move through parameter space by a process of anomalous diffusion in which distance travelled grows as a power\nlaw in the number of gradient updates with a nontrivial exponent.\"\n\n(b) \"We reveal an intricate interaction between the hyperparameters of optimization, the structure in the gradient noise, and the Hessian matrix at the end of training that explains this anomalous diffusion.\"\n\n(c) \"we show that the key ingredient driving these dynamics is not the original training loss, but rather the combination of a modified loss, which implicitly regularizes the velocity, and probability currents, which cause oscillations in phase space\"\n\n(d) \"We identify qualitative and quantitative predictions of this theory in the dynamics of a ResNet-18 model trained on ImageNet.\"\n\n(e) \"We uncover a mechanistic origin for the anomalous limiting dynamics of deep neural networks trained with SGD\"\n\n(f) \"one of the most significant results of our analysis is that, depending on the relationship between the gradient noise covariance and the Hessian, the stationary distribution in phase space will generically violate detailed balance\"\n\n(g) 'Natural intuitions, such as “the network converges in parameter space” or “the network stays within a local region”, are wrong'\n", "summary_of_the_review": "The paper presents several concepts in a very clear and readable manner. However, the novelties presented are limited/incremental (certainly not as revolutionary as suggested by the title), and some of the results are either not new (anomalous diffusion in DNNs has been known for years now), not relevant (the lack of detailed balance in results relying on a stationary state distribution), or not well-supported (space filling curves instead of limit cycles). On the other hand, I do find some value in some of the results, such as the characterization of diffusion and displacement as a function of hyperparameters, and showing that under certain conditions the modified loss is isotropic.\n\nThe authors should reduce their claims regarding what is new in their work and adapt them to what is the real state of the art of the understanding of DNNs, and remove disputable statements such as \"The major limitation of the stationary dynamics approach to analyzing SGD is its implicit assumption that the system is in detailed balance\".\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "LFMQII2Xjm6", "reviewer": "Reviewer_erMi", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strength**\n\n1. The theoretical results are solid and well-motivated;\n2. The experiments are thorough, and consistent with theoretical results;\n3. The contribution of this work are carefully clarified.\n\n**Weakness**\n\nThe only pity is that all the theoretical results are derived from a linear regression model. The impact of this work will significantly increase if these interesting theoretical results can be derived in more general cases.\n\n**Further concerns**\n\nI notice some disruptive claims are raised in the last paragraph. Please can you elaborate them in order to raise my score? The following two concern me the most:\n- “the network stays within a local region” is wrong;\n-  \"Finally, the intuition that faster learning rates would lead to faster global displacement in parameter space is also wrong; instead induced velocity anti-correlations lead to slower global displacement\"", "summary_of_the_paper": "This paper derive a continuous model for SGD in linear regression settings to build understanding on some interesting phenomenon of deep neural network trained with SGD. ", "summary_of_the_review": "In this paper, the insight from the theoretical results can match the empirical observations pretty well. Their results certainly deepens our understanding on the learning dynamics of modern neural network.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6ndJvDE9NTv", "reviewer": "Reviewer_4wwS", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I welcome the general direction of the work: trying to understand the dynamics\nof neural networks is a complex undertaking that a large community of\nresearchers is pursuing, so the study of simplified models is a promising\navenue.\n\n### Setup of the study\n\nIn studying the *limiting* dynamics however, the authors limit themselves to a\nsetup where I don't see any immediate connections to learning or representations\nof neural networks. The dynamics the authors describe happen *after* resuming\ntraining of a pre-trained neural network, thus I feel like their setup restricts\nthe potential impact of the results of this study.\n\n### Clarity / novelty of the results\n\nI found the article hard to read at times because the authors repeatedly qualify\ntheir observations as \"surprising\", \"contrary to common intuition\",\n\"nonintuitive\", etc. I think these qualifiers can be mistaken as claims of\nnovelty etc. and would hence use them more sparingly. For example, the fact that\nneural networks continue to move through their weight space has been\nwell-established for quite some time now, cf. for example [Jastrzebski et\nal. '17, Chaudhari & Soatto, 18, Baity-Jesi et al. '18, and many more]. Hence I\ndidn't find the observation in Figure 1 \"surprising\" (p 2 after the equation),\nwhich underlines the subjectiveness of these claims - or else I maybe reading\nthe figure incorrectly?\n\nAnother example: I would consider it well-established that OU processes whose\ndiffusion matrix is not isotropic do not follow the naïve Gibbs distribution,\nbut instead equilibrate in a modified potential (see for example Section 5.3\n\"Potential conditions\" of Gardiner's \"Handbook of stochastic methods\" etc.)\nFurthermore, modified losses arising through SGD dynamics have been studied in a\nnumber of recent deep learning papers, some of which are cited by the authors \n\nOther claims about the significance of the results should equally be clarified\nin my opinion, for example:\n\n> The expectation that the training trajectory would reflect the underlying \n> anisotropy of the training loss driving the dynamics is also wrong;\" (p. 9)\n\nIn my understanding, this study is only concerned with the *limiting* dynamics\nof learning, and hence conclusions about the training cannot be drawn\nimmediately?\n\n### Separation of experimental from theoretical results\n\nThe authors should be lauded for trying to connect their theoretical results\nwith empirical results on deep networks. Again though, I think the presentation\nof the results should be revised to clarify which predictions are actually\nderived from theory.  Take the exponent of anomalous diffusion (bottom of Fig\n6): it cannot be estimated of the global displacement (12), as the authors\nexplain. Instead, the authors evaluate the dependence of the diffusion constant\non learning rate, batch size and momentum parameter directly from a simulation\nby fitting a power-law to the empirical displacement. I would present this\nresult separately, as it is not a theoretical prediction, and thus presenting it\nin a section entitled \"Predicting the diffusive behaviour of the limiting\ndynamics\" could cause confusion in my opinion.\n", "summary_of_the_paper": "The starting point of the present paper is the observation that the *parameters*\nof a neural network continue to change when training the network with stochastic\ngradient descent (SGD) even after the performance of the network has stabilised\nat a final value. The authors study these \"limiting dynamics\" in linear\nregression, which they model as an underdamped Langevin equation resulting in an\nOrnstein-Uhlenbeck (OU) process and find an oscillatory dynamics in the\n(position, velocity)-phase space. The authors discuss how to connect their\nfindings in this simple model with the dynamics of Resnet18 (see\nbelow). Finally, they attempt to predict features of the diffusive behaviour of\nthe ResNet18 from their model (Sec. 8)", "summary_of_the_review": "I welcome the direction of the present study, trying to understand the dynamics\nof deep neural networks via simpler models. However, I cannot recommend\nacceptance of the present paper for three reasons: the implications of the\nstudied dynamics on learning or representations in neural network remain unclear\nto me; the results derived from the Ornstein-Uhlenbeck process do not go beyond\nthe previous work to an extent that would warrant the publication at ICLR in my\nopinion; and the unclear presentation of the results which doesn't always\nclearly separate previous work from the presented results, or theoretical\npredictions from numerical measurements (Fig. 6).\n\n**Edit Nov 14th** I thank the authors for the lengthy replies to the points raised in this review. After careful consideration of their comments, and also taking into account the valuable points made by other referees, I keep my original score for now. I have increased my confidence from 4 to 5.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3PG5T5G9LCG", "reviewer": "Reviewer_sDRt", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In the current format, the paper suffers from several limitation that should be addressed by the authors. I believe that there are interesting results and I encourage the authors to answer my comments as I would be very happy to rise my recommendation score. \nThe following comments try to follow the section in the submission:\n\n \n\n\nThe authors in their analysis of related works miss a large portion of the relevant literature. Discussing the analysis of the dynamics of neural networks: \n* The mean-field limit was introduced by three different independent works [Rotskoff, Vanden-Eijnden 2018] [Chizat, Bach 2018] [Song et al. 2018], only the ladder is cited.\n* If the mean-field limit is mentioned then also the dynamical mean field theory approach should be cited. In the context of neural networks we have: [Mignacco et al. 2020][Mannelli et al. 2020][Mignacco et al. 2021].\n* The dynamics of stochastic gradient descent and momentum were also obtained in dynamical mean field theory by the two papers by Mignacco above and [Mannelli, Urbani 2021].\n* When mentioning the modelling of SGD using SDE, it is important to discuss the results of [Simsekli et al. 2019] where they showed that the jumps are fat-tail distributed and how that is compatible with the approximation.\n* On the \"empirical exploration\" the authors refer only to [Papyan 2018], the results of [Ghorbani, Krishnan, Xiao 2019] that derived the same method independently and make thorough analysis of the dynamics.  In the same section also [Sagun et al. 2016] (ref [35]) should appear. \n\nOn page 2, the authors say \"Surprisingly, [..] the networks continue to move through parameter space\" but this very well known. It known that stochastic gradient descent keeps moving after training, indeed this is at the basis of ref. [16]. What is not known is how is moving, and indeed this finding is interesting.\n\nIn figure 1 lower panel, there is only one value on x-ax. Please add at least a second one. I can guess that the next number after 10^4 would be 2*10^4 but this is not obvious a priori.\n\nThe authors should try to avoid to add adjectives to embellish the paper in the technical sections. These parts of the paper should be factual and reporting their results in a scientific way. \nFor instance in section 8, describing figure 5 the authors say \"As predicted, the spiral appears more evident ..\", can you quantify how \"spiral\" the plot is? It is hard to conclude from that picture. Furthermore, that appears to be the result of one run. It should be good to average over many simulations to avoid the risk of observing random fluctuations.\nThe fact that the second figure appears \"less evident\" may indicate a limit in the theory. Since the authors are considering a SDE, higher order effects should appear more frequently for eigenvectors corresponding to smaller eigenvalues. \nIt would be interesting to see what happens to the 1000th eigenvectors. Since ImageNet has 1000 classes, according to [Sagun et al. 2016][Papyan 2019] that eigenvector should be associated to irrelevant information and probably this behavior disappear. You could observe this also if you have the pretrained network in CIFAR100 or CIFAR10, by looking at the 100th and 10th eigenvector respectively.\n\nAn important aspect, that has not been discuss, is the effect on generalization. What are the practical consequences of these results? Although the theoretical framework can not be applied to the test dataset, the authors can verify the practical effects. \nIs being super/sub diffusive an advantage? Does the performance change? Answering these questions will bring value to the work.\n\nThe authors affirm that the dynamics brake detailed balance. Unfortunately I must have missed this part in the paper. Could they clarify?\nAlso, what are the practical implications of braking detailed balance?\n\nIn figure 6 the authors claim that their theory predicts the behavior of the displacement. It is not clear to me the procedure used. When you say \"using this single estimate\" what do you mean? Are evaluating \\sigma^2 tr(H) in a single point for each plot and using this information in draw the line? It is not very clear.\nOn the second line, it is an overstatement to say that you predict the diffusion constant. Your model, by construction, is a correlated Brownian  motion with a drift, therefore the diffusion constant is 1. \nThis is a pity, there is no explanation for the very interesting behavior the you observe.\n\n", "summary_of_the_paper": "The paper uses a stochastic differential equation approximation of stochastic gradient descent momentum to model the neural networks dynamics. At the theoretical level the study focuses on linear regression. They showed that there is agreement between the results in the simpler model and deep neural networks. \nThe most relevant fact is that the networks can have a super or sub diffusive behavior at the end of training. Unfortunately, their theory cannot explain this finding nor the practical consequences on the test loss.", "summary_of_the_review": "The paper presents an interesting result about the presence of sub/super diffusive behaviour at the end of training. Although the result may be of limited practical interest, it is a new finding and would be interesting to understand causes and consequences of that (not discussed in the paper). \nThe paper has a large number of flaws that I believe can be amended to get a publishable result.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors study the limiting dynamics of a simple linear regression model. They use an underdamped Langevin equation which is quite common in the literature. Although the reviewers welcome the direction and the attempt to understand the dynamics of a simple model, the novelty of the paper is limited. As an example, the paper shows that the key ingredient driving these dynamics is not the original training loss, but a modified loss. As pointed out by the reviewers, this has already been pointed out in multiple papers. One important problem is the tendency of the paper to oversell the results (including the title), which makes it difficult to clearly separate the contributions made in the paper. After a discussion with the reviewers, the overall feeling did not change.

I can therefore not recommend acceptance. I strongly recommend the authors do a significant rewrite of the paper in order to clearly separate what contributions are truly novel and also improve the discussion of prior work.
Decision: Reject