A Quasistatic Derivation of Optimization Algorithms' Exploration on Minima Manifolds
Authors: Chao Ma
Daniel Kunin
Lexing Ying
Venue: Submitted to ICLR 2023 (2023)
Keywords: Implicit bias
minima manifold
time-scale separation
Adam

Abstract: A quasistatic approach is proposed to derive the optimization algorithms' effective dynamics on the manifold of minima when the iterator oscillates around the manifold. Compared with existing strict analysis, our derivation method is simple and intuitive, has wide applicability, and produces easy-to-interpret results. As examples, we derive the manifold dynamics for SGD, SGD with momentum (SGDm) and Adam with different noise covariances, and justify the closeness of the derived manifold dynamics with the true dynamics through numerical experiments. We then use minima manifold dynamics to study and compare the properties of optimization algorithms. For SGDm, we show that scaling up learning rate and batch size simultaneously accelerates exploration without affecting generalization, which confirms a benefit of large batch training. For Adam, we show that the speed of its manifold dynamics changes with the direction of the manifold, because Adam is not rotationally invariant. This may cause slow exploration in high dimensional parameter spaces.

Under review as a conference paper at ICLR 2023
A QUASISTATIC DERIVATION OF OPTIMIZATION ALGO- RITHMS' EXPLORATION ON THE MINIMA MANIFOLD
Anonymous authors Paper under double-blind review
ABSTRACT
A quasistatic approach is proposed to derive the optimization algorithms' effec- tive dynamics on the manifold of minima when the iterator oscillates around the manifold. Compared with existing strict analysis, our derivation method is sim- ple and intuitive, has wide applicability, and produces easy-to-interpret results. As examples, we derive the manifold dynamics for SGD, SGD with momentum (SGDm) and Adam with different noise covariances, and justify the closeness of the derived manifold dynamics with the true dynamics through numerical experiments. We then use minima manifold dynamics to study and compare the properties optimization algorithms. For SGDm, we show that scaling up learning rate and batch size simultaneously accelerates exploration without affecting generalization, which confirms a benefit of large batch training. For Adam, we show that the speed of its manifold dynamics changes with the direction of the manifold, because Adam is not rotationally invariant. This may cause slow exploration high dimensional parameter spaces.
INTRODUCTION
The ability of stochastic optimization algorithms explore among (global) minima is believed be one of the essential mechanisms behind the good generalization performance of stochastically trained over-parameterized neural networks. Until recently, research on this topic has focused on how the iterator jumps between the attraction basins of many isolated minima and settles down around the flattest one Xie et al. (2020); Nguyen al. (2019): Dai Zhu 20200;] Mori et al. (202T) However, for over- parameterized models, the picture of isolated minima is not accurate, since global minima usually form manifolds connected minima Cooper (018). addition to crossing barriers and jumping out of the attraction basin of one minima, the optimizer also moves along minima manifold and search for better solutions Wang et al. (021). Hence, understanding how optimization algorithms explore along the minima manifold is crucial understanding how stochastic optimization algorithms are able to find generalizing solutions for werr-parmeetrrized neural networks.
Some recent works have begun examine the exploration dynamics of Stochastic Gradient Descent (SGD) along minima manifolds. Many of these works have identified how a change of flatness in the minima manifold adds a driving force SGD oscillates around the minima. For example, Damian et al. (2021) considered an SGD training a neural network with label noise, and showed that the optimizer can find the flattest minimum among all global minima. A more recent work Li et al. (2021b) derived an effective stochastic dynamics for SGD on the manifold. The results in Li et al. (2021b) show that (when the learning rate tends zero) the changing flatness can give force the SGD iterator along the minima manifold and induce a slow dynamics on the manifold that helps the SGD move to the vicinity of flatter minima.
In this work, we study the same questions of the flatness-driven exploration along the minima manifold. Instead of searching for a strict proof, we focus on simple and intuitive ways to derive the manifold dynamics. Specifically, we propose quasistatic approach to derive the manifold dynamics for different optimization algorithms and stochastic noises. The main technique our derivation is a time-scale decomposition of the motions perpendicular to and parallel with the minima manifold, which we call the normal component and the tangent component, respectively. We treat the normal component infinitely faster than the tangent component, and thus is always equilibrium given

---

