A view of mini-batch SGD via generating functions: conditions of convergence, phase transitions,  benefit from negative momenta.
Authors: Maksim Velikanov
Denis Kuznedelev
Dmitry Yarotsky
Venue: ICLR 2023 poster (2023)
Keywords: SGD
linear models
optimization
analytic framework
NTK

Abstract: Mini-batch SGD with momentum is a fundamental algorithm for learning large predictive models. In this paper we develop a new analytic framework to analyze noise-averaged properties of mini-batch SGD for linear models at constant learning rates, momenta and sizes of batches. Our key idea is to consider the dynamics of the second moments of model parameters for a special family of "Spectrally Expressible" approximations. This allows to obtain an explicit expression for the generating function of the sequence of loss values. By analyzing this generating function, we find, in particular, that 1) the SGD dynamics exhibits several convergent and divergent regimes depending on the spectral distributions of the problem; 2) the convergent regimes admit explicit stability conditions, and explicit loss asymptotics in the case of power-law spectral distributions; 3) the optimal convergence rate can be achieved at negative momenta. We verify our theoretical predictions by extensive experiments with MNIST and synthetic problems, and find a good quantitative agreement.

Published as a conference paper at ICLR 2023
A VIEW OF MINI-BATCH SGD VIA GENERATING FUNC- TIONS: CONDITIONS OF CONVERGENCE, PHASE TRAN- SITIONS, BENEFIT FROM NEGATIVE MOMENTA
Maksim Velikanov
Denis Kuznedelev
Technology Innovation Institute, CMAP, Ecole Polytechnique maksim. velikaa velikanov@tii.aa
Skoltech,
Yandex Research
denis kuzredeleeessklttech..t
Dmitry Yarotsky
Skoltech d. d.jarotskyysskottech.cr
ABSTRACT
Mini-batch SGD with momentum fundamental algorithm for learning large predictive models. In this paper we develop a new analytic framework to analyze noise-averaged properties of mini-bater SGD for linear models at constant learn- ing rates, momenta and sizes of batches. Our key idea is to consider the dynamics the second moments of model parameters for a special family of "Spectrally Expressible" approximations. This allows obtain explicit expression for the generating function of the sequence of loss values. By analyzing this generating function, find, in particular, that the SGD dynamics exhibits several con- vergent and divergent regimes depending on the spectral distributions of the prob- lem: the convergent regimes admit explicit stability conditions, and explicit loss asymptotics in the case of power-law spectral distributions: 3) the optimal convergence rate can be achieved negative momenta. We verify our theoret- ical predictions by extensive experiments with MNIST, CIFARIO and synthetic problems, and find good quantitative agreement.
INTRODUCTION
We consider classical mini-batch Stochastic Gradient Descent (SGD) algorithm (Robbins &
Monro 195I Bottou & Bousquet] 2007) with momentum (Polyak 1964):
Wc+1 W. Ve+1,
Vi+1 -a V LB. (w.) B.V.
Here, [.((w) = Σ [(f(x,, w), v.) is the sampled loss of model 7 f(x, w). computed using a pointwise loss 1( y) on mini-batch B {(x, yx), 1 of b data points representing the target function previous iterations and well-known to significantly improve convergence both generally (Polyak) 1987) and for neural networks (Sutskever et al., 2013). Re-sampling of the mini-batch B, each SGD iteration creates a specific gradient noise, structured according to both the local geometry of the model f(x, w) and the quality of current approximation learning, f(x, w) is usually very complex, and the quantitative prediction of the SGD behavior becomes challenging task that is far from being complete at the moment.
f (x). The momentum term v.. represents information about gradients from
In the context of modern deep
Our goal to obtain explicit expressions characterizing the average case convergence of minh-bbtch SGD for the classical least-squares problem of minimizing quadratic objective I(w).. This setup is directly related to modern neural networks trained with a quadratic loss function, since networks can often be well described e.g.. in the large-width limit (Jacot et al.) 2018 Lee Lee al. 2019) or during the late stage of training (Fort et al. 2020)- by their linearization w.r.t. parameters w.
A significant part of the work done while Skoltech.

---

Published as a conference paper at ICLR 2023
A fundamental way to characterize least squares problems is through their spectral distributions: the eigenvalues 2y of the Hessian and the coefficients CK of the expansion of the optimal solution can estimate certain metrics of the problem through
w" over the Hessian eigenvectors. Then, spectral expressions, i.e. explicit formulas that operate with spectral distributions Aς. Ck but not with other details of the solution w or the Hessian. A simple example is the standard stability condition for full-batch gradient descent (GD): a < 2/2mxx. Various exact or approximate spectral expressions are available for full-batch GD-based algorithms (Fischer 1996) and ridge regression (Canatar et al. 2021] Wei et al. (022). Here, we aim at obtaining spectral expressions and associated results (stability conditions, phase structure, loss asymptotics....) for average train loss under mini-
batch SGD.
An important feature of spectral distributions in deep learning problems is that they often obey macroscopic laws- quite commonly power law with a long tail of eigenvalues converging to 0 (see Cui et al. (202T) Bahri et al. (202T); Kopitkov & Indelman (2020); Velikanov & Yarotsky (202T); Atanasov al. 2021 Basri et al. (2020) and Figs. 1 The typically simple form of macroscopic laws allows to theoretically analyze spectral expressions and obtain fine-grained results.
As an illustration, consider the full-batch GD for least squares regression on MNIST dataset. Standard optimization results [Polyak] 1987) do not take into account fine spectral details and give either non-strongly convex bound LoD(W,) O(t or strongly bound LGD(w,) I L(wo)( Amax Amin )zt Both these bounds rather crude and poorly agree with the experimen-
Amax A
tally observed (Bordelon & Pehlevan, 2021: Velikanov & Yarotsky] 2022) loss trajectory which can be approximately described as L(w,) ~ Ct E 0.25 (cf. our Fig. []. In contrast, fitting power-laws to both eigenvalues d and coefficients eg and using the spectral expression (G.( = (1 a.)x)2 ) allows to accurately predict both exponent & and constant C. Accordingly, one of the purposes of the present paper is to investigate whether similar predictions can be made for mini-batch SGD under power-law spectral distributions.
Outline and main contributions. We develop a new, spectrum based analytic approach to the study of bacc- SGD, The results obtained within this approach and its key steps are naturally divided into three parts:
1. We show that in contrast to the full-batch GD, loss trajectories of the mini-batch SGD cannot be determined merely from the spectral properties of the problem. To overcome this difficulty, we propose a natural family of Spectrally Expressible (SE) approximations for SGD dynamics that admit an analytic solution. We provide multiple justifications for these approximations, including theoretical scenarios where they are exact and empirical evidence of their accuracy for describing optimization of models on MNIST and CIFAR 10.
2. To characterize SGD dynamics under SE approximation, we derive explicit spectral expres- sions for the generating function of the sequence of loss values, L(x) = E, L(w))z', and show that it decomposes into the "signal" (2) and "noise" (()) generating functions. An- alyzing (z), we derive a novel stability condition of mini-batch SGD in terms of only the problem spectrum Ag. In the practically relevant case of large momentum parameter 8 x 1, stability condition simplifies to the restriction of effective learning rate Ocft = I-3 < Xepis with some critical value Acrit determined by the spectrum. Finally, we find the characteris- tic divergence time when stability condition is violated.
3. By assuming paw distributions for both eigenvalues A -- and coefficient partial sums Sa = Eise λreî x --* we show that SGD exhibits distinct "signal-dominated" and "noise-dominated" convergence regimes (previously known for SGD without mo- menta (Varre al. (02T)) depending on the sign of x + 2v. For both regimes we obtain pawer-law loss convergence rates and find the explicit constant in the leading term. Using these rates, we demonstrate a dynamical phase transition between the phases and find its characteristic transition time. Finally, we analyze optimal hyper parameters in both phases. In particular, we show that negative momenta can be beneficial in the "noise- dominated" phase but not in the "signal-dominater phase.
We discuss related work in Appendix A and experimental details in Appendix F
Our code: https: http:s/////oomousss.mm 4ow.:ouuuu.ep............ htt..//innn///////eraaaa////oo/e
1401/
2

---

Published as a conference paper at ICLR 2023
b=10
b=100
Eigendecomposition
105
10a
NN
NN linearized SE (T
linearized SE
SE [T=
SE (T
Ay/Ag, 1.30
Ct 022
Cr-025
Su/S.. 0.34
10M4
10-2
Optimization trajectories and spectral distributions for MNIST. Left, center: Loss tra- jectories of different training regimes for batch sizes 10 and 100 (see Sec. 3 We see good agreement of NN with SE approximation at T = 1. Right: Eigenvalues A and partial sums S₁ -u=>i -I=E A,Co characterizing H and Co. We see good agreement with assumed spectral power laws 0 Ak-",S, the right figure. The dotted lines in the left and central figures then show Lapprox (t) Ct- with C' coomputed Csignal in Eq. (28)) and x 0.25. We see good agreement between Theorem and experiment. Similar experiments for CIFARIO 10 with ResNet/MobileNet are given in Fig. 9 and also agree well with theory.
Figure I:
Kk- in Eq. 25). We estimate the values A, K, x from the curves in
//
2 THE SETTING AND SGD DYNAMICS
We consider linear model f(x, w) (w, (x)) with non-linear features +b(x) trained to ap- proximate a target function y" (x) by minimizing quadratic loss L(w) 1 Σ ((w, (bx)))
2N
y" (x;))² over training dataset D {x:]_= The model's Jacobian given by v (xb())) the Hessian H t Ei (xx,) @ (xxx))
N T and the tangent kernel K(x, x)) (xb(x). (xx')). (xx))) Assuming that the target function is representable as y*(x) = (w", xb(x)), the loss takes the form L(w)= - / w H(w w")). We allow model parameters
be either finite dimensional vectors or belong to a Hilbert space H. Similarly, dataset D can be either finite (N 20) or infinite, and in the latter case we only require finite norm of the target function y (x) but not of the solution w'. Note that this setting is quite rich and can accommodate kernel meth- ods, linearization of real neural networks, and infinitely wide neural networks in the NTK regime. See Appendix B for details. In the experimental results, we refer to this setting as linearized (to be distinguished from nonlinear NN's and the SE approximation introduced in the sequel).
Denoting deviation from the optimum as Aw w - w", we can write a single SGD step at iteration t with a randomly chosen batch B, of size b as
A Aw Ve+1 - (1- -o, H(,) -
a ((,)
I.I (Aw:). 3.I
1 H(B,) = Σ (xx,) @ (x()).
(2)
iE
important feature of the multiplicative noise introduced by the random choice of B, is that the n'th moments of parameter and momentum vectors Aw Vi-1 in dynamics (2) are fully determined from the same moments at the previous step t. As our loss function is quadratic, we focus on the second moments
C = Al & Aw], J = E/Aw v., V = Evvv,, M ( V)
(3)
Then the average loss is [[L(w)] Tr(HC) and for combined second moment matrix we have Proposition Consider SGD ①1 with learning rates at, momentum 3r, batch size b and random uniform choice of the batch B. Then the update of second moments is
CCt J+t1 Jhn V441
I H.H
31 3I)
C V.)
(--ortt
H,H
31 3.1]
T +
C E E(C,) E(C,)
=
-
J
3

---

Published as a conference paper at ICLR 2023
Here (E(C is the covariance of gradient noise due mini-batch sampling, with y (N-1)6 and
NNb
E(C) = N (x,)), (x)()))) @ <xx,) HCH.
(5)
SPECTRALLY EXPRESSIBLE APPROXIMATION
An important property of non-stochastic GD is that it can be "solved" by diagonalizing the Hessian H. Namely, typical quantities of interest in the context of optimization have the form [r] with some function Φ: e.g,, o(z) = 2 for the loss, Φ(z) 6xx) = z² for the squared gradient norm. Given a Hessian H. by Riesz-Markov theorem we can associate to a positive semi-definite C the respective spectral measure pc(dA):
for the parameter displacement,
TT((H))) - - (())P((())) VΦ E C(spec
(6)
We will assume that H has eigenbasis Ur with eigenvalues AE; in this case PC = E CRROAE where CEa (us: Cug). For vanilla GD, the spectral measure of the second moment matrix C₁ at iteration is given by Pc, (dA) př PA) (dA), where pc(x)) is polynomial fully determined by the learning algorithm (Fischer 1996). Thus, the information encoded in the initial spectral measure PCs is, in principle, sufficient compute main characteristics of the optimization trajectories.
This conclusion, however, does not hold for SGD, since the first component of the noise term ( uses finer, non-spectral details of the problem. For a problem with fixed spectrum d and measure Pc, we show (Sec. D.I that the spectral components Σεa = (ug, Σug) the noise term (5) can take any value the wide interval [Tr[HC - Cx.),, (N - 1))] so that non-spectral details can make strong impact on the optimization trajectory.
Yet, in many theoretical and practical cases (see below), we find that non-spectral details are not significant, and components Σ& can be reduced to spectral expression. Specifically, for a certain choice of parameters T1, 72 E R, TI 2 T2, the noise measure pr is determined by pc via
(E(d)) - (f (xxc((x))
())o(((())))) PH - T2X2 pddA), PH = Σ AROAN-
We refer to this as the Spectrally Expressible (SE) approximation. Under this approximation we again can, in principle, fully reconstruct the trajectory of observables Cr, from the initial Cg. There are multiple reasons justifying this approximation.
1) There are scenarios where Eq. (7) holds exactly. First, we prove this for innslation-invarian models, with T = 1, = 1 (see Sec. D.5):
Proposition 2. Consider problem with regular grid dataser D = {xi (2mi 2 )) E (Z./N x (Z./aZZ) on the d-dimensional torus Td (Re/2 and with translation invariant kernel K(x,x)) = K(x - x'). Then 77 holds exactly with 1, = 1.
2) Also, SE approximation with T2 = generally results if pointwise losses L(x) are statisti- cally independent feature eigencomponents v²2 = (Uk, (xx))² w.r.t. x ~ D (Sec. D.2).
3) On the other hand, Bordelon & Pehlevan (2021) show that when the features (x) are Gaussian w.r.t. x ~ D, the SE approximation is exact with T = 1, 72 -1
4) Eq. can be derived by considering a generalization of the map (5) and requiring it to be "spectrally expressible". Namely, replace summation over in (5) by a general linear combination with coefficients πt =(iiiggisi)) only assuming the symmetries Rapigisis Righissi = R
N
Σ CC(C)
Σ i1.44
Raiiiis (xb(x.)) (ob(x;))))(x,, (xx,)) ) HCH. (8)
Assuming the coefficients R. to be fixed and independent of V H. C. prove
4

---