Under review as a conference paper at ICLR 2023
the tangent component. The effective dynamics of the tangent component, i.e. the manifold dynamics, is obtained by taking the expectation over the equilibrium distribution of the normal component.
The main step in analysis involves deriving the equilibrium covariance of SDE. Compared with the theoretical analysis in Li et al. (2021b), our derivation and results are simpler and easier to interpret, and clearly identifies the roles played by each component of the optimization algorithm (noise covariance, learning rate, momentum). The following simple example demonstrates the main idea of the our derivations.
A simple illustrative example: Consider loss function f(x, v) h(z))², where h(x) > 0 is a differentiable function of æ. The global minima of this function lie on the x-axis, forming flat manifold, and h(z) controls the flatness of the loss function at any minimum (xr, 0). Let 2 [x. y)" We consider an SGD approximated by SDE Li et al. (0117 Li et al. (021a))
[ (z )dl (D)) dz: V
where 1) is the learning rate, D is the square root of the covariance matrix of the gradient noise, and W. is a Brownian motion. For the convenience presentation, for points (x, y) that are close to the x-axis, assume the noise covariance aligns with the Hessian of the loss function at (I, 0), i...
0 0 D
(²(2) / ((,
o2((x))
where a > 0 is a scalar. Then, the SDE equation can
be written
dx;= = -h' (1 dy: 2h(x))yedt + t / /nh(xx)dW.
with We being a 1-D Brownian motion. When yt is close to O, the speed of x, is much slower than y because of the vi in the dynamics of is much smaller than the yt in the dynamics of y. When this
separation of speed is large, the dynamics above can be approximated by the following quasistatic dynamics
dx, = lim By. ((x:) dy, = (h((:))-dd + Vnh(z,JodW,
(2)
+0000
which assumes is always at equilibrium given zy. Solving the Ornstein -bllenbek process [ we know the equilibrium distribution of yr is ~ N(0, ue" ), and hence the manifold dynamics is
dr. no²h'(x,)
4
This derivation shows the slow effective dynamics along the manifold is gradient flow minimizing the flatness f(z). This simple quasistatic derivation reveals the flatness-driven motion of SGD along the minima manifold, and recovers the same dynamics given by Li al. (2221b) in this specific case. On the left panel of Figure we show an SGD trajectory for f(x, y) (I + x >y illustrating the exploration along the manifold due the oscillation in the normal space. On the right panel we verify the closeness of the manifold dynamics with the true SGD trajectories for the same objective function. The "Hessian noise" and "Isotropic noise" represent noises whose covariance are the Hessian matrix of (as analyzed above) and the identity matrix (covered by the analysis in Section 2 respectively.
Theoretical applications of the manifold dy- namics: The minima manifold dynamics of optimization algorithms can be used tool to study and compare the behaviors of optimization algorithms. Section 3 we illustrate how our derivations can be applied to study the behav- ior of SGD on a matrix factorization problem. Two more interesting applications are discussed Section 4 and 5 In Section 4 we focus on SGDm, and study the role played by the leam- ing rate, batch size, and momentum coefficient in its manifold dynamics. Based on the analysis, we explore approaches to reliably accelerate the manifold dynamics, which may help accelerate training. Especially, we show that scaling up learning
leedann 500
ninima
Isstropic SGD
Hesdan
iestropi oin
100
200
300
400 500
T-lteratlnssstte size
Figure 1: (left) The trajectory of SGD (real dynamics) with Hessian noise initialized from (5, 0). (middle) The x-coordinate of the real dynamies and the manifold dynamics for SGD with Hessian and isotropic noises.
2

---

Under review as a conference paper at ICLR 2023
rate and batch size simultaneously accelerates exploration without affecting generalization, which confirms a benefit of large batch training. In Section 5 we study adaptive gradient methods, and show that the speed of the manifold dynamics of Adam Kingma & Ba (2014) changes with the direction of the manifold, because Adam is not rotationally invariant. When the manifold does not align well with some axis direction, the exploration of Adam along the manifold is even slower than SGD with the same learning rate. This shows the sensitivity of Adam (and other adaptive gradient methods) to the parameterization and a potential weakness of Adam on the exploration among global minima.
Additional related work: Besides the important related works discussed above, our work is related broadly with the topic of implicit regularization of optimization algorithms for overparameterized models. Such works include global minima selection Wu et al. (2018); Ma & Ying (202I). the max- margin bias of some homogeneous models Soudry et al. (2018) Lyu & Li (2019); Lyu et al. (022). and others such as Gunasekar et al. (2017): Li et al. (2018): Woodworth et al (2020); Neyshabur et al.)
(2017) (not to be comprehensive). One work that is worth mentioning is Blanc et al. (2020), in which the implicit regularization effect of SGD under label noise is studied by understanding the SGD as a
Ornstein- Ornstein-Uhlenbeck process (OU process), which also plays important role in our derivation.
For applications, we give theoretical explanations for the benefit of large batch training, which is an important and extensively studied empirical topic You et al. 2019 2017); Hoffer et al. (2017); Geiping et al. (202I). Our study on adaptive gradient methods concerns the family of widely used methods which adjusts the learning rate adaptively (and usually coordinatewise) Duchi al. [2011)
Tieleman & Hinton (2012); Kingma & Ba (2014).
2 THE QUASISTATIC DERIVATION FOR MANIFOLD DYNAMICS
In this section, we introduce our quasistatic approach for deriving minima manifold dynamics, which is an effective exploration dynamics of optimization algorithms on minima manifolds.
Let M be smooth manifold in Rd with Euclidean metric. Throughout this paper, use z to denote points in Rd, including M,, and use x and y to denote components of z used in the quasistatic derivation for speed separation. For any Let T-M- be the orthogonal complement of T- M, which is the normal space of M at z. Let PM be the projection operator onto M, i.e. for any z Rdd, P. gives the closest point on M (if exists and unique). PM is well defined when z is close to M Lee 2013). In the paper, we ignore the subscript M for PM when there no confusion.
Notations:
M,, let T&M be the tangent space of .
Problem settings We consider boundless k- k-dimensional smooth manifold M in Rd, formed by (local or global) minima of function f : IR + RR Since we are interested in the behavior of stochastic optimizers near the minima manifold M,, we ignore the landscape of when x is far away from M, and only consider a quadratic expansion of I at M. Specifically, let H()) : M + Rdxd be a function on M that gives the Hessian of the loss function on the minima manifold. For any E M, we assume (x) ) is positive semidefinite, and its O-cigenspace is exactly T.M. The loss functions that we consider take the form
f(z) = (z - P2) HP=)) Pz), = Rd.
For the optimization dynamics, we start from SGD, approximated by the following SDE Li al. 2017 Li et al. (202
dz, = - V((zz)dl (DD( -
(4)
where Wt is Brownian motion in IRd and D M -> Rdxd is the square root of the noise covariance. Later we will extend our analysis to SGDm and Adam. Strictly speaking, the noise of SGD depends on z, which is not always on M. However, in settings that we study, we assume only depends on Pz because is close to M.
2.1 FLAT MANIFOLD
We start from the case in which the minima manifold M is flat. The derivation in this case is similar to the 2-D example given in the introduction. For any z E Rd. let z = [x" y y"1", with
R and y Rd-k, Without loss of generality, we assume M = {z [xT y T y- 0).
3

---

Under review as a conference paper at ICLR 2023
i.e. M the linear subspace formed by the first & axes. Then, for any z
x
we have
y
P= = x and = - Pz 0 y
For the loss function f. since we assume its Hessian has zero
cigenvalues along the tangent space of M, the Hessian must take the form H())
(x)
with (x) E ((d=)) Here, we can treat H as function of x because it is defined on M. Hence, written as function of x and y. the loss function becomes
f(x, y) = y" H())y.
(5)
Next, we rewrite the S using x and y. Again, since the noise coefficient D in
4 is defined on E Rdxd with
1x)
122(x)
M, it can be treated as a function of x. For any x, let D(x) = D21
D22(x)
(x1 Bk Then, D11 represents the noise in the tangent space, Dzz represents the noise in the normal space, and D12, D21 represent the interaction of the tangent space and normal space noises. In many cases, the noise covariance matrix of SGD aligns with the Hessian Mori et al. (202I). Hence, D22 dominates the other three components. (Actually, when the noise covariance strictly aligns with the Hessian, only D2g is nonzero.) In the following, we assume the interactions Dig and D21 are O, while D1 can still be nonzero. Then, by the form of the loss function in 5 the SDE can be written as the following system of x and y:
dx; = - v 0x (x:) +
dy: 2H xx, yydyt + (x))(x)))W
(6)
where 0₂ (x;: is a (d - k)) x (d 1:) x tensor containing all partial derivatives of x. When is close M, y is small, in which case the dynamics of y is much faster than that of x, because the drift term for x depends quadratically with y while the drift term for only depends linearly with Therefore, we can use a quasistatic dynamies approximate the original dynamics. The quasistatic dynamics assumes that is always at equilibrium:
dx, lim y H (xt) (x)
dy, = (x.) + (n
(8)
Fixing xs, the dynamics for y, is linear SDE. We have y, as lim Ey, = 0 and lim Ey, y. = Ve, 7
where V E (dd-A) satisfies Hx,)V, + VxH((x,, DD2 (x.) (x, ) /2. The derivations
here are standard. Readers can refer to textbooks or lecture notes such as Herzog 2013). Substituting the moments into the u-dynamics in
7 we have
12 n
dx;
(V)) ( + (x)
(9)
Understanding x as a vector on M, equation D gives the effective manifold dynamics on M. This result recovers the simple example in the introduction if we take f(x, v) = h(x)y² and D2 = o² H//2
2.2 GENERAL MANIFOLD
For general smooth manifold, the manifold dynamics can be derived locally by approximating M using a flat manifold. The resulting dynamics is different from only in that the gradients taken on the manifold. To see this, consider any point za E M. Without loss of generality, we assume
x R² } Because M is a smooth manifold, around zo the projection operator onto T M, denoted by P induces an 1-1 map between M and T M, and we have =| P =I = O[I - =oll²). Let P-1 T Tr, + M be the inverse of this 1-1 map. With an abuse notations, for z = [x". O E M, we sometimes also use P-1x to denote P-12 z. Let Oz = (ei,..., ea) be the standard orthonormal basis for Rdd
there exists Xo E F. such that xn
Xo and T M
0
4

---

Under review as a conference paper at ICLR 2023
with {e1, .... C being orthonormal basis for MM Also because M is smooth, for any M close to 20, there exists an orthonormal basis O₂ = {ef .... eă} for R which is close to Ozo- such
that {eī ..... ex) form an orthonormal basis for M. Specifically, for any 1 S i S d we have llet e = O =-
Now, consider SDE with the loss function f defined in For any .M close 2o, let
4
((z) - 0 0
D11(2)
and ((2)
be the Hessian and the noise coefficient D22(2)
(=)
matrix expressed in Oz. The zeros in the expressions are due to assumptions on H and the noise, i.c. T.M is in the O-eigenspace of II, and there is interaction between the tangent space noise and the
normal space noise. Now, we define companion loss function f whose minima manifold is T M
by f(z) = Y (P= lxyy, where z = [x_, y² y² and consider a companion SDE dz. -V )dl Po,
(10)
The SDE above approximates SGD minimizing f. whose minima manifold is flat. Hence, using the results from the previous subsection, we can derive manifold dynamics on M,
dx Σ
V ()))))) (P (')WW((
(11)
WE=
where Vi is obtained by T(PP,'x)V: (H
/ /2 By the discussions above, around zo the SDE is close to because f(z) ~ f(z) and PPP Ro D(Pz). Hence, the effective manifold dynamics of is close 11 in neighborhood of zu. and this approximation is better in smaller neighborhood of zg. Therefore, at x = 2u the manifold dynamics approximately 11 taking Xp. which leads
nD P Dzs
n- Σ (f())))) + vnDn(=)dw
dz =
Since (+)())) are gradients in the tangent space T2, M, under Euclidean metric they are gradients the manifold 2g- Hence, 2a the manifold dynamics can be written as
n k Σ Vo VV + (TDn()) FJ=1
dz
(12)
Since the above analysis holds for any 2o M, the manifold dynamics of is given by 12 at any
Examples: One interesting case is when noise covariance matrix is proportional with the Hessian. In this case, (²) = (z) for any z E .M. Since H()) contains all the nonzero eigenvalues of (2)) and the nonzero eigenspace corresponds to T&M 1, we have Du = and lfl² (z:) which gives 12 we have the following effective dynamics
Dz2 (2) vv (=). In this case, V/, satisfies i1zz,VV, Vz,((z,
V, n00 T. Substituting V, into the dynamies equation
4
n1²² V (r.
dz, 4
(13)
In the case of a flat manifold, this result above corresponds the dynamics derived in the introduction. Effectively, the SGD is minimizing ge² Tr(H(z)) on the manifold using a gradient flow. (Note that (T((()))=T((((()) for z M).
For another case, we assume the noise in T&M is isotropic with constant magnitude, we have Dz(2)) a for some Hence, we have V/2 = ne H (z:). and the manifold dynamics becomes
dz, 11022 VM Tr(log (2.))
(14)
na² I Tr(log (=)) on the manifold with a gradient flow. Effectively, the SGD is minimizing
'Athhough [1 is dynamics for x Rt it can he understond dynamics for z on T M, in which the component is always
5

---

Under review as a conference paper at ICLR 2023
Remark 1. The manifold dynamics we derive is similar to that studied in Li et al. (2021b). Instead of providing a rigorous proof. our main contribution is to give a simple and intuitive quasistatic approach to derive the manifold dynamics. Our methods can be applied models, and also can be applied to other optimizers such as SGD with momentum (See Section 2.3 and Adam (See Section 5)
wide class of noise
2.3 EXTENDING ANALYSIS TO SGD WITH MOMENTUM
The quasistatic approach we take can be extended to derive the effective manifold dynamics for SGDm. We consider the following SGDm scheme the same loss function studied above:
mg+1 = jam - (Zzk, Zi+1 = 21: + M
(15)
where 1) is the learning rate, l E [0, 1) is the momentum factor, and is the momentum. By the derivation in Li et al. 2017). we consider the following SDE system that approximates [5]
1 m, + fz.))) + (nD(Px))dW dz, m,.dt 7
dm,
(16)
7
By the discussion for SGD, the manifold dynamics can be obtained by assuming that M is flat, and applying quasistatic analysis on a decomposition of tangent and normal components. We put the details in Appendix A
The resulting manifold dynamics general manifold is
P
(V ggdd Day V (H((())) dt + (TDD dz, = -m,dt,
dm,
17
11=1
(17)
where the form and derivation of V.sgdm are given in Appendix A (equation 39p Note that when 7) is small, equation 39 is close to HV VH शाーニッ D22 Diz Let Vagdm be the solution of
this equation. We have / Vsgdmm yroa previous sections. This shows that the momentum amplifies the flatness driven force by a factor of 1/ - μ). Besides this accelcration, the momentum scheme itself also accelerates the speed of manifold dynamies. To see this, when there is noise along the manifold direction, i.e. D11 by Kovachki & Stuart (2021), the ODE
I-a" where V is the matrix V for SGD used in
nI
2 = 1
(V, (Vagdim Jiy V.
(18)
is a first-order approximation of the manifold dynamics (This approximation assumes the momen- tum always at equilibrium). The term 1/(1 p)) the acceleration brought by the momentum scheme. In this case, compared with SGD, the approximate manifold dynamics for SGDm is
(verd) )ig V MM(H
(1 22²
which is 1/(1 - 2 faster than SGD. The full derivation for SGDm are put Appendix
Examples:
We still consider the example where D1 = and D22(2) av Hz)). In this case, 5sgddm
yagdm
na² 1(1 1(1-))
I. Then, the approximate effective dynamics according to equation 18 is
n/2
2, = 4(1
V (r.).
(19)
As numerical justification for the manifold dynamics derived for SGDm, for the same function fæ, y) = (x² + 1)y² tested in Figure [ we compare the true x-coordinate dynamics, the SiDm-like discretization for equation 17 and the ODE solution T9 Shown in the left panel of Figure 2 the three dynamics are close for all μ tested. The results also show that the manifold dynamics get faster for larger μ., as predicted by its expression.
APPLICATION ON MATRIX FACTORIZATION PROBLEMS
In this section, we consider an objective function f(U V) ||UT - M with U E RMxp, V E Raxp and an SGD with Hessian noise, i.e. the noise covariance is proportional with the Hessian of f.

---

Under review as a conference paper at ICLR 2023
Let HU, be the Hessian matrix of
10a
deplacement
ditannce
2.5
f. Theoretically, easy to verify that (U, V)) is
5030 distanee
5000 dipactment
Scam munifold ODC nonitald
1.5
Oetance
0.6
2:
1.0
proportional to I + || wher- M.
0.5
p
0.0
20
40
60
80
100
200
400
600
300
1000
200
430
600
800
1000
ever UV Therefore, the man- ifold dynamics of SGD minimizes || + +VV and
T-iteratio steeteg size
Twrteratonsssttpppoo size
Teiteratfonsxstep size
Figure 2: (left) SGDm and its manifold dynamics for different μ. (middle) The distance between the real optimization dynamics and the manifold dynamics SGD and SGDm with Hessian noise, compared with the displacement of the optimization dynamics' iterators. (right) The distance and displacement curves for average trajectories of SGD when there is noise long the minima manifold.
drives the iterator to the most balanced global minimum, which is also the flattest. Numerically, we take true dynamics with the manifold dynamics. The experiments are initialized from Uy- = M, Vo The middle panel of Figure 2 shows the distance between the true dynamics and the manifold dynamics, as well as the distance traveled by the true dynamics, for SGD and SGDm. On the right panel, we inject isotropic noise on the tangent space for SGD and do the same comparison. The results show good approximation of the manifold dynamics to the real dynamics. Under the Hessian noise, then, the SGD indeed moves towards minima with smaller Frobenius norm, same as the manifold dynamics.
= 5 and compare the
4 LEARNING RATE, BATCH SIZE, AND MOMENTUM
With the manifold dynamics, we can study the impact of hyperparameters on the behavior of optimizers around manifold of minima. Here, we focus on the learning rate n) batch size B, and momentum of the SGDm algorithm. Written in the SDE form, the batch size changes the covariance of the noise by a factor 1/ BB Hence, using equation 16 the SDE with batch size B is
dm, = - -1--" m, + (Z())) +
1--mmal m.d v D(P2))dW, dz, = " VB
(20)
17
To focus on the drift dynamics on the manifold and avoid the influence of the noise, we consider Hessian noise which only exists in the normal space of the manifold, i.e. we assume D1
D2n (z) = o (z). Then, by equation 18 the first-order ODE on the minima manifold representing the manifold dynamics is
2
M Tr(F(22)))
(21)
4B(1 11)
By equation 21 the manifold dynamics takes the same trajectory with different speed for different hy- perparameters. Let 2, be the trajectory of & = o² V M Tr (z)). Considering the discretization, an SGDm with learning rate 17. batch size B, and momentum μ takes TB(1-=) iterations to solve for ε, until = T. Hence, decreasing the batch size, or increasing the learning rate momentum factor, can accelerate the (discrete) manifold dynamics. We let 8(11, Jt, B)
nja
,
B(1=2² the speed factor
for the dynamics. The experiment results in the left panel of Figure 3 justify that the speed factor indeed controls the dynamics' speed.
Implications in practical cases: For the training process of over-parameterized neural networks, the exploration around the minima manifold is an important source of implicit regularization. The driven force of the movement along the manifold is still the change of flatness. However, in this more complicated case, the discussion above may face two problems: (I) The curvature in the directions perpendicular with the manifold may not be quadratic. Which leads to different manifold trajectory if the range oscillation is different. (i.e. the iterator oscillates in a larger range, the flatness driven force may change its direction.) (2) The manifold dynamics is just first-order approximation of the true dynamics, which may not accurate for a long time period. The second problem is intrinsic to all studies that use continuous dynamics approximate discrete dynamics. It will not impose
7

---

Under review as a conference paper at ICLR 2023
Adam net-0"
199
0.16
60
Adami mi0117
Adam, not-3"
10:
1-0.04
Accur
90
D150 20 to
Alam 42
Alam, no--07
0.02
expertment excprtimett
4- 0.001 .0 00000..0----e....
of
9- 1 (0003, M...B.
0.0
0.5
1.n
155
2.0
0.00
0.0
0.5
1.5
2.0
40
80
120
:80
1500 T=itetationxsttp
Benations
Epochs
Figure 3: (Left) The x-coordinate dynamics of three experiments for SGD with different hyperparameters, on function [x. v) = (1 + x² )²² with Hessian noise. Two experiments have the same s(1), j. B), while the other one has a smaller (11), μ4, B). Here the role of B is played by 1/a". (Middle left) The moving average of y-magnitude for the three experiments shown in the left panel. The three experiments have the same r(n, H4 B). hence the y- magnitudes are on the same order. (Middle right) The test accuracy curves of two neural network runs with different hyperparameters but the same s and r. Experiments are conducted on ResNet18 and CIFAR 100 dataset. (Right) The distance between the manifold projection and the origin of Adam iterators for loss function equation 23 with different 0, compared with that for SGD.
a serious problem as long as the curvature on and around the minima manifold does not change drastically. The first problem motivates us to find ways to accelerate the manifold dynamics without changing the range of oscillation. By the discussion
Section [.3] the range of the oscillation is given by Vagdmm which is 1B(1-3 Let r(n, H1 B)
- be the "range factor". Then, Combining the discussions above, we want to increase the speed factor s(n: μI, B) while keeping the range factor r(1,, Jr, B) fixed. Since the ratio 7/(1 μ1) appears in both factors, to increase the dynamics speed without changing the oscillation range we must change B. Concretely, we pick B' = cB and n', pl such that I-p² CI- then
na²
2 BO(1-10) p) B²(1- w')
17
B(1-w)² - 17
(22)
i.e. the range factor is not changed while the speed factor is multiplied by c. In the second panel of Figure 3] we show for synthetic problem that r(n, μ. B) is indeed proportional to the range of oscillation. In the third panel, show for neural networks that two experiments with the same speed and range factors but different hyperparameters indeed follow the same test loss curve with respect to epochs. On the other hand, if the oscillation range not kept, the training trajectories may go to different solutions (with different training and testing error).
By equation 22 increasing the batch size while changing other hyperparameterss ac- cordingly can accelerate the training speed without changing the training trajectory. However, this acceleration happens on the level of number of iterations. Since the batch size changes accordingly, the number of samples used during the training period does not change. This means we are training for less iterations, but the same number of epochs. This is shown in the third panel of Figure 3 Nevertheless, this usually saves time because training one big batch is faster than training several small batches with the same number of total samples. Therefore, our results reveals theoretical mechanism underpinning the empirical benefit of large butch training You el al. 2019 2017
Remark 2.
Hoffer
el al. 2017;; Geiping al. (2027).
ADAPTIVE GRADIENT METHODS AND ROTATIONAL INVARIANCE
We can also study the manifold dynamics of adaptive gradient methods. We start with experiments which show that the manifold dynamics of Adam changes according to the direction of the manifold. This is due to the fact that adaptive gradient methods are not rotational invariant. We consider the loss function,
(23) which is the counterclockwise rotation of (a² + y)y² 2 by The minima manifold of f is the line x sin 0 + y cos = O. We run Adam Kingma & Ba 2014) on f with different 0 with Hessian noise. The right panel of Figure compares
f(z, v) (a sin 0 y cos 0) 2 (Ix cos y sin 0)2 + 1),
comprees the dynamics projected onto the minima manifold. The results
"Our analysis works for any loss function with the form f(z, v) (ar sin + yoss 0) h cos y sin 0).
8

---

Under review as a conference paper at ICLR 2023
show that Adam moves very fast along the manifold when the manifold aligns well with axis. When the manifold does not align with an axis, Adam moves much slower, sometimes even slower than plain SGD. This is because in Adam the adaptive learning rate is computed for each coordinate, and hence when the manifold direction close to an axis, the learning rate along the manifold can be drastically increased due to the small gradient on this corresponding axis direction. Otherwise, all axis directions have big gradients due to the oscillation and the learning rate along the manifold is not increased desirable way.
Using the SDE approximation for Adam recently derived in Malladi et al. 2022). we can derive and compare the manifold dynamics for Adam on the loss equation 23 for different Consider an Adam algorithm with hyperparameters ($1, B2, 17, e). where 31 and B2 are momentum coefficients for the first and second order moments, respectively, 17 the learning rate, and 6 the small number that prevents division by zero Kingma & Ba (2014). By Malladi et al. (2022). let E be the gradient noise covariance matrix depending on the parameters, and be a additional noise strength (i.e. the real noise covariance is o)), define do om, 6o E7), C1 (1 31)/n² 02 (1 32)/1/² (t) - e cit and (t) = 1 oat then the Adam trajectory is approximated by the SDE:
dx, rl) P. P.'madh Mtd, (t)
dm, c1 (Vf(x)) m;)dt ooc1
(24)
du (ddigg((((xx))) Un dt,
Pt ondiag(u.)) + /
Here, f the loss function, x is the parameter, and W, is Brownian motion. The time scale of the SDE above is t = kny² which is different from the usual time scale t = kn studies for other optimization algorithms.
Using the quasistatic approach, in Appendix B we derive the approximate manifold dynamics for two cases: (1) = 0, in which the minima manifold aligns with one axis, and (2) 0 = x/4, in which the angles between the minima manifold and coordinate axes maximized. After some approximations which are detailed in Appendix B (such as (!) = (t) = 1 which happens when t is large), we have the following effective dynamics on the minima manifold:
agh'(xa)
Mxx1 dt, dx. Ma.t dt. co
0 = 0: dm =1
(25)
44/((z;)
V2mg.t dt. Mzxt dt, dx,
0 dm C1
(26)
(xx)
ao /ov/ (h(aa))
Here, x is the coordinate along the manifold direction, and m is corresponding momentum. The SDEs equation 25 and equation 26 show the difference of the manifold dynamics for different When 0 = 0, the dynamics is very fast, because of the on the denominator. When 0 x/ x/4, instead, the x dynamics slower. we further make a first-order approximation of the dynamics by assuming the momentum is always equilibrium, like we did for SGDm in equation [8] following manifold dynamics:
we have the
onh'(z) 400
h'(x)
==0 0:
I
(27)
2h(z)
Here we see that when 0 = 0 we get a gradient flow minimizing v /h(x)) on the minima manifold, while when 0 = T/ /4 get gradient flow minimizing In h(x). The former dynamics is much faster due to the co on the denominator. For the detail of the analysis please see Appendix B When the dimension of the parameter space is high, it is hard for the minima manifold to align well with coordinate directions. Hence, the exploration of Adam (as well as other adaptive gradient methods) on the minima manifold is slower than SGD and SGDm. This may be one reason that Adam does not generalize as good as SGD in many cases Keskar & Socher (2017); Wilson et al. (2017).
Remark 3. Unlike the effective dynamics for SGD and SGDm, the dynamics in equation 27 do not depend on the learning rate n7. This is because in the SDE approximation the time scale is = knp. Therefore, a n factor will appear if we transform the time scale I. kn-

---

Under review as a conference paper at ICLR 2023
REFERENCES
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an omstein-uhlenbect like process. In Conference on learning theory, pp- 483-513. PMLR, 2020.
Yaim Cooper. The loss landscape of overparameterized neural networks. arXiv preprint arXiv: 10200, 2018.
Xiaowu Dai and Yuhua Zhu. On large batch training and sharp minima: fokker-planck perspective. Journal of Statistical Theory and Practice, 14(3):1-31, 2020.
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34, 2021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Jonas Geiping, Micah Goldblum, Phillip E Pope, Michael Moeller, and Tom Goldstein. Stochastic training not necessary for generalization. arXiv preprint arXiv:2/09. /4/ 2021.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli,, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
Florian Herzog. Stochastic differential equations. lecture notes, 2013.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks. Advances in neural information processing systems, 30, 2017.
Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint arXiv.1712.07628, 2017.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 1411v:1412.6988, 2014.
Nikola B Kovachki and Andrew M Stuart. Continuous time analysis of momentum methods. Journal of Machine Learning Research, 22(17))1-400 2021.
John M Lee. Smooth manifolds. In Introduction to smooth manifolds, pp- 1-31. Springer, 2013.
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pp. 2101-2110 10. PMLR, 2017.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in avrr-parameterizee matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34:12712- 12725, 2021a.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint 1o.iv:2/10.069// 14, 2021b.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv: 1906.99899990 2019.
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34, 2021.
10

---

Under review as a conference paper at ICLR 2023
Chao Ma and Lexing Ying. On linear stability of sgd and imput-smoothness of neural networks. Advances in Neural Information Processing Systems, 34, 2021.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. arXiv preprint arXiv:2205. 2022.
Takashi Mori, Liu Ziyin, Kangqiao Liu, and Masahito Ueda. Logarithmic landscape and power-law escape rate of sgd. arXiv preprint arXiv:2105.09555) 2021.
Behnam Neyshabur, Ryota Tomioka, Ruslan Salakhutdinov, and Nathan Srebro. Geometry of optimization and implicit regularization deep learning. arXiv preprint urXiv:1705.088: 2017.
Thanh Huy Nguyen, Umut Simsekli, Mert Gurbuzbalaban, and Gaël Richard. First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. Advances in neural information processing systems, 32, 2019.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26-31, 2012.
Yuging Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect. arXiv preprint 2rXiv:21 2021.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stem, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. Advances neural information processing systems, 30, 2017.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference Learning Theory, pp. 3-35-36633 PMLR, 2020.
Lei Wu, Chao Ma, et al. How sgd selects the global minima in pee-pparameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint urXiv:2002.03495 2020.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 03888, 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 1904.00922 2019.
A SGD WITH MOMENTUM
Consider the following SDE approximation for SGDm:
1 dm, 11 P m, + (f(z)) + )iD(Pz))dW, dz,
d,tt. 7
(28)
By the discussion for SGD, the manifold dynamics can be obtained by assuming that M flat. The result for non-flat manifold is different only by gradient taken the manifold. In the flat case, still
11

---

Under review as a conference paper at ICLR 2023
letting [x", 1 with x in the tangent space and y in the normal space, the dynamics 28 be written as an SDE system for x and y.
μ m t of vn (x))()))))) dt /nDii(x))dWW)
dmg,1
dx: m rdt,
" 1 my + (x))1) di (x)
1 m.
dmy,
dy.
(29)
1
")
Here, mg. and my. denote the momentum for x and components, respectively. Again, since the y dynamics is faster than the dynamics, we take quasistatic approach by assuming yı is always at in the dynamics. Note that the two equations for my.t and y, form linear system of SDEs, we can still compute the first and second moments of at equilibrium and substitute the results into the equations for myxt and which gives the following effective dynamics for x
the equilibrium given X;, and taking expectation
( 1 17
dmx,1
m
(Vaddm) (x(H(x))))) dt + (TDn(xx)dW, dx; dtt
1
where the form and derivation of V Vagdm are given in the next subsection. Finally, replacing x by x on M and consider gradients on M, we have the following effective dynamics general minima manifold:
n-
dm, =
m
+
(V. (Vagdm) ((((((Z)))) dt + VnDiu(z)dWWs dz, = - -m,tt,
1
1.j=1
(30)
A.1
THE DERIVATION FOR Vagdm
In this section, we derive Vegam from the SDE for SGDm equation [6 Assume x; and my, is fixed, search for the equilibrium of the following system of y and my:
dmy 1 f
my. 2H(()))) dr + (x
"
dy. T = 1)
-y.
The SDE system above is linear. Let u,
my.r
the SDE can be written
yr
du, du viDdBr.
(31)
where we have
(xx)
E (R=
0 D22
E (R2(d-))xd
A
D =
0
and B, is Brownian motion.
Let C₂ be the second moment matrix Ex,, u. By Herzog 2013 C, satisfies the ODE
d C, AC, C-A" DDD
d
dr
Therefore, taking - xi, let C.. lim Cr be the moment matrix equilibrium, we have
AC.o + CTo -DDD
(32)
By the definition of u, we have
Emy m y.20 Eyy m
Emy yo
Coc
Ey
y.o0
12

---

Under review as a conference paper at ICLR 2023
We are interested in the Yoo in the above matrix. By the symmetry of Co. let
C C2 c
].
Cx -
C2
Then, want derive Cz. Substituting the blockwise Coc into equation 32 we have
(x)(x)
][
C
C2
C2
1
1
D22
D22
Dr
0
LT
c
C3
+
ch
C3
(x)
0
= -1
0
which gives
2(1 2(1- Ci - ((F - + C2H((x)
D22 D3
(33)
11
- 1 1-H C2 (x) =
(34)
11
C2 1 c 0.
(35)
C1 and C33 are symmetric. Hence, adding equa- By equation 35 C2 is skew symmetric. By definition, tion 34 with its transpose, we obtain
Lo +((x))(++C((()x))) O,
which gives
C (+(x))+
(36)
Substituting into equation 35
we have
C2 = (
(37)
Plugging equation 36 and equation
37 into equation
33
we have
2n
2(1 ([I=))(((x(C)) + Cxll(x))
I (Cali(x)² + ixx)) (1(())
=nD22D
(38)
There, denote H = -x),, Vegdm is the solution of
HV + +V
(1- VH²² + -2VVA
"1 D DT
(39)
2(1 - 4)
B ADAM AND RMSPROP
In this section, we derive the effective dynamics equation 25 and equation 26 for Adam, the 2-D problem equation 23
f(r = (rr sin + y cos ((z cos y 0)² + 1).
first describe the relation between two coordinate systems. Let aOy be coordinate system on which the loss function equation 23 is defined, and Adam is conducted. Let x' O'y' be the coordinate system obtained by rotating zOy counterclockwise by 8. Then, the w' axis aligns with the direction of the minima manifold of In this coordinate system, f has the form:
f(z, v) ((x')² + 1)(2)²
Let Rg be the rotation matrix counterclockwise by 0, i.e. Ro cos -sn 0
Then, for any
sin 0 cos
vector in rOy, let x' be its coordinate in x' z'Oy'' have x' R_xx, and Rox'.
13

---

Under review as a conference paper at ICLR 2023
Recall the SDE for Adam derived in Malladi et al. (2022).
72(1) P,lma (11)
dx. du, c2 cy(diag(E((x))) Uy uy)dl,
dm; = 01 (V](x;)
- m:)dl + onci E1/2 (x)),W
(40)
Pt= ( 1/2 v2())
where x is the parameter vector, m is the first momentum vector, is the second momentum vector, and E is the noise covariance. All these quantities are defined the coordinate system rOy. Let x', m', E' be the counterparts of x, m, E in x d'Oy'. then
= R m' R. R
(41)
We do not consider u' as in r'Oy', because is not rotationally invariant. This is the reason that Adam has different effective dynamics on the minima manifold for different 0.
By the relations equation 41 the SDE equation 40 can be written as
v12(t) (t)
dxF
R Rom dmj = (R-o mt)dt + obce +oceE'1/22 //1/2 (x, )dB,
(42)
du, =2((digg (x(x)) uz)di, where B, = R_WW,
Pt oodiag(u,) + r2(t
W/Wi,, is a Brownian motion in the r'Oy/ system. Note that R._ V fxx) is the gradient of f in the o'Oy' system, letting x [z. y/" and [x', v')", we have
22²²) 2(x + 1)y
R-V
From now on, we denote hxx) = x² + 1. Then
(aa') (h((')s)
}.
RoV =
Our analysis actually works for any positive and differentiable function h. Under the Hessian noise assumption, take
0 ).
Σ'(x)
h(2²)
then for V we have
sin? 0 sin cos 0
sin 0 cos 0
Σ(x) R (x) h((²))
cos² 0
We do not add a a before / /1(x') because in the derivation of the SDE equation 40 strength factor is included into ag-
By the discussion on V and E'(x), letting m' = [my, my u = [u, u] we can write equation 42 into the following system of x', 1/ my,, My', u, V:
dz
(t()
m dt My dy
(1)) R P.- Ro
dyl
dm ci (h' (x/))// my dt
dmy = ei (2h(2/)vl my () + Jaci / /
du = c2(h(z)) sin u, d
du, c2(h(z)) cos2 0 ddt
0 + EO 72())1.
P. =
(43)
2t
In equation 43 Next, consider two cases: and -.
B, is 1-D Brownian motion.
14

---

Under review as a conference paper at ICLR 2023
Case 0 0.
When 0 0, we have Rg R_o I. Also, since
1/2
0or co co)
P.
00
Ut
+ c0 72(1)
0ove: + eov 72(1)
we have
ooyarteo 72(1)
P
Joy / (t())
Therefore, equation
43 can be written as
rr(() My'.1
de,
dl
71(1) ou + con (2())
dy' V22(() my/it
dl
(t))) u + cu (1a(t
dm ci (h'' (x'))ul - +)
dmy', (2h((4))' my my ooei v h((/)dB.
du, -c2ud
de, - dt.
-
(44)
Using the quasistatic approach, we assume the dynamics of y/ and my is at equilibrium at any fixed Hence, fixing z/, we consider the system
dy/ (2((1)
Myyre
(1) To + r2(( dr
dmy', ei (2h(z))v( mgg.-)d + ouei V h(r))dB,
and compute lim Ey,-. This is linear SDE system of y' and my,, by techniques in Herzog (0013)
we have
o (t())
lim Eyf
(45)
T-000
4m (t) + (t)
Substituting equation 45 into equation
44
we have the following effective dynamics only on the
minima manifold:
(2(t)
dr,
Jovuu Ue + 6D (1( dt
71())
o (t)h'() (xt)
dm = Ci (t) + +00
Maye
du: dr
du; c2 (2((aa)) v dtt.
(46)
Next, we try to make the manifold dynamics equation simpler by doing some approximations. First, solving the ODEs for and Vr, we get
est + 02 L' cs(s- (
U = uge
egt
Vt = Uhe
We first assume t is big enough, such that - is close to 0. Then, we can take
U4 O, it = C2 (=2 c=(ls-t)((z)dss
and also 7/2 (t) = (2(+)= = Moreover, since the dynamics of x' is slow, we can assume the change of h(xx)) is slow compared with ccs In this case, we have
62 eslo- M = h(xl))
15

---

Under review as a conference paper at ICLR 2023
Hence, can approximate u, and by
Wr O,
- h((())
(47)
Substituting equation into equation approximate manifold dynamics:
46 and taking 71 = (2() 1, we get the following
Mx'd
dxy.
60
ogh'(zt)
dmz
Ci
mex dt.
(48)
Aaov (1(x)) + 4en
Finally, in the denominator of the dynamics of my, the 6o term is usually small compared with the do (x²)) term. Hence, can take 400 /h(x)) + 4ea 40 vh(z() and write the following approximate manifold dynamics:
dz² Me'd dt 60
onh'(a')
dm
Mate dl.
(49)
4v/(xz))
Note that M(x) (2 h(r)), the dynamics equation 49 can be understood as gradient flow with Vh(u)
momentum that minimizes vh on the minima manifold. In the region that h is close to linear, we can suppose
(oh)(at)
ma'
Av((a))
and write down the following first-order approximation of the dynamics:
dxd
onh' (x+)
(50)
dt
4Eo [h(z))
time unit of the manifold dynamics that we derive corresponds to 1/1/² steps in the SGD trajectory. i.e. each SGD step corresponds to time period of 2/2. If we change the dynamics the usual time scale, each SGD step correspond to time period 7), an additional y will appear on the numerator of the dynamics. Therefore, the dynamics still slow dynamics that is an r) factor slower than the original SGD dynamics.
Note that due to the time scale chosen in the SDE equation 40
Case 0= 7. When 0 = -. by equation 43 the dynamics of and are
M(x²))
du, C2
we dt)
2
h(x²)
dv, C2
vt)dt. dt.
2
In this case, u, and U have the same dynamics. If we assume ugg Up. then have 24 Ue for any
t 2 0. Then,
P. P-
ao toy Ta(t)
I.
0
To + Eo l2()
Payo+teo T0(1)
Hence,
1
R_PP.' Rg
R
I.
TO + (t)
To + Eo r2())
16

---

Under review as a conference paper at ICLR 2023
Therefore, equation 43 can be written as
del (2(()
mrr,t
dt
71(1)
/ + (2(1)
dyl V22(t) my's
(1)) o0 ooveu + co / 72())
dm C1 (h' (x/))/)// a)
dmy C1 (2h(z)))/
my
my +) dt + ooci v Vh(r/)dB.
h(24)
du. C2
dl.
(51)
2
The quasistatic step here still deals with the system of yf and my and the results take the same form. we have
o 72(1)
lim Ey
7-000
( th + co 72())
and the following effective dynamics the minima manifold:
r2
dæl,
dt
(1(1)
0 + co vov 72(1)
dm
-az's dt
(oioo)ooo//u + + =
(Mx²)) ue dt.
du C2
(52)
To simplify the manifold dynamics above, we take the similar approximation steps as did for the f = 0 case. We first solve the u dynamics, which gives
= uge C2t ( (es(o-t) h(x()dss.
Still assume is large, and the change of h h(x)) is slow compared with e- Then, we can approxi- mately take
h(*²))
(t) = (t) = 1, Ut
2
Substituting the above approximations into equation 52
we obtain
1x'51
da'p
A(x!) 60
To
dm Ca
ma'd dt.
(53)
2v20. /hxx)) + 4ep
Dropping the co terms the denominator, equation 53 is approximated by
V2m
dxy
aovh(z²)
oph((x1)
dm C
-mz'' dt.
(54)
2 22/))x
(al)
Finally, we assume my is close to its stationary solution, i.e. ma' 2/22(a)))
we have the
following first-order dynamics for x' that approximates the manifold dynamics:
dz 11(x()
(55)
dt (x(x)))
17

---

Under review as a conference paper at ICLR 2023
The dynamics equation 55 is gradient flow that minimizes In h(x) the minima manifold. Again, due to the time scale choice, the dynamics gets slower for smaller learning rate. Compared with the dynamics equation 50 for the case 0, this dynamics is slower because there is no Eo on the denominator.
Remark 4. The approximation steps the derivations above are conducted intuitively without rigorous proof. The goal is to unveil and compare the essential components of the dynamics. If strict theorems are to be proved, conditions and assumptions need to be imposed.
18


Reviews:
[{"review_id": "xINkiBb5rZb", "reviewer": "Reviewer_pGAs", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "The authors provide a technique to approximate the dynamics of stochastic optimizers that typically oscillate around a manifold, which corresponds to the minima of the loss. This allows to study and compare the properties of stochastic optimizers, for example the tendency to search for flatter minima. In the experiments empirically compare the behavior of the proposed approximation to the true dynamics, and demonstrate the ability to study properties of the optimizer.", "summary_of_the_review": "Questions:\n\n1. I think that some illustrations can help significantly and make the paper much more accessible, as there are a lot of technical details that can be probably summarized accordingly.\n\n2. What is the definition of the minima manifold? Is it the set in the parameter space where the loss function is zero? What is the effective/manifold dynamics? Perhaps these terms should be defined explicitly.\n\n3. I believe that a description of the quasistatic dynamics should be included at least in the appendix, while now the context is roughly presented around Eq. 2.\n\n4. For Eq. 2 the quasistatic means that we have a coupled system of differential equations, and near the minima manifold the $y_t$ moves faster than $x_t$, where \"faster\" simply means that the step along the y-axis is bigger? The implied trajectory of the derived system is for every time step $t$ a point $y_t$ sampled from the associated Gaussian while the $x_t$ follows the dynamics of the equation after Eq. 2?\n\n5. In general, what is the intuition of the dynamical system after applying the quasistatic derivation? Can we still dicretize the process and get a discrete trajectory?\n\n6. Is the high level idea of your approach that according to your approximation, when we are near the minima manifold, the optimizer moves along this manifold searching for a better minimum, while along the normal direction the optimizer moves \"faster\" (jumping along the normal direction), but in expectation we remain near the manifold?\n\n7. The derived dynamics are not for the actual loss function, but for the surrogate function $f$ (Eq. 5) that is a quadratic approximation to the true loss (using Taylor's expansion)?\n\n8. I suppose that the particular structure of the Hessian with many zeros, holds only when the minimizer is on an axis of $\\mathbb{R}^d$ e.g. [x, 0]. Because, even if we have a flat minima manifold that is not axis aligned, then the particular Hessian does not necessarily have this particular structure. How the analysis changes if the Hessian does not have all these zeros? Probably some extra linear terms $y_\\tau$ appear and their expectation goes to 0?\n\n9. How the analyis changes if the the assumption about the uncorrelated noise (between tangent and normal space) does not hold?\n\n10. I think that it is not clear where the discussion for the orthonormal tangent spaces at the end of page 4 is actually used.\n\n11. For the general manifold it is assumed that the point $z_0$ is axis aligned, so the Hessian has the structure with zeros. How the analysis changes if this is not the case?\n\n12. What is the intuition for the companion loss function $\\tilde{f}$ ? If I understood correctly this function takes points around the $T_{z_0}\\mathcal{M}$, maps them around the $\\mathcal{M}$, and then these points are projected to $\\mathcal{M}$ such that to use the corresponding Hessian of the final point?\n\n13. The actual implication for Eq. 10 and Eq. 11 is that $z_0$ lies on the axis and in a local neighborhood the minima manifold can be seen as flat? So the actual analysis for the general manifold concludes that in a local neighborhood the dynamics are exactly the same as the flat manifold, as the manifold locally is considered flat?\n\n14. Does the $\\nabla_\\mathcal{M}$ implies that we actually compute the gradient on the manifold which returns a \"tangent vector\"? Also, what is the dimensionality of Eq. 12? \n\n15. Could you provide some empirical evidences for deep networks? \n\n16. In the \"Example\" (Eq. 13) it is implied that the Hessian has this particular structure with zeros, and for this reason the $\\tilde{D}_{11} = 0$, but I think that this only holds if the minima are axis aligned?\n\nIn general, I like theoretical works that analyze the problem in an accessible manner. I think that the proposed idea has some very interesting aspects. However, in my opinion the current version of the paper is quite confusing in some parts. I acknowledge the fact that theoretical work is typically harder to explain and simplify, but I believe that some improvements can be made. In my opinion, the current version is not ready for publication mainly for clarity and accessibility issues.", "strength_and_weaknesses": "The problem of understanding the learning dynamics of stochastic optimizers, especially for over-parametrized neural networks, that the paper aims to address is indeed important in the field. The authors introduce this problem mainly based on works related to deep networks, and even if I am not an expert in this field, I think that the related work is decently presented. I find the motivation for the proposed approach ok, but I believe that perhaps this can be improved (see questions). The high level idea of the approach is interesting and seems sensible, but in some parts the report is rather confusing and I think hard to follow in details (see questions).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "inroWqpEW_g", "reviewer": "Reviewer_fp5w", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "The authors use a quasistatic approach to analyze an SDE related to SGD near minima of the loss landscape. They find a coupled SDE that has two component, one tangent and another normal to the manifold of minima. The drift terms in the equations lead to a separation of timescales where the normal component is significantly faster and is treated as being in equilibrium. Plugging in the equilibrium behavior of the normal component into the SDE for the tangent component's equation reveals an implicit regularization term that is minimized by SGD. The authors apply this analysis to SDEs modeling SGD, SGD with momentum, and Adam. ", "summary_of_the_review": "I enjoyed reading the paper and feel it contains some nice results. I don't feel all the implications of the analysis are tied back to realistic models---more could be done in this direction. The main concern is novelty, especially with respect to Li et al. I would be happy for the other reviewers to comment on this.", "strength_and_weaknesses": "**Strengths**\n\n- The paper is clearly written with nice illustrative, simple examples to explain the core ideas.\n- Fig 3 is nice. How are the s and r kept the same in (Middle right)? Which hyperparameters are changed between the two models?\n\n\n**Weakness**\n\n- The assumptions seem to be doing a lot of work in the analysis. For example, alignment of the noise covariance and the Hessian is frequently assumed. It's not clear to me how reasonable such assumptions are for real models. \n- On page 4: \"In many cases, the noise covariance matrix of SGD aligns with the Hessian Mori et al. (2021).\" Is this an empirical fact?\n- My main concern is novelty.", "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": "PyJVxGc9hR", "reviewer": "Reviewer_KFCz", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "This submission proposes a quasistatic approach to derive the optimization algorithms's behavior on the manifold of minima. It has tried to understand the role of some parameters such as learning rate and batch size. Some unrealistic assumptions have been made, the results are not supported by rigorous analysis and are not well verified by solid empirical studies. ", "summary_of_the_review": "Some unrealistic assumptions have been made, the results are not supported by rigorous analysis and are not well verified by solid empirical studies. ", "strength_and_weaknesses": "1. There are a number of assumption used without being justified. For example, it is assumed that the Hessian has zero eigenvalues along the tangent space of $\\mathcal{M}$. And in eq(18), it is assumed that the momentum is always at equilibrium. How to justify these assumptions for over-parameterized deep neural network? \n\n2. The claims are not well supported by theory, instead it is more of some intuition for some specific problems, not for over-parameterized neural network as claimed. For example, the loss function is approximated by a squared function based on Hessian matrix. But how accurate is this approximation for a deep neural network? In what range is this approximation accurate? These are all critical components of the analysis and they need rigorous analysis. \n\n3. The study in Section 3 and Section 5 are for quadratic convex problems which cannot directly provide intuition for deep neural networks which have much more complicated landscape. \n\n4. The experiments are not adequate at all to verify the proposed approach. Especially considering a rigorous theoretical analysis is lacked, it takes solid empirical studies to make the results convincing. \n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "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": {}}]

Author rebuttal:

Decision: Reject