Published as a conference paper at ICLR 2023
Theorem 1. For SGD dynamics (4) with noise term (). the following statements are equiv- alent: (1) Spectral measures pC, occurring during SGD are uniquely determined by the initial spectral measure PCo and by the eigenvalues X& of Hessian H; (2) Rizézigis
Ti Ti Sighs ( + b Otzis) for some T1, T2, and Eq. 7 holds; (3) The trajectory {Ct] is in- variant under transformations = -U of the feature matrix by any orthogonal matrix U.
5) Finally, we observe empirically that SE approximation works well on realistic problems such as MNIST and CIFARIO for predicting loss evolution (Figs. [ 9 10 and Sec. [.4), determination of convergence conditions (Fig. 2 and approximation original noise term (5) in terms of trace norm (Sec. F.4
In the remainder we focus on analytical treatment of SGD dynamics under SE approximation 7 Note that parameters 7, r enter evolution equations in combinations 772,772- This allows to lighten the notation in the sequel by setting 71 = 1 and T2 = T without restricting generality.
4 REDUCTION TO GENERATING FUNCTIONS
Gas of rank-1 operators. Iterations 4 are linear and can be compactly written as Mc+1 F,M, with a linear operator Fr. Under SE approximation (7). for all four blocks of M, we only need to consider the diagonal components the eigenbasis of H. e.g. CKR. = (ug, C,uz). Then
I(T) = HCT (((0 FT ...
Here (00) = (6: (1 C1. 9) C22 ()...) E R & are (2 x 2)- matrix-valued sequences indexed by the cigenvalues AR. and (+.-) +) is the natural scalar product in & 2= The evolution operator Ft can be represented as the sum F; = Qt + A, of a rank-one operator Qe (noise term) and an operator A, (main term) that acts independently at each eigenvalue A as a 4 x 4 matrix:
Q: mafA(1 (
(10)
A = (Ara), A = = ('-ai Be B, -a
(11)
where A (A1, 22,,..) .T. Let us expand the loss by the binomial formula, with m terms Qe chosen at positions t,...., im and T m terms A, at the remaining positions:
1 Σ
Σ
L(T) 2
UT+1A UnmAmm- Una Vas
(12)
m=0Octi Cct <T-1
where
U,, ((20), 4 A A. 1
(13)
V, ( A A1 Co 0))
(14)
Expansion 122 has suggestive interpretation as partition function of a gas of rank-1 operators interacting with nearest neighbors (via U.,,) and the origin (via V,). Interactions via U,., contain the factor so that their strength depends on the sampling noise. We expect the model to be in different phases depending on the amount of noise: in the "signal-dominated" regime L(T) is primarily determined by V, at large T. while in the "noise-dominated" regime L(T)) is primarily determined by Ur... We show later that such a phase transition indeed occurs for non-strongly convex problems.
Generating functions. Expansion
12 allows to compute T.(T) iteratively:
L(T) = 2
Vi+++ + Σ
UT+1 L(tm - 1)..
(15)
bm==1
For constant learning rates and momenta 0t = a. 0r 3r B we have F. = F. the value Ur becomes translation invariant, U.,- U.-.. Then, as we show in Sec. E.I the loss can be conveniently described by generating functions (or equivalently, Laplace transtorm):
V(2)/2
ī(x)= = Σ 1(1)2² - (2)
(16)
2=0
5

---

Published as a conference paper at ICLR 2023
where Ü and V are the "noise" and "signal" generating functions:
Ü(z) = Σ U4+12' = = (() 8), --A))
(17)
1=
= ya²
2
Σ
1B 1)
(18)
x
V(2) = Σ
Vi+ = =((0)) (1 (1-22)-1((0
(19)
32 1) - 20.3822² 20,322
Σ
(20)
+ +2,3Az + 20 - 63 + 32 -3= 3z² - -z 2 +
In the remainder we derive various properties of the loss evolution by analyzing these formulas.
5 STABILITY ANALYSIS
Starting from this section we assume non-strongly-convex scenario with an infinite-dimensional and compact H, i.e., λ + O. Also, we assume that 0 < T < 1 in SE approximations: this simplifies some statements and fits the experiment for practical problems (see Figure D and Appendix D.
Conditions of loss convergence (stability). First note that as A - 0 the "main" components Ax of the evolution operator F appearing in Eq. (IT) have their action in trial matrix M converge to Co 8 )M 1 (112)) )², and from (33 we get U² = const Σm= A2 + o(1),, for each t. This means that if >x 2 co then each U, = xo and loss diverges already at first step (T-(t = 1) o0) we call this effect immediate divergence. On the other hand, assuming Σx X < oc, loss stability can be related to the first positive singularity of the loss generating function T(z) given by (16) Indeed,
let rr. be the convergence radius of power series (16). Since L(t) Z 0, L(z) must be monotone increasing on [0, rz) and have a singularity at TL- rf < 1, then L(l) A 0 as l + xe, while if fr > 1, then L(t) - 0 exponentially fast. Thus, large-t loss stability can be characterized by the condition rz 2 1, which can be further related to the generating function (=).
Proposition 3. Let 3 E (-1,1), a < 2 ([8+11 1)/Amas and y E [0, 1]. Then r2 < <1iff 1.
This result is proved by showing that (=), (2x) have no singularities for E (0, 1) and that =ū())
is monotone increasing on (0, 1). The only source of singularity in L(=) is then the denominator in
(T).. Note that conditions 3 E (-1,1), 1), 0 < a < 2B3 + 1)/Amax exactly specify the convergence
region for the non-stochastic problem with 7 0 (see (Roy & Shynk] 1990 Tugay & Tanik 1989)
and Appendix H). This is also a necessary condition in the presence of sampling noise since the
second term in 4 is positive semi-definite matrix.
To get a better understanding of convergence condition Ú(1) we use Eq. 18) at z and find
θ(1== Σ (x
d&
ary
Σ &
(1
(21)
1-8 71+33)
+
21-8)
2(1 B)
47
This shows that Ae co iff Ū11 0c > 1 (regardless of values of a, 3, 7), in which case TL < 1 and L()) diverges exponentially as t + oc. We refer to this scenario as eventual divergence.
Next, consider effective learning rute aetf - accumulated effect of previous gradient descent iterations (Tugay & Tanik 1989 Yuan al. 2016). We show that condition (1) is closely related simple bound Offf terms of critical regularization Aerit. which is defined only through problem's 's spectrum A and SE parameter r.
+ a: + a,82 +... which reflects the
Proposition 4. Retain assumptions of Prop. 3 and define devit as the unique positive solution to
Ee FAx Aemi = 1. Then, convergence condition (1) < requires for arst to obey
Oreff Aecno
(22)

---

Published as a conference paper at ICLR 2023
Moreover, for fixed B, the convergence condition (1) < 1 is equivalent the condition aef Oef y), where the critical values ael (3,7) are upper bounded by 7dont (by Eq. (22) and
MAurt / ( off (B,7) 15 Amax 1-8 3-1 0(1-3) B).
(23)
Acrit 7(++B)
The bound (2 is much simpler to than the exact condition (()) < 1, while (33 shows that
22 is tight for practically common case of close to Both conditions Ú(1) < 1 and (22 agree
very well with experiments on MNIST and symthetic data (Fig. 2 Importantly, Eq. (22) implies
that, in contrast to the non-stochastic (y 0) GD, effective learning rate in SGD must be bounded to prevent divergence. In non-strongly convex problems, accelerated convergence of non-stochastic GD achieved by gradually adjusting / at a fixed during optimization (Nemirovskiy & Polyak 1984). We see that this mechanism is not applicable to mini-batch SGD.
Exponential divergence. Our stability analysis above shows that loss diverges at large / when ((1) > 1. This divergence is primarily characterized by the convergence radius rz determined from the equation (rL) = 1 according to (16). Indeed, assuming an asymptotically exponential ansatz. T(t) Cet/ta we immediately get the divergence time scale lavv -1 logrr. A closer inspection of [(2) near ru also allows to determine the constant C' (see appendix E.4b
V(rL)
L(t) 2(1 + (rL)())) r
(24)
Note that rz can very close to 1, which is naturally observed in eventual divergent phase by taking sufficiently small and/or 7, In this case the characteristic divergence time faiv is large and we expect divergent behavior to occur only for 2 fotowup, while for < folowup the loss converges roughly with the rate of noiseless model (2) = 0. As fooowp should have meaning of time moment where loss starts to significantly deviate from noiseless trajectory we define it by equating divergent 24 and noiseless losses. We confirm these effects experimentally in Fig, 3 (right).
SOLUTIONS FOR POWER-LAW SPECTRAL DISTRIBUTIONS
To get a more detailed picture of the loss evolution let us assume now that the eigenvalues A and the second moments Cak.o of initial (L = 0) approximation error are subject to large-k power laws
A = A o(1), Σ 12k
ACC. = Kk (1++(()))
(25)
with some exponents U, x > (also denote Ç = x/v) and coefficients A, K. Such (or similar) power laws are empirically observed many high-dimensional problems, can be derived theoretically, and are often assumed for theoretical optimization guarantees (see many references in App. AA). If v S 2, then En X2 = do and we are in the "immediate divergence" regime mentioned earlier, and if < S. 1, then E. A do and we are in "eventual divergence" regime. Away from these two divergent regimes, we precisely characterize the late time loss asymptotics:
Theorem 2. Assume spectral conditions (25) with > 1 and that parameters a, B, are as in
Proposition 3 and such that convergence condition (1) 1 is satisfied. Then
T
KT(C+1) 20A cg2-C, 2 2- 6-<>/ 1/v, 2(-0(())2 ()
Σ = ((+))=(++0(1))
(26)
V 1/v) 20A
( 1/0p2/1 2 ( <<<<< 1/v..
8(1-U(1))
The idea of the proof is to observe that the generating function for the sequence tL(t) is Ee (L(()=² = L()), where
d
[(x)
(27)
dz
2(1 2(1))
2(1 2=)))
We argue that ry. = 1 and apply the Tauberian theorem, relating the asymptotics of the cumulative
sums Etz (T())) the singularity of the generating function at rt 1. One can then show

---

Published as a conference paper at ICLR 2023
MNIST/ SE 1)
MNIST/ linerrized
v=1.5 x=3.07 SE
ve 3.0 linearized
min
min
8.59
0.30
-31
Nomeetum (83
4.48
0.0
0.00
0.00
0(3)
0111
0111
2
22+01
-2.9
20-0
20-
2-
-635
Learning lol
Leaming rate (o)
Leaming
Learning
(c)
Figure 2:
L(t = 10 for different learning rates a and momenta 3, in the SE(T= = 1) and linearized
regimes, for MNIST and synthetic data given by Eq. (25) with 1.5, x = 3. for b 10. Black lines show convergence boundaries from Proposition (3) and Eqs. (21) and (22). MNIST experiment is in the signal signal-domintted phase and synthetic in the dosse-domintted. The stars show arg min 3 L.
using (88) 19 that (2) diverges as Ú1 - E) x xE= and (2) diverges as (' - x c
for S < Depending on which of the two terms in Eq. (27) has stronger singularity, we observe two qualitatively different phases (previously known in SGD without momenta (Varre et al., 202T)):
the "signal-dominated" phase at Ç - < - 1/v and the "noise-dominated" phase at 2 See Figure (left) for the resulting full phase diagram.
In the sequel we assume that not only cumulative sums Σ tL(t), but also individual terms tL())
obey power-law asymptotics. Then L(t) Lapprox (t) Ct 5, where, by Eq. 26
Ceignalt < KT(<-1) (204)
2(1 (1)) 55
<< 1/v,
(28a)
Lupprox (t)
11 (2 1/v) 2nd 1/2 <> 1/v.
(28b)
(1))
Note that the asymptotics (28) depend on SGD hyperparameters (a,, 3, only through the constants Csignal (a, B, 7) and (a, 7), which makes them essential for describing finer details of the loss trajectories in both phases. In the following sections we give few examples of such analysis.
Transition between phases. Asymptotic 28) allows us to go further in understanding the "noise- dominated" phase > 2 1/v. Using representation 27 and linearity of Laplace transform, we write loss asymptotic as T.(() Rs Csignalt
terms given by (28a). (2bb) As Cnoise vanishes in the noiseless limit y - O. we expect that for sufficiently small the signal term here dominates the noise term up to time scale trrana given by
C 2+ with
Osignal
1/(S-2+1/4
1L-- 1/v4C 4UKTIC+)(--Un 2/(5- 200h
btrans
(29)
Cuoise
7T(2-1/v)V1)
This conclusion is fully confirmed by our experiments with simulated data, see Fig. center). Note that the time scale berans can vary widely depending on ax, 3, 7; in particular, 4rans might not be reached at all in realistic optimization time the noise y and effective learning rate I-3 are small.
A similar reasoning can be used to estimate the time (Blowup of transition from early convergence to subsequent divergence (see Fig. 3 (right)). In Sec. E.4 we consider the scenario with 8 0, 7 т = 1 and power laws 255 in which 2 < v < 1 (eventual divergence) and Ç < 1. We show that at small a we have tbtowup RC autax s3 a, [22 1/v) - 1)) w/(x-1) (2aA) ((v-), where a, a. (v,, () is the solution of the equation 1/v-1 r(1-c) a.
ca.
Hy Hyperparrmeeer optimization. Observe that the optimal learning ratse a and momenta / in Fig. differ significantly between signal-dominated ((a), (b)) and noise-dominated ((c). (d)) phases. To explain this observation, we analyze late time loss asymptotic (28) at a fixed batch size b, and show that two phases indeed have distinct patterns of optimal a. S.
"Signal- "ominated" phase. First, note that near the divergence boundary Ú(1) the loss is high due to denominator 1 - (1) in 28a). This suggests that at optimal (x, 3 the difference Ū(1) - is of order 1, which allows us to neglect dependence of loss on (1) as long as convergence condition (1) 1 holds. Then, loss becomes T.(l) (oent)-5 5 and it is beneficial to increase Ocft- However,
8

---

Published as a conference paper at ICLR 2023
1.5.
x 0.25,
0.75.
675.6-
1.
10
Immediate divergence Σ, Al 00
10"2
10a
Eventual divergence Σ, Ai co, Σ, A
10:3
6-1
100
#cHMInnataaa Signal dominated
10
Noise dominated (it) x pla-
10
Gond e
10-1
L(0 x 1
10-14 10
C
10a
10a
10a
Figure 3: Left: Phase diagram for L(t) under power-law spectral conditions 25 The MNIST
and CIFARIO 10 stars are placed according to power-fits Figs. and (top). Center: Transition from signal-dominated to noise-dominated regime during training. Asymptotic power-law lines are calculated from (28a) and (28b). Right: Divergent T.(t) in the phase & < v < 1: experimental trajectories (solid) with marked points at positions t tdiv (squares) and t dblowup (circles), theoretical late-time asymptotics (24) (dash-dotted) and noiseless early time trajectorics (dashed).
Figure 4: Theoretical and experimental values of optimal mo- menta |3 in the range (v, x) E [1,3] x [0, 6] of power- law models (55) The blue/red background shows the val-
ues =\t given by 30 and characterizing the theoretical sign of 0g SaLapprox [8) =
in the noisy regime. The four diamonds represent the pairs (v, x) for which we experi- mentally find optimal Sopt- The diamond color shows the value of Bopt- In agreement with Prop. experimental Sops. > 0 (< 0) when < 0 (> 0). The dotted yellow line separates the signal- dominated (helow) and the noise-dominated (above) regimes.
x
-03
Ocft has an upper bound 23 which becomes tight for sufficiently small 3 (see
23 Thus, in
agreement with Fig. 2 (left), the optimal a, B correspond to small B and Oeff 2 "acrit"
phase, we see from 28b that (1) can be In contrast to signal-dominated phase, Ocf a/(1 [1 - 33 appears with the positive power suggesting that high momenta 3 ~ 1 are non-optimal. At the same time, it can be shown that V(1) ~ a- - at small a, so it is neither favorable to use small a. Thus, accordance with Fig. 2 (right), the optimal a, B
"Noise-dominated" phase. Similarly to signal-dominan
neglected away from convergence boundary. Then, the loss is L(t) (1)
///1//=2, (/1)a//////-=
should be located in a general position inside the convergence region (1) <1)<1.
Additionally, for the problem of optimizing batch-size b at fixed computational budget bl, in Sec. E.3 we derive the previously observed linear scalability of SGD w.r.t. batch size (Ma et al. 2018).
Positive negative momenta. It is known that in some noisy settings SGD does not benefit from using momentum (Polyak] 1987 Kidambi et al. 2018 Paquette & Paquette| 2021). Empirically, we find that in the signal-doming dommnnedd regime momentum helps significantly, while in the noise-dominated regime its effect is weaker, and it can even be preferable to set 3 < (see Fig. 2 Our asymptotic formula 28 with explicit constants allows
confirm this theoretically and, moreover, in the case = 1 give simple explicit condition under which 3 < 0 are preferable. Specifically, we show that if we fix a at the optimal value Copt at 5 O, then we can further decrease the loss by increasing or decreasing B, depending on the regime and a spectral characteristic 3gt defined below:
Proposition 5 (short version). Assume asymptotic power laws 25 with some exponents V. > 0,5 = and and coefficients A, K. Let O'opt = arg ming Lapprox (1)T Then: (signal-dominated case) Let < 2 1/v and dominated case) Let Ç > - 1/v and = = 1, then the sign of Ög Lapprox (tooproxx()()0_ccc))) (+) equals the sign of the expression
< 1, < 7 1, then dg (t)
0. II. (noise-
= v H Tr HC, - - (v-
(30)
In Fig. we experimentally validate the E criterion for several pairs (v, x). See details in Sec. E.6

---

Published as a conference paper at ICLR 2023
7 DISCUSSION
We can summarize our approach to the theoretical analysis of mini-batch SGD as consisting of 3 key stages: 1) Perform SE approximation; 2) Use it to solve the SGD dynamics and obtain initial spectral expressions (in our case Eqs. (18) 20) combined with (1)); 3) Apply these spectral expressions study optimization stability, phase transitions, optimal parameters, etc. We have demonstrated this approach to produce various new and nontrivial theoretical results, which are also in good quantitative agreement with experiment on realistic problems such MNIST and CIFARIO learned by MLP/ResNee/MobieNNe.
In particular, we have shown that, in contrast full-batch GD, SGD has an effective learning rate limit determined by the spectral properties of the problem and batch size. Also, we have shown that in the "signal-dominare dominnteed" phase (including, e.g., the MNIST and CIFARIO tasks) momentum is always beneficial, but in the "noise-dominated" phase a momentum B < 0 can be preferable.
One natural future research direction to investigate which properties of data and features (b(x) are responsible for the accuracy of SE approximation, and how to systematically choose SE, parameters T1, T2- Second, we were able to solve SE approximation only for the constant learning rate and mo- mentum, and a more general solution might open up way to better understanding of accelerated SGD. Third, one can ask how the picture changes for spectral distributions other than power-law.
ACKNOWLEDGMENTS
We thank the anonymous reviewers for several useful comments that allowed us to improve the paper. Most numerical experiments were performed on the computational cluster Zhores (Zacharov et al. 2019). Research was supported by Russian Science Foundation, grant 21-11-00373.
8 BIBLIOGRAPHY
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. arXiv preprint vrXiv:2111.00034 2021.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2102.06701, 2021.
Ronen Basri, Meirav Galun, Amnon Geifman, David Jacobs, Yoni Kasten, and Shira Kritchman. Frequency bias in neural networks for input of non-uniform density. In International Conference on Machine Learning. pp. 685 694. PMLR, 2020.
Saugata Basu. Algorithms in real algebraic geometry: survey. arXiv preprint arXiv:1409, 1534,, 2014.
Raphaël Berthier, Francis Bach, and Pierre Gaillard. Tight nonparametric convergence rates for stochastic gradient descent under the noiseless linear model. arXiv preprint 2rXx:2000608822,, 2020.
Alberto Bietti. Approximation and learning with deep convolutional models: kernel perspective. arXiv preprini 1rXiv:2102.10033 2021.
Blake Bordelon and Cengiz Pehlevan. Learning curves for sgd on structured features. arXiv preprint arXiv:2106.0277 02711. 2021.
Léon Bottou and Olivier Bousquet. The tradeoff's of large scale learning. J. Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Neural Information Processing Systems, vol- ume 20. Curran Associates, Inc., 2007. URL httttp::://rccceednngn www.oiss.suuuiusuooummmm
Benjamin Bowman and Guido Montufar. Spectral bias outside the training set for deep networks in the kemee regime. arXiv preprint arXiv:2206.02927, 2022.
Helmut Brakhage. On iol-po problems and the method of conjugate gradients. In Inverse and ill-posed Problems, pp. 165-175. Elsevier, 1987.
10

---

Published as a conference paper at ICLR 2023
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model align- ment explain generalization in kernel regression and infinitely wide neural networks. Na- ture Communicationss 12(1), may 2021. doi: 10. 1o..11K641147--21-21000000
https::
7/doi org/10. s 4 2112121-113--
Andrea Caponnetto and Emesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics, 7(3):331-368, 2007.
Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Generalization error rates in kernel regression: The crossover from the noiseless to noisy regime. Advances in Neural Information Processing Systems, 34, 2021.
Andreas Dolzmann and Thomas Sturm. Redlog: Computer algebra meets computer logic. Acm
Sigsam Bulletin, 31(2):2 9, 1997.
Xialiang Dou and Tengyuan Liang. Training neural networks learning data-adaptive ker- nels: Provable representation and approximation benefits. Journal of the American Statisti- cal Association, 1166535): apr 2020. doi: 10.10800162145922200..0222
URL
William Feller. An infroduction to probability theory and its applications. Vol. II. Second edition. John Wiley Sons Inc., New York, 1971.
B. Fischer. Polynomial based iteration methods for symmetric linear systems. 1996.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: empirical study of loss landscape geometry and the time evolution of the neural tangent kernel, 2020. URL https::////aatt
org/abs/2010.1 15110
Sergei Farshatovich Gilyazov and Nataliă L'vovna Golndman. Regularization of ill-posed problems by iteration methods, volume 499. Springer Science & Business Media, 2013.
Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc Mézard, and Lenka Zde- borová. The gaussian equivalence of generative models for learning with shallow neural net- works. 2020a. doi: 1..45550AARXI..200..4444 URR/nttps.//aaiiv.oor//// wbw..//r///b//.......///
14709
Sebastian Goldt, Mare Mé zard, Florent Krzakala, and Lenka Zdeborová. Modeling the influence of data structure learning in neural networks: The hidden manifold model. Physical Review X, 10(4), dec 2020b. doi: 10.1 10331/hysseexx100001104 URL :ttp:://ee..rrr/11.1110000. 11033
2Fphysrevx. 10 041044
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. http: 7/ www..eeepleaannnngooo..o.o org
Martin Hanke. Accelerated landweber iterations for the solution of ill-posed equations. Numerische mathematik, 60(1):341-773, 1991.
Martin Hanke. Asymptotics of orthogonal polynomials and the numerical solution of ill-posed problems. Numerical Algorithuss, 11(1):203-213, 1996.
Nicholas JA Harvey, Christopher Liaw, Yaniv Plan, and Sikander Randhawa. Tight analyses for non- smooth stochastic gradient descent. In Conference on Learning Theory, pp. 1579-1613, PMLR, 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. arXiv preprint 17X6002806.07572, 2018.
Prateek Jain, Dheeraj Nagaraj, and Praneeth Netrapalli. Making the last iterate of sgd information theoretically optimal. In Conference on Learning Theory, pp. 1752-1755. PMLR, 2019.
Harold Jeffreys. On isotropic tensors. Mathematical Proceedings of the Cambridge Philosophical Society, 73(1):173-17) 1973. doi: 10.1017/503050041000111

---

Published as a conference paper at ICLR 2023
Hui Jin, Pradeep Kr Banerjee, and Guido Montúfar. Learning curves for gaussian process regression with power-law priors and targets. arXiv preprint arXiv:2110.122311 2021.
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the insufficiency of ex- isting momentum schemes for stochastic optimization. In 2018 Information Theory and Applica- lions Workshop (ITA), pp. 1-9. IEEE, 2018.
Jack Kiefer and Jacob Wolfowitz. Stochastic estimation of the maximum of regression function. The Annals of Mathematical Statistics, pp. 462-466, 1952.
Hermann König. Eigenvalue distribution of compact operators, volume 16. Birkhäuser, 2013.
Dmitry Kopitkov and Vadim Indelman. Neural spectrum alignment: Empirical study. In Interna- tional Conference on Artificial Neural Networks, pp. 168-179. Springer, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 444, 2015.
Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent, 2019.
Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems, 11115555-55272 2020.
Bruno Loureiro, Cédric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mézard, and Lenka Zdeborová. Learning curves of generic features maps for realistic datasets with teacher- student model. 2021. doi: 1048550/ARXI.21102.00222 URL
2102.08122
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effec- tiveness of sgd in modern over-parametrized learning. In International Conference on Machine Learning, pp. 3325-3334 PMLR, 2018.
Aaron Meurer, Christopher P. Smith, Mateusz Paprocki, Ondiej Čertik, Sergey B. Kirpichev, Matthew Rocklin, AMiT Kumar, Sergiu Ivanov, Jason K. Moore, Sartaj Singh, Thilina Rath- nayake, Sean Vig, Brian E. Granger, Richard P. Muller, Francesco Bonazzi, Harsh Gupta, Shivam Vats, Fredrik Johansson, Fabian Pedregosa, Matthew J. Curry, Andy R. Terrel, Stěpán Roučka, Ashutosh Saboo, Isuru Fernando, Sumith Kulal, Robert Cimrman, and Anthony Scopatz. Sympy: symbolic computing in python. PeerJ Computer Science, 103, January 2017. ISSN 2376-5992. Scienneeerscince... doi: 10.7717/peerj-es. 103. / wttp://///....777perrrrr...
771 7117111/e/rr.ss. 103
Eric Moulines and Francis Bach. Non-asymptotic analysis of stochastic approximation algorithms for machine learning. Advances in neural information processing systems, 24, 2011.
Arkadi S Nemirovskiy and Boris T Polyak. Iterative methods for solving linear ill-posed problems under precise information. Eng, Cyber. (4):50-56, 1984.
Yurii Evgen' evich Nesterov. A method of solving convex programming problem with convergence rate o(k2) In Doklady Akademii Nauk, volume 269, pp. 543-547. Russian Academy of Sciences, 1983.
Courtney Paquette and Elliot Paquette. Dynamics of stochastic momentum methods on large-scale, quadratic models. Advances in Neural Information Processing Systems, 34:9229-999 92400 2021.
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr compu- fational mathematics and mathematical physics, 4(5):1-17, 1964.
B.T. Polyak. Introduction Optimization. Optimization Software, New York, 1987.
12

---

Published as a conference paper at ICLR 2023
John Proakis. Channel identification for high speed digital communications. JEEE Transactions Automatic Control, 19(6)916-99) 922, 1974.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statistics, pp. 400-407, 1951.
Sumit Roy and John J Shynk. Analysis of the momentum Ims algorithm. IEEE transactions on acoustics, speech, and signal processing. 38(12):008-2099) 1990.
David E Rumelhart. Geoffrey E Hinton, and Ronald J Williams. Learning representations by back- propagating errors. nature, 523(6088):533-536 1986.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJV), (1(3):211-2555 2015. doi: 10. 10.1007/s1 12030015-8--111
Christopher Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arXiv: 1.03600, 2018
Ohad Shamir and Tong Zhang. Stochastic gradient descent for non-smooth optimization: Conver- gence results and optimal averaging schemes. In Sanjoy Dasgupta and David McAllester (eds.), Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceed- ings of Machine Learming Research, pp. 71-79, Atlanta, Georgia. USA. 17-19 Jun 2013. PMLR. URL http:tttp/ppoceeeeeeeee
mir r. prwssppsspee. wwwpsseeww///saaaiee
html
Ingo Steinwart, Don R Hush, Clint Scovel, et al. Optimal rates for regularized least squares regres- sion. In COLT, pp. 79-93, 2009.
Пya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning. pp- 1147. PMLR, 2013.
Mchmet Ali Tugay and Yalein Tanik. Properties of the momentum Ims algorithm. Signal Processing, 18(2):117-127, 1989.
Aditya Varre and Nicolas Flammarion. Accelerated sgd for con---toonlyccoweee least squares, arXiv preprint arXiv:2203.01744 2022.
Aditya Varre, Loucas Pillaud- Pillaud-Vivien, and Nicolas Flammarion. Last iterate convergence of sgd for least-squares in the interpolation regime. arXiv preprint arXiv:2102.03/88. 2021.
Maksim Velikanov and Dmitry Yarotsky. Explicit loss asymptotics in the gradient descent training of neural networks. Advances in Neural Information Processing Systems, 34, 2021.
Maksim Velikanov and Dmitry Yarotsky. Tight convergence rate bounds for optimization under power law spectral conditions, 2022.
00992
Alexander Wei, Wei Hu, and Jacob Steinhardt. More than toy: Random matrix models predict how worldworld neural representations generalize, 2022. URL https:::/axxv..oro////// 2203.06177
Hermann Weyl. The classical groups 1939.
Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the noisy gradient descent that generalizes as sgd. In International Conference Machine Learning, pp. 10367-10376. PMLR, 2020.
Kun Yuan, Bicheng Ying, and Ali Sayed. On the influence of momentum acceleration on online learning. The Journal of Machine Learning Research, 17(1 1716)6602-6667, 2016.
13

---

Published as a conference paper at ICLR 2023
Igor Zacharov, Rinat Arslanov, Maksim Gunin, Daniil Stefonishin, Andrey Bykov, Sergey Pavlov, Oleg Panarin, Anton Maliutin, Sergey Rykovanov, and Maxim Fedorov. "zhores" petaaflpss supercomputer for data-driven modeling, machine learning and artificial intelligence installed in skolkovo institute of science and technology. Open Engineering, 9()::122-520 2019.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint arXiv:1803.00190 2018.
Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Benign over- fitting constant-stepsize sgd for linear regression. arXiv preprint arXiv:2103, 12692, 2021.
14

---

Published as a conference paper at ICLR 2023
CONTENTS
Introduction
The setting and SGD dynamics
Spectrally Expressible approximation
Reduction generating functions
Stability analysis
Solutions for power-law spectral distributions
Discussion
10
8 Bibliography
10
Related work
16
Details of the setting
17
B.
I Basics
17
B.2 Eigendecompooitions and other details
18
B.3 Output and parameter spaces
20
Dynamics of second moments
20
SE approximation
23
D.I Presence of non-spectral details in the general case.
................................
23
D.2 SE approximation for uncorrelated losses and features
................................
24
D.3 Dynamics of spectral measures
25
D.4
Theorem
26
D.5
Translation invariant models
29
Generating functions and their applications
30
Reduction to generating functions
30
Convergence conditions
31
Loss asymptotic under power-
33
E,4
Exponential divergence of the loss and convergence convergence-diivegenn diwnoeeeddvveegege transition
37
Budget optimization)
40
E,6
Positive vs. negative momenta
41
Experiments
44
F1
Training regimes
44
15

---

Published as a conference paper at ICLR 2023
F2 Datasets
45
Experimental details
45
Validity of spectrally expressible approximation
46
F55
Additional experiments
48
Mini-batch SGD vs. SGD with additive noise
49
H Noiseless convergence condition
51
A
RELATED WORK
SGD.
Stochastic Gradient Descent (SGD) was introduced in Robbins & Monro (951), and its
various versions have been extensively studied since then, see e.g. Kiefer & Wolfowitz (1952);
Rumelhart et al. (1986); Bottou & Bousquet (2007); [Shamir& & Zhang (2013); Jain et al.
2011
Harvey et al. (2019;; Moulines & Bach] (2011). The version ( considered
the present paper
motivated by applications to training neural networks and is widely used (LeCun et al., 2015
Goodfellow al. 2016 Sutskever al. 2013). Its key feature the presence of momentum. which
both theoretically known to improve convergence and widespread in practice (Polyak 1964 1987
Nesterov 1983 Sutskever et al. 2013 Shallue et al. 2018).
Mini-batch SGD and spectrally expressible representation. In the present paper consider the mini-batch SGD, meaning that the stochasticity in SGD is specifically associated with the ran- dom choice of different batches used to approximate the full loss function at each iteration (not to be confused with SGD modeled as gradient descent with additive noise as e.g. in the gradient Langevin dynamics; e.g. Wu et al. (2020); Zhu al. (2018) for some discussion and comparison, well as our Section G for an illustration in our case of quadratic problems).
We study this SGD in terms of how the second moments of the errors evolve with time. Our Propo- sition [ generalizes existing evolution equations (Varre et al.) 2021) to SGD with momentum. A serious difficulty in the study of this evolution is the presence higher order mixing terms coupling different spectral components (the first term on the r.h.s. in our Eq. (5). In Varre et al. (2021). it was shown that the effects of these terms can be controlled by suitable inequalities, allowing to establish upper bounds on error rates. In contrast, our SE approximation allows us to go beyond upper bounds and, for example, obtain explicit coefficients the asymptotics of large-/ loss evolution without Gaussian assumption (see Theorem 2).
Our approach is different and close to the approach of recent paper Bordelon & Pehlevan (2021) where the authors considered a dynamic equation equivalent to our SE approximation with parameter
T = 1 but without momentum. Apart from considering a wider family of dynamic equations, our major advancement compared to Bordelon & Pehlevan (202I) is that manage to "solve" these dynamic equations using generating functions, and consequently use this solution to obtain a various results concerning stability conditions (Sec. (5) and various SGD characteristics under power-law spectral distributions (Sec.
Related assumptions of Gaussian features have also been used in recent series of works Goldt ee al. (2020a): Loureiro et al. (2021): Goldt et al. (2020b) to analytically predict performance of linear models trained with full- mini-batch SGD.
Spectral power laws. In Section 6 we adopt very convenient power law spectral assumptions (25) on the asymptotic behavior of eigenvalues and target expansion coefficients. Such power laws are empirically observed in many high-dimensional problems (Cui et al.] 2021 Bahri et al.) 2021
Lee et al. 2020 Canatar et al., 22211 Kopitkov & Indelman 2020 Dou Liang 2020 Atanasov
al. 2021 Bordelon & Pehlevan, [2021] Basri al. 2020 Bietti 2021), can be derived theoret-
ically (Basri et al. 2020: Velikanov
Yarotsky] 2021). and closely related conditions are often
assumed for theoretical optimization guarantees Nemirovskiy & Polyak| 1984 Brakhage. 1987
Hanke| 1991| 1996 Gilyazov & Gol dman 2013) Varre et al.] [021]
[002:|Capomaa Caponnetto & De Vito 2007
Steinwart al. 2009 Berthier et al. 2020 Zou
202
Jin et al. 2021 Bowman & Montufar
16

---

Published as a conference paper at ICLR 2023
2022
An important aspect of our work is that, in contrast to most other works, allow the expo- nent characterizing initial second moments Cks.o to have values in the interval (0, 1). In this case, This effectively means that the target function (xx has finite euclidean norm, while the norm of the respective optimal weight vector w is infinite (i.e., the loss does not attain minimum value). This broader assumption allows us, in partieular, to correctly describe the loss evolution on MNIST for which RJ 0.25 (see Figures [T and 3): without assuming
E CR < while Ex CEA.O
||w*|l oe, the SGD proved to have error rate (t-=)) with & clearly contradicts the experiment (see Figure [ [left).
(Varre et al. 2021). which
Phase diagram. It is known that performance of SGID can be decomposed into "bias" and "vari- ance" terms, with the dominant term determining what we call
"noise- dominated" regime (Moulines & Bach 2011 Varre et al. 2021] Varre & Flammarion, [022). In particular, see Varre et al. (2021) for respective upper bounds for SGD without momentum. The
main novelty of our paper compared to these previous results is the analytic framework allowing to obtain large-t loss asymptoties with explicit constants for SGD with momentum (see Theorem 2 and use them to derive various quantitative conclusions related to model stability, phase transitions, and parameter optimization (see Section 6
Noisy GD with momentum. GD with momentum (Heavy Ball) (Polyak 1964) was extensively studied, but mostly in noiseless setting or for generic (non-sampling) kinds of noise; see e.g. the monograph Polyak (1987). See PPrakiis (1974); Roy & Shynk (1990) Tugay & Tanik (1989) for discussion of stability on quadratic problems. Polyak (1987) analyses various methods with noisy gradients for different kinds of noise and concludes that the fast convergence advantage of Heavy Ball and Conjugate Gradients (compared to GD) are not preserved under noise unless it decreases to O as the optimization trajectory approaches the solution. See also Kidambi et al. (2018): Paquette & Paquette (202T) for further comparisons of noisy GD with or without momentum. Our work refines this picture in the case of sampling noise. Our results in Section 6 suggest that including positive momentum in SGD always improves convergence in the signal-dominated phase, but generally not in the noise-dominated phase (where may even be beneficial use negative momentum).
B
DETAILS OF THE SETTING
B. BASICS
In our problem setting we always consider regression tasks of fitting target function (x(x using linear models f'x,, w) the form
f(x, w) (w, 2b(x))
(31)
where () stands for scalar product and (xx) are some (presumably non-linear) features of inputs x. The space X' containing inputs x not important in our context. Next, the model will be always trained to minimize a quadratic loss function over training dataset DD:
L(w) Ex ((w, (xx)) (x))²
(32)
Denote by B, set of b [BB| training samples x drawn independently and uniformly without replacement from D. Then single step of SGD with learning rate a and momentum parameter S takes the form
Wr+1
I (B)
B.I
(w:)
(w:)
(orh( B.) +(ah((B))
Vi+1
=
(a.H(B.)
&I)
+
(ar((B))
(33)
H(B,) b
1
Σ (xx;) @ (xx:)
(34)
x.EB.
h (B)) = b Σ (xx))x(xx))
(35)
x.EB,
Note that expression 33 of SGD step is equivalent to more compact expression 2 when there exist optimal parameters w such that (w*, (x)) = yn (x). As we have already noted, existence of
17

---

Published as a conference paper at ICLR 2023
optimal parameters with ||w"| oo may not be the case for models relevant to practical problems (e.g. MNIIT).. Thus we will consider (33) to be the basic form of SGD step in our work. However, the formal usage of (2) instead of (33 always leads to the same conclusions and therefore we focus on the case ||w*) < oo in the main paper. In subsection B.3 we consider formulation of SGD in the
output space, which allows treat cases |ww oo and w oo withing the same framework.
is convenient to distinguish several cases of our problem setting differing by finite/infinite dimen- sionality of parameters w, and finite/infinite size [D] N of the dataset used for training. In our work we take into account all these cases, as each of them is used in either experimental or theoretical part of the paper.
B.2 EIGENDECOMPOSITION AND OTHER DETAILS
Case I: d < co. This is the simplest case and we will often repeat the respective paragraph of section 2 First, recall that matrix elements d x N Jacobian matrix are
Vi = dy (x,)) Xj E D
(36)
We will assume without loss of generality that v has full rank r = min(d, (d, N) (the general case reduces this after suitable projection). Then by (32) the Hessian of the model is
N
H
T
T
bb (x,) 2b (x²)-
(37)
N
N
j=1
Another important matrix is the kernel K (xx, x') (x)) (xx) calculated on the training dataset D
N
K
T 4
Σ vi (x()))
(38)
N
31'=11
Next, consider the SVD decomposition of Jacobian V VNUAV with U = (ux) being the matrix of left eigenvectors, V ( being the matrix of right eigenvectors, and rectangular d X N diagonal matrix with singular values Auk V k = 1...r the diagonal. The Hessian and kernel matrices are symmetric and have eigendecompositions with shared spectrum of non-zero eigenvalues: λ₁, U2 for H and λ2. V2 for K.
Now let us proceed to the target function g"(x), As in this paper we focus on the training error 32). we may also restrict ourselves to the target vector y" of target function values at training points: of JØ (x;), x; D. The normalization by VN is done here allow for direct correspondence with the rank and therefore the target function is reachable since the equation 1 w T = always has VN solution. However, if d < N, there are many possible solutions of WT T = and we choose w" to be the one with the same projection on ker H as the initial model parameters We- Later, this choice will guarantee that during SGD dynamic (33) parameter deviation always lies in the image of the Hessian: Aw, w. w E Im H. Finally, if N > d, the target may by unreachable: y & Im T In this case we consider the decomposition of the output space R. = Im K @ ker K and the respective decomposition of target vector y* = yj yi. Then y is always reachable
Do case. If d N, we recall our assumption of having full
VN
yi E Im T and in the rest we simply consider reachable part of the target yj instead of the original target
Apart from general finite regression problems, the case d, N < do described above applies to lin- carization of practical neural networks (and other parametric models). Specifically, consider a model f(w, x) initialized at wo with initial predictions fo(x) = f(wo, x). The training of linearized model xww fo(x) + (w wo. Vwf(wo, can be mapped to out basic problem setting (3] with replacements (xx)
V V V (w., w - ww-- wo. y (x) 4 (x)) - fx().
Case 2: d 00, N oo. In this case we assume that model parameters w and features (x) belong to a Hilbert space H. and (.-)) denotes the scalar product in H. For the Hessian we have
H Σ bbxx) 2 (xy), x, E
(39)
18

---

Published as a conference paper at ICLR 2023
We assume that the features xxxx), x, E D are linearly independent (as is usually the case in practice), then dim Im H = N < co. As SGD updates occur in Im H. we again use space de- composition H. Im H + ker H and respective decomposition of parameters
wI + W and features & (xx) [I + (b()) As (b(x)) X2 E D and W do not affect the loss (32)
can restrict ourselves to Im H which bring us to finite dimensional case d = < oo fully described above.
The case d = 00, N < oo is primarily used for kernel regression problems defined by a kernel function K (xx,x' x') which can be, for example, of the "classical" kernels (e.g. Gaussian) or Neural Tangent Kernel (NTK) of some infinitely wide neural network. In this case H is a repro- ducing kernel Hilbert space (RKHS) of the kernel [xx, x'), and features {b(x) the respective mappings 3 X + H. with the property K (xx,x') = (xb(x), xxx)))
Case 3: d, N = oo. We are particularly interested in the case = N = oo because in practice N is quite large and, as previously discussed, the eigenvalues λ and target expansion coefficients often are distributed according to asymptotic power laws, which all suggests working in an infinite- dimensional setting. A standard approach to rigorously accommodate N oc, which we sketch now, is based on Mercer's 's theorem.
We describe the infinite (N = oc) training dataset by probability measure μ(x) so that the loss takes the form
L(w) = / ((w, bbxx) y² dx))) =
1 = xbxx)) - y (x)))
(40)
assuming that (x)]] < 0o for the loss to be well defined. Here in addition to the Hilbert space of parameters H. w we introduced the Hilbert space F of square integrable functions [x) | = (x))) < co. The Hessian H, and the counterparts of the kernel and feature matrices K. 38 and L (36) are now operators
H: H + H, H / (xxx)x @ (xx))d((x)
(41)
Kg(x)= = / (9xx
(42)
/ (b(x))((x) dd((x)
V : F H,
Dg(x)
(43)
Again, due to our interest in convergence of the loss (32) during SGD (3). we project parameters w to H θ ker H and target function y" (x) to F e ker K. After this projection the target is reach- able in the sense that y/ yx(x E F e ker K Im ff In particular, this means that optimum such that [ww,, (x)) - y² (x) = 0 may not exist, but there is always a sequence w,, such that lim [ww, (x)) - y (x)]
Now we proceed to eigendecompositions of H and K. According to the Mercer's theorem, the restriction of the operator K given by (42) to F F ker K admits an eigendecomposition of the form
Kg(x)= = Σ (x)x(x) / dn
(44)
k=1
with 2 > ʃ Φ& (xdxxx Sei and {dx(x)) being complete basis in F F ker K, The latter allows to decompose features as
x(x)= = Σ (x)
(45)
with | 2 = 1. Substituting (4) into 42 gives (uk, ui) = 5bz and P =>x- Substituting then
into 41 gives
H
ugu e Ug
(46)
A=1
which also makes )u orthonormal basis in H ker H.
19

---

Published as a conference paper at ICLR 2023
Under mild regularity assumptions, to O, and we will assume this in the sequel. Note that if X - 0, then the range the operator L not closed so that, exactly pointed out above, for some y" (x) e F ker K. there no exact finite-norm minimizers w E HH satisfying (w*, a5(x)) = y (x) for ju-almost all x.
show (König 2013) that the eigenvalues Å converge
B.3 OUTPUT AND PARAMETER SPACES
Recall that only assume J y dp(xx) < but do guarantee existence of the optimum w in the case N = d = oc. This means that considering SGD dynamics directly in the space of model outputs fx) = (w, (xx)) may be advantageous when w* is not available. Moreover, it will allow to completely bypass construction of parameter space when original problem is formulated in terms of kernel K (xx, x') (e.g. for an infinitely wide neural network in the NTK regime). First, let
us rewrite SGD recursion in the output space by taking the scalar product of (3 with +b(x):
[x+)(()] ((+1(x))
(1 a (B.)
3.1
(1(x)) (((x))
+ (t2 K (B:)yy
(47)
K(B))
8,I)
(a,K(B,)y"(X
K(B()g(x) 1 Σ k(x, xx)g(x;)
(48)
X.EB,
Here (x(x) = (ve, (xx)) = fi(x) - fx+1(x). Considering the approximation error Af(x)
=
f(x) - y² (x) we get
(A = (x)1
I I- II
B.I) B,I)
(Af(()) v,(x)
(49)
As 47
is equivalent to 33 we also get that 49) is equivalent to 2 when w is available. Thus can considered as the most general and convenient form of writing an SGD iteration.
49
Similarly to second moments output space:
defined for parameter space, we define second moments in the
(x, x') E[Af())Af(x))
Jou (x,x') [[Af(x))(x')
you (x, x') E[u(x)u(x')
(50)
(x,x') (="(xxx))
M Xx x') = (J x) )
where the superscript out indicates the output space (to avoid confusion with the respective moment counterparts defined in the parameter space). The same moments can be written in the eigenbasis {o(x)} of the operator K according to the rule
Cou kJ
=
(x(x)) (x, ())(()))
(51)
where G stands for either C,F or V. When w is available, the moments Cout in the output space are related to the moments Cn in the parameter space simply by
Cou VArdCCaa
(52)
In particular,
Coon A4
(53)
and
Σ ARCKE = Σ CRm [[[AS]1]]] co.
(54)
Even when w is not available we can still define Cst by Eq. 52), but in this case Σn Cka = oc.
C DYNAMICS OF SECOND MOMENTS
The purpose of this section is to prove proposition However, as discussed in section B.3 the description in terms of moments in output space (50) is more widely applicable than description in terms of parameter moments (3). Thus we first prove an analogue proposition for output space moments. For convinience, let denote the moments 50) Cot Jour your F F.
20

---

Published as a conference paper at ICLR 2023
with learning rates az, momentum 3z batch size b and random uniform choice of the batch Br. Then the update of second moments (50) is
Proposition 6. Consider SGD 449
C (Jut
Jou Von [41
I a,K (-
B.I BTT)
Ca (J
Jut V
( aa
a
3,1
ai
-a, K
3,1
(55)
Equ Dq
+ taa
Dqu
The second term represents covariance function gradients induced mini-batch sampling noise and is given by
1) for finite dataset D
{xi )
(x,x'
Σ
K(x,x; (x,, Xx),, -
(56)
N
X
and the umplitude 7
N-b (N-1)6'
2) for infinite dataset DD with density dp(x)):
(x,,') = f
xx)K(x", x')dx("""dxx") - KO,,"K(x,X))
(57)
and the amplitude 8.
Proof. First, we take expression of single SGD
49 and use it to express M through M.
A fi+1 (x) (x)2
(x)))
M (x, x') E
1+1
1/ (x) fe+1 (x')
(x)
(B:)
8.I
(Af())) (Afi(x)) (xxx))
(T a.K(B.)
3.1
E
ar a(tB)
BI
((a)(
3,I
aj K(B))
8,I
C (Joun)t
Jour Vour
[--ik((Bi)
3,1 77 3,1
((,)
B,I
K₂K
811
Cou (J
Jou Vout
I a,K aeKK K
3,1
T
K.K
BII)
3,I
E
KK (K.:
K Cou
E
E
COu
(58)
Here in (1) we used that A vr are independent from B, and therefore the average factorizes into product of averages. In (2) we introduced notation 5K(B,) = K(B,) - K and used that [B))(
Now we proceed to the calculation of the average from the last line in 58
2 out = E ("*BB)
(59)
1) Finite D [xx)
21

---

Published as a conference paper at ICLR 2023
Denoting for convenience function values calculated at X, with subscripts iy. and omitting time index f, we get
your
(K(B)
Σ
N Σ
Ki jy Cou
N
KEB
JEBB
1
Σ 2.5")
Ki K Cpj
12 K'J'EB
N2
(60)
1 b
b(b - 1)
Σ "z3
Kig Ky Ozy
02
N
N
B N(N 1) N
N b
N b
Σ (サ11)
N2 Kid Kz
b(N
P=j'
Σ
(N6 1) N2 Ke Kiz Oyg.
(N 1)b v8
Here unspecified sums over dataset D. In (1) used that the fraction of batches B which contain two indices i' * j' is (
N-22 2) ( [B[(B-1])) and the fraction of batches containing
I(N-
index j' is (1 ((²) = N Taking (N- we get (56)
[Bl
N-6
2) Infinite DD with density (x).). Proceeding similarly to 600 we get
Σ (x, =E K ((B) (x,x') KC K(x,'')
b
/
K(x, K(x,x") (x", xx"""""")K"""" Xx)((x, x')dn(x")
82
+ b(b - 1) KCO x') KC x')
(61)
62
( K(xxx"))<""t (x" ,x"x"" ,3)"","")((""" Kx) 1xx""xxx)xdx" x')dx(x")) KC K((xx'))
which gives 55 after setting 7 5.
Let us write second moments dynamics for both parameter space 4 and output space (55) in eigen- basises {ur} and {Øx (x)} using decompositions (45) and (44). For parameter space we get
(HuLtI JuI,++1 Juat V
(1
aqdx
3:)
(CuL. June)
(1
an Be (e)
T
- ardR
B.)
June
Via.e)
orde
(62)
/
(1 1)
Σ Ae
VViee de Cr.
(x) (x)dq((x) CkI.
And for output space
Crout Jout Jiktt1
kL.t+
ki.e+1 Voii k3,0+1
aidi
3e
ChiL Jout
Jout
("-aash
11 acdi Be
-add
B.)
B1
ki.t
(63)
(1 1)
draddedi
Σ k.V.
Couut
(x()))r(x))d (x)) (x) (x)d((()
Clout
Now are ready to prove proposition
22

---

Published as a conference paper at ICLR 2023
Proposition 1. Consider SGD (1) with learning rates au, momentum 3t. batch size b and random
uniform choice of the batch Bt. Then the update of second moments 1
(C++ JI+1 J1t VI+1) = -rr
[ a,
B B.
( ()
C J
aH
3.I 3
T
E(C,) E E(C,) E(C,)
JJ
+70
Here TereyE((GG (CC) is the covariance of gradient noise due to mini-batch sampling, with 7 (N-1jt and
N-b
E(C) = E (xs),, (x)()))) @ (xx))) HCH.
(5)
Proof. Note that in the eigenbasis f (x)} of K and eigenbasis {ux} of H output and parameter second moments are connected as
MA M
(64)
Using this connection rule see that parameter dynamics 62) in eigenbasis of H is equivalent to output space second moments dynamics 663 in the eigenbasis of K. Finally, (63 is equivalent to 55
proved in 6
For completeness, let us also write formula of SE approximation in output space:
ORe Joul kk,
ck.n+1
Jkk kk youl
(1 -acde
ardx accx
8. 3,
ORE kk.t Joul kktt
Rkut you! KR.t
(
Be
kk,
=
-a,2
Br
(65)
(1 1).
+r(ard6)
TI Σ
Cut
T2 CkC.
D SE APPROXIMATION
D.1 PRESENCE OF NON-SPECTRAL DETAILS IN THE GENERAL CASE.
In this section we show that for a mini-batch SGD noise term Σ given by (5) one cannot exactly describe respective noise spectral components Σx (uR, Eu;) terms of only spectral distribu- tions:eigenvalues Å and second moments components Clk- In turn, for SGD dynamics (4) this would imply that spectral distributions are not sufficient to reconstruct PC++1 from Pc.. The fol- lowing proposition characterizes non-spectral variability of Σι in a stronger sense: when not only Ak: C&k are fixed, but the full Hessian and second moments C.
Proposition 7. Consider problem with fixed Hessian H and second moments C, and finite dataset size N < co. Then, all noise spectral component Σys are bounded as
Exk < (N - 1))& Tr[HC]
(66)
Next, any chosen component ΣRR can take any value in the interval [A, B] (or [B, B < A) where endpoints A, are given by
(Tr(HC AcCra), B= (1)
(67)
Let us first discuss this result. Intuitively, the value A corresponds to the case where the an- gle between different (xx) and (x(x) is small, and therefore different rank-one terms in H
N Σ (x;) & (x;) greatly overlap with each other. The value B corresponds to the opposite case where all (x:) are orthogonal.
1
Now we characterize the magnitude of non-spectral variations. For this, from SE approximation 7 we can take A Tr HC] as natural scale of ΣςR. Then, bound (66) shows that relative to this natural scale, the magnitude of non-spectral variations of Σςa is bounded by N. The values A, B in 67 actually show that this magnitude is approximately achieved for a typical scenario with large dataset size N and the most contributing term to the trace Tr HC) having the same order as the full trace: Ag. CR [C Indeed, in this case B N CR. ~ NÀ, Tr[HC], A S A HH and the ratio B/A 2/A2 2 N.
23

---

Published as a conference paper at ICLR 2023
Proof of Proposition Recall SVD decomposition (see section with = [ (ux))== being the matrix of left eigenvectors, V O
B of the Jacobian V = VNUAV
k= Also, in the proof is
always fixed and denote an index for each want analyze Egs. Then, spectral component Σn can be written as
N
ERk = AE
RR
Σ K1
D / de, Akg Ckykz
(68)
1=1
.22
From representation (68) we observe that Tr = (N TrCH by using orthogonality of both columns and rows of V (Øf k= Then, since Σ is positive semi-definite, have
A E !r] and bound (66 follows.
The key idea of the proof for the second part of the proposition is that after fixing H and C, we fixed U and A but V can be arbitrary orthogonal matrix. Now we construct two specific examples of V so that Σ&g attains values A and B given by (67)
To get the A value, for all take OR JN. Substituting this into 668 and using orthogonality
Σ, OKI = Ōkgk2 gives
= A
ōxyia Axa, Ax,DDeo Anz Ckika - AR = (Tr(HC - ^u CRR)
(69)
k2
To get the B value, we set V to the identity matrix. Then we have
B= d
N
Σ k2
Sko v Aby Ako CRika
X2C NA XAC
(70)
1=1
Finally, we demonstrate that Σes can take any value between A and B by showing that two examples can be continuously deformed from one to another. Indeed, in the first example we specified only single row of V, and the rest of the matrix can be chosen that det V 1. As the special orthogonal group SO(N) is path connected, the required continuous deformation is guaranteed to exist.
D.2 SE APPROXIMATION FOR UNCORRELATED LOSSES AND FEATURES
Recall expression 68 for the noise-spectral component ΣEa- From SVD decomposition of Jacobian V VNUAV get (x:) ) (ug, (xx))) NA Using again SVD of the Jacobian and definition Oki (uk C we get
Σ 2k22 / de, A Ckyka (xb(x;),, Cxb(x;)) B BB[(w w (bx)))))
(71)
= E [(w,)()))))) - - 2L(xx)
=
Now we can rewrite (68)
Eιs Σ (x)) A²CaR ExND (xx(()2(()) - ~220ut
(72)
If Lxx) and v are statistically independent w.r.t. x D. the expectation factorizes into the product of expectations Ex-D 2 (x) d and Ex [2L(x) Tr(HC]. In this case
En = Ex-D (x(())2((x) - X2 d TH[HC] X CuR
(73)
which is exactly SE approximation with π =2 =1
24

---

Published as a conference paper at ICLR 2023
D.3 DYNAMICS OF SPECTRAL MEASURES
In Sec. 3 we mention that SE approximation allows to reconstruct the trajectories of observables Tr((HH)C C,] from initial state Co. The way do of course. through considering dynamics of spectral measures, which we explicitly write below. Denote, for simplicity, =
2 (J.J
Assume approximation 7 holds for all Ct during SGD dynamics spectral measures of all second moment matrices pC,, PJ,, PV, can be written
Proposition
Then
pc, (pc((d)) (dX) PJ, (dX)
PCc((A) PJo (dX (s((A))H(d)) Pvo (d2)
G.(A)
(74)
(dx)
where (A is X dependent 3 matrix, and (e is X dependent 3-dimensional vector. During SGD iterations Ge(A) and sc(A) are transformed
(1 (1- (1-oa Gz+1 -0;A(1 ax)) (72(2)))2
23x(1 - a:d)
3
1 23rr2
8
Gx(A)
(75)
(ax2)² - -TT2(1)))
(1- >T2(aa)))
23r(1 acd)
87
(1) = -a,a(1 - azd) 7T2 (T,
23,ayd
8
(e(A)
(and)² (a,)
-23,a
+ roo / (oo:()))((()))
1
(76)
t / x [x(onu(A)css +
11)
Proof. Let us rewrite dynamics second moments 4 as linear transformation of on (C,, J,, V. and then take traces Tr (r))(H)# of both sides with arbitrary test function ((A). Since arbitrariness of ((A) in the obtained relation implies the same relation but for spectral measures instead of traces
|a (dd))
(1-aad) PJ++: (dA) = -a,A(1- - aad) (ard)²
23,(1- and)
82
pc, (dA)
/EE,
1 1- 28,a
87
PJ. (dx)
+ 700
PE, PE.
(77)
oV..i(dA)
-23aa
Pv. (d.(
Now we use SE approximation
to express Pz, through Pc, and PH. The result is
PCC.1 (d))
(1- axd) (2)(2))2
23, - a2
pc. (d))
PJ++1 (dA)
-ai and) (ar22
23,00d
Ps. (dx)
PV., (dX)
(ard)2 - (Tz(add)
-33,a
5?
Pv. (dA)
(78)
(H(d)) (d))
+ maa / (ax)
PH (H(d)) (d)
(H(dA)
Now obtained form of recursion of spectral measures prove representation 77 and transition formulas (76). We proceed by induction. At t representation (74) indeed holds with Go(A) I and (A) Now assume that (74) holds for iteration & and substitute it into (78)
25

---

Published as a conference paper at ICLR 2023
The result is
(C (dA)
(1 -1-ay a,13 - 7T2(a)d)2
23 022)
18)
PC. (da((d))
PJ++s (dA)
-ap(11 a4d) (T2(1)))2
1 2,,00
B B..
Gx(2)
PJo (dd) PV. (d))
PV (d))
(ar1)² (T2((a1)
-23,002
(1 ayd)² 7T2(0x))2
21 -a())
8
-add(1 and) 7T2 rr2eaad)
1 23ra
3
st (1)pHH(d))
(az1)² - (
-2aaaa
PH(dA) PH(dA) (H()))
+ f (G (dX)) Gr,12 (d.
(H(dA) PH (dA) PH (dA)
+mma ((sx(())H((a))
(79)
Now we simply observe that right-hand side (79) has exactly the form 74 with (A)) 8;41
expressed through (e(A),s(()) according to (75 (76). (76)
Note that updates of matrix GA) and vector (A representation (74) depend only on learning algorithm parameters dt 3r, combinations 7T2, and non-zero eigenvalues of Hessian H. Also, we stress that although this proposition we considered all possible initial conditions with Jo +
0, V O, in all other parts of the paper focus the case Jo V = 0.
D.4 THEOREM
In this section we provide full version of theorem I and its proof, and also elaborate on how spectral measures Pc, change during training. We keep the notations from previous sections, in particular the model Jacobian is L the Hessian H = N LTT" the kernel matrix K N y v, and output space second moments C TV However, for simplicity, in this section we consider only finite training datasets of the results to all values of d, N seems to be rather technical and we expect it to lead to effectively the same conclusions.
oo and features dimension 2 N (including
oo). The extension
We start with analogue of theorem but for the isolated noise term { containing the tensor R. turns that the SE approximation is closely related what we call the SE family of tensors It:
- 72 (dirig + (igis)
Rajigigha =
(80)
is easy to check that the respective Σ(R) satisfies not only 77. but also a stronger version
E (C) = HT (r2 HCH
(81)
Let's focus only the R-dependent term in ( and for convenience rewrite it in terms L as
E CC)
Σ
Raiizisi T
CV L
(82)
N2
where V = =>(x,) denotes the i'th column of V. Then we have
Proposition 9. Consider a fixed tensor R. and let operators E', C be related to each over as
E' E' (C). Then, the following 3 statements are equivalent:
1. For any Jacobian &e and initial state C, spectral measure P' depends on C only through its spectral measure pa, and on Yi only through non-zero eigenvalues X of Hessian H.
2. Tensor R belongs to SE family 80
26

---

Published as a conference paper at ICLR 2023
3. For any Jacobian & and initial state C, the operator E' does not change by the replacement of y with V TUT in (82 for any orthogonal matrix U.
Proof.
From (2) (3). Substituting SE form 80) of tensor R into
82 we get
Σ' = HT HCH
(83)
Then statement (3) follows from invariance of Hessian under orthogonal transformations of the Jacobian
VU' =H
(84)
From (3) to (2). Denote E' the result of (22 but rotated Jacobian Tे UUT in the right-hand hand side. is convenient to represent Σ'
E'= Σ
Raiizisia Ciyla L
i2
(85)
N
$1.82.83.1611
where R is the rotated tensor R with coordinates Rigézigis statement (3) implies ''
Unin Ujois Unis Ujata Riiiisii As
Σ' we get
0 = Σ
Σ
Cout L X L
(86)
N2
11,42,83.24
where the tensor 8ñ = R R is the difference between original and "rotated" tensor R. Note that 2 R. R and therefore 81t R R symmetric w.r.t permutation of the first two and the last two
2
indices. Taking Jacobian with the full rank rank( (V)) = N and exploiting permutation symmetry
i + ig of 6R we notice that (86) implies
RR
your iass
= 0
(87)
1g-1==1
Next, we observe that for full rank v we can choose C such that C is an arbitrary positive-definic matrix. Then, permutation symmetry ia ++ is and arbitrariness of O implies 8R 0. As U was arbitrary, we see that
U3 juis Uzai Una Risziiss Ripigisks VU E O(N),
(88)
where O(N) is the group of orthogonal matrices of size N. According to Weyl Weyl (1939
(see
Jeffreys (173) for compact reference), the equality above implies that is isotropic expressed as sum of products of dirac deltas dar and possibly one anti-symmetric symbol 6j1j.. our case we can not have tensor: for N it has higher rank N > rank(Fi) Rf); ; for N 4 does not satisfy symmetries of Rt; for N = 3 it cannot be present because 3 is odd; for N = 2 all three possible combinations d2 Eiyia 0lgs, ét, 8,,6, Cizis again do not satisfy permutation symmetries ++ in and i3 + 24. Then we left with 3 possible combinations of Dirac deltas
R ci Sigka + c2 1i + C30 8
(89)
Again, from symmetry considerations we get C3. and finally recover (8)) by setting C1 = 71 and C2 = C3 T2.
From (2) to (1). We take arbitrary test function ((A) and calculate its average with spectral measure p
f (()))=))))))
+2 HCH)
= [H [HCH [H]H
(90)
=/ ((2))22"pc(dd) + (H((()) (po(ds)
27

---

Published as a conference paper at ICLR 2023
As ((A) is arbitrary, we obtain
= PcX + (1I Xxo(ax)) pr(d))
(91)
As the right-hand side depends only on the initial spectral measure pc (dA) and the Hessian spec- tral measure (H(dA), statement (1) follows from the fact that pH((A) depends only non-zero eigenvalues of H.
From (1) to (2). We take arbitrary U ( and follow the notation of the From (3) to (2) proof.
As we already noted, original and rotated Jacobians have the same hessian H H, and therefore eigenvalues λy- Then, as we do not change C with the rotation of Jacobian, statement (3) implies that spectral measure of E' and E' are the same: PV' Ps. For arbitrary test function o(A) have
0 = / ((A)(P))))))) (dA) - PP (dA))
N2 THHHH) Σ
liniz Vi 80 122
1142
(92)
T (H) 1 (6PeC.o )2:12
N2
Tr (KK
Let us now fix some arbitrary output space matrix C As V and C were arbitrary, we can choose
them so that we get arbitrary K but C C It means that in the last line in (92) we can consider
Cout be fixed, but K, be arbitrary. Then (92 implies that SRC and we can repeat respective steps of From (3) to (2) proof to get the statement (2).
With proposition 9 as a basis, we are ready to state and prove the full version of theorem !
Theorem 3. Consider SGD dynamics 44p with noise term (8) initial condition Vo 0 and y, og > If tensor R is fixed, the following statements are equivalent:
1. For any Jacobian Ve and initial state Co, spectral measures pc, occurring during SGD are uniquely determined by the initial spectral measure Pc,, the eigenvalues As of Hessian H, and SGD parameters {ay, B:
2. Tensor R is given by Rigizisis TT 8igis Dan (Big Orgis + & Sigis ) for some Tı, T2. and Eq. 7 holds.
3. For any Jacobian We and initial state Co. the trajectory [C is invariant under trans- formations V UUT of the feature matrix V by any U E O(N).
Proof. From (1) to (2). Let us consider the first iteration of SGD. The initial condition Vo = 0
implies Ju Vo 0. Also, denote for simplicity E, E C.). Then, according (4) we have
C = (I 00 H) H) + 7a
(93)
Respective spectral measure
PC, = (1 - and)² Pc. + rapps, = -(1 a11- Pc. + tagpE
(94)
Then, according to statement (1), we get that yy; uniquely determined by Pc, and eigenvalues of H. and therefore R is given by 80 according equivalence of statements (1) and (2) of proposition 9
From (3) to (2). Considering first iteration of SGD as in the previous step, we get that E is invariant under orthogonal transformations of Jacobian V. Then, statement (2) follows from equivalence of
statements (3) and (2) of proposition 9
28

---

Published as a conference paper at ICLR 2023
From (2) to (3), We proceed by induction. First, notice that Co. Jo. Vo are (trivially) invariant under orthogonal transformations of Jacobian V
V
To make an induction step, assume that C; J,, V, are invariant under orthogonal transformations of Jacobian. Then, according to equivalence of statements (3) and (2) of proposition 9 and invariance of Hessian H under orthogonal transformations of Jacobian Yr we get that Σ, is also invariant, As first term in 4 depend on y only through we get that the whole right-hand side of (4) is invariant under orthogonal transformations of Jacobian, and so are C,+1, Vi+1-
From (2) to (1) As noted in the statement (2). SE approximation 7 is satisfied. means that
can apply proposition 8 with initial conditions PJ. PV. O. Then, according to representation 74) and update rules (75)(76). spectral measures PC, are uniquely determined by the initial spec-
tral measure PCo, the eigenvalues Å of Hessian H, SGD parameters {ar,, Br Je and parameters
TI 72. Recalling that 7, TI, 72 are considered to be fixed statement (1) finishes the proof. C
D.5 TRANSLATION INVARIANT MODELS
The purpose of this section is prove Proposition 2
Proposition Consider a problem with regular grid dataset D {xi 2mm
N1 mx i
(Z/N x (Z.)&ZZ) on the d-dimensional torus T (R/22Z2) and with translation invariant kernel K (x, x') K(x -
-x'). x'). Then holds exactly with TI 1,72
Proof. Statement with 71 = 1, 72 = equivalent to Tr(o(H)E) Tr(HC) T( VΦ.
(95) We will proving this latter statement. Denote the grid-indexing set (Z/Nı x... (Z/NaZ) by Z./ /Z. Thanks to translation invariance of the data set and kernel, the eigenvectors u, of H have an explicit representation in the x-domain in terms of Fourier modes. Specifically, let K be the x N kernel matrix with matrix elements
TH
Kjj = Ki-j = Kj-i N (((x)),((3)) i Z/NZ.
(96)
Consider the functions Øk D given by
Φ (x;) = 1 Óua e Z/NZ.
(97)
These functions form orthonormal basis in CD and diagonalize the matrix K:
Kij Σ
Akokadn
(98)
KEZ/NZ
with
A
Σ
(99)
iEZ/NZ
Operator H is unitarily equivalent to K up to its nullspace. Assuming that the Hilbert space of H is complexified, the respective normalized eigenvectors uk of H can be written as
Uk- (NNx)) 1/2
Σ (xx))
Huk
(100)
IEZ/NZ
Assume for the moment that the problem nondegenerate in the sense that all A so that all the vectors 4j are well-defined. It follows in particular that the indices & of the eigenvalues D can be identified with k E Z/NZ. By inverting relation 100).
(xxx)
Σ
(101)
KEZ/NZ
29

---

Published as a conference paper at ICLR 2023
We can then find the diagonal elements of E using Eq. (5) and Eqs. 197). [01
Lik
(x;), (x)(x;)) (xx,) HCH
(102)
kk
N
X
1/2 OKr 1/2
Ckk Ckk
(103)
k'
Ak
Ak'
1/2 \1 CKev Σ Preidri 1R
(104)
Ak
Σ A1C11 - Cxk
(105)
It follows that
(r((H)) = Σ ((xx)22 Sick
(106)
k
= Σ ((K)A E C1C Σ (AA)) k
(107)
k
Tr(Ho THCHC - Tr(H (H)C).
(108)
which is the desired Eq. 95
Now comment on the degenerate case, when Av for some k. This occurs if there is linear
dependence between some (xx)), i.c. dim Ran (H < N. The vectors uk given by (100) with
A * form a basis in (an(H). Inverse relation (001) remains valid with arbitrary vectors ug for
k such that 0k- = 0 (this can be seen, for example, by lifting the degeneracy of the problem with regularization and then letting -> 0). As result, Eq. (TO5) remains valid for k such that A > 0. On the other hand, if uN is an eigenvector corresponding to the eigenvalue A = 0, i.c. ug E ker H, then Uy is orthogonal to (xx) ) for all and then Lik = for such k. Thus, formula (105) remains valid in this case too, and so Eq. (088) holds true.
E GENERATING FUNCTIONS AND THEIR APPLICATIONS
E. REDUCTION TO GENERATING FUNCTIONS
We clarify some details of derivations in Section
Gas of rank-1 operators. Let us first discuss SE condition 7 Recall that we assumed H to be diagonalized by the eigenbasis uy with eigenvalues dy. If the eigenvalues of H are non-degenerate, SE condition can be equivalently written as relation between the diagonal matrix elements ERR, CR& of the matrices E. C:
Σμa d E'CL TX Vl
(109)
(where we have incorporated the convention Ti 1, 72 7). More generally, if several Å may be equal to some À and so the choice of cigenvectors un is not unique, conditions (109) should be replaced by weaker conditions (invariant w.r.t. this choice):
EEA Σ E
NC AC TX
Σ Cek. k=
AA spec
(110)
Note, however, that representation (109) can be used without any restriction of generality if we are interested in the observables C, (in particular, L(t)) along the optimization trajectory 5 by Proposition g such quantities are uniquely determined by initial conditions and the learning rates of SGD once the SE approximation is assumed.
Representations 9 now follow directly from the evolution law 5 by inspecting the evolution of diagonal elements submatrices C,, V, using Fqq. 109).
30

---

Published as a conference paper at ICLR 2023
Genarting functions. using Eq. (13):
Eq.
16 follows from relation between generating functions obtained
[(2) Σ
L(t)2
(111)
T
Σ
V/+++ +
Σ
UT+1- Almm - 1)) =7
(112)
T-0
lm=1
0c
00
00
Σ
V +
U2+1 U2+1L(s)) L())
(113)
T=0
1=0
(2)
(2)
(114)
where we have used the substitution T fm t, tm =
The resolvent formulas in Eqs. (17)
19 follow from the definitions of U, V, (Eqs.
20 of the noise and signal generating functions Ú(2), (=) series of scalar rational functions. Consider first expansion (18) for Ú(z). The terms of this expansion correspond to the eigenvalues A of H. For each eigenvalue, the respective term is, by Eq. (17).
We explain now the expansions 18
x²(())) (1- -
(115)
where AĄ is the linear operator acting on the space of 2 x matrices and given by Eq. 11 The scalar product in (115) can be computed as the component 21 of the solution of the linear system
(1 - 2=AAA) TL m22) (1
(116)
Since all the components of the x matrix of the operator 1 2AA are polynomial in the parameters ar 9, 7, T, A, z, the resulting æ11 rational function of these parameters. The explicit form of æn is tedious to derive by hand, but can be easily obtained with the help of a computer algebra system such as Sympy (Meurer al. 2017). We provide the respective Jupyter notebook in the supplementary material
The terms of V are obtained similarly, but by solving the equation - (1==AA)(0
#1 212 ) = (6 221 222
(117)
instead of Eq. 116).
E.2 CONVERGENCE CONDITIONS
Proof of Proposition
3
Note first that Proposition is implied by the following lemma:
Lemma 1. Assume (as in Proposition
3 that 3 E (-1,1), < 2((++)/) /max and y E [0, 1].
Then
1. 8(a,3,7, > 0 for all z E (0, 1);
2. Ü2), (2) are analytic in an open neighborhood of the interval (0, 1) and have no singu- larities there;
3. (z) is monotone increasing for all E (0,
Indeed, by statement and representation 16 singularity of Ī(2) ((3)/2(1- on the interval (0, 1) can only be pole resulting from a zero of the denominator =i). By statement 3), such a zero is present if and only if (z) > 1, as claimed. It remains to prove the lemma.
httpps//cooab. research http://wwcw.peeeeccccco///eeeeee http:///wrreeeeeeaaaaaeaeaeaeeer usp-sharing
bw?
31

---

Published as a conference paper at ICLR 2023
Proof of Lemmal
1) Our proof of statement 1 is canputer-assistdd. Note first that this statement can be written as the condition W = Ø for semi semi-algebraic set W C IR defined by several polynomial inequalities with integer coefficients:
W = (3 7,2 R (-1 < A 2(3 1)) A A S 1) (0 < 2 <
(118)
5))
(119)
where we have introduced the variable a ax. By Seidibbred theorem, the emptyness of semi-algebraic set in 8" is algorithmically decidable problem (solvable by sequential elimina- tion of existential quantifiers, via repeated standard operations such as differentiation, evaluation of polynomials, division of polynomials with remainder, etc., e.g. (Basu 2014)). This procedure implemented some computer algebra systems, e.g. in REDUCE/Redlog (Dolzmann Sturm We used this system to verify that W Ø:
1997..
1: rlset reals$
2: S a**2*bxgvz**2 + ax*2*bazxx2
+ a**2*g*2 ax*2*z 2xaxbe=2
2*a*b+2**2 + 2+a*b*z + 2*a+z bx*3*2**3 b**3+z**2 + b**2*z**2 - b
bxz+*2 b+z
1$
3: W := (, -1<b and b<1 and 0<a and a<2* (b+1) and 0<=g and g<=1
and 0<z and 2<1 and S<=0)$
4: rlqe W; false
2) Statement 2 follows from statement and the fact that S( B, x, (1-2)(1-3:111-833
A O. Indeed, since the series Σx XE and Σn CsCrsO converge, singularities Ü, outside the points with zeros of S, which are excluded by statement 1.
1,,11 f can only be poles of some of the terms of expansions (18). 20 associated
3) Our proof of statement is also computer-assisted and similar to the proof of statement 1. First we compute
(=ū(=)) -a0 Σ X Ría, C,,,3,>u,3
(120)
where
R(a, B, )x, = 22a 2x 1 + 2
+, (121)
By statement 1, sufficient to check that R 2 0 on the domain of interest. This done by verifying the emptyness of the set
Wi (3 3,
<a 2(3+11 1))
(122)
A (0 < (R(o <0)
(123)
in REDUCE/Redlog:
1: rlset reals$
2: R1 := -2+a+*2+b+2+*2 + 4+a*b* 4*a*b*2**2 + b**4*2**4
2xb+*3+2**3 2+b++3+2**2 2+b++2+2+*2 2+b+2**2 + 2+b+2 + 1$
3: w1 := ex b(a,,,, -1kb and b<1 and 0<a and a<2x (b+1) and 0<z and z<1
and R1<0)
4: rlqe Wl; false
We provide experimental verifications of statements and in the accompanying Jupyter notebook.
32

---

Published as a conference paper at ICLR 2023
Proof of Proposition 4
Recall representation 21
X
Ú(1)= Σ
(124)
)x 7(1+B) +
2(1-8)
an
and rewrite it as (1)) f( -(20-0) Y): Y
1-8
where f(z, y) is defined as
-[1+3]
f(a, -E du(T - y) x
A
(125)
Note that due assumptions a'mmax < <2(3+1 and T O, the denominator in 124 is always positive. Then, from (124) we see that (1) is monotonically increasing function of a, and there- fore at fixed 8, the convergence condition (1) < 1 can be written as < Qc. equivalently yo / < 1 //e.. where ze is the largest positive solution of f(z. Y) 1. Also, original definition of Acrit in 4 can be formulated in terms of f is a unique positive solution of f(Acrit 0) = 1. Now
observe that (x, Y) > f(x, 0) for x > Iec. and f (ar 0) is monotonically decreasing functions of for x > 0. These observations immediately imply that > Aorit and thus the statement (22).
Now we obtain the estimate (33 by abusing the fact f(Aerit: 0) f Y)
we explicitly
calculate the difference f'(eerit, 0) [(Ie, Y) and use it to bound Ic. Xarit.
Au Io
Aecit Axa
[(eerit, - - Y) Σ
(126)
Au(T - -Y) + Ic)
Tdo + -ddrit)
A )orit)
A
Σ b
Σ (Au(T- + Øe)
(127)
AulT - 1) +)+ + xe
(Ads )erit)
(Txx + derit
(128)
Next we bound hh.ss from below with TAm aamaa---ertt (by bounding denominator with Aerit and equat- ing the rest the series to ( Y)). To bound the sum in r.hss we note that it can ex- pressed Σx axdi were >x ak [(xe, Y) and E. by = f(>orit 0) Then, the upper bound Amax E. agbE < bo TAmax + is constructed by noting that redistribution of the mass dg ōλo yields the maximum r.h.s. for fixed by and under constraint as 2 0. Σx ak S 1. Thus we arrive
Aaii
Aerit
Amax
S
(129)
T Amax + Aexit
T TAmax + Verit
To - Verit 5
Amax
(130)
7(1 + B)
Noting that a (3, =
we get 23 as
2
Re -eeit
Amax 1-8 3
lo (B,
/ (8,7)=
(131)
Marit
Aeric
Aerit 7(1 + B3)
E.3 LOSS ASYMPTOTIC UNDER POWER-LAWS
The main purpose of this section is to prove theorem First, note that A - and CkK > for infinitely many l (non-strongly-comwe problems, including power-laws 25 then rz S 1,
thus excluding the possibility of exponential loss convergence. Indeed, S(a, 8, 7,A. A
(1
=1 100--322 31) - 3)2). ). It follows that for small X the respective terms in expansion (8) have three poles z4,0: 1, = converging to 1, 1,,-1,,8-
respectively. By Eq. (6)). L(=) has a removable
singularity 24.0, moreover L(z+.0) < 0 if Cik. 0 and i is large enough, meaning that rz 2k.o- We then get r < 1 by taking the limit in k.
As rr for convergent trajectories and r7 for non-strongly-conver problems, we conclude that Fr. = 1. Then, we formulate a lemma allowing to calculate the - +0, asymptotics of the generating functions (18) and 20
33

---

Published as a conference paper at ICLR 2023
Lemma 2. Suppose that Sy and Ck are sequences such that S₁ kk o(1)) and Fie
Li=b G== = (1 + o(1)) as h + bo, with some constants v, x > 0. Also denote Ç x. Then
1.
S 1 x- o(1)), - co. s+1
C =
(132)
2. Assume Ç n. for some n > 0, then
20
Ck
r(c + 1)f(n 2 65-00 (11+0(1)), 8 ++0+
(133)
E + Sk)"
[(n)
Proof.
Part 1. Fix k and consider discrete measure A4 with atoms located at points æ1, (L/k)-":
1.
(134)
By assumption on Fg as k: -> oo, the cumulative distribution function of μı converges to cumulative distribution function of measure p(dx) 1,,,jd i.e. the measure ja weakly converges to μ..
Next, consider the function fi( defined at points by fe((I.)) k' with fu(00 = 0, linearly interpolated between the points ki, and with f(x 1) f(*,,)) Then fa(z) converges uniformly on [0, 1] to fxx = x as k + co.
The above two observations on μı and f& give
lim k0000
x
Σ 1=k
C.S: = lim
fa())uu(dz) = p do
(135)
k050
C+1
Part For any > 0 ha,s = [(sa)-2/vj. Rescale the sum interest and divide into two parts separated by index ka,si
Ra,e
Ck
-- +=((a,))+l((a.e))))
(136)
+ Sp)"
k=1
k=Ra,et1
We will show that
limsup (a(aa) O(aS-a)
(137)
e-0+
lim I2(a,e) = 100 (1 (1+ J(a). 6-04
dzs
(138)
These two fact will imply the statement the lemma since ( - n < and
lim ((a) ( ÇoS-1 dr (1+x)"
(139)
1+0=1/1 {-¹² (0---111-)-- -tt
(140)
-(B(n- (.())
(141)
34

---

Published as a conference paper at ICLR 2023
To prove Eq.
[37
summation by parts and the hypotheses on Sk, Fk:
C
8"
60
Ckkin)
(142)
(c Sx)²
==1
k=1
(Fx 8" 00
F
(143)
ka.e
= =0 Fi
Fkoictiha
ken
Fa (k.t
(k. 1))))))
(144)
(145)
= =((a-")).
(146)
To prove Eq. 138 it is convenient to represent the sum I2(a,e) an integral over discrete measure Ha.s with atoms Th,s k/e/
12(0,))= = / fed
(147)
Na,t = = Crdarie
(148)
k=
Se( = SS/E))
(149)
Again, by assumption on FE. as e + 0+, the cumulative distribution functions of Pa,e converge to the cumulative distribution function of the measure (a(x)) = 1 [0.0]] dzz5, i.e., the measures Pa,e weakly converge to At the same time, the functions fe (xr converge to (1 + ->- uniformly on [0, a). This implies desired convergence (138
Now we are ready to proceed with the proof of theorem 2
Theorem 2. Assume spectral conditions
1 and that parameters a, B, in Proposition 3 and such that convergence condition
(25 with
Ú(1) is satisfied. Then
KT({+1) 21-ū(1)(-
2n St2- 2- 5-6> 1/v,
1-8
LL(t) = (1 (1++
(26)
V 1/v)
2ad 1/24//2,
2- 1/v.
8(1
Proof. As asymptotic of partial sums E +0 we recall its expression 27}
ΣT= (tt) are associated with asymptotic of İ'(1 - E) at
d T(22)
d V
(z)(z)) 2(1 =ū(2))
(150)
dz
2(1 - =U()))
Full expressions for generating functions g +0 and \ +
18 and
20 can be significantly simplified in the limit
Ú(1 8) =
(1 +(x00))+0(0
(151)
+ 2 a4g
k=1
1-3
A4
V(1 6) Σ
+000))++00
(152)
E Too
k=1
I-$
35

---

Published as a conference paper at ICLR 2023
Differentiating get
Ad 2 I-3
01( -)) = Σ
(1 O(Ae) + +()))
(153)
[ 1-B
V' =)) =
2 (1 (1()())+0((E)
(154)
E + 2
Let's make a couple of observations, From asymptotic expressions above one can show that
b
X x
8(1-8)
(155)
()(1-s)
1 I~ b
(156)
(E + )1)² 2
k=1
k=1
8
CkCkk
1, {>1
İ(1 ((1-6) 6)
Ckk
(157)
= d
(56-1, s<1
ARCRE
Cek 11, {>2 2
('(1 -6)
(158)
(E + Ax))
d
1 <<2
Here the sign means asymptotic (e - 0) equality up multiplicative constant. Using these asymptotic relations we that for T.(1 8) leading asymptotic terms are given by either U7 (1 e) when Ç > 2 1 or by V'(1 =) when
Now we set for simplicity A = K = 1 and apply lemma to 153 and 154
to get more refined
asymptotic expressions. For (1 - 8) we assume Ç < because (1 E) = O(1) for Ç > 2 and refined asymptotic is not needed.
ARCrEe1 + O(Ae) + O())) 2
= Σ A'I ()
E 2 1-3
k=1
8
8
O(CrA)))
(L+00 ((e))
5 2
E E 2
(159)
(1)
2a
CREAE
+ 0 0(emin(0,5-11)
1-
((1-9) >18)
k=1
( r(c + 111 -C)
2a
2 (1+0(1))
1-
Here in (1) we first used first part of lemma to get ΣiSR CuAi = Chク + o(1)) Then for 1 < Ç < 2 the second sum is finite (see (188) and for Ç < second part of lemma 2 allows to bound the sum O In (2) we applied second part of lemma 2 to the first sum in (1).
36

---

Published as a conference paper at ICLR 2023
For (1 c) we have
aAg 1-8
(1+ + O(Ae) + 0()))
U1 - =) =
k=
+ 2
@ 2 (1 O()))
X
7
Σ
0
(160)
4
k=1
=(1-0)
+ A1s
4
k=
=(1-0) 2n
+ Xx)
2a
2 [2 4v
- TT
2a
-1 (1+
1
In (I) the second sum can be bounded (1) after recalling that
1. In (2) we again applied
second part lemma 2 with C = X and n = 2 since Eizk X 20- -- o(1)) To get
1
22+
coefficient in (2) we used r(3-1/v)
r(2-1/w)
2w-1
Substituting (66 and (159) into 150 we get C
mw (1+ + V(1)o e (1++((1))
(1- = 2(1 1 (())
(161)
2(1 (11)2
where Cy and Ct constants in the asymptotics (159) and (60). Picking the leading term depending on the sign of Ç (2 we obtain final expressions
D(G+1)2-1))
20x
g 2 31+ o(1))
<<2-
T'(1 6)=
201 20--((1))
1-3
(162)
V
2a
e (1 + o())),
(<2-)
S0(1-0(1))
The last step is to apply Tauberian theorem
(eeller (1971), p. 445) which states that if generating function F(z) of sequence F(L) has asymptotic F(1 -) -) = Cs Cs-P(1 + +1))) then the partial sums
have the asymptotic E F()) F ('(+o))) Applying this theorem to F(I) = (L(() and
get precisely (66) with K 1. To restore the values of constants A, K recall that the part the loss coming from signal (V())) is simply proportional to K, and A always goes in combination aA.
its generating function given by (62
Finally, let us show how (26) can used to obtain asymptotics themselves. Assuming that Ltt) Ct 2, get
28b 28a for the loss values L(t)
T Σ - Σ s+1 (1)) C / 1-5t (()+)++)))==++-- 2-
C 72 o(1) (163)
Comparing this with
26 find the constants C appearing in
28a (28
E.4 EXPONENTIAL DIVERGENCE OF THE LOSS AND CONVERGENCE-DIVERGIN TRANSITION
Convergence radius r and exponential divergence rate. By Proposition 3 the convergence
radius T7 < and only if (1) 1, and in the case (1) > it determined by the condition
1 rrrūÜrr).
(164)
rz < L, loss exponentially diverges as - oc: Proposition 10. Under assumption of Proposition suppose that F < Then
T Σ
V(rx) +(1+o(1)), (T-)))
L(rt 2(1 + ++(00(())))
(165)
0=
37

---

Published as a conference paper at ICLR 2023
Proof. Consider the modified generating function L1(2) I(r[2), which is the generating fune- tion for the sequence L(t)r
[1(=) = Σ =1((1
(166)
The convergence radius for L, is 1. We can also write
((r
L c)
(167)
2(1 - r1(11 - (r((1))((r((1-c)))
((z))
2(1 + r1 (1 (1+ o(1)) (e+
(168)
The statement of the proposition then follows by the Tauberian theorem.
If we assume that with some constant C, then this constant is fixed by Eq. (165)
L(t) Cri
(169)
((rl)
L(I) 2(1 + (rT)))
(170)
"Eventual divergence" phase with small a. Recall that we characterized the "eventual diver- gence" phase by the condition (1) = oc. In this phase we necessarily have rr. < However, Tr. can be very close to 1 (and thus "divergence delayed")- for example if the learning rate is very small. In this scenario we can derive an asymptotic expression for the difference 6, = TL-
In the sequel, we assume for simplicity that 8 0 and r initial covariance coefficients Cia.o satisfy power law relations 225 with & < < (corresponding the "eventual divergence" phase).
1, and that the eigenvalues and
Proposition 11. As 0,
1 r - 1/100/11100011 w/(1-v)
6. = -=--r_- TL =
(2aA) o(1).
(171)
4v
Proof. With 3 = and T = 7 = 1, we have by Eq.
18
x
Ú(1
2
(172)
20 6) +
22/(11 22xAR)
(173)
2A
>x/(A(1 20cdx)) )) +
<((2AA)
For any now apply Lemma 2) with 22 2 1 1,5= 1/v and n = 1 and get aA
Ú(1 (1 2(2v - 1) [(3 1/38111/[(1///111 1) ( (
1-1/2 (1+)(1))
20d
(174)
1 - (1)))((//--((aa))(////
(1++( (e 0).
(175)
4v
is easy to see that o(1) here is uniform for small a. The desired e. = e. (a) is determined from
the equation 1 = - (1-e))00 e1) -.). As + O, we clearly have e. = (c + 0. It follows that
= 41
- 11
-a- 0).
(176)
This implies desired expression
171
38

---

Published as a conference paper at ICLR 2023
We can similarly use Lemma 2 also estimate the derivative Ú'(1 c)
11 20dx))
(1(1-) = a²
Σ
(177)
(2ar) 6x) + e)²
22 2adx) (An/(A(1 2adx)) + (/(2)))
(178)
4A
4(2v r)r(3
r(3- 185-1/110 -1/v)[(1//)
2001
6
-1/V
(1 + +(1))
(179)
T(2 1/000(10022a)((xx +1)1
(180)
4v
1/v 1 e)1-))(1++(()) (c
(181)
(this result can also formally be obtained by differentiating Eq. 174) in E). In particular, for 6,
we get
(1+ o(1)), (a + 0).
(0)1 ET
(182)
"Early convergent" regime. If TL 1, then, discussed above, loss eventually diverges ex- ponentially, However, in general, we expect this divergence to occur only after some interval of convergence (see Figure righht). We will give the loss asymptotic for this "early convergent" phase and give rough estimate of the time scale at which the transition between the two regimes from the previous paragraph, and assume
occurs. We retain the assumptions 8 0 and y =
power laws (25) with 2 < v < 1 and Ç < 1.
We start with finding the loss asymptotic in the early convergent phase. First we argue that this phase can be approximately described by the loss generating function
1 1 (z), 2
I(2)
(183)
ie. generating function [16) in which we approximate the denominator by 1:
1 xz)
(184)
This approximation can be justified for TT by noting that, by Eq. 174
Ú(1-)) ( U(1
1/v
e.) ( 1-1/v
(185)
with 1 1-1 < 0. We remark that approximations (14)) and (83) essentially mean that we ignore terms containing factors U2. associated with noise in expansion 12)
relation (13) held for all z E (0, 1), then we would have an asymptotic convergence of the loss for all t. Indeed, under our assumptions,
CrCRO V(I (1-6)= 2ad (1 c)
(186)
and V(1 c) + xo as e e (since ( < 1). Applying again Lemma
we find
Au AaCaa.D/ 20rAx)
v(1 -) =) 2aA Σ
(187)
Aa/(A(1 2ad+)) + (/22A)
K
I(s 1)r(1 ()(
1))((1())
2ad (1+
E)
(188)
2aA
K
r(s 1)[(1 (2(a)) - + ((1), (e ( 0).
(189)
2
39

---

Published as a conference paper at ICLR 2023
The Tauberian theorem then implies that
T
K
L(t)=
I(s 1)T -((2A))))))))) ST ((1 o(1)
(190)
21 - ()
1=0
K
r(5 + (1 (c- 0).
(191)
2(1 5)
Assuming a power-law form
L(t) Ct
(192)
we then find
K l(s + 1)(20A) 2
L(0)
(193)
This expression agrees with "signal-dominated" loss asymptotic (28a) up to the factor (1 (1- (1)) (missing because of our approximation (184).
Transition from convergence to divergence. Of course, the obtained asymptotic (993 does not hold for very large since representation (183) breaks down at Z TL. and the loss trajectory switches to exponential divergence. We can roughly estimate the switching time bblowup by equating the asymptotic expressions for L(t) in the convergent and divergent phases:
K r(c + 1)(aAA)
(rL)
folowop
(194)
2
hlowup
2(1 + r (rL))
This is transcendental equation not explicitly solvable in Untowup- We simplify it using number of approximations. Using approximation 1899 for Vir and the faet that c, small, we can simplify this equation to
1 rg. -1,00'(1 - (.) ! blowup ~ (1 - c.,)
1 c
- Inlowap ss dbolowwu
(195)
T(1 -C
Next, using Eq.
182
we can approximate the l.h.s. by
1/v -
r'1
5 blowup ~
r(1 ()
E, 56-
bbooupp
(196)
It follows that Eq. 195
be approximated
1/v - 1
r(1 (1- (c. (bblooup))) ) = e. Ablowup
(197)
Thus, we can write
tblowup P3 a.
(198)
where a, is the solution of the equation
1/v r(1
a 5 =
(199)
If faiv = -1/ log ry denotes the characteristic time scale of exponential divergence, then at small have daiv 1/eu, the characteristic transition time bboowp is related faiv by the simple relation
blowup = autdiv.
(200)
E.5 BUDGET OPTIMIZATION
It is known that for SGD with not too large batches, there is approximately linear scalability of learning w.r.t batch size, in the sense that to reach a given loss value with an increased batch size the number of iterations should be decreased proportionally (Ma et al.) 2018 Shallue et al. 2018). This law is naturally no longer valid at very large batch sizes, since in the limit of infinitely large batch size SGD just converges to the non-stochastic GD, so that further increasing batch size when
40

---

Published as a conference paper at ICLR 2023
MNIST
101
L(2)L(()
b=1
b-3
b-10
b=30
b=100
10a
10
104
307
101
102
12
tb
Figure 5: Representation of loss [t) obtained using different batch sizes as (approximately) function of the computation budged bt (regardless of how it is factorized over b and t) for MNIST. Momentum is fixed = and the ratio a/b is kept the same for curves.
it already large brings little benefit (the effect of "diminishing returns"). Figure slswws that this effect is also reproduced in our setting.
Let us now show that this effect should be observed in the signal-dominated phase if for each batch size b, learning rate and momentum are chosen optimally. In the respective part of Sec. 6 we showed that loss can be approximately written as L(t) (Qotrt) the optimal learning rate and momenta are located near the convergence boundary and the momenta is close to 1. Since the effective learning rate bounded Oeff <
7Pom (see eq. (22)) and the bound is tight for the belu close to 1 (see Eq.
approximate Cett with "dara This leads to the optimal loss L(t) z (²(t/r)= Replacing y with (which is very accurate for large dataset sizes N, see prop. we get L(t) = (b which is exactly linear scalability of SGD with batch size.
23 in the optimal loss value
E.6 POSITIVE VS. NEGATIVE MOMENTA
Recall that our analysis of the effect of momentum is based on the approximations 28a Ltt) at large
28b for
KT((+1) 20d - S
<<2-1/1
(201a)
2(1-((1)
Taaprrxx (t)
20 1/gl//u2
<> 2>2-1/ 1/v.
(201b)
80(1- (1))
We now formulate and prove the full version of Proposition 5 describing how Lapprox (t) is affected by changing the momentum from 8 = 0 at the optimal a. Note that t enters expression Lapprox (t)
only through the factor 1-S or f1/v-2, As result, the statements below hold in the same form for all é > O.
Proposition 12 (full version of Proposition Suppose that the eigenvalues dg and initial moments Ckk are subject to large-k power law asymptotics (5 with some exponents v, 0,5 2/v and coefficients A, K. Denote Clopt arg min, Lapprox (t)
I. (signal-dominated phase) Let Ç 2 and < T 1, S Then
1. At B 0 and a function of a, the function Lapprox has unique minimum in the interval (0, (mmax), where Otuax the critical stability value defined by the condition Ü(z 1, a Ommax 3 0) 1. Accordingly, Copt arg minl Lapprox (f)]a=0 is well-defined and unique.
2. 0g Iapapp (esaappox((t)33-000 (t)) < 0.
II. (noise-dominated phase) Let Ç > - 1/v and Let T = 7 Then
1. Olopt 2(v-1) (T)
41

---

Published as a conference paper at ICLR 2023
2. The sign of 0g Lapprox
equals the sign of the expression
== v
(v 1-- Tr Tr Cr)
(202)
Before giving the proof, we make a few remarks.
1. The "noise-diminatee
part of our result is proved under much more restrictive assumptions (r y = 1) than the "signal- digmal-dominated" part (0 < T < 1,0 < S 1). The reason is that the effect of is much more subtle in the noise-dominated case, and is described by relatively cumbersome formulas. Setting T = allows us to substantially simplify these formulas. contrast, improvement from positive in the signal-dominated case appears to be general and robust phenomenon.
2. It can be observed from the proof in the noise-dominated case that the second, negative term in Eq. (202) results from the "signal" generating function (1) in Eq. (011), while the first, positive term results from the "noise" generating function ( A similar obser- vation can be made in the signal-dominated case, but Eq. (201a) only includes the "signal" generating function (1), so the counterpart of expression E in this case is always negative.
3. r = = 1, then in the signal-dominated case Coopt has the simple form Copt =
26 c+1 (Tr[H)
4. If
1, then the critical stability value 0'ax which Ü 1, a, B 0) has the simple form Qmax = ((Tr(H)) In particular, we see from the explicit formulas for Opopt that in both signal-dominated and noise-dominated regime Cappt Ommxx expected.
Proof of Proposition 12
Let U1 V12 denote the values 1), İ(2 1). Recall from Eqs.
18
20 that
(a
(i(a,)) Σ
(203)
aBT - aBAx + atT>; a. - 2,3 + +2'
CRa,n(2331 + 83 B3 -3+1 +1
V)(a, Σ alaBrydx + aBAr aTT/A= aadx 2,3 2
(204)
As preliminary step, note that at B 0 we have:
0) = Σ aT 1 1))+2
yadx
(205)
0,ú1(a,0) = E adx(TT 11+
7/1
yadx An -)
(206)
(a)u(ry-1)- 2)2
27dx
Σ (adulTT - 1) -1)
(207)
yadx
Jadz axx
0gūi(a.)) Σ adu(TY - 1)
(208)
(ade(TT - +
277 X2 + 22ya4x
(209)
(adu(TT + +2
270 2
-E (adu(Ty - 2)2
(210)
By our assumptions, 1 < 0. As result, the function Ū (a, 0) is monotone increasing in and there is unique critical stability value Omax at which Ú, (a 0) 1 and Lapprox (t) becomes infinite by Eq. (00))
I. (signal-dominated phase)
42

---

Published as a conference paper at ICLR 2023
1. It is convenient to consider the logarithmic derivative of Lapprox (t) B=0. By Eq. (01a).
Og In Lapprox()) (t)l==0
(211)
= =a5 + (1 - (1)
(212)
(1 E (TA - 1) + -
7adx
Σ k
27
(213)
(adx(T7 1) 2)
(1 0))=
(214)
yadss
(T) -1 2)3 - c]
2
(215)
adx(TT - +2
x xa (1-
(216)
Note that the last expression in square brackets monotone increasing in a. because TV < 1. If
= O, this expression negative. On the other hand, at a Omax such that U1 (Amaxx 0) = 1 this
expression is positive. It follows that there is unique point Oopt E (0, Omax) where
a Tapprox (1) (1)0==am..000 =0
(217)
and (approx(t)) (+) attains its minimum.
Now observe that
0g In Lpproxx (f)
(218)
[1 (0-5)-50((1-999 (1- ((-)))))=((-((())))))((((((( 0]][a=
(219)
2ya X3
= (1 -i,(a, (((a))))(-))))))))))))) -0, 0) Σ (adu(Ty - 1) + 2)
(220)
x (2,
(221)
2ya 2
aa In Lapprox (1)
(1 (1-U (a;
(222)
(ad7<77 1) + 2)²
At Omax the first term in Eq. 222 vanishes, and the remaining term is negative, proving that
3g L apprrox (t) (t))==0,a=oopt 0.
(223)
II. (noise-dominated phase) A 1, formulas 2055
210
admit a simpler form:
(i(,,0) Σ
Tr
(224)
2
0. 0.ū 0)
1
Σ d
1
Tr
(225)
2
2
h
0aū1(a,0) 0) = a
Σ
As(1 - ada)
(226)
2
k
a²
= = U, (a, 0) 2 Tr[H').
(227)
Snce Eq. (201)) for the noise-dominated phase involves Vi, we will also need similar formulas for we have
this expression. At
Vi(a,
1
Σ
Ckk Tr(Co).
(228)
2a
2a
b
0, (,Vi
Σ
CRA.D
1
Tr(Co),
(229)
2n
20
CCk,,o (2crh 1)
0 V,Vi Σ
(230)
2a
1
1
Tr
Tr(Co]-
(231)
2
2ax
43

---

Published as a conference paper at ICLR 2023
1. Differentiating Eq.
201b)
O, In Lapprox (1)
(232)
= a 1/20
Ba 1/2 + V1- (a, 0)0 (a 0) + 2(1 -
0)
(233)
1
1
+ 22(1
Tr
-11
Tr(H]
(234)
av
a
2
a)
(1
[r[H]
Tr H (1 Tr Tr
(235)
av
2
-1) +
1
Tr (1 TH
(236)
2
2v/
The expression in brackets monotone increasing in a and vanishes at
2(v 1)
Cppp
(Tr(H))
(237)
3v
proving statement
2. We have
0g anLappoox(t)))
(238)
[1 01 14 + 02000000022000000000000001101110
8-0
1 = - 21 T-
a² a T-
- T-
-
At Opope we can simplify this expression by recalling that this Eq.
234 vanishes, and
using Eq. (237)
approx dg In L (4) (000000m
(239)
aopt (THCo + 22(11 21
2 - Tr(H] -1 ( (
Copt
appt Tr
2 (Tr(Col) (Tr[H)) TH T TC - (v- - 1) Tr(Co] Tr(H3]
Coopt
(
proving statement 2.
In Figure 6 we illustrate this result for range of models with exact (not asymptotic) power laws (25 different and x. In this case, in the noise-dominated regime, characteristic E is positive for sufficiently large but negative for smaller x. In agreement with Proposition T2 we observe the
experimentally found optimal be negative or positive, respectively.
F EXPERIMENTS
F. TRAINING REGIMES
We have considered four regimes of training with the increasing level of approximation for real network dynamics:
1. Real neural network (NN).
We used fully-con fully-comnected neural network with one hidden layer the NTK parametrization:
1 wi ('RLLUWW + bz)
f(x, w)
Here w (1,2) are the weight matrices, by are the biases for neuron I, and is the number of units in the hidden layer. We have used D 1000 in our experiments. For the sake of simplicity, the network outputs a single value for each input, ie the classification task is treated as regression to the label with the MSE loss. The model is trained via SGD in the standard way: at each step a subset of size is sampled without repetition from the training dataset and the weights updated according to the equation (T)
44

---

Published as a conference paper at ICLR 2023
26-1
0.X1m115.3.5 0.81=228.35)
64
600
Figure 6: Theoretical and experimental study of optimal B. Left: Dependence of the loss the learning rate a at 3 = 0, for different v and x (synthetic data) the "noisy" phase. The stars show the experimentally determined optimal learning rates arg ming L(a), while the dashed vertical lines show the respective theoretical values Copt = EEx A 30-1 The experimental and theoretical optima agree well for two pairs (v, x) deep in the noisy phase, but not so much for (v. x) near the
2(p-1)
phase boundary (ef. right figure). Center: Dependence of the loss log-ratio log ming I.(A) B
L(3)
experimental Ctopt- The stars denote a (B3). In two scenarios the minimum attained at < 0, and in the other improvement from 3 < 0 at large by Proposition 12 computed for the range (v, x) E [1,3] x [0, 6 of power-law models (25). Diamonds show the pairs (v, x) in the experiments
at 33 > 0. Right: The characteristic E given by (30) and characterizing
the left and central
figures. In agreement with our theoretical prediction, the characteristic =\t is positive in the scenarios
where the optimal 3 is negative, and vice versa. The dashed yellow line x 2v 1 separates the noise-dominated (top) and the signal- degna-ddm (bottom) regimes.
2. Linearized regime with sampled batches (linearized or sampled). Model outputs (xx evolve on the dataset of interest according to the equation (47 At each step a batch of size is sampled without replacement from the whole dataset of size N and SGD step is performed.
3. All second matrix moments dynamics (all sm).
We simulate the dynamics of the whole second moment matrix that is obtained after taking expectation over all possible noise configurations, defined by (58) with the noise term given in the last line of (60).
4. spectrally expressible regime (SE).
In the spectrally expressible regime, the diagonal of the second moment matrix is invariant under SGD dynamics, i.e diagonal elements evolve independently from off-diagonal terms. Therefore, one needs to keep track only of the diagonal terms. The dynamics in the output space follows Eq. (65
F.2 DATASETS
In the experiments we have used synthetic data with the NTK spectrum and second moment matrix obeying a power-law decay rule d = k- and AςC*R. = k --=-1 with specific v and x and the subset of size N of the MNIST dataset. The whole train dataset size 60000 requires much memory for the computation of NTK and is too costly from computational side for running large series of experiments, therefore we have used only part of the data size 1000 10000 depending on the experiment. The digits are distributed uniformly in the dataset, i.e there are in samples of each class. We have observed that this subset of data is sufficient to approximate the power- law asymptotic for NTK and partial sums since the exponents do not change significantly with the increase of the amount of data.
N
F.3 EXPERIMENTAL DETAILS
For Fig. one we have taken subset of N = 3000 digits from MNIST and ran 10 SGD updates for the real neural network with hidden layer, linearized dynamics, and in the spectrally expressible
45

---

Published as a conference paper at ICLR 2023
appoximation with 5 =1. One can see that 7 1 provides accurate approximation for smaller batch size b = 10, whereas for = 100 both J = =1 are rather accurate.
For Fig. 2 we ran 10000 gradient descent updates for every point the grid of learning rates a and B, where we have taken 100 values of a uniformly selected on [0, 4] and 50 values for 8 uniformly on [0, 1]. The grey regions on the 2d plots correspond the values of and /3 for which the optimization procedure has diverged. The size of the dataset is 1000 both in the synthetic and MNIST case.
For Figs 5 6 (left, center and right) we have taken the dataset of size N = 10000 both for synthetic and MNIST I case and performed 10000 steps of gradient descent. In Fig. 6 (center) we have taken
100 points uniformly on the interval a E [0, 0.5], and 100 points uniformly on the interval B e [-0.3, 0.3] for Fig. (right). Concerning the correspondence of optimal learning rates 6 with the analytical formula (30) we observe that quite accurate for large ratio K/VV since the ratio Guas 100 for xu,x) (1.5, 3.5), (2.0, 5.0), whereas for Caise
Cagaal
for (V, x)
(.5),
therefore both the signal and noise asymptotic are of same order for quite a long time, since the difference between the 'signal' exponent & and "noisy" 2 1 is relatively small.
F.4
VALIDITY OF SPECTRALLY EXPRESSIBLE APPROXIMATION
MNIST
Bike Sharing
sgemm GPU kernel performance
1.5
5r -0
2.00
53 -1.00
-1.00
- 8:01
Se 2:01
0.9
ST 1.2)
1.51
2.5)
05
2.51
0.6
sampled
2.0)
0.4
1.5)
03
Figure 7: Loss trajectories for linearized stochastic loss dynamics (solid yellow), exact simulation of the loss dynamics with full second matrix (dashed black) (55), and the loss trajectories for several values of r (shown with different colors). Linearized dynamics with stochastic sampling of batches averaged over 1000 runs.
In section 3 we introduced several scenarios where SE approximation is exact. However, more interesting and at the same time much more difficult question is how well SE approximation well suited to describe the practical problems, where we are given a particular Jacobian L, initial state Co and then generate the sequence of states C, according to with true noise term (5).
We can approach this question in a different ways. First approach is to consider individually the true dynamics with (5) and approximated dynamics with noise term 77) Then we can look at cer- tain quantities of practical interest and check whether they are close to each other. In figure 2 we compared stability regions in (a, 3) plane for MNIST and found a good agreement between actual dynamics and SE approximation with TI = 12 1. In figure [ we did it for the loss values on MNIST, but instead of noise-averaged loss values we compared with loss on single stochastic trajectory. To make the comparison for loss values more thorough, in figure F.4 we add average of several stochastic trajectories (sampled), true second moments dynamics (all sm), and several SE approximations with 71 = 1 but different 72. For MNIST we find excellent agreement with 72 1, for Bike Sharing dataset agreement is good with 72 = 0, and for GPU performance dataset the agreement is worse with optimal 72 changing during optimization from 0 to -1
The second approach is to run true dynamics of second moments, and on each step compare how true noise covariance Σ, given by 55 with SE noise covariance > (Ti, 72), which we take to be given by
(5) HT. 72 HC,H
(240)
46

---

Published as a conference paper at ICLR 2023
Approximation errer Eg(r) 13) att=g
Condition number of Ez(tu, T2)
Optimal approximation error Esto) .51
14.4
glbbal
11.2
9.6
4.8
-3
0.0
Optimal 12 Ty m-
Sharpness PEy(l, 121-
Leading E2 digerveccto angle
1.400
201
307
10
absolute
relative
0.355
Figure 8: Various characteristics relative approximation error E2(TI, 72) for SGD optimization on MNIST dataset. Left top: Dependence of B2 (11, T2) on at initialization. Red star denotes the global minimum, while light red star denotes the minimum on the line 71 Center top: The ratio of smaller to bigger eigenvalue of the 2 x 2 matrix E which determines purely quadratic part of E2(TI. 72). Right top: Comparison of globally optimal approximation errors and approximation error optimal on the line = 1. Left bottom: Value 72 minimizing E2 T2) on the line
1. Center bottom: length scale on line T = 1 on which cither approximation error increases by (absolute) when moving from optimum 72. increases by optimal approximation error (relative) when moving from optimum 7.. Right bottom:
Angle between T2 axis and leading egenvector of
E.
We measure closeness of true and approximated noise covariances with relative quadratic trace dis- lance
E2(T1,72) Tr (E,
E (TI: r2)) (Ee
E (T,)
(241)
Tr
Σ Σ
Note that linearity of ESE (T, 72) implies that E2(TI, 72) is quadratic in 71, 72. We analyze quadratic form Eg(T1, T2) in Figure F.4 From the bottom left part we see that the ellipses corresponding to level sets of (2(T1, 72) are extremely stretched. meaning that there exist a direction along which approximation еттог changes very weekly. To check that conclusion for during the whole SGD optimization, we plot in center top part the ratio of smaller bigger principal axes of level sets ellipses and see that during optimization the stretching only increases. One particular consequence that there a ambiguity in choosing the optimal Tı, 72, which means that we can safely put right top part where both global and on
1 and vary only 72. This conclusion is illustrated Ti = optimal aprroximation error are very small S 10-3 during the whole optimization. Finally, observe left bottom part of F4 that optimal 72 is extremely close to 1, which well agrees with the conclusion of the first approach based on predicting the evolution loss values and stability regions.
47

---

Published as a conference paper at ICLR 2023
F.5 ADDITIONAL EXPERIMENTS
To study the validity of the theoretical framework developed in this work in more practically interest- of ing setting we computed empirical NTKs of the pretrained models from torchvision and conducted experiments with model training in the linearized regime and spectrally expressible approximation with r = +1 on the CIFAR10 dataset Krizhevsky et al. (2009). One can observe from Fig that spectrally expressible approximation can accurately describe the averaged dynamic of the linearized model. The NTK eigendecomposition and target expansion in the NTK eigenvectors follow again power-law decay rule. In the experiments below we used popular ResNet-18 and MobileNet-V2 models pretrained on the ImageNet dataset Russakovsky et al. (015)).
b=10
b=100
Eigendecomposition
100
10-1
10:
0.3
linearized
linearized
SE (t= -1)
SE (r= -1)
10a
SE (T- 1)
SE (T - 1) C2-925
Au/Ao. v- 1.12
Cr-923
10-5
5w/So. x 0.25
102
101
107
20a
20
10a
10a
10a
107
10F
10a
201
107
20*
201
x
b=10
b=100
Eigendecomposition
100
1.0
1.0
10a4
linearized SE [T= -1)
linearized ST=
SE (t= 1)
(T 1)
10t*
Ac/Ac, == 1.28
Cr-070
Cr-020
Sy/Sp. 0.26
0.1
0.1
10"0
Figure 9: Loss trajectories and spectrum for linearized regression of prertrained ResNet-18 (Top) and MobileNet-V2 (Bottom). All notations follow Fig The power-law exponents are v = 1.12, x ~ 0.25, Ç ~ 0.23 for 18 and y ~ 1.28, x 0.26, Ç ~ 0.20 for MobileNet-V2.
In addition to the image and synthetic data we ran experiments with one hidden layer fully calle-connected neural network and with the approximate dynamics for two tabular datasets from UCI Machine Learning Repository: Bike Sharing dataset ( and SGEMM GPU kernel performance dataset F The former is regression problem, namely, prediction of the count of rental bikes per hour or per day, and we've ve selected the per hour dataset. There are 17389 instances in the dataset and 16 features per imput sample. The latter is regression problem well, there are 14 features, describing the different properties of the computation hardware and software and the outputs are the execution times of 4 runs of matrix-matrix product for two matrices of size 2048 2048. We predict the average of four runs. There are 241600 instances in the dataset. See results in Fig. 10
In order investigate the problem of the loss convergence and determination of the critical learning rate we performed experiments with different batch sizes = 1, 10, 1000, 1000 on MNIST and 1.5,
synthetic data with 1.5, x 0.75 (the 'signal' regime) and with "noisy' regime). For each value of momentum 3 we plot the critical value of learning rate Cceitt such that the dynamics starts to diverge. See results in Figures 12 and 14
3.0 (the
Moreover, we carried out
We adopted the open source https://archt es. http:///ww....aaaaooa/oaaaaaaaaa https://archive. iss usi http:///wwac//a///////////a//a/a
48

---

Published as a conference paper at ICLR 2023
b=10
b=100
Elgendecomposition
10"4
10a4
30"
linearized
linearized
(T-
SE (T- +-1)
SE (T- 1)
SE (T- 1]
Ag/Ag, 1.43
Cz-020
Cr-0.20
Sx/Sc. 0.29
30"4
10-1
b=10
b= 100
Eigendecomposition
10P
10t1
6x10
6x:10
10-3
10-5
30
NN
NN
10'
linearized
linearized
(T= + 1)
(T + 1)
(r== 1)
SE (T= -]
As/Ao, v 1.52 5y/5,, x 0.23
0.15
0.15
10-11
10a
10
10
10
10" 10 101 10
15
10
10a
15
10'
10'
Figure 10: Loss trajectories and spectrum for Bike Sharing dataset (Top) and Sgemm GPU kernel performance dataset (Bottom)... All notations follow [ig[1 The power-law exponents are y ~ 1.43, x z 0.29. Ç 0.20 for Bike Sharing dataset and 1.52, x = 0.23, Ç = 0.15 for Sgemm GPU kernel performance dataset.
the same experiment with empirical NTK computed on CIFAR10 for pretrained ResNet 8 and batch sizes b = 10, 100 (see Figure
CIFAR10
SF(7- 1), b-10
linearized,
SEIT 11,
linearized,
075
0.5
Momertium (1)
2.35
025
025
011)
D(1)
012)
(11)
-2.9
-125
-1.%
-635
(0.28, 0.94)
-375
(0.60, 0.86)
-375
(1:12, 0.98)
-075
11.04, 0.98)
plut
-97
Leamiig rate lo!
Leaming rate (e)
Leaming rate (a)
Learning rate
Figure 1I: L(t = 10t) for different learning rates a, momenta 3 and batch sizes b on CIFARIO.
Legend and notations follow Fig. 2
G MINI-BATCH SGD VS. SGD WITH ADDITIVE NOISE
In the paper we considered SGD with the natural noise structure induced by random sampling of batches. However, it is instructive to consider SGD with an additive noise model and observe how this surrogate noise model qualitatively changes optimization. Specifically, we will assume that on each iteration, a Gaussian vector with covariance G is added to the true loss gradient: v L(w) = V L(W(w) + gi, g: (O G). For quadratic problems the SGD iteration becomes
Aw = V++1
I O,
- H
3.1
I,I
Aw: V.
O2:
g: g:
(242)
49

---

Published as a conference paper at ICLR 2023
MNIST
1.00
1.00
b 10
1.00
b- 100
1.00
- 1000
66
6.15
0.73
0.75
3.5
0.00
010
93
Nomentum
6:3
0.25
035
125
600
0.00
000
3.00
D111 m
-3.5
011) =
0(1)
-12
011) -1
-012
23.33
2
0= 20-33
-190
1011
-375
-106
2
Leatning
Figure 12: L(t 10t) for different learning rates a, momenta and batch sizes on MNIST. Legend and notations follow Fig. 2
1000
Homerram.pp (3)
0.2
-33
6.2
023
0121
011) 1
=1
011] [11]=1
201 1.
20-10
Learning
Leaming rate (c)
Learning rate to)
Learning
Figure 13: L(t 10*) for different learning rates a, momenta / and batch sizes b on synthetic dataset with // = 1.5 and x = 0.75. Legend and notations follow Fig. 2
3.0
-24
Momortan 03
0111-
011) -1
0111-
0111
20
21x8
20
21.
20-33
20-31
Tearnin tate in)
Teaning rate (el
searinn rate 10)
Tewnnig mete (r)
Figure 14: L(t = 10a) for different learning rates a, momenta ,8 and batch sizes on synthetic
dataset with 1.5 and 3.0. Legend and notations follow Fig. 2
Then, due to Egg]] = 0, the dynamics of first moments is identical to that of noiseless GD, like in the case of sampling noise. However, dynamics of second moments different:
M++1 (1 (T- O. H
3I BI
- oz G. H
3. 8,I
T + a (1 !)
-
M.
(243)
Considering again constant learning rates a₁ Bt = B, we see that by making a linear shift of second moments we get the noiseless dynamics:
M - M = (1 (1- H,
-
8.1 (M. - M
-
II H 3,1 T 3,1
(244)
8II
H
where 0 is the "residual uncertainty" determined from the linear equation
M (1 (1_ H.H
8 BT
I
H
8. BI
G (1 !
1
(245)
o
H
For simplicity, let's restrict ourselves to SGD without momentum (B = 0). Then in the eigenbasis of H we have
CAkI411 - Cr (1 - a (C - Cgt),
C
a Gaa
(246)
Ae adx)
50

---

Published as a conference paper at ICLR 2023
Thus, optimization converges with the speed of noiseless GD, but the limiting loss is nonzero:
1 Σ 2 C > O.
lim L(wi) L (000
(247)
The key difference of SGD with sampling noise (4) 5 is the multiplicative character of the noise, in the sense that it vanishes as the optimization converges the true solution and M, becomes close to O. To summarize, the two largest differences between SGD with additive and multiplicative noises
1. Loss converges to value _ > 0 for additive noise and to 0 for multiplicative noise.
2. Convergence rate independent of noise strength for additive noise, but significantly de- pends on it for multiplicative noise.
H NOISELESS CONVERGENCE CONDITION
The convergence condition < 2(++8) for noiseless GD is known (see (Tugay & Tanik 1989)).
Ames
For completeness, below we present its derivation in our setting.
First observe that, in the noiseless GD, different eigenspaces of H evolve independently from each other. Thus need to separately require convergence in each eigenspace of H. Denoting the eigenvalue of the chosen subspace by X, we write the GD iteration as
Awi+1 Vi+1
(A Vi+1
S = (1
as ax
9 3)
=
(248)
-ax
The characteristic equation for eigenvalues s of matrix is
s² (1 ad B)s + 3 0,
(249)
yielding solutions
$1,2 - ax + B + (1 ad 3)2 - 43
(250)
2
The GD dynamics converges iff [s1] < 1 and |22 < 1. First note that |s1||52| = |8| and therefore
for [8] > 1 at least one of |1,2 will be 2 1 and GD will not converge. Therefore only need to
consider 3 E (-)
When 3 E ((1 2 (1 V) the eigenvalues 81,2 are complex and |s1] = |82l The convergence condition in this region is |@8| < which is automatically satisfied.
Now consider the region where 81,2 real. As 51 82, the convergence conditions in this region are 81 and 82 > -1. Note that for 3 S the eigenvalues 8122 are always real, while for 3 they are real for [0, oo) (A1, A2) with X (1 /B)) /a and X2 = (1 + /B)//a.
Let us first check the condition < 1. The equation 81 (X) gives X = 0. For 3 > 0 we have sI = 0) = 1, 8 (A = A1) = /B (c) (Ai, )2) due to continuity of s (A). For < the condition (X) is always satisfied
and 81 (A = A₂) 3 < 1 and therefore 81 (A) on
because $1 (A) is decreasing with X.
Proceeding to condition $2 > -1 we find that equation s²() -1 gives ax* 2(3 + 1). Note that x > X2. For B 0 we have s2(A = 0) = > -1, s2(A = A1) = /3 -1,
82(A- = A²) -VB
-1 and s2() + so) +
Therefore, due continuity of (2))
have 82( 82(A) A) > -1 on (0, X*) />)/ X2) and (2) < - on (X*, Similarly, for B < 0 we again
have (A) -1 on (0) and (A) -1 on oc)
Combining all the observations above, we get a single convergence condition for non-stochastic GD:
3)
BE(-1,1)
(251)
Amax
where Amax is the largest eigenvalues of H.
51


Reviews:
[{"review_id": "R85D9EiXEL", "reviewer": "Reviewer_LoFN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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 work investigates stochastic gradient descent with momentum for the least squares problem at fixed design. In particular, it derives closed-form equations for the evolution of second moments of the weights and momentum parameters under a \"spectrally expressible\" (SE) approximation, which allow one to compute the trajectory of observables throughout the dynamics from the initial weight covariance matrix and covariance of the features (i.e. the Hessian). In particular, this approximation holds for translational invariant kernels in a grid and for Gaussian features, and Theorem 1 provides different equivalent characterisations of the proposed SE approximation. By rewriting the evolution and identifying it with an interacting gas of rank-one operators, the authors separate the contributions of the signal and sampling noise.  The authors study this evolution equations under the SE approximation and derive different regimes for the dynamics as a function of the learning rate, the momentum strength and the spectrum of the features covariance. In particular, it is observed that an unusual choice of small negative momentum can be desirable in the noise dominated phase.\n", "summary_of_the_review": "This work provides a characterisation of the SGD dynamics with momentum for the least squares problems with fixed design. The introduction overstates the importance of the studied setting. The manuscript is dense and hard to parse, and I feel that developing better the discussion and highlighting better the main contributions would improve the reading. For a non-expert, it is not very clear what are the challenges and the progress made the manuscript with regard to previous literature. Therefore, I believe this work could benefit from a revision. ", "strength_and_weaknesses": "**Comments and Questions**:\n\n**[Q1]**: A series of recent works have shown that the behaviour of one-pass SGD on small neural networks [Goldt et al. '20] and random features [Goldt et al. '22] as well as full-batch learning [Loureiro et al. '21] and mini-batch SGD [Bordelon, Pehlevan '21] on linear models trained on fixed feature maps are well approximated by an equivalent model with Gaussian features in the high-dimensional limit. As mentioned in the manuscript, in this case the \"Spectrally expressible\" approximation is exact. By the way, all of the aforementioned works contain experiments of Gaussian equivalence on MNIST, which is consistent with the authors observation that their SE approximatiom works well in this case. I wonder whether Gaussian equivalence is only sufficient for SE to hold or also necessary. In other words: do the authors have an example where the SE approximation holds but the phenomenology is different from having Gaussian features?\n\n**[Q2]**: The sentence in the introduction:\n> The goal of the present work is to derive explicit and easy-to-use analytic evolution laws for the SGD (1) applicable to modern deep neural networks trained on real world datasets.\n\nis a huge overstating, not only given the theoretical setting considered in this work (least-squares regression) but also the lack of empirical evidence for that (i.e. checking the theory on MNIST). Moreover, the following sentence:\n\n**[Q3]**: Moreover, in the next sentence in the introduction:\n>First, in many (though not all) cases the nonlinear network training is reasonably well approximated (Lee et al., 2020) by the linearized NTK regime (Jacot et al., 2018), in which the loss L is quadratic with respect to parameters w.\n\nWhere in (Jacot et al., 2018) it is shown that NTK = square loss?\n\n**[Q4]**: I am not very familiar with the related literature. But browsing some papers for this review I stumbled upon [Paquette, Paquette '21], who studies SGD with momentum in least squares with Gaussian random design. Their conclusion is that in the strong convex case momentum does not improve the performance of SGD, while in the non-strongly convex case it is possible to get considerable improvement. Is this consistent with taking a Marchenko-Pastur distribution for the Hessian matrix here? How does this conclusion changes as a function of the choice of feature map?\n\n**References**:\n\n[[Bordelon, Pehlevan '21]](https://arxiv.org/abs/2106.02713) B Bordelon, C Pehlevan. *Learning Curves for SGD on Structured Features*. arXiv: 2106.02713 [cs.LG]\n\n[[Loureiro et al. '21]](https://proceedings.neurips.cc/paper/2021/hash/9704a4fc48ae88598dcbdcdf57f3fdef-Abstract.html) B Loureiro, C Gerbelot, H Cui, S Goldt, F Krzakala, M Mezard, L Zdeborová. *Learning curves of generic features maps for realistic datasets with a teacher-student model*. Part of Advances in Neural Information Processing Systems 34 (NeurIPS 2021).\n\n[[Goldt et al. '20]](https://journals.aps.org/prx/abstract/10.1103/PhysRevX.10.041044) S Goldt, M Mézard, F Krzakala, L Zdeborová. *Modelling the influence of data structure on learning in neural networks: the hidden manifold model*. Phys. Rev. X 10, 041044 – Published 3 December 2020.\n\n[[Goldt et al. '22]](https://proceedings.mlr.press/v145/goldt22a.html) S Goldt, B Loureiro, G Reeves, F Krzakala, M Mezard, L Zdeborova. *The Gaussian equivalence of generative models for learning with shallow neural networks*. Proceedings of the 2nd Mathematical and Scientific Machine Learning Conference, PMLR 145:426-471, 2022.\n\n[[Paquette, Paquette '21]](https://papers.nips.cc/paper/2021/hash/4cf0ed8641cfcbbf46784e620a0316fb-Abstract.html) C Paquette, E Paquette. *Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models*.  Part of Advances in Neural Information Processing Systems 34 (NeurIPS 2021).", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "H30_SDJ1vZO", "reviewer": "Reviewer_KUvw", "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": "", "summary_of_the_paper": "This paper develops a theoretical framework to analysize the noise averged properties of  for mini-batch SGD with momentum (bSGDM) for the linear models when learning rate, mini-batch size and momentum parameter is fixed. To handle the noise terms appearing in the analysis, they approximate that variance term with specterally exprissible approximations. This approximation allows them to form a generative function for the loss function trajectory. Analysing this generating function 1- they specify the noise dominant and signal dominant phase in loss trajectory 2- when Hessian spectral distribution follows power low distribution, they analyze the stability of loss function and show when divergence and convergence happens 3- show that negative momentum sometimes can helps. They give some theoretical approximation for optimal parameter values and the phase transition times. Empirically they show that their analysis matches the practice for MNIST and some artificial data sets. \n", "summary_of_the_review": "The major problem I have with the paper is its writing. I believe it is hard to understand and read the paper easily. Here are some examples: \n\n1- there isn’t any part in the paper to explains the notation and operations used in the formula. For example there U \\in O(N) where U is a matrix. \n\n2- You claim that you drive exat loss asymptotics for bSGDM however you use a lot of approximations in your analysis. So making such a general claim is not accurate. \n\n3- You have some terms which are not defined explicitly. For example, \\lambda_crit or regularization parameter is not defined anywhere. \n\n4-  The idea behind approximating eq 5 by eq 6 isn’t explained well. The paper says this is due to the existence of non-spectral details in eq 5 but this is not enough. It would be helpful to show these details and show how eq 7 would solve it. \n\n5- In eq 8 which is an approximate for eq 5, there are cross terms between feature maps of two different samples that don’t exist in eq 5 and it can remove the stochasticity from the dynamics. So the question is how good is this approximation i.e. the upper bound of the difference of these two eqs. \n\n6- For your approximation of L(t) in “Transitions between phases”, it adds eq 27a and eq 27b. However, depending on \\zetta and \\nu one of these equations can happen for L(t). So why the sum of them is also a good approx for L(t)? \n\n7- One of the claimed contributions is that “the optimal convergence rate can be achieved at negative momenta”. However, the analysis in the “negative momenta” section is hand-wavy and for a very limiting case. \n", "strength_and_weaknesses": "The major strength of the paper is that its analytical framework matches well enough with the empirical observation. \n\nThe major weakness of the paper is that the provided framework is applicable for too limited cases due to their assumptions and approximations. Besides the paper is not well-written and easy to follow.\n", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "4rQLwpJiq1", "reviewer": "Reviewer_NWpJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "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 work builds a new convergence analysis framework for SGD algorithm (with momentum). It proposes a special family functions of ”Spectrally Expressible” approximations, which provides a new perspective to understand the behavior of classical SGD. A specific senario where a negative momenta can be the optimal choice is constructed.", "summary_of_the_review": "I recommend to accept this paper as its results are impressive and really new to the community.", "strength_and_weaknesses": "***Strenght***\nThis paper applies the generating function to analyze the higher moments of SGD noises. The new tool reveals the explicit stablility condition, which answer when SGD can converge and diverge. Most intersting part is that this paper introduces a situation where the negative momenta is better.  Also, the theoretical result perfectly explains the dynamic behavior of SGD over MNIST dataset.\n\n***Weaknesses***\nThe loss function is relatively special, though I agree that this loss is widely used in modern machine learning tasks. \n", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster